From af8d8413ea029c35dcdba301a67ebbc2507c9552 Mon Sep 17 00:00:00 2001 From: Chris Feijoo Date: Mon, 3 Aug 2026 15:47:04 +0200 Subject: [PATCH] FE-1222, FE-1226: Classify optimization transport errors and auto-reconnect detached runs - Classify Petrinaut optimization transport failures in the UI so the user sees whether a run was rejected, dropped, or is retrying, instead of one opaque error. - Auto-reconnect detached optimization runs by run id and event `seq` cursor, re-attaching after connection drops and (where storage allows) page reloads; report an honest connection state from the new `onAttached` signal. - Make the detached contract the only one: the host capability is now `createOptimizationRun`/`attachOptimizationRun`/`cancelOptimizationRun`, and the legacy single-connection `optimize` method is removed from petrinaut-core. - Remove the legacy optimization path end to end: the optimizer's `/optimize/all` and `/optimize/best` routes and their `stream_all`/`stream_best` generators, the client's `openPetrinautOptimizationStream`, and the provider's legacy fork. - Migrate the website optimization demo and the SimulateView stories to the detached-run API, and update the Petrinaut optimization user docs. --- .changeset/eight-streams-return.md | 6 + .../processes/[uuid].page/process-editor.tsx | 316 +++++-- ...dge-petrinaut-optimization.browser.test.ts | 248 ++++- ...eate-bridge-petrinaut-optimization.test.ts | 94 +- .../create-bridge-petrinaut-optimization.ts | 344 ++++++- .../pages/processes/shared/messages.test.ts | 11 + .../src/pages/processes/shared/messages.ts | 82 +- .../pages/processes/shared/use-host-bridge.ts | 20 +- apps/petrinaut-opt/README.md | 9 +- apps/petrinaut-opt/openapi/openapi.json | 116 --- apps/petrinaut-opt/src/optimization_api.py | 87 +- apps/petrinaut-opt/src/petrinaut_optimizer.py | 288 +----- .../tests/test_optimization_api.py | 257 +---- .../tests/test_petrinaut_optimizer.py | 310 +----- .../petrinaut-opt-optimization.test.ts | 90 +- .../petrinaut-opt-optimization.ts | 130 ++- libs/@hashintel/petrinaut-core/src/index.ts | 1 + .../petrinaut-core/src/optimization.test.ts | 60 +- .../petrinaut-core/src/optimization.ts | 57 +- .../@hashintel/petrinaut/docs/optimization.md | 23 +- libs/@hashintel/petrinaut/src/react/index.ts | 1 + .../src/react/optimizations/context.ts | 42 + .../src/react/optimizations/provider.test.tsx | 779 +++++++++++++++- .../src/react/optimizations/provider.tsx | 879 ++++++++++++++++-- .../create-optimization-drawer.test.tsx | 1 + .../view-optimization-drawer.tsx | 25 +- .../SimulateView/simulate-view.stories.tsx | 57 +- .../SimulateView/simulate-view.test.tsx | 6 +- .../src/attach-optimization-run.test.ts | 31 - .../src/attach-optimization-run.ts | 42 +- .../src/decode-optimization-stream.test.ts | 102 +- .../src/decode-optimization-stream.ts | 71 +- .../petrinaut-optimizer-client/src/index.ts | 10 +- .../src/open-optimization-stream.test.ts | 208 ----- .../src/open-optimization-stream.ts | 78 -- .../src/optimizer-http.ts | 16 + 36 files changed, 3129 insertions(+), 1768 deletions(-) create mode 100644 .changeset/eight-streams-return.md delete mode 100644 libs/@local/petrinaut-optimizer-client/src/open-optimization-stream.test.ts delete mode 100644 libs/@local/petrinaut-optimizer-client/src/open-optimization-stream.ts diff --git a/.changeset/eight-streams-return.md b/.changeset/eight-streams-return.md new file mode 100644 index 00000000000..65ff4fce567 --- /dev/null +++ b/.changeset/eight-streams-return.md @@ -0,0 +1,6 @@ +--- +"@hashintel/petrinaut-core": patch +"@hashintel/petrinaut": patch +--- + +Make optimization runs detached and resumable — and make that the only contract. Optimization events carry a server-issued `seq`; the host optimization capability is now `createOptimizationRun`/`attachOptimizationRun`/`cancelOptimizationRun` (all required; attachments accept an `onAttached` callback so UIs can report an honest connection state), and the legacy single-connection `optimize` method is removed. The optimizations UI auto-reconnects dropped event streams by run id and cursor, re-attaching after page reloads where storage allows. diff --git a/apps/hash-frontend/src/pages/processes/[uuid].page/process-editor.tsx b/apps/hash-frontend/src/pages/processes/[uuid].page/process-editor.tsx index e2bb16c1336..d23bbd3128b 100644 --- a/apps/hash-frontend/src/pages/processes/[uuid].page/process-editor.tsx +++ b/apps/hash-frontend/src/pages/processes/[uuid].page/process-editor.tsx @@ -11,6 +11,7 @@ import { apiOrigin } from "@local/hash-isomorphic-utils/environment"; import { type HostNetMode, + type HostToIframeMessage, type PetrinautAiMessage, type PetrinautHostCapabilities, type RevisionSummary, @@ -57,12 +58,162 @@ const PETRINAUT_EMBED_SRC = "/processes/draft/embed"; */ const PETRINAUT_AI_CHAT_API = "/api/petrinaut-ai-chat"; -/** Authenticated NodeAPI endpoint that proxies the optimizer container. */ -const PETRINAUT_OPTIMIZATION_API = `${apiOrigin}/api/petrinaut-optimizer/optimize`; +/** + * Authenticated NodeAPI endpoint for detached optimization runs, proxying + * the optimizer container. `POST` creates a run and replies + * `201 {"runId": ...}` without streaming. + */ +const PETRINAUT_OPTIMIZATION_RUNS_API = `${apiOrigin}/api/petrinaut-optimizer/optimize/runs`; + +/** + * NodeAPI validates the replay cursor against `^\d{1,15}$` (and 400s + * otherwise), so anything malformed or out of that range falls back to 0 + * (replay everything) rather than letting iframe-supplied data shape the URL. + */ +const MAX_OPTIMIZATION_CURSOR = 999_999_999_999_999; + +/** + * Cursor-resumable NDJSON event stream of a detached optimization run. + * Replays events with `seq` greater than `cursor`, then tails live events. + */ +const petrinautOptimizationEventsUrl = ( + runId: string, + cursor: number, +): string => + `${PETRINAUT_OPTIMIZATION_RUNS_API}/${encodeURIComponent(runId)}/events?cursor=${ + Number.isSafeInteger(cursor) && + cursor >= 0 && + cursor <= MAX_OPTIMIZATION_CURSOR + ? cursor + : 0 + }`; + +/** A single detached optimization run; `DELETE` cancels it idempotently. */ +const petrinautOptimizationRunUrl = (runId: string): string => + `${PETRINAUT_OPTIMIZATION_RUNS_API}/${encodeURIComponent(runId)}`; + +/** + * Fire-and-forget cancellation. Idempotent server-side, and every caller has + * already given up on the run, so a failure is only logged for debugging. + */ +const cancelOptimizationRun = (runId: string): void => { + void fetch(petrinautOptimizationRunUrl(runId), { + method: "DELETE", + credentials: "include", + }).catch((error: unknown) => { + // eslint-disable-next-line no-console + console.error("Failed to cancel Petrinaut optimization run", error); + }); +}; /** Authenticated NodeAPI endpoint reporting deployment configuration only. */ const PETRINAUT_CAPABILITIES_API = `${apiOrigin}/api/petrinaut-optimizer/capabilities`; +/** Parse a `Retry-After` header's delay-seconds form; `undefined` otherwise. */ +const parseRetryAfterSeconds = (header: string | null): number | undefined => { + if (header === null) { + return undefined; + } + const seconds = Number.parseInt(header, 10); + return Number.isFinite(seconds) && seconds >= 0 ? seconds : undefined; +}; + +/** + * Extract NodeAPI's own error message from a failed optimizer response body. + * NodeAPI error bodies are server-authored (`{"error": ...}`), never user + * content, so they are safe to forward into the iframe; anything unparseable + * falls back to a generic status message. + */ +const readOptimizationErrorMessage = async ( + response: Response, +): Promise => { + try { + const body = (await response.json()) as { + error?: unknown; + message?: unknown; + }; + if (typeof body.error === "string" && body.error) { + return body.error; + } + if (typeof body.message === "string" && body.message) { + return body.message; + } + } catch { + // Non-JSON body; fall through to the generic message. + } + return `The optimization request failed with status ${response.status}`; +}; + +/** + * Relay an optimizer NDJSON response into the iframe byte-for-byte over the + * response-start/chunk/end/error message family, keyed by the iframe's + * request id. + * + * HTTP and protocol classification happen in the iframe bridge; the catch + * here only sees transport-level failures of the host's own fetch/read (a + * reset connection surfaces as a `TypeError`), which it classifies as + * `network`. + */ +const relayOptimizationStream = async ({ + requestId, + controller, + send, + fetchResponse, +}: { + requestId: string; + controller: AbortController; + send: (message: HostToIframeMessage) => void; + fetchResponse: () => Promise; +}): Promise => { + // Captured from the response headers so a transport failure that happens + // mid-stream stays traceable to the NodeAPI/optimizer logs. + let hashRequestId: string | null = null; + let optimizationRunId: string | null = null; + try { + const response = await fetchResponse(); + + hashRequestId = response.headers.get("x-hash-request-id"); + optimizationRunId = response.headers.get("x-optimization-run-id"); + + send({ + kind: "optimizationResponseStart", + requestId, + ok: response.ok, + status: response.status, + statusText: response.statusText, + hashRequestId, + optimizationRunId, + }); + + if (response.body) { + const reader = response.body.getReader(); + let result = await reader.read(); + while (!result.done) { + send({ + kind: "optimizationChunk", + requestId, + bytes: result.value, + }); + result = await reader.read(); + } + } + + send({ kind: "optimizationEnd", requestId }); + } catch { + // An iframe-initiated abort is expected control flow. + if (!controller.signal.aborted) { + send({ + kind: "optimizationError", + requestId, + category: "network", + message: "The optimization service connection was interrupted", + hashRequestId, + optimizationRunId, + }); + } + } +}; + /** * URL-derived view that the editor renders. The host page resolves this from * `router.query` and passes it in; the editor reconciles its internal state @@ -494,89 +645,150 @@ export const ProcessEditor = ({ controller?.abort(); aiChatAbortControllersRef.current.delete(requestId); }, - onOptimizationRequest: ({ requestId, input }) => { + onOptimizationCreate: ({ requestId, input }) => { const parsedInput = petrinautOptimizationInputSchema.safeParse(input); if (!parsedInput.success) { bridge.send({ - kind: "optimizationError", + kind: "optimizationCreateResult", requestId, + ok: false, + category: "protocol", message: "The optimization request is invalid", }); return; } - if (optimizationAbortControllersRef.current.has(requestId)) { - bridge.send({ - kind: "optimizationError", - requestId, - message: "An optimization with this request id is already running", - }); - return; - } - - const controller = new AbortController(); - optimizationAbortControllersRef.current.set(requestId, controller); /** - * The sandboxed iframe has no credentials or network access. Validate - * its structured-cloned request, call the one hard-coded NodeAPI route - * with HASH's session, then relay the NDJSON response byte-for-byte. + * Create a detached run against the hard-coded NodeAPI runs route. A + * short-lived POST (the `201` reply carries only the run id, no + * stream), so there is no abort tracking — the iframe times the + * round-trip out on its side. */ void (async () => { + let response: Response; try { - const response = await fetch(PETRINAUT_OPTIMIZATION_API, { + response = await fetch(PETRINAUT_OPTIMIZATION_RUNS_API, { method: "POST", credentials: "include", headers: { "content-type": "application/json" }, body: JSON.stringify(parsedInput.data), - signal: controller.signal, }); + } catch { + bridge.send({ + kind: "optimizationCreateResult", + requestId, + ok: false, + category: "network", + message: "The optimization service could not be reached", + }); + return; + } + /** + * NodeAPI echoes its request id, and forwards the optimizer's run id + * even when creation fails upstream. Carry both back to the iframe so + * a failure the user reports can be found in the logs. + */ + const correlation = { + hashRequestId: response.headers.get("x-hash-request-id"), + optimizationRunId: response.headers.get("x-optimization-run-id"), + }; + + if (!response.ok) { bridge.send({ - kind: "optimizationResponseStart", + kind: "optimizationCreateResult", requestId, - ok: response.ok, + ok: false, + category: "http", status: response.status, - statusText: response.statusText, + retryAfter: parseRetryAfterSeconds( + response.headers.get("retry-after"), + ), + message: await readOptimizationErrorMessage(response), + ...correlation, }); + return; + } - if (response.body) { - const reader = response.body.getReader(); - let result = await reader.read(); - while (!result.done) { - bridge.send({ - kind: "optimizationChunk", - requestId, - bytes: result.value, - }); - result = await reader.read(); - } - } - - bridge.send({ kind: "optimizationEnd", requestId }); - } catch (error) { - // An iframe-initiated abort is expected control flow. - if (!controller.signal.aborted) { - bridge.send({ - kind: "optimizationError", - requestId, - message: error instanceof Error ? error.message : String(error), - }); - } - } finally { - if ( - optimizationAbortControllersRef.current.get(requestId) === - controller - ) { - optimizationAbortControllersRef.current.delete(requestId); + const body = (await response.json().catch(() => null)) as { + runId?: unknown; + } | null; + if (typeof body?.runId !== "string" || body.runId === "") { + /** + * The run was created — the status said so — but its id never + * reached us in the body. When the correlation header carries it, + * cancel the run nobody can now own; otherwise it holds the + * account's single-flight slot until the reaper takes it. + */ + if (correlation.optimizationRunId !== null) { + cancelOptimizationRun(correlation.optimizationRunId); } + bridge.send({ + kind: "optimizationCreateResult", + requestId, + ok: false, + category: "protocol", + message: + "The optimization service returned an unexpected response", + ...correlation, + }); + return; } + + bridge.send({ + kind: "optimizationCreateResult", + requestId, + ok: true, + runId: body.runId, + ...correlation, + }); })(); }, + onOptimizationAttach: ({ requestId, runId, cursor }) => { + if (optimizationAbortControllersRef.current.has(requestId)) { + bridge.send({ + kind: "optimizationError", + requestId, + category: "protocol", + message: "An optimization with this request id is already running", + }); + return; + } + + const controller = new AbortController(); + optimizationAbortControllersRef.current.set(requestId, controller); + + /** + * The iframe only names a run id and cursor; the URL is built here + * against the one hard-coded NodeAPI route, so the sandboxed iframe + * can never make the host fetch an arbitrary target. + */ + void relayOptimizationStream({ + requestId, + controller, + send: bridge.send, + fetchResponse: () => + fetch(petrinautOptimizationEventsUrl(runId, cursor), { + credentials: "include", + signal: controller.signal, + }), + }).finally(() => { + if ( + optimizationAbortControllersRef.current.get(requestId) === + controller + ) { + optimizationAbortControllersRef.current.delete(requestId); + } + }); + }, onOptimizationAbort: ({ requestId }) => { const controller = optimizationAbortControllersRef.current.get(requestId); controller?.abort(); }, + onOptimizationCancel: ({ runId }) => { + cancelOptimizationRun(runId); + }, onAiMessagesChanged: ({ messages }) => { if (!loadedView) { return; diff --git a/apps/hash-frontend/src/pages/processes/[uuid]/embed.page/create-bridge-petrinaut-optimization.browser.test.ts b/apps/hash-frontend/src/pages/processes/[uuid]/embed.page/create-bridge-petrinaut-optimization.browser.test.ts index a9f37e455eb..93668364682 100644 --- a/apps/hash-frontend/src/pages/processes/[uuid]/embed.page/create-bridge-petrinaut-optimization.browser.test.ts +++ b/apps/hash-frontend/src/pages/processes/[uuid]/embed.page/create-bridge-petrinaut-optimization.browser.test.ts @@ -53,13 +53,25 @@ const input = petrinautOptimizationInputSchema.parse({ study: { trials: 1, sampler: "tpe" }, }); -const getOptimizationRequest = (calls: readonly (readonly unknown[])[]) => +const getPostedMessage = ( + calls: readonly (readonly unknown[])[], + kind: string, +) => calls .map( ([message]) => - message as { kind?: string; requestId?: string; input?: unknown }, + message as { + kind?: string; + requestId?: string; + input?: unknown; + runId?: string; + cursor?: number; + }, ) - .find(({ kind }) => kind === "optimizationRequest"); + .find((message) => message.kind === kind); + +const getOptimizationAttach = (calls: readonly (readonly unknown[])[]) => + getPostedMessage(calls, "optimizationAttach"); const sendFromHost = (data: unknown) => { window.dispatchEvent( @@ -77,17 +89,17 @@ describe("createBridgePetrinautOptimization", () => { .spyOn(window.parent, "postMessage") .mockImplementation(() => undefined); const iterator = createBridgePetrinautOptimization() - .optimize(input) + .attachOptimizationRun("run-1") [Symbol.asyncIterator](); const firstEvent = iterator.next(); await vi.waitFor(() => { - expect(getOptimizationRequest(postMessage.mock.calls)).toBeDefined(); + expect(getOptimizationAttach(postMessage.mock.calls)).toBeDefined(); }); - const optimizationRequest = getOptimizationRequest(postMessage.mock.calls); - const requestId = optimizationRequest?.requestId; + const attachMessage = getOptimizationAttach(postMessage.mock.calls); + const requestId = attachMessage?.requestId; expect(requestId).toBeDefined(); - expect(optimizationRequest?.input).toEqual(input); + expect(attachMessage).toMatchObject({ runId: "run-1", cursor: 0 }); sendFromHost({ kind: "optimizationResponseStart", @@ -95,6 +107,8 @@ describe("createBridgePetrinautOptimization", () => { ok: true, status: 200, statusText: "OK", + hashRequestId: "req-1", + optimizationRunId: "run-1", }); sendFromHost({ kind: "optimizationChunk", @@ -126,14 +140,14 @@ describe("createBridgePetrinautOptimization", () => { .mockImplementation(() => undefined); const abortController = new AbortController(); const iterator = createBridgePetrinautOptimization() - .optimize(input, { signal: abortController.signal }) + .attachOptimizationRun("run-2", { signal: abortController.signal }) [Symbol.asyncIterator](); const firstEvent = iterator.next(); await vi.waitFor(() => { - expect(getOptimizationRequest(postMessage.mock.calls)).toBeDefined(); + expect(getOptimizationAttach(postMessage.mock.calls)).toBeDefined(); }); - const requestId = getOptimizationRequest(postMessage.mock.calls)?.requestId; + const requestId = getOptimizationAttach(postMessage.mock.calls)?.requestId; abortController.abort(); await expect(firstEvent).rejects.toMatchObject({ name: "AbortError" }); @@ -146,4 +160,216 @@ describe("createBridgePetrinautOptimization", () => { ), ).toBe(true); }); + + it("classifies a mid-stream host error with its correlation ids", async () => { + const postMessage = vi + .spyOn(window.parent, "postMessage") + .mockImplementation(() => undefined); + const iterator = createBridgePetrinautOptimization() + .attachOptimizationRun("run-9") + [Symbol.asyncIterator](); + + const firstEvent = iterator.next(); + await vi.waitFor(() => { + expect(getOptimizationAttach(postMessage.mock.calls)).toBeDefined(); + }); + const requestId = getOptimizationAttach(postMessage.mock.calls)?.requestId; + + sendFromHost({ + kind: "optimizationResponseStart", + requestId, + ok: true, + status: 200, + statusText: "OK", + hashRequestId: "req-9", + optimizationRunId: "run-9", + }); + sendFromHost({ + kind: "optimizationError", + requestId, + category: "network", + message: "The optimization service connection was interrupted", + }); + + // The bridge reclassifies the host error and backfills the correlation ids + // captured from the earlier response-start message. + await expect(firstEvent).rejects.toMatchObject({ + name: "PetrinautOptimizationTransportError", + category: "network", + hashRequestId: "req-9", + optimizationRunId: "run-9", + }); + }); + + it("resolves a created detached run's id", async () => { + const postMessage = vi + .spyOn(window.parent, "postMessage") + .mockImplementation(() => undefined); + const created = + createBridgePetrinautOptimization().createOptimizationRun(input); + + await vi.waitFor(() => { + expect( + getPostedMessage(postMessage.mock.calls, "optimizationCreate"), + ).toBeDefined(); + }); + const createMessage = getPostedMessage( + postMessage.mock.calls, + "optimizationCreate", + ); + expect(createMessage?.input).toEqual(input); + + sendFromHost({ + kind: "optimizationCreateResult", + requestId: createMessage?.requestId, + ok: true, + runId: "run-42", + }); + + await expect(created).resolves.toEqual({ runId: "run-42" }); + }); + + it("rejects a failed creation with its classification, status, and retry delay", async () => { + const postMessage = vi + .spyOn(window.parent, "postMessage") + .mockImplementation(() => undefined); + const created = + createBridgePetrinautOptimization().createOptimizationRun(input); + + await vi.waitFor(() => { + expect( + getPostedMessage(postMessage.mock.calls, "optimizationCreate"), + ).toBeDefined(); + }); + sendFromHost({ + kind: "optimizationCreateResult", + requestId: getPostedMessage(postMessage.mock.calls, "optimizationCreate") + ?.requestId, + ok: false, + category: "http", + status: 429, + retryAfter: 30, + message: "Too many concurrent optimizations", + }); + + await expect(created).rejects.toMatchObject({ + name: "PetrinautOptimizationTransportError", + category: "http", + httpStatus: 429, + retryAfter: 30, + message: "Too many concurrent optimizations", + }); + }); + + it("cancels a run created after the local create was aborted", async () => { + const postMessage = vi + .spyOn(window.parent, "postMessage") + .mockImplementation(() => undefined); + const abortController = new AbortController(); + const created = createBridgePetrinautOptimization().createOptimizationRun( + input, + { signal: abortController.signal }, + ); + + await vi.waitFor(() => { + expect( + getPostedMessage(postMessage.mock.calls, "optimizationCreate"), + ).toBeDefined(); + }); + abortController.abort(); + await expect(created).rejects.toMatchObject({ name: "AbortError" }); + + // The host's reply arrives after the local promise already settled: the + // run exists server-side but nobody will ever own it, so the bridge asks + // the host to cancel it. + sendFromHost({ + kind: "optimizationCreateResult", + requestId: getPostedMessage(postMessage.mock.calls, "optimizationCreate") + ?.requestId, + ok: true, + runId: "run-late", + }); + + await vi.waitFor(() => { + expect( + getPostedMessage(postMessage.mock.calls, "optimizationCancel"), + ).toMatchObject({ runId: "run-late" }); + }); + }); + + it("attaches to a run with the resume cursor and relays replayed and live events", async () => { + const postMessage = vi + .spyOn(window.parent, "postMessage") + .mockImplementation(() => undefined); + const onAttached = vi.fn(); + const iterator = createBridgePetrinautOptimization() + .attachOptimizationRun("run-7", { cursor: 2, onAttached }) + [Symbol.asyncIterator](); + + const firstEvent = iterator.next(); + await vi.waitFor(() => { + expect( + getPostedMessage(postMessage.mock.calls, "optimizationAttach"), + ).toBeDefined(); + }); + const attachMessage = getPostedMessage( + postMessage.mock.calls, + "optimizationAttach", + ); + expect(attachMessage).toMatchObject({ runId: "run-7", cursor: 2 }); + + const requestId = attachMessage?.requestId; + sendFromHost({ + kind: "optimizationResponseStart", + requestId, + ok: true, + status: 200, + statusText: "OK", + hashRequestId: "req-7", + optimizationRunId: "run-7", + }); + // A replayed event (seq 3) followed by a live terminal event (seq 4). + sendFromHost({ + kind: "optimizationChunk", + requestId, + bytes: new TextEncoder().encode( + '{"type":"trial","trial":2,"parameters":{"rate":0.4},"objective":1,"state":"complete","best":null,"seq":3}\n', + ), + }); + sendFromHost({ + kind: "optimizationChunk", + requestId, + bytes: new TextEncoder().encode( + '{"type":"complete","requestedTrials":3,"completedTrials":3,"prunedTrials":0,"failedTrials":0,"best":null,"seq":4}\n', + ), + }); + sendFromHost({ kind: "optimizationEnd", requestId }); + + await expect(firstEvent).resolves.toMatchObject({ + done: false, + value: { type: "trial", trial: 2, seq: 3 }, + }); + await expect(iterator.next()).resolves.toMatchObject({ + done: false, + value: { type: "complete", completedTrials: 3, seq: 4 }, + }); + await expect(iterator.next()).resolves.toEqual({ + done: true, + value: undefined, + }); + // The accepted response fired the "connected" signal exactly once. + expect(onAttached).toHaveBeenCalledTimes(1); + }); + + it("posts a fire-and-forget cancel for a detached run", async () => { + const postMessage = vi + .spyOn(window.parent, "postMessage") + .mockImplementation(() => undefined); + + await createBridgePetrinautOptimization().cancelOptimizationRun("run-9"); + + expect( + getPostedMessage(postMessage.mock.calls, "optimizationCancel"), + ).toMatchObject({ runId: "run-9" }); + }); }); diff --git a/apps/hash-frontend/src/pages/processes/[uuid]/embed.page/create-bridge-petrinaut-optimization.test.ts b/apps/hash-frontend/src/pages/processes/[uuid]/embed.page/create-bridge-petrinaut-optimization.test.ts index cc9e75122f0..f579f30890e 100644 --- a/apps/hash-frontend/src/pages/processes/[uuid]/embed.page/create-bridge-petrinaut-optimization.test.ts +++ b/apps/hash-frontend/src/pages/processes/[uuid]/embed.page/create-bridge-petrinaut-optimization.test.ts @@ -1,8 +1,14 @@ import { describe, expect, it } from "vitest"; -import { parsePetrinautOptimizationResponse } from "./create-bridge-petrinaut-optimization"; +import { + parsePetrinautOptimizationResponse, + PetrinautOptimizationTransportError, +} from "./create-bridge-petrinaut-optimization"; -const responseFromChunks = (chunks: string[], status = 200): Response => { +const responseFromChunks = ( + chunks: string[], + { status = 200, headers }: { status?: number; headers?: HeadersInit } = {}, +): Response => { const encoder = new TextEncoder(); return new Response( new ReadableStream({ @@ -13,7 +19,7 @@ const responseFromChunks = (chunks: string[], status = 200): Response => { controller.close(); }, }), - { status }, + { status, ...(headers ? { headers } : {}) }, ); }; @@ -41,28 +47,98 @@ describe("parsePetrinautOptimizationResponse", () => { ]); }); - it("rejects a stream with no terminal event", async () => { + it("passes server-issued sequence numbers through decoded events", async () => { + const response = responseFromChunks([ + '{"type":"started","requestedTrials":1,"seq":1}\n', + '{"type":"complete","requestedTrials":1,"completedTrials":1,"prunedTrials":0,"failedTrials":0,"best":null,"seq":2}\n', + ]); + + await expect(collect(response)).resolves.toMatchObject([ + { type: "started", seq: 1 }, + { type: "complete", seq: 2 }, + ]); + }); + + it("rejects a stream with no terminal event as a protocol error", async () => { const response = responseFromChunks([ '{"type":"started","requestedTrials":2}\n', ]); - await expect(collect(response)).rejects.toThrow("without a terminal event"); + await expect(collect(response)).rejects.toMatchObject({ + category: "protocol", + message: expect.stringContaining("without a terminal event"), + }); }); - it("rejects data after a terminal event", async () => { + it("rejects data after a terminal event as a protocol error", async () => { const response = responseFromChunks([ '{"type":"error","code":"failed","message":"nope","retryable":false}\n', '{"type":"started","requestedTrials":2}\n', ]); - await expect(collect(response)).rejects.toThrow("after a terminal event"); + await expect(collect(response)).rejects.toMatchObject({ + category: "protocol", + message: expect.stringContaining("after a terminal event"), + }); + }); + + it("rejects a schema-invalid event as a protocol error, not a raw ZodError", async () => { + // Syntactically valid JSON in the wrong shape. It must take the same + // classified path as malformed NDJSON, so the provider reconnects and the + // schema's validation detail never reaches the user. + const response = responseFromChunks(['{"type":"trial","trial":"first"}\n']); + + const error = await collect(response).then( + () => null, + (caught: unknown) => caught, + ); + + expect(error).toBeInstanceOf(PetrinautOptimizationTransportError); + expect(error).toMatchObject({ + category: "protocol", + message: "The optimizer returned an unrecognized event", + }); }); - it("surfaces a structured HTTP error", async () => { + it("carries the response's correlation ids on a mid-stream protocol error", async () => { + const response = responseFromChunks( + ['{"type":"started","requestedTrials":2}\n'], + { + headers: { + "x-hash-request-id": "req-2", + "x-optimization-run-id": "run-2", + }, + }, + ); + + await expect(collect(response)).rejects.toMatchObject({ + category: "protocol", + hashRequestId: "req-2", + optimizationRunId: "run-2", + }); + }); + + it("surfaces a structured HTTP error with status and correlation ids", async () => { const response = new Response(JSON.stringify({ error: "Not configured" }), { status: 503, + headers: { + "x-hash-request-id": "req-1", + "x-optimization-run-id": "run-1", + }, }); - await expect(collect(response)).rejects.toThrow("Not configured"); + const error = await collect(response).then( + () => null, + (caught: unknown) => caught, + ); + + expect(error).toBeInstanceOf(PetrinautOptimizationTransportError); + expect(error).toMatchObject({ + category: "http", + httpStatus: 503, + hashRequestId: "req-1", + optimizationRunId: "run-1", + message: "Not configured", + }); }); }); diff --git a/apps/hash-frontend/src/pages/processes/[uuid]/embed.page/create-bridge-petrinaut-optimization.ts b/apps/hash-frontend/src/pages/processes/[uuid]/embed.page/create-bridge-petrinaut-optimization.ts index 8ff5f77c509..ad823f809d8 100644 --- a/apps/hash-frontend/src/pages/processes/[uuid]/embed.page/create-bridge-petrinaut-optimization.ts +++ b/apps/hash-frontend/src/pages/processes/[uuid]/embed.page/create-bridge-petrinaut-optimization.ts @@ -11,23 +11,93 @@ import { nextRequestId, } from "../../shared/messages"; +/** Classification of an optimization transport failure. */ +export type OptimizationErrorCategory = + | "network" + | "http" + | "protocol" + | "aborted"; + +/** + * A classified optimization transport failure carrying the correlation ids + * needed to trace it to the NodeAPI and optimizer logs. Consumers build a + * user-facing message from `category` and progress rather than surfacing the + * raw `message`. + */ +export class PetrinautOptimizationTransportError extends Error { + readonly category: OptimizationErrorCategory; + readonly hashRequestId: string | null; + readonly optimizationRunId: string | null; + readonly httpStatus: number | null; + /** Seconds from a `Retry-After` header, when the service sent one (429). */ + readonly retryAfter: number | null; + + constructor( + message: string, + options: { + category: OptimizationErrorCategory; + hashRequestId?: string | null; + optimizationRunId?: string | null; + httpStatus?: number | null; + retryAfter?: number | null; + }, + ) { + super(message); + this.name = "PetrinautOptimizationTransportError"; + this.category = options.category; + this.hashRequestId = options.hashRequestId ?? null; + this.optimizationRunId = options.optimizationRunId ?? null; + this.httpStatus = options.httpStatus ?? null; + this.retryAfter = options.retryAfter ?? null; + } +} + type PendingRequest = { stream: ReadableStream; controller: ReadableStreamDefaultController; resolveResponse: (response: Response) => void; rejectResponse: (error: Error) => void; responded: boolean; + /** Correlation ids from the response-start header, for later stream errors. */ + hashRequestId: string | null; + optimizationRunId: string | null; clearResponseStartTimeout: () => void; cleanup: () => void; }; type OptimizationSignal = NonNullable< - Parameters[1] + Parameters[1] >["signal"]; +/** A `createOptimizationRun` round-trip awaiting its `optimizationCreateResult`. */ +type PendingCreate = { + resolve: (result: { runId: string }) => void; + reject: (error: Error) => void; + cleanup: () => void; +}; + const pendingRequests = new Map(); +const pendingCreates = new Map(); const RESPONSE_START_TIMEOUT_MS = 45_000; +/** + * Outcomes of already-settled create round-trips, kept briefly so a late + * `optimizationCreateResult` can be told apart from an unknown one. A late + * success for a create that was locally aborted or timed out names a live + * run nobody will ever own — it must be cancelled — while a duplicated + * success for an accepted create must NOT cancel the live run. + */ +const settledCreates = new Map(); + +const rememberSettledCreate = ( + requestId: string, + outcome: "accepted" | "rejected", +) => { + settledCreates.set(requestId, outcome); + // Expire the tombstone once a late reply can no longer be expected. + setTimeout(() => settledCreates.delete(requestId), RESPONSE_START_TIMEOUT_MS); +}; + const postToHost = (message: IframeToHostMessage) => { // The sandboxed iframe has an opaque origin. This still targets only its // parent window; the host independently verifies `event.source`. @@ -57,6 +127,17 @@ const rejectPendingRequest = (requestId: string, error: Error) => { pendingRequests.delete(requestId); }; +const rejectPendingCreate = (requestId: string, error: Error) => { + const pending = pendingCreates.get(requestId); + if (!pending) { + return; + } + pending.cleanup(); + pendingCreates.delete(requestId); + rememberSettledCreate(requestId, "rejected"); + pending.reject(error); +}; + let listenerInstalled = false; const ensureListener = () => { @@ -83,6 +164,44 @@ const ensureListener = () => { } const message = data as HostToIframeMessage; + if (message.kind === "optimizationCreateResult") { + const pendingCreate = pendingCreates.get(message.requestId); + if (!pendingCreate) { + // A late reply for a create that already timed out or aborted + // locally (or an unknown request id): a successful one names a live + // run nobody will ever own, so ask the host to cancel it. + if ( + message.ok && + typeof message.runId === "string" && + settledCreates.get(message.requestId) !== "accepted" + ) { + postToHost({ kind: "optimizationCancel", runId: message.runId }); + } + return; + } + pendingCreate.cleanup(); + pendingCreates.delete(message.requestId); + if (message.ok && typeof message.runId === "string") { + rememberSettledCreate(message.requestId, "accepted"); + pendingCreate.resolve({ runId: message.runId }); + } else { + rememberSettledCreate(message.requestId, "rejected"); + pendingCreate.reject( + new PetrinautOptimizationTransportError( + message.message ?? "The optimization request failed", + { + category: message.category ?? "http", + httpStatus: message.status, + retryAfter: message.retryAfter, + hashRequestId: message.hashRequestId, + optimizationRunId: message.optimizationRunId, + }, + ), + ); + } + return; + } + if ( message.kind !== "optimizationResponseStart" && message.kind !== "optimizationChunk" && @@ -102,15 +221,31 @@ const ensureListener = () => { if (pending.responded) { rejectPendingRequest( message.requestId, - new Error("The optimizer sent more than one response header"), + new PetrinautOptimizationTransportError( + "The optimizer sent more than one response header", + { category: "protocol" }, + ), ); return; } pending.responded = true; + pending.hashRequestId = message.hashRequestId; + pending.optimizationRunId = message.optimizationRunId; pending.clearResponseStartTimeout(); + const headers = new Headers({ + "content-type": "application/x-ndjson", + }); + // Carry the correlation ids on the synthesized response so the HTTP + // error path (`readHttpError`) can attach them too. + if (message.hashRequestId !== null) { + headers.set("x-hash-request-id", message.hashRequestId); + } + if (message.optimizationRunId !== null) { + headers.set("x-optimization-run-id", message.optimizationRunId); + } pending.resolveResponse( new Response(pending.stream, { - headers: { "content-type": "application/x-ndjson" }, + headers, status: message.status, statusText: message.statusText, }), @@ -129,7 +264,10 @@ const ensureListener = () => { if (!pending.responded) { rejectPendingRequest( message.requestId, - new Error("The optimizer ended before sending a response"), + new PetrinautOptimizationTransportError( + "The optimizer ended before sending a response", + { category: "protocol" }, + ), ); return; } @@ -143,19 +281,34 @@ const ensureListener = () => { break; } case "optimizationError": - rejectPendingRequest(message.requestId, new Error(message.message)); + rejectPendingRequest( + message.requestId, + new PetrinautOptimizationTransportError(message.message, { + category: message.category, + hashRequestId: message.hashRequestId ?? pending.hashRequestId, + optimizationRunId: + message.optimizationRunId ?? pending.optimizationRunId, + httpStatus: message.httpStatus, + }), + ); break; } }); }; -const bridgeFetch = ( - input: PetrinautOptimizationInput, +/** + * Ask the host to open an optimizer NDJSON stream and synthesize a `Response` + * from the relayed response-start/chunk/end/error messages. `initiate` is the + * `optimizationAttach` message naming the detached run's event stream, and + * must carry `requestId` so the relayed replies correlate back to this call. + */ +const openBridgeStream = ( + requestId: string, + initiate: IframeToHostMessage, signal?: OptimizationSignal, ): Promise => { ensureListener(); - const requestId = nextRequestId(); let streamController!: ReadableStreamDefaultController; const stream = new ReadableStream({ @@ -178,7 +331,10 @@ const bridgeFetch = ( postToHost({ kind: "optimizationAbort", requestId }); rejectPendingRequest( requestId, - new Error("The optimization service did not respond in time"), + new PetrinautOptimizationTransportError( + "The optimization service did not respond in time", + { category: "network" }, + ), ); }, RESPONSE_START_TIMEOUT_MS); const clearResponseStartTimeout = () => clearTimeout(responseStartTimeout); @@ -190,6 +346,8 @@ const bridgeFetch = ( resolveResponse: resolve, rejectResponse: reject, responded: false, + hashRequestId: null, + optimizationRunId: null, clearResponseStartTimeout, cleanup: () => { clearResponseStartTimeout(); @@ -204,11 +362,75 @@ const bridgeFetch = ( } signal?.addEventListener("abort", onAbort, { once: true }); - postToHost({ kind: "optimizationRequest", requestId, input }); + postToHost(initiate); return response; }; -const readHttpError = async (response: Response): Promise => { +/** + * Ask the host to create a detached optimization run and resolve its + * server-issued run id. Rejects with a classified + * {@link PetrinautOptimizationTransportError} (carrying the HTTP status and + * any `Retry-After` seconds) when the host reports a failure, or with an + * `AbortError` when `signal` fires first. + */ +const requestOptimizationRunCreation = ( + input: PetrinautOptimizationInput, + signal?: OptimizationSignal, +): Promise<{ runId: string }> => { + ensureListener(); + + const requestId = nextRequestId(); + + return new Promise<{ runId: string }>((resolve, reject) => { + const timeout = setTimeout(() => { + rejectPendingCreate( + requestId, + new PetrinautOptimizationTransportError( + "The optimization service did not respond in time", + { category: "network" }, + ), + ); + }, RESPONSE_START_TIMEOUT_MS); + const onAbort = () => rejectPendingCreate(requestId, abortError()); + + pendingCreates.set(requestId, { + resolve, + reject, + cleanup: () => { + clearTimeout(timeout); + signal?.removeEventListener("abort", onAbort); + }, + }); + + if (signal?.aborted) { + onAbort(); + return; + } + signal?.addEventListener("abort", onAbort, { once: true }); + + postToHost({ kind: "optimizationCreate", requestId, input }); + }); +}; + +/** The correlation ids NodeAPI mirrors onto every proxied response. */ +type ResponseCorrelation = { + hashRequestId: string | null; + optimizationRunId: string | null; +}; + +const responseCorrelation = (response: Response): ResponseCorrelation => ({ + hashRequestId: response.headers.get("x-hash-request-id"), + optimizationRunId: response.headers.get("x-optimization-run-id"), +}); + +const readHttpError = async ( + response: Response, +): Promise => { + const correlation = { + category: "http" as const, + ...responseCorrelation(response), + httpStatus: response.status, + }; const body = await response.text(); if (body) { try { @@ -220,26 +442,51 @@ const readHttpError = async (response: Response): Promise => { ? json.message : null; if (message) { - return new Error(message); + return new PetrinautOptimizationTransportError(message, correlation); } } catch { // Fall through and include the plain response body. } } - return new Error( + return new PetrinautOptimizationTransportError( body || `Optimization request failed with status ${response.status} ${response.statusText}`, + correlation, ); }; -const parseEventLine = (line: string): PetrinautOptimizationEvent => { +/** + * A protocol violation while decoding the optimizer's NDJSON stream. The + * correlation ids come from the response being decoded, so a stream that goes + * wrong mid-flight is as traceable as one that fails its status line. + */ +const protocolError = (message: string, correlation?: ResponseCorrelation) => + new PetrinautOptimizationTransportError(message, { + category: "protocol", + ...correlation, + }); + +const parseEventLine = ( + line: string, + correlation: ResponseCorrelation, +): PetrinautOptimizationEvent => { let parsed: unknown; try { parsed = JSON.parse(line); } catch { - throw new Error("The optimizer returned malformed NDJSON"); + throw protocolError("The optimizer returned malformed NDJSON", correlation); + } + // A syntactically valid line can still be the wrong shape. Classify that as + // a protocol violation too, so it takes the same reconnect path and the + // schema's validation detail never reaches the user. + const event = petrinautOptimizationEventSchema.safeParse(parsed); + if (!event.success) { + throw protocolError( + "The optimizer returned an unrecognized event", + correlation, + ); } - return petrinautOptimizationEventSchema.parse(parsed); + return event.data; }; /** Validate and decode the optimizer's public NDJSON protocol. */ @@ -249,8 +496,12 @@ export async function* parsePetrinautOptimizationResponse( if (!response.ok) { throw await readHttpError(response); } + const correlation = responseCorrelation(response); if (!response.body) { - throw new Error("The optimizer returned an empty response body"); + throw protocolError( + "The optimizer returned an empty response body", + correlation, + ); } const reader = response.body.getReader(); @@ -261,9 +512,12 @@ export async function* parsePetrinautOptimizationResponse( const parseAndTrack = (line: string, terminalSeen: boolean) => { if (terminalSeen) { - throw new Error("The optimizer returned data after a terminal event"); + throw protocolError( + "The optimizer returned data after a terminal event", + correlation, + ); } - const event = parseEventLine(line); + const event = parseEventLine(line, correlation); return { event, terminal: event.type === "complete" || event.type === "error", @@ -298,7 +552,10 @@ export async function* parsePetrinautOptimizationResponse( yield parsed.event; } if (!terminalEventSeen) { - throw new Error("The optimizer stream ended without a terminal event"); + throw protocolError( + "The optimizer stream ended without a terminal event", + correlation, + ); } } finally { if (!reachedEnd) { @@ -307,11 +564,25 @@ export async function* parsePetrinautOptimizationResponse( } } -async function* streamOptimization( - input: PetrinautOptimizationInput, - signal?: OptimizationSignal, +async function* streamOptimizationRun( + runId: string, + options: { + cursor: number; + signal?: OptimizationSignal; + onAttached?: () => void; + }, ): AsyncGenerator { - const response = await bridgeFetch(input, signal); + const requestId = nextRequestId(); + const response = await openBridgeStream( + requestId, + { kind: "optimizationAttach", requestId, runId, cursor: options.cursor }, + options.signal, + ); + if (response.ok) { + // The attachment was accepted; events may still be a long way off on a + // quiet run, so consumers get an explicit "connected" signal. + options.onAttached?.(); + } yield* parsePetrinautOptimizationResponse(response); } @@ -319,6 +590,25 @@ async function* streamOptimization( * HASH implementation of Petrinaut's host capability. The sandboxed editor * never receives API credentials or an upstream URL; its parent owns both. */ -export const createBridgePetrinautOptimization = (): PetrinautOptimization => ({ - optimize: (input, options) => streamOptimization(input, options?.signal), -}); +export const createBridgePetrinautOptimization = (): PetrinautOptimization => { + // Installed at bridge creation, not lazily on the first optimization call: + // a late `optimizationCreateResult` arriving right after an iframe reload + // must be heard so its orphaned run can be cancelled. + ensureListener(); + return { + createOptimizationRun: (input, options) => + requestOptimizationRunCreation(input, options?.signal), + attachOptimizationRun: (runId, options) => + streamOptimizationRun(runId, { + cursor: options?.cursor ?? 0, + signal: options?.signal, + onAttached: options?.onAttached, + }), + cancelOptimizationRun: (runId) => { + // Fire-and-forget by design: the host DELETEs the run and only logs + // failures, so there is no reply to await. + postToHost({ kind: "optimizationCancel", runId }); + return Promise.resolve(); + }, + }; +}; diff --git a/apps/hash-frontend/src/pages/processes/shared/messages.test.ts b/apps/hash-frontend/src/pages/processes/shared/messages.test.ts index 76e8e270c3c..e3f2e5067ce 100644 --- a/apps/hash-frontend/src/pages/processes/shared/messages.test.ts +++ b/apps/hash-frontend/src/pages/processes/shared/messages.test.ts @@ -24,6 +24,17 @@ describe("isHostToIframeMessage", () => { ).toBe(false); }); + it("accepts detached-run creation replies", () => { + expect( + isHostToIframeMessage({ + kind: "optimizationCreateResult", + requestId: "req-1", + ok: true, + runId: "run-1", + }), + ).toBe(true); + }); + it("rejects unknown host message kinds", () => { expect(isHostToIframeMessage({ kind: "notAHostMessage" })).toBe(false); }); diff --git a/apps/hash-frontend/src/pages/processes/shared/messages.ts b/apps/hash-frontend/src/pages/processes/shared/messages.ts index 1466e8f16cd..49a47732a08 100644 --- a/apps/hash-frontend/src/pages/processes/shared/messages.ts +++ b/apps/hash-frontend/src/pages/processes/shared/messages.ts @@ -179,12 +179,49 @@ export type HostToIframeMessage = message: string; } | { - /** First reply to an `optimizationRequest`. */ + /** + * Reply to an `optimizationCreate`. On success carries the run id the + * NodeAPI issued for the detached run. On failure `category` classifies + * the problem, `status` carries the HTTP status (when one was received), + * `retryAfter` the parsed `Retry-After` header in seconds (when sent, + * e.g. on 429), and `message` is a safe, human-readable summary that + * never contains internal or user-authored content. + * + * `hashRequestId` and `optimizationRunId` mirror the correlation headers + * of the create response, so a failed creation can still be traced + * through NodeAPI and the optimizer's logs. NodeAPI forwards the + * upstream run id even when creation fails, and that is exactly the case + * worth correlating, so both are reported on failures too. They are + * absent when no response was received at all (a network failure). + */ + kind: "optimizationCreateResult"; + requestId: string; + ok: boolean; + runId?: string; + status?: number; + retryAfter?: number; + message?: string; + category?: "network" | "http" | "protocol" | "aborted"; + hashRequestId?: string | null; + optimizationRunId?: string | null; + } + | { + /** + * First reply to an `optimizationAttach`, mirroring the proxied HTTP + * response's status line. + */ kind: "optimizationResponseStart"; requestId: string; ok: boolean; status: number; statusText: string; + /** + * NodeAPI's `x-hash-request-id` / `X-Optimization-Run-ID` response + * headers, forwarded so a later transport failure remains traceable to + * the NodeAPI and optimizer logs. Null when the header is absent. + */ + hashRequestId: string | null; + optimizationRunId: string | null; } | { /** A verbatim chunk of the optimizer's NDJSON response body. */ @@ -198,10 +235,19 @@ export type HostToIframeMessage = requestId: string; } | { - /** The optimization fetch failed before or while streaming. */ + /** + * The optimization fetch failed before or while streaming. `category` + * classifies the failure so the UI can show an actionable message + * instead of a raw exception string; `message` is a safe, human-readable + * summary that never contains internal or user-authored content. + */ kind: "optimizationError"; requestId: string; + category: "network" | "http" | "protocol" | "aborted"; message: string; + hashRequestId?: string | null; + optimizationRunId?: string | null; + httpStatus?: number; }; /** @@ -303,18 +349,43 @@ export type IframeToHostMessage = } | { /** - * Ask the authenticated host to start a Petrinaut optimization. The - * host validates this public request before forwarding it to NodeAPI. + * Ask the authenticated host to create a detached optimization run. + * No events flow on this request id — the host replies once with + * `optimizationCreateResult` and the iframe then attaches to the run's + * event stream via `optimizationAttach`. The host validates this + * public request before forwarding it to NodeAPI. */ - kind: "optimizationRequest"; + kind: "optimizationCreate"; requestId: string; input: PetrinautOptimizationInput; } + | { + /** + * Attach to a detached optimization run's NDJSON event stream, + * replaying events with `seq` greater than `cursor` (0 replays + * everything) before tailing live events. The host relays the stream + * via the `optimizationResponseStart`/`optimizationChunk`/ + * `optimizationEnd`/`optimizationError` family, keyed by `requestId`. + */ + kind: "optimizationAttach"; + requestId: string; + runId: string; + cursor: number; + } | { /** Abort the matching in-flight optimization all the way upstream. */ kind: "optimizationAbort"; requestId: string; } + | { + /** + * Idempotently cancel a detached optimization run server-side. Aborting + * a local `optimizationAttach` only drops the connection; this stops + * the run itself. Fire-and-forget: the host sends no reply. + */ + kind: "optimizationCancel"; + runId: string; + } | { /** * The AI-assistant conversation changed (a turn finished, or the @@ -346,6 +417,7 @@ const hostToIframeMessageKinds: ReadonlySet = new Set< "aiChatChunk", "aiChatEnd", "aiChatError", + "optimizationCreateResult", "optimizationResponseStart", "optimizationChunk", "optimizationEnd", diff --git a/apps/hash-frontend/src/pages/processes/shared/use-host-bridge.ts b/apps/hash-frontend/src/pages/processes/shared/use-host-bridge.ts index 0bff1c6d25f..22dd9dc9e4f 100644 --- a/apps/hash-frontend/src/pages/processes/shared/use-host-bridge.ts +++ b/apps/hash-frontend/src/pages/processes/shared/use-host-bridge.ts @@ -31,12 +31,18 @@ type HostBridgeHandlers = { onAiChatAbort?: ( payload: Extract, ) => void; - onOptimizationRequest?: ( - payload: Extract, + onOptimizationCreate?: ( + payload: Extract, + ) => void; + onOptimizationAttach?: ( + payload: Extract, ) => void; onOptimizationAbort?: ( payload: Extract, ) => void; + onOptimizationCancel?: ( + payload: Extract, + ) => void; onAiMessagesChanged?: ( payload: Extract, ) => void; @@ -121,12 +127,18 @@ export const useHostBridge = ({ case "aiChatAbort": current.onAiChatAbort?.(data); break; - case "optimizationRequest": - current.onOptimizationRequest?.(data); + case "optimizationCreate": + current.onOptimizationCreate?.(data); + break; + case "optimizationAttach": + current.onOptimizationAttach?.(data); break; case "optimizationAbort": current.onOptimizationAbort?.(data); break; + case "optimizationCancel": + current.onOptimizationCancel?.(data); + break; case "aiMessagesChanged": current.onAiMessagesChanged?.(data); break; diff --git a/apps/petrinaut-opt/README.md b/apps/petrinaut-opt/README.md index 4304956367d..53a450c86df 100644 --- a/apps/petrinaut-opt/README.md +++ b/apps/petrinaut-opt/README.md @@ -10,13 +10,10 @@ read Petrinaut models, scenario bindings, metrics, or the Petrinaut type system. ## API -Both optimization endpoints accept the complete optimization manifest as their -JSON request body. The manifest is produced by the Petrinaut UI/Node API and is -forwarded unchanged to the CLI. +Run creation accepts the complete optimization manifest as its JSON request +body. The manifest is produced by the Petrinaut UI/Node API and is forwarded +unchanged to the CLI. -- `POST /optimize/all` streams every finished trial. -- `POST /optimize/best` streams the best-so-far result after each finished - trial. No data frame is emitted until at least one trial completes. - `POST /optimize/runs` starts a detached run and returns its id. Attach or reattach to its replayable event stream with `GET /optimize/runs/{run_id}/events`; `DELETE /optimize/runs/{run_id}` diff --git a/apps/petrinaut-opt/openapi/openapi.json b/apps/petrinaut-opt/openapi/openapi.json index 36c8aaa4e8e..2340b1416ce 100644 --- a/apps/petrinaut-opt/openapi/openapi.json +++ b/apps/petrinaut-opt/openapi/openapi.json @@ -166,122 +166,6 @@ "summary": "Health" } }, - "/optimize/all": { - "post": { - "description": "Stream one SSE data frame for every completed Optuna trial.", - "operationId": "post_optimize_all_optimize_all_post", - "requestBody": { - "content": { - "application/json": { - "schema": { - "additionalProperties": true, - "title": "Optimization Manifest", - "type": "object" - } - } - }, - "required": true - }, - "responses": { - "200": { - "content": { - "text/event-stream": { - "schema": { - "type": "string" - } - } - }, - "description": "Server-Sent Events optimization stream" - }, - "413": { - "description": "The optimization manifest exceeds 8 MiB" - }, - "422": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - }, - "description": "Validation Error" - }, - "429": { - "description": "The service is already at its study limit", - "headers": { - "Retry-After": { - "description": "Seconds to wait before retrying the study", - "schema": { - "type": "string" - } - } - } - }, - "500": { - "description": "The CLI or optimization study could not initialize" - } - }, - "summary": "Post Optimize All" - } - }, - "/optimize/best": { - "post": { - "description": "Stream the best-so-far SSE data frame after every completed trial.", - "operationId": "post_optimize_best_optimize_best_post", - "requestBody": { - "content": { - "application/json": { - "schema": { - "additionalProperties": true, - "title": "Optimization Manifest", - "type": "object" - } - } - }, - "required": true - }, - "responses": { - "200": { - "content": { - "text/event-stream": { - "schema": { - "type": "string" - } - } - }, - "description": "Server-Sent Events optimization stream" - }, - "413": { - "description": "The optimization manifest exceeds 8 MiB" - }, - "422": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - }, - "description": "Validation Error" - }, - "429": { - "description": "The service is already at its study limit", - "headers": { - "Retry-After": { - "description": "Seconds to wait before retrying the study", - "schema": { - "type": "string" - } - } - } - }, - "500": { - "description": "The CLI or optimization study could not initialize" - } - }, - "summary": "Post Optimize Best" - } - }, "/optimize/runs": { "post": { "description": "Start a detached run that remains available for later SSE attachment.\n\nWhen the caller stamps ``x-hash-account-id`` (the authenticated proxy\ndoes), the run is owned: the account is single-flight while it lives, and\nonly requests carrying the same tag may attach to or cancel it.", diff --git a/apps/petrinaut-opt/src/optimization_api.py b/apps/petrinaut-opt/src/optimization_api.py index 6cd97b989fc..93c1bdbd805 100644 --- a/apps/petrinaut-opt/src/optimization_api.py +++ b/apps/petrinaut-opt/src/optimization_api.py @@ -7,7 +7,7 @@ import logging import os import threading -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from collections.abc import AsyncIterator, Awaitable, Callable from contextlib import asynccontextmanager, suppress from pathlib import Path from typing import Any @@ -36,12 +36,9 @@ MAX_REQUEST_BODY_BYTES = 8 * 1024 * 1024 MAX_ACTIVE_OPTIMIZATIONS = 4 RETRY_AFTER_SECONDS = 30 -_OPTIMIZATION_PATHS = {"/optimize/all", "/optimize/best", "/optimize/runs"} -_SSE_RESPONSES = { - 200: { - "description": "Server-Sent Events optimization stream", - "content": {"text/event-stream": {"schema": {"type": "string"}}}, - }, +_OPTIMIZATION_PATHS = {"/optimize/runs"} +_CREATE_RUN_RESPONSES = { + 201: {"description": "A detached optimization run was started"}, 413: {"description": "The optimization manifest exceeds 8 MiB"}, 429: { "description": "The service is already at its study limit", @@ -54,12 +51,6 @@ }, 500: {"description": "The CLI or optimization study could not initialize"}, } -_CREATE_RUN_RESPONSES = { - 201: {"description": "A detached optimization run was started"}, - 413: _SSE_RESPONSES[413], - 429: _SSE_RESPONSES[429], - 500: _SSE_RESPONSES[500], -} _RUN_EVENTS_RESPONSES = { 200: { "description": ( @@ -340,16 +331,6 @@ async def cleanup() -> None: return cleanup -async def _stream_with_cleanup( - stream: AsyncIterator[str], cleanup: Callable[[], Awaitable[None]] -) -> AsyncIterator[str]: - try: - async for frame in stream: - yield frame - finally: - await asyncio.shield(cleanup()) - - def _initialization_error( app: FastAPI, run_id: str, @@ -427,66 +408,6 @@ async def _admit_and_initialize_run( return run_id, optimizer, correlation -@app.post( - "/optimize/all", - response_class=StreamingResponse, - responses=_SSE_RESPONSES, -) -async def post_optimize_all( - request: Request, - optimization_manifest: dict[str, Any], -) -> StreamingResponse: - """Stream one SSE data frame for every completed Optuna trial.""" - run_id, optimizer, _correlation = await _admit_and_initialize_run( - request, optimization_manifest - ) - - cleanup = _create_admitted_run_cleanup(request.app, optimizer) - return StreamingResponse( - _stream_with_cleanup( - optimizer.stream_all(request, run_id=run_id, n_trials=optimizer.n_trials), - cleanup, - ), - media_type="text/event-stream", - headers={ - "Cache-Control": "no-cache", - "X-Accel-Buffering": "no", - "X-Optimization-Run-ID": run_id, - }, - background=BackgroundTask(cleanup), - ) - - -@app.post( - "/optimize/best", - response_class=StreamingResponse, - responses=_SSE_RESPONSES, -) -async def post_optimize_best( - request: Request, - optimization_manifest: dict[str, Any], -) -> StreamingResponse: - """Stream the best-so-far SSE data frame after every completed trial.""" - run_id, optimizer, _correlation = await _admit_and_initialize_run( - request, optimization_manifest - ) - - cleanup = _create_admitted_run_cleanup(request.app, optimizer) - return StreamingResponse( - _stream_with_cleanup( - optimizer.stream_best(request, run_id=run_id, n_trials=optimizer.n_trials), - cleanup, - ), - media_type="text/event-stream", - headers={ - "Cache-Control": "no-cache", - "X-Accel-Buffering": "no", - "X-Optimization-Run-ID": run_id, - }, - background=BackgroundTask(cleanup), - ) - - @app.post("/optimize/runs", status_code=201, responses=_CREATE_RUN_RESPONSES) async def post_optimize_runs( request: Request, diff --git a/apps/petrinaut-opt/src/petrinaut_optimizer.py b/apps/petrinaut-opt/src/petrinaut_optimizer.py index 20500afaa09..578759e136b 100644 --- a/apps/petrinaut-opt/src/petrinaut_optimizer.py +++ b/apps/petrinaut-opt/src/petrinaut_optimizer.py @@ -8,14 +8,12 @@ import logging import math import os -import queue import threading -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import Callable, Mapping from datetime import datetime from typing import Any, Literal, TypeAlias, cast import optuna -from fastapi import Request from opentelemetry import context as otel_context, trace from opentelemetry.trace import Span, Status, StatusCode @@ -32,7 +30,6 @@ "random": optuna.samplers.RandomSampler, } DEFAULT_STUDY_NAME = "opt_study" -SSE_HEARTBEAT_SECONDS = 30 # The service-side mirror of the optimization manifest's trial cap; it also # bounds every run's in-memory event log to one frame per trial plus a # handful of control frames, even against a CLI reporting a huge study. @@ -465,286 +462,3 @@ async def pump_events( except ValueError: pass study_span.end() - - async def stream_all( - self, request: Request, run_id: str, n_trials: int - ) -> AsyncIterator[str]: - """Stream Yannis's per-trial SSE frames, followed by the done frame.""" - app = request.app - log_context = { - "request_id": request.headers.get("x-hash-request-id"), - "run_id": run_id, - } - if not self.lock.acquire(blocking=False): - yield 'event: error\ndata: {"message": "already running"}\n\n' - return - - set_status(app, run_id, phase=Phase.running, detail="optimization running") - log.info( - "optimization study started", - extra={"event": "study_started", "trials": n_trials, **log_context}, - ) - loop = asyncio.get_running_loop() - events: asyncio.Queue[dict[str, Any] | object] = asyncio.Queue() - stop_flag = threading.Event() - - def callback(study: optuna.Study, trial: optuna.trial.FrozenTrial) -> None: - payload = { - "step": trial.number, - "params": dict(trial.params), - "init_state": {}, - "metric": trial.value, - "state": trial.state.name, - } - loop.call_soon_threadsafe(events.put_nowait, payload) - if stop_flag.is_set(): - study.stop() - - worker, study_span = self._start_study_worker( - loop, events, n_trials=n_trials, callback=callback - ) - next_heartbeat = loop.time() + SSE_HEARTBEAT_SECONDS - completed = False - - try: - while True: - if await request.is_disconnected(): - stop_flag.set() - set_status( - app, - run_id, - phase=Phase.idle, - detail="client disconnected, stopped", - ) - log.info( - "client disconnected, stopping optimization study", - extra={"event": "client_disconnected", **log_context}, - ) - break - heartbeat_wait = max(0.0, next_heartbeat - loop.time()) - try: - item = await asyncio.wait_for( - events.get(), - timeout=min(_DISCONNECT_POLL_SECONDS, heartbeat_wait), - ) - except asyncio.TimeoutError: - if loop.time() >= next_heartbeat: - yield ": heartbeat\n\n" - next_heartbeat = loop.time() + SSE_HEARTBEAT_SECONDS - continue - if item is _SENTINEL: - set_status( - app, - run_id, - phase=Phase.done, - detail="optimization completed", - ) - completed = True - log.info( - "optimization study completed", - extra={ - "event": "study_completed", - "trials": n_trials, - **log_context, - }, - ) - yield "event: done\ndata: {}\n\n" - break - event = cast(dict[str, Any], item) - if event.get("state") == "ERROR": - set_status( - app, - run_id, - phase=Phase.error, - detail=cast(str, event.get("message")), - ) - log.warning( - "optimization study failed", - extra={ - "event": "study_failed", - **log_context, - }, - ) - yield f"data: {json.dumps(event)}\n\n" - if event.get("state") == "ERROR": - break - finally: - stop_flag.set() - try: - # Only a study that ran to completion left the CLI idle enough - # for the graceful EOF shutdown; every other exit (disconnect, - # error, cancellation) terminates the process group promptly. - await asyncio.to_thread(self.pn_model.close, graceful=completed) - await asyncio.to_thread(worker.join, _WORKER_SHUTDOWN_TIMEOUT_SECONDS) - if worker.is_alive(): - log.error( - "Petrinaut optimizer worker did not stop after CLI shutdown", - extra={"event": "worker_join_timeout", **log_context}, - ) - finally: - self.lock.release() - try: - study_span.set_attribute( - "optuna.study.best_value", self.study.best_value - ) - except ValueError: - # No trial completed (immediate disconnect, or all pruned), - # so there is no best value to record. - pass - study_span.end() - - async def stream_best( - self, request: Request, run_id: str, n_trials: int - ) -> AsyncIterator[str]: - """Stream Yannis's best-so-far SSE frames, followed by the done frame.""" - app = request.app - log_context = { - "request_id": request.headers.get("x-hash-request-id"), - "run_id": run_id, - } - if not self.lock.acquire(blocking=False): - yield 'event: error\ndata: {"message": "already running"}\n\n' - return - - set_status(app, run_id, phase=Phase.running, detail="optimization running") - log.info( - "optimization study started", - extra={"event": "study_started", "trials": n_trials, **log_context}, - ) - loop = asyncio.get_running_loop() - events: asyncio.Queue[dict[str, Any] | object] = asyncio.Queue() - stop_flag = threading.Event() - - def callback(study: optuna.Study, trial: optuna.trial.FrozenTrial) -> None: - has_completed = any( - candidate.state is optuna.trial.TrialState.COMPLETE - for candidate in study.get_trials(deepcopy=False) - ) - if has_completed: - payload = { - "step": trial.number, - "params": dict(study.best_params), - "init_state": {}, - "metric": study.best_value, - "state": "COMPLETE", - } - loop.call_soon_threadsafe(events.put_nowait, payload) - if stop_flag.is_set(): - study.stop() - - worker, study_span = self._start_study_worker( - loop, events, n_trials=n_trials, callback=callback - ) - next_heartbeat = loop.time() + SSE_HEARTBEAT_SECONDS - completed = False - - try: - while True: - if await request.is_disconnected(): - stop_flag.set() - set_status( - app, - run_id, - phase=Phase.idle, - detail="client disconnected, stopped", - ) - log.info( - "client disconnected, stopping optimization study", - extra={"event": "client_disconnected", **log_context}, - ) - break - heartbeat_wait = max(0.0, next_heartbeat - loop.time()) - try: - item = await asyncio.wait_for( - events.get(), - timeout=min(_DISCONNECT_POLL_SECONDS, heartbeat_wait), - ) - except asyncio.TimeoutError: - if loop.time() >= next_heartbeat: - yield ": heartbeat\n\n" - next_heartbeat = loop.time() + SSE_HEARTBEAT_SECONDS - continue - if item is _SENTINEL: - set_status( - app, - run_id, - phase=Phase.done, - detail="optimization completed", - ) - completed = True - log.info( - "optimization study completed", - extra={ - "event": "study_completed", - "trials": n_trials, - **log_context, - }, - ) - yield "event: done\ndata: {}\n\n" - break - event = cast(dict[str, Any], item) - if event.get("state") == "ERROR": - set_status( - app, - run_id, - phase=Phase.error, - detail=cast(str, event.get("message")), - ) - log.warning( - "optimization study failed", - extra={ - "event": "study_failed", - **log_context, - }, - ) - yield f"data: {json.dumps(event)}\n\n" - if event.get("state") == "ERROR": - break - finally: - stop_flag.set() - try: - # Only a study that ran to completion left the CLI idle enough - # for the graceful EOF shutdown; every other exit (disconnect, - # error, cancellation) terminates the process group promptly. - await asyncio.to_thread(self.pn_model.close, graceful=completed) - await asyncio.to_thread(worker.join, _WORKER_SHUTDOWN_TIMEOUT_SECONDS) - if worker.is_alive(): - log.error( - "Petrinaut optimizer worker did not stop after CLI shutdown", - extra={"event": "worker_join_timeout", **log_context}, - ) - finally: - self.lock.release() - try: - study_span.set_attribute( - "optuna.study.best_value", self.study.best_value - ) - except ValueError: - # No trial completed (immediate disconnect, or all pruned), - # so there is no best value to record. - pass - study_span.end() - - def run_stream(self, study: optuna.Study, objective: Any, n_trials: int) -> Any: - """Run a study synchronously, retaining the original local-test shape.""" - events: queue.Queue[Any] = queue.Queue() - done = object() - - def callback(_study: optuna.Study, trial: optuna.trial.FrozenTrial) -> None: - events.put( - ( - str(trial.state), - trial.number, - dict(trial.params), - {}, - trial.value, - ) - ) - - def run() -> None: - study.optimize(objective, n_trials=n_trials, callbacks=[callback]) - events.put(done) - - threading.Thread(target=run, daemon=True).start() - while (item := events.get()) is not done: - yield item diff --git a/apps/petrinaut-opt/tests/test_optimization_api.py b/apps/petrinaut-opt/tests/test_optimization_api.py index 0f8180e2794..4db35cfce04 100644 --- a/apps/petrinaut-opt/tests/test_optimization_api.py +++ b/apps/petrinaut-opt/tests/test_optimization_api.py @@ -15,86 +15,40 @@ from src.optimization_runs import CANCELLED_FRAME, RunState -class FakeOptimizer: - n_trials = 1 - - async def stream_all(self, *_args: Any, **_kwargs: Any): - yield ( - 'data: {"step": 0, "params": {"rate": 1.0}, ' - '"init_state": {}, "metric": 2.0, "state": "COMPLETE"}\n\n' - ) - yield "event: done\ndata: {}\n\n" - - async def stream_best(self, *_args: Any, **_kwargs: Any): - yield ( - 'data: {"step": 0, "params": {"rate": 1.0}, ' - '"init_state": {}, "metric": 2.0, "state": "COMPLETE"}\n\n' - ) - yield "event: done\ndata: {}\n\n" - - -def test_posts_an_opaque_manifest_to_the_all_sse_route( - optimization_manifest: dict, - monkeypatch, -) -> None: - received: list[dict[str, Any]] = [] - - def initialize(manifest: dict[str, Any], **_kwargs: Any) -> FakeOptimizer: - received.append(manifest) - return FakeOptimizer() - - monkeypatch.setattr(optimization_api, "initialize_optimizer", initialize) - - with TestClient(optimization_api.app) as client: - response = client.post("/optimize/all", json=optimization_manifest) - assert optimization_api.app.state.active_optimizations == 0 - - assert response.status_code == 200 - assert response.headers["content-type"].startswith("text/event-stream") - assert response.headers["cache-control"] == "no-cache" - assert response.headers["x-accel-buffering"] == "no" - assert response.headers["x-optimization-run-id"] - assert received == [optimization_manifest] - assert response.text.endswith("event: done\ndata: {}\n\n") - - -def test_posts_an_opaque_manifest_to_the_best_sse_route( - optimization_manifest: dict, - monkeypatch, -) -> None: - received: list[dict[str, Any]] = [] - - def initialize(manifest: dict[str, Any], **_kwargs: Any) -> FakeOptimizer: - received.append(manifest) - return FakeOptimizer() - - monkeypatch.setattr(optimization_api, "initialize_optimizer", initialize) - - with TestClient(optimization_api.app) as client: - response = client.post("/optimize/best", json=optimization_manifest) - assert optimization_api.app.state.active_optimizations == 0 +class RecordingModel: + """Track how the CLI adapter is shut down.""" - assert response.status_code == 200 - assert response.headers["content-type"].startswith("text/event-stream") - assert received == [optimization_manifest] - assert response.text.endswith("event: done\ndata: {}\n\n") + def __init__(self) -> None: + self.close_calls: list[bool] = [] + def close(self, *, graceful: bool = True) -> None: + self.close_calls.append(graceful) -def test_get_is_not_retained_for_manifest_routes() -> None: - with TestClient(optimization_api.app) as client: - assert client.get("/optimize/all").status_code == 405 - assert client.get("/optimize/best").status_code == 405 +class FakeOptimizer: + """Minimal pump-driven optimizer double for create-route tests.""" -def test_rejects_oversized_manifests_on_both_routes(monkeypatch) -> None: - monkeypatch.setattr(optimization_api, "MAX_REQUEST_BODY_BYTES", 8) + n_trials = 1 - with TestClient(optimization_api.app) as client: - all_response = client.post("/optimize/all", content=b'{"long":true}') - best_response = client.post("/optimize/best", content=b'{"long":true}') + def __init__(self) -> None: + self.pn_model = RecordingModel() - assert all_response.status_code == 413 - assert best_response.status_code == 413 + async def pump_events( + self, + _app: Any, + _run_id: str, + _n_trials: int, + *, + on_event: Callable[[str], Any], + cancel_event: asyncio.Event, + on_outcome: Callable[[str], Any] | None = None, + correlation: Any = None, + ) -> str: + on_event("event: done\ndata: {}\n\n") + self.pn_model.close(graceful=True) + if on_outcome is not None: + on_outcome("completed") + return "completed" def test_rejects_an_oversized_chunked_manifest(monkeypatch) -> None: @@ -123,8 +77,8 @@ async def downstream(_scope, receive_body, _send) -> None: "http_version": "1.1", "method": "POST", "scheme": "http", - "path": "/optimize/all", - "raw_path": b"/optimize/all", + "path": "/optimize/runs", + "raw_path": b"/optimize/runs", "query_string": b"", "root_path": "", "headers": [], @@ -139,35 +93,6 @@ async def downstream(_scope, receive_body, _send) -> None: assert outgoing[0]["status"] == 413 -def test_reports_initialization_failure_with_the_run_id( - optimization_manifest: dict, - monkeypatch, -) -> None: - def initialize(_manifest: dict[str, Any], **_kwargs: Any) -> FakeOptimizer: - raise RuntimeError("manifest rejected by CLI") - - monkeypatch.setattr(optimization_api, "initialize_optimizer", initialize) - - with TestClient(optimization_api.app) as client: - response = client.post("/optimize/all", json=optimization_manifest) - run_id = response.headers["x-optimization-run-id"] - assert optimization_api.app.state.active_optimizations == 0 - statuses = client.get("/status") - run_status = client.get(f"/status/{run_id}") - - assert response.status_code == 500 - assert "manifest rejected by CLI" in response.json()["detail"] - assert statuses.json() == [ - { - "phase": "error", - "detail": "Petrinaut CLI and Optimization Model could NOT be initialized", - "updated_at": statuses.json()[0]["updated_at"], - "run_id": run_id, - } - ] - assert run_status.json() == statuses.json()[0] - - def test_initialization_failure_log_omits_the_raw_error_message( optimization_manifest: dict, monkeypatch, @@ -183,7 +108,7 @@ def initialize(_manifest: dict[str, Any], **_kwargs: Any) -> FakeOptimizer: with caplog.at_level("ERROR", logger="pn_api"): with TestClient(optimization_api.app) as client: - response = client.post("/optimize/all", json=optimization_manifest) + response = client.post("/optimize/runs", json=optimization_manifest) failures = [r for r in caplog.records if r.event == "initialization_failed"] assert failures @@ -210,26 +135,13 @@ def initialize(_manifest: dict[str, Any], **_kwargs: Any) -> FakeOptimizer: with TestClient(optimization_api.app) as client: event_loop_thread_id = client.portal.call(threading.get_ident) - response = client.post("/optimize/all", json=optimization_manifest) + response = client.post("/optimize/runs", json=optimization_manifest) - assert response.status_code == 200 + assert response.status_code == 201 assert initializer_thread_ids assert initializer_thread_ids[0] != event_loop_thread_id -def test_rejects_studies_above_the_process_local_limit( - optimization_manifest: dict, -) -> None: - with TestClient(optimization_api.app) as client: - optimization_api.app.state.active_optimizations = ( - optimization_api.MAX_ACTIVE_OPTIMIZATIONS - ) - response = client.post("/optimize/all", json=optimization_manifest) - - assert response.status_code == 429 - assert response.headers["retry-after"] == str(optimization_api.RETRY_AFTER_SECONDS) - - def test_capacity_rejection_is_logged_with_the_request_id( optimization_manifest: dict, caplog: pytest.LogCaptureFixture, @@ -240,7 +152,7 @@ def test_capacity_rejection_is_logged_with_the_request_id( optimization_api.MAX_ACTIVE_OPTIMIZATIONS ) response = client.post( - "/optimize/all", + "/optimize/runs", json=optimization_manifest, headers={"x-hash-request-id": "request-cap-1"}, ) @@ -256,16 +168,6 @@ def test_capacity_rejection_is_logged_with_the_request_id( assert "manifest" not in rejection.getMessage() -class RecordingModel: - """Track how the CLI adapter is shut down.""" - - def __init__(self) -> None: - self.close_calls: list[bool] = [] - - def close(self, *, graceful: bool = True) -> None: - self.close_calls.append(graceful) - - def _admitted_test_app(active_optimizations: int) -> FastAPI: test_app = FastAPI() test_app.state.optimization_admission_lock = asyncio.Lock() @@ -279,31 +181,6 @@ def _optimizer_with_recording_model() -> FakeOptimizer: return optimizer -def test_releases_admission_slot_when_a_stream_fails() -> None: - test_app = _admitted_test_app(active_optimizations=1) - optimizer = _optimizer_with_recording_model() - cleanup = optimization_api._create_admitted_run_cleanup( - test_app, - optimizer, # type: ignore[arg-type] - ) - - async def failing_stream(): - raise RuntimeError("stream failed") - yield "unreachable" # pragma: no cover - - async def consume() -> None: - with pytest.raises(RuntimeError, match="stream failed"): - async for _frame in optimization_api._stream_with_cleanup( - failing_stream(), cleanup - ): - pass - - asyncio.run(consume()) - - assert test_app.state.active_optimizations == 0 - assert optimizer.pn_model.close_calls == [False] # type: ignore[attr-defined] - - def test_admitted_run_cleanup_releases_the_slot_exactly_once() -> None: test_app = _admitted_test_app(active_optimizations=1) optimizer = _optimizer_with_recording_model() @@ -322,46 +199,6 @@ async def run_twice() -> None: assert optimizer.pn_model.close_calls == [False] # type: ignore[attr-defined] -def test_background_cleanup_covers_a_stream_that_never_starts( - optimization_manifest: dict, - monkeypatch: pytest.MonkeyPatch, -) -> None: - """An aborted response may never pull the body, skipping generator finallys.""" - optimizer = _optimizer_with_recording_model() - monkeypatch.setattr( - optimization_api, "initialize_optimizer", lambda _manifest, **_kwargs: optimizer - ) - - async def abandon_response() -> None: - test_app = optimization_api.app - test_app.state.statuses = optimization_api.StatusStore() - test_app.state.optimization_admission_lock = asyncio.Lock() - test_app.state.active_optimizations = 0 - scope = { - "type": "http", - "app": test_app, - "method": "POST", - "path": "/optimize/all", - "headers": [], - "query_string": b"", - } - request = optimization_api.Request(scope) - response = await optimization_api.post_optimize_all( - request, optimization_manifest - ) - - assert test_app.state.active_optimizations == 1 - assert response.background is not None - # The client is gone before the body iterator is ever started; only - # the background task remains to release the slot and the CLI. - await response.background() - assert test_app.state.active_optimizations == 0 - - asyncio.run(abandon_response()) - - assert optimizer.pn_model.close_calls == [False] # type: ignore[attr-defined] - - def test_cancellation_during_initialization_closes_cli_and_releases_slot( optimization_manifest: dict, monkeypatch: pytest.MonkeyPatch, @@ -451,23 +288,6 @@ async def cancel_twice() -> None: assert test_app.state.active_optimizations == 0 -def test_openapi_exposes_post_sse_paths_with_an_untyped_json_body() -> None: - schema = optimization_api.app.openapi() - - for path in ("/optimize/all", "/optimize/best"): - operation = schema["paths"][path] - assert "post" in operation - assert "get" not in operation - request_schema = operation["post"]["requestBody"]["content"][ - "application/json" - ]["schema"] - assert request_schema["type"] == "object" - stream_schema = operation["post"]["responses"]["200"]["content"][ - "text/event-stream" - ]["schema"] - assert stream_schema["type"] == "string" - - class ScriptedDetachedOptimizer: """Pump-driven double for detached-run endpoint tests. @@ -978,9 +798,20 @@ def initialize(_manifest: dict[str, Any], **_kwargs: Any) -> Any: run_id = response.headers["x-optimization-run-id"] assert optimization_api.app.state.active_optimizations == 0 assert optimization_api.app.state.optimization_runs.get(run_id) is None + statuses = client.get("/status") + run_status = client.get(f"/status/{run_id}") assert response.status_code == 500 assert "manifest rejected by CLI" in response.json()["detail"] + assert statuses.json() == [ + { + "phase": "error", + "detail": "Petrinaut CLI and Optimization Model could NOT be initialized", + "updated_at": statuses.json()[0]["updated_at"], + "run_id": run_id, + } + ] + assert run_status.json() == statuses.json()[0] def test_owned_runs_enforce_account_single_flight_and_visibility( diff --git a/apps/petrinaut-opt/tests/test_petrinaut_optimizer.py b/apps/petrinaut-opt/tests/test_petrinaut_optimizer.py index 087789d9547..57b5c73c179 100644 --- a/apps/petrinaut-opt/tests/test_petrinaut_optimizer.py +++ b/apps/petrinaut-opt/tests/test_petrinaut_optimizer.py @@ -42,12 +42,6 @@ def close(self, *, graceful: bool = True) -> None: self.close_calls.append(graceful) -class SlowModel(FakeModel): - def objective(self, parameter_values: dict[str, Any]) -> float: - time.sleep(0.04) - return super().objective(parameter_values) - - class FailingModel(FakeModel): def __init__(self, description: dict[str, Any], error: Exception) -> None: super().__init__(description) @@ -70,30 +64,6 @@ def objective(self, parameter_values: dict[str, Any]) -> float: raise PetrinautClientError("CLI closed") -class ConnectedRequest: - def __init__(self) -> None: - self.app = FastAPI() - self.app.state.statuses = StatusStore() - self.headers: dict[str, str] = {} - - async def is_disconnected(self) -> bool: - return False - - -class DisconnectedAfterWorkerStarts(ConnectedRequest): - def __init__(self, model: StubbornModel) -> None: - super().__init__() - self.model = model - - async def is_disconnected(self) -> bool: - await asyncio.to_thread(self.model.entered.wait, 1) - return True - - -def _run_id(request: ConnectedRequest) -> str: - return request.app.state.statuses.create().run_id - - def test_maps_float_integer_step_and_boolean_descriptors_to_optuna( optimization_description: dict, ) -> None: @@ -232,217 +202,6 @@ def test_rejects_invalid_cli_descriptions( ) -def test_stream_all_logs_the_study_lifecycle_with_correlation( - optimization_description: dict, - caplog: pytest.LogCaptureFixture, -) -> None: - model = FakeModel(optimization_description) - optimizer = PetrinautOptimizer(model) # type: ignore[arg-type] - request = ConnectedRequest() - request.headers = {"x-hash-request-id": "request-s1"} # type: ignore[attr-defined] - run_id = _run_id(request) - - async def consume() -> None: - async for _frame in optimizer.stream_all( - request, - run_id, - optimizer.n_trials, # type: ignore[arg-type] - ): - pass - - with caplog.at_level(logging.INFO, logger="pn_optimize"): - asyncio.run(consume()) - - events = { - getattr(record, "event", None): record - for record in caplog.records - if record.name == "pn_optimize" - } - for expected in ("study_started", "study_completed"): - record = events[expected] - assert record.run_id == run_id - assert record.request_id == "request-s1" - assert record.trials == optimizer.n_trials - - -def test_disconnect_is_logged_with_the_run_id( - optimization_description: dict, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, -) -> None: - optimization_description["study"]["trials"] = 1 - monkeypatch.setattr(petrinaut_optimizer, "_WORKER_SHUTDOWN_TIMEOUT_SECONDS", 0.01) - model = StubbornModel(optimization_description) - optimizer = PetrinautOptimizer(model) # type: ignore[arg-type] - request = DisconnectedAfterWorkerStarts(model) - run_id = _run_id(request) - - async def consume() -> None: - async for _frame in optimizer.stream_all( - request, - run_id, - optimizer.n_trials, # type: ignore[arg-type] - ): - pass - model.release.set() - await asyncio.sleep(0.05) - - with caplog.at_level(logging.INFO, logger="pn_optimize"): - asyncio.run(consume()) - - disconnected = next( - record - for record in caplog.records - if getattr(record, "event", None) == "client_disconnected" - ) - assert disconnected.run_id == run_id - - -def test_stream_all_preserves_the_existing_sse_frame_shape( - optimization_description: dict, -) -> None: - model = FakeModel(optimization_description) - optimizer = PetrinautOptimizer(model) # type: ignore[arg-type] - request = ConnectedRequest() - - async def collect() -> list[str]: - return [ - frame - async for frame in optimizer.stream_all( - request, - _run_id(request), - optimizer.n_trials, # type: ignore[arg-type] - ) - ] - - frames = asyncio.run(collect()) - data = [ - json.loads(frame.removeprefix("data: ")) - for frame in frames - if frame.startswith("data: ") - ] - - assert frames[-1] == "event: done\ndata: {}\n\n" - assert len(data) == 3 - assert all( - set(payload) == {"step", "params", "init_state", "metric", "state"} - for payload in data - ) - assert all(payload["init_state"] == {} for payload in data) - assert all( - set(payload["params"]) == {"rate", "count", "enabled"} for payload in data - ) - assert model.closed is True - assert model.close_calls == [True] - - -def test_stream_best_preserves_the_existing_sse_frame_shape( - optimization_description: dict, -) -> None: - model = FakeModel(optimization_description) - optimizer = PetrinautOptimizer(model) # type: ignore[arg-type] - request = ConnectedRequest() - - async def collect() -> list[str]: - return [ - frame - async for frame in optimizer.stream_best( - request, - _run_id(request), - optimizer.n_trials, # type: ignore[arg-type] - ) - ] - - frames = asyncio.run(collect()) - data = [ - json.loads(frame.removeprefix("data: ")) - for frame in frames - if frame.startswith("data: ") - ] - - assert frames[-1] == "event: done\ndata: {}\n\n" - assert len(data) == 3 - assert all(payload["state"] == "COMPLETE" for payload in data) - assert all(payload["init_state"] == {} for payload in data) - assert model.closed is True - assert model.close_calls == [True] - - -def test_stream_sends_comment_heartbeats_while_a_trial_is_running( - optimization_description: dict, - monkeypatch: pytest.MonkeyPatch, -) -> None: - optimization_description["study"]["trials"] = 1 - monkeypatch.setattr(petrinaut_optimizer, "SSE_HEARTBEAT_SECONDS", 0.01) - model = SlowModel(optimization_description) - optimizer = PetrinautOptimizer(model) # type: ignore[arg-type] - request = ConnectedRequest() - - async def collect() -> list[str]: - return [ - frame - async for frame in optimizer.stream_all( - request, - _run_id(request), - optimizer.n_trials, # type: ignore[arg-type] - ) - ] - - frames = asyncio.run(collect()) - - assert ": heartbeat\n\n" in frames - assert frames[-1] == "event: done\ndata: {}\n\n" - - -@pytest.mark.parametrize("stream_name", ["stream_all", "stream_best"]) -def test_stream_error_is_terminal_and_is_not_followed_by_done( - optimization_description: dict, - stream_name: str, - caplog: pytest.LogCaptureFixture, -) -> None: - model = FailingModel( - optimization_description, PetrinautClientError("transport failed") - ) - optimizer = PetrinautOptimizer(model) # type: ignore[arg-type] - request = ConnectedRequest() - run_id = _run_id(request) - - async def collect() -> list[str]: - stream = getattr(optimizer, stream_name) - return [ - frame - async for frame in stream( - request, - run_id, - optimizer.n_trials, # type: ignore[arg-type] - ) - ] - - with caplog.at_level(logging.WARNING, logger="pn_optimize"): - frames = asyncio.run(collect()) - status = request.app.state.statuses.get(run_id) - - assert any( - json.loads(frame.removeprefix("data: ")) - == {"state": "ERROR", "message": "transport failed"} - for frame in frames - if frame.startswith("data: ") - ) - assert "event: done\ndata: {}\n\n" not in frames - assert status is not None - assert status.phase is Phase.error - assert model.closed is True - assert model.close_calls == [False] - failure = next( - record - for record in caplog.records - if getattr(record, "event", None) == "study_failed" - ) - assert failure.run_id == run_id - assert "transport failed" not in failure.getMessage() - assert not hasattr(failure, "detail") - - def _status_app_with_run() -> tuple[FastAPI, str]: app = FastAPI() app.state.statuses = StatusStore() @@ -486,6 +245,39 @@ async def drive() -> str: assert status.phase is Phase.done +def test_pump_events_logs_the_study_lifecycle_with_correlation( + optimization_description: dict, + caplog: pytest.LogCaptureFixture, +) -> None: + model = FakeModel(optimization_description) + optimizer = PetrinautOptimizer(model) # type: ignore[arg-type] + app, run_id = _status_app_with_run() + + async def drive() -> str: + return await optimizer.pump_events( + app, + run_id, + optimizer.n_trials, + on_event=lambda _frame: None, + cancel_event=asyncio.Event(), + correlation={"request_id": "request-s1"}, + ) + + with caplog.at_level(logging.INFO, logger="pn_optimize"): + asyncio.run(drive()) + + events = { + getattr(record, "event", None): record + for record in caplog.records + if record.name == "pn_optimize" + } + for expected in ("study_started", "study_completed"): + record = events[expected] + assert record.run_id == run_id + assert record.request_id == "request-s1" + assert record.trials == optimizer.n_trials + + def test_pump_events_reports_a_study_failure_without_done( optimization_description: dict, ) -> None: @@ -732,39 +524,3 @@ async def drive() -> str: assert all('"state": "ERROR"' not in frame for frame in frames) assert status is not None assert status.phase is Phase.done - - -def test_disconnect_closes_cli_before_a_bounded_worker_join( - optimization_description: dict, - monkeypatch: pytest.MonkeyPatch, -) -> None: - optimization_description["study"]["trials"] = 1 - monkeypatch.setattr(petrinaut_optimizer, "_WORKER_SHUTDOWN_TIMEOUT_SECONDS", 0.01) - model = StubbornModel(optimization_description) - optimizer = PetrinautOptimizer(model) # type: ignore[arg-type] - request = DisconnectedAfterWorkerStarts(model) - run_id = _run_id(request) - - async def collect() -> list[str]: - started_at = time.monotonic() - frames = [ - frame - async for frame in optimizer.stream_all( - request, - run_id, - optimizer.n_trials, # type: ignore[arg-type] - ) - ] - assert time.monotonic() - started_at < 0.5 - model.release.set() - await asyncio.sleep(0.05) - return frames - - frames = asyncio.run(collect()) - status = request.app.state.statuses.get(run_id) - - assert frames == [] - assert model.closed is True - assert model.close_calls == [False] - assert status is not None - assert status.phase is Phase.idle diff --git a/apps/petrinaut-website/src/main/app/optimization-demo/petrinaut-opt-optimization.test.ts b/apps/petrinaut-website/src/main/app/optimization-demo/petrinaut-opt-optimization.test.ts index 5255df8035a..996541b48d3 100644 --- a/apps/petrinaut-website/src/main/app/optimization-demo/petrinaut-opt-optimization.test.ts +++ b/apps/petrinaut-website/src/main/app/optimization-demo/petrinaut-opt-optimization.test.ts @@ -4,29 +4,89 @@ import { createPetrinautOptOptimization } from "./petrinaut-opt-optimization"; import type { PetrinautOptimizationInput } from "@hashintel/petrinaut-core"; +const input = { + objective: { direction: "maximize" }, + study: { trials: 2 }, +} as PetrinautOptimizationInput; + describe("createPetrinautOptOptimization", () => { - it("configures the development proxy endpoint", async () => { - const input = { - objective: { direction: "maximize" }, - study: { trials: 2 }, - } as PetrinautOptimizationInput; - const fetchImpl = vi.fn(async () => - Promise.resolve( - new Response("event: done\ndata: {}\n\n", { + it("creates, attaches to, and cancels runs via the development proxy", async () => { + const fetchImpl = vi.fn(async (_url: string | URL, init?: RequestInit) => { + if (init?.method === "POST") { + return Promise.resolve( + Response.json({ run_id: "run-1" }, { status: 201 }), + ); + } + if (init?.method === "DELETE") { + return Promise.resolve(new Response(null, { status: 204 })); + } + return Promise.resolve( + new Response("id: 1\nevent: done\ndata: {}\n\n", { status: 200, - headers: { "content-type": "text/event-stream" }, + headers: { + "content-type": "text/event-stream", + "x-requested-trials": "2", + }, }), - ), - ); + ); + }); const optimization = createPetrinautOptOptimization(fetchImpl); - for await (const _event of optimization.optimize(input)) { + const { runId } = await optimization.createOptimizationRun(input); + expect(runId).toBe("run-1"); + + for await (const _event of optimization.attachOptimizationRun(runId)) { // Exhaust the stream so the shared client performs the request. } - expect(fetchImpl).toHaveBeenCalledWith( - "/api/petrinaut-opt/optimize/all", - expect.anything(), + await optimization.cancelOptimizationRun(runId); + + const calledUrls = fetchImpl.mock.calls.map(([url]) => url.toString()); + expect(calledUrls).toEqual([ + expect.stringContaining("/api/petrinaut-opt/optimize/runs"), + expect.stringContaining("/api/petrinaut-opt/optimize/runs/run-1/events"), + expect.stringContaining("/api/petrinaut-opt/optimize/runs/run-1"), + ]); + }); + + it("classifies a body that dies mid-stream as a network failure", async () => { + // A dropped connection surfaces as a `TypeError` from the body reader, + // after the response headers already arrived. Classifying that as + // `protocol` would blame the optimizer for a transport problem. + const fetchImpl = vi.fn(() => + Promise.resolve( + new Response( + new ReadableStream({ + start: (controller) => { + controller.enqueue(new TextEncoder().encode(": heartbeat\n\n")); + controller.error(new TypeError("network error")); + }, + }), + { + status: 200, + headers: { "content-type": "text/event-stream" }, + }, + ), + ), ); + const optimization = createPetrinautOptOptimization(fetchImpl); + + const onAttached = vi.fn(); + const error = await (async () => { + try { + for await (const _event of optimization.attachOptimizationRun("run-1", { + onAttached, + })) { + // The stream fails before producing any event. + } + return null; + } catch (caught: unknown) { + return caught; + } + })(); + + expect(onAttached).toHaveBeenCalled(); + expect(error).toBeInstanceOf(TypeError); + expect(error).toHaveProperty("category", "network"); }); }); diff --git a/apps/petrinaut-website/src/main/app/optimization-demo/petrinaut-opt-optimization.ts b/apps/petrinaut-website/src/main/app/optimization-demo/petrinaut-opt-optimization.ts index 43460902c3d..a8b959da96b 100644 --- a/apps/petrinaut-website/src/main/app/optimization-demo/petrinaut-opt-optimization.ts +++ b/apps/petrinaut-website/src/main/app/optimization-demo/petrinaut-opt-optimization.ts @@ -1,25 +1,125 @@ -import { openPetrinautOptimizationStream } from "@local/petrinaut-optimizer-client"; +import { + attachPetrinautOptimizationRunStream, + createPetrinautOptimizerClient, + PetrinautOptimizerHttpError, + petrinautOptimizerHttpErrorFromResponse, +} from "@local/petrinaut-optimizer-client"; import type { PetrinautOptimization, - PetrinautOptimizationInput, + PetrinautOptimizationEvent, } from "@hashintel/petrinaut-core"; import type { PetrinautOptimizerFetch } from "@local/petrinaut-optimizer-client"; -const PETRINAUT_OPTIMIZE_ENDPOINT = "/api/petrinaut-opt/optimize/all"; +/** + * Dev-proxy base for the local Petrinaut Optimizer: `vite.config.ts` rewrites + * `/api/petrinaut-opt/*` to the Python service. Resolved against the current + * document at call time; the client's URL builder keeps the path prefix. + */ +const petrinautOptEndpoint = (): URL => + new URL( + "/api/petrinaut-opt/", + // Tests run without a DOM; the browser always resolves from the page. + typeof location === "undefined" ? "http://localhost/" : location.href, + ); + +/** + * Stamp the duck-typed classification fields Petrinaut's optimization + * provider reads (`category`, `httpStatus`, `retryAfter`) onto the client's + * HTTP error, so e.g. a 404 on re-attaching to an expired run silently drops + * the record instead of surfacing a raw error message. + */ +const classifyHttpError = (error: unknown): unknown => + error instanceof PetrinautOptimizerHttpError + ? Object.assign(error, { + category: "http", + httpStatus: error.status, + ...(error.retryAfter === null + ? {} + : { retryAfter: Number.parseInt(error.retryAfter, 10) }), + }) + : error; + +/** + * Classify a mid-stream failure so the provider reconnects with its cursor + * instead of failing the run on the first dropped connection. Aborts pass + * through untouched. A response body that dies mid-stream rejects the reader + * with a `TypeError`, which is a transport failure rather than a malformed + * frame — the remaining non-abort errors are the decoder's own validation + * failures, which stay `protocol`. + */ +const classifyStreamError = (error: unknown): unknown => + error instanceof Error && error.name !== "AbortError" + ? Object.assign(error, { + category: error instanceof TypeError ? "network" : "protocol", + }) + : error; + +/** + * Classify a request-time failure: HTTP errors keep their status semantics, + * and anything else non-abort (a fetch `TypeError` from a dropped + * connection) is `network` — so an attach that dies before responding + * reconnects with backoff exactly like a mid-stream drop, instead of + * definitively failing a possibly-live run. + */ +const classifyRequestError = (error: unknown): unknown => + error instanceof PetrinautOptimizerHttpError + ? classifyHttpError(error) + : error instanceof Error && error.name !== "AbortError" + ? Object.assign(error, { category: "network" }) + : error; /** Create the local-only Petrinaut capability backed directly by Python. */ export const createPetrinautOptOptimization = ( fetchImpl: PetrinautOptimizerFetch = fetch, -): PetrinautOptimization => ({ - /** Post one manifest and stream its canonical optimization events. */ - async *optimize(input: PetrinautOptimizationInput, options) { - const { events } = await openPetrinautOptimizationStream({ - endpoint: PETRINAUT_OPTIMIZE_ENDPOINT, - fetchImpl, - input, - ...(options?.signal ? { signal: options.signal } : {}), - }); - yield* events; - }, -}); +): PetrinautOptimization => { + const client = createPetrinautOptimizerClient( + petrinautOptEndpoint(), + fetchImpl, + ); + // openapi-fetch names its verb methods in caps; alias them so call sites + // don't read as constructor calls (oxlint's new-cap). + const { DELETE: deleteRun, POST: postRun } = client; + + return { + async createOptimizationRun(input, options) { + const created = await postRun("/optimize/runs", { + body: input, + ...(options?.signal ? { signal: options.signal as AbortSignal } : {}), + }).catch((error: unknown) => { + throw classifyRequestError(error); + }); + if (!created.response.ok || !created.data?.run_id) { + throw classifyHttpError( + await petrinautOptimizerHttpErrorFromResponse(created.response), + ); + } + return { runId: created.data.run_id }; + }, + async *attachOptimizationRun(runId, options) { + let events: AsyncIterable; + try { + ({ events } = await attachPetrinautOptimizationRunStream({ + endpoint: petrinautOptEndpoint(), + fetchImpl, + runId, + ...(options?.cursor === undefined ? {} : { cursor: options.cursor }), + ...(options?.signal ? { signal: options.signal } : {}), + })); + } catch (error) { + throw classifyRequestError(error); + } + options?.onAttached?.(); + try { + yield* events; + } catch (error) { + throw classifyStreamError(error); + } + }, + async cancelOptimizationRun(runId) { + await deleteRun("/optimize/runs/{run_id}", { + params: { path: { run_id: runId } }, + }); + }, + }; +}; diff --git a/libs/@hashintel/petrinaut-core/src/index.ts b/libs/@hashintel/petrinaut-core/src/index.ts index fe9ca5d0889..3e493eb241f 100644 --- a/libs/@hashintel/petrinaut-core/src/index.ts +++ b/libs/@hashintel/petrinaut-core/src/index.ts @@ -79,6 +79,7 @@ export type { PetrinautMutations, } from "./instance"; export { + PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, PETRINAUT_OPTIMIZATION_MAX_SEED, PETRINAUT_OPTIMIZATION_MAX_STEPS_PER_TRIAL, PETRINAUT_OPTIMIZATION_MAX_TOTAL_STEPS, diff --git a/libs/@hashintel/petrinaut-core/src/optimization.test.ts b/libs/@hashintel/petrinaut-core/src/optimization.test.ts index eac2e136ddb..15469d5b52d 100644 --- a/libs/@hashintel/petrinaut-core/src/optimization.test.ts +++ b/libs/@hashintel/petrinaut-core/src/optimization.test.ts @@ -1,6 +1,9 @@ import { describe, expect, it } from "vitest"; -import { petrinautOptimizationManifestSchema } from "./optimization"; +import { + petrinautOptimizationEventSchema, + petrinautOptimizationManifestSchema, +} from "./optimization"; const scenario = { id: "baseline", @@ -361,3 +364,58 @@ describe("petrinautOptimizationManifestSchema", () => { expect(tooMuchTotalWork.success).toBe(false); }); }); + +describe("petrinautOptimizationEventSchema", () => { + const events = [ + { type: "started", requestedTrials: 2 }, + { + type: "trial", + trial: 0, + parameters: { rate: 0.5 }, + objective: 1, + state: "complete", + best: { trial: 0, parameters: { rate: 0.5 }, objective: 1 }, + }, + { + type: "complete", + requestedTrials: 2, + completedTrials: 2, + prunedTrials: 0, + failedTrials: 0, + best: null, + }, + { type: "error", code: "failed", message: "nope", retryable: false }, + ]; + + it("accepts events with and without a sequence number on every variant", () => { + for (const [index, event] of events.entries()) { + const withoutSeq = petrinautOptimizationEventSchema.safeParse(event); + const withSeq = petrinautOptimizationEventSchema.safeParse({ + ...event, + seq: index + 1, + }); + + expect(withoutSeq.success).toBe(true); + expect(withSeq.success).toBe(true); + if (withSeq.success) { + expect(withSeq.data.seq).toBe(index + 1); + } + } + }); + + it("rejects negative or fractional sequence numbers", () => { + const negative = petrinautOptimizationEventSchema.safeParse({ + type: "started", + requestedTrials: 2, + seq: -1, + }); + const fractional = petrinautOptimizationEventSchema.safeParse({ + type: "started", + requestedTrials: 2, + seq: 1.5, + }); + + expect(negative.success).toBe(false); + expect(fractional.success).toBe(false); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/optimization.ts b/libs/@hashintel/petrinaut-core/src/optimization.ts index f599c8f095d..d049ffbebbd 100644 --- a/libs/@hashintel/petrinaut-core/src/optimization.ts +++ b/libs/@hashintel/petrinaut-core/src/optimization.ts @@ -491,11 +491,20 @@ const optimizationBestSchema = z }) .meta({ description: "The best completed trial so far." }); +/** + * Server-authoritative, strictly increasing sequence number attached to each + * event of a detached optimization run. A client resuming a run asks for the + * events with `seq` greater than the last one it applied, and skips any + * replayed event at or below that cursor. Optional so streams from hosts that + * predate detached runs keep validating. + */ +const optimizationEventSeqSchema = z.number().int().nonnegative().optional(); + export const petrinautOptimizationStartedEventSchema = z .strictObject({ type: z.literal("started"), requestedTrials: z.number().int().positive(), - seq: z.number().int().nonnegative().optional(), + seq: optimizationEventSeqSchema, }) .meta({ description: "The optimizer accepted and started the study." }); @@ -507,7 +516,7 @@ export const petrinautOptimizationTrialEventSchema = z objective: z.number().nullable(), state: z.enum(["complete", "pruned", "failed"]), best: optimizationBestSchema.nullable(), - seq: z.number().int().nonnegative().optional(), + seq: optimizationEventSeqSchema, }) .meta({ description: "One completed Optuna trial and the running best." }); @@ -519,17 +528,28 @@ export const petrinautOptimizationCompleteEventSchema = z prunedTrials: z.number().int().nonnegative(), failedTrials: z.number().int().nonnegative(), best: optimizationBestSchema.nullable(), - seq: z.number().int().nonnegative().optional(), + seq: optimizationEventSeqSchema, }) .meta({ description: "The final optimization summary." }); +/** + * The `code` of the terminal error event that reports a cancellation rather + * than a failure. A detached run is cancelled out-of-band — an explicit + * `DELETE`, orphan reaping, or optimizer shutdown — and the stream has no + * event type of its own for that, so it arrives as a non-retryable error. + * Consumers presenting run outcomes should treat this code as "cancelled", + * not "failed". + */ +export const PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE = + "optimization_cancelled"; + export const petrinautOptimizationErrorEventSchema = z .strictObject({ type: z.literal("error"), code: z.string(), message: z.string(), retryable: z.boolean(), - seq: z.number().int().nonnegative().optional(), + seq: optimizationEventSeqSchema, }) .meta({ description: "A terminal optimizer error." }); @@ -577,10 +597,35 @@ export type PetrinautOptimizationTrialEvent = z.infer< typeof petrinautOptimizationTrialEventSchema >; -/** Host-provided optimization capability for Petrinaut. */ +/** + * Host-provided optimization capability for Petrinaut. + * + * A run is detached from any one connection: it is created by id, its event + * stream can be (re-)attached with a `seq` cursor, and it is cancelled + * explicitly — which lets the UI survive connection drops and page reloads. + */ export type PetrinautOptimization = { - optimize( + /** Start a detached run and resolve its server-issued run id. */ + createOptimizationRun( input: PetrinautOptimizationInput, options?: { signal?: AbortSignalLike }, + ): Promise<{ runId: string }>; + /** + * Stream a detached run's events, replaying those with `seq` greater than + * `cursor` (0 replays everything) before tailing live events. The stream + * ends after a terminal `complete`/`error` event. `onAttached` fires once + * the attachment is accepted (the response headers arrived OK), which may + * be long before the first event on a quiet run — UIs use it to report an + * honest connection state while reconnecting. + */ + attachOptimizationRun( + runId: string, + options?: { + cursor?: number; + signal?: AbortSignalLike; + onAttached?: () => void; + }, ): AsyncIterable; + /** Idempotently stop a detached run server-side. */ + cancelOptimizationRun(runId: string): Promise; }; diff --git a/libs/@hashintel/petrinaut/docs/optimization.md b/libs/@hashintel/petrinaut/docs/optimization.md index ce625dde231..ca10f61634a 100644 --- a/libs/@hashintel/petrinaut/docs/optimization.md +++ b/libs/@hashintel/petrinaut/docs/optimization.md @@ -77,6 +77,23 @@ Closing the drawer does not stop the optimization. Use **Cancel** to abort an active run. Completed, cancelled, and failed records can be removed from their result drawer. -For the initial integration, an optimization is tied to its browser connection. -Closing or reloading the page cancels the active request rather than creating a -persistent background job. +If a run fails, the drawer explains what happened — for example, a lost +connection reports how many of the requested trials had completed and includes +a diagnostic identifier for support. Trials received before the failure are +kept, and a **Retry** action starts a fresh run with the same settings. + +## Connection drops and reloads + +An optimization runs on the server, not in your browser tab. If the connection +drops while you watch one, Petrinaut reconnects automatically and resumes from +the last result it received — the status shows **(reconnecting…)** while it +retries, and no trials are lost or double-counted. Only if reconnecting keeps +failing does the run report a connection error, which keeps the received +trials and offers **Retry**. + +Reloading or closing the page is different: the page loses its view of a +still-running optimization. The run itself continues on the server until it +finishes or is cleaned up, and it can block you from starting a new +optimization until then — so use **Cancel** first if you intend to reload and +run something else. **Cancel** is also what actually stops a run: it ends the +optimization on the server, not just your view of it. diff --git a/libs/@hashintel/petrinaut/src/react/index.ts b/libs/@hashintel/petrinaut/src/react/index.ts index e0dd8302fc6..67ebe507cde 100644 --- a/libs/@hashintel/petrinaut/src/react/index.ts +++ b/libs/@hashintel/petrinaut/src/react/index.ts @@ -25,6 +25,7 @@ export { } from "./optimizations/context"; export type { OptimizationBest, + OptimizationConnectionState, OptimizationRecord, OptimizationStatus, OptimizationsContextValue, diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/context.ts b/libs/@hashintel/petrinaut/src/react/optimizations/context.ts index 538ab5b6b0d..ee0474595a5 100644 --- a/libs/@hashintel/petrinaut/src/react/optimizations/context.ts +++ b/libs/@hashintel/petrinaut/src/react/optimizations/context.ts @@ -13,6 +13,28 @@ export type OptimizationStatus = | "error" | "cancelled"; +/** How an optimization transport failure was classified. */ +export type OptimizationErrorCategory = + | "network" + | "http" + | "protocol" + | "aborted"; + +/** Correlation ids for tracing a failure to the NodeAPI/optimizer logs. */ +export type OptimizationErrorDiagnostics = { + hashRequestId: string | null; + optimizationRunId: string | null; + httpStatus: number | null; +}; + +/** + * Live transport state of a detached run's event stream. `streaming` while + * events are flowing; `reconnecting` while a dropped connection is being + * re-established with backoff. `null` for legacy single-connection runs and + * once a run reaches a terminal status. + */ +export type OptimizationConnectionState = "streaming" | "reconnecting"; + export type OptimizationBest = NonNullable< Extract["best"] >; @@ -23,6 +45,20 @@ export type OptimizationRecord = { createdAt: number; status: OptimizationStatus; error: string | null; + /** Set when a transport failure was classified; null otherwise. */ + errorCategory: OptimizationErrorCategory | null; + /** Correlation ids for a classified failure, for the diagnostic UI. */ + errorDiagnostics: OptimizationErrorDiagnostics | null; + /** Server-issued id of a detached run; null for legacy streaming runs. */ + runId: string | null; + /** + * Highest server-issued event sequence number applied to this record. A + * reconnect resumes the event stream from this cursor, and replayed events + * at or below it are skipped so trials are never double-counted. + */ + lastSeq: number; + /** Transport state of a detached run's event stream; null otherwise. */ + connectionState: OptimizationConnectionState | null; requestedTrials: number; completedTrials: number; prunedTrials: number; @@ -47,6 +83,11 @@ export type OptimizationsContextValue = { createOptimization: (input: PetrinautOptimizationInput) => Promise; cancelOptimization: (optimizationId: string) => void; removeOptimization: (optimizationId: string) => void; + /** + * Start a fresh optimization from a prior one's input (e.g. after a + * transport failure). Returns the new id, or null if the record is gone. + */ + retryOptimization: (optimizationId: string) => Promise; }; const DEFAULT_CONTEXT_VALUE: OptimizationsContextValue = { @@ -58,6 +99,7 @@ const DEFAULT_CONTEXT_VALUE: OptimizationsContextValue = { Promise.reject(new Error("Optimization is unavailable")), cancelOptimization: () => {}, removeOptimization: () => {}, + retryOptimization: () => Promise.resolve(null), }; export const OptimizationsContext = createContext( diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/provider.test.tsx b/libs/@hashintel/petrinaut/src/react/optimizations/provider.test.tsx index efeb86edd01..44810f6e935 100644 --- a/libs/@hashintel/petrinaut/src/react/optimizations/provider.test.tsx +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider.test.tsx @@ -2,10 +2,11 @@ * @vitest-environment jsdom */ import { act, cleanup, render, waitFor } from "@testing-library/react"; -import { use } from "react"; -import { afterEach, describe, expect, it } from "vitest"; +import { StrictMode, use } from "react"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { + PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, petrinautOptimizationInputSchema, type PetrinautOptimization, } from "@hashintel/petrinaut-core"; @@ -94,50 +95,145 @@ function renderProvider(capability: PetrinautOptimization) { }; } -afterEach(cleanup); +beforeEach(() => { + sessionStorage.clear(); +}); + +afterEach(() => { + cleanup(); + vi.useRealTimers(); +}); + +/** + * A trial event as a detached attachment delivers it — `best: null`, since + * the service no longer knows the objective direction after the creating + * request ends; the provider computes the running best itself. `overrides` + * typically sets `seq` and `objective`. + */ +const trialEvent = ( + trial: number, + overrides: Record = {}, +) => ({ + type: "trial" as const, + trial, + parameters: { infected_ratio: 0.01 * (trial + 1) }, + objective: 0.4 - trial * 0.1, + state: "complete" as const, + best: null, + ...overrides, +}); + +/** + * NodeAPI's per-attachment timeout line: terminal for the attachment window, + * not for the run, so the provider must reconnect. Carries no `seq`. + */ +const retryableErrorEvent = { + type: "error" as const, + code: "optimization_timeout", + message: "The optimization attachment timed out", + retryable: true, +}; + +class FakeClassifiedError extends Error { + category: string; + hashRequestId: string | null; + optimizationRunId: string | null; + httpStatus: number | null; + retryAfter: number | null; + + constructor( + message: string, + options: { + category: string; + hashRequestId?: string | null; + optimizationRunId?: string | null; + httpStatus?: number | null; + retryAfter?: number | null; + }, + ) { + super(message); + this.category = options.category; + this.hashRequestId = options.hashRequestId ?? null; + this.optimizationRunId = options.optimizationRunId ?? null; + this.httpStatus = options.httpStatus ?? null; + this.retryAfter = options.retryAfter ?? null; + } +} describe("OptimizationsProvider", () => { - it("collects streamed trials and the final best result", async () => { + it("retries a failed optimization from its original input", async () => { + let call = 0; const capability: PetrinautOptimization = { - async *optimize(request) { - yield { type: "started", requestedTrials: 2 }; - yield { - type: "trial", - trial: 0, - parameters: { infected_ratio: 0.01 }, - objective: 0.4, - state: "complete", - best: { - trial: 0, - parameters: { infected_ratio: 0.01 }, - objective: 0.4, - }, - }; + createOptimizationRun: () => { + call += 1; + return call === 1 + ? Promise.reject( + new FakeClassifiedError("connection interrupted", { + category: "network", + }), + ) + : Promise.resolve({ runId: `run-retry-${call}` }); + }, + async *attachOptimizationRun() { yield { - type: "trial", - trial: 1, - parameters: { infected_ratio: 0.02 }, - objective: 0.2, - state: "complete", - best: { - trial: 1, - parameters: { infected_ratio: 0.02 }, - objective: 0.2, - }, + type: "complete", + requestedTrials: 2, + completedTrials: 0, + prunedTrials: 0, + failedTrials: 0, + best: null, + seq: 1, }; + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + const getValue = renderProvider(capability); + let failedId = ""; + + await act(async () => { + failedId = await getValue().createOptimization(input); + }); + await waitFor(() => + expect(getValue().optimizations[0]?.status).toBe("error"), + ); + + let retriedId: string | null = null; + await act(async () => { + retriedId = await getValue().retryOptimization(failedId); + }); + + expect(retriedId).not.toBeNull(); + expect(retriedId).not.toBe(failedId); + await waitFor(() => + expect( + getValue().optimizations.find((o) => o.id === retriedId)?.status, + ).toBe("complete"), + ); + // The retry reuses the failed run's input, so the failed record remains. + expect(getValue().optimizations).toHaveLength(2); + }); + + it("runs detached create + attach when the capability supports it", async () => { + const cursors: number[] = []; + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "run-1" }), + // Attachments emit no `started` event: the first line is a trial (or + // terminal) event. + async *attachOptimizationRun(_runId, options) { + cursors.push(options?.cursor ?? -1); + yield trialEvent(0, { seq: 1 }); + yield trialEvent(1, { seq: 2 }); yield { type: "complete", - requestedTrials: request.study.trials, + requestedTrials: 2, completedTrials: 2, prunedTrials: 0, failedTrials: 0, - best: { - trial: 1, - parameters: { infected_ratio: 0.02 }, - objective: 0.2, - }, + best: null, + seq: 3, }; }, + cancelOptimizationRun: () => Promise.resolve(), }; const getValue = renderProvider(capability); @@ -149,23 +245,229 @@ describe("OptimizationsProvider", () => { expect(getValue().optimizations[0]?.status).toBe("complete"), ); const optimization = getValue().optimizations[0]!; + expect(cursors).toEqual([0]); + expect(optimization.runId).toBe("run-1"); + expect(optimization.lastSeq).toBe(3); expect(optimization.trials).toHaveLength(2); expect(optimization.completedTrials).toBe(2); + // The objective is minimized and no event carried `best`, so the + // provider computed the running best itself. expect(optimization.best).toEqual({ trial: 1, - parameters: { infected_ratio: 0.02 }, + parameters: trialEvent(1).parameters, + objective: trialEvent(1).objective, + }); + }); + + it("reports a created run as running before any event arrives", async () => { + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "run-quiet" }), + // A quiet run: the attachment is accepted but no event arrives for a + // long time (attachments emit no `started` event by design). + // eslint-disable-next-line require-yield -- the run stays quiet until aborted + async *attachOptimizationRun(_runId, options) { + options?.onAttached?.(); + await new Promise((resolve) => { + options?.signal?.addEventListener("abort", resolve, { once: true }); + }); + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + const getValue = renderProvider(capability); + + await act(async () => { + await getValue().createOptimization(input); + }); + + await waitFor(() => + expect(getValue().optimizations[0]?.status).toBe("running"), + ); + expect(getValue().optimizations[0]?.connectionState).toBe("streaming"); + expect(getValue().optimizations[0]?.trials).toHaveLength(0); + }); + + it("reconnects after a network failure, resuming from the last applied seq without duplicating trials or clobbering totals", async () => { + vi.useFakeTimers(); + const cursors: number[] = []; + let attachCalls = 0; + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "run-2" }), + async *attachOptimizationRun(_runId, options) { + attachCalls += 1; + cursors.push(options?.cursor ?? -1); + if (attachCalls === 1) { + yield trialEvent(0, { seq: 1, objective: 0.4 }); + throw new FakeClassifiedError("connection interrupted", { + category: "network", + }); + } + // Overlapping replay: the run must skip the already-applied seq 1. + yield trialEvent(0, { seq: 1, objective: 0.4 }); + // A better post-reconnect objective updates the running best... + yield trialEvent(1, { seq: 2, objective: 0.2 }); + // ...and a worse one does not (the objective is minimized). + yield trialEvent(2, { seq: 3, objective: 0.5 }); + yield { + type: "complete", + requestedTrials: 3, + // Attachment summaries are since-cursor, not run totals: the + // provider must keep its own accumulated counters. + completedTrials: 2, + prunedTrials: 0, + failedTrials: 0, + best: null, + seq: 4, + }; + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + const getValue = renderProvider(capability); + + await act(async () => { + await getValue().createOptimization(input); + }); + // Flush the create + first (failing) attachment. + await act(async () => { + await vi.advanceTimersByTimeAsync(0); + }); + + const interrupted = getValue().optimizations[0]!; + expect(interrupted.status).toBe("running"); + expect(interrupted.connectionState).toBe("reconnecting"); + expect(interrupted.trials).toHaveLength(1); + expect(interrupted.best?.trial).toBe(0); + + // The first backoff delay elapses and the second attachment completes. + await act(async () => { + await vi.advanceTimersByTimeAsync(1_000); + }); + + const optimization = getValue().optimizations[0]!; + expect(cursors).toEqual([0, 1]); + expect(optimization.status).toBe("complete"); + expect(optimization.connectionState).toBeNull(); + // The replayed trial at seq 1 was deduplicated. + expect(optimization.trials).toHaveLength(3); + expect(optimization.trials.map((trial) => trial.trial)).toEqual([0, 1, 2]); + // All trials applied across both attachments, not the since-cursor 2. + expect(optimization.completedTrials).toBe(3); + expect(optimization.requestedTrials).toBe(3); + // The running best crossed the reconnect: trial 1 (0.2) beat trial 0 + // (0.4) and survived trial 2 (0.5). + expect(optimization.best).toEqual({ + trial: 1, + parameters: trialEvent(1).parameters, objective: 0.2, }); + expect(optimization.error).toBeNull(); + }); + + it("surfaces the classified failure after repeated reconnects fail and cancels the orphaned run", async () => { + vi.useFakeTimers(); + let attachCalls = 0; + const cancelledRunIds: string[] = []; + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "run-3" }), + // eslint-disable-next-line require-yield -- every attachment fails before yielding + async *attachOptimizationRun() { + attachCalls += 1; + throw new FakeClassifiedError("connection interrupted", { + category: "network", + optimizationRunId: "run-3", + }); + }, + cancelOptimizationRun: (runId) => { + cancelledRunIds.push(runId); + return Promise.resolve(); + }, + }; + const getValue = renderProvider(capability); + + await act(async () => { + await getValue().createOptimization(input); + }); + // Walk through every backoff delay (1s, 2s, 4s, ... capped at 30s) until + // the 8th consecutive failure stops the reconnection loop. + for (const delayMs of [ + 1_000, 2_000, 4_000, 8_000, 16_000, 30_000, 30_000, + ]) { + await act(async () => { + await vi.advanceTimersByTimeAsync(delayMs); + }); + } + + const optimization = getValue().optimizations[0]!; + expect(attachCalls).toBe(8); + expect(optimization.status).toBe("error"); + expect(optimization.connectionState).toBeNull(); + expect(optimization.errorCategory).toBe("network"); + expect(optimization.error).toBe( + "Connection to the optimization service was interrupted after 0 of 2 trials. Retry the optimization. (diagnostic id: run-3)", + ); + expect(optimization.errorDiagnostics).toEqual({ + hashRequestId: null, + optimizationRunId: "run-3", + httpStatus: null, + }); + // The possibly-live run was cancelled so the account's single-flight + // frees up. Its stored entry survives on purpose: a cancel's resolution + // does not prove the server acted (some hosts fire-and-forget), so the + // next reload's re-attach settles the run's true fate instead. + expect(cancelledRunIds).toEqual(["run-3"]); + expect( + sessionStorage.getItem("petrinaut:active-optimization-runs"), + ).toContain("run-3"); }); - it("aborts and marks an active optimization as cancelled", async () => { + it("lets Remove cancel a possibly-live run after a terminal error", async () => { + const cancelledRunIds: string[] = []; const capability: PetrinautOptimization = { - async *optimize(_request, options) { - yield { type: "started", requestedTrials: 2 }; + createOptimizationRun: () => Promise.resolve({ runId: "run-10" }), + // eslint-disable-next-line require-yield -- the attachment is rejected outright + async *attachOptimizationRun() { + throw new FakeClassifiedError("Optimization run not found", { + category: "http", + httpStatus: 404, + }); + }, + cancelOptimizationRun: (runId) => { + cancelledRunIds.push(runId); + return Promise.resolve(); + }, + }; + const getValue = renderProvider(capability); + let optimizationId = ""; + + await act(async () => { + optimizationId = await getValue().createOptimization(input); + }); + await waitFor(() => + expect(getValue().optimizations[0]?.status).toBe("error"), + ); + // The attach loop has ended (its live-loop map entry is gone); Remove + // must still find the run id on the record itself. + act(() => getValue().removeOptimization(optimizationId)); + + expect(cancelledRunIds.at(-1)).toBe("run-10"); + // Once via the give-up path, once via the explicit Remove. + expect(cancelledRunIds).toHaveLength(2); + expect(getValue().optimizations).toHaveLength(0); + }); + + it("cancels a detached run server-side and aborts its attachment", async () => { + const cancelledRunIds: string[] = []; + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "run-4" }), + async *attachOptimizationRun(_runId, options) { + yield trialEvent(0, { seq: 1 }); await new Promise((resolve) => { options?.signal?.addEventListener("abort", resolve, { once: true }); }); }, + cancelOptimizationRun: (runId) => { + cancelledRunIds.push(runId); + return Promise.resolve(); + }, }; const getValue = renderProvider(capability); let optimizationId = ""; @@ -179,6 +481,407 @@ describe("OptimizationsProvider", () => { act(() => getValue().cancelOptimization(optimizationId)); + expect(cancelledRunIds).toEqual(["run-4"]); expect(getValue().optimizations[0]?.status).toBe("cancelled"); }); + + it("re-attaches to stored runs after a reload, rebuilding from a full replay", async () => { + sessionStorage.setItem( + "petrinaut:active-optimization-runs", + JSON.stringify({ "run-5": { input, createdAt: 123 } }), + ); + const cursors: number[] = []; + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "unused" }), + async *attachOptimizationRun(_runId, options) { + cursors.push(options?.cursor ?? -1); + yield trialEvent(0, { seq: 1 }); + yield { + type: "complete", + requestedTrials: 2, + completedTrials: 1, + prunedTrials: 0, + failedTrials: 0, + best: null, + seq: 2, + }; + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + const getValue = renderProvider(capability); + + await waitFor(() => + expect(getValue().optimizations[0]?.status).toBe("complete"), + ); + const optimization = getValue().optimizations[0]!; + expect(cursors).toEqual([0]); + expect(optimization.runId).toBe("run-5"); + expect(optimization.createdAt).toBe(123); + expect(optimization.trials).toHaveLength(1); + expect(optimization.completedTrials).toBe(1); + // The best was rebuilt locally from the replayed trial. + expect(optimization.best?.trial).toBe(0); + // The settled run was forgotten so the next reload doesn't re-attach. + expect(sessionStorage.getItem("petrinaut:active-optimization-runs")).toBe( + "{}", + ); + }); + + it("settles a replayed cancellation as cancelled rather than failed", async () => { + // The give-up path cancels a possibly-live run and deliberately keeps its + // stored entry, expecting the next reload to settle it. That replay must + // report Cancelled — not a failed run offering Retry. + sessionStorage.setItem( + "petrinaut:active-optimization-runs", + JSON.stringify({ "run-6": { input, createdAt: 123 } }), + ); + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "unused" }), + async *attachOptimizationRun() { + yield trialEvent(0, { seq: 1 }); + yield { + type: "error", + code: PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, + message: "The optimization was cancelled", + retryable: false, + seq: 2, + }; + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + const getValue = renderProvider(capability); + + await waitFor(() => + expect(getValue().optimizations[0]?.status).toBe("cancelled"), + ); + const optimization = getValue().optimizations[0]!; + expect(optimization.error).toBeNull(); + expect(optimization.errorCategory).toBeNull(); + // The trial applied before the cancellation is still part of the record. + expect(optimization.trials).toHaveLength(1); + expect(sessionStorage.getItem("petrinaut:active-optimization-runs")).toBe( + "{}", + ); + }); + + it("silently drops a stored run the service no longer knows", async () => { + sessionStorage.setItem( + "petrinaut:active-optimization-runs", + JSON.stringify({ "run-6": { input, createdAt: 123 } }), + ); + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "unused" }), + // eslint-disable-next-line require-yield -- the run is gone server-side + async *attachOptimizationRun() { + throw new FakeClassifiedError("Run not found", { + category: "http", + httpStatus: 404, + }); + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + const getValue = renderProvider(capability); + + await waitFor(() => expect(getValue().optimizations).toHaveLength(0)); + expect(sessionStorage.getItem("petrinaut:active-optimization-runs")).toBe( + "{}", + ); + }); + + it("treats a retryable NodeAPI error event as a dropped connection and reconnects", async () => { + vi.useFakeTimers(); + const cursors: number[] = []; + let attachCalls = 0; + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "run-7" }), + async *attachOptimizationRun(_runId, options) { + attachCalls += 1; + cursors.push(options?.cursor ?? -1); + if (attachCalls === 1) { + yield trialEvent(0, { seq: 1 }); + // NodeAPI's attachment window died; the run itself continues. + yield retryableErrorEvent; + return; + } + yield trialEvent(1, { seq: 2 }); + yield { + type: "complete", + requestedTrials: 2, + completedTrials: 1, + prunedTrials: 0, + failedTrials: 0, + best: null, + seq: 3, + }; + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + const getValue = renderProvider(capability); + + await act(async () => { + await getValue().createOptimization(input); + }); + await act(async () => { + await vi.advanceTimersByTimeAsync(0); + }); + + expect(getValue().optimizations[0]?.connectionState).toBe("reconnecting"); + + await act(async () => { + await vi.advanceTimersByTimeAsync(1_000); + }); + + const optimization = getValue().optimizations[0]!; + expect(cursors).toEqual([0, 1]); + expect(optimization.status).toBe("complete"); + expect(optimization.trials).toHaveLength(2); + expect(optimization.completedTrials).toBe(2); + expect(optimization.error).toBeNull(); + }); + + it("surfaces NodeAPI's terminal message after retryable error events exhaust the reconnect cap", async () => { + vi.useFakeTimers(); + let attachCalls = 0; + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "run-8" }), + async *attachOptimizationRun() { + attachCalls += 1; + // Every attachment window dies without yielding any progress. + yield retryableErrorEvent; + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + const getValue = renderProvider(capability); + + await act(async () => { + await getValue().createOptimization(input); + }); + for (const delayMs of [ + 1_000, 2_000, 4_000, 8_000, 16_000, 30_000, 30_000, + ]) { + await act(async () => { + await vi.advanceTimersByTimeAsync(delayMs); + }); + } + + const optimization = getValue().optimizations[0]!; + expect(attachCalls).toBe(8); + expect(optimization.status).toBe("error"); + expect(optimization.connectionState).toBeNull(); + expect(optimization.error).toBe("The optimization attachment timed out"); + }); + + it("surfaces a mid-run 404 as a classified error without retrying", async () => { + vi.useFakeTimers(); + let attachCalls = 0; + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "run-9" }), + async *attachOptimizationRun() { + attachCalls += 1; + if (attachCalls === 1) { + yield trialEvent(0, { seq: 1 }); + throw new FakeClassifiedError("connection interrupted", { + category: "network", + }); + } + // The run is gone by the time the reconnect lands (e.g. NodeAPI + // dropped ownership after forwarding the terminal event elsewhere). + throw new FakeClassifiedError("Optimization run not found", { + category: "http", + httpStatus: 404, + }); + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + const getValue = renderProvider(capability); + + await act(async () => { + await getValue().createOptimization(input); + }); + await act(async () => { + await vi.advanceTimersByTimeAsync(1_000); + }); + + const optimization = getValue().optimizations[0]!; + expect(optimization.status).toBe("error"); + expect(optimization.errorCategory).toBe("http"); + expect(optimization.error).toContain("(status 404)"); + + // No further reconnects are scheduled for the definitive 404. + await act(async () => { + await vi.advanceTimersByTimeAsync(120_000); + }); + expect(attachCalls).toBe(2); + }); + + it("reconnects through a transient gateway error", async () => { + vi.useFakeTimers(); + let attachCalls = 0; + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "run-11" }), + async *attachOptimizationRun() { + attachCalls += 1; + if (attachCalls === 1) { + // NodeAPI is restarting or deploying. + throw new FakeClassifiedError("Bad gateway", { + category: "http", + httpStatus: 503, + }); + } + yield trialEvent(0, { seq: 1 }); + yield { + type: "complete", + requestedTrials: 2, + completedTrials: 1, + prunedTrials: 0, + failedTrials: 0, + best: null, + seq: 2, + }; + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + const getValue = renderProvider(capability); + + await act(async () => { + await getValue().createOptimization(input); + }); + await act(async () => { + await vi.advanceTimersByTimeAsync(1_000); + }); + + const optimization = getValue().optimizations[0]!; + expect(attachCalls).toBe(2); + expect(optimization.status).toBe("complete"); + expect(optimization.trials).toHaveLength(1); + expect(optimization.error).toBeNull(); + }); + + it("explains a busy service when creation is rejected with 429", async () => { + const capability: PetrinautOptimization = { + createOptimizationRun: () => + Promise.reject( + new FakeClassifiedError("Too many optimization requests", { + category: "http", + httpStatus: 429, + retryAfter: 30, + }), + ), + // eslint-disable-next-line require-yield -- creation is rejected before any attachment + async *attachOptimizationRun() { + throw new Error("Nothing to attach to"); + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + const getValue = renderProvider(capability); + + await act(async () => { + await getValue().createOptimization(input); + }); + + await waitFor(() => + expect(getValue().optimizations[0]?.status).toBe("error"), + ); + expect(getValue().optimizations[0]?.error).toBe( + "The optimization service is busy — another optimization may already be running for your account. Try again in ~30s.", + ); + }); + + it("does not duplicate restored runs under StrictMode double-mounting", async () => { + sessionStorage.setItem( + "petrinaut:active-optimization-runs", + JSON.stringify({ "run-12": { input, createdAt: 123 } }), + ); + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "unused" }), + async *attachOptimizationRun() { + yield trialEvent(0, { seq: 1 }); + yield { + type: "complete", + requestedTrials: 2, + completedTrials: 1, + prunedTrials: 0, + failedTrials: 0, + best: null, + seq: 2, + }; + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + + let latest: OptimizationsContextValue | null = null; + render( + + + + { + latest = value; + }} + /> + + + , + ); + const getValue = () => { + if (!latest) { + throw new Error("Optimization context was not captured"); + } + return latest; + }; + + await waitFor(() => + expect(getValue().optimizations[0]?.status).toBe("complete"), + ); + // The double-invoked effect cleaned its first pass up instead of + // re-attaching the same stored run twice. + expect(getValue().optimizations).toHaveLength(1); + expect(getValue().optimizations[0]?.runId).toBe("run-12"); + }); + + it("restores the streaming state as soon as a quiet reattachment is accepted", async () => { + vi.useFakeTimers(); + let attachCalls = 0; + const capability: PetrinautOptimization = { + createOptimizationRun: () => Promise.resolve({ runId: "run-13" }), + async *attachOptimizationRun(_runId, options) { + attachCalls += 1; + if (attachCalls === 1) { + options?.onAttached?.(); + yield trialEvent(0, { seq: 1 }); + throw new FakeClassifiedError("connection interrupted", { + category: "network", + }); + } + // The reattachment is accepted but the run stays quiet: no events. + options?.onAttached?.(); + await new Promise((resolve) => { + options?.signal?.addEventListener("abort", resolve, { once: true }); + }); + }, + cancelOptimizationRun: () => Promise.resolve(), + }; + const getValue = renderProvider(capability); + + await act(async () => { + await getValue().createOptimization(input); + }); + await act(async () => { + await vi.advanceTimersByTimeAsync(0); + }); + + expect(getValue().optimizations[0]?.connectionState).toBe("reconnecting"); + + await act(async () => { + await vi.advanceTimersByTimeAsync(1_000); + }); + + const optimization = getValue().optimizations[0]!; + expect(attachCalls).toBe(2); + // No event has arrived yet, but the accepted attachment already cleared + // the reconnecting indicator. + expect(optimization.connectionState).toBe("streaming"); + expect(optimization.status).toBe("running"); + expect(optimization.error).toBeNull(); + }); }); diff --git a/libs/@hashintel/petrinaut/src/react/optimizations/provider.tsx b/libs/@hashintel/petrinaut/src/react/optimizations/provider.tsx index bab30c69e90..b965d71790a 100644 --- a/libs/@hashintel/petrinaut/src/react/optimizations/provider.tsx +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider.tsx @@ -1,10 +1,19 @@ -import { use, useEffect, useRef, useState } from "react"; +import { use, useCallback, useEffect, useRef, useState } from "react"; -import { petrinautOptimizationInputSchema } from "@hashintel/petrinaut-core"; +import { + PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, + petrinautOptimizationInputSchema, + type PetrinautOptimization, + type PetrinautOptimizationEvent, + type PetrinautOptimizationInput, +} from "@hashintel/petrinaut-core"; import { useBlockWindowClose } from "../hooks/use-block-window-close"; import { PetrinautOptimizationContext } from "../optimization-context"; import { + type OptimizationBest, + type OptimizationErrorCategory, + type OptimizationErrorDiagnostics, isOptimizationActive, type OptimizationRecord, OptimizationsContext, @@ -13,6 +22,133 @@ import { import type { PropsWithChildren } from "react"; +const ERROR_CATEGORIES = new Set([ + "network", + "http", + "protocol", + "aborted", +]); + +/** First reconnect delay after a dropped detached-run event stream. */ +const RECONNECT_BASE_DELAY_MS = 1_000; +/** Ceiling for the exponential reconnect backoff. */ +const RECONNECT_MAX_DELAY_MS = 30_000; +/** + * Consecutive failed attachments (no event received in between) after which + * reconnecting stops and the classified failure is surfaced instead. + */ +const MAX_CONSECUTIVE_RECONNECT_FAILURES = 8; + +/** + * Gateway statuses a re-attach may transiently hit while the service + * restarts or deploys; they reconnect within the same failure cap. Every + * other http status (404 unknown run, other 4xx) is definitive. + */ +const RECONNECTABLE_HTTP_STATUSES = new Set([502, 503, 504]); + +/** Exponential backoff: 1s, 2s, 4s, ... capped at 30s. */ +const reconnectDelayMs = (consecutiveFailures: number): number => + Math.min( + RECONNECT_BASE_DELAY_MS * 2 ** (consecutiveFailures - 1), + RECONNECT_MAX_DELAY_MS, + ); + +/** Resolve after `ms`, or immediately once `signal` aborts. */ +const abortableDelay = (ms: number, signal: AbortSignal): Promise => + new Promise((resolve) => { + if (signal.aborted) { + resolve(); + return; + } + const timer = setTimeout(resolve, ms); + // The listener stays attached when the delay elapses normally: at most a + // handful accumulate per run, and they die with the run's controller. + signal.addEventListener( + "abort", + () => { + clearTimeout(timer); + resolve(); + }, + { once: true }, + ); + }); + +/** + * sessionStorage key recording the detached runs this tab may re-attach to + * after a reload: a JSON object mapping run id to its manifest and creation + * time. Session-scoped on purpose — a run belongs to the tab that started it. + * + * When storage is unavailable (e.g. Petrinaut runs in a sandboxed iframe with + * an opaque origin) every helper degrades to a no-op: reload re-attachment is + * lost, while in-page reconnection keeps working. + */ +const ACTIVE_RUNS_STORAGE_KEY = "petrinaut:active-optimization-runs"; + +type StoredActiveRun = { input: unknown; createdAt: number }; + +const readStoredActiveRuns = (): Record => { + try { + const raw = sessionStorage.getItem(ACTIVE_RUNS_STORAGE_KEY); + if (!raw) { + return {}; + } + const parsed: unknown = JSON.parse(raw); + if ( + typeof parsed !== "object" || + parsed === null || + Array.isArray(parsed) + ) { + return {}; + } + const runs: Record = {}; + for (const [runId, value] of Object.entries(parsed)) { + if (typeof value === "object" && value !== null && "input" in value) { + const createdAt = (value as { createdAt?: unknown }).createdAt; + runs[runId] = { + input: (value as { input: unknown }).input, + createdAt: typeof createdAt === "number" ? createdAt : Date.now(), + }; + } + } + return runs; + } catch { + // Unavailable or corrupted storage; see ACTIVE_RUNS_STORAGE_KEY. + return {}; + } +}; + +const writeStoredActiveRuns = (runs: Record): void => { + try { + sessionStorage.setItem(ACTIVE_RUNS_STORAGE_KEY, JSON.stringify(runs)); + } catch { + // Unavailable storage or exceeded quota; see ACTIVE_RUNS_STORAGE_KEY. + } +}; + +const storeActiveRun = ( + runId: string, + input: PetrinautOptimizationInput, +): void => { + const runs = readStoredActiveRuns(); + runs[runId] = { input, createdAt: Date.now() }; + writeStoredActiveRuns(runs); +}; + +const removeStoredActiveRun = (runId: string): void => { + const runs = readStoredActiveRuns(); + if (runId in runs) { + delete runs[runId]; + writeStoredActiveRuns(runs); + } +}; + +type ClassifiedError = { + category: OptimizationErrorCategory; + /** Seconds from a `Retry-After` header, when the service sent one (429). */ + retryAfter: number | null; + diagnostics: OptimizationErrorDiagnostics; +}; + function isAbortError(error: unknown): boolean { return ( (error instanceof DOMException && error.name === "AbortError") || @@ -20,9 +156,147 @@ function isAbortError(error: unknown): boolean { ); } +/** + * Read the structured fields off a classified transport error without + * depending on the host bridge's class: the error crosses from the app into + * this library, so it is duck-typed rather than matched with `instanceof`. + */ +function classifyError(error: unknown): ClassifiedError | null { + if (typeof error !== "object" || error === null) { + return null; + } + const candidate = error as Record; + if ( + typeof candidate.category !== "string" || + !ERROR_CATEGORIES.has(candidate.category as OptimizationErrorCategory) + ) { + return null; + } + return { + category: candidate.category as OptimizationErrorCategory, + retryAfter: + typeof candidate.retryAfter === "number" ? candidate.retryAfter : null, + diagnostics: { + hashRequestId: + typeof candidate.hashRequestId === "string" + ? candidate.hashRequestId + : null, + optimizationRunId: + typeof candidate.optimizationRunId === "string" + ? candidate.optimizationRunId + : null, + httpStatus: + typeof candidate.httpStatus === "number" ? candidate.httpStatus : null, + }, + }; +} + +/** Build a safe, actionable message from a classified failure. */ +function buildErrorMessage( + classified: ClassifiedError, + progress: { completedTrials: number; requestedTrials: number }, +): string { + const after = `after ${progress.completedTrials} of ${progress.requestedTrials} trials`; + const { httpStatus, optimizationRunId, hashRequestId } = + classified.diagnostics; + const diagnosticId = optimizationRunId ?? hashRequestId; + const diagnostic = diagnosticId ? ` (diagnostic id: ${diagnosticId})` : ""; + + switch (classified.category) { + case "http": + if (httpStatus === 429) { + return `The optimization service is busy — another optimization may already be running for your account.${ + classified.retryAfter === null + ? "" + : ` Try again in ~${classified.retryAfter}s.` + }${diagnostic}`; + } + return `The optimization service rejected the request${ + httpStatus === null ? "" : ` (status ${httpStatus})` + } ${after}. Retry the optimization.${diagnostic}`; + case "protocol": + return `The optimization stream ended unexpectedly ${after}. Retry the optimization.${diagnostic}`; + case "aborted": + return "The optimization was cancelled."; + case "network": + default: + return `Connection to the optimization service was interrupted ${after}. Retry the optimization.${diagnostic}`; + } +} + +/** + * Fold a completed trial into the running best. Attachments deliver + * `best: null` (the service no longer knows the objective direction after + * the creating request ends), so the provider maintains the best itself from + * every trial it applies; `event.best` is still preferred when present. + */ +const computeRunningBest = ( + current: OptimizationRecord, + event: Extract, +): OptimizationBest | null => { + if (event.state !== "complete" || event.objective === null) { + return current.best; + } + const isBetter = + current.best === null || + (current.input.objective.direction === "maximize" + ? event.objective > current.best.objective + : event.objective < current.best.objective); + return isBetter + ? { + trial: event.trial, + parameters: event.parameters, + objective: event.objective, + } + : current.best; +}; + +/** + * A NodeAPI-authored terminal error event with `retryable: true`: the + * per-attachment window died (overall or idle timeout) while the run itself + * may still be live. Thrown inside the attach loop so the shared + * reconnect-with-cursor path handles it like a dropped connection; only if + * reconnecting is exhausted is the event applied as the run's terminal error. + */ +class RetryableRunInterruption extends Error { + readonly event: Extract; + + constructor(event: Extract) { + super(event.message); + this.name = "RetryableRunInterruption"; + this.event = event; + } +} + +const createOptimizationRecord = ( + id: string, + input: PetrinautOptimizationInput, + overrides: Partial = {}, +): OptimizationRecord => ({ + id, + input, + createdAt: Date.now(), + status: "initializing", + error: null, + errorCategory: null, + errorDiagnostics: null, + runId: null, + lastSeq: 0, + connectionState: null, + requestedTrials: input.study.trials, + completedTrials: 0, + prunedTrials: 0, + failedTrials: 0, + trials: [], + best: null, + ...overrides, +}); + export const OptimizationsProvider = ({ children }: PropsWithChildren) => { const capability = use(PetrinautOptimizationContext); const abortControllersRef = useRef(new Map()); + /** Server run ids of active detached runs, keyed by record id. */ + const runIdsRef = useRef(new Map()); const [optimizations, setOptimizations] = useState([]); const [selectedOptimizationId, setSelectedOptimizationId] = useState< string | null @@ -42,18 +316,345 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { }; }, []); - const patchOptimization = ( - optimizationId: string, - updater: (optimization: OptimizationRecord) => OptimizationRecord, - ) => { + const patchOptimization = useCallback( + ( + optimizationId: string, + updater: (optimization: OptimizationRecord) => OptimizationRecord, + ) => { + setOptimizations((current) => + current.map((optimization) => + optimization.id === optimizationId + ? updater(optimization) + : optimization, + ), + ); + }, + [], + ); + + const dropOptimizationRecord = useCallback((optimizationId: string) => { setOptimizations((current) => - current.map((optimization) => - optimization.id === optimizationId - ? updater(optimization) - : optimization, - ), + current.filter((optimization) => optimization.id !== optimizationId), ); - }; + setSelectedOptimizationId((current) => + current === optimizationId ? null : current, + ); + }, []); + + const markOptimizationCancelled = useCallback( + (optimizationId: string) => { + patchOptimization(optimizationId, (current) => ({ + ...current, + status: "cancelled", + error: null, + errorCategory: null, + errorDiagnostics: null, + connectionState: null, + })); + }, + [patchOptimization], + ); + + const markOptimizationFailed = useCallback( + ( + optimizationId: string, + error: unknown, + classified: ClassifiedError | null, + ) => { + patchOptimization(optimizationId, (current) => ({ + ...current, + status: "error", + connectionState: null, + // A classified transport failure yields a safe, actionable message + // and correlation ids; anything else keeps its message. + error: classified + ? buildErrorMessage(classified, current) + : error instanceof Error + ? error.message + : String(error), + errorCategory: classified?.category ?? null, + errorDiagnostics: classified?.diagnostics ?? null, + })); + }, + [patchOptimization], + ); + + /** + * Fold one canonical optimizer event into the record, causing a single + * state update per event. + */ + const applyOptimizationEvent = useCallback( + ( + optimizationId: string, + event: PetrinautOptimizationEvent, + options: { + /** Stream-level fields (resume cursor, connection state). */ + extra?: Partial; + } = {}, + ) => { + const { extra = {} } = options; + switch (event.type) { + case "started": + patchOptimization(optimizationId, (current) => ({ + ...current, + ...extra, + status: "running", + requestedTrials: event.requestedTrials, + })); + break; + case "trial": + patchOptimization(optimizationId, (current) => ({ + ...current, + ...extra, + status: "running", + completedTrials: + current.completedTrials + (event.state === "complete" ? 1 : 0), + prunedTrials: + current.prunedTrials + (event.state === "pruned" ? 1 : 0), + failedTrials: + current.failedTrials + (event.state === "failed" ? 1 : 0), + trials: [...current.trials, event], + best: event.best ?? computeRunningBest(current, event), + })); + break; + case "complete": + patchOptimization(optimizationId, (current) => ({ + ...current, + ...extra, + status: "complete", + connectionState: null, + // The complete event's requested-trial count is the true total, + // but its completed/pruned/failed counts only cover the frames + // this attachment observed (everything past its cursor), so the + // record's own accumulated counters and running best stay + // authoritative. + requestedTrials: event.requestedTrials, + best: event.best ?? current.best, + })); + break; + case "error": + patchOptimization(optimizationId, (current) => ({ + ...current, + ...extra, + connectionState: null, + /** + * A cancellation reaches us as a non-retryable error event — the + * stream has no type of its own for it. It is an outcome, not a + * failure, so settle it exactly as a locally-driven cancel does: + * otherwise re-attaching after a give-up cancel, a reaped orphan, + * or a cancel issued elsewhere shows a failed run offering Retry. + */ + ...(event.code === PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE + ? { + status: "cancelled" as const, + error: null, + errorCategory: null, + errorDiagnostics: null, + } + : { status: "error" as const, error: event.message }), + })); + break; + } + }, + [patchOptimization], + ); + + /** + * Consume a detached run's event stream, re-attaching with exponential + * backoff when the connection drops. Every reconnect resumes from the last + * applied `seq`, and replayed events at or below that cursor are skipped so + * trials are never double-counted. Reconnecting stops after + * {@link MAX_CONSECUTIVE_RECONNECT_FAILURES} attachments in a row that + * failed before yielding an event; the classified failure is surfaced then. + * + * Four kinds of interruption reconnect, all sharing the failure cap: + * `network` failures, `protocol` failures (a proxy tearing an idle + * connection down cleanly surfaces as a `protocol` "stream ended without a + * terminal event"), NodeAPI-authored `retryable: true` error events (its + * per-attachment window died while the run continues), and gateway + * `http` statuses (502/503/504 — NodeAPI restarting or deploying). + * Resuming from the cursor is safe in every case because replayed events + * are deduplicated. Every other `http` failure (404 unknown run, other + * 4xx) is definitive and fails immediately, as do `retryable: false` + * error events. + * + * On every give-up path the run — which may still be live server-side — + * is cancelled fire-and-forget: releasing NodeAPI's per-account ownership + * slot means a follow-up run (e.g. the drawer's Retry) isn't rejected as + * busy for the rest of the ownership TTL. + */ + const runAttachLoop = useCallback( + async ({ + optimizationId, + runId, + attach, + cancel, + abortController, + dropRecordOnNotFound = false, + }: { + optimizationId: string; + runId: string; + attach: PetrinautOptimization["attachOptimizationRun"]; + cancel: PetrinautOptimization["cancelOptimizationRun"]; + abortController: AbortController; + /** + * Silently drop the record when the very first attachment 404s — used + * when re-attaching to a stored run that may have expired server-side. + */ + dropRecordOnNotFound?: boolean; + }): Promise => { + const { signal } = abortController; + // Read through a call so the abort flag is re-checked after each await + // (a plain property read would be control-flow-narrowed to `false`). + const isCancelled = () => signal.aborted; + let lastSeq = 0; + let sawTerminalEvent = false; + let consecutiveFailures = 0; + let receivedAnyEvent = false; + + while (!isCancelled()) { + try { + for await (const event of attach(runId, { + cursor: lastSeq, + signal, + /** + * Restore the honest connection state as soon as the attachment + * is accepted — a quiet run may not produce an event for a long + * time, and "(reconnecting…)" would otherwise stick until one + * arrives. Deliberate trade-off: only received EVENTS reset the + * failure counter, so NodeAPI attachment windows that keep + * dying without yielding progress still exhaust the reconnect + * cap even though each of them attached successfully. + */ + onAttached: () => { + patchOptimization(optimizationId, (current) => ({ + ...current, + connectionState: "streaming", + })); + }, + })) { + if (isCancelled()) { + break; + } + if (typeof event.seq === "number") { + if (event.seq <= lastSeq) { + // A replayed event the record already contains. + continue; + } + lastSeq = event.seq; + } + if (event.type === "error" && event.retryable) { + // NodeAPI closed its attachment window (overall/idle timeout) + // while the run may still be live. Deliberately checked before + // the failure-counter reset: a window that keeps dying without + // yielding progress must still exhaust the cap. + throw new RetryableRunInterruption(event); + } + consecutiveFailures = 0; + receivedAnyEvent = true; + if (event.type === "complete" || event.type === "error") { + sawTerminalEvent = true; + } + applyOptimizationEvent(optimizationId, event, { + extra: { lastSeq, connectionState: "streaming" }, + }); + } + if (isCancelled() && !sawTerminalEvent) { + markOptimizationCancelled(optimizationId); + return; + } + // A normal end implies a terminal event was decoded (the stream + // parser rejects endings without one), so the record is settled. + removeStoredActiveRun(runId); + return; + } catch (error) { + const classified = classifyError(error); + const retryableInterruption = + error instanceof RetryableRunInterruption ? error : null; + if ( + isCancelled() || + isAbortError(error) || + classified?.category === "aborted" + ) { + markOptimizationCancelled(optimizationId); + return; + } + if (sawTerminalEvent) { + // The run already settled; a trailing transport hiccup after the + // terminal event changes nothing. + removeStoredActiveRun(runId); + return; + } + if ( + dropRecordOnNotFound && + !receivedAnyEvent && + classified?.category === "http" && + classified.diagnostics.httpStatus === 404 + ) { + removeStoredActiveRun(runId); + dropOptimizationRecord(optimizationId); + return; + } + consecutiveFailures += 1; + const reconnectable = + retryableInterruption !== null || + classified?.category === "network" || + classified?.category === "protocol" || + (classified?.category === "http" && + classified.diagnostics.httpStatus !== null && + RECONNECTABLE_HTTP_STATUSES.has( + classified.diagnostics.httpStatus, + )); + if ( + reconnectable && + consecutiveFailures < MAX_CONSECUTIVE_RECONNECT_FAILURES + ) { + patchOptimization(optimizationId, (current) => ({ + ...current, + connectionState: "reconnecting", + })); + await abortableDelay(reconnectDelayMs(consecutiveFailures), signal); + if (isCancelled()) { + markOptimizationCancelled(optimizationId); + return; + } + continue; + } + // Give up. The run may still be live server-side; cancelling it + // frees the account's single-flight so a fresh run (e.g. the + // drawer's Retry) isn't rejected as busy. The stored entry is + // deliberately kept — some hosts' cancel resolves before the + // server acted, so resolution proves nothing. The next reload's + // re-attach settles it: a delivered cancel replays the cancelled + // terminal, a reaped run 404s (silently dropped), and a run the + // cancel never reached is recovered live. + void cancel(runId).catch(() => undefined); + if (retryableInterruption) { + // Reconnection is exhausted: NodeAPI's own terminal error event + // (a safe, server-authored message) becomes the run's outcome. + applyOptimizationEvent( + optimizationId, + retryableInterruption.event, + { extra: { lastSeq, connectionState: null } }, + ); + } else { + markOptimizationFailed(optimizationId, error, classified); + } + return; + } + } + // Aborted between attachments (e.g. while waiting to reconnect). The + // stored entry is kept: only an explicit cancel forgets a live run. + markOptimizationCancelled(optimizationId); + }, + [ + applyOptimizationEvent, + dropOptimizationRecord, + markOptimizationCancelled, + markOptimizationFailed, + patchOptimization, + ], + ); const createOptimization: OptimizationsContextValue["createOptimization"] = async (rawInput) => { @@ -64,124 +665,195 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { const input = petrinautOptimizationInputSchema.parse(rawInput); const optimizationId = crypto.randomUUID(); const abortController = new AbortController(); - const optimization: OptimizationRecord = { - id: optimizationId, - input, - createdAt: Date.now(), - status: "initializing", - error: null, - requestedTrials: input.study.trials, - completedTrials: 0, - prunedTrials: 0, - failedTrials: 0, - trials: [], - best: null, - }; abortControllersRef.current.set(optimizationId, abortController); - setOptimizations((current) => [optimization, ...current]); + setOptimizations((current) => [ + createOptimizationRecord(optimizationId, input), + ...current, + ]); setSelectedOptimizationId(optimizationId); - const consumeEvents = async () => { + const consumeRun = async () => { + let runId: string; try { - for await (const event of capability.optimize(input, { + ({ runId } = await capability.createOptimizationRun(input, { signal: abortController.signal, - })) { - if (abortController.signal.aborted) { - break; - } - - switch (event.type) { - case "started": - patchOptimization(optimizationId, (current) => ({ - ...current, - status: "running", - requestedTrials: event.requestedTrials, - })); - break; - case "trial": - patchOptimization(optimizationId, (current) => ({ - ...current, - status: "running", - completedTrials: - current.completedTrials + - (event.state === "complete" ? 1 : 0), - prunedTrials: - current.prunedTrials + (event.state === "pruned" ? 1 : 0), - failedTrials: - current.failedTrials + (event.state === "failed" ? 1 : 0), - trials: [...current.trials, event], - best: event.best ?? current.best, - })); - break; - case "complete": - patchOptimization(optimizationId, (current) => ({ - ...current, - status: "complete", - requestedTrials: event.requestedTrials, - completedTrials: event.completedTrials, - prunedTrials: event.prunedTrials, - failedTrials: event.failedTrials, - best: event.best, - })); - break; - case "error": - patchOptimization(optimizationId, (current) => ({ - ...current, - status: "error", - error: event.message, - })); - break; - } - } - } catch (error) { - patchOptimization(optimizationId, (current) => ({ - ...current, - status: - abortController.signal.aborted || isAbortError(error) - ? "cancelled" - : "error", - error: - abortController.signal.aborted || isAbortError(error) - ? null - : error instanceof Error - ? error.message - : String(error), })); - } finally { - abortControllersRef.current.delete(optimizationId); + } catch (error) { + const classified = classifyError(error); + if ( + abortController.signal.aborted || + isAbortError(error) || + classified?.category === "aborted" + ) { + markOptimizationCancelled(optimizationId); + } else { + markOptimizationFailed(optimizationId, error, classified); + } + return; + } + + if (abortController.signal.aborted) { + // Cancelled while the run was being created: stop it server-side + // too, since the cancel action couldn't know its id yet. + void capability.cancelOptimizationRun(runId).catch(() => undefined); + markOptimizationCancelled(optimizationId); + return; } + + runIdsRef.current.set(optimizationId, runId); + storeActiveRun(runId, input); + patchOptimization(optimizationId, (current) => ({ + ...current, + runId, + // Creation only resolves once the study is running server-side, + // and attachments emit no `started` event — without this a quiet + // run would show "initializing" until its first trial. + status: "running", + connectionState: "streaming", + })); + + await runAttachLoop({ + optimizationId, + runId, + attach: capability.attachOptimizationRun.bind(capability), + cancel: capability.cancelOptimizationRun.bind(capability), + abortController, + }); }; - void consumeEvents(); + void consumeRun().finally(() => { + abortControllersRef.current.delete(optimizationId); + runIdsRef.current.delete(optimizationId); + }); + return optimizationId; }; + /** + * Re-attach to the detached runs a previous document in this tab recorded + * (sessionStorage survives reloads but not new tabs). Each restored run is + * rebuilt from a full replay (cursor 0). + * + * The cleanup aborts the loops and drops the records this invocation + * created, so a re-run (React StrictMode double-invokes effects; a swapped + * capability) rebuilds them cleanly instead of duplicating records. + * Aborting keeps the sessionStorage entries, so the re-run finds them + * again. + */ + useEffect(() => { + if (!capability) { + return; + } + + // Snapshot the (provider-lifetime) maps so the cleanup below operates on + // the same instances it registered into. + const abortControllers = abortControllersRef.current; + const runIds = runIdsRef.current; + + const startedIds: string[] = []; + for (const [runId, storedRun] of Object.entries(readStoredActiveRuns())) { + const parsedInput = petrinautOptimizationInputSchema.safeParse( + storedRun.input, + ); + if (!parsedInput.success) { + removeStoredActiveRun(runId); + continue; + } + + const optimizationId = crypto.randomUUID(); + startedIds.push(optimizationId); + const abortController = new AbortController(); + abortControllers.set(optimizationId, abortController); + runIds.set(optimizationId, runId); + setOptimizations((current) => [ + createOptimizationRecord(optimizationId, parsedInput.data, { + createdAt: storedRun.createdAt, + status: "running", + runId, + connectionState: "streaming", + }), + ...current, + ]); + + void runAttachLoop({ + optimizationId, + runId, + attach: capability.attachOptimizationRun.bind(capability), + cancel: capability.cancelOptimizationRun.bind(capability), + abortController, + dropRecordOnNotFound: true, + }).finally(() => { + abortControllers.delete(optimizationId); + runIds.delete(optimizationId); + }); + } + + return () => { + for (const optimizationId of startedIds) { + abortControllers.get(optimizationId)?.abort(); + abortControllers.delete(optimizationId); + runIds.delete(optimizationId); + } + setOptimizations((current) => + current.filter((optimization) => !startedIds.includes(optimization.id)), + ); + }; + }, [capability, runAttachLoop]); + + /** + * The run id of a detached record: from the live-loop map while its attach + * loop runs, falling back to the record itself once the loop has ended + * (e.g. after a surfaced terminal error, when the run may still be live + * server-side and an explicit cancel/remove must still DELETE it). + */ + const resolveRunId = (optimizationId: string): string | undefined => + runIdsRef.current.get(optimizationId) ?? + optimizations.find((optimization) => optimization.id === optimizationId) + ?.runId ?? + undefined; + const cancelOptimization: OptimizationsContextValue["cancelOptimization"] = ( optimizationId, ) => { + const runId = resolveRunId(optimizationId); + if (runId !== undefined) { + runIdsRef.current.delete(optimizationId); + removeStoredActiveRun(runId); + // Stop the detached run server-side; aborting the local attachment + // below only drops this tab's connection to it. + void capability?.cancelOptimizationRun(runId).catch(() => undefined); + } abortControllersRef.current.get(optimizationId)?.abort(); abortControllersRef.current.delete(optimizationId); - patchOptimization(optimizationId, (current) => ({ - ...current, - status: "cancelled", - error: null, - })); + markOptimizationCancelled(optimizationId); }; const removeOptimization: OptimizationsContextValue["removeOptimization"] = ( optimizationId, ) => { + const runId = resolveRunId(optimizationId); + if (runId !== undefined) { + runIdsRef.current.delete(optimizationId); + removeStoredActiveRun(runId); + void capability?.cancelOptimizationRun(runId).catch(() => undefined); + } abortControllersRef.current.get(optimizationId)?.abort(); abortControllersRef.current.delete(optimizationId); - setOptimizations((current) => - current.filter((optimization) => optimization.id !== optimizationId), - ); - setSelectedOptimizationId((current) => - current === optimizationId ? null : current, - ); + dropOptimizationRecord(optimizationId); }; + const retryOptimization: OptimizationsContextValue["retryOptimization"] = + async (optimizationId) => { + const existing = optimizations.find( + (optimization) => optimization.id === optimizationId, + ); + if (!existing) { + return null; + } + return createOptimization(existing.input); + }; + const selectedOptimization = optimizations.find( (optimization) => optimization.id === selectedOptimizationId, @@ -195,6 +867,7 @@ export const OptimizationsProvider = ({ children }: PropsWithChildren) => { createOptimization, cancelOptimization, removeOptimization, + retryOptimization, }; return {children}; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/create-optimization-drawer.test.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/create-optimization-drawer.test.tsx index 3eb1eafa6f6..7da0d7c0e3d 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/create-optimization-drawer.test.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/create-optimization-drawer.test.tsx @@ -209,6 +209,7 @@ const TestProviders = ({ createOptimization, cancelOptimization: () => {}, removeOptimization: () => {}, + retryOptimization: () => Promise.resolve(null), }; const drawer = ( diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/view-optimization-drawer.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/view-optimization-drawer.tsx index dfda6825847..4c70eb522cd 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/view-optimization-drawer.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/optimizations/view-optimization-drawer.tsx @@ -223,6 +223,9 @@ const OptimizationSummary = ({ Status {formatStatus(optimization.status)} + {optimization.connectionState === "reconnecting" + ? " (reconnecting…)" + : ""}
@@ -283,7 +286,8 @@ export const ViewOptimizationDrawer = ({ onClose: () => void; optimization: OptimizationRecord | undefined; }) => { - const { cancelOptimization, removeOptimization } = use(OptimizationsContext); + const { cancelOptimization, removeOptimization, retryOptimization } = + use(OptimizationsContext); if (!open || !optimization) { return null; @@ -372,6 +376,25 @@ export const ViewOptimizationDrawer = ({ Cancel ) : null} + {optimization.status === "error" ? ( + + ) : null} diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/simulate-view.stories.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/simulate-view.stories.tsx index 1d72abea27b..a485f473eb1 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/simulate-view.stories.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/simulate-view.stories.tsx @@ -7,6 +7,7 @@ import { DEFAULT_PETRINAUT_EXTENSIONS, type PetrinautOptimization, type PetrinautOptimizationEvent, + type PetrinautOptimizationInput, type PetrinautOptimizationParameterBinding, type SDCPN, } from "@hashintel/petrinaut-core"; @@ -177,8 +178,33 @@ const getFakeTrialState = (trial: number, seed: number): FakeTrialState => { return roll < 82 ? "complete" : roll < 94 ? "pruned" : "failed"; }; +/** Inputs of the fake detached runs created in this story session. */ +const fakeRuns = new Map(); +let nextFakeRunId = 1; + const fakeOptimization: PetrinautOptimization = { - async *optimize(input, options) { + createOptimizationRun: (input) => { + const runId = `story-run-${nextFakeRunId++}`; + fakeRuns.set(runId, input); + return Promise.resolve({ runId }); + }, + cancelOptimizationRun: (runId) => { + fakeRuns.delete(runId); + return Promise.resolve(); + }, + async *attachOptimizationRun(runId, options) { + const input = fakeRuns.get(runId); + if (!input) { + // Shaped like a classified transport 404 so the provider silently + // drops records restored from a previous story session. + throw Object.assign(new Error(`Unknown story run ${runId}`), { + category: "http", + httpStatus: 404, + }); + } + options?.onAttached?.(); + + let seq = 0; const requestedTrials = input.study.trials; let completedTrials = 0; let prunedTrials = 0; @@ -187,7 +213,12 @@ const fakeOptimization: PetrinautOptimization = { Extract["best"] > | null = null; - yield { type: "started", requestedTrials }; + const cursor = options?.cursor ?? 0; + + seq += 1; + if (seq > cursor) { + yield { type: "started", requestedTrials, seq }; + } for (let trial = 0; trial < requestedTrials; trial += 1) { await wait(250, options?.signal); @@ -231,16 +262,21 @@ const fakeOptimization: PetrinautOptimization = { failedTrials += 1; } - yield { - type: "trial", - trial, - parameters, - objective, - state, - best, - }; + seq += 1; + if (seq > cursor) { + yield { + type: "trial", + trial, + parameters, + objective, + state, + best, + seq, + }; + } } + seq += 1; yield { type: "complete", requestedTrials, @@ -248,6 +284,7 @@ const fakeOptimization: PetrinautOptimization = { prunedTrials, failedTrials, best, + seq, }; }, }; diff --git a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/simulate-view.test.tsx b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/simulate-view.test.tsx index 26790773140..f03ce8ec97c 100644 --- a/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/simulate-view.test.tsx +++ b/libs/@hashintel/petrinaut/src/ui/views/Editor/panels/SimulateView/simulate-view.test.tsx @@ -38,9 +38,11 @@ vi.mock("./scenarios/scenarios-view", () => ({ })); const capability: PetrinautOptimization = { - async *optimize() { - yield { type: "started", requestedTrials: 1 }; + createOptimizationRun: () => Promise.resolve({ runId: "run-test" }), + async *attachOptimizationRun() { + yield { type: "started", requestedTrials: 1, seq: 1 }; }, + cancelOptimizationRun: () => Promise.resolve(), }; afterEach(cleanup); diff --git a/libs/@local/petrinaut-optimizer-client/src/attach-optimization-run.test.ts b/libs/@local/petrinaut-optimizer-client/src/attach-optimization-run.test.ts index c4b2aaa0bb9..0ff4470a4d3 100644 --- a/libs/@local/petrinaut-optimizer-client/src/attach-optimization-run.test.ts +++ b/libs/@local/petrinaut-optimizer-client/src/attach-optimization-run.test.ts @@ -105,37 +105,6 @@ describe("attachPetrinautOptimizationRunStream", () => { ); }); - it("aggregates a best-so-far when a direction is supplied", async () => { - const { events } = await attachPetrinautOptimizationRunStream({ - endpoint: "http://petrinaut-opt.test", - runId: "run-42", - direction: "minimize", - fetchImpl: async () => - new Response( - 'id: 1\ndata: {"step":0,"params":{"rate":0.8},"init_state":{},"metric":4,"state":"COMPLETE"}\n\n' + - 'id: 2\ndata: {"step":1,"params":{"rate":0.4},"init_state":{},"metric":2,"state":"COMPLETE"}\n\n' + - "id: 3\nevent: done\ndata: {}\n\n", - { - headers: { - "content-type": "text/event-stream", - "x-requested-trials": "3", - }, - }, - ), - }); - - const collected = await collect(events); - expect(collected.at(-1)).toEqual({ - type: "complete", - requestedTrials: 3, - completedTrials: 2, - prunedTrials: 0, - failedTrials: 0, - best: { trial: 1, parameters: { rate: 0.4 }, objective: 2 }, - seq: 3, - }); - }); - it("adapts a cancelled run's terminal frame", async () => { const { events } = await attachPetrinautOptimizationRunStream({ endpoint: "http://petrinaut-opt.test", diff --git a/libs/@local/petrinaut-optimizer-client/src/attach-optimization-run.ts b/libs/@local/petrinaut-optimizer-client/src/attach-optimization-run.ts index 7a09d204e4a..aebbc7173ed 100644 --- a/libs/@local/petrinaut-optimizer-client/src/attach-optimization-run.ts +++ b/libs/@local/petrinaut-optimizer-client/src/attach-optimization-run.ts @@ -1,13 +1,23 @@ import { decodePetrinautOptimizerStream } from "./decode-optimization-stream.js"; -import { petrinautOptimizerHttpErrorFromResponse } from "./optimizer-http.js"; +import { + petrinautOptimizerHttpErrorFromResponse, + petrinautOptimizerUrl, +} from "./optimizer-http.js"; -import type { PetrinautOptimizationStreamHandle } from "./open-optimization-stream.js"; import type { PetrinautOptimizerFetch } from "./optimizer-http.js"; import type { AbortSignalLike, - PetrinautOptimizationInput, + PetrinautOptimizationEvent, } from "@hashintel/petrinaut-core"; +/** One opened optimization event stream plus its upstream correlation id. */ +export type PetrinautOptimizationStreamHandle = { + /** Canonical optimization events decoded from the upstream stream. */ + events: AsyncIterable; + /** The optimizer's `X-Optimization-Run-ID` header, when provided. */ + optimizationRunId: string | null; +}; + /** Configuration for attaching to one detached Petrinaut Optimizer run. */ export type AttachPetrinautOptimizationRunStreamOptions = { /** Base URL (origin) of the Petrinaut Optimizer service. */ @@ -19,17 +29,6 @@ export type AttachPetrinautOptimizationRunStreamOptions = { * live-tailing. Omitted or `0` requests a full replay. */ cursor?: number; - /** - * Whether lower or higher objective values are considered better. - * - * Deliberately optional: an attachment that resumes past a cursor never - * observes the whole study, so a best-so-far computed here would silently - * disagree with the true running best. When omitted (the recommended - * attachment configuration), the decoder skips best-so-far aggregation and - * every trial and complete event carries `best: null` — the consumer keeps - * its own running best across reconnections instead. - */ - direction?: PetrinautOptimizationInput["objective"]["direction"]; /** Fetch implementation supplied by the current runtime or a test. */ fetchImpl?: PetrinautOptimizerFetch; /** Optional maximum UTF-8 size of one upstream SSE event. */ @@ -50,16 +49,15 @@ export type AttachPetrinautOptimizationRunStreamOptions = { * * Buffered frames with seq > cursor are replayed, then new frames are * live-tailed; each adapted event carries the frame's sequence number as - * `seq` so the consumer can re-attach from where it stopped. Unlike the - * legacy study stream, no synthetic `started` event is emitted: the study - * started when the run was created, not when this consumer attached. - * Disconnecting never affects the run itself. + * `seq` so the consumer can re-attach from where it stopped. No synthetic + * `started` event is emitted: the study started when the run was created, + * not when this consumer attached. Disconnecting never affects the run + * itself. */ export const attachPetrinautOptimizationRunStream = async ({ endpoint, runId, cursor, - direction, fetchImpl = fetch, headers, maxEventBytes, @@ -67,9 +65,9 @@ export const attachPetrinautOptimizationRunStream = async ({ requestId, signal, }: AttachPetrinautOptimizationRunStreamOptions): Promise => { - const url = new URL( - `/optimize/runs/${encodeURIComponent(runId)}/events`, + const url = petrinautOptimizerUrl( endpoint, + `optimize/runs/${encodeURIComponent(runId)}/events`, ); if (cursor !== undefined) { url.searchParams.set("cursor", String(cursor)); @@ -104,9 +102,7 @@ export const attachPetrinautOptimizationRunStream = async ({ return { events: decodePetrinautOptimizerStream(response.body, { - emitSyntheticStarted: false, requestedTrials, - ...(direction === undefined ? {} : { direction }), ...(maxEventBytes === undefined ? {} : { maxEventBytes }), ...(onActivity ? { onActivity } : {}), }), diff --git a/libs/@local/petrinaut-optimizer-client/src/decode-optimization-stream.test.ts b/libs/@local/petrinaut-optimizer-client/src/decode-optimization-stream.test.ts index bb7de3387d3..d1056f1fb77 100644 --- a/libs/@local/petrinaut-optimizer-client/src/decode-optimization-stream.test.ts +++ b/libs/@local/petrinaut-optimizer-client/src/decode-optimization-stream.test.ts @@ -20,22 +20,13 @@ const streamChunks = (...chunks: string[]) => { const collect = async ( stream: ReadableStream, options: { - direction?: "maximize" | "minimize" | null; - emitSyntheticStarted?: boolean; maxEventBytes?: number; onActivity?: () => void; } = {}, ) => { const events = []; for await (const event of decodePetrinautOptimizerStream(stream, { - // `null` requests attachment-style decoding without a direction. - ...(options.direction === null - ? {} - : { direction: options.direction ?? "maximize" }), requestedTrials: 2, - ...(options.emitSyntheticStarted === undefined - ? {} - : { emitSyntheticStarted: options.emitSyntheticStarted }), ...(options.maxEventBytes === undefined ? {} : { maxEventBytes: options.maxEventBytes }), @@ -58,14 +49,16 @@ describe("decodePetrinautOptimizerStream", () => { ); expect(events).toEqual([ - { type: "started", requestedTrials: 2 }, { type: "trial", trial: 0, parameters: { workers: 2 }, objective: 10, state: "complete", - best: { trial: 0, parameters: { workers: 2 }, objective: 10 }, + // A consumer may attach past a cursor and so never observe the whole + // study; it retains its own running best, and the decoder never + // aggregates one. + best: null, }, { type: "trial", @@ -73,7 +66,7 @@ describe("decodePetrinautOptimizerStream", () => { parameters: { workers: 3 }, objective: null, state: "pruned", - best: { trial: 0, parameters: { workers: 2 }, objective: 10 }, + best: null, }, { type: "complete", @@ -81,38 +74,17 @@ describe("decodePetrinautOptimizerStream", () => { completedTrials: 1, prunedTrials: 1, failedTrials: 0, - best: { trial: 0, parameters: { workers: 2 }, objective: 10 }, + best: null, }, ]); }); - it("selects the lowest completed objective for minimization", async () => { - const events = await collect( - streamChunks( - 'data: {"step":0,"params":{"rate":0.8},"metric":4,"state":"COMPLETE"}\n\n', - 'data: {"step":1,"params":{"rate":0.4},"metric":2,"state":"COMPLETE"}\n\n', - "event: done\ndata: {}\n\n", - ), - { direction: "minimize" }, - ); - - expect(events.at(-1)).toEqual({ - type: "complete", - requestedTrials: 2, - completedTrials: 2, - prunedTrials: 0, - failedTrials: 0, - best: { trial: 1, parameters: { rate: 0.4 }, objective: 2 }, - }); - }); - it("adapts named and state-based terminal optimizer errors", async () => { await expect( collect( streamChunks('event: error\ndata: {"message":"study failed"}\n\n'), ), ).resolves.toEqual([ - { type: "started", requestedTrials: 2 }, { type: "error", code: "optimization_failed", @@ -125,7 +97,6 @@ describe("decodePetrinautOptimizerStream", () => { streamChunks('data: {"state":"ERROR","message":"scenario failed"}\n\n'), ), ).resolves.toEqual([ - { type: "started", requestedTrials: 2 }, { type: "error", code: "optimization_failed", @@ -176,17 +147,27 @@ describe("decodePetrinautOptimizerStream", () => { ).rejects.toThrow("after a terminal event"); }); - it("cancels upstream when its consumer stops after the started event", async () => { + it("cancels upstream when its consumer stops mid-stream", async () => { const cancel = vi.fn(); - const stream = new ReadableStream({ cancel }); + const encoder = new TextEncoder(); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue( + encoder.encode( + 'data: {"step":0,"params":{},"metric":1,"state":"COMPLETE"}\n\n', + ), + ); + // Deliberately left open: the consumer walks away mid-stream. + }, + cancel, + }); const events = decodePetrinautOptimizerStream(stream, { - direction: "maximize", requestedTrials: 1, })[Symbol.asyncIterator](); - await expect(events.next()).resolves.toEqual({ + await expect(events.next()).resolves.toMatchObject({ done: false, - value: { type: "started", requestedTrials: 1 }, + value: { type: "trial", trial: 0 }, }); await events.return?.(); @@ -208,15 +189,13 @@ describe("decodePetrinautOptimizerStream", () => { ); expect(events).toEqual([ - // The synthetic started event predates upstream bytes, so it has no seq. - { type: "started", requestedTrials: 2 }, { type: "trial", trial: 0, parameters: { workers: 2 }, objective: 10, state: "complete", - best: { trial: 0, parameters: { workers: 2 }, objective: 10 }, + best: null, seq: 1, }, { @@ -225,7 +204,7 @@ describe("decodePetrinautOptimizerStream", () => { completedTrials: 1, prunedTrials: 0, failedTrials: 0, - best: { trial: 0, parameters: { workers: 2 }, objective: 10 }, + best: null, seq: 2, }, ]); @@ -265,7 +244,6 @@ describe("decodePetrinautOptimizerStream", () => { it("adapts a superseded frame to a terminal, non-retryable attachment error", async () => { const events = await collect( streamChunks("event: superseded\ndata: {}\n\n"), - { direction: null, emitSyntheticStarted: false }, ); expect(events).toEqual([ @@ -288,38 +266,4 @@ describe("decodePetrinautOptimizerStream", () => { ), ).rejects.toThrow("after a terminal event"); }); - - it("decodes attachments without a synthetic started event or a best", async () => { - const events = await collect( - streamChunks( - 'id: 3\ndata: {"step":2,"params":{"rate":0.4},"metric":2,"state":"COMPLETE"}\n\n', - "id: 4\nevent: done\ndata: {}\n\n", - ), - { direction: null, emitSyntheticStarted: false }, - ); - - expect(events).toEqual([ - { - type: "trial", - trial: 2, - parameters: { rate: 0.4 }, - objective: 2, - state: "complete", - // Without a direction, best-so-far aggregation is skipped: a replay - // past a cursor cannot know the true running best, so the consumer's - // retained best stays authoritative. - best: null, - seq: 3, - }, - { - type: "complete", - requestedTrials: 2, - completedTrials: 1, - prunedTrials: 0, - failedTrials: 0, - best: null, - seq: 4, - }, - ]); - }); }); diff --git a/libs/@local/petrinaut-optimizer-client/src/decode-optimization-stream.ts b/libs/@local/petrinaut-optimizer-client/src/decode-optimization-stream.ts index 2ff018b4238..a3aae81d412 100644 --- a/libs/@local/petrinaut-optimizer-client/src/decode-optimization-stream.ts +++ b/libs/@local/petrinaut-optimizer-client/src/decode-optimization-stream.ts @@ -1,9 +1,9 @@ import { createParser } from "eventsource-parser"; import { + PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, petrinautOptimizationEventSchema, type PetrinautOptimizationEvent, - type PetrinautOptimizationInput, } from "@hashintel/petrinaut-core"; import type { EventSourceMessage, ParseError } from "eventsource-parser"; @@ -12,47 +12,26 @@ type JsonRecord = Record; type StreamState = { requestedTrials: number; - direction: PetrinautOptimizationInput["objective"]["direction"] | undefined; completedTrials: number; prunedTrials: number; failedTrials: number; - best: Extract["best"]; terminal: boolean; }; /** * Configuration needed to adapt one upstream optimization stream. * - * The decoder has two modes: - * - * - **Study mode** (the default, used by `POST /optimize/all`): the caller - * knows the manifest, passes `direction`, and a synthetic `started` event is - * emitted before any upstream bytes are read. Best-so-far aggregation runs - * client-side from the trials observed on this stream. - * - **Attachment mode** (used by `GET /optimize/runs/{run_id}/events`): the - * caller re-attaches mid-run, so a replay that starts past the cursor never - * represents the whole study. Pass `emitSyntheticStarted: false` and omit - * `direction`: no `started` event is emitted and best-so-far aggregation is - * skipped — every trial and complete event then carries `best: null` (legal - * per the canonical schema) and the consumer, which retains its own running - * best across reconnections, remains the single source of truth for it. + * The decoder serves attachments to detached runs + * (`GET /optimize/runs/{run_id}/events`): a consumer may re-attach mid-run, + * so a replay that starts past the cursor never represents the whole study. + * No best-so-far aggregation happens here — every trial and complete event + * carries `best: null` (legal per the canonical schema) and the consumer, + * which retains its own running best across reconnections, remains the + * single source of truth for it. */ export type DecodePetrinautOptimizerStreamOptions = { - /** - * Whether lower or higher objective values are considered better. - * - * When omitted, best-so-far aggregation is skipped entirely and every - * emitted trial and complete event carries `best: null`. - */ - direction?: PetrinautOptimizationInput["objective"]["direction"]; /** Number of trials requested by the optimization manifest. */ requestedTrials: number; - /** - * Whether to emit the synthetic client-side `started` event before reading - * upstream bytes. Defaults to `true`; attachments to an already-running - * detached run pass `false` because the study started long before them. - */ - emitSyntheticStarted?: boolean; /** Optional UTF-8 byte limit applied to each complete upstream event. */ maxEventBytes?: number; /** Called whenever bytes arrive, including heartbeat-only chunks. */ @@ -82,8 +61,7 @@ const utf8ByteLength = (value: string): number => * Parse an SSE frame's `id:` line into a canonical sequence number. * * Detached-run attachments stamp every frame with `id: ` so consumers - * can resume from a cursor; the legacy study stream sends no ids at all, in - * which case the adapted events simply carry no `seq`. + * can resume from a cursor; frames without an id simply carry no `seq`. * * Only plain bounded decimal ids are accepted (mirroring the cursor NodeAPI * accepts): `Number()` would otherwise coerce empty strings, exponent or hex @@ -223,7 +201,7 @@ const adaptSseEvent = ( return { event: petrinautOptimizationEventSchema.parse({ type: "error", - code: "optimization_cancelled", + code: PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, message: "The optimization was cancelled", retryable: false, ...sequenceField, @@ -239,7 +217,7 @@ const adaptSseEvent = ( completedTrials: state.completedTrials, prunedTrials: state.prunedTrials, failedTrials: state.failedTrials, - best: state.best, + best: null, ...sequenceField, }), state: { ...state, terminal: true }, @@ -266,34 +244,19 @@ const adaptSseEvent = ( } const trial = parseTrial(value); - const best = - state.direction !== undefined && - trial.state === "complete" && - trial.objective !== null && - (state.best === null || - (state.direction === "maximize" - ? trial.objective > state.best.objective - : trial.objective < state.best.objective)) - ? { - trial: trial.trial, - parameters: trial.parameters, - objective: trial.objective, - } - : state.best; const nextState: StreamState = { ...state, completedTrials: state.completedTrials + (trial.state === "complete" ? 1 : 0), prunedTrials: state.prunedTrials + (trial.state === "pruned" ? 1 : 0), failedTrials: state.failedTrials + (trial.state === "failed" ? 1 : 0), - best, }; return { event: petrinautOptimizationEventSchema.parse({ type: "trial", ...trial, - best, + best: null, ...sequenceField, }), state: nextState, @@ -314,11 +277,9 @@ export async function* decodePetrinautOptimizerStream( ): AsyncIterable { let state: StreamState = { requestedTrials: options.requestedTrials, - direction: options.direction, completedTrials: 0, prunedTrials: 0, failedTrials: 0, - best: null, terminal: false, }; const events: EventSourceMessage[] = []; @@ -374,14 +335,6 @@ export async function* decodePetrinautOptimizerStream( }; try { - if (options.emitSyntheticStarted ?? true) { - // Synthesized before any upstream bytes exist, so it never has a seq. - yield petrinautOptimizationEventSchema.parse({ - type: "started", - requestedTrials: state.requestedTrials, - }); - } - let result = await reader.read(); while (!result.done) { options.onActivity?.(); diff --git a/libs/@local/petrinaut-optimizer-client/src/index.ts b/libs/@local/petrinaut-optimizer-client/src/index.ts index fd16e9a1fcf..fea8d04224d 100644 --- a/libs/@local/petrinaut-optimizer-client/src/index.ts +++ b/libs/@local/petrinaut-optimizer-client/src/index.ts @@ -1,16 +1,14 @@ export { attachPetrinautOptimizationRunStream } from "./attach-optimization-run.js"; -export type { AttachPetrinautOptimizationRunStreamOptions } from "./attach-optimization-run.js"; +export type { + AttachPetrinautOptimizationRunStreamOptions, + PetrinautOptimizationStreamHandle, +} from "./attach-optimization-run.js"; export { createPetrinautOptimizerClient, type PetrinautOptimizerClient, } from "./client.js"; export { decodePetrinautOptimizerStream } from "./decode-optimization-stream.js"; export type { DecodePetrinautOptimizerStreamOptions } from "./decode-optimization-stream.js"; -export { openPetrinautOptimizationStream } from "./open-optimization-stream.js"; -export type { - OpenPetrinautOptimizationStreamOptions, - PetrinautOptimizationStreamHandle, -} from "./open-optimization-stream.js"; export { PetrinautOptimizerHttpError, petrinautOptimizerHttpErrorFromResponse, diff --git a/libs/@local/petrinaut-optimizer-client/src/open-optimization-stream.test.ts b/libs/@local/petrinaut-optimizer-client/src/open-optimization-stream.test.ts deleted file mode 100644 index 69a77b9a45f..00000000000 --- a/libs/@local/petrinaut-optimizer-client/src/open-optimization-stream.test.ts +++ /dev/null @@ -1,208 +0,0 @@ -import { describe, expect, it, vi } from "vitest"; - -import { openPetrinautOptimizationStream } from "./open-optimization-stream.js"; -import { PetrinautOptimizerHttpError } from "./optimizer-http.js"; - -import type { PetrinautOptimizationInput } from "@hashintel/petrinaut-core"; - -const input = { - objective: { direction: "maximize" }, - study: { trials: 2 }, -} as PetrinautOptimizationInput; - -/** Collect every event returned by one opened optimization stream. */ -const collect = async (events: AsyncIterable): Promise => { - const collected = []; - for await (const event of events) { - collected.push(event); - } - return collected; -}; - -describe("openPetrinautOptimizationStream", () => { - it("posts the manifest and returns canonical optimization events", async () => { - const onActivity = vi.fn(); - const signal = new AbortController().signal; - const fetchImpl = vi.fn(async () => - Promise.resolve( - new Response( - 'data: {"step":0,"params":{"rate":0.4},"metric":2,"state":"COMPLETE"}\n\n' + - "event: done\ndata: {}\n\n", - { - headers: { - "content-type": "text/event-stream", - "x-optimization-run-id": "run-42", - }, - }, - ), - ), - ); - - const { events, optimizationRunId } = await openPetrinautOptimizationStream( - { - endpoint: "http://petrinaut-opt.test/optimize/all", - fetchImpl, - input, - onActivity, - signal, - }, - ); - - expect(optimizationRunId).toBe("run-42"); - await expect(collect(events)).resolves.toEqual([ - { type: "started", requestedTrials: 2 }, - { - type: "trial", - trial: 0, - parameters: { rate: 0.4 }, - objective: 2, - state: "complete", - best: { trial: 0, parameters: { rate: 0.4 }, objective: 2 }, - }, - { - type: "complete", - requestedTrials: 2, - completedTrials: 1, - prunedTrials: 0, - failedTrials: 0, - best: { trial: 0, parameters: { rate: 0.4 }, objective: 2 }, - }, - ]); - expect(fetchImpl).toHaveBeenCalledWith( - "http://petrinaut-opt.test/optimize/all", - { - method: "POST", - headers: { - accept: "text/event-stream", - "content-type": "application/json", - }, - body: JSON.stringify(input), - signal, - }, - ); - expect(onActivity).toHaveBeenCalledOnce(); - }); - - it("forwards the request id header and reads the missing run id as null", async () => { - const fetchImpl = vi.fn(async () => - Promise.resolve( - new Response("event: done\ndata: {}\n\n", { - headers: { "content-type": "text/event-stream" }, - }), - ), - ); - - const { events, optimizationRunId } = await openPetrinautOptimizationStream( - { - endpoint: "http://petrinaut-opt.test/optimize/all", - fetchImpl, - input, - requestId: "request-123", - }, - ); - await collect(events); - - expect(optimizationRunId).toBeNull(); - expect(fetchImpl).toHaveBeenCalledWith( - "http://petrinaut-opt.test/optimize/all", - expect.objectContaining({ - headers: { - accept: "text/event-stream", - "content-type": "application/json", - "x-hash-request-id": "request-123", - }, - }), - ); - }); - - it("surfaces a FastAPI error message", async () => { - const result = openPetrinautOptimizationStream({ - endpoint: "/optimize/all", - fetchImpl: async () => - Response.json( - { detail: "Invalid optimization manifest" }, - { status: 422, headers: { "retry-after": "5" } }, - ), - input, - }); - - await expect(result).rejects.toBeInstanceOf(PetrinautOptimizerHttpError); - await expect(result).rejects.toMatchObject({ - message: "Invalid optimization manifest", - retryAfter: "5", - status: 422, - }); - }); - - it("captures the run id from a failed optimizer response", async () => { - const result = openPetrinautOptimizationStream({ - endpoint: "/optimize/all", - fetchImpl: async () => - Response.json( - { detail: "failed to initialise optimization" }, - { status: 500, headers: { "x-optimization-run-id": "run-err-7" } }, - ), - input, - }); - - await expect(result).rejects.toMatchObject({ - optimizationRunId: "run-err-7", - status: 500, - }); - }); - - it("preserves the busy status and Retry-After of an optimizer 429", async () => { - const result = openPetrinautOptimizationStream({ - endpoint: "/optimize/all", - fetchImpl: async () => - Response.json( - { - detail: - "The optimizer is already running its maximum number of studies", - }, - { status: 429, headers: { "retry-after": "30" } }, - ), - input, - }); - - await expect(result).rejects.toBeInstanceOf(PetrinautOptimizerHttpError); - await expect(result).rejects.toMatchObject({ - message: "The optimizer is already running its maximum number of studies", - retryAfter: "30", - status: 429, - }); - }); - - it("reports a null Retry-After when the optimizer 429 omits the header", async () => { - const result = openPetrinautOptimizationStream({ - endpoint: "/optimize/all", - fetchImpl: async () => Response.json({ detail: "busy" }, { status: 429 }), - input, - }); - - await expect(result).rejects.toMatchObject({ - retryAfter: null, - status: 429, - }); - }); - - it("falls back to the upstream status for an unstructured error", async () => { - await expect( - openPetrinautOptimizationStream({ - endpoint: "/optimize/all", - fetchImpl: async () => new Response("failure", { status: 500 }), - input, - }), - ).rejects.toThrow("Petrinaut optimizer returned status 500"); - }); - - it("rejects a successful response without a body", async () => { - await expect( - openPetrinautOptimizationStream({ - endpoint: "/optimize/all", - fetchImpl: async () => new Response(null, { status: 200 }), - input, - }), - ).rejects.toThrow("Petrinaut optimizer returned an empty response"); - }); -}); diff --git a/libs/@local/petrinaut-optimizer-client/src/open-optimization-stream.ts b/libs/@local/petrinaut-optimizer-client/src/open-optimization-stream.ts deleted file mode 100644 index b71c409258c..00000000000 --- a/libs/@local/petrinaut-optimizer-client/src/open-optimization-stream.ts +++ /dev/null @@ -1,78 +0,0 @@ -import { decodePetrinautOptimizerStream } from "./decode-optimization-stream.js"; -import { petrinautOptimizerHttpErrorFromResponse } from "./optimizer-http.js"; - -import type { PetrinautOptimizerFetch } from "./optimizer-http.js"; -import type { - AbortSignalLike, - PetrinautOptimizationEvent, - PetrinautOptimizationInput, -} from "@hashintel/petrinaut-core"; - -/** Configuration for opening one Petrinaut Optimizer study stream. */ -export type OpenPetrinautOptimizationStreamOptions = { - /** URL of Petrinaut Optimizer's `/optimize/all` endpoint. */ - endpoint: string | URL; - /** Fetch implementation supplied by the current runtime or a test. */ - fetchImpl?: PetrinautOptimizerFetch; - /** Complete optimization manifest sent to Petrinaut Optimizer. */ - input: PetrinautOptimizationInput; - /** Optional maximum UTF-8 size of one upstream SSE event. */ - maxEventBytes?: number; - /** Called whenever upstream bytes arrive, including heartbeats. */ - onActivity?: () => void; - /** Correlation id forwarded upstream as the `x-hash-request-id` header. */ - requestId?: string; - /** Signal used to cancel the request and its response stream. */ - signal?: AbortSignalLike; -}; - -/** One opened optimization stream plus its upstream correlation id. */ -export type PetrinautOptimizationStreamHandle = { - /** Canonical optimization events decoded from the upstream stream. */ - events: AsyncIterable; - /** The optimizer's `X-Optimization-Run-ID` header, when provided. */ - optimizationRunId: string | null; -}; - -/** - * Post an optimization manifest and open its canonical event stream. - * - * This isomorphic transport boundary is shared by NodeAPI and direct browser - * development integrations so they use identical request and error handling. - */ -export const openPetrinautOptimizationStream = async ({ - endpoint, - fetchImpl = fetch, - input, - maxEventBytes, - onActivity, - requestId, - signal, -}: OpenPetrinautOptimizationStreamOptions): Promise => { - const response = await fetchImpl(endpoint, { - method: "POST", - headers: { - accept: "text/event-stream", - "content-type": "application/json", - ...(requestId === undefined ? {} : { "x-hash-request-id": requestId }), - }, - body: JSON.stringify(input), - signal: signal as AbortSignal | undefined, - }); - if (!response.ok) { - throw await petrinautOptimizerHttpErrorFromResponse(response); - } - if (!response.body) { - throw new Error("Petrinaut optimizer returned an empty response"); - } - - return { - events: decodePetrinautOptimizerStream(response.body, { - direction: input.objective.direction, - requestedTrials: input.study.trials, - ...(maxEventBytes === undefined ? {} : { maxEventBytes }), - ...(onActivity ? { onActivity } : {}), - }), - optimizationRunId: response.headers.get("x-optimization-run-id"), - }; -}; diff --git a/libs/@local/petrinaut-optimizer-client/src/optimizer-http.ts b/libs/@local/petrinaut-optimizer-client/src/optimizer-http.ts index b1b891e6251..e2b0953f6ca 100644 --- a/libs/@local/petrinaut-optimizer-client/src/optimizer-http.ts +++ b/libs/@local/petrinaut-optimizer-client/src/optimizer-http.ts @@ -20,6 +20,22 @@ export class PetrinautOptimizerHttpError extends Error { } } +/** + * Resolve a service-relative path against the optimizer endpoint, keeping any + * path prefix the endpoint carries (e.g. a dev proxy mounting the service + * under `/api/petrinaut-opt`). + */ +export const petrinautOptimizerUrl = ( + endpoint: string | URL, + path: string, +): URL => { + const base = new URL(endpoint); + if (!base.pathname.endsWith("/")) { + base.pathname = `${base.pathname}/`; + } + return new URL(path, base); +}; + /** Return whether an unknown value is a non-array JSON object. */ const isJsonRecord = (value: unknown): value is JsonRecord => typeof value === "object" && value !== null && !Array.isArray(value);