v0.0.5: context compaction, minimal config, retry/search/code symbols/diff/status/test commands

2026-06-24 02:20:19 -07:00
parent e4c75c7c0e
commit 950c63adaa
24 changed files with 1939 additions and 220 deletions
@@ -10,8 +10,13 @@ import (
 	"io"
 	"net/http"
 	"strings"
+	"time"
 )

+const maxChatAttempts = 3
+
+const retryDelay = 100 * time.Millisecond
+
 type OpenAICompatibleClient struct {
 	baseURL    string
 	apiKey     string
@@ -31,34 +36,82 @@ func NewOpenAICompatibleClient(baseURL, apiKey string, httpClient *http.Client)

 func (c *OpenAICompatibleClient) ChatStream(ctx context.Context, req ChatRequest, emit func(StreamEvent) error) error {
 	req.Stream = true
+	extra := make(map[string]any, len(req.Extra)+1)
+	for key, value := range req.Extra {
+		extra[key] = value
+	}
+	if _, exists := extra["stream_options"]; !exists {
+		extra["stream_options"] = map[string]any{"include_usage": true}
+	}
+	req.Extra = extra
 	body, err := json.Marshal(req)
 	if err != nil {
 		return fmt.Errorf("marshal chat request: %w", err)
 	}

-	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
-	if err != nil {
-		return fmt.Errorf("create chat request: %w", err)
-	}
-	httpReq.Header.Set("Content-Type", "application/json")
-	httpReq.Header.Set("Accept", "text/event-stream")
-	if c.apiKey != "" {
-		httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
-	}
+	var lastErr error
+	for attempt := 1; attempt <= maxChatAttempts; attempt++ {
+		httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
+		if err != nil {
+			return fmt.Errorf("create chat request: %w", err)
+		}
+		httpReq.Header.Set("Content-Type", "application/json")
+		httpReq.Header.Set("Accept", "text/event-stream")
+		if c.apiKey != "" {
+			httpReq.Header.Set("Authorization", "Bearer "+c.apiKey)
+		}

-	resp, err := c.httpClient.Do(httpReq)
-	if err != nil {
-		return fmt.Errorf("send chat request: %w", err)
-	}
-	defer resp.Body.Close()
+		resp, err := c.httpClient.Do(httpReq)
+		if err != nil {
+			lastErr = fmt.Errorf("send chat request: %w", err)
+			if !shouldRetryRequest(ctx, attempt, 0) {
+				return lastErr
+			}
+			if waitErr := waitBeforeRetry(ctx); waitErr != nil {
+				return waitErr
+			}
+			continue
+		}

-	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-		limit := io.LimitReader(resp.Body, 4096)
-		data, _ := io.ReadAll(limit)
-		return fmt.Errorf("chat request failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(data)))
-	}
+		if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+			limit := io.LimitReader(resp.Body, 4096)
+			data, _ := io.ReadAll(limit)
+			_ = resp.Body.Close()
+			lastErr = fmt.Errorf("chat request failed: status=%d body=%s", resp.StatusCode, strings.TrimSpace(string(data)))
+			if !shouldRetryRequest(ctx, attempt, resp.StatusCode) {
+				return lastErr
+			}
+			if waitErr := waitBeforeRetry(ctx); waitErr != nil {
+				return waitErr
+			}
+			continue
+		}

-	return readSSE(resp.Body, emit)
+		defer resp.Body.Close()
+		return readSSE(resp.Body, emit)
+	}
+	return lastErr
+}
+
+func shouldRetryRequest(ctx context.Context, attempt int, status int) bool {
+	if ctx.Err() != nil || attempt >= maxChatAttempts {
+		return false
+	}
+	if status == 0 {
+		return true
+	}
+	return status == http.StatusRequestTimeout || status == http.StatusTooManyRequests || status >= 500
+}
+
+func waitBeforeRetry(ctx context.Context) error {
+	timer := time.NewTimer(retryDelay)
+	defer timer.Stop()
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		return nil
+	}
 }

 func readSSE(r io.Reader, emit func(StreamEvent) error) error {
@@ -79,7 +132,7 @@ func readSSE(r io.Reader, emit func(StreamEvent) error) error {
 			if parseErr != nil {
 				return parseErr
 			}
-			if event.Content != "" || len(event.ToolCalls) > 0 || event.FinishReason != "" {
+			if event.Content != "" || len(event.ToolCalls) > 0 || event.FinishReason != "" || event.Usage != nil {
 				if emitErr := emit(event); emitErr != nil {
 					return emitErr
 				}
@@ -100,6 +153,7 @@ type streamPayload struct {
 		} `json:"delta"`
 		FinishReason string `json:"finish_reason"`
 	} `json:"choices"`
+	Usage *Usage `json:"usage,omitempty"`
 	Error *struct {
 		Message string `json:"message"`
 		Type    string `json:"type"`
@@ -127,15 +181,14 @@ func parseStreamPayload(payload string) (StreamEvent, error) {
 		}
 		return StreamEvent{}, fmt.Errorf("provider error: %s", decoded.Error.Message)
 	}
+	event := StreamEvent{Usage: decoded.Usage}
 	if len(decoded.Choices) == 0 {
-		return StreamEvent{}, nil
+		return event, nil
 	}

 	choice := decoded.Choices[0]
-	event := StreamEvent{
-		Content:      choice.Delta.Content,
-		FinishReason: choice.FinishReason,
-	}
+	event.Content = choice.Delta.Content
+	event.FinishReason = choice.FinishReason
 	for _, tc := range choice.Delta.ToolCalls {
 		event.ToolCalls = append(event.ToolCalls, ToolCallDelta{
 			Index:     tc.Index,
@@ -30,6 +30,10 @@ func TestOpenAICompatibleClientStreamsContentAndToolCalls(t *testing.T) {
 		if raw["thinking"] != "high" {
 			t.Fatalf("thinking = %#v", raw["thinking"])
 		}
+		streamOptions, ok := raw["stream_options"].(map[string]any)
+		if !ok || streamOptions["include_usage"] != true {
+			t.Fatalf("stream_options = %#v", raw["stream_options"])
+		}
 		if !req.Stream {
 			t.Fatal("request did not enable stream")
 		}
@@ -66,6 +70,75 @@ func TestOpenAICompatibleClientStreamsContentAndToolCalls(t *testing.T) {
 	}
 }

+func TestChatStreamCapturesUsage(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body := mustReadBody(t, r)
+		var raw map[string]any
+		if err := json.NewDecoder(strings.NewReader(body)).Decode(&raw); err != nil {
+			t.Fatal(err)
+		}
+		streamOptions, ok := raw["stream_options"].(map[string]any)
+		if !ok || streamOptions["include_usage"] != true {
+			t.Fatalf("stream_options = %#v", raw["stream_options"])
+		}
+
+		w.Header().Set("Content-Type", "text/event-stream")
+		_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"content":"ok"},"finish_reason":""}]}` + "\n\n"))
+		_, _ = w.Write([]byte(`data: {"choices":[],"usage":{"prompt_tokens":12,"completion_tokens":3,"total_tokens":15}}` + "\n\n"))
+		_, _ = w.Write([]byte("data: [DONE]\n\n"))
+	}))
+	defer server.Close()
+
+	client := NewOpenAICompatibleClient(server.URL, "sk-test", server.Client())
+	var usage *Usage
+	err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(event StreamEvent) error {
+		if event.Usage != nil {
+			usage = event.Usage
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if usage == nil {
+		t.Fatal("usage was not emitted")
+	}
+	if usage.PromptTokens != 12 || usage.CompletionTokens != 3 || usage.TotalTokens != 15 {
+		t.Fatalf("usage = %#v", usage)
+	}
+}
+
+func TestOpenAICompatibleClientRetriesTransientHTTPError(t *testing.T) {
+	attempts := 0
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		attempts++
+		if attempts == 1 {
+			http.Error(w, "temporary", http.StatusBadGateway)
+			return
+		}
+		w.Header().Set("Content-Type", "text/event-stream")
+		_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"content":"ok"},"finish_reason":""}]}` + "\n\n"))
+		_, _ = w.Write([]byte("data: [DONE]\n\n"))
+	}))
+	defer server.Close()
+
+	client := NewOpenAICompatibleClient(server.URL, "sk-test", server.Client())
+	var content strings.Builder
+	err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(event StreamEvent) error {
+		content.WriteString(event.Content)
+		return nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if attempts != 2 {
+		t.Fatalf("attempts = %d, want 2", attempts)
+	}
+	if content.String() != "ok" {
+		t.Fatalf("content = %q", content.String())
+	}
+}
+
 func mustReadBody(t *testing.T, r *http.Request) string {
 	t.Helper()
 	data, err := io.ReadAll(r.Body)
@@ -0,0 +1,43 @@
+package llm
+
+import "encoding/json"
+
+const tokenEstimateCharRatio = 4
+
+// EstimateTokens estimates the token cost of messages and tool definitions in a
+// chat completion request. It intentionally uses a simple JSON-size heuristic:
+// roughly four UTF-8 bytes per token plus a small per-message framing overhead.
+func EstimateTokens(messages []Message, tools []Tool) int {
+	if len(messages) == 0 && len(tools) == 0 {
+		return 0
+	}
+
+	payload := struct {
+		Messages []Message `json:"messages,omitempty"`
+		Tools    []Tool    `json:"tools,omitempty"`
+	}{
+		Messages: messages,
+		Tools:    tools,
+	}
+	data, err := json.Marshal(payload)
+	if err != nil {
+		return 0
+	}
+
+	return estimateBytesAsTokens(len(data)) + len(messages)*4
+}
+
+func estimateMessageTokens(msg Message) int {
+	data, err := json.Marshal(msg)
+	if err != nil {
+		return 0
+	}
+	return estimateBytesAsTokens(len(data)) + 4
+}
+
+func estimateBytesAsTokens(n int) int {
+	if n <= 0 {
+		return 0
+	}
+	return (n + tokenEstimateCharRatio - 1) / tokenEstimateCharRatio
+}
@@ -0,0 +1,45 @@
+package llm
+
+import "testing"
+
+func TestEstimateTokensEmpty(t *testing.T) {
+	if got := EstimateTokens(nil, nil); got != 0 {
+		t.Fatalf("EstimateTokens(nil, nil) = %d, want 0", got)
+	}
+}
+
+func TestEstimateTokens(t *testing.T) {
+	shortMessages := []Message{{Role: RoleUser, Content: "hello"}}
+	short := EstimateTokens(shortMessages, nil)
+	if short <= 4 {
+		t.Fatalf("short estimate = %d, want message framing plus content", short)
+	}
+	if short > 50 {
+		t.Fatalf("short estimate = %d, looks too high for one short message", short)
+	}
+
+	longMessages := []Message{
+		{Role: RoleUser, Content: "hello"},
+		{Role: RoleAssistant, Content: "This is a longer response with enough content to exceed the short estimate."},
+	}
+	long := EstimateTokens(longMessages, nil)
+	if long <= short {
+		t.Fatalf("long estimate = %d, want > short estimate %d", long, short)
+	}
+
+	withTools := EstimateTokens(shortMessages, []Tool{{
+		Type: "function",
+		Function: ToolFunction{
+			Name:        "file_read",
+			Description: "Read a file",
+			Parameters:  []byte(`{"type":"object"}`),
+		},
+	}})
+	if withTools <= short {
+		t.Fatalf("withTools estimate = %d, want > short estimate %d", withTools, short)
+	}
+
+	if msgTokens := estimateMessageTokens(shortMessages[0]); msgTokens <= 4 {
+		t.Fatalf("estimateMessageTokens = %d, want message framing plus content", msgTokens)
+	}
+}
@@ -79,6 +79,13 @@ type StreamEvent struct {
 	Content      string
 	ToolCalls    []ToolCallDelta
 	FinishReason string
+	Usage        *Usage
+}
+
+type Usage struct {
+	PromptTokens     int `json:"prompt_tokens"`
+	CompletionTokens int `json:"completion_tokens"`
+	TotalTokens      int `json:"total_tokens"`
 }

 type ToolCallDelta struct {