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/.changeset/fe-1225-optimization-event-seq.md b/.changeset/fe-1225-optimization-event-seq.md new file mode 100644 index 00000000000..b928dbe7022 --- /dev/null +++ b/.changeset/fe-1225-optimization-event-seq.md @@ -0,0 +1,5 @@ +--- +"@hashintel/petrinaut-core": patch +--- + +Add an optional `seq` sequence number to every optimization event so detached, reconnectable optimization runs can be resumed from a cursor. diff --git a/apps/hash-api/README.md b/apps/hash-api/README.md index cfdd3733872..c76b5fe0f62 100644 --- a/apps/hash-api/README.md +++ b/apps/hash-api/README.md @@ -37,6 +37,15 @@ The HASH Backend API service is configured using the following environment varia together to enable Petrinaut optimization. The authenticated capabilities endpoint reports this configuration independently of service health, so a temporary optimizer outage does not make the feature disappear from HASH. + - NodeAPI proxies detached, reconnectable optimization runs: + `POST /api/petrinaut-optimizer/optimize/runs` starts a run and returns its + id, `GET /api/petrinaut-optimizer/optimize/runs/:runId/events?cursor=N` + attaches to (and resumes) its NDJSON event stream without affecting the + run, and `DELETE /api/petrinaut-optimizer/optimize/runs/:runId` cancels + it. NodeAPI is a thin proxy: it authenticates, rate-limits, and + validates, then forwards with the account's `x-hash-account-id` tag — + the optimizer itself enforces per-account single-flight and owner-only + run visibility. - `FRONTEND_URL`: The URL the frontend is hosted on. - Vault - `HASH_VAULT_HOST`: The host address (including protocol) that the Vault server is running on, e.g. `http://127.0.0.1` diff --git a/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-handler.test.ts b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-handler.test.ts deleted file mode 100644 index 1c537ad64a9..00000000000 --- a/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-handler.test.ts +++ /dev/null @@ -1,827 +0,0 @@ -import { EventEmitter } from "node:events"; - -import { describe, expect, it, vi } from "vitest"; - -import { createPetrinautOptimizationHandler } from "./create-petrinaut-optimization-handler"; - -import type { Logger } from "@local/hash-backend-utils/logger"; -import type { Request, Response as ExpressResponse } from "express"; - -const validOptimizationInput = { - kind: "petrinaut-optimization", - version: 1, - name: "Optimize rate", - model: { - title: "Example", - definition: { - places: [], - transitions: [], - types: [], - differentialEquations: [], - parameters: [], - subnets: [], - componentInstances: [], - scenarios: [ - { - id: "baseline", - name: "Baseline", - scenarioParameters: [ - { identifier: "rate", type: "real", default: 0.5 }, - ], - parameterOverrides: {}, - initialState: { type: "per_place", content: {} }, - }, - ], - metrics: [{ id: "profit", name: "Profit", code: "return 1;" }], - }, - }, - scenario: { - id: "baseline", - parameterBindings: { - rate: { - kind: "optimize", - domain: { - kind: "continuous", - minimum: 0.1, - maximum: 1, - scale: "linear", - }, - }, - }, - }, - objective: { metricId: "profit", direction: "maximize" }, - execution: { seed: 42, dt: 0.1, maxTime: 10 }, - study: { trials: 2, sampler: "tpe" }, -}; - -type RecordedLog = { - level: "info" | "warn"; - message: string; - metadata: Record; -}; - -/** Build a logger fake that records structured logs. */ -const createRecordingLogger = () => { - const entries: RecordedLog[] = []; - const record = - (level: "info" | "warn") => - (message: string, metadata?: Record) => { - entries.push({ - level, - message, - metadata: metadata ?? {}, - }); - }; - const logger = { - info: record("info"), - warn: record("warn"), - } as unknown as Pick; - return { entries, logger }; -}; - -const unexpectedFetch = async (): Promise => { - throw new Error("Unexpected upstream request"); -}; - -/** Mutable fake Express response the backpressure tests can drive. */ -type FakeResponse = EventEmitter & ExpressResponse & { destroyed: boolean }; - -const callHandler = async ({ - authenticated = true, - body, - fetchImpl = unexpectedFetch, - handler, - logger, - onRequest, - onResponse, - writeReturns, -}: { - authenticated?: boolean; - body: unknown; - fetchImpl?: (input: string | URL, init?: RequestInit) => Promise; - /** Reuse one handler across calls to observe its slot bookkeeping. */ - handler?: ReturnType; - logger?: Pick; - onRequest?: (request: EventEmitter) => void; - onResponse?: (response: FakeResponse) => void; - /** Decide each write's backpressure result; defaults to no backpressure. */ - writeReturns?: (value: string) => boolean; -}) => { - let statusCode = 200; - let bodyResult: unknown; - const headers: Record = {}; - const output: string[] = []; - let headersSent = false; - let writableEnded = false; - let writableNeedDrain = false; - const responseEmitter = new EventEmitter(); - // `Object.assign` would copy getter *values*, freezing them at `false`. - Object.defineProperties(responseEmitter, { - headersSent: { get: () => headersSent }, - writableEnded: { get: () => writableEnded }, - writableNeedDrain: { get: () => writableNeedDrain }, - }); - // Clear the drain flag without registering a listener, so the tests can - // assert that the handler leaves no listeners of its own behind. - const originalEmit = responseEmitter.emit.bind(responseEmitter); - responseEmitter.emit = (eventName: string | symbol, ...args: unknown[]) => { - if (eventName === "drain") { - writableNeedDrain = false; - } - return originalEmit(eventName, ...args); - }; - const response = Object.assign(responseEmitter, { - destroyed: false, - end: () => { - writableEnded = true; - }, - flushHeaders: () => { - headersSent = true; - }, - get: (name: string) => - name.toLowerCase() === "x-hash-request-id" ? "request-id-1" : undefined, - json: (value: unknown) => { - bodyResult = value; - headersSent = true; - writableEnded = true; - return response; - }, - set: (value: Record) => { - Object.assign(headers, value); - return response; - }, - status: (value: number) => { - statusCode = value; - return response; - }, - write: (value: string) => { - headersSent = true; - output.push(value); - const flushed = writeReturns?.(value) ?? true; - if (!flushed) { - writableNeedDrain = true; - } - return flushed; - }, - }) as unknown as FakeResponse; - const request = Object.assign(new EventEmitter(), { - body, - user: authenticated - ? ({ accountId: "user-1" } as NonNullable) - : undefined, - }) as unknown as Request; - const activeHandler = - handler ?? - createPetrinautOptimizationHandler({ - fetchImpl, - logger: logger ?? createRecordingLogger().logger, - origin: new URL("http://petrinaut-opt:4004"), - }); - - const handlerPromise = activeHandler(request, response, () => undefined); - onRequest?.(request); - onResponse?.(response); - await handlerPromise; - - return { body: bodyResult, headers, output, response, statusCode }; -}; - -describe("createPetrinautOptimizationHandler", () => { - it("requires authentication", async () => { - const result = await callHandler({ authenticated: false, body: {} }); - - expect(result).toMatchObject({ - body: { error: "Authentication required" }, - statusCode: 401, - }); - }); - - it("validates the public optimization request", async () => { - const { entries, logger } = createRecordingLogger(); - const result = await callHandler({ - body: { - ...validOptimizationInput, - scenario: { - ...validOptimizationInput.scenario, - parameterBindings: { - rate: { - kind: "optimize", - domain: { - kind: "continuous", - minimum: 0.1, - maximum: 1, - scale: "sqrt", - }, - }, - }, - }, - }, - logger, - }); - - expect(result.statusCode).toBe(400); - expect(result.body).toEqual({ - code: "invalid_optimization_request", - details: { - issues: [ - { - message: expect.any(String), - path: "scenario.parameterBindings.rate.domain.scale", - }, - ], - truncated: false, - }, - error: "Invalid optimization request", - }); - - // HTTP tracing captures the 400; avoid duplicating validation data in - // application logs because the manifest can embed user-authored code. - expect(entries).toEqual([]); - const loggedText = JSON.stringify(entries); - expect(loggedText).not.toContain("return 1;"); - expect(loggedText).not.toContain("petrinaut-optimization"); - }); - - it("preserves an upstream optimizer busy response", async () => { - const result = await callHandler({ - body: validOptimizationInput, - fetchImpl: async () => - Response.json( - { detail: "The optimizer is busy" }, - { status: 429, headers: { "retry-after": "10" } }, - ), - }); - - expect(result).toMatchObject({ - body: { error: "Petrinaut optimizer is busy" }, - headers: { "Retry-After": "10" }, - statusCode: 429, - }); - }); - - it("keeps a quiet downstream optimization stream alive", async () => { - vi.useFakeTimers(); - let activeRequest: EventEmitter | undefined; - - try { - const resultPromise = callHandler({ - body: validOptimizationInput, - fetchImpl: async (_input, init) => - new Response( - new ReadableStream({ - start(controller) { - init?.signal?.addEventListener( - "abort", - () => - controller.error(new DOMException("Aborted", "AbortError")), - { once: true }, - ); - }, - }), - { headers: { "content-type": "text/event-stream" } }, - ), - onRequest: (request) => { - activeRequest = request; - }, - }); - - await vi.advanceTimersByTimeAsync(0); - await vi.advanceTimersByTimeAsync(25_000); - activeRequest?.emit("aborted"); - - const result = await resultPromise; - expect(result.output).toContain("\n"); - } finally { - vi.useRealTimers(); - } - }); - - it("proxies optimizer SSE as canonical NDJSON", async () => { - let upstreamRequest: - | { body: string | undefined; requestId: unknown; url: string } - | undefined; - const upstream = [ - 'data: {"step":0,"params":{"rate":0.4},"init_state":{},"metric":2,"state":"COMPLETE"}\n\n', - ": heartbeat\n\n", - 'data: {"step":1,"params":{"rate":0.8},"init_state":{},"metric":4,"state":"COMPLETE"}\n\n', - "event: done\ndata: {}\n\n", - ].join(""); - const { entries, logger } = createRecordingLogger(); - - const result = await callHandler({ - body: validOptimizationInput, - fetchImpl: async (input, init) => { - upstreamRequest = { - body: typeof init?.body === "string" ? init.body : undefined, - requestId: new Headers(init?.headers).get("x-hash-request-id"), - url: input.toString(), - }; - return new Response(upstream, { - headers: { - "content-type": "text/event-stream", - "x-optimization-run-id": "run-1", - }, - }); - }, - logger, - }); - - expect(result.statusCode).toBe(200); - expect(result.headers).toMatchObject({ - "Content-Type": "application/x-ndjson; charset=utf-8", - "X-Accel-Buffering": "no", - // The optimizer's run id reaches the browser for correlation. - "X-Optimization-Run-ID": "run-1", - }); - expect(result.output).toEqual([ - '{"type":"started","requestedTrials":2}\n', - '{"type":"trial","trial":0,"parameters":{"rate":0.4},"objective":2,"state":"complete","best":{"trial":0,"parameters":{"rate":0.4},"objective":2}}\n', - '{"type":"trial","trial":1,"parameters":{"rate":0.8},"objective":4,"state":"complete","best":{"trial":1,"parameters":{"rate":0.8},"objective":4}}\n', - '{"type":"complete","requestedTrials":2,"completedTrials":2,"prunedTrials":0,"failedTrials":0,"best":{"trial":1,"parameters":{"rate":0.8},"objective":4}}\n', - ]); - expect(upstreamRequest?.url).toBe("http://petrinaut-opt:4004/optimize/all"); - expect(upstreamRequest?.requestId).toBe("request-id-1"); - expect(JSON.parse(upstreamRequest?.body ?? "null")).toEqual( - validOptimizationInput, - ); - - // The full request lifecycle is logged with its correlation identifiers. - expect(entries).toEqual([ - { - level: "info", - message: "Petrinaut optimization stream opened", - metadata: { - optimizationRunId: "run-1", - requestId: "request-id-1", - }, - }, - { - level: "info", - message: "Petrinaut optimization finished", - metadata: { - durationMs: expect.any(Number), - optimizationRunId: "run-1", - outcome: "completed", - requestId: "request-id-1", - }, - }, - ]); - // The manifest itself — including user-authored code — is never logged. - const loggedText = JSON.stringify(entries); - expect(loggedText).not.toContain("return 1;"); - expect(loggedText).not.toContain("petrinaut-optimization"); - }); - - it("logs an upstream terminal error as failed", async () => { - const { entries, logger } = createRecordingLogger(); - const result = await callHandler({ - body: validOptimizationInput, - fetchImpl: async () => - new Response('event: error\ndata: {"message":"failed"}\n\n', { - headers: { - "content-type": "text/event-stream", - "x-optimization-run-id": "run-failed-1", - }, - }), - logger, - }); - - expect(result.statusCode).toBe(200); - expect(result.output.at(-1)).toContain('"type":"error"'); - expect(entries.at(-1)).toEqual({ - level: "warn", - message: "Petrinaut optimization failed", - metadata: { - durationMs: expect.any(Number), - optimizationRunId: "run-failed-1", - outcome: "upstream-error", - requestId: "request-id-1", - }, - }); - }); - - it("preserves an upstream busy response without a Retry-After header", async () => { - const result = await callHandler({ - body: validOptimizationInput, - fetchImpl: async () => - Response.json({ detail: "The optimizer is busy" }, { status: 429 }), - }); - - expect(result.statusCode).toBe(429); - expect(result.body).toEqual({ error: "Petrinaut optimizer is busy" }); - expect(result.headers).not.toHaveProperty("Retry-After"); - }); - - it("correlates a pre-stream upstream failure with the optimizer run id", async () => { - const { entries, logger } = createRecordingLogger(); - const result = await callHandler({ - body: validOptimizationInput, - fetchImpl: async () => - Response.json( - { detail: "failed to initialise optimization" }, - { status: 500, headers: { "x-optimization-run-id": "run-init-9" } }, - ), - logger, - }); - - expect(result.statusCode).toBe(502); - expect(result.headers).toMatchObject({ - "X-Optimization-Run-ID": "run-init-9", - }); - const failure = entries.find( - (entry) => entry.message === "Petrinaut optimization failed", - ); - expect(failure?.metadata).toMatchObject({ - errorType: "PetrinautOptimizerHttpError", - optimizationRunId: "run-init-9", - outcome: "upstream-error", - upstreamStatus: 500, - }); - expect(JSON.stringify(failure)).not.toContain( - "failed to initialise optimization", - ); - }); - - it("cleans up backpressure listeners and survives a late close", async () => { - const unhandledRejections: unknown[] = []; - const onUnhandledRejection = (reason: unknown) => - unhandledRejections.push(reason); - process.on("unhandledRejection", onUnhandledRejection); - let activeResponse: FakeResponse | undefined; - - try { - const upstream = [ - 'data: {"step":0,"params":{"rate":0.4},"init_state":{},"metric":2,"state":"COMPLETE"}\n\n', - 'data: {"step":1,"params":{"rate":0.8},"init_state":{},"metric":4,"state":"COMPLETE"}\n\n', - "event: done\ndata: {}\n\n", - ].join(""); - - const result = await callHandler({ - body: validOptimizationInput, - fetchImpl: async () => - new Response(upstream, { - headers: { "content-type": "text/event-stream" }, - }), - onResponse: (response) => { - activeResponse = response; - }, - writeReturns: (value) => { - if (value.includes('"trial":0')) { - // Backpressure this write, then let the buffer drain shortly - // after the handler has started waiting. - setImmediate(() => activeResponse?.emit("drain")); - return false; - } - return true; - }, - }); - - expect(result.statusCode).toBe(200); - expect(result.output).toHaveLength(4); - expect(result.response.listenerCount("drain")).toBe(0); - expect(result.response.listenerCount("close")).toBe(0); - - // The response closing after completion must not trip the listener - // that lost the earlier drain/close race. - result.response.emit("close"); - await new Promise(setImmediate); - expect(unhandledRejections).toEqual([]); - } finally { - process.off("unhandledRejection", onUnhandledRejection); - } - }); - - it("stops a backpressured stream when the client closes and frees the slot", async () => { - const handler = createPetrinautOptimizationHandler({ - fetchImpl: async (_input, init) => - new Response( - new ReadableStream({ - start(controller) { - controller.enqueue( - new TextEncoder().encode( - 'data: {"step":0,"params":{"rate":0.4},"init_state":{},"metric":2,"state":"COMPLETE"}\n\n', - ), - ); - init?.signal?.addEventListener( - "abort", - () => - controller.error(new DOMException("Aborted", "AbortError")), - { once: true }, - ); - }, - }), - { headers: { "content-type": "text/event-stream" } }, - ), - logger: createRecordingLogger().logger, - origin: new URL("http://petrinaut-opt:4004"), - }); - let activeResponse: FakeResponse | undefined; - - const first = await callHandler({ - body: validOptimizationInput, - handler, - onResponse: (response) => { - activeResponse = response; - }, - writeReturns: (value) => { - if (value.includes('"trial":0')) { - setImmediate(() => { - if (activeResponse) { - activeResponse.destroyed = true; - activeResponse.emit("close"); - } - }); - return false; - } - return true; - }, - }); - - expect(first.statusCode).toBe(200); - expect(first.response.listenerCount("drain")).toBe(0); - expect(first.response.listenerCount("close")).toBe(0); - - // The same user must be admitted again: the slot was released. - const second = await callHandler({ - body: validOptimizationInput, - handler, - writeReturns: (value) => { - if (value.includes('"trial":0')) { - setImmediate(() => { - activeResponse!.destroyed = true; - activeResponse!.emit("close"); - }); - return false; - } - return true; - }, - onResponse: (response) => { - activeResponse = response; - }, - }); - expect(second.statusCode).toBe(200); - expect(second.body).toBeUndefined(); - }); - - it("breaks a backpressured wait on timeout instead of holding the slot", async () => { - vi.useFakeTimers(); - - try { - const handler = createPetrinautOptimizationHandler({ - fetchImpl: async (_input, init) => - new Response( - new ReadableStream({ - start(controller) { - controller.enqueue( - new TextEncoder().encode( - 'data: {"step":0,"params":{"rate":0.4},"init_state":{},"metric":2,"state":"COMPLETE"}\n\n', - ), - ); - init?.signal?.addEventListener( - "abort", - () => - controller.error(new DOMException("Aborted", "AbortError")), - { once: true }, - ); - }, - }), - { headers: { "content-type": "text/event-stream" } }, - ), - logger: createRecordingLogger().logger, - origin: new URL("http://petrinaut-opt:4004"), - }); - - const firstPromise = callHandler({ - body: validOptimizationInput, - handler, - // The client never drains and never disconnects. - writeReturns: (value) => !value.includes('"trial":0'), - }); - - // The 5 minute idle timeout aborts the stalled stream. - await vi.advanceTimersByTimeAsync(0); - await vi.advanceTimersByTimeAsync(5 * 60_000 + 1_000); - const first = await firstPromise; - - expect(first.statusCode).toBe(200); - expect(first.output.at(-1)).toContain('"code":"optimization_timeout"'); - expect(first.response.writableEnded).toBe(true); - expect(first.response.listenerCount("drain")).toBe(0); - expect(first.response.listenerCount("close")).toBe(0); - // Heartbeats must not stuff a buffer that already needs draining. - expect(first.output.filter((frame) => frame === "\n")).toHaveLength(0); - - // The user slot must be free again after the timeout teardown. - const secondPromise = callHandler({ - body: validOptimizationInput, - handler, - writeReturns: (value) => !value.includes('"trial":0'), - }); - await vi.advanceTimersByTimeAsync(0); - await vi.advanceTimersByTimeAsync(5 * 60_000 + 1_000); - const second = await secondPromise; - expect(second.statusCode).toBe(200); - expect(second.body).toBeUndefined(); - } finally { - vi.useRealTimers(); - } - }); - - it("keeps a single terminal event when a timeout hits the final write", async () => { - vi.useFakeTimers(); - - try { - const upstream = [ - 'data: {"step":0,"params":{"rate":0.4},"init_state":{},"metric":2,"state":"COMPLETE"}\n\n', - "event: done\ndata: {}\n\n", - ].join(""); - - const resultPromise = callHandler({ - body: validOptimizationInput, - fetchImpl: async () => - new Response(upstream, { - headers: { "content-type": "text/event-stream" }, - }), - // The final complete event is committed to the buffer, but the - // client never drains it before the idle timeout fires. - writeReturns: (value) => !value.includes('"type":"complete"'), - }); - - await vi.advanceTimersByTimeAsync(0); - await vi.advanceTimersByTimeAsync(5 * 60_000 + 1_000); - const result = await resultPromise; - - const terminalEvents = result.output.filter( - (frame) => - frame.includes('"type":"complete"') || - frame.includes('"type":"error"'), - ); - expect(terminalEvents).toHaveLength(1); - expect(terminalEvents[0]).toContain('"type":"complete"'); - expect(result.response.writableEnded).toBe(true); - } finally { - vi.useRealTimers(); - } - }); - - it.each([ - { - expectedLevel: "info" as const, - expectedMessage: "Petrinaut optimization finished", - expectedOutcome: "completed", - terminalType: "complete", - upstream: [ - 'data: {"step":0,"params":{"rate":0.4},"init_state":{},"metric":2,"state":"COMPLETE"}\n\n', - "event: done\ndata: {}\n\n", - ].join(""), - }, - { - expectedLevel: "warn" as const, - expectedMessage: "Petrinaut optimization failed", - expectedOutcome: "upstream-error", - terminalType: "error", - upstream: 'event: error\ndata: {"message":"failed"}\n\n', - }, - ])( - "preserves $expectedOutcome when the client disconnects during terminal backpressure", - async ({ - expectedLevel, - expectedMessage, - expectedOutcome, - terminalType, - upstream, - }) => { - const { entries, logger } = createRecordingLogger(); - let activeResponse: FakeResponse | undefined; - - const result = await callHandler({ - body: validOptimizationInput, - fetchImpl: async () => - new Response(upstream, { - headers: { - "content-type": "text/event-stream", - "x-optimization-run-id": "run-terminal-1", - }, - }), - logger, - onResponse: (response) => { - activeResponse = response; - }, - writeReturns: (value) => { - if (value.includes(`"type":"${terminalType}"`)) { - setImmediate(() => { - activeResponse!.destroyed = true; - activeResponse!.emit("close"); - }); - return false; - } - return true; - }, - }); - - const terminalEvents = result.output.filter( - (frame) => - frame.includes('"type":"complete"') || - frame.includes('"type":"error"'), - ); - expect(terminalEvents).toHaveLength(1); - expect(terminalEvents[0]).toContain(`"type":"${terminalType}"`); - expect(entries.at(-1)).toEqual({ - level: expectedLevel, - message: expectedMessage, - metadata: { - durationMs: expect.any(Number), - optimizationRunId: "run-terminal-1", - outcome: expectedOutcome, - requestId: "request-id-1", - }, - }); - }, - ); - - it("aborts disconnected requests and releases the user slot", async () => { - let abortObserved = false; - const { entries, logger } = createRecordingLogger(); - await callHandler({ - body: validOptimizationInput, - fetchImpl: async (_input, init) => - new Promise((_resolve, reject) => { - init?.signal?.addEventListener( - "abort", - () => { - abortObserved = true; - reject(new DOMException("Aborted", "AbortError")); - }, - { once: true }, - ); - }), - logger, - onRequest: (request) => request.emit("aborted"), - }); - - expect(abortObserved).toBe(true); - // A silent client disconnect still leaves a terminal outcome log. - expect(entries.at(-1)).toEqual({ - level: "info", - message: "Petrinaut optimization finished", - metadata: { - durationMs: expect.any(Number), - optimizationRunId: null, - outcome: "client-disconnected", - requestId: "request-id-1", - }, - }); - - const second = await callHandler({ - body: validOptimizationInput, - fetchImpl: async () => - new Response('event: error\ndata: {"message":"failed"}\n\n', { - headers: { "content-type": "text/event-stream" }, - }), - }); - expect(second.statusCode).toBe(200); - expect(second.output).toHaveLength(2); - }); - - it("logs capacity rejections for a user with a running optimization", async () => { - const { entries, logger } = createRecordingLogger(); - let releaseFirst: (() => void) | undefined; - const handler = createPetrinautOptimizationHandler({ - fetchImpl: async (_input, init) => - new Promise((_resolve, reject) => { - releaseFirst = () => - reject(new DOMException("Aborted", "AbortError")); - init?.signal?.addEventListener("abort", () => releaseFirst?.(), { - once: true, - }); - }), - logger, - origin: new URL("http://petrinaut-opt:4004"), - }); - let firstRequest: EventEmitter | undefined; - - const firstPromise = callHandler({ - body: validOptimizationInput, - handler, - onRequest: (request) => { - firstRequest = request; - }, - }); - // Let the first request reach the pending upstream fetch. - await new Promise(setImmediate); - - const second = await callHandler({ body: validOptimizationInput, handler }); - expect(second.statusCode).toBe(429); - expect(entries).toContainEqual({ - level: "warn", - message: "Petrinaut optimization rejected: account busy", - metadata: { - activeOptimizationCount: 1, - requestId: "request-id-1", - }, - }); - - firstRequest?.emit("aborted"); - await firstPromise; - }); -}); diff --git a/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-handler.ts b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-handler.ts deleted file mode 100644 index 763e5dd8442..00000000000 --- a/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-handler.ts +++ /dev/null @@ -1,443 +0,0 @@ -import { once } from "node:events"; - -import { petrinautOptimizationInputSchema } from "@hashintel/petrinaut-core"; -import { - openPetrinautOptimizationStream, - PetrinautOptimizerHttpError, -} from "@local/petrinaut-optimizer-client"; - -import type { PetrinautOptimizationEvent } from "@hashintel/petrinaut-core"; -import type { Logger } from "@local/hash-backend-utils/logger"; -import type { PetrinautOptimizerFetch } from "@local/petrinaut-optimizer-client"; -import type { - Request, - RequestHandler, - Response as ExpressResponse, -} from "express"; - -const RESPONSE_START_TIMEOUT_MS = 30_000; -const DOWNSTREAM_HEARTBEAT_INTERVAL_MS = 25_000; -const IDLE_TIMEOUT_MS = 5 * 60_000; -const OVERALL_TIMEOUT_MS = 15 * 60_000; -const MAX_REQUEST_BYTES = 8 * 1024 * 1024; -const MAX_VALIDATION_ISSUES = 5; -const MAX_VALIDATION_MESSAGE_LENGTH = 300; -// A trial repeats optimized parameter identifiers from the accepted manifest, -// so use the same bound rather than rejecting a valid large search space. -const MAX_EVENT_BYTES = MAX_REQUEST_BYTES; -const MAX_CONCURRENT_OPTIMIZATIONS = 4; -const INVALID_OPTIMIZATION_REQUEST = { - code: "invalid_optimization_request", - error: "Invalid optimization request", -} as const; - -type OptimizationTimeoutKind = "response_start" | "idle" | "overall"; -type OptimizationTerminalOutcome = "completed" | "upstream-error"; - -type OptimizationRequestLifecycle = { - /** Signal that cancels the upstream optimizer request. */ - signal: AbortSignal; - /** Mutable outcome needed when translating a stream failure. */ - state: { - clientDisconnected: boolean; - terminalOutcome: OptimizationTerminalOutcome | null; - timeoutKind: OptimizationTimeoutKind | null; - }; - /** Stop timers, heartbeats, and request/response listeners. */ - cleanup: () => void; - /** Clear the response-start deadline and begin idle/heartbeat tracking. */ - markResponseStarted: () => void; - /** Remember the domain outcome committed to the public stream. */ - markTerminalOutcome: (outcome: OptimizationTerminalOutcome) => void; - /** Restart the deadline for receiving upstream bytes. */ - resetIdleTimeout: () => void; -}; - -/** Track cancellation, deadlines, and cleanup for one streamed request. */ -const createOptimizationRequestLifecycle = ( - request: Request, - response: ExpressResponse, -): OptimizationRequestLifecycle => { - const abortController = new AbortController(); - const state: OptimizationRequestLifecycle["state"] = { - clientDisconnected: false, - terminalOutcome: null, - timeoutKind: null, - }; - let downstreamHeartbeat: ReturnType | undefined; - let idleTimeout: ReturnType | undefined; - - /** Abort the upstream request after the HASH client disconnects. */ - const abortForClientDisconnect = () => { - state.clientDisconnected = true; - abortController.abort(); - }; - - /** Record a timeout category and abort the upstream request. */ - const abortForTimeout = (kind: OptimizationTimeoutKind) => { - state.timeoutKind ??= kind; - abortController.abort(); - }; - - const responseStartTimeout = setTimeout( - () => abortForTimeout("response_start"), - RESPONSE_START_TIMEOUT_MS, - ); - const overallTimeout = setTimeout( - () => abortForTimeout("overall"), - OVERALL_TIMEOUT_MS, - ); - - /** Restart the inactivity deadline after any upstream bytes arrive. */ - const resetIdleTimeout = () => { - clearTimeout(idleTimeout); - idleTimeout = setTimeout(() => abortForTimeout("idle"), IDLE_TIMEOUT_MS); - }; - - request.once("aborted", abortForClientDisconnect); - response.once("close", abortForClientDisconnect); - - return { - signal: abortController.signal, - state, - cleanup: () => { - clearTimeout(responseStartTimeout); - clearInterval(downstreamHeartbeat); - clearTimeout(idleTimeout); - clearTimeout(overallTimeout); - request.off("aborted", abortForClientDisconnect); - response.off("close", abortForClientDisconnect); - }, - markResponseStarted: () => { - clearTimeout(responseStartTimeout); - resetIdleTimeout(); - downstreamHeartbeat = setInterval(() => { - if ( - !response.destroyed && - !response.writableEnded && - !response.writableNeedDrain - ) { - // Blank NDJSON lines are transport heartbeats, not domain events. - response.write("\n"); - } - }, DOWNSTREAM_HEARTBEAT_INTERVAL_MS); - downstreamHeartbeat.unref(); - }, - markTerminalOutcome: (outcome) => { - state.terminalOutcome = outcome; - }, - resetIdleTimeout, - }; -}; - -/** Return a bounded, transport-stable summary of manifest validation errors. */ -const summarizeValidationIssues = ( - issues: readonly { message: string; path: readonly PropertyKey[] }[], -) => ({ - issues: issues.slice(0, MAX_VALIDATION_ISSUES).map(({ message, path }) => ({ - path: path.map(String).join(".") || "$", - message: message.slice(0, MAX_VALIDATION_MESSAGE_LENGTH), - })), - truncated: issues.length > MAX_VALIDATION_ISSUES, -}); - -/** Write one canonical optimization event as NDJSON with backpressure. */ -const writeOptimizationEvent = async ( - response: ExpressResponse, - event: PetrinautOptimizationEvent, - signal: AbortSignal, -): Promise => { - // This is schema-validated NDJSON served with `nosniff`, never HTML. - // nosemgrep: javascript.express.security.audit.xss.direct-response-write.direct-response-write - if (response.write(`${JSON.stringify(event)}\n`)) { - return; - } - if (response.destroyed || response.writableEnded) { - throw new Error("The optimization client disconnected"); - } - if (signal.aborted) { - throw new Error("The optimization request was aborted"); - } - // Aborting this controller removes whichever listeners lost the race, so a - // backpressured stream cannot accumulate `drain`/`close` listeners and the - // loser can never surface a late unhandled rejection. - const waitCleanup = new AbortController(); - try { - await Promise.race([ - once(response, "drain", { signal: waitCleanup.signal }), - once(response, "close", { signal: waitCleanup.signal }).then(() => { - throw new Error("The optimization client disconnected"); - }), - once(signal, "abort", { signal: waitCleanup.signal }).then(() => { - throw new Error("The optimization request was aborted"); - }), - ]); - } finally { - waitCleanup.abort(); - } -}; - -/** - * Write a final event without waiting for backpressure to clear. - * - * Teardown must never block on a stalled client, so this write is - * best-effort: the response ends immediately afterwards either way. - */ -const writeTerminalOptimizationEventBestEffort = ( - response: ExpressResponse, - event: PetrinautOptimizationEvent, -): void => { - if (response.destroyed || response.writableEnded) { - return; - } - // This is schema-validated NDJSON served with `nosniff`, never HTML. - // nosemgrep: javascript.express.security.audit.xss.direct-response-write.direct-response-write - response.write(`${JSON.stringify(event)}\n`); -}; - -/** Forward canonical optimization events as NDJSON. */ -const forwardOptimizationEvents = async ( - events: AsyncIterable, - response: ExpressResponse, - markTerminalOutcome: (outcome: OptimizationTerminalOutcome) => void, - signal: AbortSignal, -): Promise => { - for await (const event of events) { - const terminalOutcome = - event.type === "complete" - ? "completed" - : event.type === "error" - ? "upstream-error" - : null; - const backpressureWait = writeOptimizationEvent(response, event, signal); - if (terminalOutcome !== null) { - // The write is committed synchronously even when the buffer is full, - // so preserve its domain outcome even if a later backpressure wait is - // interrupted by a timeout or downstream disconnect. - markTerminalOutcome(terminalOutcome); - } - await backpressureWait; - if (terminalOutcome !== null) { - return terminalOutcome; - } - } - throw new Error("The optimizer stream ended without a terminal event"); -}; - -/** Create the authenticated endpoint that proxies optimization studies. */ -export const createPetrinautOptimizationHandler = ({ - fetchImpl, - logger, - origin, -}: { - fetchImpl: PetrinautOptimizerFetch; - logger: Pick; - origin: URL | null; -}): RequestHandler => { - let activeOptimizationCount = 0; - const activeUserIds = new Set(); - - return async (request, response) => { - const startedAt = Date.now(); - const requestId = response.get("x-hash-request-id") ?? ""; - const logContext = requestId ? { requestId } : {}; - - if (!request.user) { - response.status(401).json({ error: "Authentication required" }); - return; - } - if (!origin) { - response - .status(503) - .json({ error: "Petrinaut optimizer is not configured" }); - return; - } - - const userId = request.user.accountId; - if (activeUserIds.has(userId)) { - logger.warn("Petrinaut optimization rejected: account busy", { - ...logContext, - activeOptimizationCount, - }); - response.status(429).json({ - error: "An optimization is already running for this account", - }); - return; - } - if (activeOptimizationCount >= MAX_CONCURRENT_OPTIMIZATIONS) { - logger.warn("Petrinaut optimization rejected: at capacity", { - ...logContext, - activeOptimizationCount, - maxConcurrentOptimizations: MAX_CONCURRENT_OPTIMIZATIONS, - }); - response.status(429).json({ error: "Petrinaut optimizer is busy" }); - return; - } - - let serializedBody: unknown; - try { - serializedBody = JSON.stringify(request.body); - } catch { - response.status(400).json(INVALID_OPTIMIZATION_REQUEST); - return; - } - if (typeof serializedBody !== "string") { - response.status(400).json(INVALID_OPTIMIZATION_REQUEST); - return; - } - if (Buffer.byteLength(serializedBody, "utf8") > MAX_REQUEST_BYTES) { - response.status(413).json({ error: "Optimization request is too large" }); - return; - } - - const input = petrinautOptimizationInputSchema.safeParse(request.body); - if (!input.success) { - response.status(400).json({ - ...INVALID_OPTIMIZATION_REQUEST, - details: summarizeValidationIssues(input.error.issues), - }); - return; - } - - activeOptimizationCount += 1; - activeUserIds.add(userId); - const lifecycle = createOptimizationRequestLifecycle(request, response); - let optimizationRunId: string | null = null; - const logTerminalOutcome = (outcome: OptimizationTerminalOutcome) => { - const terminalMetadata = { - ...logContext, - durationMs: Date.now() - startedAt, - optimizationRunId, - outcome, - }; - if (outcome === "completed") { - logger.info("Petrinaut optimization finished", terminalMetadata); - } else { - logger.warn("Petrinaut optimization failed", terminalMetadata); - } - }; - - try { - const upstream = await openPetrinautOptimizationStream({ - endpoint: new URL("/optimize/all", origin), - fetchImpl, - input: input.data, - maxEventBytes: MAX_EVENT_BYTES, - onActivity: lifecycle.resetIdleTimeout, - ...(requestId ? { requestId } : {}), - signal: lifecycle.signal, - }); - optimizationRunId = upstream.optimizationRunId; - logger.info("Petrinaut optimization stream opened", { - ...logContext, - optimizationRunId, - }); - - response.status(200); - response.set({ - "Cache-Control": "no-cache, no-store", - "Content-Type": "application/x-ndjson; charset=utf-8", - "X-Accel-Buffering": "no", - // Expose the optimizer's run id so a browser can correlate this - // response with NodeAPI and optimizer logs when the stream fails. - ...(optimizationRunId === null - ? {} - : { "X-Optimization-Run-ID": optimizationRunId }), - }); - response.flushHeaders(); - lifecycle.markResponseStarted(); - const outcome = await forwardOptimizationEvents( - upstream.events, - response, - lifecycle.markTerminalOutcome, - lifecycle.signal, - ); - response.end(); - logTerminalOutcome(outcome); - } catch (error) { - const durationMs = Date.now() - startedAt; - // A study that fails before its stream opens still has a run id on the - // optimizer's error response; keep it so the failure log correlates. - if ( - optimizationRunId === null && - error instanceof PetrinautOptimizerHttpError - ) { - optimizationRunId = error.optimizationRunId; - } - if (lifecycle.state.terminalOutcome !== null) { - logTerminalOutcome(lifecycle.state.terminalOutcome); - if (!response.destroyed && !response.writableEnded) { - response.end(); - } - return; - } - if (lifecycle.state.clientDisconnected || response.destroyed) { - logger.info("Petrinaut optimization finished", { - ...logContext, - durationMs, - optimizationRunId, - outcome: "client-disconnected", - }); - return; - } - logger.warn("Petrinaut optimization failed", { - ...logContext, - durationMs, - errorType: error instanceof Error ? error.name : typeof error, - optimizationRunId, - outcome: lifecycle.state.timeoutKind - ? `timeout:${lifecycle.state.timeoutKind}` - : "upstream-error", - timeoutKind: lifecycle.state.timeoutKind, - upstreamStatus: - error instanceof PetrinautOptimizerHttpError ? error.status : null, - }); - if (!response.headersSent) { - if (optimizationRunId !== null) { - response.set({ "X-Optimization-Run-ID": optimizationRunId }); - } - if ( - error instanceof PetrinautOptimizerHttpError && - error.status === 429 - ) { - if (error.retryAfter) { - response.set({ "Retry-After": error.retryAfter }); - } - response.status(429).json({ error: "Petrinaut optimizer is busy" }); - } else { - response.status(lifecycle.state.timeoutKind ? 504 : 502).json({ - error: lifecycle.state.timeoutKind - ? "Petrinaut optimization timed out" - : "Petrinaut optimization failed", - }); - } - return; - } - try { - writeTerminalOptimizationEventBestEffort(response, { - type: "error", - code: lifecycle.state.timeoutKind - ? "optimization_timeout" - : "upstream_stream_error", - message: lifecycle.state.timeoutKind - ? "The optimization exceeded its execution time limit" - : "The optimizer stream ended unexpectedly", - retryable: true, - } satisfies PetrinautOptimizationEvent); - } catch (writeError) { - logger.warn("Could not report Petrinaut optimization failure", { - ...logContext, - error: writeError, - optimizationRunId, - }); - } - if (!response.writableEnded) { - response.end(); - } - } finally { - lifecycle.cleanup(); - activeOptimizationCount -= 1; - activeUserIds.delete(userId); - } - }; -}; diff --git a/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-cancel-handler.test.ts b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-cancel-handler.test.ts new file mode 100644 index 00000000000..95c49fe9f57 --- /dev/null +++ b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-cancel-handler.test.ts @@ -0,0 +1,135 @@ +import { describe, expect, it } from "vitest"; + +import { createPetrinautOptimizerClient } from "@local/petrinaut-optimizer-client"; + +import { createPetrinautOptimizationRunCancelHandler } from "./create-petrinaut-optimization-run-cancel-handler"; +import { + callOptimizationRunHandler, + createRecordingLogger, + unexpectedFetch, +} from "./shared/optimization-run-test-harness"; + +import type { PetrinautOptimizerFetch } from "@local/petrinaut-optimizer-client"; + +const createHandler = ({ + fetchImpl = unexpectedFetch, + logger = createRecordingLogger().logger, + origin = new URL("http://petrinaut-opt:4004"), +}: { + fetchImpl?: PetrinautOptimizerFetch; + logger?: ReturnType["logger"]; + origin?: URL | null; +} = {}) => + createPetrinautOptimizationRunCancelHandler({ + client: createPetrinautOptimizerClient( + origin ?? "http://petrinaut-optimizer.invalid", + fetchImpl, + ), + logger, + origin, + }); + +describe("createPetrinautOptimizationRunCancelHandler", () => { + it("requires authentication", async () => { + const result = await callOptimizationRunHandler({ + authenticated: false, + handler: createHandler(), + params: { runId: "run-1" }, + }); + + expect(result).toMatchObject({ + body: { error: "Authentication required" }, + statusCode: 401, + }); + }); + + it("forwards the cancel with the account tag and answers 204", async () => { + let upstreamRequest: + | { + accountId: unknown; + method: string | undefined; + requestId: unknown; + url: string; + } + | undefined; + const { entries, logger } = createRecordingLogger(); + const result = await callOptimizationRunHandler({ + handler: createHandler({ + fetchImpl: async (input, init) => { + const headers = new Headers(init?.headers); + upstreamRequest = { + accountId: headers.get("x-hash-account-id"), + method: init?.method, + requestId: headers.get("x-hash-request-id"), + url: input.toString(), + }; + return new Response(null, { status: 204 }); + }, + logger, + }), + params: { runId: "run-1" }, + }); + + expect(result.statusCode).toBe(204); + expect(upstreamRequest).toEqual({ + accountId: "user-1", + method: "DELETE", + requestId: "request-id-1", + url: "http://petrinaut-opt:4004/optimize/runs/run-1", + }); + expect(entries).toContainEqual({ + level: "info", + message: "Petrinaut optimization run cancelled", + metadata: { + durationMs: expect.any(Number), + optimizationRunId: "run-1", + requestId: "request-id-1", + userId: "user-1", + }, + }); + }); + + it("passes through the optimizer's 404 for unknown or foreign runs", async () => { + const result = await callOptimizationRunHandler({ + handler: createHandler({ + fetchImpl: async () => + Response.json( + { detail: "optimization run not found: run-1" }, + { status: 404 }, + ), + }), + params: { runId: "run-1" }, + }); + + expect(result).toMatchObject({ + body: { error: "Optimization run not found" }, + statusCode: 404, + }); + }); + + it("maps an upstream cancellation failure to 502", async () => { + const { entries, logger } = createRecordingLogger(); + const result = await callOptimizationRunHandler({ + handler: createHandler({ + fetchImpl: async () => + Response.json({ detail: "cancellation failed" }, { status: 500 }), + logger, + }), + params: { runId: "run-1" }, + }); + + expect(result).toMatchObject({ + body: { error: "Petrinaut optimization failed" }, + statusCode: 502, + }); + expect(entries).toContainEqual({ + level: "warn", + message: "Petrinaut optimization run cancellation failed", + metadata: expect.objectContaining({ + optimizationRunId: "run-1", + outcome: "upstream-error", + userId: "user-1", + }), + }); + }); +}); diff --git a/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-cancel-handler.ts b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-cancel-handler.ts new file mode 100644 index 00000000000..9bdfc4cda6a --- /dev/null +++ b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-cancel-handler.ts @@ -0,0 +1,96 @@ +import { RESPONSE_START_TIMEOUT_MS } from "./shared/optimization-request-lifecycle"; +import { + resolveOptimizationRouteContext, + respondUpstreamFailure, + RUN_NOT_FOUND, +} from "./shared/optimization-route-context"; +import { capPathSegment } from "./shared/validate-optimization-request"; + +import type { Logger } from "@local/hash-backend-utils/logger"; +import type { PetrinautOptimizerClient } from "@local/petrinaut-optimizer-client"; +import type { RequestHandler } from "express"; + +/** + * Create the authenticated endpoint that cancels a detached optimization run. + * + * The optimizer enforces ownership: the account tag is forwarded and a + * foreign or unknown run answers 404 (never 403, so run ids cannot be + * probed). Cancellation is idempotent upstream — a run that already + * finished, was reaped, or expired still answers 204 while it is retained. + */ +export const createPetrinautOptimizationRunCancelHandler = ({ + client, + logger, + origin, +}: { + client: PetrinautOptimizerClient; + logger: Pick; + origin: URL | null; +}): RequestHandler => { + return async (request, response) => { + const context = resolveOptimizationRouteContext(request, response, { + logger, + origin, + }); + if (!context) { + return; + } + const { requestId, requestLogger, startedAt, userId } = context; + const runId = request.params.runId ?? ""; + // The run id is user-controlled; the optimizer decides whether it names + // a run this account owns. + const loggedRunId = capPathSegment(runId); + + const deadline = { timedOut: false }; + const cancelTimeout = AbortSignal.timeout(RESPONSE_START_TIMEOUT_MS); + cancelTimeout.addEventListener("abort", () => { + deadline.timedOut = true; + }); + + try { + const cancelled = await client.DELETE("/optimize/runs/{run_id}", { + params: { path: { run_id: runId } }, + headers: { + "x-hash-account-id": userId, + ...(requestId ? { "x-hash-request-id": requestId } : {}), + }, + signal: cancelTimeout, + }); + if (cancelled.response.status === 404) { + requestLogger.warn("Petrinaut optimization run cancel rejected", { + optimizationRunId: loggedRunId, + reason: "unknown-or-foreign-run", + userId, + }); + response.status(404).json(RUN_NOT_FOUND); + return; + } + if (!cancelled.response.ok) { + requestLogger.warn("Petrinaut optimization run cancellation failed", { + durationMs: Date.now() - startedAt, + optimizationRunId: loggedRunId, + outcome: "upstream-error", + upstreamStatus: cancelled.response.status, + userId, + }); + respondUpstreamFailure(response, false); + return; + } + requestLogger.info("Petrinaut optimization run cancelled", { + durationMs: Date.now() - startedAt, + optimizationRunId: runId, + userId, + }); + response.status(204).end(); + } catch (error) { + requestLogger.warn("Petrinaut optimization run cancellation failed", { + durationMs: Date.now() - startedAt, + error, + optimizationRunId: loggedRunId, + outcome: deadline.timedOut ? "timeout" : "upstream-error", + userId, + }); + respondUpstreamFailure(response, deadline.timedOut); + } + }; +}; diff --git a/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-events-handler.test.ts b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-events-handler.test.ts new file mode 100644 index 00000000000..dae54b44399 --- /dev/null +++ b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-events-handler.test.ts @@ -0,0 +1,270 @@ +import { describe, expect, it } from "vitest"; + +import { createPetrinautOptimizationRunEventsHandler } from "./create-petrinaut-optimization-run-events-handler"; +import { + callOptimizationRunHandler, + createRecordingLogger, + unexpectedFetch, +} from "./shared/optimization-run-test-harness"; + +import type { FakeResponse } from "./shared/optimization-run-test-harness"; +import type { PetrinautOptimizerFetch } from "@local/petrinaut-optimizer-client"; + +const createHandler = ({ + fetchImpl = unexpectedFetch, + logger = createRecordingLogger().logger, + origin = new URL("http://petrinaut-opt:4004"), +}: { + fetchImpl?: PetrinautOptimizerFetch; + logger?: ReturnType["logger"]; + origin?: URL | null; +} = {}) => + createPetrinautOptimizationRunEventsHandler({ + fetchImpl, + logger, + origin, + }); + +describe("createPetrinautOptimizationRunEventsHandler", () => { + it("requires authentication", async () => { + const result = await callOptimizationRunHandler({ + authenticated: false, + handler: createHandler(), + params: { runId: "run-1" }, + }); + + expect(result).toMatchObject({ + body: { error: "Authentication required" }, + statusCode: 401, + }); + }); + + it("rejects an invalid replay cursor", async () => { + const handler = createHandler(); + + for (const cursor of ["-1", "abc", "1.5", ["1", "2"]]) { + const result = await callOptimizationRunHandler({ + handler, + params: { runId: "run-1" }, + query: { cursor }, + }); + expect(result).toMatchObject({ + body: { error: "Invalid optimization cursor" }, + statusCode: 400, + }); + } + }); + + it("replays sequenced events as NDJSON with the account tag forwarded", async () => { + let upstreamRequest: + | { accountId: unknown; requestId: unknown; url: string } + | undefined; + const upstream = [ + 'id: 3\ndata: {"step":1,"params":{"rate":0.8},"init_state":{},"metric":4,"state":"COMPLETE"}\n\n', + ": heartbeat\n\n", + "id: 4\nevent: done\ndata: {}\n\n", + ].join(""); + const { entries, logger } = createRecordingLogger(); + const handler = createHandler({ + fetchImpl: async (input, init) => { + const headers = new Headers(init?.headers); + upstreamRequest = { + accountId: headers.get("x-hash-account-id"), + requestId: headers.get("x-hash-request-id"), + url: input.toString(), + }; + return new Response(upstream, { + headers: { + "content-type": "text/event-stream", + "x-optimization-run-id": "run-1", + "x-requested-trials": "2", + }, + }); + }, + logger, + }); + + const result = await callOptimizationRunHandler({ + handler, + params: { runId: "run-1" }, + query: { cursor: "2" }, + }); + + expect(result.statusCode).toBe(200); + expect(result.headers).toMatchObject({ + "Content-Type": "application/x-ndjson; charset=utf-8", + "X-Accel-Buffering": "no", + "X-Optimization-Run-ID": "run-1", + }); + // No synthetic started event, upstream sequence numbers visible as + // `seq`, and `best` stays null: the browser retains its own best. + expect(result.output).toEqual([ + '{"type":"trial","trial":1,"parameters":{"rate":0.8},"objective":4,"state":"complete","best":null,"seq":3}\n', + '{"type":"complete","requestedTrials":2,"completedTrials":1,"prunedTrials":0,"failedTrials":0,"best":null,"seq":4}\n', + ]); + expect(upstreamRequest?.url).toBe( + "http://petrinaut-opt:4004/optimize/runs/run-1/events?cursor=2", + ); + expect(upstreamRequest?.requestId).toBe("request-id-1"); + expect(upstreamRequest?.accountId).toBe("user-1"); + + // The attachment lifecycle is logged with its correlation identifiers. + expect(entries).toContainEqual({ + level: "info", + message: "Petrinaut optimization run attachment finished", + metadata: { + durationMs: expect.any(Number), + optimizationRunId: "run-1", + outcome: "completed", + requestId: "request-id-1", + userId: "user-1", + }, + }); + }); + + it("defaults to a full replay when no cursor is supplied", async () => { + let upstreamUrl: string | undefined; + const handler = createHandler({ + fetchImpl: async (input) => { + upstreamUrl = input.toString(); + return new Response("id: 1\nevent: done\ndata: {}\n\n", { + headers: { + "content-type": "text/event-stream", + "x-requested-trials": "2", + }, + }); + }, + }); + + const result = await callOptimizationRunHandler({ + handler, + params: { runId: "run-1" }, + }); + + expect(result.statusCode).toBe(200); + expect(upstreamUrl).toBe( + "http://petrinaut-opt:4004/optimize/runs/run-1/events?cursor=0", + ); + }); + + it("forwards a cancelled run's terminal frame", async () => { + const result = await callOptimizationRunHandler({ + handler: createHandler({ + fetchImpl: async () => + new Response("id: 5\nevent: cancelled\ndata: {}\n\n", { + headers: { "content-type": "text/event-stream" }, + }), + }), + params: { runId: "run-1" }, + }); + + expect(result.statusCode).toBe(200); + expect(result.output).toEqual([ + '{"type":"error","code":"optimization_cancelled","message":"The optimization was cancelled","retryable":false,"seq":5}\n', + ]); + }); + + it("forwards a superseded attachment's terminal event", async () => { + const result = await callOptimizationRunHandler({ + handler: createHandler({ + fetchImpl: async () => + new Response("event: superseded\ndata: {}\n\n", { + headers: { "content-type": "text/event-stream" }, + }), + }), + params: { runId: "run-1" }, + }); + + expect(result.statusCode).toBe(200); + expect(result.output).toEqual([ + '{"type":"error","code":"attachment_superseded","message":"Another consumer attached to this optimization run","retryable":false}\n', + ]); + }); + + it("passes through the optimizer's 404 for unknown or foreign runs", async () => { + const result = await callOptimizationRunHandler({ + handler: createHandler({ + fetchImpl: async () => + Response.json( + { detail: "optimization run not found: run-1" }, + { status: 404 }, + ), + }), + params: { runId: "run-1" }, + }); + + expect(result).toMatchObject({ + body: { error: "Optimization run not found" }, + statusCode: 404, + }); + }); + + it("commits the terminal write even when the client vanishes mid-drain", async () => { + const terminalLine = + '{"type":"complete","requestedTrials":2,"completedTrials":0,"prunedTrials":0,"failedTrials":0,"best":null,"seq":4}\n'; + const handler = createHandler({ + fetchImpl: async () => + new Response("id: 4\nevent: done\ndata: {}\n\n", { + headers: { + "content-type": "text/event-stream", + "x-requested-trials": "2", + }, + }), + }); + + let activeResponse: FakeResponse | undefined; + const first = await callOptimizationRunHandler({ + handler, + params: { runId: "run-1" }, + onResponse: (response) => { + activeResponse = response; + }, + writeReturns: (value) => { + if (value.includes('"type":"complete"')) { + // The terminal write is committed to the buffer, but the client + // disconnects before it ever drains. + setImmediate(() => { + activeResponse!.destroyed = true; + activeResponse!.emit("close"); + }); + return false; + } + return true; + }, + }); + + // The terminal line was committed exactly once — the failure path must + // not append a second terminal event after the client vanished. + expect(first.output).toEqual([terminalLine]); + + // A later attachment simply replays from the optimizer's retained log. + const second = await callOptimizationRunHandler({ + handler, + params: { runId: "run-1" }, + }); + expect(second.statusCode).toBe(200); + expect(second.output).toEqual([terminalLine]); + }); + + it("writes a retryable terminal line after a mid-stream failure", async () => { + const result = await callOptimizationRunHandler({ + handler: createHandler({ + fetchImpl: async () => + // The stream ends without a terminal event (e.g. the attachment + // was superseded or the optimizer restarted mid-flight). + new Response( + 'id: 3\ndata: {"step":1,"params":{"rate":0.8},"init_state":{},"metric":4,"state":"COMPLETE"}\n\n', + { headers: { "content-type": "text/event-stream" } }, + ), + }), + params: { runId: "run-1" }, + }); + + expect(result.statusCode).toBe(200); + // A retryable NodeAPI-authored terminal line tells the browser the run + // may still be live: re-attach using the last seen `seq` as the cursor. + expect(result.output.at(-1)).toBe( + '{"type":"error","code":"upstream_stream_error","message":"The optimizer stream ended unexpectedly","retryable":true}\n', + ); + }); +}); diff --git a/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-events-handler.ts b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-events-handler.ts new file mode 100644 index 00000000000..6e27aaed6c8 --- /dev/null +++ b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-events-handler.ts @@ -0,0 +1,200 @@ +import { + attachPetrinautOptimizationRunStream, + PetrinautOptimizerHttpError, +} from "@local/petrinaut-optimizer-client"; + +import { + forwardOptimizationEvents, + writeTerminalOptimizationEventBestEffort, +} from "./shared/forward-optimization-events"; +import { createOptimizationRequestLifecycle } from "./shared/optimization-request-lifecycle"; +import { + resolveOptimizationRouteContext, + respondUpstreamFailure, + RUN_NOT_FOUND, +} from "./shared/optimization-route-context"; +import { + capPathSegment, + MAX_EVENT_BYTES, +} from "./shared/validate-optimization-request"; + +import type { PetrinautOptimizationEvent } from "@hashintel/petrinaut-core"; +import type { Logger } from "@local/hash-backend-utils/logger"; +import type { PetrinautOptimizerFetch } from "@local/petrinaut-optimizer-client"; +import type { RequestHandler } from "express"; + +/** Parse the replay cursor; `undefined` means the query was invalid. */ +const parseCursor = (value: unknown): number | undefined => { + if (value === undefined) { + return 0; + } + if (typeof value !== "string" || !/^\d{1,15}$/.test(value)) { + return undefined; + } + return Number(value); +}; + +/** + * Create the authenticated endpoint that attaches to a detached run's events. + * + * The response replays buffered events past `?cursor=` as NDJSON, then + * live-tails new ones; each replayed line carries the upstream sequence + * number as `seq` so the browser can re-attach from where it stopped. + * The optimizer enforces ownership from the forwarded account tag: unknown + * run ids — including runs another account owns — answer 404 without + * revealing whether they exist. + */ +export const createPetrinautOptimizationRunEventsHandler = ({ + fetchImpl, + logger, + origin, +}: { + fetchImpl: PetrinautOptimizerFetch; + logger: Pick; + origin: URL | null; +}): RequestHandler => { + return async (request, response) => { + const context = resolveOptimizationRouteContext(request, response, { + logger, + origin, + }); + if (!context) { + return; + } + const { requestId, requestLogger, startedAt, userId } = context; + + const runId = request.params.runId ?? ""; + const cursor = parseCursor(request.query.cursor); + if (cursor === undefined) { + requestLogger.warn("Petrinaut optimization run attach rejected", { + // The run id is user-controlled; the optimizer owns the check. + optimizationRunId: capPathSegment(runId), + reason: "invalid-cursor", + userId, + }); + response.status(400).json({ error: "Invalid optimization cursor" }); + return; + } + + requestLogger.info("Petrinaut optimization run attach requested", { + cursor, + optimizationRunId: runId, + userId, + }); + + const lifecycle = createOptimizationRequestLifecycle(request, response); + + try { + const upstream = await attachPetrinautOptimizationRunStream({ + endpoint: context.origin, + runId, + cursor, + fetchImpl, + // The optimizer enforces that this account owns the run; it also + // reports the manifest's requested trial count on the response so + // the synthesized summary is sized without NodeAPI remembering the + // manifest. No direction is passed, so `best` stays null (the + // browser keeps its own running best across reconnections). + headers: { "x-hash-account-id": userId }, + maxEventBytes: MAX_EVENT_BYTES, + onActivity: lifecycle.resetIdleTimeout, + ...(requestId ? { requestId } : {}), + signal: lifecycle.signal, + }); + requestLogger.info("Petrinaut optimization run attached", { + cursor, + optimizationRunId: runId, + userId, + }); + + response.status(200); + response.set({ + "Cache-Control": "no-cache, no-store", + "Content-Type": "application/x-ndjson; charset=utf-8", + "X-Accel-Buffering": "no", + "X-Optimization-Run-ID": runId, + }); + response.flushHeaders(); + lifecycle.markResponseStarted(); + await forwardOptimizationEvents( + upstream.events, + response, + lifecycle.markTerminalEvent, + lifecycle.signal, + ); + response.end(); + requestLogger.info("Petrinaut optimization run attachment finished", { + durationMs: Date.now() - startedAt, + optimizationRunId: runId, + outcome: "completed", + userId, + }); + } catch (error) { + const durationMs = Date.now() - startedAt; + if (lifecycle.state.clientDisconnected || response.destroyed) { + // The run itself is unaffected: the browser can simply re-attach. + requestLogger.info("Petrinaut optimization run attachment finished", { + durationMs, + optimizationRunId: runId, + outcome: "client-disconnected", + userId, + }); + return; + } + const runGoneUpstream = + error instanceof PetrinautOptimizerHttpError && error.status === 404; + requestLogger.warn("Petrinaut optimization run attachment failed", { + durationMs, + error, + optimizationRunId: runId, + outcome: lifecycle.state.timeoutKind + ? `timeout:${lifecycle.state.timeoutKind}` + : runGoneUpstream + ? "run-gone" + : "upstream-error", + timeoutKind: lifecycle.state.timeoutKind, + userId, + }); + if (!response.headersSent) { + if (runGoneUpstream) { + response.status(404).json(RUN_NOT_FOUND); + } else { + respondUpstreamFailure( + response, + Boolean(lifecycle.state.timeoutKind), + ); + } + return; + } + if (!lifecycle.state.terminalEventSent) { + try { + // `retryable: true` tells the browser the run may still be live: + // re-attach with the last seen `seq` as the cursor. + writeTerminalOptimizationEventBestEffort(response, { + type: "error", + code: lifecycle.state.timeoutKind + ? "optimization_timeout" + : "upstream_stream_error", + message: lifecycle.state.timeoutKind + ? "The optimization stream attachment timed out" + : "The optimizer stream ended unexpectedly", + retryable: true, + } satisfies PetrinautOptimizationEvent); + } catch (writeError) { + requestLogger.warn( + "Could not report Petrinaut optimization run attachment failure", + { + error: writeError, + optimizationRunId: runId, + }, + ); + } + } + if (!response.writableEnded) { + response.end(); + } + } finally { + lifecycle.cleanup(); + } + }; +}; diff --git a/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-handler.test.ts b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-handler.test.ts new file mode 100644 index 00000000000..940b9840b0e --- /dev/null +++ b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-handler.test.ts @@ -0,0 +1,352 @@ +import { describe, expect, it } from "vitest"; + +import { createPetrinautOptimizerClient } from "@local/petrinaut-optimizer-client"; + +import { createPetrinautOptimizationRunHandler } from "./create-petrinaut-optimization-run-handler"; +import { + callOptimizationRunHandler, + createRecordingLogger, + unexpectedFetch, + validOptimizationInput, +} from "./shared/optimization-run-test-harness"; + +import type { PetrinautOptimizerFetch } from "@local/petrinaut-optimizer-client"; +import type { EventEmitter } from "node:events"; + +const createHandler = ({ + fetchImpl = unexpectedFetch, + logger = createRecordingLogger().logger, + origin = new URL("http://petrinaut-opt:4004"), +}: { + fetchImpl?: PetrinautOptimizerFetch; + logger?: ReturnType["logger"]; + origin?: URL | null; +} = {}) => + createPetrinautOptimizationRunHandler({ + client: createPetrinautOptimizerClient( + origin ?? "http://petrinaut-optimizer.invalid", + fetchImpl, + ), + logger, + origin, + }); + +describe("createPetrinautOptimizationRunHandler", () => { + it("requires authentication", async () => { + const result = await callOptimizationRunHandler({ + authenticated: false, + body: {}, + handler: createHandler(), + }); + + expect(result).toMatchObject({ + body: { error: "Authentication required" }, + statusCode: 401, + }); + }); + + it("reports an unconfigured optimizer", async () => { + const result = await callOptimizationRunHandler({ + body: validOptimizationInput, + handler: createHandler({ origin: null }), + }); + + expect(result).toMatchObject({ + body: { error: "Petrinaut optimizer is not configured" }, + statusCode: 503, + }); + }); + + it("validates the public optimization request without logging it", async () => { + const { entries, logger } = createRecordingLogger(); + const result = await callOptimizationRunHandler({ + body: { + ...validOptimizationInput, + scenario: { + ...validOptimizationInput.scenario, + parameterBindings: { + rate: { + kind: "optimize", + domain: { + kind: "continuous", + minimum: 0.1, + maximum: 1, + scale: "sqrt", + }, + }, + }, + }, + }, + handler: createHandler({ logger }), + }); + + expect(result.statusCode).toBe(400); + expect(result.body).toEqual({ + code: "invalid_optimization_request", + details: { + issues: [ + { + message: expect.any(String), + path: "scenario.parameterBindings.rate.domain.scale", + }, + ], + truncated: false, + }, + error: "Invalid optimization request", + }); + + // The failure is logged with only issue counts and schema paths — never + // the validation messages or the manifest, which can embed user code. + expect(entries).toEqual([ + { + level: "warn", + message: "Petrinaut optimization request failed validation", + metadata: { + issueCount: 1, + issuePaths: ["scenario.parameterBindings.rate.domain.scale"], + issuePathsTruncated: false, + requestId: "request-id-1", + userId: "user-1", + }, + }, + ]); + const loggedText = JSON.stringify(entries); + expect(loggedText).not.toContain("return 1;"); + expect(loggedText).not.toContain("petrinaut-optimization"); + }); + + it("forwards the manifest with the account tag and returns the run id", async () => { + let upstreamRequest: + | { + accountId: unknown; + body: string | undefined; + requestId: unknown; + url: string; + } + | undefined; + const { entries, logger } = createRecordingLogger(); + const handler = createHandler({ + fetchImpl: async (input, init) => { + const headers = new Headers(init?.headers); + upstreamRequest = { + accountId: headers.get("x-hash-account-id"), + body: typeof init?.body === "string" ? init.body : undefined, + requestId: headers.get("x-hash-request-id"), + url: input.toString(), + }; + return Response.json( + { run_id: "run-1" }, + { status: 201, headers: { "x-optimization-run-id": "run-1" } }, + ); + }, + logger, + }); + + const result = await callOptimizationRunHandler({ + body: validOptimizationInput, + handler, + }); + + expect(result.statusCode).toBe(201); + expect(result.body).toEqual({ runId: "run-1" }); + expect(result.headers).toMatchObject({ "X-Optimization-Run-ID": "run-1" }); + expect(upstreamRequest?.url).toBe( + "http://petrinaut-opt:4004/optimize/runs", + ); + expect(upstreamRequest?.accountId).toBe("user-1"); + expect(upstreamRequest?.requestId).toBe("request-id-1"); + expect(JSON.parse(upstreamRequest?.body ?? "null")).toEqual( + validOptimizationInput, + ); + expect(entries).toEqual([ + { + level: "info", + message: "Petrinaut optimization run requested", + metadata: { + bodyBytes: expect.any(Number), + requestId: "request-id-1", + requestedTrials: 2, + userId: "user-1", + }, + }, + { + level: "info", + message: "Petrinaut optimization run created", + metadata: { + durationMs: expect.any(Number), + optimizationRunId: "run-1", + requestId: "request-id-1", + userId: "user-1", + }, + }, + ]); + // The manifest itself — including user-authored code — is never logged. + const loggedText = JSON.stringify(entries); + expect(loggedText).not.toContain("return 1;"); + expect(loggedText).not.toContain("petrinaut-optimization"); + }); + + it("forwards the optimizer's 429 detail and Retry-After", async () => { + const result = await callOptimizationRunHandler({ + body: validOptimizationInput, + handler: createHandler({ + fetchImpl: async () => + Response.json( + { detail: "An optimization is already running for this account" }, + { status: 429, headers: { "retry-after": "30" } }, + ), + }), + }); + + // The optimizer's server-authored detail distinguishes account + // single-flight from service capacity, so it is forwarded verbatim. + expect(result).toMatchObject({ + body: { error: "An optimization is already running for this account" }, + headers: { "Retry-After": "30" }, + statusCode: 429, + }); + }); + + it("correlates an upstream failure with the optimizer run id", async () => { + const { entries, logger } = createRecordingLogger(); + const result = await callOptimizationRunHandler({ + body: validOptimizationInput, + handler: createHandler({ + fetchImpl: async () => + Response.json( + { detail: "failed to initialise optimization" }, + { status: 500, headers: { "x-optimization-run-id": "run-err-9" } }, + ), + logger, + }), + }); + + expect(result.statusCode).toBe(502); + expect(result.body).toEqual({ error: "Petrinaut optimization failed" }); + expect(result.headers).toMatchObject({ + "X-Optimization-Run-ID": "run-err-9", + }); + const failure = entries.find( + (entry) => entry.message === "Petrinaut optimization run creation failed", + ); + expect(failure?.metadata).toMatchObject({ + optimizationRunId: "run-err-9", + outcome: "upstream-error", + }); + }); + + describe("client disconnects", () => { + it("abandons a create whose client disconnects mid-round-trip", async () => { + const { entries, logger } = createRecordingLogger(); + let abortObserved = false; + const handler = createHandler({ + fetchImpl: async (_input, init) => + new Promise((_resolve, reject) => { + const abort = () => { + abortObserved = true; + reject(new DOMException("Aborted", "AbortError")); + }; + // Real fetch rejects a signal that aborted before the call too. + if (init?.signal?.aborted) { + abort(); + return; + } + init?.signal?.addEventListener("abort", abort, { once: true }); + }), + logger, + }); + + const result = await callOptimizationRunHandler({ + body: validOptimizationInput, + handler, + onRequest: (request) => request.emit("aborted"), + }); + + // No response is written to the vanished client… + expect(abortObserved).toBe(true); + expect(result.body).toBeUndefined(); + expect(result.response.headersSent).toBe(false); + // …and the outcome is a disconnect, not an upstream error. + expect(entries.at(-1)).toEqual({ + level: "info", + message: "Petrinaut optimization run creation finished", + metadata: { + durationMs: expect.any(Number), + outcome: "client-disconnected", + requestId: "request-id-1", + userId: "user-1", + }, + }); + }); + + it("cancels a run created for a client that already disconnected", async () => { + const { entries, logger } = createRecordingLogger(); + const calls: { accountId: unknown; method: string; url: string }[] = []; + let resolveCreate: (() => void) | undefined; + const handler = createHandler({ + fetchImpl: async (input, init) => { + const method = init?.method ?? "GET"; + calls.push({ + accountId: new Headers(init?.headers).get("x-hash-account-id"), + method, + url: input.toString(), + }); + if (method === "DELETE") { + return new Response(null, { status: 204 }); + } + // Admit the run only after the client has already vanished. + await new Promise((resolve) => { + resolveCreate = resolve; + }); + return Response.json( + { run_id: "run-orphan" }, + { status: 201, headers: { "x-optimization-run-id": "run-orphan" } }, + ); + }, + logger, + }); + + let activeRequest: EventEmitter | undefined; + const resultPromise = callOptimizationRunHandler({ + body: validOptimizationInput, + handler, + onRequest: (request) => { + activeRequest = request; + }, + }); + // Reach the pending upstream create, disconnect, then let it resolve. + await new Promise(setImmediate); + activeRequest!.emit("aborted"); + resolveCreate?.(); + const result = await resultPromise; + + // The orphaned run is cancelled upstream on the account's behalf, so + // the optimizer frees its single-flight immediately. + expect(calls).toEqual([ + { + accountId: "user-1", + method: "POST", + url: "http://petrinaut-opt:4004/optimize/runs", + }, + { + accountId: "user-1", + method: "DELETE", + url: "http://petrinaut-opt:4004/optimize/runs/run-orphan", + }, + ]); + expect(result.body).toBeUndefined(); + expect(result.response.headersSent).toBe(false); + expect(entries.at(-1)).toEqual({ + level: "info", + message: "Petrinaut optimization run creation finished", + metadata: { + durationMs: expect.any(Number), + optimizationRunId: "run-orphan", + outcome: "client-disconnected", + requestId: "request-id-1", + userId: "user-1", + }, + }); + }); + }); +}); diff --git a/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-handler.ts b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-handler.ts new file mode 100644 index 00000000000..31742b42f63 --- /dev/null +++ b/apps/hash-api/src/petrinaut-optimizer/create-petrinaut-optimization-run-handler.ts @@ -0,0 +1,189 @@ +import { RESPONSE_START_TIMEOUT_MS } from "./shared/optimization-request-lifecycle"; +import { + resolveOptimizationRouteContext, + respondUpstreamFailure, +} from "./shared/optimization-route-context"; +import { validateOptimizationRequest } from "./shared/validate-optimization-request"; + +import type { Logger } from "@local/hash-backend-utils/logger"; +import type { PetrinautOptimizerClient } from "@local/petrinaut-optimizer-client"; +import type { RequestHandler } from "express"; + +/** + * Create the authenticated endpoint that starts a detached optimization run. + * + * NodeAPI only authenticates, rate-limits, and validates: the manifest is + * forwarded with the account's tag (`x-hash-account-id`), and the optimizer + * itself owns run state — per-account single-flight at admission, and + * owner-only visibility on attach/cancel. The response carries no events: + * the browser attaches to `GET …/optimize/runs/:runId/events` (and + * re-attaches after a disconnect) to consume them. + */ +export const createPetrinautOptimizationRunHandler = ({ + client, + logger, + origin, +}: { + client: PetrinautOptimizerClient; + logger: Pick; + origin: URL | null; +}): RequestHandler => { + return async (request, response) => { + const context = resolveOptimizationRouteContext(request, response, { + logger, + origin, + }); + if (!context) { + return; + } + const { requestId, requestLogger, startedAt, userId } = context; + + const validation = validateOptimizationRequest({ + body: request.body, + requestLogger, + userId, + }); + if (!validation.ok) { + response.status(validation.status).json(validation.body); + return; + } + const { bodyBytes, input } = validation; + + requestLogger.info("Petrinaut optimization run requested", { + bodyBytes, + requestedTrials: input.study.trials, + userId, + }); + + // Creation is a single round-trip: bound its time to the upstream + // answer. The controller also fires when the HASH client disconnects + // while the round-trip is still in flight. + const abortController = new AbortController(); + const outcome = { clientDisconnected: false, timedOut: false }; + const abortForClientDisconnect = () => { + outcome.clientDisconnected = true; + abortController.abort(); + }; + request.once("aborted", abortForClientDisconnect); + response.once("close", abortForClientDisconnect); + const createTimeout = setTimeout(() => { + outcome.timedOut = true; + abortController.abort(); + }, RESPONSE_START_TIMEOUT_MS); + + try { + const created = await client.POST("/optimize/runs", { + body: input, + headers: { + "x-hash-account-id": userId, + ...(requestId ? { "x-hash-request-id": requestId } : {}), + }, + signal: abortController.signal, + }); + const runId = created.data?.run_id; + if (!created.response.ok || !runId) { + const durationMs = Date.now() - startedAt; + const upstreamRunId = created.response.headers.get( + "x-optimization-run-id", + ); + requestLogger.warn("Petrinaut optimization run creation failed", { + durationMs, + optimizationRunId: upstreamRunId, + outcome: "upstream-error", + upstreamStatus: created.response.status, + userId, + }); + if (upstreamRunId !== null) { + response.set({ "X-Optimization-Run-ID": upstreamRunId }); + } + if (created.response.status === 429) { + const retryAfter = created.response.headers.get("retry-after"); + if (retryAfter) { + response.set({ "Retry-After": retryAfter }); + } + // The optimizer's detail is server-authored and distinguishes + // "your account already runs one" from "the service is at + // capacity", so it is forwarded rather than flattened. + response.status(429).json({ + error: + typeof created.error?.detail === "string" + ? created.error.detail + : "Petrinaut optimizer is busy", + }); + } else { + respondUpstreamFailure(response, false); + } + return; + } + + if (outcome.clientDisconnected || response.destroyed) { + // The run was admitted upstream but nobody will ever learn its id. + // Cancel it best-effort so the account is not blocked for the rest + // of the orphan grace period; the optimizer's reaper is the backstop. + const cancelled = await client + .DELETE("/optimize/runs/{run_id}", { + params: { path: { run_id: runId } }, + headers: { + "x-hash-account-id": userId, + ...(requestId ? { "x-hash-request-id": requestId } : {}), + }, + signal: AbortSignal.timeout(RESPONSE_START_TIMEOUT_MS), + }) + .catch((error: unknown) => { + requestLogger.warn( + "Could not cancel abandoned Petrinaut optimization run", + { error, optimizationRunId: runId, userId }, + ); + return null; + }); + if (cancelled && !cancelled.response.ok) { + requestLogger.warn( + "Could not cancel abandoned Petrinaut optimization run", + { + optimizationRunId: runId, + upstreamStatus: cancelled.response.status, + userId, + }, + ); + } + requestLogger.info("Petrinaut optimization run creation finished", { + durationMs: Date.now() - startedAt, + optimizationRunId: runId, + outcome: "client-disconnected", + userId, + }); + return; + } + + requestLogger.info("Petrinaut optimization run created", { + durationMs: Date.now() - startedAt, + optimizationRunId: runId, + userId, + }); + response.status(201); + response.set({ "X-Optimization-Run-ID": runId }); + response.json({ runId }); + } catch (error) { + const durationMs = Date.now() - startedAt; + if (outcome.clientDisconnected || response.destroyed) { + requestLogger.info("Petrinaut optimization run creation finished", { + durationMs, + outcome: "client-disconnected", + userId, + }); + return; + } + requestLogger.warn("Petrinaut optimization run creation failed", { + durationMs, + error, + outcome: outcome.timedOut ? "timeout" : "upstream-error", + userId, + }); + respondUpstreamFailure(response, outcome.timedOut); + } finally { + clearTimeout(createTimeout); + request.off("aborted", abortForClientDisconnect); + response.off("close", abortForClientDisconnect); + } + }; +}; diff --git a/apps/hash-api/src/petrinaut-optimizer/setup-petrinaut-optimizer-handler.test.ts b/apps/hash-api/src/petrinaut-optimizer/setup-petrinaut-optimizer-handler.test.ts index 0d966abfd4d..027ebf99642 100644 --- a/apps/hash-api/src/petrinaut-optimizer/setup-petrinaut-optimizer-handler.test.ts +++ b/apps/hash-api/src/petrinaut-optimizer/setup-petrinaut-optimizer-handler.test.ts @@ -3,7 +3,9 @@ import { describe, expect, it } from "vitest"; import { getPetrinautOptimizerOrigin, PETRINAUT_OPTIMIZER_CAPABILITIES_PATH, - PETRINAUT_OPTIMIZER_OPTIMIZE_PATH, + PETRINAUT_OPTIMIZER_OPTIMIZE_RUN_EVENTS_PATH, + PETRINAUT_OPTIMIZER_OPTIMIZE_RUN_PATH, + PETRINAUT_OPTIMIZER_OPTIMIZE_RUNS_PATH, setupPetrinautOptimizerHandler, } from "./setup-petrinaut-optimizer-handler"; @@ -24,6 +26,30 @@ type Handler = ( response: ExpressResponse, ) => Promise | void; +type RegisteredRoute = { + handlerCount: number; + method: "delete" | "get" | "post"; + path: string; +}; + +/** Build an Express fake that records every route registration in order. */ +const createRecordingApp = () => { + const routes: RegisteredRoute[] = []; + const handlers = new Map(); + const register = + (method: RegisteredRoute["method"]) => + (path: string, ...routeHandlers: unknown[]) => { + routes.push({ handlerCount: routeHandlers.length, method, path }); + handlers.set(`${method} ${path}`, routeHandlers.at(-1) as Handler); + }; + const app = { + delete: register("delete"), + get: register("get"), + post: register("post"), + } as unknown as Express; + return { app, handlers, routes }; +}; + const callGetHandler = async ({ authenticated = true, origin, @@ -31,17 +57,9 @@ const callGetHandler = async ({ authenticated?: boolean; origin: URL | null; }) => { - let handler: Handler | undefined; - const app = { - get: (routePath: string, routeHandler: Handler) => { - if (routePath === PETRINAUT_OPTIMIZER_CAPABILITIES_PATH) { - handler = routeHandler; - } - }, - post: () => undefined, - } as unknown as Express; - + const { app, handlers } = createRecordingApp(); setupPetrinautOptimizerHandler(app, { logger, origin }); + const handler = handlers.get(`get ${PETRINAUT_OPTIMIZER_CAPABILITIES_PATH}`); if (!handler) { throw new Error("The capabilities route was not registered"); } @@ -88,21 +106,33 @@ describe("getPetrinautOptimizerOrigin", () => { }); }); -it("mounts the optimization endpoint", () => { - let registeredRoute: { handlerCount: number; path: string } | undefined; - const app = { - get: () => undefined, - post: (path: string, ...handlers: unknown[]) => { - registeredRoute = { handlerCount: handlers.length, path }; - }, - } as unknown as Express; +it("mounts every optimization endpoint behind the rate limiter", () => { + const { app, routes } = createRecordingApp(); setupPetrinautOptimizerHandler(app, { logger, origin: null }); - expect(registeredRoute).toEqual({ - handlerCount: 2, - path: PETRINAUT_OPTIMIZER_OPTIMIZE_PATH, - }); + expect(routes).toEqual([ + { + handlerCount: 1, + method: "get", + path: PETRINAUT_OPTIMIZER_CAPABILITIES_PATH, + }, + { + handlerCount: 2, + method: "post", + path: PETRINAUT_OPTIMIZER_OPTIMIZE_RUNS_PATH, + }, + { + handlerCount: 2, + method: "get", + path: PETRINAUT_OPTIMIZER_OPTIMIZE_RUN_EVENTS_PATH, + }, + { + handlerCount: 2, + method: "delete", + path: PETRINAUT_OPTIMIZER_OPTIMIZE_RUN_PATH, + }, + ]); }); describe(PETRINAUT_OPTIMIZER_CAPABILITIES_PATH, () => { diff --git a/apps/hash-api/src/petrinaut-optimizer/setup-petrinaut-optimizer-handler.ts b/apps/hash-api/src/petrinaut-optimizer/setup-petrinaut-optimizer-handler.ts index 810badd70b5..2b522660bb9 100644 --- a/apps/hash-api/src/petrinaut-optimizer/setup-petrinaut-optimizer-handler.ts +++ b/apps/hash-api/src/petrinaut-optimizer/setup-petrinaut-optimizer-handler.ts @@ -1,31 +1,58 @@ import { ipKeyGenerator, rateLimit } from "express-rate-limit"; -import { createPetrinautOptimizationHandler } from "./create-petrinaut-optimization-handler"; +import { createPetrinautOptimizerClient } from "@local/petrinaut-optimizer-client"; + +import { createPetrinautOptimizationRunCancelHandler } from "./create-petrinaut-optimization-run-cancel-handler"; +import { createPetrinautOptimizationRunEventsHandler } from "./create-petrinaut-optimization-run-events-handler"; +import { createPetrinautOptimizationRunHandler } from "./create-petrinaut-optimization-run-handler"; import type { Logger } from "@local/hash-backend-utils/logger"; import type { PetrinautOptimizerFetch } from "@local/petrinaut-optimizer-client"; -import type { Express } from "express"; +import type { Express, Request } from "express"; export const PETRINAUT_OPTIMIZER_CAPABILITIES_PATH = "/api/petrinaut-optimizer/capabilities"; -export const PETRINAUT_OPTIMIZER_OPTIMIZE_PATH = - "/api/petrinaut-optimizer/optimize"; +export const PETRINAUT_OPTIMIZER_OPTIMIZE_RUNS_PATH = + "/api/petrinaut-optimizer/optimize/runs"; +export const PETRINAUT_OPTIMIZER_OPTIMIZE_RUN_PATH = + "/api/petrinaut-optimizer/optimize/runs/:runId"; +export const PETRINAUT_OPTIMIZER_OPTIMIZE_RUN_EVENTS_PATH = + "/api/petrinaut-optimizer/optimize/runs/:runId/events"; type PetrinautOptimizerHandlerOptions = { origin: URL | null; fetchImpl?: PetrinautOptimizerFetch; - logger: Pick; + logger: Pick; }; +const rateLimitKeyGenerator = (request: Request) => + request.user?.accountId ?? + (request.ip ? ipKeyGenerator(request.ip) : "ip-unavailable"); + /** Bound expensive optimization attempts per authenticated account or IP. */ const optimizationRateLimiter = rateLimit({ windowMs: process.env.NODE_ENV === "test" ? 10 : 60_000, limit: 10, standardHeaders: true, legacyHeaders: false, - keyGenerator: (request) => - request.user?.accountId ?? - (request.ip ? ipKeyGenerator(request.ip) : "ip-unavailable"), + keyGenerator: rateLimitKeyGenerator, + message: { error: "Too many optimization requests" }, +}); + +/** + * Attaching and cancelling do no expensive upstream work, and both must stay + * available to a busy account: the auto-reconnect loop may legitimately + * re-attach many times behind a flaky connection, and cancel is how an + * account frees its own single-flight slot. Sharing creation's bucket would + * 429 an account off the event stream of — or out of cancelling — its own + * live run. + */ +const nonAdmissionRateLimiter = rateLimit({ + windowMs: process.env.NODE_ENV === "test" ? 10 : 60_000, + limit: 60, + standardHeaders: true, + legacyHeaders: false, + keyGenerator: rateLimitKeyGenerator, message: { error: "Too many optimization requests" }, }); @@ -57,7 +84,19 @@ export const getPetrinautOptimizerOrigin = ( return new URL(`http://${urlHost}:${port}`); }; -/** Mount authenticated NodeAPI routes for Petrinaut optimization. */ +/** + * Mount authenticated NodeAPI routes for Petrinaut optimization. + * + * NodeAPI is a thin, stateless proxy: it authenticates, rate-limits, and + * validates, then forwards to the optimizer with the account's tag + * (`x-hash-account-id`). The optimizer owns all run state — per-account + * single-flight at admission, and owner-only visibility on attach/cancel. + * + * `POST …/optimize/runs` returns a run id immediately; events are consumed + * (and resumed via `?cursor=`) through `GET …/optimize/runs/:runId/events` + * attachments that never affect the run; `DELETE …/optimize/runs/:runId` + * cancels explicitly. + */ export const setupPetrinautOptimizerHandler = ( app: Express, { origin, fetchImpl = fetch, logger }: PetrinautOptimizerHandlerOptions, @@ -78,16 +117,30 @@ export const setupPetrinautOptimizerHandler = ( res.json({ optimization: origin !== null }); }); - /** - * Validate and proxy one authenticated optimization request. - * - * The route enforces request and concurrency limits, owns cancellation and - * execution deadlines, and converts the shared client's canonical events to - * the NDJSON protocol consumed by the HASH frontend. - */ + // A dummy base keeps the client constructible when the optimizer is not + // configured; the handlers answer 503 before ever using it then. + const client = createPetrinautOptimizerClient( + origin ?? "http://petrinaut-optimizer.invalid", + fetchImpl, + ); + app.post( - PETRINAUT_OPTIMIZER_OPTIMIZE_PATH, + PETRINAUT_OPTIMIZER_OPTIMIZE_RUNS_PATH, optimizationRateLimiter, - createPetrinautOptimizationHandler({ fetchImpl, logger, origin }), + createPetrinautOptimizationRunHandler({ client, logger, origin }), + ); + app.get( + PETRINAUT_OPTIMIZER_OPTIMIZE_RUN_EVENTS_PATH, + nonAdmissionRateLimiter, + createPetrinautOptimizationRunEventsHandler({ + fetchImpl, + logger, + origin, + }), + ); + app.delete( + PETRINAUT_OPTIMIZER_OPTIMIZE_RUN_PATH, + nonAdmissionRateLimiter, + createPetrinautOptimizationRunCancelHandler({ client, logger, origin }), ); }; diff --git a/apps/hash-api/src/petrinaut-optimizer/shared/forward-optimization-events.ts b/apps/hash-api/src/petrinaut-optimizer/shared/forward-optimization-events.ts new file mode 100644 index 00000000000..26a273b58d5 --- /dev/null +++ b/apps/hash-api/src/petrinaut-optimizer/shared/forward-optimization-events.ts @@ -0,0 +1,77 @@ +import { once } from "node:events"; + +import type { PetrinautOptimizationEvent } from "@hashintel/petrinaut-core"; +import type { Response as ExpressResponse } from "express"; + +/** Write one canonical optimization event as NDJSON with backpressure. */ +export const writeOptimizationEvent = async ( + response: ExpressResponse, + event: PetrinautOptimizationEvent, + signal: AbortSignal, +): Promise => { + // This is schema-validated NDJSON served with `nosniff`, never HTML. + // nosemgrep: javascript.express.security.audit.xss.direct-response-write.direct-response-write + if (response.write(`${JSON.stringify(event)}\n`)) { + return; + } + if (response.destroyed || response.writableEnded) { + throw new Error("The optimization client disconnected"); + } + if (signal.aborted) { + throw new Error("The optimization request was aborted"); + } + // Aborting this controller removes whichever listeners lost the race, so a + // backpressured stream cannot accumulate `drain`/`close` listeners and the + // loser can never surface a late unhandled rejection. + const waitCleanup = new AbortController(); + try { + await Promise.race([ + once(response, "drain", { signal: waitCleanup.signal }), + once(response, "close", { signal: waitCleanup.signal }).then(() => { + throw new Error("The optimization client disconnected"); + }), + once(signal, "abort", { signal: waitCleanup.signal }).then(() => { + throw new Error("The optimization request was aborted"); + }), + ]); + } finally { + waitCleanup.abort(); + } +}; + +/** + * Write a final event without waiting for backpressure to clear. + * + * Teardown must never block on a stalled client, so this write is + * best-effort: the response ends immediately afterwards either way. + */ +export const writeTerminalOptimizationEventBestEffort = ( + response: ExpressResponse, + event: PetrinautOptimizationEvent, +): void => { + if (response.destroyed || response.writableEnded) { + return; + } + // This is schema-validated NDJSON served with `nosniff`, never HTML. + // nosemgrep: javascript.express.security.audit.xss.direct-response-write.direct-response-write + response.write(`${JSON.stringify(event)}\n`); +}; + +/** Forward canonical optimization events as NDJSON. */ +export const forwardOptimizationEvents = async ( + events: AsyncIterable, + response: ExpressResponse, + markTerminalEvent: () => void, + signal: AbortSignal, +): Promise => { + for await (const event of events) { + const backpressureWait = writeOptimizationEvent(response, event, signal); + if (event.type === "complete" || event.type === "error") { + // The write is committed synchronously even when the buffer is full, + // so the stream contains its terminal event regardless of how the + // backpressure wait settles — an abort must not append a second one. + markTerminalEvent(); + } + await backpressureWait; + } +}; diff --git a/apps/hash-api/src/petrinaut-optimizer/shared/optimization-request-lifecycle.ts b/apps/hash-api/src/petrinaut-optimizer/shared/optimization-request-lifecycle.ts new file mode 100644 index 00000000000..df068ff4e3f --- /dev/null +++ b/apps/hash-api/src/petrinaut-optimizer/shared/optimization-request-lifecycle.ts @@ -0,0 +1,104 @@ +import type { Request, Response as ExpressResponse } from "express"; + +export const RESPONSE_START_TIMEOUT_MS = 30_000; +export const DOWNSTREAM_HEARTBEAT_INTERVAL_MS = 25_000; +export const IDLE_TIMEOUT_MS = 5 * 60_000; +export const OVERALL_TIMEOUT_MS = 15 * 60_000; + +export type OptimizationTimeoutKind = "response_start" | "idle" | "overall"; + +export type OptimizationRequestLifecycle = { + /** Signal that cancels the upstream optimizer request. */ + signal: AbortSignal; + /** Mutable outcome needed when translating a stream failure. */ + state: { + clientDisconnected: boolean; + terminalEventSent: boolean; + timeoutKind: OptimizationTimeoutKind | null; + }; + /** Stop timers, heartbeats, and request/response listeners. */ + cleanup: () => void; + /** Clear the response-start deadline and begin idle/heartbeat tracking. */ + markResponseStarted: () => void; + /** Remember that the public stream already contains its terminal event. */ + markTerminalEvent: () => void; + /** Restart the deadline for receiving upstream bytes. */ + resetIdleTimeout: () => void; +}; + +/** Track cancellation, deadlines, and cleanup for one streamed request. */ +export const createOptimizationRequestLifecycle = ( + request: Request, + response: ExpressResponse, +): OptimizationRequestLifecycle => { + const abortController = new AbortController(); + const state: OptimizationRequestLifecycle["state"] = { + clientDisconnected: false, + terminalEventSent: false, + timeoutKind: null, + }; + let downstreamHeartbeat: ReturnType | undefined; + let idleTimeout: ReturnType | undefined; + + /** Abort the upstream request after the HASH client disconnects. */ + const abortForClientDisconnect = () => { + state.clientDisconnected = true; + abortController.abort(); + }; + + /** Record a timeout category and abort the upstream request. */ + const abortForTimeout = (kind: OptimizationTimeoutKind) => { + state.timeoutKind ??= kind; + abortController.abort(); + }; + + const responseStartTimeout = setTimeout( + () => abortForTimeout("response_start"), + RESPONSE_START_TIMEOUT_MS, + ); + const overallTimeout = setTimeout( + () => abortForTimeout("overall"), + OVERALL_TIMEOUT_MS, + ); + + /** Restart the inactivity deadline after any upstream bytes arrive. */ + const resetIdleTimeout = () => { + clearTimeout(idleTimeout); + idleTimeout = setTimeout(() => abortForTimeout("idle"), IDLE_TIMEOUT_MS); + }; + + request.once("aborted", abortForClientDisconnect); + response.once("close", abortForClientDisconnect); + + return { + signal: abortController.signal, + state, + cleanup: () => { + clearTimeout(responseStartTimeout); + clearInterval(downstreamHeartbeat); + clearTimeout(idleTimeout); + clearTimeout(overallTimeout); + request.off("aborted", abortForClientDisconnect); + response.off("close", abortForClientDisconnect); + }, + markResponseStarted: () => { + clearTimeout(responseStartTimeout); + resetIdleTimeout(); + downstreamHeartbeat = setInterval(() => { + if ( + !response.destroyed && + !response.writableEnded && + !response.writableNeedDrain + ) { + // Blank NDJSON lines are transport heartbeats, not domain events. + response.write("\n"); + } + }, DOWNSTREAM_HEARTBEAT_INTERVAL_MS); + downstreamHeartbeat.unref(); + }, + markTerminalEvent: () => { + state.terminalEventSent = true; + }, + resetIdleTimeout, + }; +}; diff --git a/apps/hash-api/src/petrinaut-optimizer/shared/optimization-route-context.ts b/apps/hash-api/src/petrinaut-optimizer/shared/optimization-route-context.ts new file mode 100644 index 00000000000..299803900b9 --- /dev/null +++ b/apps/hash-api/src/petrinaut-optimizer/shared/optimization-route-context.ts @@ -0,0 +1,69 @@ +import type { Logger } from "@local/hash-backend-utils/logger"; +import type { Request, Response } from "express"; + +export const RUN_NOT_FOUND = { error: "Optimization run not found" } as const; + +/** Per-request logger correlated with the `x-hash-request-id` header. */ +export type OptimizationRequestLogger = Pick; + +export type OptimizationRouteContext = { + /** The configured optimizer origin (the 503 precondition has passed). */ + origin: URL; + requestId: string; + requestLogger: OptimizationRequestLogger; + startedAt: number; + userId: string; +}; + +/** + * Resolve the preamble shared by every optimization route: the correlated + * request logger, the authenticated account, and the configured upstream + * origin. Writes the 401/503 response and returns `null` when a + * precondition fails, so callers can simply bail out. + */ +export const resolveOptimizationRouteContext = ( + request: Request, + response: Response, + { + logger, + origin, + }: { + logger: Pick; + origin: URL | null; + }, +): OptimizationRouteContext | null => { + const startedAt = Date.now(); + const requestId = response.get("x-hash-request-id") ?? ""; + const requestLogger = logger.child({ requestId }); + + if (!request.user) { + response.status(401).json({ error: "Authentication required" }); + return null; + } + if (!origin) { + response + .status(503) + .json({ error: "Petrinaut optimizer is not configured" }); + return null; + } + + return { + origin, + requestId, + requestLogger, + startedAt, + userId: request.user.accountId, + }; +}; + +/** Answer a failed upstream round-trip: 504 for our own deadline, 502 else. */ +export const respondUpstreamFailure = ( + response: Response, + timedOut: boolean, +): void => { + response.status(timedOut ? 504 : 502).json({ + error: timedOut + ? "Petrinaut optimization timed out" + : "Petrinaut optimization failed", + }); +}; diff --git a/apps/hash-api/src/petrinaut-optimizer/shared/optimization-run-test-harness.ts b/apps/hash-api/src/petrinaut-optimizer/shared/optimization-run-test-harness.ts new file mode 100644 index 00000000000..a2f404ece9f --- /dev/null +++ b/apps/hash-api/src/petrinaut-optimizer/shared/optimization-run-test-harness.ts @@ -0,0 +1,196 @@ +/** + * Test-only fakes shared by the detached-run handler suites. + * + * The harness provides a + * recording logger, an EventEmitter-based Express request/response pair with + * caller-controlled backpressure, and a driver that runs one handler call. + */ +import { EventEmitter } from "node:events"; + +import type { Logger } from "@local/hash-backend-utils/logger"; +import type { + Request, + RequestHandler, + Response as ExpressResponse, +} from "express"; + +export const validOptimizationInput = { + kind: "petrinaut-optimization", + version: 1, + name: "Optimize rate", + model: { + title: "Example", + definition: { + places: [], + transitions: [], + types: [], + differentialEquations: [], + parameters: [], + subnets: [], + componentInstances: [], + scenarios: [ + { + id: "baseline", + name: "Baseline", + scenarioParameters: [ + { identifier: "rate", type: "real", default: 0.5 }, + ], + parameterOverrides: {}, + initialState: { type: "per_place", content: {} }, + }, + ], + metrics: [{ id: "profit", name: "Profit", code: "return 1;" }], + }, + }, + scenario: { + id: "baseline", + parameterBindings: { + rate: { + kind: "optimize", + domain: { + kind: "continuous", + minimum: 0.1, + maximum: 1, + scale: "linear", + }, + }, + }, + }, + objective: { metricId: "profit", direction: "maximize" }, + execution: { seed: 42, dt: 0.1, maxTime: 10 }, + study: { trials: 2, sampler: "tpe" }, +}; + +export type RecordedLog = { + level: "info" | "warn"; + message: string; + metadata: Record; +}; + +/** Build a logger fake that records request-scoped structured logs. */ +export const createRecordingLogger = () => { + const entries: RecordedLog[] = []; + const record = + (level: "info" | "warn", childMetadata: Record) => + (message: string, metadata?: Record) => { + entries.push({ + level, + message, + metadata: { ...childMetadata, ...metadata }, + }); + }; + const logger = { + child: (childMetadata: Record) => ({ + info: record("info", childMetadata), + warn: record("warn", childMetadata), + }), + info: record("info", {}), + warn: record("warn", {}), + } as unknown as Pick; + return { entries, logger }; +}; + +export const unexpectedFetch = async (): Promise => { + throw new Error("Unexpected upstream request"); +}; + +/** Mutable fake Express response the backpressure tests can drive. */ +export type FakeResponse = EventEmitter & + ExpressResponse & { destroyed: boolean }; + +/** Drive one handler invocation against fake Express request/response. */ +export const callOptimizationRunHandler = async ({ + accountId = "user-1", + authenticated = true, + body, + handler, + onRequest, + onResponse, + params = {}, + query = {}, + writeReturns, +}: { + accountId?: string; + authenticated?: boolean; + body?: unknown; + handler: RequestHandler; + onRequest?: (request: EventEmitter) => void; + onResponse?: (response: FakeResponse) => void; + params?: Record; + query?: Record; + /** Decide each write's backpressure result; defaults to no backpressure. */ + writeReturns?: (value: string) => boolean; +}) => { + let statusCode = 200; + let bodyResult: unknown; + const headers: Record = {}; + const output: string[] = []; + let headersSent = false; + let writableEnded = false; + let writableNeedDrain = false; + const responseEmitter = new EventEmitter(); + // `Object.assign` would copy getter *values*, freezing them at `false`. + Object.defineProperties(responseEmitter, { + headersSent: { get: () => headersSent }, + writableEnded: { get: () => writableEnded }, + writableNeedDrain: { get: () => writableNeedDrain }, + }); + // Clear the drain flag without registering a listener, so the tests can + // assert that the handler leaves no listeners of its own behind. + const originalEmit = responseEmitter.emit.bind(responseEmitter); + responseEmitter.emit = (eventName: string | symbol, ...args: unknown[]) => { + if (eventName === "drain") { + writableNeedDrain = false; + } + return originalEmit(eventName, ...args); + }; + const response = Object.assign(responseEmitter, { + destroyed: false, + end: () => { + writableEnded = true; + }, + flushHeaders: () => { + headersSent = true; + }, + get: (name: string) => + name.toLowerCase() === "x-hash-request-id" ? "request-id-1" : undefined, + json: (value: unknown) => { + bodyResult = value; + headersSent = true; + writableEnded = true; + return response; + }, + set: (value: Record) => { + Object.assign(headers, value); + return response; + }, + status: (value: number) => { + statusCode = value; + return response; + }, + write: (value: string) => { + headersSent = true; + output.push(value); + const flushed = writeReturns?.(value) ?? true; + if (!flushed) { + writableNeedDrain = true; + } + return flushed; + }, + }) as unknown as FakeResponse; + const request = Object.assign(new EventEmitter(), { + body, + params, + query, + user: authenticated + ? ({ accountId } as NonNullable) + : undefined, + }) as unknown as Request; + + const handlerPromise = handler(request, response, () => undefined); + onRequest?.(request); + onResponse?.(response); + await handlerPromise; + + return { body: bodyResult, headers, output, response, statusCode }; +}; diff --git a/apps/hash-api/src/petrinaut-optimizer/shared/validate-optimization-request.ts b/apps/hash-api/src/petrinaut-optimizer/shared/validate-optimization-request.ts new file mode 100644 index 00000000000..aba2cbf44c4 --- /dev/null +++ b/apps/hash-api/src/petrinaut-optimizer/shared/validate-optimization-request.ts @@ -0,0 +1,103 @@ +import { petrinautOptimizationInputSchema } from "@hashintel/petrinaut-core"; + +import type { PetrinautOptimizationInput } from "@hashintel/petrinaut-core"; +import type { Logger } from "@local/hash-backend-utils/logger"; + +export const MAX_REQUEST_BYTES = 8 * 1024 * 1024; +export const MAX_VALIDATION_ISSUES = 5; +export const MAX_VALIDATION_MESSAGE_LENGTH = 300; +// Manifest keys (e.g. scenario parameter identifiers) are user-controlled and +// unbounded, so cap each path segment before it reaches a response or a log. +export const MAX_VALIDATION_PATH_SEGMENT_LENGTH = 100; +// A trial repeats optimized parameter identifiers from the accepted manifest, +// so use the same bound rather than rejecting a valid large search space. +export const MAX_EVENT_BYTES = MAX_REQUEST_BYTES; +export const INVALID_OPTIMIZATION_REQUEST = { + code: "invalid_optimization_request", + error: "Invalid optimization request", +} as const; + +/** Cap one user-controlled value before it reaches a response or a log. */ +export const capPathSegment = (segment: string): string => + segment.slice(0, MAX_VALIDATION_PATH_SEGMENT_LENGTH); + +/** Return a bounded, transport-stable summary of manifest validation errors. */ +export const summarizeValidationIssues = ( + issues: readonly { message: string; path: readonly PropertyKey[] }[], +) => ({ + issues: issues.slice(0, MAX_VALIDATION_ISSUES).map(({ message, path }) => ({ + path: + path.map((segment) => capPathSegment(String(segment))).join(".") || "$", + message: message.slice(0, MAX_VALIDATION_MESSAGE_LENGTH), + })), + truncated: issues.length > MAX_VALIDATION_ISSUES, +}); + +export type OptimizationRequestValidation = + | { + ok: true; + bodyBytes: number; + input: PetrinautOptimizationInput; + } + | { + ok: false; + status: 400 | 413; + body: Record; + }; + +/** + * Enforce the size cap and manifest schema shared by every create route. + * + * Rejections are logged here with the same discipline as the legacy handler: + * only sizes, issue counts, and capped schema paths — never the validation + * messages or the manifest itself, which can embed user-authored code. + */ +export const validateOptimizationRequest = ({ + body, + requestLogger, + userId, +}: { + body: unknown; + requestLogger: Pick; + userId: string; +}): OptimizationRequestValidation => { + let serializedBody: unknown; + try { + serializedBody = JSON.stringify(body); + } catch { + return { ok: false, status: 400, body: INVALID_OPTIMIZATION_REQUEST }; + } + if (typeof serializedBody !== "string") { + return { ok: false, status: 400, body: INVALID_OPTIMIZATION_REQUEST }; + } + const bodyBytes = Buffer.byteLength(serializedBody, "utf8"); + if (bodyBytes > MAX_REQUEST_BYTES) { + requestLogger.warn("Petrinaut optimization rejected: body too large", { + bodyBytes, + userId, + }); + return { + ok: false, + status: 413, + body: { error: "Optimization request is too large" }, + }; + } + + const input = petrinautOptimizationInputSchema.safeParse(body); + if (!input.success) { + const details = summarizeValidationIssues(input.error.issues); + requestLogger.warn("Petrinaut optimization request failed validation", { + issueCount: input.error.issues.length, + issuePaths: details.issues.map((issue) => issue.path), + issuePathsTruncated: details.truncated, + userId, + }); + return { + ok: false, + status: 400, + body: { ...INVALID_OPTIMIZATION_REQUEST, details }, + }; + } + + return { ok: true, bodyBytes, input: input.data }; +}; 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..fe2e79aef80 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,148 @@ 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)}`; /** 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 +631,139 @@ 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; + } + 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), }); + 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 === "") { + bridge.send({ + kind: "optimizationCreateResult", + requestId, + ok: false, + category: "protocol", + message: + "The optimization service returned an unexpected response", + }); + return; } + + bridge.send({ + kind: "optimizationCreateResult", + requestId, + ok: true, + runId: body.runId, + }); })(); }, + 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 }) => { + /** + * Fire-and-forget: cancellation is idempotent server-side and the + * iframe has already torn down its own state, so a failure is only + * logged for debugging. + */ + 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); + }); + }, 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..61d1e32428c 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,6 +1,9 @@ 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 encoder = new TextEncoder(); @@ -41,28 +44,62 @@ 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("surfaces a structured HTTP error", async () => { + 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..ad2bc710d85 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,42 @@ 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, + }, + ), + ); + } + return; + } + if ( message.kind !== "optimizationResponseStart" && message.kind !== "optimizationChunk" && @@ -102,15 +219,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 +262,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 +279,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 +329,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 +344,8 @@ const bridgeFetch = ( resolveResponse: resolve, rejectResponse: reject, responded: false, + hashRequestId: null, + optimizationRunId: null, clearResponseStartTimeout, cleanup: () => { clearResponseStartTimeout(); @@ -204,11 +360,65 @@ 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 }); + }); +}; + +const readHttpError = async ( + response: Response, +): Promise => { + const correlation = { + category: "http" as const, + hashRequestId: response.headers.get("x-hash-request-id"), + optimizationRunId: response.headers.get("x-optimization-run-id"), + httpStatus: response.status, + }; const body = await response.text(); if (body) { try { @@ -220,24 +430,29 @@ 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, ); }; +/** A protocol violation while decoding the optimizer's NDJSON stream. */ +const protocolError = (message: string) => + new PetrinautOptimizationTransportError(message, { category: "protocol" }); + const parseEventLine = (line: string): 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"); } return petrinautOptimizationEventSchema.parse(parsed); }; @@ -250,7 +465,7 @@ export async function* parsePetrinautOptimizationResponse( throw await readHttpError(response); } if (!response.body) { - throw new Error("The optimizer returned an empty response body"); + throw protocolError("The optimizer returned an empty response body"); } const reader = response.body.getReader(); @@ -261,7 +476,7 @@ 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"); } const event = parseEventLine(line); return { @@ -298,7 +513,9 @@ 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", + ); } } finally { if (!reachedEnd) { @@ -307,11 +524,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 +550,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..df770c0624a 100644 --- a/apps/hash-frontend/src/pages/processes/shared/messages.ts +++ b/apps/hash-frontend/src/pages/processes/shared/messages.ts @@ -179,12 +179,40 @@ 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. + */ + kind: "optimizationCreateResult"; + requestId: string; + ok: boolean; + runId?: string; + status?: number; + retryAfter?: number; + message?: string; + category?: "network" | "http" | "protocol" | "aborted"; + } + | { + /** + * 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 +226,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 +340,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 +408,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 068c456688f..53a450c86df 100644 --- a/apps/petrinaut-opt/README.md +++ b/apps/petrinaut-opt/README.md @@ -10,13 +10,21 @@ 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. - -- `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. +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/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}` + cancels it. + +When run creation carries an `x-hash-account-id` header (the authenticated +NodeAPI proxy stamps it), the run is owned: the account may drive only one +live run at a time (429 otherwise), and attach/cancel answer 404 unless the +same tag is presented — identical to an unknown run, so foreign run ids +cannot be probed. Requests without the header (local development, the +website demo) create ownerless, openly attachable runs. The response is `text/event-stream`. Existing frame bodies are preserved: @@ -47,10 +55,10 @@ every 30 seconds: SSE clients ignore comment frames, while load balancers and proxies see traffic before their idle timeout. -Streams are not resumable: they do not emit event IDs or replay missed trials. -If the caller disconnects, the service stops that study and releases its CLI; -the caller must submit a new optimization request rather than reconnect with -`Last-Event-ID`. +The streaming endpoints are not resumable: disconnecting stops that study and +releases its CLI. Detached-run event streams are resumable: every frame has an +event id, buffered frames can be replayed using `Last-Event-ID` (or `cursor`), +and disconnecting an attachment does not stop the run. Each response has an `X-Optimization-Run-ID` header for status queries: @@ -77,8 +85,15 @@ optimization manifests, user-authored code, or raw request bodies. The process admits at most four active optimizations. Additional requests receive HTTP 429, and slots are released after initialization failures, stream -failures, completion, or disconnect. `GET /status` retains the 100 most recent -runs so process memory cannot grow without bound. +failures, completion, disconnect, or detached-run cancellation/reaping. +`GET /status` retains the 100 most recent runs so process memory cannot grow +without bound. + +Detached runs reject descriptions above 1,000 trials and, by default, stop +after 900 seconds (`HASH_PETRINAUT_OPT_MAX_STUDY_SECONDS`); invalid values use +the default and zero disables the wall-clock limit. Their event log is retained +for the detach-grace period and an attachment cursor is clamped to the current +log, so malformed resume requests cannot suppress later terminal events. Optimization request bodies are limited to 8 MiB, including chunked bodies. diff --git a/apps/petrinaut-opt/openapi/openapi.json b/apps/petrinaut-opt/openapi/openapi.json index 1c6baa01e1f..877d4b5672d 100644 --- a/apps/petrinaut-opt/openapi/openapi.json +++ b/apps/petrinaut-opt/openapi/openapi.json @@ -14,6 +14,20 @@ "title": "HTTPValidationError", "type": "object" }, + "OptimizationRunCreated": { + "description": "Response body of a successful detached-run creation.", + "properties": { + "run_id": { + "title": "Run Id", + "type": "string" + } + }, + "required": [ + "run_id" + ], + "title": "OptimizationRunCreated", + "type": "object" + }, "Phase": { "enum": [ "idle", @@ -135,10 +149,10 @@ "summary": "Root" } }, - "/optimize/all": { + "/optimize/runs": { "post": { - "description": "Stream one SSE data frame for every completed Optuna trial.", - "operationId": "post_optimize_all_optimize_all_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.", + "operationId": "post_optimize_runs_optimize_runs_post", "requestBody": { "content": { "application/json": { @@ -152,15 +166,15 @@ "required": true }, "responses": { - "200": { + "201": { "content": { - "text/event-stream": { + "application/json": { "schema": { - "type": "string" + "$ref": "#/components/schemas/OptimizationRunCreated" } } }, - "description": "Server-Sent Events optimization stream" + "description": "A detached optimization run was started" }, "413": { "description": "The optimization manifest exceeds 8 MiB" @@ -190,25 +204,77 @@ "description": "The CLI or optimization study could not initialize" } }, - "summary": "Post Optimize All" + "summary": "Post Optimize Runs" } }, - "/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" - } + "/optimize/runs/{run_id}": { + "delete": { + "description": "Cancel a detached run and wait for its resources to be released.", + "operationId": "delete_optimize_run_optimize_runs__run_id__delete", + "parameters": [ + { + "in": "path", + "name": "run_id", + "required": true, + "schema": { + "title": "Run Id", + "type": "string" } + } + ], + "responses": { + "204": { + "description": "The run was cancelled (or had already reached a terminal state); when this call stopped it, the event log's terminal frame is `event: cancelled`" }, - "required": true + "404": { + "description": "No optimization run with this id is registered" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } }, + "summary": "Delete Optimize Run" + } + }, + "/optimize/runs/{run_id}/events": { + "get": { + "description": "Attach to a detached run and replay events after the supplied cursor.\n\nThe route is excluded from ASGI auto-instrumentation (see\n``telemetry.py``): its response live-tails until the consumer detaches,\nand a SERVER span that long would read as worst-case latency in the RED\nSLIs and only export on disconnect. A short manual SERVER span covers\njust the attach itself \u2014 run lookup, cursor resolution, and attachment\nregistration \u2014 and ends when the tail starts, so the latency SLI\nmeasures \"was attaching fast\" while the client\u2192optimizer service-graph\nedge is preserved.", + "operationId": "get_optimize_run_events_optimize_runs__run_id__events_get", + "parameters": [ + { + "in": "path", + "name": "run_id", + "required": true, + "schema": { + "title": "Run Id", + "type": "string" + } + }, + { + "in": "query", + "name": "cursor", + "required": false, + "schema": { + "anyOf": [ + { + "minimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Cursor" + } + } + ], "responses": { "200": { "content": { @@ -218,10 +284,18 @@ } } }, - "description": "Server-Sent Events optimization stream" + "description": "Server-Sent Events attachment to a detached optimization run. Every frame carries an `id: ` line (seq starts at 1); buffered frames with seq > cursor are replayed, then new frames are live-tailed with `: heartbeat` comments roughly every 30 seconds. The terminal frame is `event: done` (completed), an ERROR data frame (study failure), or `event: cancelled` (cancelled or reaped). If the run is already terminal the response closes after the replay. Disconnecting does not affect the run; a newer attachment supersedes this one.", + "headers": { + "X-Requested-Trials": { + "description": "Trial count requested by the run's manifest, for sizing synthesized summaries", + "schema": { + "type": "string" + } + } + } }, - "413": { - "description": "The optimization manifest exceeds 8 MiB" + "404": { + "description": "No optimization run with this id is registered" }, "422": { "content": { @@ -232,23 +306,9 @@ } }, "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" + "summary": "Get Optimize Run Events" } }, "/status": { diff --git a/apps/petrinaut-opt/src/optimization_api.py b/apps/petrinaut-opt/src/optimization_api.py index 4b808ae4571..7f3090c9358 100644 --- a/apps/petrinaut-opt/src/optimization_api.py +++ b/apps/petrinaut-opt/src/optimization_api.py @@ -13,11 +13,16 @@ from typing import Any from dotenv import load_dotenv -from fastapi import FastAPI, HTTPException, Request +from fastapi import FastAPI, HTTPException, Query, Request, Response +from pydantic import BaseModel from fastapi.responses import JSONResponse, StreamingResponse +from opentelemetry import trace +from opentelemetry.propagate import extract +from opentelemetry.trace import SpanKind, Status, StatusCode from starlette.background import BackgroundTask from starlette.types import ASGIApp, Message, Receive, Scope, Send +from src.optimization_runs import OptimizationRunRegistry, attachment_event_stream from src.petrinaut_client import PetrinautModel from src.petrinaut_optimizer import PetrinautOptimizer from src.telemetry import flush_telemetry, setup_telemetry @@ -27,15 +32,13 @@ REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent load_dotenv(REPO_ROOT / ".env") log = logging.getLogger("pn_api") +tracer = trace.get_tracer("pn_api") MAX_REQUEST_BODY_BYTES = 8 * 1024 * 1024 MAX_ACTIVE_OPTIMIZATIONS = 4 RETRY_AFTER_SECONDS = 30 -_OPTIMIZATION_PATHS = {"/optimize/all", "/optimize/best"} -_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", @@ -48,6 +51,48 @@ }, 500: {"description": "The CLI or optimization study could not initialize"}, } +_RUN_EVENTS_RESPONSES = { + 200: { + "description": ( + "Server-Sent Events attachment to a detached optimization run. " + "Every frame carries an `id: ` line (seq starts at 1); " + "buffered frames with seq > cursor are replayed, then new frames " + "are live-tailed with `: heartbeat` comments roughly every 30 " + "seconds. The terminal frame is `event: done` (completed), an " + "ERROR data frame (study failure), or `event: cancelled` " + "(cancelled or reaped). If the run is already terminal the " + "response closes after the replay. Disconnecting does not affect " + "the run; a newer attachment supersedes this one." + ), + "content": {"text/event-stream": {"schema": {"type": "string"}}}, + "headers": { + "X-Requested-Trials": { + "description": ( + "Trial count requested by the run's manifest, for " + "sizing synthesized summaries" + ), + "schema": {"type": "string"}, + }, + }, + }, + 404: {"description": "No optimization run with this id is registered"}, +} +_DELETE_RUN_RESPONSES = { + 204: { + "description": ( + "The run was cancelled (or had already reached a terminal " + "state); when this call stopped it, the event log's terminal " + "frame is `event: cancelled`" + ), + }, + 404: {"description": "No optimization run with this id is registered"}, +} + + +class OptimizationRunCreated(BaseModel): + """Response body of a successful detached-run creation.""" + + run_id: str class _RequestBodyTooLarge(Exception): @@ -106,17 +151,22 @@ async def _reject(scope: Scope, receive: Receive, send: Send) -> None: @asynccontextmanager async def lifespan(app: FastAPI) -> AsyncIterator[None]: - """Initialize the run-scoped status registry.""" + """Initialize the status registry and detached-run engine.""" app.state.statuses = StatusStore() app.state.optimization_admission_lock = asyncio.Lock() app.state.active_optimizations = 0 + app.state.optimization_pending_accounts = set() + app.state.optimization_runs = OptimizationRunRegistry() try: yield finally: - # Flush buffered spans/metrics/logs on graceful shutdown so a SIGTERM - # (e.g. `docker stop`) does not drop whatever the batch processors have - # queued. - flush_telemetry() + try: + await app.state.optimization_runs.shutdown() + finally: + # Flush buffered spans/metrics/logs on graceful shutdown so a SIGTERM + # (e.g. `docker stop`) does not drop whatever the batch processors have + # queued. + flush_telemetry() app = FastAPI(title="Petrinaut optimization Python API", lifespan=lifespan) @@ -146,8 +196,27 @@ async def _acquire_optimization_slot( app: FastAPI, *, request_id: str | None = None, + account_id: str | None = None, ) -> None: async with app.state.optimization_admission_lock: + # One account drives at most one optimization at a time. Pending + # creations are tracked separately because the run only appears in + # the registry once its CLI has initialized. + if account_id is not None and ( + account_id in app.state.optimization_pending_accounts + or app.state.optimization_runs.has_live_run_for_account(account_id) + ): + log.warning( + "optimization request rejected: account busy", + extra={ + "event": "account_busy_rejected", + "request_id": request_id, + }, + ) + raise HTTPException( + status_code=429, + detail="An optimization is already running for this account", + ) if app.state.active_optimizations >= MAX_ACTIVE_OPTIMIZATIONS: log.warning( "optimization request rejected: study limit reached", @@ -166,6 +235,27 @@ async def _acquire_optimization_slot( headers={"Retry-After": str(RETRY_AFTER_SECONDS)}, ) app.state.active_optimizations += 1 + if account_id is not None: + app.state.optimization_pending_accounts.add(account_id) + + +def _request_correlation(request: Request) -> dict[str, str | None]: + return {"request_id": request.headers.get("x-hash-request-id")} + + +def _account_id(request: Request) -> str | None: + """The opaque owner tag stamped by the authenticated proxy, if any. + + Direct callers (local development, the website demo) send none: their + runs are created ownerless and stay attachable without the header. + """ + value = (request.headers.get("x-hash-account-id") or "").strip() + return value or None + + +def _discard_pending_account(app: FastAPI, account_id: str | None) -> None: + if account_id is not None: + app.state.optimization_pending_accounts.discard(account_id) async def _release_optimization_slot(app: FastAPI) -> None: @@ -241,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, @@ -283,18 +363,24 @@ def _initialization_error( ) -@app.post( - "/optimize/all", - response_class=StreamingResponse, - responses=_SSE_RESPONSES, -) -async def post_optimize_all( +async def _admit_and_initialize_run( request: Request, optimization_manifest: dict[str, Any], -) -> StreamingResponse: - """Stream one SSE data frame for every completed Optuna trial.""" - request_id = request.headers.get("x-hash-request-id") - await _acquire_optimization_slot(request.app, request_id=request_id) + *, + account_id: str | None = None, +) -> tuple[str, PetrinautOptimizer, dict[str, str | None]]: + """Admit and initialize one optimizer, releasing its slot on failure. + + The caller owns clearing the account's pending-creation mark (in a + ``finally`` around the whole creation flow), so every failure path here + only needs to release the slot itself. + """ + correlation = _request_correlation(request) + await _acquire_optimization_slot( + request.app, + request_id=correlation["request_id"], + account_id=account_id, + ) try: run_id = request.app.state.statuses.create().run_id except BaseException: @@ -317,78 +403,160 @@ async def post_optimize_all( await asyncio.to_thread(optimizer.pn_model.close, graceful=False) await asyncio.shield(_release_optimization_slot(request.app)) raise _initialization_error( - request.app, run_id, error, request_id=request_id + request.app, run_id, error, request_id=correlation["request_id"] ) from error - - 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), - ) + return run_id, optimizer, correlation -@app.post( - "/optimize/best", - response_class=StreamingResponse, - responses=_SSE_RESPONSES, -) -async def post_optimize_best( +@app.post("/optimize/runs", status_code=201, responses=_CREATE_RUN_RESPONSES) +async def post_optimize_runs( request: Request, optimization_manifest: dict[str, Any], -) -> StreamingResponse: - """Stream the best-so-far SSE data frame after every completed trial.""" - request_id = request.headers.get("x-hash-request-id") - await _acquire_optimization_slot(request.app, request_id=request_id) - try: - run_id = request.app.state.statuses.create().run_id - except BaseException: - await asyncio.shield(_release_optimization_slot(request.app)) - raise - optimizer: PetrinautOptimizer | None = None + response: Response, +) -> OptimizationRunCreated: + """Start a detached run that remains available for later SSE attachment. + + When the caller stamps ``x-hash-account-id`` (the authenticated proxy + does), the run is owned: the account is single-flight while it lives, and + only requests carrying the same tag may attach to or cancel it. + """ + account_id = _account_id(request) try: - optimizer = await _initialize_admitted_optimizer( - request.app, optimization_manifest + run_id, optimizer, correlation = await _admit_and_initialize_run( + request, optimization_manifest, account_id=account_id ) - set_status( + cleanup = _create_admitted_run_cleanup(request.app, optimizer) + request.app.state.optimization_runs.create_run( request.app, - run_id, - phase=Phase.running, - detail="Petrinaut CLI and Optimization Model initialized", + run_id=run_id, + optimizer=optimizer, + cleanup=cleanup, + correlation=correlation, + account_id=account_id, ) - except Exception as error: - if optimizer is not None: - with suppress(Exception): - await asyncio.to_thread(optimizer.pn_model.close, graceful=False) - await asyncio.shield(_release_optimization_slot(request.app)) - raise _initialization_error( - request.app, run_id, error, request_id=request_id - ) from error + finally: + # The registry entry (or the failure) now carries the truth; the + # pending-creation mark has done its job on every path, including + # cancellation mid-initialization. + _discard_pending_account(request.app, account_id) + response.headers["X-Optimization-Run-ID"] = run_id + return OptimizationRunCreated(run_id=run_id) + + +def _resolve_owned_run(request: Request, run_id: str) -> Any: + """Look up a run, enforcing its owner tag when it has one. + + Answers 404 for a mismatch — identical to an unknown run, so id-guessing + callers cannot learn that a foreign run exists. + """ + run = request.app.state.optimization_runs.get(run_id) + if run is None or ( + run.account_id is not None and _account_id(request) != run.account_id + ): + raise HTTPException(404, f"optimization run not found: {run_id}") + return run + + +def _cursor_from_last_event_id(request: Request) -> int: + raw = request.headers.get("last-event-id") + if raw is None: + return 0 + try: + return max(0, int(raw.strip())) + except ValueError: + return 0 + + +@app.get( + "/optimize/runs/{run_id}/events", + response_class=StreamingResponse, + responses=_RUN_EVENTS_RESPONSES, +) +async def get_optimize_run_events( + run_id: str, + request: Request, + cursor: int | None = Query(default=None, ge=0), +) -> StreamingResponse: + """Attach to a detached run and replay events after the supplied cursor. + + The route is excluded from ASGI auto-instrumentation (see + ``telemetry.py``): its response live-tails until the consumer detaches, + and a SERVER span that long would read as worst-case latency in the RED + SLIs and only export on disconnect. A short manual SERVER span covers + just the attach itself — run lookup, cursor resolution, and attachment + registration — and ends when the tail starts, so the latency SLI + measures "was attaching fast" while the client→optimizer service-graph + edge is preserved. + """ + with tracer.start_as_current_span( + "GET /optimize/runs/{run_id}/events", + context=extract(dict(request.headers)), + kind=SpanKind.SERVER, + attributes={ + "http.method": "GET", + "http.route": "/optimize/runs/{run_id}/events", + "http.target": f"/optimize/runs/{run_id}/events", + }, + # Status and attributes are set explicitly below; an expected 404 + # must not surface as a recorded exception on the span. + record_exception=False, + set_status_on_exception=False, + end_on_exit=True, + ) as attach_span: + try: + run = _resolve_owned_run(request, run_id) + except HTTPException: + attach_span.set_attribute("http.status_code", 404) + attach_span.set_status(Status(StatusCode.ERROR)) + raise + if cursor is None: + cursor = _cursor_from_last_event_id(request) + epoch = run.begin_attachment() + attach_span.set_attribute("http.status_code", 200) + + async def detach_if_never_started() -> None: + # A consumer that aborts before the body iterator ever starts would + # otherwise stay marked attached and shield the run from the reaper. + # This must be async: Starlette runs sync background callables in a + # threadpool, where end_attachment's get_running_loop() would fail. + run.end_attachment(epoch) - 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, + attachment_event_stream( + run, + cursor=cursor, + epoch=epoch, + request=request, + correlation=_request_correlation(request), ), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", "X-Optimization-Run-ID": run_id, + # Sizes the consumer's synthesized summary without the proxy + # having to remember the manifest. + "X-Requested-Trials": str(run.requested_trials), }, - background=BackgroundTask(cleanup), + background=BackgroundTask(detach_if_never_started), ) +@app.delete( + "/optimize/runs/{run_id}", status_code=204, responses=_DELETE_RUN_RESPONSES +) +async def delete_optimize_run(run_id: str, request: Request) -> Response: + """Cancel a detached run and wait for its resources to be released.""" + run = _resolve_owned_run(request, run_id) + if run.request_cancel("cancelled by client request"): + log.info( + "optimization run cancellation requested", + extra={"event": "run_cancel_requested", "run_id": run_id, **_request_correlation(request)}, + ) + await run.finished.wait() + return Response(status_code=204) + + @app.get("/status") def get_status() -> list[RunStatus]: """Return a snapshot of every optimization run's status.""" diff --git a/apps/petrinaut-opt/src/optimization_runs.py b/apps/petrinaut-opt/src/optimization_runs.py new file mode 100644 index 00000000000..1a478d319d7 --- /dev/null +++ b/apps/petrinaut-opt/src/optimization_runs.py @@ -0,0 +1,454 @@ +#!/usr/bin/env python3 +"""Detached, reconnectable optimization runs. + +A detached run owns one admitted optimizer/CLI pair and drives its study from +a background task, appending every SSE frame body to an in-memory event log +instead of an HTTP response. Consumers attach — and re-attach — over +``GET /optimize/runs/{run_id}/events`` with a cursor; disconnecting a consumer +does not affect the run. The admission slot's lifetime equals the run's +lifetime: the idempotent cleanup (prompt CLI close plus slot release) runs +when the run reaches a terminal state or is cancelled/reaped, never when a +consumer detaches. + +The event log holds one frame per completed trial plus a handful of control +frames; it is bounded because the optimizer itself rejects study descriptions +above ``MAX_STUDY_TRIALS`` (1000) trials — mirroring the optimization +manifest contract — rather than trusting the CLI's reported trial count. +""" + +from __future__ import annotations + +import asyncio +import logging +import math +import os +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from contextlib import suppress +from enum import Enum +from typing import Any + +from src.utils import Phase, set_status + +log = logging.getLogger("pn_runs") + +DETACH_GRACE_ENVIRONMENT_VARIABLE = "HASH_PETRINAUT_OPT_DETACH_GRACE_SECONDS" +DEFAULT_DETACH_GRACE_SECONDS = 300.0 +SSE_HEARTBEAT_SECONDS = 30 +_TAIL_POLL_SECONDS = 0.1 +# The single terminal frame appended when a run is cancelled (client DELETE, +# orphan reaping, or service shutdown). Documented in the README and OpenAPI +# responses; NodeAPI distinguishes it from a study failure's ERROR data frame. +CANCELLED_FRAME = "event: cancelled\ndata: {}\n\n" + + +def detach_grace_seconds_from_environment() -> float: + """Read the orphan grace period; invalid values fall back to the default. + + A value of zero or below disables the orphan reaper (terminal-run log + retention then falls back to the default period). Invalid and non-finite + values (``inf``, ``nan``, overflowing literals) fall back to the default. + """ + raw = os.environ.get(DETACH_GRACE_ENVIRONMENT_VARIABLE, "").strip() + if not raw: + return DEFAULT_DETACH_GRACE_SECONDS + try: + value = float(raw) + except ValueError: + return DEFAULT_DETACH_GRACE_SECONDS + if not math.isfinite(value): + return DEFAULT_DETACH_GRACE_SECONDS + return value + + +class RunState(str, Enum): + running = "running" + completed = "completed" + failed = "failed" + cancelled = "cancelled" + + +class OptimizationRun: + """One detached optimization run and its replayable event log. + + Every mutation happens on the event loop thread (the pump's ``on_event`` + callback, the endpoints, and the registry), so plain attributes suffice; + the worker thread only reaches the loop via ``call_soon_threadsafe`` + inside the optimizer. + """ + + def __init__( + self, + *, + run_id: str, + optimizer: Any, + cleanup: Callable[[], Awaitable[None]], + correlation: Mapping[str, str | None] | None = None, + account_id: str | None = None, + requested_trials: int = 0, + ) -> None: + loop = asyncio.get_running_loop() + self.run_id = run_id + self.optimizer = optimizer + self.cleanup = cleanup + self.correlation: dict[str, str | None] = dict(correlation or {}) + # Opaque owner tag stamped by the authenticated proxy at creation. + # When set, only requests carrying the same tag may attach to or + # cancel the run, and the account is single-flight while it lives. + self.account_id = account_id + # Recorded so attachments can size their synthesized summary without + # the proxy having to remember the manifest. + self.requested_trials = requested_trials + self.state = RunState.running + # (seq, SSE frame body); seq starts at 1 and increments by one, so the + # frame with sequence number ``seq`` lives at index ``seq - 1``. + self.events: list[tuple[int, str]] = [] + self.attached = False + self.attachment_epoch = 0 + self.created_at = loop.time() + # Grace also counts from creation, before the first attachment. + self.last_detached_at = loop.time() + self.terminal_at: float | None = None + self.cancel_requested = asyncio.Event() + self.cancel_reason = "optimization run cancelled" + self.finished = asyncio.Event() + self.task: asyncio.Task[None] | None = None + self._changed = asyncio.Event() + + def append_event(self, frame: str) -> int: + """Append one SSE frame body and wake waiting tails.""" + seq = len(self.events) + 1 + self.events.append((seq, frame)) + self._wake() + return seq + + def request_cancel(self, reason: str) -> bool: + """Ask the pump to stop; False when already requested or terminal.""" + if self.cancel_requested.is_set() or self.state is not RunState.running: + return False + self.cancel_reason = reason + self.cancel_requested.set() + return True + + def begin_attachment(self) -> int: + """Register a new consumer, superseding any previous attachment.""" + self.attachment_epoch += 1 + self.attached = True + # Wake the superseded tail so it observes the epoch change promptly. + self._wake() + return self.attachment_epoch + + def end_attachment(self, epoch: int) -> None: + """Mark the consumer gone; stale epochs are ignored. Idempotent.""" + if epoch != self.attachment_epoch or not self.attached: + return + self.attached = False + self.last_detached_at = asyncio.get_running_loop().time() + + def mark_terminal(self, state: RunState) -> None: + self.state = state + self.terminal_at = asyncio.get_running_loop().time() + self.finished.set() + self._wake() + + def _wake(self) -> None: + """Wake every waiting tail; each wait round gets a fresh event.""" + changed = self._changed + self._changed = asyncio.Event() + changed.set() + + async def wait_for_change(self, timeout: float) -> None: + """Wait until the run appends, terminates, or is superseded.""" + waiter = self._changed + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(waiter.wait(), timeout) + + +async def attachment_event_stream( + run: OptimizationRun, + *, + cursor: int, + epoch: int, + request: Any, + correlation: Mapping[str, str | None] | None = None, +) -> AsyncIterator[str]: + """Replay buffered frames with seq > cursor, then live-tail new ones. + + Every frame carries an ``id: `` line so a consumer can re-attach + with ``?cursor=`` (or ``Last-Event-ID``). The stream ends when the + run's log is terminal and fully replayed, when a newer attachment + supersedes this one, or when the consumer disconnects — none of which + affect the run itself. Comment heartbeats are sent while tailing. + """ + log_context = {**(correlation or {}), "run_id": run.run_id} + loop = asyncio.get_running_loop() + # A cursor pointing past the end of the log would otherwise also filter + # out every frame appended after attaching — including the terminal one. + # Clamping is safe: sequence numbers are dense, append-only, and never + # truncated, so nothing past the current length can have been delivered. + cursor = min(cursor, len(run.events)) + next_index = 0 + next_heartbeat = loop.time() + SSE_HEARTBEAT_SECONDS + log.info( + "consumer attached to optimization run", + extra={"event": "run_attached", "cursor": cursor, **log_context}, + ) + try: + while True: + while next_index < len(run.events): + seq, frame = run.events[next_index] + next_index += 1 + if seq <= cursor: + continue + yield f"id: {seq}\n{frame}" + next_heartbeat = loop.time() + SSE_HEARTBEAT_SECONDS + if run.state is not RunState.running and next_index >= len(run.events): + return + if epoch != run.attachment_epoch: + log.info( + "optimization run attachment superseded", + extra={"event": "run_attachment_superseded", **log_context}, + ) + # Terminal for this ATTACHMENT, not the run: without a + # sentinel the consumer's decoder would report a truncated + # stream and reconnect — superseding the newer attachment in + # turn, ping-ponging forever. No `id:` line on purpose; the + # frame is attachment-scoped, not part of the run's log. + yield "event: superseded\ndata: {}\n\n" + return + if await request.is_disconnected(): + return + await run.wait_for_change(_TAIL_POLL_SECONDS) + if loop.time() >= next_heartbeat: + yield ": heartbeat\n\n" + next_heartbeat = loop.time() + SSE_HEARTBEAT_SECONDS + finally: + run.end_attachment(epoch) + log.info( + "consumer detached from optimization run", + extra={"event": "run_detached", **log_context}, + ) + + +class OptimizationRunRegistry: + """All detached runs, plus the reaper that bounds their lifetimes. + + The reaper cancels a running run with no attached consumer for the detach + grace period, and drops a terminal run's log after the retention period + (the StatusStore keeps the status row as before). + """ + + def __init__( + self, + *, + detach_grace_seconds: float | None = None, + retention_seconds: float | None = None, + ) -> None: + self.detach_grace_seconds = ( + detach_grace_seconds_from_environment() + if detach_grace_seconds is None + else detach_grace_seconds + ) + if retention_seconds is not None: + self.retention_seconds = retention_seconds + elif self.detach_grace_seconds > 0: + self.retention_seconds = self.detach_grace_seconds + else: + self.retention_seconds = DEFAULT_DETACH_GRACE_SECONDS + self._runs: dict[str, OptimizationRun] = {} + self._reaper_task: asyncio.Task[None] | None = None + + def get(self, run_id: str) -> OptimizationRun | None: + return self._runs.get(run_id) + + def runs(self) -> list[OptimizationRun]: + return list(self._runs.values()) + + def has_live_run_for_account(self, account_id: str) -> bool: + """Whether the account owns a run that is still running.""" + return any( + run.account_id == account_id and run.state is RunState.running + for run in self._runs.values() + ) + + def create_run( + self, + app: Any, + *, + run_id: str, + optimizer: Any, + cleanup: Callable[[], Awaitable[None]], + correlation: Mapping[str, str | None] | None = None, + account_id: str | None = None, + ) -> OptimizationRun: + """Register a run and start the background pump driving its study.""" + run = OptimizationRun( + run_id=run_id, + optimizer=optimizer, + cleanup=cleanup, + correlation=correlation, + account_id=account_id, + requested_trials=int(getattr(optimizer, "n_trials", 0) or 0), + ) + self._runs[run_id] = run + run.task = asyncio.create_task( + self._drive_run(app, run), name=f"petrinaut-run-{run_id}" + ) + self._ensure_reaper() + return run + + async def _drive_run(self, app: Any, run: OptimizationRun) -> None: + """Pump one run to a terminal state, then release its resources.""" + state = RunState.failed + cancellation: asyncio.CancelledError | None = None + # The pump records its outcome right before returning — ahead of its + # cancellable CLI-shutdown finally — so a cancellation landing in + # that teardown window cannot relabel an already-decided study. + recorded_outcomes: list[str] = [] + try: + state = RunState( + await run.optimizer.pump_events( + app, + run.run_id, + run.optimizer.n_trials, + on_event=run.append_event, + cancel_event=run.cancel_requested, + on_outcome=recorded_outcomes.append, + correlation=run.correlation, + ) + ) + except asyncio.CancelledError as error: + # Service shutdown cancels the pump task; the pump's own finally + # already closed the CLI on its way out. Keep a decided outcome: + # its terminal frame is already in the log, and appending a + # cancelled frame after it would corrupt the replay. + state = ( + RunState(recorded_outcomes[0]) + if recorded_outcomes + else RunState.cancelled + ) + cancellation = error + except Exception as error: + # Backstop only: the pump reports study failures itself. The raw + # message may quote user content, so log its type only. + log.error( + "optimization run pump failed", + extra={ + "event": "run_pump_failed", + "run_id": run.run_id, + "error_type": type(error).__name__, + **run.correlation, + }, + ) + run.append_event( + 'data: {"state": "ERROR", "message": "optimization run failed"}\n\n' + ) + with suppress(Exception): + set_status( + app, + run.run_id, + phase=Phase.error, + detail="optimization run failed", + ) + finally: + try: + if state is RunState.cancelled: + run.append_event(CANCELLED_FRAME) + with suppress(Exception): + set_status( + app, + run.run_id, + phase=Phase.idle, + detail=run.cancel_reason, + ) + log.info( + "optimization run cancelled", + extra={ + "event": "run_cancelled", + "run_id": run.run_id, + "detail": run.cancel_reason, + **run.correlation, + }, + ) + await asyncio.shield(run.cleanup()) + finally: + run.mark_terminal(state) + if cancellation is not None: + raise cancellation + + def _ensure_reaper(self) -> None: + if self._reaper_task is None or self._reaper_task.done(): + self._reaper_task = asyncio.create_task( + self._reap_loop(), name="petrinaut-run-reaper" + ) + + @property + def _tick_seconds(self) -> float: + horizons = [self.retention_seconds] + if self.detach_grace_seconds > 0: + horizons.append(self.detach_grace_seconds) + return max(0.01, min(1.0, min(horizons) / 10)) + + async def _reap_loop(self) -> None: + while True: + await asyncio.sleep(self._tick_seconds) + now = asyncio.get_running_loop().time() + for run in self.runs(): + try: + self._reap_one(run, now) + except Exception: + # One bad run must not stop the reaper: an unreaped + # orphan would otherwise hold its slot forever (the loop + # is only recreated by the next create_run). + log.exception( + "optimization run reaping failed", + extra={"event": "run_reap_failed", "run_id": run.run_id}, + ) + + def _reap_one(self, run: OptimizationRun, now: float) -> None: + if run.state is RunState.running: + if ( + self.detach_grace_seconds > 0 + and not run.attached + and now - run.last_detached_at >= self.detach_grace_seconds + and run.request_cancel( + "no attached consumer within the detach grace period" + ) + ): + log.warning( + "optimization run reaped: no attached consumer", + extra={ + "event": "run_reaped", + "run_id": run.run_id, + "detach_grace_seconds": self.detach_grace_seconds, + **run.correlation, + }, + ) + elif ( + run.terminal_at is not None + and now - run.terminal_at >= self.retention_seconds + ): + del self._runs[run.run_id] + log.info( + "optimization run log expired", + extra={ + "event": "run_expired", + "run_id": run.run_id, + **run.correlation, + }, + ) + + async def shutdown(self) -> None: + """Cancel the reaper and every pump; used by the app's lifespan.""" + tasks: list[asyncio.Task[None]] = [] + if self._reaper_task is not None: + self._reaper_task.cancel() + tasks.append(self._reaper_task) + self._reaper_task = None + for run in self.runs(): + if run.task is not None and not run.task.done(): + if not run.cancel_requested.is_set(): + run.cancel_reason = "service shutting down" + run.task.cancel() + tasks.append(run.task) + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) diff --git a/apps/petrinaut-opt/src/petrinaut_optimizer.py b/apps/petrinaut-opt/src/petrinaut_optimizer.py index 162c9478acf..578759e136b 100644 --- a/apps/petrinaut-opt/src/petrinaut_optimizer.py +++ b/apps/petrinaut-opt/src/petrinaut_optimizer.py @@ -7,14 +7,13 @@ import json import logging import math -import queue +import os 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 @@ -31,7 +30,12 @@ "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. +MAX_STUDY_TRIALS = 1000 +MAX_STUDY_SECONDS_ENVIRONMENT_VARIABLE = "HASH_PETRINAUT_OPT_MAX_STUDY_SECONDS" +DEFAULT_MAX_STUDY_SECONDS = 900.0 _DISCONNECT_POLL_SECONDS = 0.1 _WORKER_SHUTDOWN_TIMEOUT_SECONDS = 12 _SENTINEL = object() @@ -40,6 +44,25 @@ ParameterDescriptor: TypeAlias = Mapping[str, Any] +def max_study_seconds_from_environment() -> float: + """Read the detached-study wall-clock ceiling in seconds. + + Defaults to ``DEFAULT_MAX_STUDY_SECONDS``; a value of zero or below + disables the ceiling. Invalid and non-finite values (``inf``, ``nan``, + overflowing literals) fall back to the default. + """ + raw = os.environ.get(MAX_STUDY_SECONDS_ENVIRONMENT_VARIABLE, "").strip() + if not raw: + return DEFAULT_MAX_STUDY_SECONDS + try: + value = float(raw) + except ValueError: + return DEFAULT_MAX_STUDY_SECONDS + if not math.isfinite(value): + return DEFAULT_MAX_STUDY_SECONDS + return value + + def _finite_number(value: Any, name: str) -> float: if ( isinstance(value, bool) @@ -73,6 +96,10 @@ def _parse_description( n_trials = study.get("trials") if isinstance(n_trials, bool) or not isinstance(n_trials, int) or n_trials < 1: raise ValueError("optimization.describe study.trials must be positive") + if n_trials > MAX_STUDY_TRIALS: + raise ValueError( + f"optimization.describe study.trials must not exceed {MAX_STUDY_TRIALS}" + ) seed = study.get("seed") if isinstance(seed, bool) or not isinstance(seed, int) or seed < 0: raise ValueError( @@ -236,11 +263,16 @@ def objective(self, trial: optuna.Trial) -> float: def _start_study_worker( self, - *, - n_trials: int, - callback: Callable[[optuna.Study, optuna.trial.FrozenTrial], None], - events: asyncio.Queue[dict[str, Any] | object], loop: asyncio.AbstractEventLoop, + events: asyncio.Queue[dict[str, Any] | object], + stop_flag: threading.Event | None = None, + n_trials: int | None = None, + payload_builder: ( + Callable[[optuna.Study, optuna.trial.FrozenTrial], dict[str, Any] | None] + | None + ) = None, + *, + callback: Callable[[optuna.Study, optuna.trial.FrozenTrial], None] | None = None, ) -> tuple[threading.Thread, Span]: """Run the study on a worker thread that inherits the request's context. @@ -250,6 +282,21 @@ def _start_study_worker( span. The returned ``optimization.study`` span is the parent of those trial spans; the caller must ``end()`` it once the stream is torn down. """ + if n_trials is None: + raise ValueError("n_trials is required") + if callback is None: + if stop_flag is None or payload_builder is None: + raise ValueError("callback or detached-run callback inputs are required") + + def callback( + study: optuna.Study, trial: optuna.trial.FrozenTrial + ) -> None: + payload = payload_builder(study, trial) + if payload is not None: + loop.call_soon_threadsafe(events.put_nowait, payload) + if stop_flag.is_set(): + study.stop() + study_span = tracer.start_span("optimization.study") study_span.set_attribute("optuna.study.trials", n_trials) study_span.set_attribute("optuna.study.direction", self.direction) @@ -281,145 +328,36 @@ def run() -> None: worker.start() return worker, study_span - 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, + @staticmethod + def _trial_payload( + _study: optuna.Study, trial: optuna.trial.FrozenTrial + ) -> dict[str, Any]: + return { + "step": trial.number, + "params": dict(trial.params), + "init_state": {}, + "metric": trial.value, + "state": trial.state.name, } - 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( - n_trials=n_trials, callback=callback, events=events, loop=loop - ) - 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, - } + 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: Mapping[str, str | None] | None = None, + ) -> str: + """Run a bounded detached study and append its frames to the event log.""" + log_context = {**(correlation or {}), "run_id": run_id} + record_outcome = on_outcome if on_outcome is not None else lambda _outcome: None if not self.lock.acquire(blocking=False): - yield 'event: error\ndata: {"message": "already running"}\n\n' - return + on_event('event: error\ndata: {"message": "already running"}\n\n') + record_outcome("failed") + return "failed" set_status(app, run_id, phase=Phase.running, detail="optimization running") log.info( @@ -430,73 +368,65 @@ async def stream_best( 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( - n_trials=n_trials, callback=callback, events=events, loop=loop + loop, events, stop_flag, n_trials, self._trial_payload + ) + max_study_seconds = max_study_seconds_from_environment() + study_deadline = ( + loop.time() + max_study_seconds if max_study_seconds > 0 else None ) - next_heartbeat = loop.time() + SSE_HEARTBEAT_SECONDS completed = False - try: while True: - if await request.is_disconnected(): + if cancel_event.is_set(): 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}, + "optimization run cancelled, stopping study", + extra={"event": "study_cancelled", **log_context}, ) - break - heartbeat_wait = max(0.0, next_heartbeat - loop.time()) + record_outcome("cancelled") + return "cancelled" + if study_deadline is not None and loop.time() >= study_deadline: + stop_flag.set() + if events.empty() and worker.is_alive(): + message = ( + "optimization study exceeded its " + f"{max_study_seconds:g} second execution limit" + ) + set_status(app, run_id, phase=Phase.error, detail=message) + log.warning( + "optimization study timed out", + extra={ + "event": "study_timeout", + "max_study_seconds": max_study_seconds, + "trials": n_trials, + **log_context, + }, + ) + payload = {"state": "ERROR", "message": message} + on_event(f"data: {json.dumps(payload)}\n\n") + record_outcome("failed") + return "failed" + # Items are still queued — or the worker already + # finished and its completion sentinel is in flight — so + # keep draining: a study that actually finished within + # the limit is reported as completed. try: item = await asyncio.wait_for( - events.get(), - timeout=min(_DISCONNECT_POLL_SECONDS, heartbeat_wait), + events.get(), timeout=_DISCONNECT_POLL_SECONDS ) 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", - ) + 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, - }, + extra={"event": "study_completed", "trials": n_trials, **log_context}, ) - yield "event: done\ndata: {}\n\n" - break + on_event("event: done\ndata: {}\n\n") + record_outcome("completed") + return "completed" event = cast(dict[str, Any], item) if event.get("state") == "ERROR": set_status( @@ -507,20 +437,15 @@ def callback(study: optuna.Study, trial: optuna.trial.FrozenTrial) -> None: ) log.warning( "optimization study failed", - extra={ - "event": "study_failed", - **log_context, - }, + extra={"event": "study_failed", **log_context}, ) - yield f"data: {json.dumps(event)}\n\n" - if event.get("state") == "ERROR": - break + on_event(f"data: {json.dumps(event)}\n\n") + record_outcome("failed") + return "failed" + on_event(f"data: {json.dumps(event)}\n\n") 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(): @@ -535,31 +460,5 @@ def callback(study: optuna.Study, trial: optuna.trial.FrozenTrial) -> None: "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/src/telemetry.py b/apps/petrinaut-opt/src/telemetry.py index 9e0bd0d2b61..2f27e1fd7ea 100644 --- a/apps/petrinaut-opt/src/telemetry.py +++ b/apps/petrinaut-opt/src/telemetry.py @@ -42,8 +42,14 @@ _DEFAULT_SERVICE_NAME = "Petrinaut Optimizer" _DEFAULT_PROTOCOL = "grpc" -_SERVICE_LOGGER_NAMES = ("pn_api", "pn_optimize", "pn_telemetry") -_FASTAPI_EXCLUDED_URLS = r"status$" +_SERVICE_LOGGER_NAMES = ("pn_api", "pn_optimize", "pn_runs", "pn_telemetry") +# `/status*` is probe noise. The run events route is excluded because its +# response live-tails for as long as the consumer stays attached: an +# auto-instrumented SERVER span would last minutes (skewing the RED latency +# SLIs and only exporting on disconnect). The route opens its own short +# manual SERVER span over the attach itself instead — see +# `get_optimize_run_events` in `optimization_api.py`. +_FASTAPI_EXCLUDED_URLS = r"/status(/|$),/optimize/runs/[^/]+/events" log = logging.getLogger("pn_telemetry") diff --git a/apps/petrinaut-opt/tests/test_optimization_api.py b/apps/petrinaut-opt/tests/test_optimization_api.py index 0d057d010b3..4db35cfce04 100644 --- a/apps/petrinaut-opt/tests/test_optimization_api.py +++ b/apps/petrinaut-opt/tests/test_optimization_api.py @@ -3,95 +3,52 @@ import asyncio import logging import threading +import time +from collections.abc import Callable from typing import Any import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -from src import optimization_api +from src import optimization_api, optimization_runs +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 - - 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") +class RecordingModel: + """Track how the CLI adapter is shut down.""" + def __init__(self) -> None: + self.close_calls: list[bool] = [] -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 + def close(self, *, graceful: bool = True) -> None: + self.close_calls.append(graceful) -def test_rejects_oversized_manifests_on_both_routes(monkeypatch) -> None: - monkeypatch.setattr(optimization_api, "MAX_REQUEST_BODY_BYTES", 8) +class FakeOptimizer: + """Minimal pump-driven optimizer double for create-route tests.""" - 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}') + n_trials = 1 - assert all_response.status_code == 413 - assert best_response.status_code == 413 + def __init__(self) -> None: + self.pn_model = RecordingModel() + + 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: @@ -120,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": [], @@ -136,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, @@ -180,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 @@ -207,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, @@ -237,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"}, ) @@ -253,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() @@ -276,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() @@ -319,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, @@ -448,18 +288,671 @@ 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: +class ScriptedDetachedOptimizer: + """Pump-driven double for detached-run endpoint tests. + + Emits one data frame, optionally holds (releasable from the test thread, + cancellable through the run), then finishes with a second data frame and + the done frame — mirroring the real engine's frame bodies. + """ + + n_trials = 2 + + def __init__(self, *, hold: bool = False) -> None: + self.pn_model = RecordingModel() + self.hold = hold + self.release = threading.Event() + + 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: + record_outcome = on_outcome if on_outcome is not None else lambda _o: None + on_event( + 'data: {"step": 0, "params": {"rate": 1.0}, ' + '"init_state": {}, "metric": 2.0, "state": "COMPLETE"}\n\n' + ) + while self.hold and not self.release.is_set(): + if cancel_event.is_set(): + self.pn_model.close(graceful=False) + record_outcome("cancelled") + return "cancelled" + await asyncio.sleep(0.005) + on_event( + 'data: {"step": 1, "params": {"rate": 1.5}, ' + '"init_state": {}, "metric": 3.0, "state": "COMPLETE"}\n\n' + ) + on_event("event: done\ndata: {}\n\n") + self.pn_model.close(graceful=True) + record_outcome("completed") + return "completed" + + +def _wait_until(predicate: Callable[[], bool], timeout: float = 2.0) -> None: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(0.01) + raise AssertionError("condition not met within the timeout") + + +def _event_ids(text: str) -> list[int]: + return [ + int(line.removeprefix("id: ")) + for line in text.splitlines() + if line.startswith("id: ") + ] + + +class _AttachedConsumerRequest: + """Request double for direct attachment calls; connected until closed.""" + + def __init__(self, app: FastAPI) -> None: + self.app = app + self.headers: dict[str, str] = {} + + async def is_disconnected(self) -> bool: + return False + + +async def _start_detached_run( + optimization_manifest: dict, + *, + detach_grace_seconds: float = 300, +) -> tuple[FastAPI, str]: + """Create a detached run through the real endpoint, direct-call style. + + The TestClient transport buffers whole response bodies, so tests that + must read a live SSE attachment incrementally call the endpoints + directly and drive ``body_iterator`` themselves. + """ + 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 + test_app.state.optimization_runs = optimization_runs.OptimizationRunRegistry( + detach_grace_seconds=detach_grace_seconds + ) + scope = { + "type": "http", + "app": test_app, + "method": "POST", + "path": "/optimize/runs", + "headers": [], + "query_string": b"", + } + request = optimization_api.Request(scope) + payload = await optimization_api.post_optimize_runs( + request, optimization_manifest, optimization_api.Response() + ) + return test_app, payload.run_id + + +async def _attach(test_app: FastAPI, run_id: str, cursor: int | None = None) -> Any: + response = await optimization_api.get_optimize_run_events( + run_id, + _AttachedConsumerRequest(test_app), # type: ignore[arg-type] + cursor=cursor, + ) + return response.body_iterator + + +def _frame_ids(frames: list[str]) -> list[int]: + return [ + int(frame.split("\n", 1)[0].removeprefix("id: ")) + for frame in frames + if frame.startswith("id: ") + ] + + +def test_create_detached_run_returns_201_and_holds_the_slot_until_terminal( + optimization_manifest: dict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + optimizer = ScriptedDetachedOptimizer(hold=True) + monkeypatch.setattr( + optimization_api, "initialize_optimizer", lambda _manifest, **_kwargs: optimizer + ) + + with TestClient(optimization_api.app) as client: + response = client.post("/optimize/runs", json=optimization_manifest) + assert response.status_code == 201 + run_id = response.json()["run_id"] + assert response.headers["x-optimization-run-id"] == run_id + # No consumer is attached, yet the slot stays held by the run. + assert optimization_api.app.state.active_optimizations == 1 + run = optimization_api.app.state.optimization_runs.get(run_id) + assert run is not None + assert run.state is RunState.running + + optimizer.release.set() + _wait_until( + lambda: optimization_api.app.state.active_optimizations == 0 + ) + _wait_until(lambda: run.state is RunState.completed) + # The pump closed the CLI gracefully; the idempotent run cleanup adds + # its prompt close, which the real adapter treats as a no-op. + assert optimizer.pn_model.close_calls == [True, False] + + +def test_attach_replays_buffered_events_then_tails_live_ones( + optimization_manifest: dict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + optimizer = ScriptedDetachedOptimizer(hold=True) + monkeypatch.setattr( + optimization_api, "initialize_optimizer", lambda _manifest, **_kwargs: optimizer + ) + + async def scenario() -> list[str]: + test_app, run_id = await _start_detached_run(optimization_manifest) + registry = test_app.state.optimization_runs + try: + response = await optimization_api.get_optimize_run_events( + run_id, + _AttachedConsumerRequest(test_app), # type: ignore[arg-type] + cursor=None, + ) + assert response.headers["x-optimization-run-id"] == run_id + assert response.media_type == "text/event-stream" + stream = response.body_iterator + frames = [await anext(stream)] + # The buffered first trial replays before any live event. + assert frames[0].startswith("id: 1\n") + optimizer.release.set() + async for frame in stream: + frames.append(frame) + run = registry.get(run_id) + assert run is not None + assert run.state is RunState.completed + assert test_app.state.active_optimizations == 0 + return frames + finally: + await registry.shutdown() + + frames = asyncio.run(scenario()) + + assert _frame_ids(frames) == [1, 2, 3] + assert frames[-1] == "id: 3\nevent: done\ndata: {}\n\n" + + +def test_cursor_reattach_skips_already_seen_events( + optimization_manifest: dict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + optimizer = ScriptedDetachedOptimizer() + monkeypatch.setattr( + optimization_api, "initialize_optimizer", lambda _manifest, **_kwargs: optimizer + ) + + with TestClient(optimization_api.app) as client: + run_id = client.post("/optimize/runs", json=optimization_manifest).json()[ + "run_id" + ] + run = optimization_api.app.state.optimization_runs.get(run_id) + assert run is not None + _wait_until(lambda: run.state is RunState.completed) + + from_cursor = client.get(f"/optimize/runs/{run_id}/events?cursor=1") + from_header = client.get( + f"/optimize/runs/{run_id}/events", headers={"Last-Event-ID": "2"} + ) + cursor_wins = client.get( + f"/optimize/runs/{run_id}/events?cursor=0", + headers={"Last-Event-ID": "2"}, + ) + + assert _event_ids(from_cursor.text) == [2, 3] + assert _event_ids(from_header.text) == [3] + assert _event_ids(cursor_wins.text) == [1, 2, 3] + assert cursor_wins.text.endswith("id: 3\nevent: done\ndata: {}\n\n") + + +def test_detaching_a_consumer_does_not_cancel_the_run( + optimization_manifest: dict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + optimizer = ScriptedDetachedOptimizer(hold=True) + monkeypatch.setattr( + optimization_api, "initialize_optimizer", lambda _manifest, **_kwargs: optimizer + ) + + async def scenario() -> list[str]: + test_app, run_id = await _start_detached_run(optimization_manifest) + registry = test_app.state.optimization_runs + try: + run = registry.get(run_id) + assert run is not None + stream = await _attach(test_app, run_id) + assert (await anext(stream)).startswith("id: 1\n") + # The consumer drops mid-run; the run must keep its slot and + # keep optimizing. + await stream.aclose() + assert run.attached is False + assert run.state is RunState.running + assert test_app.state.active_optimizations == 1 + + optimizer.release.set() + await asyncio.wait_for(run.finished.wait(), 2) + assert run.state is RunState.completed + assert test_app.state.active_optimizations == 0 + + # A later attachment replays the full log, then closes. + replay = await _attach(test_app, run_id) + return [frame async for frame in replay] + finally: + await registry.shutdown() + + frames = asyncio.run(scenario()) + + assert _frame_ids(frames) == [1, 2, 3] + assert frames[-1] == "id: 3\nevent: done\ndata: {}\n\n" + + +def test_a_second_attachment_supersedes_the_first( + optimization_manifest: dict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + optimizer = ScriptedDetachedOptimizer(hold=True) + monkeypatch.setattr( + optimization_api, "initialize_optimizer", lambda _manifest, **_kwargs: optimizer + ) + + async def scenario() -> None: + test_app, run_id = await _start_detached_run(optimization_manifest) + registry = test_app.state.optimization_runs + try: + first = await _attach(test_app, run_id) + first_frame = await anext(first) + assert first_frame.startswith("id: 1\n") + + second = await _attach(test_app, run_id) + # The superseded first stream ends with the attachment-scoped + # sentinel (no id line: it is not part of the run's log), so its + # consumer sees a clean, terminal end instead of a truncated + # stream it would reconnect after; bounded so a supersession + # regression fails fast instead of stalling until the ~30s + # heartbeat. + assert ( + await asyncio.wait_for(anext(first), timeout=1) + == "event: superseded\ndata: {}\n\n" + ) + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(anext(first), timeout=1) + run = registry.get(run_id) + assert run is not None + assert run.attached is True + + assert await anext(second) == first_frame + optimizer.release.set() + remaining = [frame async for frame in second] + assert _frame_ids(remaining) == [2, 3] + finally: + await registry.shutdown() + + asyncio.run(scenario()) + + +def test_delete_cancels_a_detached_run_and_is_idempotent( + optimization_manifest: dict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + optimizer = ScriptedDetachedOptimizer(hold=True) + monkeypatch.setattr( + optimization_api, "initialize_optimizer", lambda _manifest, **_kwargs: optimizer + ) + + with TestClient(optimization_api.app) as client: + run_id = client.post("/optimize/runs", json=optimization_manifest).json()[ + "run_id" + ] + assert optimization_api.app.state.active_optimizations == 1 + + response = client.delete(f"/optimize/runs/{run_id}") + assert response.status_code == 204 + run = optimization_api.app.state.optimization_runs.get(run_id) + assert run is not None + assert run.state is RunState.cancelled + assert run.events[-1] == (2, CANCELLED_FRAME) + # The pump closed the CLI promptly and the cleanup added its + # idempotent prompt close; neither waited for a graceful EOF. + assert optimizer.pn_model.close_calls == [False, False] + assert optimization_api.app.state.active_optimizations == 0 + + second = client.delete(f"/optimize/runs/{run_id}") + assert second.status_code == 204 + assert optimization_api.app.state.active_optimizations == 0 + assert optimizer.pn_model.close_calls == [False, False] + + replay = client.get(f"/optimize/runs/{run_id}/events") + + assert _event_ids(replay.text) == [1, 2] + assert replay.text.endswith(f"id: 2\n{CANCELLED_FRAME}") + + +def test_the_reaper_cancels_a_run_nobody_attached_to( + optimization_manifest: dict, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setenv( + optimization_runs.DETACH_GRACE_ENVIRONMENT_VARIABLE, "0.05" + ) + optimizer = ScriptedDetachedOptimizer(hold=True) + monkeypatch.setattr( + optimization_api, "initialize_optimizer", lambda _manifest, **_kwargs: optimizer + ) + + with caplog.at_level(logging.WARNING, logger="pn_runs"): + with TestClient(optimization_api.app) as client: + run_id = client.post( + "/optimize/runs", json=optimization_manifest + ).json()["run_id"] + run = optimization_api.app.state.optimization_runs.get(run_id) + assert run is not None + _wait_until(lambda: run.state is RunState.cancelled) + _wait_until( + lambda: optimization_api.app.state.active_optimizations == 0 + ) + assert run.events[-1] == (2, CANCELLED_FRAME) + + reaped = next( + record + for record in caplog.records + if getattr(record, "event", None) == "run_reaped" + ) + assert reaped.run_id == run_id + + +def test_the_reaper_spares_a_run_with_an_attached_consumer( + optimization_manifest: dict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + optimizer = ScriptedDetachedOptimizer(hold=True) + monkeypatch.setattr( + optimization_api, "initialize_optimizer", lambda _manifest, **_kwargs: optimizer + ) + + async def scenario() -> None: + test_app, run_id = await _start_detached_run( + optimization_manifest, detach_grace_seconds=0.05 + ) + registry = test_app.state.optimization_runs + try: + run = registry.get(run_id) + assert run is not None + stream = await _attach(test_app, run_id) + assert (await anext(stream)).startswith("id: 1\n") + # Stay attached well past several grace periods; the reaper must + # leave the run alone. + await asyncio.sleep(0.2) + assert run.state is RunState.running + + optimizer.release.set() + remaining = [frame async for frame in stream] + assert _frame_ids(remaining) == [2, 3] + assert run.state is RunState.completed + finally: + await registry.shutdown() + + asyncio.run(scenario()) + + +def test_background_detach_covers_an_attachment_that_never_starts( + optimization_manifest: dict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An aborted attach may never pull the body, skipping generator finallys. + + Only the response's background task remains to end the attachment, and it + must run on the event loop (a sync callable would land in Starlette's + threadpool, where ``end_attachment``'s ``get_running_loop()`` raises) — + otherwise the run stays marked attached and is shielded from the reaper. + """ + optimizer = ScriptedDetachedOptimizer(hold=True) + monkeypatch.setattr( + optimization_api, "initialize_optimizer", lambda _manifest, **_kwargs: optimizer + ) + + async def scenario() -> None: + test_app, run_id = await _start_detached_run(optimization_manifest) + registry = test_app.state.optimization_runs + try: + run = registry.get(run_id) + assert run is not None + response = await optimization_api.get_optimize_run_events( + run_id, + _AttachedConsumerRequest(test_app), # type: ignore[arg-type] + cursor=None, + ) + assert run.attached is True + assert response.background is not None + + # The client is gone before the body iterator ever starts. + await response.background() + + assert run.attached is False + # The reaper's grace clock restarts from the backstop detach. + assert run.last_detached_at > run.created_at + # A later attachment is unaffected by the stale epoch. + replay = await _attach(test_app, run_id) + assert (await anext(replay)).startswith("id: 1\n") + await replay.aclose() + await response.body_iterator.aclose() + finally: + optimizer.release.set() + await registry.shutdown() + + asyncio.run(scenario()) + + +def test_unknown_detached_runs_return_404() -> None: + with TestClient(optimization_api.app) as client: + assert client.get("/optimize/runs/missing/events").status_code == 404 + assert client.delete("/optimize/runs/missing").status_code == 404 + + +def test_create_detached_run_preserves_the_capacity_rejection( + 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/runs", json=optimization_manifest) + + assert response.status_code == 429 + assert response.headers["retry-after"] == str(optimization_api.RETRY_AFTER_SECONDS) + + +def test_create_detached_run_rejects_an_oversized_manifest( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(optimization_api, "MAX_REQUEST_BODY_BYTES", 8) + + with TestClient(optimization_api.app) as client: + response = client.post("/optimize/runs", content=b'{"long":true}') + + assert response.status_code == 413 + + +def test_create_detached_run_reports_initialization_failure_with_the_run_id( + optimization_manifest: dict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + def initialize(_manifest: dict[str, Any], **_kwargs: Any) -> Any: + raise RuntimeError("manifest rejected by CLI") + + monkeypatch.setattr(optimization_api, "initialize_optimizer", initialize) + + with TestClient(optimization_api.app) as client: + response = client.post("/optimize/runs", json=optimization_manifest) + 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( + optimization_manifest: dict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An owner tag makes the run single-flight and invisible to others.""" + optimizer = ScriptedDetachedOptimizer(hold=True) + second_optimizer = ScriptedDetachedOptimizer(hold=True) + optimizers = iter([optimizer, second_optimizer]) + monkeypatch.setattr( + optimization_api, + "initialize_optimizer", + lambda _manifest, **_kwargs: next(optimizers), + ) + + with TestClient(optimization_api.app) as client: + created = client.post( + "/optimize/runs", + json=optimization_manifest, + headers={"x-hash-account-id": "user-1"}, + ) + assert created.status_code == 201 + run_id = created.json()["run_id"] + + # The same account cannot start a second run while this one lives. + busy = client.post( + "/optimize/runs", + json=optimization_manifest, + headers={"x-hash-account-id": "user-1"}, + ) + assert busy.status_code == 429 + assert ( + busy.json()["detail"] + == "An optimization is already running for this account" + ) + + # Another account is admitted normally (global slots permitting). + other = client.post( + "/optimize/runs", + json=optimization_manifest, + headers={"x-hash-account-id": "user-2"}, + ) + assert other.status_code == 201 + + # A foreign or absent tag cannot see the owned run — 404, identical + # to an unknown id, on both attach and cancel. + for headers in ({}, {"x-hash-account-id": "user-2"}): + assert ( + client.get( + f"/optimize/runs/{run_id}/events", headers=headers + ).status_code + == 404 + ) + assert ( + client.delete( + f"/optimize/runs/{run_id}", headers=headers + ).status_code + == 404 + ) + + # The owner attaches and cancels normally; the attachment reports + # the manifest's requested trial count for synthesized summaries. + optimizer.release.set() + second_optimizer.release.set() + attached = client.get( + f"/optimize/runs/{run_id}/events", + headers={"x-hash-account-id": "user-1"}, + ) + assert attached.status_code == 200 + assert attached.headers["x-requested-trials"] == "2" + assert ( + client.delete( + f"/optimize/runs/{run_id}", + headers={"x-hash-account-id": "user-1"}, + ).status_code + == 204 + ) + + +def test_account_single_flight_frees_after_terminal_and_ignores_ownerless_runs( + optimization_manifest: dict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + optimizer = ScriptedDetachedOptimizer() + monkeypatch.setattr( + optimization_api, "initialize_optimizer", lambda _manifest, **_kwargs: optimizer + ) + + with TestClient(optimization_api.app) as client: + # Runs created without a tag stay open and never block an account. + untagged = client.post("/optimize/runs", json=optimization_manifest) + assert untagged.status_code == 201 + run_id = untagged.json()["run_id"] + run = optimization_api.app.state.optimization_runs.get(run_id) + assert run is not None + assert run.account_id is None + _wait_until(lambda: run.state is not RunState.running) + + # A terminal run stops blocking its account: the tagged create is + # admitted even though the (finished) untagged run is retained. + tagged = client.post( + "/optimize/runs", + json=optimization_manifest, + headers={"x-hash-account-id": "user-1"}, + ) + assert tagged.status_code == 201 + tagged_run = optimization_api.app.state.optimization_runs.get( + tagged.json()["run_id"] + ) + assert tagged_run is not None + _wait_until(lambda: tagged_run.state is not RunState.running) + + again = client.post( + "/optimize/runs", + json=optimization_manifest, + headers={"x-hash-account-id": "user-1"}, + ) + assert again.status_code == 201 + + +def test_openapi_exposes_the_detached_run_paths() -> 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" + create = schema["paths"]["/optimize/runs"]["post"] + request_schema = create["requestBody"]["content"]["application/json"]["schema"] + assert request_schema["type"] == "object" + assert "201" in create["responses"] + assert "Retry-After" in create["responses"]["429"]["headers"] + assert "413" in create["responses"] + assert "500" in create["responses"] + + events = schema["paths"]["/optimize/runs/{run_id}/events"]["get"] + stream_schema = events["responses"]["200"]["content"]["text/event-stream"][ + "schema" + ] + assert stream_schema["type"] == "string" + assert "event: cancelled" in events["responses"]["200"]["description"] + assert "404" in events["responses"] + + run_path = schema["paths"]["/optimize/runs/{run_id}"] + assert set(run_path) == {"delete"} + assert "204" in run_path["delete"]["responses"] + assert "404" in run_path["delete"]["responses"] diff --git a/apps/petrinaut-opt/tests/test_optimization_runs.py b/apps/petrinaut-opt/tests/test_optimization_runs.py new file mode 100644 index 00000000000..6cd4987a9cd --- /dev/null +++ b/apps/petrinaut-opt/tests/test_optimization_runs.py @@ -0,0 +1,597 @@ +from __future__ import annotations + +import asyncio +import logging +import threading +from typing import Any + +import pytest +from fastapi import FastAPI + +from src import optimization_runs +from src.petrinaut_optimizer import PetrinautOptimizer +from src.optimization_runs import ( + CANCELLED_FRAME, + DETACH_GRACE_ENVIRONMENT_VARIABLE, + OptimizationRun, + OptimizationRunRegistry, + RunState, + attachment_event_stream, + detach_grace_seconds_from_environment, +) +from src.utils import Phase, StatusStore + + +FIRST_FRAME = ( + 'data: {"step": 0, "params": {"rate": 1.0}, ' + '"init_state": {}, "metric": 2.0, "state": "COMPLETE"}\n\n' +) +DONE_FRAME = "event: done\ndata: {}\n\n" + + +class ConnectedRequest: + headers: dict[str, str] = {} + + async def is_disconnected(self) -> bool: + return False + + +class DisconnectingRequest(ConnectedRequest): + def __init__(self) -> None: + self.polls = 0 + + async def is_disconnected(self) -> bool: + self.polls += 1 + return self.polls > 1 + + +class HoldingPumpOptimizer: + """Pump double that emits one frame, then waits for cancellation.""" + + n_trials = 1 + + def __init__(self) -> None: + self.pump_started = asyncio.Event() + + async def pump_events( + self, + _app: Any, + _run_id: str, + _n_trials: int, + *, + on_event: Any, + cancel_event: asyncio.Event, + on_outcome: Any = None, + correlation: Any = None, + ) -> str: + on_event(FIRST_FRAME) + self.pump_started.set() + await cancel_event.wait() + if on_outcome is not None: + on_outcome("cancelled") + return "cancelled" + + +class CompletingPumpOptimizer: + """Pump double that finishes immediately with a done frame.""" + + n_trials = 1 + + async def pump_events( + self, + _app: Any, + _run_id: str, + _n_trials: int, + *, + on_event: Any, + cancel_event: asyncio.Event, + on_outcome: Any = None, + correlation: Any = None, + ) -> str: + on_event(FIRST_FRAME) + on_event(DONE_FRAME) + if on_outcome is not None: + on_outcome("completed") + return "completed" + + +class CompletedThenBlockedPumpOptimizer: + """Pump double stuck in its (cancellable) teardown after deciding. + + Mirrors the real engine's shape: the done frame is appended and the + outcome recorded before the CLI-shutdown ``finally`` — a cancellable + window in which ``registry.shutdown()`` may land its cancellation. + """ + + n_trials = 1 + + def __init__(self) -> None: + self.teardown_entered = asyncio.Event() + + async def pump_events( + self, + _app: Any, + _run_id: str, + _n_trials: int, + *, + on_event: Any, + cancel_event: asyncio.Event, + on_outcome: Any = None, + correlation: Any = None, + ) -> str: + on_event(FIRST_FRAME) + on_event(DONE_FRAME) + if on_outcome is not None: + on_outcome("completed") + self.teardown_entered.set() + await asyncio.Event().wait() + return "completed" + + +class CrashingPumpOptimizer: + """Pump double whose engine fails outside the study contract.""" + + n_trials = 1 + + async def pump_events(self, *_args: Any, **_kwargs: Any) -> str: + raise RuntimeError("pump crashed: user_secret_xyz") + + +def _cleanup_recorder() -> tuple[list[int], Any]: + calls: list[int] = [] + + async def cleanup() -> None: + calls.append(1) + + return calls, cleanup + + +def _status_app_with_run() -> tuple[FastAPI, str]: + app = FastAPI() + app.state.statuses = StatusStore() + return app, app.state.statuses.create().run_id + + +def test_detach_grace_environment_parsing(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(DETACH_GRACE_ENVIRONMENT_VARIABLE, raising=False) + assert detach_grace_seconds_from_environment() == 300.0 + + monkeypatch.setenv(DETACH_GRACE_ENVIRONMENT_VARIABLE, "12.5") + assert detach_grace_seconds_from_environment() == 12.5 + + monkeypatch.setenv(DETACH_GRACE_ENVIRONMENT_VARIABLE, "not-a-number") + assert detach_grace_seconds_from_environment() == 300.0 + + monkeypatch.setenv(DETACH_GRACE_ENVIRONMENT_VARIABLE, "0") + assert detach_grace_seconds_from_environment() == 0.0 + + +@pytest.mark.parametrize("raw", ["inf", "-inf", "nan", "1e400"]) +def test_detach_grace_rejects_non_finite_values( + monkeypatch: pytest.MonkeyPatch, raw: str +) -> None: + """`float()` accepts inf/nan spellings, which must not become a period.""" + monkeypatch.setenv(DETACH_GRACE_ENVIRONMENT_VARIABLE, raw) + assert detach_grace_seconds_from_environment() == 300.0 + + +def test_event_log_sequences_start_at_one_and_increment() -> None: + async def scenario() -> None: + _calls, cleanup = _cleanup_recorder() + run = OptimizationRun(run_id="r1", optimizer=None, cleanup=cleanup) + assert run.append_event(FIRST_FRAME) == 1 + assert run.append_event(DONE_FRAME) == 2 + assert run.events == [(1, FIRST_FRAME), (2, DONE_FRAME)] + + asyncio.run(scenario()) + + +def test_attachment_replays_past_the_cursor_and_closes_on_terminal() -> None: + async def scenario() -> list[str]: + _calls, cleanup = _cleanup_recorder() + run = OptimizationRun(run_id="r1", optimizer=None, cleanup=cleanup) + run.append_event(FIRST_FRAME) + run.append_event(FIRST_FRAME) + run.append_event(DONE_FRAME) + run.mark_terminal(RunState.completed) + + epoch = run.begin_attachment() + frames = [ + frame + async for frame in attachment_event_stream( + run, cursor=2, epoch=epoch, request=ConnectedRequest() + ) + ] + assert run.attached is False + return frames + + frames = asyncio.run(scenario()) + + assert frames == [f"id: 3\n{DONE_FRAME}"] + + +def test_attachment_tails_live_appends_after_the_replay() -> None: + async def scenario() -> list[str]: + _calls, cleanup = _cleanup_recorder() + run = OptimizationRun(run_id="r1", optimizer=None, cleanup=cleanup) + run.append_event(FIRST_FRAME) + epoch = run.begin_attachment() + stream = attachment_event_stream( + run, cursor=0, epoch=epoch, request=ConnectedRequest() + ) + frames = [await anext(stream)] + + async def append_rest() -> None: + run.append_event(FIRST_FRAME) + await asyncio.sleep(0) + run.append_event(DONE_FRAME) + run.mark_terminal(RunState.completed) + + appender = asyncio.create_task(append_rest()) + async for frame in stream: + frames.append(frame) + await appender + return frames + + frames = asyncio.run(scenario()) + + assert frames == [ + f"id: 1\n{FIRST_FRAME}", + f"id: 2\n{FIRST_FRAME}", + f"id: 3\n{DONE_FRAME}", + ] + + +def test_attachment_with_a_cursor_past_the_log_end_still_tails_live_events() -> None: + """An out-of-range cursor must not swallow later frames or the terminal.""" + + async def scenario() -> list[str]: + _calls, cleanup = _cleanup_recorder() + run = OptimizationRun(run_id="r1", optimizer=None, cleanup=cleanup) + run.append_event(FIRST_FRAME) + epoch = run.begin_attachment() + stream = attachment_event_stream( + run, cursor=999, epoch=epoch, request=ConnectedRequest() + ) + + async def append_rest() -> None: + await asyncio.sleep(0.01) + run.append_event(FIRST_FRAME) + run.append_event(DONE_FRAME) + run.mark_terminal(RunState.completed) + + appender = asyncio.create_task(append_rest()) + frames = [ + frame async for frame in stream if not frame.startswith(": heartbeat") + ] + await appender + return frames + + frames = asyncio.run(scenario()) + + # The clamped cursor skips the already-buffered frame but delivers every + # frame appended after attaching, including the terminal one. + assert frames == [f"id: 2\n{FIRST_FRAME}", f"id: 3\n{DONE_FRAME}"] + + +def test_a_second_attachment_supersedes_the_first() -> None: + async def scenario() -> None: + _calls, cleanup = _cleanup_recorder() + run = OptimizationRun(run_id="r1", optimizer=None, cleanup=cleanup) + run.append_event(FIRST_FRAME) + + first_epoch = run.begin_attachment() + first_stream = attachment_event_stream( + run, cursor=0, epoch=first_epoch, request=ConnectedRequest() + ) + assert await anext(first_stream) == f"id: 1\n{FIRST_FRAME}" + + second_epoch = run.begin_attachment() + # The superseded stream ends with the attachment-scoped sentinel (no + # id line), then closes. Bounded so a supersession regression fails + # fast instead of stalling until the ~30s heartbeat. + assert ( + await asyncio.wait_for(anext(first_stream), timeout=1) + == "event: superseded\ndata: {}\n\n" + ) + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(anext(first_stream), timeout=1) + # The stale attachment must not clear the newer one's attached mark. + assert run.attached is True + + second_stream = attachment_event_stream( + run, cursor=0, epoch=second_epoch, request=ConnectedRequest() + ) + assert await anext(second_stream) == f"id: 1\n{FIRST_FRAME}" + await second_stream.aclose() + assert run.attached is False + + asyncio.run(scenario()) + + +def test_a_disconnected_attachment_marks_the_run_detached() -> None: + async def scenario() -> None: + _calls, cleanup = _cleanup_recorder() + run = OptimizationRun(run_id="r1", optimizer=None, cleanup=cleanup) + run.append_event(FIRST_FRAME) + epoch = run.begin_attachment() + frames = [ + frame + async for frame in attachment_event_stream( + run, cursor=0, epoch=epoch, request=DisconnectingRequest() + ) + ] + assert frames == [f"id: 1\n{FIRST_FRAME}"] + assert run.attached is False + assert run.state is RunState.running + + asyncio.run(scenario()) + + +def test_registry_reaps_a_run_never_attached_within_the_grace_period( + caplog: pytest.LogCaptureFixture, +) -> None: + async def scenario() -> None: + app, run_id = _status_app_with_run() + registry = OptimizationRunRegistry( + detach_grace_seconds=0.05, retention_seconds=60 + ) + calls, cleanup = _cleanup_recorder() + run = registry.create_run( + app, + run_id=run_id, + optimizer=HoldingPumpOptimizer(), + cleanup=cleanup, + correlation={"request_id": "request-reap-1"}, + ) + try: + await asyncio.wait_for(run.finished.wait(), 2) + assert run.state is RunState.cancelled + assert run.events[-1] == (2, CANCELLED_FRAME) + assert calls == [1] + status = app.state.statuses.get(run_id) + assert status is not None + assert status.phase is Phase.idle + finally: + await registry.shutdown() + + with caplog.at_level(logging.WARNING, logger="pn_runs"): + asyncio.run(scenario()) + + reaped = next( + record + for record in caplog.records + if getattr(record, "event", None) == "run_reaped" + ) + assert reaped.request_id == "request-reap-1" + + +def test_registry_does_not_reap_an_attached_run() -> None: + async def scenario() -> None: + app, run_id = _status_app_with_run() + registry = OptimizationRunRegistry( + detach_grace_seconds=0.05, retention_seconds=60 + ) + _calls, cleanup = _cleanup_recorder() + optimizer = HoldingPumpOptimizer() + run = registry.create_run( + app, run_id=run_id, optimizer=optimizer, cleanup=cleanup + ) + try: + await asyncio.wait_for(optimizer.pump_started.wait(), 1) + run.begin_attachment() + await asyncio.sleep(0.2) + assert run.state is RunState.running + finally: + await registry.shutdown() + + asyncio.run(scenario()) + + +def test_registry_keeps_disabled_grace_runs_alive() -> None: + async def scenario() -> None: + app, run_id = _status_app_with_run() + registry = OptimizationRunRegistry(detach_grace_seconds=0) + _calls, cleanup = _cleanup_recorder() + optimizer = HoldingPumpOptimizer() + run = registry.create_run( + app, run_id=run_id, optimizer=optimizer, cleanup=cleanup + ) + try: + await asyncio.wait_for(optimizer.pump_started.wait(), 1) + await asyncio.sleep(0.1) + assert run.state is RunState.running + finally: + await registry.shutdown() + + asyncio.run(scenario()) + + +def test_registry_drops_terminal_run_logs_after_the_retention_window() -> None: + async def scenario() -> None: + app, run_id = _status_app_with_run() + registry = OptimizationRunRegistry( + detach_grace_seconds=60, retention_seconds=0.05 + ) + _calls, cleanup = _cleanup_recorder() + run = registry.create_run( + app, run_id=run_id, optimizer=CompletingPumpOptimizer(), cleanup=cleanup + ) + try: + await asyncio.wait_for(run.finished.wait(), 2) + for _ in range(200): + if registry.get(run_id) is None: + break + await asyncio.sleep(0.01) + assert registry.get(run_id) is None + # The status row outlives the replay log, as before. + assert app.state.statuses.get(run_id) is not None + finally: + await registry.shutdown() + + asyncio.run(scenario()) + + +def test_registry_shutdown_cancels_running_pumps() -> None: + async def scenario() -> None: + app, run_id = _status_app_with_run() + registry = OptimizationRunRegistry(detach_grace_seconds=0) + calls, cleanup = _cleanup_recorder() + optimizer = HoldingPumpOptimizer() + run = registry.create_run( + app, run_id=run_id, optimizer=optimizer, cleanup=cleanup + ) + await asyncio.wait_for(optimizer.pump_started.wait(), 1) + + await registry.shutdown() + + assert run.state is RunState.cancelled + assert run.cancel_reason == "service shutting down" + assert run.events[-1] == (2, CANCELLED_FRAME) + assert calls == [1] + + asyncio.run(scenario()) + + +def test_shutdown_during_pump_teardown_keeps_the_decided_outcome() -> None: + """A cancel landing after `done` must not relabel the run as cancelled.""" + + async def scenario() -> None: + app, run_id = _status_app_with_run() + registry = OptimizationRunRegistry(detach_grace_seconds=0) + calls, cleanup = _cleanup_recorder() + optimizer = CompletedThenBlockedPumpOptimizer() + run = registry.create_run( + app, run_id=run_id, optimizer=optimizer, cleanup=cleanup + ) + await asyncio.wait_for(optimizer.teardown_entered.wait(), 1) + + await registry.shutdown() + + assert run.state is RunState.completed + # The event log ends with the single decided terminal frame; no + # cancelled frame is appended after it. + assert run.events == [(1, FIRST_FRAME), (2, DONE_FRAME)] + assert calls == [1] + + asyncio.run(scenario()) + + +class SlowCloseStudyModel: + """Real-optimizer model whose CLI close blocks until released.""" + + def __init__(self, description: dict[str, Any]) -> None: + self.description = description + self.close_started = threading.Event() + self.close_release = threading.Event() + self.close_calls: list[bool] = [] + + def describe_optimization(self) -> dict[str, Any]: + return self.description + + def objective(self, _parameter_values: dict[str, Any]) -> float: + return 1.0 + + def close(self, *, graceful: bool = True) -> None: + self.close_started.set() + self.close_release.wait(timeout=2) + self.close_calls.append(graceful) + + +def test_cancel_during_the_real_pumps_cli_close_keeps_completed( + optimization_description: dict, +) -> None: + """The real engine records its outcome before the cancellable close.""" + + async def scenario() -> None: + app, run_id = _status_app_with_run() + registry = OptimizationRunRegistry(detach_grace_seconds=0) + calls, cleanup = _cleanup_recorder() + model = SlowCloseStudyModel(optimization_description) + optimizer = PetrinautOptimizer(model) # type: ignore[arg-type] + run = registry.create_run( + app, run_id=run_id, optimizer=optimizer, cleanup=cleanup + ) + try: + # The study finishes, the done frame lands, then the pump blocks + # in its CLI-shutdown finally; cancel it right there. + await asyncio.to_thread(model.close_started.wait, 2) + assert run.task is not None + run.task.cancel() + model.close_release.set() + await asyncio.gather(run.task, return_exceptions=True) + + assert run.state is RunState.completed + assert run.events[-1][1] == DONE_FRAME + assert CANCELLED_FRAME not in [frame for _seq, frame in run.events] + assert calls == [1] + status = app.state.statuses.get(run_id) + assert status is not None + assert status.phase is Phase.done + finally: + await registry.shutdown() + + asyncio.run(scenario()) + + +def test_a_pump_crash_appends_a_bounded_backstop_error_frame( + caplog: pytest.LogCaptureFixture, +) -> None: + async def scenario() -> None: + app, run_id = _status_app_with_run() + registry = OptimizationRunRegistry(detach_grace_seconds=60) + calls, cleanup = _cleanup_recorder() + run = registry.create_run( + app, run_id=run_id, optimizer=CrashingPumpOptimizer(), cleanup=cleanup + ) + try: + await asyncio.wait_for(run.finished.wait(), 2) + assert run.state is RunState.failed + assert run.events == [ + ( + 1, + 'data: {"state": "ERROR", ' + '"message": "optimization run failed"}\n\n', + ) + ] + assert calls == [1] + finally: + await registry.shutdown() + + with caplog.at_level(logging.ERROR, logger="pn_runs"): + asyncio.run(scenario()) + + failure = next( + record + for record in caplog.records + if getattr(record, "event", None) == "run_pump_failed" + ) + assert failure.error_type == "RuntimeError" + # The raw pump error may quote user content, so it is never logged or + # replayed to consumers. + assert "user_secret_xyz" not in failure.getMessage() + + +def test_attachment_heartbeats_while_tailing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(optimization_runs, "SSE_HEARTBEAT_SECONDS", 0.01) + + async def scenario() -> list[str]: + _calls, cleanup = _cleanup_recorder() + run = OptimizationRun(run_id="r1", optimizer=None, cleanup=cleanup) + epoch = run.begin_attachment() + stream = attachment_event_stream( + run, cursor=0, epoch=epoch, request=ConnectedRequest() + ) + frames = [await anext(stream)] + run.append_event(DONE_FRAME) + run.mark_terminal(RunState.completed) + async for frame in stream: + frames.append(frame) + return frames + + frames = asyncio.run(scenario()) + + assert frames[0] == ": heartbeat\n\n" + assert frames[-1] == f"id: 1\n{DONE_FRAME}" diff --git a/apps/petrinaut-opt/tests/test_petrinaut_optimizer.py b/apps/petrinaut-opt/tests/test_petrinaut_optimizer.py index 79d0f6d6e13..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: @@ -184,6 +154,15 @@ def test_uses_the_cli_supplied_seed_for_deterministic_sampling( [ {"direction": "up"}, {"study": {"trials": 0, "sampler": "random", "seed": 42}}, + # The service-side trial cap bounds every run's event log even when + # the CLI's reported study is huge. + { + "study": { + "trials": petrinaut_optimizer.MAX_STUDY_TRIALS + 1, + "sampler": "random", + "seed": 42, + } + }, {"study": {"trials": 1, "sampler": "unknown", "seed": 42}}, {"study": {"trials": 1, "sampler": "random", "seed": -1}}, { @@ -223,26 +202,69 @@ def test_rejects_invalid_cli_descriptions( ) -def test_stream_all_logs_the_study_lifecycle_with_correlation( +def _status_app_with_run() -> tuple[FastAPI, str]: + app = FastAPI() + app.state.statuses = StatusStore() + return app, app.state.statuses.create().run_id + + +def test_pump_events_appends_frames_and_completes( + optimization_description: dict, +) -> None: + model = FakeModel(optimization_description) + optimizer = PetrinautOptimizer(model) # type: ignore[arg-type] + app, run_id = _status_app_with_run() + frames: list[str] = [] + + async def drive() -> str: + return await optimizer.pump_events( + app, + run_id, + optimizer.n_trials, + on_event=frames.append, + cancel_event=asyncio.Event(), + ) + + state = asyncio.run(drive()) + data = [ + json.loads(frame.removeprefix("data: ")) + for frame in frames + if frame.startswith("data: ") + ] + + assert state == "completed" + 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 model.close_calls == [True] + status = app.state.statuses.get(run_id) + assert status is not None + 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] - request = ConnectedRequest() - request.headers = {"x-hash-request-id": "request-s1"} # type: ignore[attr-defined] - run_id = _run_id(request) + app, run_id = _status_app_with_run() - async def consume() -> None: - async for _frame in optimizer.stream_all( - request, + async def drive() -> str: + return await optimizer.pump_events( + app, run_id, - optimizer.n_trials, # type: ignore[arg-type] - ): - pass + 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(consume()) + asyncio.run(drive()) events = { getattr(record, "event", None): record @@ -256,215 +278,249 @@ async def consume() -> None: assert record.trials == optimizer.n_trials -def test_disconnect_is_logged_with_the_run_id( +def test_pump_events_reports_a_study_failure_without_done( 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) + model = FailingModel( + optimization_description, PetrinautClientError("transport failed") + ) optimizer = PetrinautOptimizer(model) # type: ignore[arg-type] - request = DisconnectedAfterWorkerStarts(model) - run_id = _run_id(request) + app, run_id = _status_app_with_run() + frames: list[str] = [] - async def consume() -> None: - async for _frame in optimizer.stream_all( - request, + async def drive() -> str: + return await optimizer.pump_events( + app, run_id, - optimizer.n_trials, # type: ignore[arg-type] - ): - pass - model.release.set() - await asyncio.sleep(0.05) + optimizer.n_trials, + on_event=frames.append, + cancel_event=asyncio.Event(), + ) - with caplog.at_level(logging.INFO, logger="pn_optimize"): - asyncio.run(consume()) + state = asyncio.run(drive()) + status = app.state.statuses.get(run_id) - disconnected = next( - record - for record in caplog.records - if getattr(record, "event", None) == "client_disconnected" - ) - assert disconnected.run_id == run_id + assert state == "failed" + assert json.loads(frames[-1].removeprefix("data: ")) == { + "state": "ERROR", + "message": "transport failed", + } + assert "event: done\ndata: {}\n\n" not in frames + assert model.close_calls == [False] + assert status is not None + assert status.phase is Phase.error -def test_stream_all_preserves_the_existing_sse_frame_shape( +def test_pump_events_cancellation_stops_the_study_promptly( optimization_description: dict, + monkeypatch: pytest.MonkeyPatch, ) -> None: - model = FakeModel(optimization_description) + 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 = 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: ") - ] + app, run_id = _status_app_with_run() + frames: list[str] = [] + cancel_event = asyncio.Event() - 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] + async def drive() -> str: + async def cancel_after_entry() -> None: + await asyncio.to_thread(model.entered.wait, 1) + cancel_event.set() + canceller = asyncio.create_task(cancel_after_entry()) + started_at = time.monotonic() + state = await optimizer.pump_events( + app, + run_id, + optimizer.n_trials, + on_event=frames.append, + cancel_event=cancel_event, + ) + assert time.monotonic() - started_at < 0.5 + await canceller + model.release.set() + await asyncio.sleep(0.05) + return state -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: ") - ] + state = asyncio.run(drive()) - 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 state == "cancelled" + # The caller owns the terminal cancelled frame, so none is appended here. + assert frames == [] assert model.closed is True - assert model.close_calls == [True] + assert model.close_calls == [False] -def test_stream_sends_comment_heartbeats_while_a_trial_is_running( - optimization_description: dict, +def test_max_study_seconds_environment_parsing( 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() + variable = petrinaut_optimizer.MAX_STUDY_SECONDS_ENVIRONMENT_VARIABLE + monkeypatch.delenv(variable, raising=False) + assert petrinaut_optimizer.max_study_seconds_from_environment() == 900.0 - 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] - ) - ] + monkeypatch.setenv(variable, "12.5") + assert petrinaut_optimizer.max_study_seconds_from_environment() == 12.5 - frames = asyncio.run(collect()) + monkeypatch.setenv(variable, "not-a-number") + assert petrinaut_optimizer.max_study_seconds_from_environment() == 900.0 - assert ": heartbeat\n\n" in frames - assert frames[-1] == "event: done\ndata: {}\n\n" + monkeypatch.setenv(variable, "0") + assert petrinaut_optimizer.max_study_seconds_from_environment() == 0.0 -@pytest.mark.parametrize("stream_name", ["stream_all", "stream_best"]) -def test_stream_error_is_terminal_and_is_not_followed_by_done( +@pytest.mark.parametrize("raw", ["inf", "-inf", "nan", "1e400"]) +def test_max_study_seconds_rejects_non_finite_values( + monkeypatch: pytest.MonkeyPatch, raw: str +) -> None: + monkeypatch.setenv( + petrinaut_optimizer.MAX_STUDY_SECONDS_ENVIRONMENT_VARIABLE, raw + ) + assert petrinaut_optimizer.max_study_seconds_from_environment() == 900.0 + + +def test_pump_events_enforces_the_study_wall_clock_ceiling( optimization_description: dict, - stream_name: str, + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, ) -> None: - model = FailingModel( - optimization_description, PetrinautClientError("transport failed") + optimization_description["study"]["trials"] = 1 + monkeypatch.setenv( + petrinaut_optimizer.MAX_STUDY_SECONDS_ENVIRONMENT_VARIABLE, "0.05" ) + monkeypatch.setattr(petrinaut_optimizer, "_WORKER_SHUTDOWN_TIMEOUT_SECONDS", 0.01) + model = StubbornModel(optimization_description) 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] - ) - ] + app, run_id = _status_app_with_run() + frames: list[str] = [] + outcomes: list[str] = [] - with caplog.at_level(logging.WARNING, logger="pn_optimize"): - frames = asyncio.run(collect()) - status = request.app.state.statuses.get(run_id) + async def drive() -> str: + started_at = time.monotonic() + state = await optimizer.pump_events( + app, + run_id, + optimizer.n_trials, + on_event=frames.append, + cancel_event=asyncio.Event(), + on_outcome=outcomes.append, + ) + assert time.monotonic() - started_at < 0.5 + model.release.set() + await asyncio.sleep(0.05) + return state - 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 + with caplog.at_level(logging.WARNING, logger="pn_optimize"): + state = asyncio.run(drive()) + status = app.state.statuses.get(run_id) + + assert state == "failed" + assert outcomes == ["failed"] + assert frames == [ + 'data: {"state": "ERROR", "message": "optimization study exceeded ' + 'its 0.05 second execution limit"}\n\n' + ] assert model.closed is True assert model.close_calls == [False] - failure = next( + assert status is not None + assert status.phase is Phase.error + timeout = next( record for record in caplog.records - if getattr(record, "event", None) == "study_failed" + if getattr(record, "event", None) == "study_timeout" ) - assert failure.run_id == run_id - assert "transport failed" not in failure.getMessage() - assert not hasattr(failure, "detail") + assert timeout.run_id == run_id + assert timeout.max_study_seconds == 0.05 -def test_disconnect_closes_cli_before_a_bounded_worker_join( +def test_pump_events_completed_before_the_ceiling_is_not_misreported( optimization_description: dict, monkeypatch: pytest.MonkeyPatch, ) -> None: + monkeypatch.setenv( + petrinaut_optimizer.MAX_STUDY_SECONDS_ENVIRONMENT_VARIABLE, "5" + ) + model = FakeModel(optimization_description) + optimizer = PetrinautOptimizer(model) # type: ignore[arg-type] + app, run_id = _status_app_with_run() + frames: list[str] = [] + outcomes: list[str] = [] + + async def drive() -> str: + return await optimizer.pump_events( + app, + run_id, + optimizer.n_trials, + on_event=frames.append, + cancel_event=asyncio.Event(), + on_outcome=outcomes.append, + ) + + state = asyncio.run(drive()) + + assert state == "completed" + assert outcomes == ["completed"] + assert frames[-1] == "event: done\ndata: {}\n\n" + assert all('"state": "ERROR"' not in frame for frame in frames) + assert model.close_calls == [True] + + +def test_pump_events_drains_a_queued_completion_after_the_deadline( + optimization_description: dict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A completion already queued when the ceiling lapses is still reported.""" optimization_description["study"]["trials"] = 1 - monkeypatch.setattr(petrinaut_optimizer, "_WORKER_SHUTDOWN_TIMEOUT_SECONDS", 0.01) - model = StubbornModel(optimization_description) + monkeypatch.setenv( + petrinaut_optimizer.MAX_STUDY_SECONDS_ENVIRONMENT_VARIABLE, "0.000001" + ) + + def preloaded_worker( + self: PetrinautOptimizer, + _loop: asyncio.AbstractEventLoop, + events: asyncio.Queue, + _stop_flag: threading.Event, + _n_trials: int, + _payload_builder: Any, + ) -> tuple[threading.Thread, Any]: + # The study "finished" before the pump's first deadline check: both + # the trial payload and the sentinel are already queued. + events.put_nowait( + { + "step": 0, + "params": {"rate": 1.0}, + "init_state": {}, + "metric": 2.0, + "state": "COMPLETE", + } + ) + events.put_nowait(petrinaut_optimizer._SENTINEL) + worker = threading.Thread(target=lambda: None, daemon=True) + worker.start() + return worker, petrinaut_optimizer.tracer.start_span("test-study") + + monkeypatch.setattr( + PetrinautOptimizer, "_start_study_worker", preloaded_worker + ) + model = FakeModel(optimization_description) optimizer = PetrinautOptimizer(model) # type: ignore[arg-type] - request = DisconnectedAfterWorkerStarts(model) - run_id = _run_id(request) + app, run_id = _status_app_with_run() + frames: list[str] = [] - 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 + async def drive() -> str: + return await optimizer.pump_events( + app, + run_id, + optimizer.n_trials, + on_event=frames.append, + cancel_event=asyncio.Event(), + ) - frames = asyncio.run(collect()) - status = request.app.state.statuses.get(run_id) + state = asyncio.run(drive()) + status = app.state.statuses.get(run_id) - assert frames == [] - assert model.closed is True - assert model.close_calls == [False] + assert state == "completed" + assert frames[-1] == "event: done\ndata: {}\n\n" + assert all('"state": "ERROR"' not in frame for frame in frames) assert status is not None - assert status.phase is Phase.idle + assert status.phase is Phase.done diff --git a/apps/petrinaut-opt/tests/test_telemetry.py b/apps/petrinaut-opt/tests/test_telemetry.py index 981ce36f92d..10634fb1256 100644 --- a/apps/petrinaut-opt/tests/test_telemetry.py +++ b/apps/petrinaut-opt/tests/test_telemetry.py @@ -78,7 +78,8 @@ def instrument(_app: FastAPI, **kwargs: object) -> None: assert meter_providers == [telemetry._providers[1]] assert instrumentation["tracer_provider"] is telemetry._providers[0] assert instrumentation["meter_provider"] is telemetry._providers[1] - assert instrumentation["excluded_urls"] == "status$" + assert instrumentation["excluded_urls"] == telemetry._FASTAPI_EXCLUDED_URLS + assert "/optimize/runs/[^/]+/events" in instrumentation["excluded_urls"] new_handlers = [ handler for handler in logging.getLogger("pn_api").handlers 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..b466273df7e 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,48 @@ 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"), + ]); }); }); 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..fb0e8b0417e 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,120 @@ -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 decode failure as `protocol` (aborts pass through + * untouched) so the provider reconnects with its cursor instead of failing + * the run on the first dropped connection. + */ +const classifyStreamError = (error: unknown): unknown => + error instanceof Error && error.name !== "AbortError" + ? Object.assign(error, { category: "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/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 c1bba3060e0..632a5c570aa 100644 --- a/libs/@hashintel/petrinaut-core/src/optimization.ts +++ b/libs/@hashintel/petrinaut-core/src/optimization.ts @@ -491,10 +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: optimizationEventSeqSchema, }) .meta({ description: "The optimizer accepted and started the study." }); @@ -506,6 +516,7 @@ export const petrinautOptimizationTrialEventSchema = z objective: z.number().nullable(), state: z.enum(["complete", "pruned", "failed"]), best: optimizationBestSchema.nullable(), + seq: optimizationEventSeqSchema, }) .meta({ description: "One completed Optuna trial and the running best." }); @@ -517,6 +528,7 @@ export const petrinautOptimizationCompleteEventSchema = z prunedTrials: z.number().int().nonnegative(), failedTrials: z.number().int().nonnegative(), best: optimizationBestSchema.nullable(), + seq: optimizationEventSeqSchema, }) .meta({ description: "The final optimization summary." }); @@ -526,6 +538,7 @@ export const petrinautOptimizationErrorEventSchema = z code: z.string(), message: z.string(), retryable: z.boolean(), + seq: optimizationEventSeqSchema, }) .meta({ description: "A terminal optimizer error." }); @@ -573,10 +586,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..dc98b6cd5c4 100644 --- a/libs/@hashintel/petrinaut/src/react/optimizations/provider.test.tsx +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider.test.tsx @@ -2,8 +2,8 @@ * @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 { petrinautOptimizationInputSchema, @@ -94,50 +94,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 +244,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: 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: { infected_ratio: 0.02 }, + 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("lets Remove cancel a possibly-live run after a terminal error", async () => { + const cancelledRunIds: string[] = []; + const capability: PetrinautOptimization = { + 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("aborts and marks an active optimization as cancelled", async () => { + it("cancels a detached run server-side and aborts its attachment", async () => { + const cancelledRunIds: string[] = []; const capability: PetrinautOptimization = { - async *optimize(_request, options) { - yield { type: "started", requestedTrials: 2 }; + 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 +480,370 @@ 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("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..33d2a085d89 100644 --- a/libs/@hashintel/petrinaut/src/react/optimizations/provider.tsx +++ b/libs/@hashintel/petrinaut/src/react/optimizations/provider.tsx @@ -1,10 +1,18 @@ -import { use, useEffect, useRef, useState } from "react"; +import { use, useCallback, useEffect, useRef, useState } from "react"; -import { petrinautOptimizationInputSchema } from "@hashintel/petrinaut-core"; +import { + 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 +21,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 +155,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 +315,332 @@ 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, + status: "error", + connectionState: null, + 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 +651,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 +853,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..2c91142d3f9 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,23 @@ 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/package.json b/libs/@local/petrinaut-optimizer-client/package.json index a4a01510818..a92f472f80b 100644 --- a/libs/@local/petrinaut-optimizer-client/package.json +++ b/libs/@local/petrinaut-optimizer-client/package.json @@ -21,7 +21,8 @@ }, "dependencies": { "@hashintel/petrinaut-core": "workspace:*", - "eventsource-parser": "3.0.8" + "eventsource-parser": "3.0.8", + "openapi-fetch": "0.17.0" }, "devDependencies": { "@apps/petrinaut-opt": "workspace:*", 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 new file mode 100644 index 00000000000..0ff4470a4d3 --- /dev/null +++ b/libs/@local/petrinaut-optimizer-client/src/attach-optimization-run.test.ts @@ -0,0 +1,181 @@ +import { describe, expect, it, vi } from "vitest"; + +import { attachPetrinautOptimizationRunStream } from "./attach-optimization-run.js"; +import { PetrinautOptimizerHttpError } from "./optimizer-http.js"; + +/** Collect every event returned by one attached optimization stream. */ +const collect = async (events: AsyncIterable): Promise => { + const collected = []; + for await (const event of events) { + collected.push(event); + } + return collected; +}; + +describe("attachPetrinautOptimizationRunStream", () => { + it("attaches with a cursor and returns sequenced canonical events", async () => { + const onActivity = vi.fn(); + const signal = new AbortController().signal; + const fetchImpl = vi.fn(async () => + Promise.resolve( + new Response( + 'id: 3\ndata: {"step":2,"params":{"rate":0.4},"init_state":{},"metric":2,"state":"COMPLETE"}\n\n' + + "id: 4\nevent: done\ndata: {}\n\n", + { + headers: { + "content-type": "text/event-stream", + "x-optimization-run-id": "run-42", + "x-requested-trials": "3", + }, + }, + ), + ), + ); + + const { events, optimizationRunId } = + await attachPetrinautOptimizationRunStream({ + endpoint: "http://petrinaut-opt.test", + runId: "run-42", + cursor: 2, + fetchImpl, + onActivity, + requestId: "request-123", + signal, + }); + + expect(optimizationRunId).toBe("run-42"); + // No synthetic started event: the run started before this attachment, + // and `best` stays null because no direction was supplied. + await expect(collect(events)).resolves.toEqual([ + { + type: "trial", + trial: 2, + parameters: { rate: 0.4 }, + objective: 2, + state: "complete", + best: null, + seq: 3, + }, + { + type: "complete", + requestedTrials: 3, + completedTrials: 1, + prunedTrials: 0, + failedTrials: 0, + best: null, + seq: 4, + }, + ]); + expect(fetchImpl).toHaveBeenCalledWith( + new URL("http://petrinaut-opt.test/optimize/runs/run-42/events?cursor=2"), + { + method: "GET", + headers: { + accept: "text/event-stream", + "x-hash-request-id": "request-123", + }, + signal, + }, + ); + expect(onActivity).toHaveBeenCalledOnce(); + }); + + it("omits the cursor parameter for a full replay", async () => { + const fetchImpl = vi.fn(async () => + Promise.resolve( + new Response("id: 1\nevent: done\ndata: {}\n\n", { + headers: { + "content-type": "text/event-stream", + "x-requested-trials": "3", + }, + }), + ), + ); + + const { events } = await attachPetrinautOptimizationRunStream({ + endpoint: "http://petrinaut-opt.test", + runId: "run-42", + fetchImpl, + }); + await collect(events); + + expect(fetchImpl).toHaveBeenCalledWith( + new URL("http://petrinaut-opt.test/optimize/runs/run-42/events"), + expect.anything(), + ); + }); + + it("adapts a cancelled run's terminal frame", async () => { + const { events } = await attachPetrinautOptimizationRunStream({ + endpoint: "http://petrinaut-opt.test", + runId: "run-42", + fetchImpl: async () => + new Response("id: 1\nevent: cancelled\ndata: {}\n\n", { + headers: { + "content-type": "text/event-stream", + "x-requested-trials": "3", + }, + }), + }); + + await expect(collect(events)).resolves.toEqual([ + { + type: "error", + code: "optimization_cancelled", + message: "The optimization was cancelled", + retryable: false, + seq: 1, + }, + ]); + }); + + it("escapes the run id in the events URL", async () => { + const fetchImpl = vi.fn(async () => + Promise.resolve( + new Response("id: 1\nevent: done\ndata: {}\n\n", { + headers: { + "content-type": "text/event-stream", + "x-requested-trials": "3", + }, + }), + ), + ); + + const { events } = await attachPetrinautOptimizationRunStream({ + endpoint: "http://petrinaut-opt.test", + runId: "run/../42", + fetchImpl, + }); + await collect(events); + + expect(fetchImpl).toHaveBeenCalledWith( + new URL("http://petrinaut-opt.test/optimize/runs/run%2F..%2F42/events"), + expect.anything(), + ); + }); + + it("surfaces an unknown or expired run as an HTTP error", async () => { + const result = attachPetrinautOptimizationRunStream({ + endpoint: "http://petrinaut-opt.test", + runId: "run-unknown", + fetchImpl: async () => + Response.json( + { detail: "optimization run not found: run-unknown" }, + { status: 404 }, + ), + }); + + await expect(result).rejects.toBeInstanceOf(PetrinautOptimizerHttpError); + await expect(result).rejects.toMatchObject({ status: 404 }); + }); + + it("rejects a successful response without a body", async () => { + await expect( + attachPetrinautOptimizationRunStream({ + endpoint: "http://petrinaut-opt.test", + runId: "run-42", + fetchImpl: async () => new Response(null, { status: 200 }), + }), + ).rejects.toThrow("Petrinaut optimizer returned an empty response"); + }); +}); diff --git a/libs/@local/petrinaut-optimizer-client/src/attach-optimization-run.ts b/libs/@local/petrinaut-optimizer-client/src/attach-optimization-run.ts new file mode 100644 index 00000000000..aebbc7173ed --- /dev/null +++ b/libs/@local/petrinaut-optimizer-client/src/attach-optimization-run.ts @@ -0,0 +1,111 @@ +import { decodePetrinautOptimizerStream } from "./decode-optimization-stream.js"; +import { + petrinautOptimizerHttpErrorFromResponse, + petrinautOptimizerUrl, +} from "./optimizer-http.js"; + +import type { PetrinautOptimizerFetch } from "./optimizer-http.js"; +import type { + AbortSignalLike, + 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. */ + endpoint: string | URL; + /** Identifier of the detached run to attach to. */ + runId: string; + /** + * Replay frames with sequence numbers greater than this cursor before + * live-tailing. Omitted or `0` requests a full replay. + */ + cursor?: number; + /** Fetch implementation supplied by the current runtime or a test. */ + fetchImpl?: PetrinautOptimizerFetch; + /** Optional maximum UTF-8 size of one upstream SSE event. */ + maxEventBytes?: number; + /** Called whenever upstream bytes arrive, including heartbeats. */ + onActivity?: () => void; + /** Extra headers forwarded upstream (e.g. the owner's account tag). */ + headers?: Record; + /** 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; +}; + +/** + * Attach to a detached run's event stream via + * `GET /optimize/runs/{run_id}/events`. + * + * 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. 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, + fetchImpl = fetch, + headers, + maxEventBytes, + onActivity, + requestId, + signal, +}: AttachPetrinautOptimizationRunStreamOptions): Promise => { + const url = petrinautOptimizerUrl( + endpoint, + `optimize/runs/${encodeURIComponent(runId)}/events`, + ); + if (cursor !== undefined) { + url.searchParams.set("cursor", String(cursor)); + } + + const response = await fetchImpl(url, { + method: "GET", + headers: { + accept: "text/event-stream", + ...(requestId === undefined ? {} : { "x-hash-request-id": requestId }), + ...headers, + }, + signal: signal as AbortSignal | undefined, + }); + if (!response.ok) { + throw await petrinautOptimizerHttpErrorFromResponse(response); + } + if (!response.body) { + throw new Error("Petrinaut optimizer returned an empty response"); + } + + // Recorded by the optimizer at creation; it only shapes the synthesized + // `complete` summary event — the completed/pruned/failed counts in that + // summary reflect the frames observed by *this* attachment (everything + // past the cursor), not necessarily the whole study. Clamped to 1 when the + // header is missing (an older optimizer): a cosmetic undercount must not + // fail canonical-event validation and kill the attachment. + const requestedTrials = Math.max( + 1, + Number(response.headers.get("x-requested-trials")) || 0, + ); + + return { + events: decodePetrinautOptimizerStream(response.body, { + requestedTrials, + ...(maxEventBytes === undefined ? {} : { maxEventBytes }), + ...(onActivity ? { onActivity } : {}), + }), + optimizationRunId: response.headers.get("x-optimization-run-id"), + }; +}; diff --git a/libs/@local/petrinaut-optimizer-client/src/client.test.ts b/libs/@local/petrinaut-optimizer-client/src/client.test.ts new file mode 100644 index 00000000000..75ddf6092b9 --- /dev/null +++ b/libs/@local/petrinaut-optimizer-client/src/client.test.ts @@ -0,0 +1,51 @@ +import { describe, expect, it, vi } from "vitest"; + +import { createPetrinautOptimizerClient } from "./client.js"; + +describe("createPetrinautOptimizerClient", () => { + it("decomposes requests into the injected (url, init) fetch shape", async () => { + const fetchImpl = vi.fn(async (_url: string | URL, _init?: RequestInit) => + Promise.resolve(Response.json({ run_id: "run-1" }, { status: 201 })), + ); + const client = createPetrinautOptimizerClient( + "http://petrinaut-opt.test", + fetchImpl, + ); + + const { data, response } = await client.POST("/optimize/runs", { + body: { name: "study" }, + headers: { + "x-hash-account-id": "user-1", + "x-hash-request-id": "request-1", + }, + }); + + expect(response.status).toBe(201); + expect(data?.run_id).toBe("run-1"); + const [url, init] = fetchImpl.mock.calls[0]!; + expect(url).toBe("http://petrinaut-opt.test/optimize/runs"); + expect(init?.method).toBe("POST"); + expect(JSON.parse(init?.body as string)).toEqual({ name: "study" }); + const headers = new Headers(init?.headers); + expect(headers.get("x-hash-account-id")).toBe("user-1"); + expect(headers.get("x-hash-request-id")).toBe("request-1"); + }); + + it("keeps an endpoint's path prefix when building URLs", async () => { + const fetchImpl = vi.fn(async (_url: string | URL, _init?: RequestInit) => + Promise.resolve(new Response(null, { status: 204 })), + ); + const client = createPetrinautOptimizerClient( + "http://localhost:5173/api/petrinaut-opt", + fetchImpl, + ); + + await client.DELETE("/optimize/runs/{run_id}", { + params: { path: { run_id: "run 1" } }, + }); + + expect(fetchImpl.mock.calls[0]![0]).toBe( + "http://localhost:5173/api/petrinaut-opt/optimize/runs/run%201", + ); + }); +}); diff --git a/libs/@local/petrinaut-optimizer-client/src/client.ts b/libs/@local/petrinaut-optimizer-client/src/client.ts new file mode 100644 index 00000000000..96fffd41b4b --- /dev/null +++ b/libs/@local/petrinaut-optimizer-client/src/client.ts @@ -0,0 +1,43 @@ +import createOpenApiClient from "openapi-fetch"; + +import type { paths } from "./openapi.gen.js"; +import type { PetrinautOptimizerFetch } from "./optimizer-http.js"; + +/** + * Typed Petrinaut Optimizer client generated from its OpenAPI schema. + * + * Covers the plain JSON endpoints (create, cancel, status); the SSE event + * stream keeps its handwritten adapter in `attach-optimization-run.ts` + * because the frame protocol is not expressible in OpenAPI. Non-2xx + * responses come back as `{ error, response }` — use + * `petrinautOptimizerHttpErrorFromResponse(response)` for throw-style + * handling with `Retry-After`/`X-Optimization-Run-ID` semantics. + */ +export type PetrinautOptimizerClient = ReturnType< + typeof createPetrinautOptimizerClient +>; + +/** Create a typed client for one optimizer endpoint (path prefixes kept). */ +export const createPetrinautOptimizerClient = ( + endpoint: string | URL, + fetchImpl?: PetrinautOptimizerFetch, +) => + createOpenApiClient({ + baseUrl: String(endpoint), + ...(fetchImpl + ? { + // openapi-fetch hands over one assembled `Request`; decompose it + // back into the `(url, init)` shape the injected implementations + // (and their test fakes) speak. + fetch: async (request: Request) => + fetchImpl(request.url, { + method: request.method, + headers: request.headers, + ...(request.method === "GET" || request.method === "HEAD" + ? {} + : { body: await request.text() }), + signal: request.signal, + }), + } + : {}), + }); 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 2bb6c85c8eb..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,14 +20,12 @@ const streamChunks = (...chunks: string[]) => { const collect = async ( stream: ReadableStream, options: { - direction?: "maximize" | "minimize"; maxEventBytes?: number; onActivity?: () => void; } = {}, ) => { const events = []; for await (const event of decodePetrinautOptimizerStream(stream, { - direction: options.direction ?? "maximize", requestedTrials: 2, ...(options.maxEventBytes === undefined ? {} @@ -51,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", @@ -66,7 +66,7 @@ describe("decodePetrinautOptimizerStream", () => { parameters: { workers: 3 }, objective: null, state: "pruned", - best: { trial: 0, parameters: { workers: 2 }, objective: 10 }, + best: null, }, { type: "complete", @@ -74,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", @@ -118,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", @@ -169,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?.(); @@ -191,4 +179,91 @@ describe("decodePetrinautOptimizerStream", () => { collect(streamChunks(`data: ${"x".repeat(20)}`), { maxEventBytes: 8 }), ).rejects.toThrow("oversized event"); }); + + it("surfaces SSE frame ids as canonical sequence numbers", async () => { + const events = await collect( + streamChunks( + 'id: 1\ndata: {"step":0,"params":{"workers":2},"metric":10,"state":"COMPLETE"}\n\n', + "id: 2\nevent: done\ndata: {}\n\n", + ), + ); + + expect(events).toEqual([ + { + type: "trial", + trial: 0, + parameters: { workers: 2 }, + objective: 10, + state: "complete", + best: null, + seq: 1, + }, + { + type: "complete", + requestedTrials: 2, + completedTrials: 1, + prunedTrials: 0, + failedTrials: 0, + best: null, + seq: 2, + }, + ]); + }); + + it("rejects non-decimal SSE frame ids", async () => { + await expect( + collect(streamChunks("id: not-a-number\nevent: done\ndata: {}\n\n")), + ).rejects.toThrow("invalid event id"); + // An explicit empty id line must not silently become NaN or 0. + await expect( + collect(streamChunks("id:\nevent: done\ndata: {}\n\n")), + ).rejects.toThrow("invalid event id"); + // `Number()` would coerce exponent notation; the protocol never uses it. + await expect( + collect(streamChunks("id: 1e2\nevent: done\ndata: {}\n\n")), + ).rejects.toThrow("invalid event id"); + }); + + it("adapts a cancelled frame to a terminal cancellation error", async () => { + const events = await collect( + streamChunks( + 'id: 1\ndata: {"step":0,"params":{"workers":2},"metric":10,"state":"COMPLETE"}\n\n', + "id: 2\nevent: cancelled\ndata: {}\n\n", + ), + ); + + expect(events.at(-1)).toEqual({ + type: "error", + code: "optimization_cancelled", + message: "The optimization was cancelled", + retryable: false, + seq: 2, + }); + }); + + it("adapts a superseded frame to a terminal, non-retryable attachment error", async () => { + const events = await collect( + streamChunks("event: superseded\ndata: {}\n\n"), + ); + + expect(events).toEqual([ + { + type: "error", + code: "attachment_superseded", + message: "Another consumer attached to this optimization run", + retryable: false, + }, + ]); + }); + + it("treats a cancelled frame as terminal", async () => { + await expect( + collect( + streamChunks( + "id: 1\nevent: cancelled\ndata: {}\n\n" + + 'id: 2\ndata: {"step":0,"params":{},"metric":1,"state":"COMPLETE"}\n\n', + ), + ), + ).rejects.toThrow("after a terminal event"); + }); }); 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 fd08d8aa1a7..b8910e8ac8d 100644 --- a/libs/@local/petrinaut-optimizer-client/src/decode-optimization-stream.ts +++ b/libs/@local/petrinaut-optimizer-client/src/decode-optimization-stream.ts @@ -3,7 +3,6 @@ import { createParser } from "eventsource-parser"; import { petrinautOptimizationEventSchema, type PetrinautOptimizationEvent, - type PetrinautOptimizationInput, } from "@hashintel/petrinaut-core"; import type { EventSourceMessage, ParseError } from "eventsource-parser"; @@ -12,18 +11,24 @@ type JsonRecord = Record; type StreamState = { requestedTrials: number; - direction: PetrinautOptimizationInput["objective"]["direction"]; completedTrials: number; prunedTrials: number; failedTrials: number; - best: Extract["best"]; terminal: boolean; }; -/** Configuration needed to adapt one upstream optimization stream. */ +/** + * Configuration needed to adapt one upstream optimization stream. + * + * 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. */ - direction: PetrinautOptimizationInput["objective"]["direction"]; /** Number of trials requested by the optimization manifest. */ requestedTrials: number; /** Optional UTF-8 byte limit applied to each complete upstream event. */ @@ -51,6 +56,26 @@ const parseJson = (data: string): unknown => { const utf8ByteLength = (value: string): number => textEncoder.encode(value).byteLength; +/** + * 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; 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 + * notation, and signs, and ids past 15 digits would lose integer precision. + */ +const parseEventSequence = (id: string | undefined): number | undefined => { + if (id === undefined) { + return undefined; + } + if (!/^\d{1,15}$/.test(id)) { + throw new Error("Petrinaut optimizer returned an invalid event id"); + } + return Number(id); +}; + /** Validate the flat parameter values returned for one Optuna trial. */ const parseParameters = (value: unknown): Record => { if (!isJsonRecord(value)) { @@ -75,7 +100,7 @@ const parseTrial = ( value: unknown, ): Omit< Extract, - "type" | "best" + "type" | "best" | "seq" > => { if (!isJsonRecord(value)) { throw new Error("Petrinaut optimizer returned an invalid SSE event"); @@ -134,6 +159,8 @@ const adaptSseEvent = ( throw new Error("Petrinaut optimizer returned an oversized event"); } + const sequence = parseEventSequence(event.id); + const sequenceField = sequence === undefined ? {} : { seq: sequence }; const value = parseJson(event.data); if (event.event === "error") { return { @@ -145,6 +172,38 @@ const adaptSseEvent = ( ? value.message : "Petrinaut optimizer reported an error", retryable: false, + ...sequenceField, + }), + state: { ...state, terminal: true }, + }; + } + if (event.event === "superseded") { + // A newer attachment took over this run's stream. Terminal for THIS + // attachment only — the run lives on under the newer consumer — so the + // consumer must not treat it as a dropped connection and reconnect + // (that would supersede the newer attachment right back, forever). + return { + event: petrinautOptimizationEventSchema.parse({ + type: "error", + code: "attachment_superseded", + message: "Another consumer attached to this optimization run", + retryable: false, + ...sequenceField, + }), + state: { ...state, terminal: true }, + }; + } + if (event.event === "cancelled") { + // A detached run's terminal cancellation frame (client DELETE, orphan + // reaping, or optimizer shutdown). It is terminal and not retryable: the + // run is gone, so re-attaching cannot resume it. + return { + event: petrinautOptimizationEventSchema.parse({ + type: "error", + code: "optimization_cancelled", + message: "The optimization was cancelled", + retryable: false, + ...sequenceField, }), state: { ...state, terminal: true }, }; @@ -157,7 +216,8 @@ const adaptSseEvent = ( completedTrials: state.completedTrials, prunedTrials: state.prunedTrials, failedTrials: state.failedTrials, - best: state.best, + best: null, + ...sequenceField, }), state: { ...state, terminal: true }, }; @@ -176,39 +236,27 @@ const adaptSseEvent = ( ? value.message : "Petrinaut optimizer reported an error", retryable: false, + ...sequenceField, }), state: { ...state, terminal: true }, }; } const trial = parseTrial(value); - const best = - 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, }; @@ -219,6 +267,8 @@ const adaptSseEvent = ( * * The decoder is browser-safe so NodeAPI and local browser integrations use * exactly the same protocol validation, aggregation, and terminal semantics. + * Frames stamped with an SSE `id:` line surface it as the canonical `seq` + * field so consumers of detached runs can resume from a cursor. */ export async function* decodePetrinautOptimizerStream( stream: ReadableStream, @@ -226,11 +276,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[] = []; @@ -286,11 +334,6 @@ export async function* decodePetrinautOptimizerStream( }; try { - 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 9dba985373c..fea8d04224d 100644 --- a/libs/@local/petrinaut-optimizer-client/src/index.ts +++ b/libs/@local/petrinaut-optimizer-client/src/index.ts @@ -1,12 +1,17 @@ +export { attachPetrinautOptimizationRunStream } 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, PetrinautOptimizerHttpError, -} from "./open-optimization-stream.js"; -export type { - OpenPetrinautOptimizationStreamOptions, - PetrinautOptimizationStreamHandle, - PetrinautOptimizerFetch, -} from "./open-optimization-stream.js"; + petrinautOptimizerHttpErrorFromResponse, +} from "./optimizer-http.js"; +export type { PetrinautOptimizerFetch } from "./optimizer-http.js"; export type { components, operations, paths, webhooks } from "./openapi.gen.js"; 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 73fcffd91c6..00000000000 --- a/libs/@local/petrinaut-optimizer-client/src/open-optimization-stream.test.ts +++ /dev/null @@ -1,210 +0,0 @@ -import { describe, expect, it, vi } from "vitest"; - -import { - openPetrinautOptimizationStream, - PetrinautOptimizerHttpError, -} from "./open-optimization-stream.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 485827225ed..00000000000 --- a/libs/@local/petrinaut-optimizer-client/src/open-optimization-stream.ts +++ /dev/null @@ -1,125 +0,0 @@ -import { decodePetrinautOptimizerStream } from "./decode-optimization-stream.js"; - -import type { - AbortSignalLike, - PetrinautOptimizationEvent, - PetrinautOptimizationInput, -} from "@hashintel/petrinaut-core"; - -type JsonRecord = Record; - -/** Fetch-compatible function used to call Petrinaut Optimizer. */ -export type PetrinautOptimizerFetch = ( - input: string | URL, - init?: RequestInit, -) => Promise; - -/** Error returned when Petrinaut Optimizer rejects an HTTP request. */ -export class PetrinautOptimizerHttpError extends Error { - /** Create an optimizer HTTP error while retaining transport metadata. */ - constructor( - message: string, - readonly status: number, - readonly retryAfter: string | null, - readonly optimizationRunId: string | null = null, - ) { - super(message); - this.name = "PetrinautOptimizerHttpError"; - } -} - -/** 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; -}; - -/** 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); - -/** Read the most useful safe message from a failed optimizer response. */ -const responseErrorMessage = async (response: Response): Promise => { - try { - const payload: unknown = await response.json(); - if (isJsonRecord(payload)) { - if (typeof payload.detail === "string") { - return payload.detail; - } - if (typeof payload.message === "string") { - return payload.message; - } - } - } catch { - // Fall back to the status when the service did not return JSON. - } - return `Petrinaut optimizer returned status ${response.status}`; -}; - -/** - * 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 new PetrinautOptimizerHttpError( - await responseErrorMessage(response), - response.status, - response.headers.get("retry-after"), - response.headers.get("x-optimization-run-id"), - ); - } - 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/openapi.gen.ts b/libs/@local/petrinaut-optimizer-client/src/openapi.gen.ts index 18aec27ebb0..2af2b923bd2 100644 --- a/libs/@local/petrinaut-optimizer-client/src/openapi.gen.ts +++ b/libs/@local/petrinaut-optimizer-client/src/openapi.gen.ts @@ -24,7 +24,7 @@ export interface paths { patch?: never; trace?: never; }; - "/optimize/all": { + "/optimize/runs": { parameters: { query?: never; header?: never; @@ -34,17 +34,21 @@ export interface paths { get?: never; put?: never; /** - * Post Optimize All - * @description Stream one SSE data frame for every completed Optuna trial. + * Post Optimize Runs + * @description Start a detached run that remains available for later SSE attachment. + * + * When the caller stamps ``x-hash-account-id`` (the authenticated proxy + * does), the run is owned: the account is single-flight while it lives, and + * only requests carrying the same tag may attach to or cancel it. */ - post: operations["post_optimize_all_optimize_all_post"]; + post: operations["post_optimize_runs_optimize_runs_post"]; delete?: never; options?: never; head?: never; patch?: never; trace?: never; }; - "/optimize/best": { + "/optimize/runs/{run_id}": { parameters: { query?: never; header?: never; @@ -53,11 +57,40 @@ export interface paths { }; get?: never; put?: never; + post?: never; /** - * Post Optimize Best - * @description Stream the best-so-far SSE data frame after every completed trial. + * Delete Optimize Run + * @description Cancel a detached run and wait for its resources to be released. */ - post: operations["post_optimize_best_optimize_best_post"]; + delete: operations["delete_optimize_run_optimize_runs__run_id__delete"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/optimize/runs/{run_id}/events": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Optimize Run Events + * @description Attach to a detached run and replay events after the supplied cursor. + * + * The route is excluded from ASGI auto-instrumentation (see + * ``telemetry.py``): its response live-tails until the consumer detaches, + * and a SERVER span that long would read as worst-case latency in the RED + * SLIs and only export on disconnect. A short manual SERVER span covers + * just the attach itself — run lookup, cursor resolution, and attachment + * registration — and ends when the tail starts, so the latency SLI + * measures "was attaching fast" while the client→optimizer service-graph + * edge is preserved. + */ + get: operations["get_optimize_run_events_optimize_runs__run_id__events_get"]; + put?: never; + post?: never; delete?: never; options?: never; head?: never; @@ -113,6 +146,14 @@ export interface components { /** Detail */ detail?: components["schemas"]["ValidationError"][]; }; + /** + * OptimizationRunCreated + * @description Response body of a successful detached-run creation. + */ + OptimizationRunCreated: { + /** Run Id */ + run_id: string; + }; /** * Phase * @enum {string} @@ -173,7 +214,7 @@ export interface operations { }; }; }; - post_optimize_all_optimize_all_post: { + post_optimize_runs_optimize_runs_post: { parameters: { query?: never; header?: never; @@ -188,13 +229,13 @@ export interface operations { }; }; responses: { - /** @description Server-Sent Events optimization stream */ - 200: { + /** @description A detached optimization run was started */ + 201: { headers: { [name: string]: unknown; }; content: { - "text/event-stream": string; + "application/json": components["schemas"]["OptimizationRunCreated"]; }; }; /** @description The optimization manifest exceeds 8 MiB */ @@ -231,32 +272,26 @@ export interface operations { }; }; }; - post_optimize_best_optimize_best_post: { + delete_optimize_run_optimize_runs__run_id__delete: { parameters: { query?: never; header?: never; - path?: never; - cookie?: never; - }; - requestBody: { - content: { - "application/json": { - [key: string]: unknown; - }; + path: { + run_id: string; }; + cookie?: never; }; + requestBody?: never; responses: { - /** @description Server-Sent Events optimization stream */ - 200: { + /** @description The run was cancelled (or had already reached a terminal state); when this call stopped it, the event log's terminal frame is `event: cancelled` */ + 204: { headers: { [name: string]: unknown; }; - content: { - "text/event-stream": string; - }; + content?: never; }; - /** @description The optimization manifest exceeds 8 MiB */ - 413: { + /** @description No optimization run with this id is registered */ + 404: { headers: { [name: string]: unknown; }; @@ -271,22 +306,48 @@ export interface operations { "application/json": components["schemas"]["HTTPValidationError"]; }; }; - /** @description The service is already at its study limit */ - 429: { + }; + }; + get_optimize_run_events_optimize_runs__run_id__events_get: { + parameters: { + query?: { + cursor?: number | null; + }; + header?: never; + path: { + run_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Server-Sent Events attachment to a detached optimization run. Every frame carries an `id: ` line (seq starts at 1); buffered frames with seq > cursor are replayed, then new frames are live-tailed with `: heartbeat` comments roughly every 30 seconds. The terminal frame is `event: done` (completed), an ERROR data frame (study failure), or `event: cancelled` (cancelled or reaped). If the run is already terminal the response closes after the replay. Disconnecting does not affect the run; a newer attachment supersedes this one. */ + 200: { headers: { - /** @description Seconds to wait before retrying the study */ - "Retry-After"?: string; + /** @description Trial count requested by the run's manifest, for sizing synthesized summaries */ + "X-Requested-Trials"?: string; [name: string]: unknown; }; - content?: never; + content: { + "text/event-stream": string; + }; }; - /** @description The CLI or optimization study could not initialize */ - 500: { + /** @description No optimization run with this id is registered */ + 404: { headers: { [name: string]: unknown; }; content?: never; }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; }; }; get_status_status_get: { diff --git a/libs/@local/petrinaut-optimizer-client/src/optimizer-http.ts b/libs/@local/petrinaut-optimizer-client/src/optimizer-http.ts new file mode 100644 index 00000000000..e2b0953f6ca --- /dev/null +++ b/libs/@local/petrinaut-optimizer-client/src/optimizer-http.ts @@ -0,0 +1,70 @@ +type JsonRecord = Record; + +/** Fetch-compatible function used to call Petrinaut Optimizer. */ +export type PetrinautOptimizerFetch = ( + input: string | URL, + init?: RequestInit, +) => Promise; + +/** Error returned when Petrinaut Optimizer rejects an HTTP request. */ +export class PetrinautOptimizerHttpError extends Error { + /** Create an optimizer HTTP error while retaining transport metadata. */ + constructor( + message: string, + readonly status: number, + readonly retryAfter: string | null, + readonly optimizationRunId: string | null = null, + ) { + super(message); + this.name = "PetrinautOptimizerHttpError"; + } +} + +/** + * 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); + +/** Read the most useful safe message from a failed optimizer response. */ +const responseErrorMessage = async (response: Response): Promise => { + try { + const payload: unknown = await response.json(); + if (isJsonRecord(payload)) { + if (typeof payload.detail === "string") { + return payload.detail; + } + if (typeof payload.message === "string") { + return payload.message; + } + } + } catch { + // Fall back to the status when the service did not return JSON. + } + return `Petrinaut optimizer returned status ${response.status}`; +}; + +/** Build the canonical error for a non-ok optimizer response. */ +export const petrinautOptimizerHttpErrorFromResponse = async ( + response: Response, +): Promise => + new PetrinautOptimizerHttpError( + await responseErrorMessage(response), + response.status, + response.headers.get("retry-after"), + response.headers.get("x-optimization-run-id"), + ); diff --git a/yarn.lock b/yarn.lock index 7bcac4a4d2e..52119493df6 100644 --- a/yarn.lock +++ b/yarn.lock @@ -8707,6 +8707,7 @@ __metadata: "@hashintel/petrinaut-core": "workspace:*" "@local/tsconfig": "workspace:*" eventsource-parser: "npm:3.0.8" + openapi-fetch: "npm:0.17.0" openapi-typescript: "npm:7.13.0" oxfmt: "npm:0.50.0" rimraf: "npm:6.1.3" @@ -34262,6 +34263,15 @@ __metadata: languageName: node linkType: hard +"openapi-fetch@npm:0.17.0": + version: 0.17.0 + resolution: "openapi-fetch@npm:0.17.0" + dependencies: + openapi-typescript-helpers: "npm:^0.1.0" + checksum: 10c0/7e981d28c7839469662a73d44f3a85e05c1a5842a8afd5d6e36c75cbf085aa7f705862871cc7ce07c04605ddc936af3d5850c2cc66d193ee1fea07e7135da192 + languageName: node + linkType: hard + "openapi-fetch@npm:^0.14.1": version: 0.14.1 resolution: "openapi-fetch@npm:0.14.1" @@ -34278,6 +34288,13 @@ __metadata: languageName: node linkType: hard +"openapi-typescript-helpers@npm:^0.1.0": + version: 0.1.0 + resolution: "openapi-typescript-helpers@npm:0.1.0" + checksum: 10c0/3a148cf7d6a7f51124966a0fef64c2fbcf1a5a0a1e6320dd627f0733787ac4259877be43f7a7a2910684e3d5d238c6a9655c6f9e1f3adaa9a2a61477e3b8481f + languageName: node + linkType: hard + "openapi-typescript@npm:7.13.0": version: 7.13.0 resolution: "openapi-typescript@npm:7.13.0"