Merge pull request #2298 from dgageot/board/remove-trailing-headers-rate-limit-handl-873dc5c9

Remove trailing headers handling for rate limit headers

Author: dgageot

pkg/chat/chat.go

@@ -146,13 +146,12 @@ type MessageStreamChoice struct {
 
 // MessageStreamResponse represents a streaming response from the model
 type MessageStreamResponse struct {
-	ID        string                `json:"id"`
-	Object    string                `json:"object"`
-	Created   int64                 `json:"created"`
-	Model     string                `json:"model"`
-	Choices   []MessageStreamChoice `json:"choices"`
-	Usage     *Usage                `json:"usage,omitempty"`
-	RateLimit *RateLimit            `json:"rate_limit,omitempty"`
+	ID      string                `json:"id"`
+	Object  string                `json:"object"`
+	Created int64                 `json:"created"`
+	Model   string                `json:"model"`
+	Choices []MessageStreamChoice `json:"choices"`
+	Usage   *Usage                `json:"usage,omitempty"`
 }
 
 type Usage struct {
@@ -163,13 +162,6 @@ type Usage struct {
 	ReasoningTokens   int64 `json:"reasoning_tokens,omitempty"`
 }
 
-type RateLimit struct {
-	Limit      int64 `json:"limit,omitempty"`
-	Remaining  int64 `json:"remaining,omitempty"`
-	Reset      int64 `json:"reset,omitempty"`
-	RetryAfter int64 `json:"retry_after,omitempty"`
-}
-
 // MessageStream interface represents a stream of chat completions
 type MessageStream interface {
 	// Recv gets the next completion chunk

pkg/model/provider/anthropic/adapter.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
-	"strconv"
 	"strings"
 
 	"github.com/anthropics/anthropic-sdk-go"
@@ -20,17 +19,15 @@ import (
 type streamAdapter struct {
 	retryableStream[anthropic.MessageStreamEventUnion]
 
-	trackUsage         bool
-	toolCall           bool
-	toolID             string
-	getResponseTrailer func() http.Header
+	trackUsage bool
+	toolCall   bool
+	toolID     string
 }
 
 func (c *Client) newStreamAdapter(stream *ssestream.Stream[anthropic.MessageStreamEventUnion], trackUsage bool) *streamAdapter {
 	return &streamAdapter{
-		retryableStream:    retryableStream[anthropic.MessageStreamEventUnion]{stream: stream},
-		trackUsage:         trackUsage,
-		getResponseTrailer: c.getResponseTrailer,
+		retryableStream: retryableStream[anthropic.MessageStreamEventUnion]{stream: stream},
+		trackUsage:      trackUsage,
 	}
 }
 
@@ -151,30 +148,11 @@ func (a *streamAdapter) Recv() (chat.MessageStreamResponse, error) {
 		} else {
 			response.Choices[0].FinishReason = chat.FinishReasonStop
 		}
-
-		// MessageStopEvent is the last event. Let's drain the response to get the trailing headers.
-		trailers := a.getResponseTrailer()
-		if trailers.Get("X-RateLimit-Limit") != "" {
-			response.RateLimit = &chat.RateLimit{
-				Limit:      parseHeaderInt64(trailers.Get("X-RateLimit-Limit")),
-				Remaining:  parseHeaderInt64(trailers.Get("X-RateLimit-Remaining")),
-				Reset:      parseHeaderInt64(trailers.Get("X-RateLimit-Reset")),
-				RetryAfter: parseHeaderInt64(trailers.Get("Retry-After")),
-			}
-		}
 	}
 
 	return response, nil
 }
 
-func parseHeaderInt64(headerValue string) int64 {
-	value, err := strconv.ParseInt(headerValue, 10, 64)
-	if err != nil {
-		return 0
-	}
-	return value
-}
-
 // Close closes the stream
 func (a *streamAdapter) Close() {
 	a.stream.Close()

pkg/model/provider/anthropic/beta_adapter.go

@@ -3,7 +3,6 @@ package anthropic
 import (
 	"fmt"
 	"log/slog"
-	"net/http"
 
 	"github.com/anthropics/anthropic-sdk-go"
 	"github.com/anthropics/anthropic-sdk-go/packages/ssestream"
@@ -16,18 +15,16 @@ import (
 type betaStreamAdapter struct {
 	retryableStream[anthropic.BetaRawMessageStreamEventUnion]
 
-	trackUsage         bool
-	toolCall           bool
-	toolID             string
-	getResponseTrailer func() http.Header
+	trackUsage bool
+	toolCall   bool
+	toolID     string
 }
 
 // newBetaStreamAdapter creates a new Beta stream adapter
 func (c *Client) newBetaStreamAdapter(stream *ssestream.Stream[anthropic.BetaRawMessageStreamEventUnion], trackUsage bool) *betaStreamAdapter {
 	return &betaStreamAdapter{
-		retryableStream:    retryableStream[anthropic.BetaRawMessageStreamEventUnion]{stream: stream},
-		trackUsage:         trackUsage,
-		getResponseTrailer: c.getResponseTrailer,
+		retryableStream: retryableStream[anthropic.BetaRawMessageStreamEventUnion]{stream: stream},
+		trackUsage:      trackUsage,
 	}
 }
 

pkg/model/provider/anthropic/client.go

@@ -6,9 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io"
 	"log/slog"
-	"net/http"
 	"net/url"
 	"strings"
 
@@ -33,21 +31,8 @@ import (
 type Client struct {
 	base.Config
 
-	clientFn         func(context.Context) (anthropic.Client, error)
-	lastHTTPResponse *http.Response
-	fileManager      *FileManager
-}
-
-func (c *Client) getResponseTrailer() http.Header {
-	if c.lastHTTPResponse == nil {
-		return nil
-	}
-
-	if c.lastHTTPResponse.Body != nil {
-		_, _ = io.Copy(io.Discard, c.lastHTTPResponse.Body)
-	}
-
-	return c.lastHTTPResponse.Trailer
+	clientFn    func(context.Context) (anthropic.Client, error)
+	fileManager *FileManager
 }
 
 // adjustMaxTokensForThinking checks if max_tokens needs adjustment for thinking_budget.
@@ -207,7 +192,6 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro
 			}
 
 			client := anthropic.NewClient(
-				option.WithResponseInto(&anthropicClient.lastHTTPResponse),
 				option.WithAuthToken(authToken),
 				option.WithAPIKey(authToken),
 				option.WithBaseURL(baseURL),

pkg/runtime/event.go

@@ -286,7 +286,6 @@ type Usage struct {
 // It embeds chat.Usage and adds Cost, Model, and FinishReason fields.
 type MessageUsage struct {
 	chat.Usage
-	chat.RateLimit
 
 	Cost         float64
 	Model        string

pkg/runtime/loop.go

@@ -471,9 +471,6 @@ func (r *LocalRuntime) recordAssistantMessage(
 		Model:        messageModel,
 		FinishReason: res.FinishReason,
 	}
-	if res.RateLimit != nil {
-		msgUsage.RateLimit = *res.RateLimit
-	}
 	return msgUsage
 }
 

pkg/runtime/streaming.go

@@ -28,7 +28,6 @@ type streamResult struct {
 	Stopped           bool
 	FinishReason      chat.FinishReason
 	Usage             *chat.Usage
-	RateLimit         *chat.RateLimit
 }
 
 // handleStream reads a chat.MessageStream to completion, emitting streaming
@@ -44,7 +43,6 @@ func (r *LocalRuntime) handleStream(ctx context.Context, stream chat.MessageStre
 	var thoughtSignature []byte
 	var toolCalls []tools.ToolCall
 	var messageUsage *chat.Usage
-	var messageRateLimit *chat.RateLimit
 	var providerFinishReason chat.FinishReason
 
 	toolCallIndex := make(map[string]int)   // toolCallID -> index in toolCalls slice
@@ -89,10 +87,6 @@ func (r *LocalRuntime) handleStream(ctx context.Context, stream chat.MessageStre
 			messageUsage = response.Usage
 		}
 
-		if response.RateLimit != nil {
-			messageRateLimit = response.RateLimit
-		}
-
 		if len(response.Choices) == 0 {
 			continue
 		}
@@ -113,7 +107,6 @@ func (r *LocalRuntime) handleStream(ctx context.Context, stream chat.MessageStre
 				Stopped:           true,
 				FinishReason:      choice.FinishReason,
 				Usage:             messageUsage,
-				RateLimit:         messageRateLimit,
 			}, nil
 		}
 
@@ -236,7 +229,6 @@ func (r *LocalRuntime) handleStream(ctx context.Context, stream chat.MessageStre
 		Stopped:           stoppedDueToNoOutput,
 		FinishReason:      finishReason,
 		Usage:             messageUsage,
-		RateLimit:         messageRateLimit,
 	}, nil
 }