Fix provider retry waiting feedback

src/main/agent/deepchat/runtime/deepChatLoopRunner.ts

@@ -1096,6 +1096,8 @@ export class DeepChatLoopRunner {
     }
     let toolSurfaceController: ToolSurfaceRunController | null = null
     let frozenSkillRequirementByName: ReadonlyMap<string, RunSkillToolRequirements> | null = null
+    let removeProviderRetryAbortListener: (() => void) | undefined
+    let clearProviderRetryWaitingMessage: (() => void) | undefined
     try {
       if (toolSurfaceMode !== 'legacy') {
       const universe = await awaitWithAbort(
@@ -1560,6 +1562,7 @@ export class DeepChatLoopRunner {
       const commitTapeProviderView = this.commitTapeProviderView.bind(this)
       const persistMessageTrace = this.persistMessageTrace.bind(this)
       const emitRateLimitWaitingMessage = this.emitRateLimitWaitingMessage.bind(this)
+      const emitProviderRetryWaitingMessage = this.emitProviderRetryWaitingMessage.bind(this)
       const clearRateLimitWaitingMessage = this.clearRateLimitWaitingMessage.bind(this)
       const toolSurfaceAdapterHistory = this.toolSurfaceAdapterHistory
       const hooks = this.ports.hookSink.scope({
@@ -1574,6 +1577,23 @@ export class DeepChatLoopRunner {
 
       let reviewConversationMessages = messages
       let activeProviderAttemptIdentity: DeepChatProviderAttemptIdentity | null = null
+      let providerRetryWaiting = false
+      const clearProviderRetryWaiting = () => {
+        if (!providerRetryWaiting) return
+        clearRateLimitWaitingMessage(sessionId, rateLimitMessageId, activeGeneration.runId)
+        providerRetryWaiting = false
+      }
+      clearProviderRetryWaitingMessage = clearProviderRetryWaiting
+      activeGeneration.abortController.signal.addEventListener(
+        'abort',
+        clearProviderRetryWaiting,
+        { once: true }
+      )
+      removeProviderRetryAbortListener = () =>
+        activeGeneration.abortController.signal.removeEventListener(
+          'abort',
+          clearProviderRetryWaiting
+        )
       const result = await processStream({
         run: loopRun,
         onConversationMessagesChange: (nextMessages) => {
@@ -2053,6 +2073,21 @@ export class DeepChatLoopRunner {
                 )
             },
             retryObserver: (event) => {
+              if (event.type === 'retry_scheduled') {
+                providerRetryWaiting = true
+                emitProviderRetryWaitingMessage(
+                  sessionId,
+                  rateLimitMessageId,
+                  activeGeneration.runId,
+                  state.providerId,
+                  event.delayMs
+                )
+              } else if (
+                event.type === 'retry_started' ||
+                (event.type === 'retry_finished' && event.retryDecision !== 'retry_scheduled')
+              ) {
+                clearProviderRetryWaiting()
+              }
               logger.info('[DeepChatAgent] Provider retry lifecycle', {
                 sessionId,
                 messageId,
@@ -2397,6 +2432,8 @@ export class DeepChatLoopRunner {
       }
       throw errorToPropagate
     } finally {
+      clearProviderRetryWaitingMessage?.()
+      removeProviderRetryAbortListener?.()
       if (
         toolSurfaceCanaryIdentity &&
         toolSurfaceMode !== 'legacy' &&
@@ -2589,6 +2626,29 @@ export class DeepChatLoopRunner {
     messageId: string,
     requestId: string,
     snapshot: RateLimitQueueSnapshot
+  ): void {
+    this.emitRateLimitWaitingBlock(sessionId, messageId, requestId, snapshot)
+  }
+
+  private emitProviderRetryWaitingMessage(
+    sessionId: string,
+    messageId: string,
+    requestId: string,
+    providerId: string,
+    estimatedWaitTime: number
+  ): void {
+    this.emitRateLimitWaitingBlock(sessionId, messageId, requestId, {
+      providerId,
+      estimatedWaitTime
+    })
+  }
+
+  private emitRateLimitWaitingBlock(
+    sessionId: string,
+    messageId: string,
+    requestId: string,
+    snapshot: Pick<RateLimitQueueSnapshot, 'providerId' | 'estimatedWaitTime'> &
+      Partial<Pick<RateLimitQueueSnapshot, 'qpsLimit' | 'currentQps' | 'queueLength'>>
   ): void {
     const block: AssistantMessageBlock = {
       type: 'action',
@@ -2598,9 +2658,11 @@ export class DeepChatLoopRunner {
       timestamp: Date.now(),
       extra: {
         providerId: snapshot.providerId,
-        qpsLimit: snapshot.qpsLimit,
-        currentQps: snapshot.currentQps,
-        queueLength: snapshot.queueLength,
+        ...(snapshot.qpsLimit === undefined ? {} : { qpsLimit: snapshot.qpsLimit }),
+        ...(snapshot.currentQps === undefined
+          ? {}
+          : { currentQps: snapshot.currentQps }),
+        ...(snapshot.queueLength === undefined ? {} : { queueLength: snapshot.queueLength }),
         estimatedWaitTime: snapshot.estimatedWaitTime
       }
     }

test/main/agent/deepchat/harness/deepChatAgentHarness.test.ts

@@ -7036,6 +7036,143 @@ describe('DeepChatAgentHarness', () => {
       })
     })
 
+    it('keeps the retry status continuous across consecutive provider retries', async () => {
+      let providerAttempt = 0
+      llmProvider.providerInstance.coreStream.mockImplementation(async function* () {
+        providerAttempt += 1
+        if (providerAttempt <= 2) {
+          yield {
+            type: 'error',
+            error_message: 'temporarily unavailable',
+            failure: {
+              statusCode: 503,
+              retryable: true,
+              retryHeaders: { 'retry-after-ms': '0' }
+            }
+          }
+          return
+        }
+        yield { type: 'stop', stop_reason: 'complete' }
+      })
+      ;(processStream as ReturnType<typeof vi.fn>).mockImplementationOnce(async (params) => {
+        for await (const _event of params.coreStream(
+          params.run.messages,
+          params.modelId,
+          params.modelConfig,
+          params.temperature,
+          params.maxTokens,
+          params.run.resources.toolDefinitions
+        )) {
+        }
+        return { status: 'completed', stopReason: 'complete' }
+      })
+
+      await agent.initSession('s1', { providerId: 'openai', modelId: 'gpt-4' })
+      await agent.processMessage('s1', 'Hello')
+
+      const typedStreamUpdates = getPublishedPayloads('chat.stream.updated').filter(
+        (payload) => typeof payload?.messageId === 'string'
+      )
+      const rateLimitUpdates = typedStreamUpdates.filter((payload) =>
+        payload.messageId.startsWith('__rate_limit__:')
+      )
+      const rateLimitShowUpdates = rateLimitUpdates.filter(
+        (payload) => Array.isArray(payload.blocks) && payload.blocks.length === 1
+      )
+      const rateLimitClearUpdates = rateLimitUpdates.filter(
+        (payload) => Array.isArray(payload.blocks) && payload.blocks.length === 0
+      )
+
+      expect(rateLimitUpdates.map((payload) => payload.blocks.length)).toEqual([1, 0, 1, 0])
+      expect(rateLimitShowUpdates).toHaveLength(2)
+      expect(rateLimitClearUpdates).toHaveLength(2)
+      const typedRateLimitShow = rateLimitShowUpdates[0]
+      const typedRateLimitClear = rateLimitClearUpdates[1]
+      expect(typedRateLimitShow).toMatchObject({
+        sessionId: 's1',
+        requestId: expect.any(String),
+        blocks: [
+          expect.objectContaining({
+            type: 'action',
+            action_type: 'rate_limit',
+            status: 'pending',
+            extra: {
+              providerId: 'openai',
+              estimatedWaitTime: expect.any(Number)
+            }
+          })
+        ]
+      })
+      expect(typedRateLimitShow.blocks[0].extra).not.toHaveProperty('qpsLimit')
+      expect(typedRateLimitShow.blocks[0].extra).not.toHaveProperty('currentQps')
+      expect(typedRateLimitShow.blocks[0].extra).not.toHaveProperty('queueLength')
+      expect(typedRateLimitClear).toMatchObject({
+        sessionId: 's1',
+        requestId: typedRateLimitShow.requestId,
+        blocks: []
+      })
+      expect(llmProvider.providerInstance.coreStream).toHaveBeenCalledTimes(3)
+    })
+
+    it('clears the retry status when cancellation interrupts the retry wait', async () => {
+      llmProvider.providerInstance.coreStream.mockImplementation(async function* () {
+        yield {
+          type: 'error',
+          error_message: 'temporarily unavailable',
+          failure: {
+            statusCode: 503,
+            retryable: true,
+            retryHeaders: { 'retry-after-ms': '10000' }
+          }
+        }
+      })
+      ;(processStream as ReturnType<typeof vi.fn>).mockImplementationOnce(async (params) => {
+        try {
+          for await (const _event of params.coreStream(
+            params.run.messages,
+            params.modelId,
+            params.modelConfig,
+            params.temperature,
+            params.maxTokens,
+            params.run.resources.toolDefinitions
+          )) {
+          }
+          return { status: 'completed', stopReason: 'complete' }
+        } catch (error) {
+          return {
+            status: 'aborted',
+            stopReason: 'user_stop',
+            errorMessage: error instanceof Error ? error.message : String(error)
+          }
+        }
+      })
+
+      await agent.initSession('s1', { providerId: 'openai', modelId: 'gpt-4' })
+      const processing = agent.processMessage('s1', 'Hello')
+      await vi.waitFor(() => {
+        expect(
+          getPublishedPayloads('chat.stream.updated').some(
+            (payload) =>
+              payload?.messageId?.startsWith('__rate_limit__:') &&
+              Array.isArray(payload.blocks) &&
+              payload.blocks.length === 1
+          )
+        ).toBe(true)
+      })
+
+      await agent.cancelGeneration('s1')
+      await processing
+
+      const rateLimitUpdates = getPublishedPayloads('chat.stream.updated').filter((payload) =>
+        payload?.messageId?.startsWith('__rate_limit__:')
+      )
+      expect(rateLimitUpdates.at(-1)).toMatchObject({
+        sessionId: 's1',
+        blocks: []
+      })
+      expect(llmProvider.providerInstance.coreStream).toHaveBeenCalledTimes(1)
+    })
+
     it('does not call provider.coreStream when a queued request is canceled', async () => {
       const abortError = new Error('Aborted')
       abortError.name = 'AbortError'