diff --git a/AGENTS.md b/AGENTS.md index 9e89f7b..90245e5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -48,7 +48,7 @@ The canonical type system is OpenAI-shaped; `gemini.go`/`anthropic.go` translate ## Testing conventions - **RED-first TDD**: failing test first, then the fix. Table-driven tests; `httptest` servers for hermetic coverage; `newTestClient` helper pins `backoffUnit` to 1ms — restore package vars in `t.Cleanup`. -- Two timing knobs are package vars for tests: `backoffUnit`, `streamIdleTimeout`. +- Two timing knobs are package vars for tests: `backoffUnit`, `streamIdleTimeout`. Operators override the idle watchdog via `SetStreamIdleTimeout` (positive values only). - The e2e suite (`e2e_test.go`) is behind a `e2e` build tag and hits **live APIs**. It must never lose that tag. Keys come from env or a gitignored `.env`; contents are never logged; tests skip when a key is absent. Adding a provider = one `e2eTarget` entry; models overridable via `_E2E_MODEL`. - Live-provider behavior (e.g. DeepSeek eliding `reasoning_content`) is **not an SDK contract** — probe softly, assert only what the SDK guarantees (call success, parsing, canonical finish). Model answer correctness is never an assertion. - Timer hygiene: since Go 1.23 no drain-before-`Reset` is needed for `time.Timer`. Note: on some dev machines `time.After` + select-default spin loops have hung — prefer deadline loops in tests. diff --git a/README.md b/README.md index a365666..7db5604 100644 --- a/README.md +++ b/README.md @@ -84,10 +84,10 @@ Requests and results are provider-neutral. Unknown message roles are rejected at ```go type ChatRequest struct { Model string // optional; ChatClient's model wins when both set - Messages []Message // RoleUser | RoleAssistant | RoleSystem | RoleTool + Messages []Message // RoleUser | RoleAssistant | RoleSystem | RoleTool; Message.Cache → Anthropic user-block cache_control System []SystemBlock // {Text, Cache} — Cache marks Anthropic prompt-cache blocks Tools []ToolDef // {Name, Description, Parameters json.RawMessage} - Thinking string // "", "enabled", "disabled", "low", "medium", "high" + Thinking string // "", "enabled", "disabled", "low", "medium", "high", "max" ThinkingBudget int // explicit token budget where the provider supports it MaxTokens int // routed to max_completion_tokens on o-series/gpt-5 Temperature float64 // 0 = provider default; negative = explicit 0 @@ -95,11 +95,11 @@ type ChatRequest struct { type ChatResult struct { Content string - ReasoningContent string // provider thinking text (advisory) + ReasoningContent string // provider thinking text; replayed as reasoning_content on OpenAI-format assistant turns ThinkingSignature string // Anthropic: replay via Message.ThinkingSignature ToolCalls []ToolCall // {ID, Name, Arguments} FinishReason string // stop | length | tool_calls | content_filter | "" - Usage Usage // {PromptTokens, CompletionTokens, ReasoningTokens} + Usage Usage // PromptTokens is uncached-only; cache volumes in CacheRead/Creation/CachedTokens } ``` @@ -109,7 +109,7 @@ Finish reasons are canonical: anything a provider reports outside the vocabulary `CallStream` enforces four guarantees, each covered by regression tests: -1. **Idle watchdog** — a stream silent longer than `streamIdleTimeout` (120s default) fails with `ErrIdleTimeout`. Keepalive comments reset it. +1. **Idle watchdog** — a stream silent longer than `StreamIdleTimeout()` (120s default; override with `SetStreamIdleTimeout`, positive values only) fails with `ErrIdleTimeout`. Keepalive comments reset it. 2. **Hard wall-clock deadline** — the whole stream is bounded by the per-request timeout (`WithRequestTimeout`, default 120s; per-client via `SetRequestTimeout`). 3. **Retries never duplicate output** — retries happen only before the first emitted delta. A failure after partial output returns the partial `*ChatResult` plus a wrapped error and is never retried. 4. **No silent empty successes** — a provider that closes the stream before its completion signal (before `[DONE]` / `message_stop`) yields a retryable error, not an empty result. Gemini, whose streams legitimately end at EOF, is exempt. @@ -143,7 +143,7 @@ On Gemini, a tool result's `ToolName` may be omitted — the SDK recovers the fu ## Extended thinking - **Anthropic** — `thinking` blocks are parsed in both buffered and streaming modes. `ChatResult.ThinkingSignature` carries the provider signature; for tool loops, replay it on the assistant message (`Message.ReasoningContent` + `Message.ThinkingSignature`) — the SDK re-serializes it as the first block, as Anthropic's API requires. Omitting it makes extended-thinking tool loops fail mid-conversation. -- **DeepSeek / GLM** — reasoning streams as `DeltaReasoning` fragments and lands in `ReasoningContent` (advisory; not replayed). +- **DeepSeek / GLM** — reasoning streams as `DeltaReasoning` fragments and lands in `ReasoningContent`. Assistant-turn replay echoes it as `reasoning_content` (required for DeepSeek/GLM tool loops). GLM maps thinking `medium` → `reasoning_effort` `high` (no medium level) and `max` → `max`. - **Gemini** — `thought: true` parts map to reasoning deltas; `thinkingConfig` is derived from `Thinking`/`ThinkingBudget`. ## Learn-once fallbacks @@ -159,10 +159,11 @@ When a provider rejects a request pattern, the SDK learns the constraint **once ## Retry policy -8 attempts, exponential backoff capped at 30s with ±20% jitter, `Retry-After` (seconds or HTTP-date) honored, context cancellation honored between and during attempts. Persistent 429s surface as `*llm.RateLimitError{Attempts, RetryAfter}` — including when the retry sleep is cut short by a deadline, so the caller never loses the retry signal. `RateLimitError` unwraps to `*APIError` for `errors.As` access to `Status`/`Retryable`. +8 attempts, exponential backoff capped at 30s with ±20% jitter, `Retry-After` (seconds or HTTP-date) honored and capped at 120s, context cancellation honored between and during attempts. Retryable statuses include 408/429/5xx plus Cloudflare 520–524 and Anthropic 529. Persistent 429s surface as `*llm.RateLimitError{Attempts, RetryAfter}` — including when the retry sleep is cut short by a deadline, so the caller never loses the retry signal. A 429 whose body is billing exhaustion (`insufficient_quota`, `exceeded your current quota`, `insufficient balance`, `no resource package`) is not retryable and fails on the first attempt. `RateLimitError` unwraps to `*APIError` for `errors.As` access to `Status`/`Retryable`. ## Timeouts & cancellation +- The pooled transport honors `HTTP_PROXY` / `HTTPS_PROXY` / `NO_PROXY` via `http.ProxyFromEnvironment`. - Buffered calls: per-request timeout on the HTTP client (`WithRequestTimeout`, per-client `SetRequestTimeout` — race-safe, swap is atomic). - Streaming: the same timeout becomes the hard wall-clock deadline via context; per-attempt SSE reads are additionally bounded by the idle watchdog. - Every wait (backoff, Retry-After, stream reads) selects on the caller's context — cancellation propagates everywhere, and a cancelled call never misreports as "retry exhausted". diff --git a/anthropic.go b/anthropic.go index 0f131d1..5bd9759 100644 --- a/anthropic.go +++ b/anthropic.go @@ -28,6 +28,8 @@ type anSysBlock struct { type anBlock struct { Type string `json:"type"` // "text" | "tool_use" | "tool_result" Text string `json:"text,omitempty"` + // CacheControl is set on user text blocks when Message.Cache is true. + CacheControl *anCacheControl `json:"cache_control,omitempty"` // thinking (replayed assistant turns; must be the FIRST block and // carry the provider signature) Thinking string `json:"thinking,omitempty"` @@ -151,9 +153,13 @@ func buildAnthropicRequest(req *ChatRequest, model string, stream bool) ([]byte, case RoleSystem: out.System = append(out.System, anSysBlock{Type: "text", Text: m.Content}) case RoleUser: + blk := anBlock{Type: "text", Text: m.Content} + if m.Cache { + blk.CacheControl = &anCacheControl{Type: "ephemeral"} + } out.Messages = append(out.Messages, anMessage{ Role: "user", - Content: []anBlock{{Type: "text", Text: m.Content}}, + Content: []anBlock{blk}, }) case RoleAssistant: var blocks []anBlock @@ -213,13 +219,33 @@ type anRespBlock struct { Input json.RawMessage `json:"input"` } +type anUsage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationTokens int `json:"cache_creation_input_tokens"` + CacheReadTokens int `json:"cache_read_input_tokens"` +} + type anResponse struct { Content []anRespBlock `json:"content"` StopReason string `json:"stop_reason"` - Usage struct { - InputTokens int `json:"input_tokens"` - OutputTokens int `json:"output_tokens"` - } `json:"usage"` + Usage anUsage `json:"usage"` +} + +// usageFromAnthropic maps Messages API usage onto canonical Usage. +// Anthropic reports cache volumes exclusively — input_tokens is already +// uncached-only — so PromptTokens is left alone. +func usageFromAnthropic(u anUsage) Usage { + out := Usage{ + PromptTokens: u.InputTokens, + CompletionTokens: u.OutputTokens, + CacheCreationTokens: u.CacheCreationTokens, + CacheReadTokens: u.CacheReadTokens, + } + if u.CacheCreationTokens > 0 || u.CacheReadTokens > 0 { + out.CacheReported = true + } + return out } // mapAnthropicStopReason maps stop_reason to canonical values. @@ -275,10 +301,7 @@ func parseAnthropicResponse(body []byte) (*ChatResult, error) { res.ThinkingSignature = b.Signature } } - res.Usage = Usage{ - PromptTokens: r.Usage.InputTokens, - CompletionTokens: r.Usage.OutputTokens, - } + res.Usage = usageFromAnthropic(r.Usage) return res, nil } @@ -288,9 +311,7 @@ type anStreamEvent struct { Type string `json:"type"` // message_start Message struct { - Usage struct { - InputTokens int `json:"input_tokens"` - } `json:"usage"` + Usage anUsage `json:"usage"` } `json:"message"` // content_block_start / content_block_stop Index int `json:"index"` @@ -326,7 +347,7 @@ func mapAnthropicStreamEvent(data []byte, acc *streamAccum) ([]Delta, bool, erro var deltas []Delta switch ev.Type { case "message_start": - acc.usage.PromptTokens = ev.Message.Usage.InputTokens + acc.usage = usageFromAnthropic(ev.Message.Usage) case "content_block_start": if ev.ContentBlock.Type == "tool_use" { c := acc.call(ev.Index) diff --git a/anthropic_test.go b/anthropic_test.go index b0da9a0..4f95648 100644 --- a/anthropic_test.go +++ b/anthropic_test.go @@ -87,6 +87,37 @@ func TestBuildAnthropicRequest_Golden(t *testing.T) { } } +func TestBuildAnthropicRequest_UserCacheMarker(t *testing.T) { + req := &ChatRequest{Messages: []Message{ + {Role: RoleUser, Content: "cached turn", Cache: true}, + {Role: RoleUser, Content: "plain turn"}, + }} + body, err := buildAnthropicRequest(req, "claude", false) + if err != nil { + t.Fatal(err) + } + var m map[string]any + if err := json.Unmarshal(body, &m); err != nil { + t.Fatalf("decode: %v\n%s", err, body) + } + msgs := m["messages"].([]any) + if len(msgs) != 2 { + t.Fatalf("messages = %d, want 2", len(msgs)) + } + cached := msgs[0].(map[string]any)["content"].([]any)[0].(map[string]any) + if cached["text"] != "cached turn" { + t.Errorf("cached text = %v", cached["text"]) + } + cc, ok := cached["cache_control"].(map[string]any) + if !ok || cc["type"] != "ephemeral" { + t.Errorf("cached user block cache_control = %v, want {type:ephemeral}", cached["cache_control"]) + } + plain := msgs[1].(map[string]any)["content"].([]any)[0].(map[string]any) + if _, ok := plain["cache_control"]; ok { + t.Errorf("plain user block must omit cache_control, got %v", plain["cache_control"]) + } +} + func TestBuildAnthropicRequest_EmptyAssistantPlaceholder(t *testing.T) { req := &ChatRequest{Messages: []Message{ {Role: RoleUser, Content: "q"}, diff --git a/chat.go b/chat.go index 5d25601..fc09907 100644 --- a/chat.go +++ b/chat.go @@ -80,9 +80,23 @@ const ( // streamIdleTimeout bounds the silence between SSE events. Thinking models // can legitimately spend minutes before their first event, so the default -// is generous. Package var so tests can shorten it. +// is generous. Package var so tests can shorten it; operators override +// via SetStreamIdleTimeout. var streamIdleTimeout = 120 * time.Second +// SetStreamIdleTimeout overrides the SSE idle watchdog. Call at startup, +// before the first request; non-positive values are ignored. +func SetStreamIdleTimeout(d time.Duration) { + if d > 0 { + streamIdleTimeout = d + } +} + +// StreamIdleTimeout reports the active idle watchdog (introspection/tests). +func StreamIdleTimeout() time.Duration { + return streamIdleTimeout +} + // errStreamStop is the internal sentinel for a clean stream end. var errStreamStop = errors.New("llm: stream complete") @@ -315,7 +329,8 @@ func billingExhausted(e *APIError) bool { m := strings.ToLower(e.Message) return strings.Contains(m, "insufficient balance") || strings.Contains(m, "insufficient_quota") || - strings.Contains(m, "no resource package") + strings.Contains(m, "no resource package") || + strings.Contains(m, "exceeded your current quota") } // retryDelay picks Retry-After when present, else exponential backoff. diff --git a/dispatch_edges_test.go b/dispatch_edges_test.go index 6e656d1..f15c51e 100644 --- a/dispatch_edges_test.go +++ b/dispatch_edges_test.go @@ -813,8 +813,8 @@ func TestBuildOpenAIRequestSystemAndToolRoles(t *testing.T) { t.Fatal(err) } s := string(b) - if !strings.Contains(s, `"role":"system","content":"sys\nband"`) { - t.Errorf("system fold missing: %s", s) + if !strings.Contains(s, `"role":"system","content":"sys"`) || !strings.Contains(s, `"role":"system","content":"band"`) { + t.Errorf("separate system messages missing: %s", s) } if !strings.Contains(s, `"role":"tool","content":"r","tool_call_id":"t1"`) { t.Errorf("tool role missing: %s", s) @@ -1003,6 +1003,31 @@ func TestListModelsGeminiMidPageError(t *testing.T) { } } +func TestBillingExhausted_Markers(t *testing.T) { + for _, tc := range []struct { + name string + status int + message string + want bool + }{ + {"z.ai insufficient balance", 429, "Insufficient balance or no resource package. Please recharge.", true}, + {"openai insufficient_quota code", 429, "insufficient_quota", true}, + {"openai quota message", 429, "You exceeded your current quota, please check your plan and billing details.", true}, + {"deepseek balance", 429, "Insufficient Balance", true}, + {"plain rate limit", 429, "rate limit exceeded, retry later", false}, + {"billing text on non-429", 500, "Insufficient balance", false}, + {"nil error", 429, "", false}, + } { + var e *APIError + if tc.name != "nil error" { + e = &APIError{Status: tc.status, Message: tc.message} + } + if got := billingExhausted(e); got != tc.want { + t.Errorf("%s: billingExhausted = %v, want %v", tc.name, got, tc.want) + } + } +} + // ── billing exhaustion: 429 that is really a permanent billing failure ─── // Billing/resource exhaustion signalled as 429 is permanent — the SDK diff --git a/idle_timeout_test.go b/idle_timeout_test.go new file mode 100644 index 0000000..d5416cb --- /dev/null +++ b/idle_timeout_test.go @@ -0,0 +1,25 @@ +package llm + +import ( + "testing" + "time" +) + +// SetStreamIdleTimeout pins the setter contract: positive values apply, +// zero and negative are ignored so a misconfigured caller cannot disable +// the watchdog by accident. +func TestSetStreamIdleTimeout(t *testing.T) { + orig := StreamIdleTimeout() + t.Cleanup(func() { streamIdleTimeout = orig }) + + SetStreamIdleTimeout(5 * time.Second) + if got := StreamIdleTimeout(); got != 5*time.Second { + t.Fatalf("StreamIdleTimeout() = %v, want 5s", got) + } + + SetStreamIdleTimeout(0) + SetStreamIdleTimeout(-1 * time.Second) + if got := StreamIdleTimeout(); got != 5*time.Second { + t.Fatalf("non-positive override applied; StreamIdleTimeout() = %v, want 5s", got) + } +} diff --git a/message.go b/message.go index 6129f22..deab76c 100644 --- a/message.go +++ b/message.go @@ -47,11 +47,12 @@ type ToolCall struct { // Message is one canonical chat message. For RoleTool messages, ToolCallID // and ToolName identify the call being answered and Content carries the // tool result. ReasoningContent is provider-reported thinking text -// (deepseek-reasoner, anthropic thinking, gemini thoughts). It is consumer -// metadata; the SDK replays it back only where a provider requires it for -// conversation continuity — Anthropic, and only when ThinkingSignature is -// also set (extended-thinking tool loops mandate the signed thinking block -// as the first block of the replayed assistant turn). +// (deepseek-reasoner, anthropic thinking, gemini thoughts). The SDK +// replays it where a provider requires conversation continuity: +// OpenAI-format assistant messages echo it as reasoning_content +// (DeepSeek/GLM tool loops), and Anthropic re-serializes a signed +// thinking block as the first content block when ThinkingSignature is +// also set. type Message struct { Role Role Content string @@ -62,6 +63,10 @@ type Message struct { ToolCalls []ToolCall ToolCallID string ToolName string + // Cache marks this user message for Anthropic prompt caching + // (cache_control ephemeral on the text block). Ignored on other + // formats and on non-user roles. + Cache bool } // SystemBlock is one system-prompt segment. On Anthropic each block maps to @@ -82,15 +87,24 @@ type ToolDef struct { } // Usage reports token accounting. Fields the provider does not report stay 0. +// PromptTokens is exclusive (uncached-only) after provider-specific +// normalization: OpenAI cached_tokens and DeepSeek hit/miss are subsets of +// prompt_tokens and are subtracted; Anthropic reports cache volumes +// exclusively and is left alone. Cache volumes live in the cache fields so +// budget enforcement can sum without double-counting. type Usage struct { - PromptTokens int - CompletionTokens int - ReasoningTokens int + PromptTokens int + CompletionTokens int + ReasoningTokens int + CacheReadTokens int + CacheCreationTokens int + CachedTokens int + CacheReported bool } // ChatRequest is the canonical request. Model is filled from the ChatClient // when empty. Thinking accepts "", "enabled", "disabled", "low", "medium", -// "high" and is translated per provider format. Temperature: 0 means use +// "high", "max" and is translated per provider format. Temperature: 0 means use // the provider default (field omitted); use a negative value to explicitly // send 0. type ChatRequest struct { diff --git a/openai.go b/openai.go index 52d80bc..7519704 100644 --- a/openai.go +++ b/openai.go @@ -28,10 +28,11 @@ type oaToolCall struct { } type oaMessage struct { - Role string `json:"role"` - Content *string `json:"content"` // nil keeps JSON null for tool calls - ToolCalls []oaToolCall `json:"tool_calls,omitempty"` - ToolCallID string `json:"tool_call_id,omitempty"` + Role string `json:"role"` + Content *string `json:"content"` // nil keeps JSON null for tool calls + ReasoningContent string `json:"reasoning_content,omitempty"` + ToolCalls []oaToolCall `json:"tool_calls,omitempty"` + ToolCallID string `json:"tool_call_id,omitempty"` } type oaToolDef struct { @@ -77,24 +78,24 @@ func buildOpenAIRequest(cfg ProviderConfig, req *ChatRequest, model string, stre out.MaxTokens = req.MaxTokens } - // System prompt: canonical blocks (+ any in-band system messages) - // concatenate into a leading system message. - var sys strings.Builder + // One OpenAI system message per SystemBlock and per in-band system + // role. Concatenating would collapse prompt-tiering (stable base + + // volatile memory) and bust prefix cache on every memory refresh. + msgs := make([]oaMessage, 0, len(req.Messages)+len(req.System)+1) + appendSystem := func(text string) { + if t := strings.TrimRight(text, "\n"); t != "" { + s := t + msgs = append(msgs, oaMessage{Role: "system", Content: &s}) + } + } for _, b := range req.System { - sys.WriteString(b.Text) - sys.WriteString("\n") + appendSystem(b.Text) } for _, m := range req.Messages { if m.Role == RoleSystem { - sys.WriteString(m.Content) - sys.WriteString("\n") + appendSystem(m.Content) } } - msgs := make([]oaMessage, 0, len(req.Messages)+1) - if sys.Len() > 0 { - s := strings.TrimRight(sys.String(), "\n") - msgs = append(msgs, oaMessage{Role: "system", Content: &s}) - } for _, m := range req.Messages { switch m.Role { case RoleSystem: @@ -107,6 +108,7 @@ func buildOpenAIRequest(cfg ProviderConfig, req *ChatRequest, model string, stre om := oaMessage{Role: "assistant"} c := m.Content om.Content = &c + om.ReasoningContent = m.ReasoningContent for _, tc := range m.ToolCalls { om.ToolCalls = append(om.ToolCalls, oaToolCall{ ID: tc.ID, @@ -161,10 +163,15 @@ func buildOpenAIRequest(cfg ProviderConfig, req *ChatRequest, model string, stre } else { out.Thinking = &oaThinking{Type: "disabled"} } - case "low", "medium", "high": + case "low", "medium", "high", "max": out.Thinking = &oaThinking{Type: "enabled"} if q.ReasoningEffort { - out.ReasoningEffort = req.Thinking + // GLM has no "medium" effort level; odek's medium maps to high. + if req.Thinking == "medium" { + out.ReasoningEffort = "high" + } else { + out.ReasoningEffort = req.Thinking + } } } case q.ReasoningEffort: @@ -220,10 +227,62 @@ type oaUsageDetails struct { ReasoningTokens int `json:"reasoning_tokens"` } +type oaPromptDetails struct { + CachedTokens int `json:"cached_tokens"` +} + type oaRespUsage struct { - PromptTokens int `json:"prompt_tokens"` - CompletionTokens int `json:"completion_tokens"` - CompletionTokensDetails oaUsageDetails `json:"completion_tokens_details"` + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + CompletionTokensDetails oaUsageDetails `json:"completion_tokens_details"` + CacheCreationTokens int `json:"cache_creation_input_tokens"` + CacheReadTokens int `json:"cache_read_input_tokens"` + PromptTokensDetails *oaPromptDetails `json:"prompt_tokens_details"` + PromptCacheHitTokens int `json:"prompt_cache_hit_tokens"` + PromptCacheMissTokens int `json:"prompt_cache_miss_tokens"` +} + +// usageFromOpenAI maps a chat-completions usage object onto canonical +// Usage with odek applyUsage exclusive normalization: OpenAI +// cached_tokens and DeepSeek hit/miss are subsets of prompt_tokens and +// are subtracted so PromptTokens is uncached-only; Anthropic-named +// cache fields are exclusive and are not subtracted. Guards keep +// hostile payloads from driving PromptTokens negative. +func usageFromOpenAI(u *oaRespUsage) Usage { + if u == nil { + return Usage{} + } + out := Usage{ + PromptTokens: u.PromptTokens, + CompletionTokens: u.CompletionTokens, + ReasoningTokens: u.CompletionTokensDetails.ReasoningTokens, + CacheCreationTokens: u.CacheCreationTokens, + CacheReadTokens: u.CacheReadTokens, + } + if u.PromptTokensDetails != nil { + out.CachedTokens = u.PromptTokensDetails.CachedTokens + out.CacheReported = true + } + if u.CacheCreationTokens > 0 || u.CacheReadTokens > 0 { + out.CacheReported = true + } + // DeepSeek native fields: a hit is prompt content read from cache; a + // miss is newly processed content that DeepSeek then caches for + // future requests, i.e. a cache write. + if u.PromptCacheHitTokens > 0 || u.PromptCacheMissTokens > 0 { + out.CacheReadTokens += u.PromptCacheHitTokens + out.CacheCreationTokens += u.PromptCacheMissTokens + out.CacheReported = true + } + if u.PromptTokensDetails != nil && u.PromptTokensDetails.CachedTokens > 0 && + u.PromptTokensDetails.CachedTokens <= out.PromptTokens { + out.PromptTokens -= u.PromptTokensDetails.CachedTokens + out.CacheReadTokens += u.PromptTokensDetails.CachedTokens + } + if total := u.PromptCacheHitTokens + u.PromptCacheMissTokens; total > 0 && total <= out.PromptTokens { + out.PromptTokens -= total + } + return out } type oaResponse struct { @@ -283,11 +342,7 @@ func parseOpenAIResponse(body []byte) (*ChatResult, error) { }) } if r.Usage != nil { - res.Usage = Usage{ - PromptTokens: r.Usage.PromptTokens, - CompletionTokens: r.Usage.CompletionTokens, - ReasoningTokens: r.Usage.CompletionTokensDetails.ReasoningTokens, - } + res.Usage = usageFromOpenAI(r.Usage) } return res, nil } @@ -327,11 +382,7 @@ func mapOpenAIStreamEvent(data []byte, acc *streamAccum) (deltas []Delta, done b return nil, false, fmt.Errorf("llm: parse stream chunk: %w", err) } if c.Usage != nil { - acc.usage = Usage{ - PromptTokens: c.Usage.PromptTokens, - CompletionTokens: c.Usage.CompletionTokens, - ReasoningTokens: c.Usage.CompletionTokensDetails.ReasoningTokens, - } + acc.usage = usageFromOpenAI(c.Usage) } for _, ch := range c.Choices { d := ch.Delta diff --git a/openai_test.go b/openai_test.go index 27992a8..36759ff 100644 --- a/openai_test.go +++ b/openai_test.go @@ -78,6 +78,38 @@ func TestBuildOpenAIRequest_Golden(t *testing.T) { } } +func TestBuildOpenAIRequest_SeparateSystemMessages(t *testing.T) { + cfg := ProviderConfig{ID: "openai", Format: FormatOpenAI} + req := &ChatRequest{ + System: []SystemBlock{{Text: "stable-base"}, {Text: "volatile-memory"}}, + Messages: []Message{ + {Role: RoleSystem, Content: "skill-block"}, + {Role: RoleUser, Content: "Hi"}, + }, + } + oa := buildOpenAIRequest(cfg, req, "gpt-4o", false, false) + body, err := json.Marshal(oa) + if err != nil { + t.Fatal(err) + } + m := openaiReqMap(t, body) + msgs := m["messages"].([]any) + if len(msgs) != 4 { + t.Fatalf("messages len = %d, want 4 (3 system + user)", len(msgs)) + } + want := []string{"stable-base", "volatile-memory", "skill-block", "Hi"} + for i, w := range want { + got := msgs[i].(map[string]any) + role := "system" + if i == 3 { + role = "user" + } + if got["role"] != role || got["content"] != w { + t.Errorf("messages[%d] = %v, want role=%s content=%q", i, got, role, w) + } + } +} + func TestBuildOpenAIRequest_TemperatureForbiddenModels(t *testing.T) { cfg := ProviderConfig{ID: "openai", Format: FormatOpenAI, Quirks: Quirks{ReasoningEffort: true}} req := sampleRequest() @@ -159,6 +191,74 @@ func TestBuildOpenAIRequest_ThinkingVariants(t *testing.T) { } } +// GLM has no "medium" effort level and accepts "max". Mapping applies only +// when both ThinkingObject and ReasoningEffort are set. +func TestBuildOpenAIRequest_GLMThinkingMediumAndMax(t *testing.T) { + zai := ProviderConfig{ID: "zai", Format: FormatOpenAI, Quirks: Quirks{ThinkingObject: true, ReasoningEffort: true, ForceThinking: []string{"glm-5.3"}}} + base := sampleRequest() + + r := *base + r.Thinking = "medium" + oa := buildOpenAIRequest(zai, &r, "glm-5.3", false, true) + body, err := json.Marshal(oa) + if err != nil { + t.Fatal(err) + } + m := openaiReqMap(t, body) + if m["reasoning_effort"] != "high" { + t.Errorf("thinking medium → reasoning_effort %v, want high (GLM has no medium)", m["reasoning_effort"]) + } + th, _ := m["thinking"].(map[string]any) + if th == nil || th["type"] != "enabled" { + t.Errorf("thinking medium → thinking %v, want {type:enabled}", m["thinking"]) + } + + r = *base + r.Thinking = "max" + oa = buildOpenAIRequest(zai, &r, "glm-5.3", false, true) + body, err = json.Marshal(oa) + if err != nil { + t.Fatal(err) + } + m = openaiReqMap(t, body) + if m["reasoning_effort"] != "max" { + t.Errorf("thinking max → reasoning_effort %v, want max", m["reasoning_effort"]) + } + th, _ = m["thinking"].(map[string]any) + if th == nil || th["type"] != "enabled" { + t.Errorf("thinking max → thinking %v, want {type:enabled}", m["thinking"]) + } +} + +func TestBuildOpenAIRequest_ReasoningContentReplay(t *testing.T) { + cfg := ProviderConfig{ID: "deepseek", Format: FormatOpenAI, Quirks: Quirks{ThinkingObject: true}} + req := &ChatRequest{Messages: []Message{ + {Role: RoleUser, Content: "q"}, + {Role: RoleAssistant, Content: "a", ReasoningContent: "internal thoughts"}, + {Role: RoleUser, Content: "again"}, + }} + oa := buildOpenAIRequest(cfg, req, "deepseek-reasoner", false, false) + body, err := json.Marshal(oa) + if err != nil { + t.Fatal(err) + } + m := openaiReqMap(t, body) + msgs := m["messages"].([]any) + if len(msgs) != 3 { + t.Fatalf("messages = %d, want 3", len(msgs)) + } + asst := msgs[1].(map[string]any) + if asst["role"] != "assistant" { + t.Errorf("role = %v, want assistant", asst["role"]) + } + if asst["reasoning_content"] != "internal thoughts" { + t.Errorf("reasoning_content = %v, want %q (DeepSeek/GLM tool loops require echo)", asst["reasoning_content"], "internal thoughts") + } + if asst["content"] != "a" { + t.Errorf("content = %v, want a", asst["content"]) + } +} + func TestBuildOpenAIRequest_ToolMessages(t *testing.T) { cfg := ProviderConfig{ID: "kimi", Format: FormatOpenAI} req := &ChatRequest{Messages: []Message{ diff --git a/retry.go b/retry.go index 7db9c51..f2f8073 100644 --- a/retry.go +++ b/retry.go @@ -16,10 +16,16 @@ const ( maxRetries = 7 maxRetryBackoff = 30 * time.Second retryJitterFactor = 0.2 + // maxRetryAfter caps how long a server's Retry-After is honored. A + // pathological or hostile value (e.g. "Retry-After: 86400") must not + // wedge a call for hours; context cancellation can still break the + // wait sooner. + maxRetryAfter = 120 * time.Second ) // retryableStatus reports whether an HTTP status should be retried. -// 529 is Anthropic's "overloaded" status. +// 529 is Anthropic's "overloaded" status. 520–524 are Cloudflare-origin +// incidents (CF-fronted providers emit these during origin hiccups). func retryableStatus(code int) bool { switch code { case http.StatusRequestTimeout, http.StatusTooManyRequests, @@ -27,27 +33,33 @@ func retryableStatus(code int) bool { http.StatusServiceUnavailable, http.StatusGatewayTimeout, 529: return true } - return false + return code >= 520 && code <= 524 } // parseRetryAfter parses a Retry-After header value: either delay-seconds -// or an HTTP-date. Returns 0 when unparseable. +// or an HTTP-date. Returns 0 when unparseable. The result is capped at +// maxRetryAfter. func parseRetryAfter(v string, now time.Time) time.Duration { if v == "" { return 0 } + var d time.Duration if secs, err := strconv.Atoi(v); err == nil { if secs < 0 { return 0 } - return time.Duration(secs) * time.Second - } - if t, err := http.ParseTime(v); err == nil { - if d := t.Sub(now); d > 0 { - return d + d = time.Duration(secs) * time.Second + } else if t, err := http.ParseTime(v); err == nil { + if delta := t.Sub(now); delta > 0 { + d = delta } + } else { + return 0 } - return 0 + if d > maxRetryAfter { + d = maxRetryAfter + } + return d } // backoffUnit is the exponential base; a package var so tests can shrink diff --git a/retry_test.go b/retry_test.go index 1fbc729..da90f45 100644 --- a/retry_test.go +++ b/retry_test.go @@ -10,12 +10,12 @@ import ( ) func TestRetryableStatus(t *testing.T) { - for _, s := range []int{408, 429, 500, 502, 503, 504, 529} { + for _, s := range []int{408, 429, 500, 502, 503, 504, 520, 521, 522, 523, 524, 529} { if !retryableStatus(s) { t.Errorf("retryableStatus(%d) = false, want true", s) } } - for _, s := range []int{400, 401, 403, 404, 422, 200} { + for _, s := range []int{400, 401, 403, 404, 422, 200, 525, 530} { if retryableStatus(s) { t.Errorf("retryableStatus(%d) = true, want false", s) } @@ -59,6 +59,20 @@ func TestParseRetryAfter_Garbage(t *testing.T) { } } +func TestParseRetryAfter_CappedAt120s(t *testing.T) { + now := time.Now() + if d := parseRetryAfter("100000", now); d != maxRetryAfter { + t.Fatalf("parseRetryAfter(\"100000\") = %v, want cap %v", d, maxRetryAfter) + } + if d := parseRetryAfter("120", now); d != 120*time.Second { + t.Fatalf("parseRetryAfter(\"120\") = %v, want 120s (exactly the cap)", d) + } + future := now.UTC().Add(time.Hour).Format(http.TimeFormat) + if d := parseRetryAfter(future, now); d != maxRetryAfter { + t.Fatalf("parseRetryAfter(HTTP-date +1h) = %v, want cap %v", d, maxRetryAfter) + } +} + func TestBackoffDelay_CappedAndNonNegative(t *testing.T) { for attempt := 1; attempt <= 10; attempt++ { d := backoffDelay(attempt) diff --git a/stream_failure_test.go b/stream_failure_test.go index 62fb19d..5a204ec 100644 --- a/stream_failure_test.go +++ b/stream_failure_test.go @@ -205,8 +205,8 @@ func TestCallBuffered429DeadlineKeepsRateLimitError(t *testing.T) { if rl.Attempts != 1 { t.Errorf("Attempts = %d, want 1", rl.Attempts) } - if rl.RetryAfter != 3600*time.Second { - t.Errorf("RetryAfter = %v, want 1h", rl.RetryAfter) + if rl.RetryAfter != maxRetryAfter { + t.Errorf("RetryAfter = %v, want cap %v (header was 3600s)", rl.RetryAfter, maxRetryAfter) } if rl.Status != http.StatusTooManyRequests { t.Errorf("Status = %d, want 429", rl.Status) diff --git a/transport.go b/transport.go index bbe7e7e..18a170f 100644 --- a/transport.go +++ b/transport.go @@ -21,6 +21,7 @@ const ( // calls reuse TCP+TLS connections instead of handshaking every request. func newPooledTransport() *http.Transport { return &http.Transport{ + Proxy: http.ProxyFromEnvironment, MaxIdleConns: defaultMaxIdleConns, MaxIdleConnsPerHost: defaultMaxIdlePerHost, IdleConnTimeout: defaultIdleTimeout, diff --git a/transport_test.go b/transport_test.go new file mode 100644 index 0000000..214a511 --- /dev/null +++ b/transport_test.go @@ -0,0 +1,12 @@ +package llm + +import ( + "testing" +) + +func TestNewPooledTransport_HonorsHTTPProxy(t *testing.T) { + tr := newPooledTransport() + if tr.Proxy == nil { + t.Fatal("Proxy is nil; want http.ProxyFromEnvironment so HTTP_PROXY/HTTPS_PROXY are honored") + } +} diff --git a/usage_cache_test.go b/usage_cache_test.go new file mode 100644 index 0000000..428538d --- /dev/null +++ b/usage_cache_test.go @@ -0,0 +1,186 @@ +package llm + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "testing" +) + +// Cache-token exclusive normalization (odek applyUsage parity). +// +// Anthropic reports cache tokens exclusively (input_tokens excludes them). +// OpenAI (prompt_tokens_details.cached_tokens) and DeepSeek +// (prompt_cache_hit_tokens + prompt_cache_miss_tokens = prompt_tokens) +// report them inclusively, as subsets of prompt_tokens. +// +// Usage.PromptTokens must be exclusive ("uncached" input) on every +// provider, with cache volumes carried in CacheReadTokens / +// CacheCreationTokens, so budget enforcement can sum them without +// double-counting. CacheReported is true when any cache field was present. + +func TestUsageCache_OpenAICachedTokens(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{ + "choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 300, + "completion_tokens": 30, + "prompt_tokens_details": {"cached_tokens": 200} + } + }`) + })) + defer srv.Close() + + cc := newTestClient(t, ProviderConfig{ID: "openai", Format: FormatOpenAI, BaseURL: srv.URL, APIKey: "k"}, srv) + res, err := cc.Call(context.Background(), &ChatRequest{Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err != nil { + t.Fatal(err) + } + if res.Usage.PromptTokens != 100 { + t.Errorf("PromptTokens = %d, want 100 (300 prompt − 200 cached; exclusive)", res.Usage.PromptTokens) + } + if res.Usage.CacheReadTokens != 200 { + t.Errorf("CacheReadTokens = %d, want 200 (OpenAI cached_tokens)", res.Usage.CacheReadTokens) + } + if res.Usage.CachedTokens != 200 { + t.Errorf("CachedTokens = %d, want 200 (display field unchanged)", res.Usage.CachedTokens) + } + if !res.Usage.CacheReported { + t.Error("CacheReported = false, want true (prompt_tokens_details present)") + } + if res.Usage.CompletionTokens != 30 { + t.Errorf("CompletionTokens = %d, want 30", res.Usage.CompletionTokens) + } +} + +func TestUsageCache_AnthropicExclusive(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{ + "content": [{"type":"text","text":"ok"}], + "stop_reason":"end_turn", + "usage":{ + "input_tokens": 500, + "output_tokens": 50, + "cache_creation_input_tokens": 400, + "cache_read_input_tokens": 100 + } + }`) + })) + defer srv.Close() + + cc := newTestClient(t, ProviderConfig{ID: "anthropic", Format: FormatAnthropic, BaseURL: srv.URL, APIKey: "k", Quirks: Quirks{AnthropicVersion: "2023-06-01"}}, srv) + res, err := cc.Call(context.Background(), &ChatRequest{Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err != nil { + t.Fatal(err) + } + // Anthropic input_tokens is already uncached-only: no subtraction. + if res.Usage.PromptTokens != 500 { + t.Errorf("PromptTokens = %d, want 500 (Anthropic is exclusive already)", res.Usage.PromptTokens) + } + if res.Usage.CacheCreationTokens != 400 || res.Usage.CacheReadTokens != 100 { + t.Errorf("cache fields = %d/%d, want 400/100", res.Usage.CacheCreationTokens, res.Usage.CacheReadTokens) + } + if !res.Usage.CacheReported { + t.Error("CacheReported = false, want true (Anthropic cache fields present)") + } + if res.Usage.CompletionTokens != 50 { + t.Errorf("CompletionTokens = %d, want 50", res.Usage.CompletionTokens) + } +} + +func TestUsageCache_DeepSeekHitMiss(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{ + "choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 1000, + "completion_tokens": 40, + "prompt_cache_hit_tokens": 750, + "prompt_cache_miss_tokens": 250 + } + }`) + })) + defer srv.Close() + + cc := newTestClient(t, ProviderConfig{ID: "deepseek", Format: FormatOpenAI, BaseURL: srv.URL, APIKey: "k"}, srv) + res, err := cc.Call(context.Background(), &ChatRequest{Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err != nil { + t.Fatal(err) + } + if res.Usage.PromptTokens != 0 { + t.Errorf("PromptTokens = %d, want 0 (prompt 1000 = hit 750 + miss 250; every token is cache-accounted)", res.Usage.PromptTokens) + } + if res.Usage.CacheReadTokens != 750 { + t.Errorf("CacheReadTokens = %d, want 750", res.Usage.CacheReadTokens) + } + if res.Usage.CacheCreationTokens != 250 { + t.Errorf("CacheCreationTokens = %d, want 250", res.Usage.CacheCreationTokens) + } + if !res.Usage.CacheReported { + t.Error("CacheReported = false, want true (DeepSeek hit/miss present)") + } +} + +func TestUsageCache_HostileCachedTokensNeverNegative(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{ + "choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 50, + "completion_tokens": 5, + "prompt_tokens_details": {"cached_tokens": 500} + } + }`) + })) + defer srv.Close() + + cc := newTestClient(t, ProviderConfig{ID: "openai", Format: FormatOpenAI, BaseURL: srv.URL, APIKey: "k"}, srv) + res, err := cc.Call(context.Background(), &ChatRequest{Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err != nil { + t.Fatal(err) + } + if res.Usage.PromptTokens < 0 { + t.Errorf("PromptTokens = %d, must never go negative", res.Usage.PromptTokens) + } + if res.Usage.PromptTokens != 50 { + t.Errorf("PromptTokens = %d, want 50 (hostile cached_tokens skipped; no subtraction)", res.Usage.PromptTokens) + } + if !res.Usage.CacheReported { + t.Error("CacheReported = false, want true (details object present)") + } +} + +func TestUsageCache_OpenAIFormatAnthropicCacheFields(t *testing.T) { + // Some OpenAI-compatible gateways forward Anthropic-named cache fields + // on the chat-completions usage object. They are exclusive — do not + // subtract from prompt_tokens. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprint(w, `{ + "choices": [{"message": {"content": "ok"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": 500, + "completion_tokens": 50, + "cache_creation_input_tokens": 400, + "cache_read_input_tokens": 100 + } + }`) + })) + defer srv.Close() + + cc := newTestClient(t, ProviderConfig{ID: "openai", Format: FormatOpenAI, BaseURL: srv.URL, APIKey: "k"}, srv) + res, err := cc.Call(context.Background(), &ChatRequest{Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err != nil { + t.Fatal(err) + } + if res.Usage.PromptTokens != 500 { + t.Errorf("PromptTokens = %d, want 500 (Anthropic-named fields are exclusive)", res.Usage.PromptTokens) + } + if res.Usage.CacheCreationTokens != 400 || res.Usage.CacheReadTokens != 100 { + t.Errorf("cache fields = %d/%d, want 400/100", res.Usage.CacheCreationTokens, res.Usage.CacheReadTokens) + } + if !res.Usage.CacheReported { + t.Error("CacheReported = false, want true") + } +}