From f81856666c28bb1b8ea17c229875e816665bd0ab Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 31 Aug 2026 13:44:52 +0800 Subject: [PATCH] fix: refresh Knowhere JWTs during long QStash workflows Expired Dashboard snapshots were reported as workflow auth failures even when Knowhere parsing later succeeded. Co-authored-by: Cursor --- .../parsed-sync-route-workflow.test.ts | 68 ++++++- .../sources/parsed-sync-route-workflow.ts | 29 ++- .../source-reconcile-route-workflow.test.ts | 99 +++++++++ .../source-reconcile-route-workflow.ts | 47 +++-- .../dashboard/api-key-service.test.ts | 191 ++++++++++++++++++ src/integrations/dashboard/api-key-service.ts | 159 ++++++++++++++- src/integrations/dashboard/orpc-request.ts | 11 + 7 files changed, 576 insertions(+), 28 deletions(-) diff --git a/src/domains/sources/parsed-sync-route-workflow.test.ts b/src/domains/sources/parsed-sync-route-workflow.test.ts index 670ee67..9a2bcff 100644 --- a/src/domains/sources/parsed-sync-route-workflow.test.ts +++ b/src/domains/sources/parsed-sync-route-workflow.test.ts @@ -7,12 +7,24 @@ const mocks = vi.hoisted(() => ({ updateSyncStatus: vi.fn(), loggerInfo: vi.fn(), loggerError: vi.fn(), + ensureFreshKnowhereApiKey: vi.fn(async (apiKey: string) => apiKey), + withFreshKnowhereApiKey: vi.fn( + async (apiKey: string, run: (apiKey: string) => Promise) => ({ + result: await run(apiKey), + apiKey, + }), + ), })) vi.mock("@/integrations/knowhere", () => ({ makeKnowhereClientWithParsedStorage: mocks.makeKnowhereClientWithParsedStorage, })) +vi.mock("@/integrations/dashboard/api-key-service", () => ({ + ensureFreshKnowhereApiKey: mocks.ensureFreshKnowhereApiKey, + withFreshKnowhereApiKey: mocks.withFreshKnowhereApiKey, +})) + vi.mock("./workflow-runtime", () => ({ sourceWorkflowRuntime: { updateSyncStatus: mocks.updateSyncStatus, @@ -62,6 +74,13 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { }, }) mocks.releaseSyncCapacity.mockResolvedValue(undefined) + mocks.ensureFreshKnowhereApiKey.mockImplementation(async (apiKey: string) => apiKey) + mocks.withFreshKnowhereApiKey.mockImplementation( + async (apiKey: string, run: (apiKey: string) => Promise) => ({ + result: await run(apiKey), + apiKey, + }), + ) }) afterEach(() => { @@ -107,12 +126,17 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { client: {}, knowledge: { syncParsedDocument }, }) - const triggered: Array<{ workflowRunId: string; segmentIndex?: number }> = [] + const triggered: Array<{ + workflowRunId: string + segmentIndex?: number + apiKey?: string + }> = [] const restore = parsedSyncRouteWorkflow.setContinuationTriggerForTesting( async (input) => { triggered.push({ workflowRunId: input.workflowRunId, segmentIndex: input.payload.segmentIndex, + apiKey: input.payload.apiKey, }) }, ) @@ -129,6 +153,7 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { expect(triggered).toHaveLength(1) expect(triggered[0]?.segmentIndex).toBe(1) expect(triggered[0]?.workflowRunId).toBe("doc_1-sync-rev_1-1") + expect(triggered[0]?.apiKey).toBe("key_1") expect(mocks.updateSyncStatus).toHaveBeenLastCalledWith( "workspace_1", "source_1", @@ -186,6 +211,7 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { readonly workflowRunId: string readonly segmentIndex?: number readonly delaySeconds?: number + readonly apiKey?: string }> = [] const restore = parsedSyncRouteWorkflow.setContinuationTriggerForTesting( async (input) => { @@ -193,6 +219,7 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { workflowRunId: input.workflowRunId, segmentIndex: input.payload.segmentIndex, delaySeconds: input.delaySeconds, + apiKey: input.payload.apiKey, }) }, ) @@ -221,11 +248,50 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => { workflowRunId: "doc_1-sync-rev_1-1", segmentIndex: 1, delaySeconds: 60, + apiKey: "key_1", }, ]) expect(mocks.releaseSyncCapacity).not.toHaveBeenCalled() }) + it("forwards a refreshed Knowhere JWT on sync continuation and capacity retry", async () => { + mocks.ensureFreshKnowhereApiKey.mockResolvedValue("jwt_refreshed") + mocks.withFreshKnowhereApiKey.mockImplementation( + async (_apiKey: string, run: (apiKey: string) => Promise) => ({ + result: await run("jwt_refreshed"), + apiKey: "jwt_refreshed", + }), + ) + const syncParsedDocument = vi.fn(async () => ({ + documentId: "doc_1", + revisionKey: "rev_1", + completed: false, + })) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: {}, + knowledge: { syncParsedDocument }, + }) + const triggered: Array<{ apiKey?: string }> = [] + const restore = parsedSyncRouteWorkflow.setContinuationTriggerForTesting( + async (input) => { + triggered.push({ apiKey: input.payload.apiKey }) + }, + ) + + try { + await parsedSyncRouteWorkflow.runParsedSyncWorkflow({ + context: createContext(), + payload: basePayload, + }) + } finally { + restore() + } + + expect(mocks.ensureFreshKnowhereApiKey).toHaveBeenCalledWith("key_1") + expect(syncParsedDocument).toHaveBeenCalled() + expect(triggered[0]?.apiKey).toBe("jwt_refreshed") + }) + it("does not release capacity when Upstash aborts during a planned step", async () => { const workflowAbort = new Error("planned workflow step") workflowAbort.name = "WorkflowAbort" diff --git a/src/domains/sources/parsed-sync-route-workflow.ts b/src/domains/sources/parsed-sync-route-workflow.ts index e832b01..6634fb8 100644 --- a/src/domains/sources/parsed-sync-route-workflow.ts +++ b/src/domains/sources/parsed-sync-route-workflow.ts @@ -4,6 +4,10 @@ import { Client, WorkflowAbort, type WorkflowContext } from "@upstash/workflow" import type { KnowledgeSyncParsedDocumentResponse } from "@ontos-ai/knowhere-sdk" import { makeKnowhereClientWithParsedStorage } from "@/integrations/knowhere" +import { + ensureFreshKnowhereApiKey, + withFreshKnowhereApiKey, +} from "@/integrations/dashboard/api-key-service" import { logger } from "@/lib/logger" import { getParsedSyncWorkflowRunId, @@ -66,10 +70,11 @@ async function runParsedSyncWorkflow(input: { readonly payload: NormalizedParsedSyncPayload }): Promise { const { context, payload } = input - const { workspaceId, sourceId, documentId, apiKey } = payload - const { knowledge } = makeKnowhereClientWithParsedStorage(apiKey, { - workspaceId, - }) + const { workspaceId, sourceId, documentId } = payload + let apiKey = await context.run( + `refresh-knowhere-jwt-${payload.segmentIndex}`, + async () => ensureFreshKnowhereApiKey(payload.apiKey), + ) let revisionKey = payload.revisionKey let completed = false @@ -131,14 +136,22 @@ async function runParsedSyncWorkflow(input: { try { for (let step = 0; step < maxSyncStepsPerSegment; step++) { - const result: KnowledgeSyncParsedDocumentResponse = await context.run( + const stepResult = await context.run( `sync-${payload.segmentIndex}-${step}`, async () => - knowledge.syncParsedDocument({ - documentId, - ...(revisionKey ? { revisionKey } : {}), + withFreshKnowhereApiKey(apiKey, async (freshKey) => { + const { knowledge } = makeKnowhereClientWithParsedStorage( + freshKey, + { workspaceId }, + ) + return knowledge.syncParsedDocument({ + documentId, + ...(revisionKey ? { revisionKey } : {}), + }) }), ) + apiKey = stepResult.apiKey + const result: KnowledgeSyncParsedDocumentResponse = stepResult.result revisionKey = result.revisionKey await context.run(`record-progress-${payload.segmentIndex}-${step}`, () => diff --git a/src/domains/sources/source-reconcile-route-workflow.test.ts b/src/domains/sources/source-reconcile-route-workflow.test.ts index 730fcf0..1498fed 100644 --- a/src/domains/sources/source-reconcile-route-workflow.test.ts +++ b/src/domains/sources/source-reconcile-route-workflow.test.ts @@ -13,6 +13,12 @@ const mocks = vi.hoisted(() => ({ makeKnowhereClientWithParsedStorage: vi.fn(), markSourceReadyAfterReconciliation: vi.fn(), pollSourceReconciliation: vi.fn(), + withFreshKnowhereApiKey: vi.fn( + async (apiKey: string, run: (apiKey: string) => Promise) => ({ + result: await run(apiKey), + apiKey, + }), + ), })) vi.mock("@/domains/sources/source-reconcile-workflow", () => ({ @@ -44,6 +50,10 @@ vi.mock("@/integrations/knowhere", () => ({ mocks.makeKnowhereClientWithParsedStorage, })) +vi.mock("@/integrations/dashboard/api-key-service", () => ({ + withFreshKnowhereApiKey: mocks.withFreshKnowhereApiKey, +})) + vi.mock("@/lib/logger", () => ({ logger: { error: mocks.loggerError, @@ -97,6 +107,12 @@ describe("sourceReconcileRouteWorkflow", () => { status: "ready", }) mocks.updateRevisionKey.mockResolvedValue({ id: "source_1" }) + mocks.withFreshKnowhereApiKey.mockImplementation( + async (apiKey: string, run: (apiKey: string) => Promise) => ({ + result: await run(apiKey), + apiKey, + }), + ) }) afterEach(() => { @@ -273,6 +289,89 @@ describe("sourceReconcileRouteWorkflow", () => { ]) }) + it("forwards a refreshed Knowhere JWT on poll continuation and parsed-sync enqueue", async () => { + const context = createWorkflowContext() + const continuations: ContinuationTriggerInput[] = [] + const restore = + sourceReconcileRouteWorkflow.setContinuationTriggerForTesting( + async (input) => { + continuations.push(input) + }, + ) + mocks.withFreshKnowhereApiKey.mockImplementation( + async (_apiKey: string, run: (apiKey: string) => Promise) => ({ + result: await run("jwt_refreshed"), + apiKey: "jwt_refreshed", + }), + ) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: { jobs: {}, documents: { listChunks: vi.fn() } }, + knowledge: { syncParsedDocument: vi.fn() }, + }) + mocks.pollSourceReconciliation.mockResolvedValue({ + kind: "waiting", + jobId: "job_1", + jobStatus: "running", + }) + + try { + await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context, + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_expired", + segmentIndex: 0, + }), + }) + } finally { + restore() + } + + expect(mocks.makeKnowhereClientWithParsedStorage).toHaveBeenCalledWith( + "jwt_refreshed", + { workspaceId: "workspace_1" }, + ) + expect(continuations[0]?.payload.apiKey).toBe("jwt_refreshed") + }) + + it("enqueues parsed-sync with a refreshed Knowhere JWT", async () => { + const context = createWorkflowContext() + mocks.withFreshKnowhereApiKey.mockImplementation( + async (_apiKey: string, run: (apiKey: string) => Promise) => ({ + result: await run("jwt_refreshed"), + apiKey: "jwt_refreshed", + }), + ) + const wired = createClient({}) + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: wired.client, + knowledge: wired.knowledge, + }) + mocks.pollSourceReconciliation.mockResolvedValue({ + kind: "ready-to-prepare", + jobId: "job_1", + documentId: "doc_1", + }) + + await sourceReconcileRouteWorkflow.runPollAndMirrorWorkflow({ + context, + payload: sourceReconcileRouteWorkflow.normalizeReconcilePayload({ + workspaceId: "workspace_1", + sourceId: "source_1", + apiKey: "jwt_expired", + }), + }) + + expect(mocks.enqueueParsedDocumentSync).toHaveBeenCalledWith({ + workspaceId: "workspace_1", + sourceId: "source_1", + documentId: "doc_1", + apiKey: "jwt_refreshed", + revisionKey: "rev_1", + }) + }) + it("marks a parsing source failed after workflow retry exhaustion", async () => { mocks.markFailed.mockResolvedValue({ id: "source_1" }) diff --git a/src/domains/sources/source-reconcile-route-workflow.ts b/src/domains/sources/source-reconcile-route-workflow.ts index 5d2cb2c..a1a316b 100644 --- a/src/domains/sources/source-reconcile-route-workflow.ts +++ b/src/domains/sources/source-reconcile-route-workflow.ts @@ -7,6 +7,7 @@ import { pollSourceReconciliation, } from "@/domains/sources/source-reconcile-workflow" import { makeKnowhereClientWithParsedStorage } from "@/integrations/knowhere" +import { withFreshKnowhereApiKey } from "@/integrations/dashboard/api-key-service" import { logger } from "@/lib/logger" import { enqueueParsedDocumentSync } from "./parsed-document-sync-scheduler" import { sourceWorkflowRuntime } from "./workflow-runtime" @@ -68,10 +69,8 @@ async function runPollAndMirrorWorkflow(input: { readonly payload: NormalizedReconcilePayload }): Promise { const { context, payload } = input - const { workspaceId, sourceId, apiKey } = payload - const { client } = makeKnowhereClientWithParsedStorage(apiKey, { - workspaceId, - }) + const { workspaceId, sourceId } = payload + let apiKey = payload.apiKey let delay = initialDelaySeconds let completedJob: { readonly jobId: string @@ -79,13 +78,20 @@ async function runPollAndMirrorWorkflow(input: { } | null = null for (let attempt = 0; attempt < maxPollAttempts; attempt++) { - const poll = await context.run(`poll-${attempt}`, async () => { - return pollSourceReconciliation({ - workspaceId, - sourceId, - client, - }) - }) + const step = await context.run(`poll-${attempt}`, async () => + withFreshKnowhereApiKey(apiKey, async (freshKey) => { + const { client } = makeKnowhereClientWithParsedStorage(freshKey, { + workspaceId, + }) + return pollSourceReconciliation({ + workspaceId, + sourceId, + client, + }) + }), + ) + apiKey = step.apiKey + const poll = step.result if (poll.kind === "ready-to-prepare") { completedJob = { @@ -146,14 +152,21 @@ async function runPollAndMirrorWorkflow(input: { ) if (ready.status === "gone") return - const revisionKey = await context.run("resolve-revision-key", async () => - resolveParsedRevisionKey({ - client, - sourceId, - documentId: jobToPrepare.documentId, - fallbackRevisionKey: jobToPrepare.jobId, + const revision = await context.run("resolve-revision-key", async () => + withFreshKnowhereApiKey(apiKey, async (freshKey) => { + const { client } = makeKnowhereClientWithParsedStorage(freshKey, { + workspaceId, + }) + return resolveParsedRevisionKey({ + client, + sourceId, + documentId: jobToPrepare.documentId, + fallbackRevisionKey: jobToPrepare.jobId, + }) }), ) + apiKey = revision.apiKey + const revisionKey = revision.result await context.run("record-source-revision-key", async () => sourceWorkflowRuntime.updateRevisionKey(workspaceId, sourceId, revisionKey), diff --git a/src/integrations/dashboard/api-key-service.test.ts b/src/integrations/dashboard/api-key-service.test.ts index 88347e0..74e9dae 100644 --- a/src/integrations/dashboard/api-key-service.test.ts +++ b/src/integrations/dashboard/api-key-service.test.ts @@ -9,8 +9,13 @@ vi.mock("next/cache", () => nextCacheMocks) import { ensureApiKeyForWorkspace, + ensureFreshKnowhereApiKey, fetchKnowhereJwt, isAuthError, + readJwtExpirySeconds, + refreshKnowhereJwt, + shouldRefreshKnowhereJwt, + withFreshKnowhereApiKey, } from "./api-key-service" function getHeaderValue(headers: HeadersInit | undefined, name: string): string | null { @@ -70,6 +75,12 @@ describe("isAuthError", () => { expect(isAuthError({ message: "Invalid API key" })).toBe(true) expect(isAuthError({ message: "Auth error occurred" })).toBe(true) expect(isAuthError({ message: "unauthenticated" })).toBe(true) + expect(isAuthError({ message: "Authentication required" })).toBe(true) + }) + + it("detects auth-related error names", () => { + expect(isAuthError({ name: "AuthenticationError" })).toBe(true) + expect(isAuthError({ name: "UnauthorizedError" })).toBe(true) }) it("detects 401/403 in error message string", () => { @@ -240,3 +251,183 @@ describe("ensureApiKeyForWorkspace", () => { expect(fetchSpy).not.toHaveBeenCalled() }) }) + +const REFRESH_JWT_PATH = "/api/orpc/knowhereServiceJwt/refresh" + +function makeJwt(payload: Record): string { + const header = Buffer.from( + JSON.stringify({ alg: "none", typ: "JWT" }), + ).toString("base64url") + const body = Buffer.from(JSON.stringify(payload)).toString("base64url") + return `${header}.${body}.sig` +} + +describe("refreshKnowhereJwt", () => { + const originalFetch = globalThis.fetch + const originalOrigin = process.env.DASHBOARD_ORIGIN + + afterEach(() => { + globalThis.fetch = originalFetch + if (originalOrigin === undefined) + delete process.env.DASHBOARD_ORIGIN + else process.env.DASHBOARD_ORIGIN = originalOrigin + }) + + it("POSTs the snapshot token to the Dashboard refresh endpoint", async () => { + process.env.DASHBOARD_ORIGIN = "https://dashboard.example" + const expectedUrl = `https://dashboard.example${REFRESH_JWT_PATH}` + const fetchSpy = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + json: { token: "refreshed.jwt.token", expiresInSeconds: 3600 }, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ) + globalThis.fetch = fetchSpy + + const token = await refreshKnowhereJwt("expired.jwt.snapshot") + + expect(token).toBe("refreshed.jwt.token") + expect(fetchSpy).toHaveBeenCalledOnce() + const [url, init] = fetchSpy.mock.calls[0]! + expect(url instanceof URL ? url.href : url).toBe(expectedUrl) + expect((init as RequestInit)?.method).toBe("POST") + expect(getHeaderValue((init as RequestInit)?.headers, "content-type")).toContain( + "application/json", + ) + expect(await readBodyText((init as RequestInit)?.body)).toBe( + JSON.stringify({ json: { token: "expired.jwt.snapshot" } }), + ) + }) + + it("throws on non-2xx from Dashboard", async () => { + process.env.DASHBOARD_ORIGIN = "https://dashboard.example" + globalThis.fetch = vi + .fn() + .mockResolvedValue(new Response("oops", { status: 401 })) + await expect(refreshKnowhereJwt("expired.jwt.snapshot")).rejects.toThrow( + /Dashboard JWT refresh: non-2xx/, + ) + }) +}) + +describe("ensureFreshKnowhereApiKey", () => { + const originalFetch = globalThis.fetch + const originalApiKey = process.env.KNOWHERE_API_KEY + const originalOrigin = process.env.DASHBOARD_ORIGIN + + afterEach(() => { + globalThis.fetch = originalFetch + if (originalApiKey === undefined) delete process.env.KNOWHERE_API_KEY + else process.env.KNOWHERE_API_KEY = originalApiKey + if (originalOrigin === undefined) + delete process.env.DASHBOARD_ORIGIN + else process.env.DASHBOARD_ORIGIN = originalOrigin + }) + + it("does not refresh persistent API keys or the KNOWHERE_API_KEY override", async () => { + process.env.KNOWHERE_API_KEY = "sk_dev_key" + const fetchSpy = vi.fn() + globalThis.fetch = fetchSpy + + await expect(ensureFreshKnowhereApiKey("sk_dev_key", { force: true })).resolves.toBe( + "sk_dev_key", + ) + await expect(ensureFreshKnowhereApiKey("sk_other")).resolves.toBe("sk_other") + expect(fetchSpy).not.toHaveBeenCalled() + }) + + it("refreshes a JWT that is within the safety window", async () => { + process.env.DASHBOARD_ORIGIN = "https://dashboard.example" + const nearExpiry = makeJwt({ exp: Math.floor(Date.now() / 1000) + 5 }) + globalThis.fetch = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + json: { token: "fresh.jwt.token", expiresInSeconds: 3600 }, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ) + + await expect(ensureFreshKnowhereApiKey(nearExpiry)).resolves.toBe( + "fresh.jwt.token", + ) + }) + + it("returns a still-valid JWT without calling Dashboard", async () => { + const fresh = makeJwt({ exp: Math.floor(Date.now() / 1000) + 3_600 }) + const fetchSpy = vi.fn() + globalThis.fetch = fetchSpy + + await expect(ensureFreshKnowhereApiKey(fresh)).resolves.toBe(fresh) + expect(fetchSpy).not.toHaveBeenCalled() + }) +}) + +describe("shouldRefreshKnowhereJwt", () => { + it("is true when exp is inside the safety window or already past", () => { + const now = 1_700_000_000 + expect(shouldRefreshKnowhereJwt(makeJwt({ exp: now + 15 }), now)).toBe(true) + expect(shouldRefreshKnowhereJwt(makeJwt({ exp: now - 1 }), now)).toBe(true) + expect(shouldRefreshKnowhereJwt(makeJwt({ exp: now + 16 }), now)).toBe(false) + expect(shouldRefreshKnowhereJwt("sk_persistent", now)).toBe(false) + }) +}) + +describe("readJwtExpirySeconds", () => { + it("reads exp from a three-part JWT and ignores other secrets", () => { + expect(readJwtExpirySeconds(makeJwt({ exp: 42 }))).toBe(42) + expect(readJwtExpirySeconds("sk_key")).toBeNull() + expect(readJwtExpirySeconds(makeJwt({ sub: "user" }))).toBeNull() + }) +}) + +describe("withFreshKnowhereApiKey", () => { + const originalFetch = globalThis.fetch + const originalOrigin = process.env.DASHBOARD_ORIGIN + + afterEach(() => { + globalThis.fetch = originalFetch + if (originalOrigin === undefined) + delete process.env.DASHBOARD_ORIGIN + else process.env.DASHBOARD_ORIGIN = originalOrigin + }) + + it("retries once after Authentication required, then uses the refreshed token", async () => { + process.env.DASHBOARD_ORIGIN = "https://dashboard.example" + const stillValid = makeJwt({ exp: Math.floor(Date.now() / 1000) + 3_600 }) + globalThis.fetch = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + json: { token: "retry.jwt.token", expiresInSeconds: 3600 }, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ) + const run = vi + .fn<(apiKey: string) => Promise>() + .mockRejectedValueOnce(new Error("Authentication required")) + .mockResolvedValueOnce("ok") + + const result = await withFreshKnowhereApiKey(stillValid, run) + + expect(result).toEqual({ result: "ok", apiKey: "retry.jwt.token" }) + expect(run).toHaveBeenCalledTimes(2) + expect(run).toHaveBeenNthCalledWith(1, stillValid) + expect(run).toHaveBeenNthCalledWith(2, "retry.jwt.token") + }) + + it("does not retry non-auth errors", async () => { + const fresh = makeJwt({ exp: Math.floor(Date.now() / 1000) + 3_600 }) + const fetchSpy = vi.fn() + globalThis.fetch = fetchSpy + const run = vi.fn().mockRejectedValue(new Error("Parser rejected this document.")) + + await expect(withFreshKnowhereApiKey(fresh, run)).rejects.toThrow( + /Parser rejected/, + ) + expect(run).toHaveBeenCalledOnce() + expect(fetchSpy).not.toHaveBeenCalled() + }) +}) diff --git a/src/integrations/dashboard/api-key-service.ts b/src/integrations/dashboard/api-key-service.ts index 743c99d..225b90f 100644 --- a/src/integrations/dashboard/api-key-service.ts +++ b/src/integrations/dashboard/api-key-service.ts @@ -9,7 +9,7 @@ import { } from "@effect/platform" import { logger } from "@/lib/logger" import { knowhereApiKeyOverride } from "@/integrations/knowhere-api-key" -import { setEmptyJsonBody } from "./orpc-request" +import { setEmptyJsonBody, setJsonBody } from "./orpc-request" import { formatUnknownForLog } from "@/lib/format-log-value" /** @@ -23,6 +23,9 @@ const JwtResponse = Schema.Struct({ }), }) +const ISSUE_JWT_PATH = "/api/orpc/users/issueServiceJwt" +const REFRESH_JWT_PATH = "/api/orpc/knowhereServiceJwt/refresh" + /** * Request a short-lived Knowhere JWT from Dashboard's generic issuance * endpoint. The returned token is passed directly to the Knowhere SDK @@ -44,7 +47,7 @@ export const fetchKnowhereJwtEffect = (cookieHeader: string) => } const http = yield* HttpClient.HttpClient - const url = `${origin}/api/orpc/users/issueServiceJwt` + const url = `${origin}${ISSUE_JWT_PATH}` const body = yield* HttpClientRequest.post(url).pipe( HttpClientRequest.setHeader("cookie", cookieHeader), setEmptyJsonBody, @@ -138,6 +141,151 @@ export async function fetchKnowhereJwt( } } +export const refreshKnowhereJwtEffect = (token: string) => + Effect.gen(function* () { + const origin = process.env.DASHBOARD_ORIGIN + if (!origin) { + return yield* Effect.die( + new Error( + "DASHBOARD_ORIGIN must be set. " + + "It should point to the Dashboard origin (see .env.local.example).", + ), + ) + } + + const http = yield* HttpClient.HttpClient + const url = `${origin}${REFRESH_JWT_PATH}` + const body = yield* setJsonBody(HttpClientRequest.post(url), { + json: { token }, + }).pipe( + http.execute, + Effect.flatMap((response) => + Effect.gen(function* () { + const status = response.status + + if (status < 200 || status >= 300) { + const rawText = yield* Effect.either(response.text) + return yield* Effect.die( + new Error( + `Dashboard JWT refresh: non-2xx (status=${status}) body=${Either.getOrElse(rawText, () => "").slice(0, 1000)}`, + ), + ) + } + + const parsed = yield* Effect.either(response.json) + if (Either.isLeft(parsed)) { + return yield* Effect.die( + new Error( + `Dashboard JWT refresh: invalid JSON (status=${status}) error=${String(parsed.left)}`, + ), + ) + } + + const result = Schema.decodeUnknownEither(JwtResponse)(parsed.right) + if (Either.isLeft(result)) { + return yield* Effect.die( + new Error( + `Dashboard JWT refresh: schema mismatch (status=${status}) body=${formatUnknownForLog(parsed.right).slice(0, 1000)}`, + ), + ) + } + + return result.right.json.token + }), + ), + ) + + return body + }) + +export async function refreshKnowhereJwt(token: string): Promise { + const start = Date.now() + try { + const refreshed = await Effect.runPromise( + refreshKnowhereJwtEffect(token).pipe( + Effect.provide(FetchHttpClient.layer), + ), + ) + logger.info("dashboard: POST /api/orpc/knowhereServiceJwt/refresh ok", { + durationMs: Date.now() - start, + }) + return refreshed + } catch (error) { + logger.error("dashboard: POST /api/orpc/knowhereServiceJwt/refresh failed", { + durationMs: Date.now() - start, + error: error instanceof Error ? error.message : String(error), + }) + throw error + } +} + +export type EnsureFreshKnowhereApiKeyOptions = { + readonly force?: boolean +} + +/** + * Keep a QStash-stored Knowhere credential usable across long workflows. + * Dashboard JWTs last one hour; this refreshes when the snapshot is within + * the safety window, already expired, or `force` is set. Persistent API keys + * are returned unchanged. + */ +export async function ensureFreshKnowhereApiKey( + apiKey: string, + options: EnsureFreshKnowhereApiKeyOptions = {}, +): Promise { + const override = knowhereApiKeyOverride.getApiKey() + if (override && apiKey === override) return apiKey + if (!looksLikeJwt(apiKey)) return apiKey + if (!options.force && !shouldRefreshKnowhereJwt(apiKey)) return apiKey + return refreshKnowhereJwt(apiKey) +} + +export async function withFreshKnowhereApiKey( + apiKey: string, + run: (apiKey: string) => Promise, +): Promise<{ readonly result: T; readonly apiKey: string }> { + let current = await ensureFreshKnowhereApiKey(apiKey) + try { + return { result: await run(current), apiKey: current } + } catch (error) { + if (!isAuthError(error)) throw error + current = await ensureFreshKnowhereApiKey(current, { force: true }) + return { result: await run(current), apiKey: current } + } +} + +export function shouldRefreshKnowhereJwt( + apiKey: string, + nowSeconds: number = Math.floor(Date.now() / 1000), +): boolean { + const expiresAt = readJwtExpirySeconds(apiKey) + if (expiresAt === null) return false + return expiresAt - nowSeconds <= jwtExpirationSafetySeconds +} + +export function looksLikeJwt(apiKey: string): boolean { + return apiKey.split(".").length === 3 +} + +const JwtExpiryPayload = Schema.Struct({ + exp: Schema.Number, +}) + +export function readJwtExpirySeconds(apiKey: string): number | null { + const parts = apiKey.split(".") + const payloadSegment = parts[1] + if (parts.length !== 3 || payloadSegment === undefined) return null + try { + const parsed: unknown = JSON.parse( + Buffer.from(payloadSegment, "base64url").toString("utf8"), + ) + const decoded = Schema.decodeUnknownEither(JwtExpiryPayload)(parsed) + return Either.isRight(decoded) ? decoded.right.exp : null + } catch { + return null + } +} + /** * Resolve the credential used for Knowhere SDK calls. Development can * short-circuit Dashboard JWT issuance by setting KNOWHERE_API_KEY. @@ -175,12 +323,19 @@ export function isAuthError(error: unknown): boolean { if ( msg.includes("unauthorized") || msg.includes("unauthenticated") || + msg.includes("authentication required") || msg.includes("forbidden") || msg.includes("invalid api key") || msg.includes("auth error") ) return true } + if (typeof err.name === "string") { + const name = err.name.toLowerCase() + if (name.includes("authentication") || name.includes("unauthorized")) { + return true + } + } return false } diff --git a/src/integrations/dashboard/orpc-request.ts b/src/integrations/dashboard/orpc-request.ts index ea73d57..a6fd5ab 100644 --- a/src/integrations/dashboard/orpc-request.ts +++ b/src/integrations/dashboard/orpc-request.ts @@ -17,3 +17,14 @@ export function setEmptyJsonBody( JSON_CONTENT_TYPE, ) } + +export function setJsonBody( + request: HttpClientRequest.HttpClientRequest, + body: unknown, +): HttpClientRequest.HttpClientRequest { + return HttpClientRequest.bodyText( + request, + JSON.stringify(body), + JSON_CONTENT_TYPE, + ) +}