Merge pull request #22 from adamdottv/adam/retries

fix(anthropic): better 429/529 handling

cmd/root.go

@@ -107,6 +107,16 @@ func setupSubscriptions(app *app.App) (chan tea.Msg, func()) {
 			wg.Done()
 		}()
 	}
+	{
+		sub := app.Status.Subscribe(ctx)
+		wg.Add(1)
+		go func() {
+			for ev := range sub {
+				ch <- ev
+			}
+			wg.Done()
+		}()
+	}
 	return ch, func() {
 		cancel()
 		wg.Wait()

internal/app/services.go

@@ -11,7 +11,9 @@ import (
 	"github.com/kujtimiihoxha/termai/internal/lsp/watcher"
 	"github.com/kujtimiihoxha/termai/internal/message"
 	"github.com/kujtimiihoxha/termai/internal/permission"
+	"github.com/kujtimiihoxha/termai/internal/pubsub"
 	"github.com/kujtimiihoxha/termai/internal/session"
+	"github.com/kujtimiihoxha/termai/internal/tui/util"
 )
 
 type App struct {
@@ -25,6 +27,7 @@ type App struct {
 
 	Logger logging.Interface
 
+	Status  *pubsub.Broker[util.InfoMsg]
 	ceanups []func()
 }
 
@@ -43,6 +46,7 @@ func New(ctx context.Context, conn *sql.DB) *App {
 		Messages:    messages,
 		Permissions: permission.NewPermissionService(),
 		Logger:      log,
+		Status:      pubsub.NewBroker[util.InfoMsg](),
 		LSPClients:  make(map[string]*lsp.Client),
 	}
 

internal/llm/agent/agent.go

@@ -15,6 +15,8 @@ import (
 	"github.com/kujtimiihoxha/termai/internal/llm/provider"
 	"github.com/kujtimiihoxha/termai/internal/llm/tools"
 	"github.com/kujtimiihoxha/termai/internal/message"
+	"github.com/kujtimiihoxha/termai/internal/pubsub"
+	"github.com/kujtimiihoxha/termai/internal/tui/util"
 )
 
 type Agent interface {
@@ -92,9 +94,24 @@ func (c *agent) processEvent(
 		assistantMsg.AppendContent(event.Content)
 		return c.Messages.Update(*assistantMsg)
 	case provider.EventError:
+		// TODO: remove when realease
 		log.Println("error", event.Error)
+		c.App.Status.Publish(pubsub.UpdatedEvent, util.InfoMsg{
+			Type: util.InfoTypeError,
+			Msg:  event.Error.Error(),
+		})
 		return event.Error
-
+	case provider.EventWarning:
+		c.App.Status.Publish(pubsub.UpdatedEvent, util.InfoMsg{
+			Type: util.InfoTypeWarn,
+			Msg:  event.Info,
+		})
+		return nil
+	case provider.EventInfo:
+		c.App.Status.Publish(pubsub.UpdatedEvent, util.InfoMsg{
+			Type: util.InfoTypeInfo,
+			Msg:  event.Info,
+		})
 	case provider.EventComplete:
 		assistantMsg.SetToolCalls(event.Response.ToolCalls)
 		assistantMsg.AddFinish(event.Response.FinishReason)

internal/llm/provider/anthropic.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"strings"
+	"time"
 
 	"github.com/anthropics/anthropic-sdk-go"
 	"github.com/anthropics/anthropic-sdk-go/option"
@@ -68,7 +70,9 @@ func (a *anthropicProvider) SendMessages(ctx context.Context, messages []message
 	anthropicMessages := a.convertToAnthropicMessages(messages)
 	anthropicTools := a.convertToAnthropicTools(tools)
 
-	response, err := a.client.Messages.New(ctx, anthropic.MessageNewParams{
+	response, err := a.client.Messages.New(
+		ctx,
+		anthropic.MessageNewParams{
 			Model:       anthropic.Model(a.model.APIModel),
 			MaxTokens:   a.maxTokens,
 			Temperature: anthropic.Float(0),
@@ -82,7 +86,8 @@ func (a *anthropicProvider) SendMessages(ctx context.Context, messages []message
 					},
 				},
 			},
-	})
+		},
+	)
 	if err != nil {
 		return nil, err
 	}
@@ -121,7 +126,52 @@ func (a *anthropicProvider) StreamResponse(ctx context.Context, messages []messa
 		temperature = anthropic.Float(1)
 	}
 
-	stream := a.client.Messages.NewStreaming(ctx, anthropic.MessageNewParams{
+	eventChan := make(chan ProviderEvent)
+
+	go func() {
+		defer close(eventChan)
+
+		const maxRetries = 8
+		attempts := 0
+
+		for {
+			// If this isn't the first attempt, we're retrying
+			if attempts > 0 {
+				if attempts > maxRetries {
+					eventChan <- ProviderEvent{
+						Type:  EventError,
+						Error: errors.New("maximum retry attempts reached for rate limit (429)"),
+					}
+					return
+				}
+
+				// Inform user we're retrying with attempt number
+				eventChan <- ProviderEvent{
+					Type: EventWarning,
+					Info: fmt.Sprintf("[Retrying due to rate limit... attempt %d of %d]", attempts, maxRetries),
+				}
+
+				// Calculate backoff with exponential backoff and jitter
+				backoffMs := 2000 * (1 << (attempts - 1)) // 2s, 4s, 8s, 16s, 32s
+				jitterMs := int(float64(backoffMs) * 0.2)
+				totalBackoffMs := backoffMs + jitterMs
+
+				// Sleep with backoff, respecting context cancellation
+				select {
+				case <-ctx.Done():
+					eventChan <- ProviderEvent{Type: EventError, Error: ctx.Err()}
+					return
+				case <-time.After(time.Duration(totalBackoffMs) * time.Millisecond):
+					// Continue with retry
+				}
+			}
+
+			attempts++
+
+			// Create new streaming request
+			stream := a.client.Messages.NewStreaming(
+				ctx,
+				anthropic.MessageNewParams{
 					Model:       anthropic.Model(a.model.APIModel),
 					MaxTokens:   a.maxTokens,
 					Temperature: temperature,
@@ -136,21 +186,20 @@ func (a *anthropicProvider) StreamResponse(ctx context.Context, messages []messa
 							},
 						},
 					},
-	})
-
-	eventChan := make(chan ProviderEvent)
-
-	go func() {
-		defer close(eventChan)
+				},
+			)
 
+			// Process stream events
 			accumulatedMessage := anthropic.Message{}
+			streamSuccess := false
 
+			// Process the stream until completion or error
 			for stream.Next() {
 				event := stream.Current()
 				err := accumulatedMessage.Accumulate(event)
 				if err != nil {
 					eventChan <- ProviderEvent{Type: EventError, Error: err}
-				return
+					return // Don't retry on accumulation errors
 				}
 
 				switch event := event.AsAny().(type) {
@@ -174,6 +223,7 @@ func (a *anthropicProvider) StreamResponse(ctx context.Context, messages []messa
 					eventChan <- ProviderEvent{Type: EventContentStop}
 
 				case anthropic.MessageStopEvent:
+					streamSuccess = true
 					content := ""
 					for _, block := range accumulatedMessage.Content {
 						if text, ok := block.AsAny().(anthropic.TextBlock); ok {
@@ -196,8 +246,51 @@ func (a *anthropicProvider) StreamResponse(ctx context.Context, messages []messa
 				}
 			}
 
-		if stream.Err() != nil {
-			eventChan <- ProviderEvent{Type: EventError, Error: stream.Err()}
+			// If the stream completed successfully, we're done
+			if streamSuccess {
+				return
+			}
+
+			// Check for stream errors
+			err := stream.Err()
+			if err != nil {
+				var apierr *anthropic.Error
+				if errors.As(err, &apierr) {
+					if apierr.StatusCode == 429 || apierr.StatusCode == 529 {
+						// Check for Retry-After header
+						if retryAfterValues := apierr.Response.Header.Values("Retry-After"); len(retryAfterValues) > 0 {
+							// Parse the retry after value (seconds)
+							var retryAfterSec int
+							if _, err := fmt.Sscanf(retryAfterValues[0], "%d", &retryAfterSec); err == nil {
+								retryMs := retryAfterSec * 1000
+
+								// Inform user of retry with specific wait time
+								eventChan <- ProviderEvent{
+									Type: EventWarning,
+									Info: fmt.Sprintf("[Rate limited: waiting %d seconds as specified by API]", retryAfterSec),
+								}
+
+								// Sleep respecting context cancellation
+								select {
+								case <-ctx.Done():
+									eventChan <- ProviderEvent{Type: EventError, Error: ctx.Err()}
+									return
+								case <-time.After(time.Duration(retryMs) * time.Millisecond):
+									// Continue with retry after specified delay
+									continue
+								}
+							}
+						}
+
+						// Fall back to exponential backoff if Retry-After parsing failed
+						continue
+					}
+				}
+
+				// For non-rate limit errors, report and exit
+				eventChan <- ProviderEvent{Type: EventError, Error: err}
+				return
+			}
 		}
 	}()
 
@@ -311,3 +404,4 @@ func (a *anthropicProvider) convertToAnthropicMessages(messages []message.Messag
 
 	return anthropicMessages
 }
+

internal/llm/provider/provider.go

@@ -17,6 +17,8 @@ const (
 	EventContentStop   EventType = "content_stop"
 	EventComplete      EventType = "complete"
 	EventError         EventType = "error"
+	EventWarning       EventType = "warning"
+	EventInfo          EventType = "info"
 )
 
 type TokenUsage struct {
@@ -40,6 +42,9 @@ type ProviderEvent struct {
 	ToolCall *message.ToolCall
 	Error    error
 	Response *ProviderResponse
+
+	// Used for giving users info on e.x retry
+	Info string
 }
 
 type Provider interface {

internal/tui/components/core/status.go

@@ -7,14 +7,14 @@ import (
 	"github.com/charmbracelet/lipgloss"
 	"github.com/kujtimiihoxha/termai/internal/config"
 	"github.com/kujtimiihoxha/termai/internal/llm/models"
+	"github.com/kujtimiihoxha/termai/internal/pubsub"
 	"github.com/kujtimiihoxha/termai/internal/tui/styles"
 	"github.com/kujtimiihoxha/termai/internal/tui/util"
 	"github.com/kujtimiihoxha/termai/internal/version"
 )
 
 type statusCmp struct {
-	err         error
-	info        string
+	info       *util.InfoMsg
 	width      int
 	messageTTL time.Duration
 }
@@ -34,17 +34,15 @@ func (m statusCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	switch msg := msg.(type) {
 	case tea.WindowSizeMsg:
 		m.width = msg.Width
-	case util.ErrorMsg:
-		m.err = msg
-		m.info = ""
+		return m, m.clearMessageCmd()
+	case pubsub.Event[util.InfoMsg]:
+		m.info = &msg.Payload
 		return m, m.clearMessageCmd()
 	case util.InfoMsg:
-		m.info = string(msg)
-		m.err = nil
+		m.info = &msg
 		return m, m.clearMessageCmd()
 	case util.ClearStatusMsg:
-		m.info = ""
-		m.err = nil
+		m.info = nil
 	}
 	return m, nil
 }
@@ -56,25 +54,25 @@ var (
 
 func (m statusCmp) View() string {
 	status := styles.Padded.Background(styles.Grey).Foreground(styles.Text).Render("? help")
-
-	if m.err != nil {
-		status += styles.Regular.Padding(0, 1).
-			Background(styles.Red).
-			Foreground(styles.Text).
-			Width(m.availableFooterMsgWidth()).
-			Render(m.err.Error())
-	} else if m.info != "" {
-		status += styles.Padded.
+	if m.info != nil {
+		infoStyle := styles.Padded.
 			Foreground(styles.Base).
-			Background(styles.Green).
-			Width(m.availableFooterMsgWidth()).
-			Render(m.info)
+			Width(m.availableFooterMsgWidth())
+		switch m.info.Type {
+		case util.InfoTypeInfo:
+			infoStyle = infoStyle.Background(styles.Blue)
+		case util.InfoTypeWarn:
+			infoStyle = infoStyle.Background(styles.Peach)
+		case util.InfoTypeError:
+			infoStyle = infoStyle.Background(styles.Red)
+		}
+		status += infoStyle.Render(m.info.Msg)
 	} else {
 		status += styles.Padded.
 			Foreground(styles.Base).
 			Background(styles.LightGrey).
 			Width(m.availableFooterMsgWidth()).
-			Render(m.info)
+			Render("")
 	}
 	status += m.model()
 	status += versionWidget
@@ -93,6 +91,6 @@ func (m statusCmp) model() string {
 
 func NewStatusCmp() tea.Model {
 	return &statusCmp{
-		messageTTL: 5 * time.Second,
+		messageTTL: 15 * time.Second,
 	}
 }

internal/tui/components/dialog/permission.go

@@ -112,13 +112,13 @@ func (p *permissionDialogCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				p.selectOption.Blur()
 				// Add a visual indicator for focus change
 				cmds = append(cmds, tea.Batch(
-					util.CmdHandler(util.InfoMsg("Viewing content - use arrow keys to scroll")),
+					util.ReportInfo("Viewing content - use arrow keys to scroll"),
 				))
 			} else {
 				p.selectOption.Focus()
 				// Add a visual indicator for focus change
 				cmds = append(cmds, tea.Batch(
-					util.CmdHandler(util.InfoMsg("Select an action")),
+					util.CmdHandler(util.ReportInfo("Select an action")),
 				))
 			}
 			return p, tea.Batch(cmds...)

internal/tui/components/repl/editor.go

@@ -140,7 +140,7 @@ func (m *editorCmp) Send() tea.Cmd {
 	return func() tea.Msg {
 		messages, _ := m.app.Messages.List(m.sessionID)
 		if hasUnfinishedMessages(messages) {
-			return util.InfoMsg("Assistant is still working on the previous message")
+			return util.ReportWarn("Assistant is still working on the previous message")
 		}
 		a, _ := agent.NewCoderAgent(m.app)
 

internal/tui/tui.go

@@ -77,6 +77,8 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	switch msg := msg.(type) {
 	case pubsub.Event[permission.PermissionRequest]:
 		return a, dialog.NewPermissionDialogCmd(msg.Payload)
+	case pubsub.Event[util.InfoMsg]:
+		a.status, _ = a.status.Update(msg)
 	case dialog.PermissionResponseMsg:
 		switch msg.Action {
 		case dialog.PermissionAllow:
@@ -121,8 +123,6 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		return a, a.moveToPage(msg.ID)
 	case util.InfoMsg:
 		a.status, _ = a.status.Update(msg)
-	case util.ErrorMsg:
-		a.status, _ = a.status.Update(msg)
 	case tea.KeyMsg:
 		if a.editorMode == vimtea.ModeNormal {
 			switch {
@@ -141,7 +141,7 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				if a.currentPage == page.ReplPage {
 					sessions, err := a.app.Sessions.List()
 					if err != nil {
-						return a, util.CmdHandler(util.ErrorMsg(err))
+						return a, util.CmdHandler(util.ReportError(err))
 					}
 					lastSession := sessions[0]
 					if lastSession.MessageCount == 0 {
@@ -149,7 +149,7 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 					}
 					s, err := a.app.Sessions.Create("New Session")
 					if err != nil {
-						return a, util.CmdHandler(util.ErrorMsg(err))
+						return a, util.CmdHandler(util.ReportError(err))
 					}
 					return a, util.CmdHandler(repl.SelectedSessionMsg{SessionID: s.ID})
 				}

internal/tui/util/util.go

@@ -9,12 +9,39 @@ func CmdHandler(msg tea.Msg) tea.Cmd {
 }
 
 func ReportError(err error) tea.Cmd {
-	return CmdHandler(ErrorMsg(err))
+	return CmdHandler(InfoMsg{
+		Type: InfoTypeError,
+		Msg:  err.Error(),
+	})
+}
+
+type InfoType int
+
+const (
+	InfoTypeInfo InfoType = iota
+	InfoTypeWarn
+	InfoTypeError
+)
+
+func ReportInfo(info string) tea.Cmd {
+	return CmdHandler(InfoMsg{
+		Type: InfoTypeInfo,
+		Msg:  info,
+	})
+}
+
+func ReportWarn(warn string) tea.Cmd {
+	return CmdHandler(InfoMsg{
+		Type: InfoTypeWarn,
+		Msg:  warn,
+	})
 }
 
 type (
-	InfoMsg       string
-	ErrorMsg      error
+	InfoMsg struct {
+		Type InfoType
+		Msg  string
+	}
 	ClearStatusMsg struct{}
 )