@@ -555,6 +555,9 @@ func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msg
555555 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
556556
557557 // Process each event in the stream.
558+ timer := time.NewTimer(streamChunkTimeout)
559+ defer timer.Stop()
560+
558561loop:
559562 for {
560563 select {
@@ -562,6 +565,9 @@ loop:
562565 if !ok {
563566 break loop
564567 }
568+ // Reset the timeout timer since we received a chunk
569+ timer.Reset(streamChunkTimeout)
570+
565571 if processErr := a.processEvent(ctx, sessionID, &assistantMsg, event); processErr != nil {
566572 if errors.Is(processErr, context.Canceled) {
567573 a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")
@@ -570,7 +576,7 @@ loop:
570576 }
571577 return assistantMsg, nil, processErr
572578 }
573− case <-time.After(streamChunkTimeout):
579+ case <-timer.C:
574580 a.finishMessage(ctx, &assistantMsg, message.FinishReasonError, "Stream timeout", "No chunk received within timeout")
575581 return assistantMsg, nil, fmt.Errorf("stream chunk timeout")
576582 case <-ctx.Done():