fix: refresh Knowhere JWTs during long QStash workflows

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<unknown>) => ({
+      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<unknown>) => ({
+        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,13 +211,15 @@ describe("parsedSyncRouteWorkflow.runParsedSyncWorkflow", () => {
       readonly workflowRunId: string
       readonly segmentIndex?: number
       readonly delaySeconds?: number
+      readonly apiKey?: string
     }> = []
     const restore = parsedSyncRouteWorkflow.setContinuationTriggerForTesting(
       async (input) => {
         triggered.push({
           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<unknown>) => ({
+        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"

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<void> {
   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}`, () =>

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<unknown>) => ({
+      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<unknown>) => ({
+        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<unknown>) => ({
+        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<unknown>) => ({
+        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" })
 

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,24 +69,29 @@ async function runPollAndMirrorWorkflow(input: {
   readonly payload: NormalizedReconcilePayload
 }): Promise<void> {
   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
     readonly documentId: string
   } | 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),