From 2f6b58cfc0791c9558c98a2e0d7e1f5cea94f062 Mon Sep 17 00:00:00 2001 From: "lutfi.ikbal" Date: Sun, 25 Oct 2026 17:06:41 +0700 Subject: [PATCH] feat(cache): stable prompt prefix, cache accounting and TUI diffs - split the system prompt into a stable prefix and a volatile trailing (git status, budget, todos, notifications, summaries) - stop sliding the transcript window per step; cut in large chunks only on context overflow - deterministic prompt assembly: sorted plugin prompts and summary files - stream_options.include_usage; parse OpenAI cached_tokens, Anthropic cache_read/cache_creation and DeepSeek prompt_cache_hit tokens; normalize into CachedTokens (and stop double counting) - Anthropic: four cache breakpoints (tools, system, previous tail, current tail) - OpenAI prompt_cache_key; add Zhipu/GLM providers - TUI: hydrate transcript from session history, render colored unified diffs - tests: cache breakpoints, normalization, retry, diff rendering Co-authored-by: CommandCodeBot --- CHANGELOG.md | 17 ++++ agent/agent.go | 124 +++++++++++++++++------ cmd/rakit/tui/diff_test.go | 29 ++++++ cmd/rakit/tui/transcript.go | 132 ++++++++++++++++++++++++ cmd/rakit/tui/tui.go | 6 ++ llm/cache_test.go | 79 +++++++++++++++ llm/provider.go | 196 ++++++++++++++++++++++++++---------- 7 files changed, 498 insertions(+), 85 deletions(-) create mode 100644 cmd/rakit/tui/diff_test.go create mode 100644 llm/cache_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a8caaa..53e5e28 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,23 @@ from this file's releases; see CONTRACT.md for the stability policy. ## v0.3.0 — unreleased ### Added +- Prompt-cache foundation: the system prompt is split into a stable prefix + (identity, project conventions, skills, memory) and a volatile + `` tail (git status, budget, todos, notifications, + summaries). The transcript is no longer slid one message per step; windows + only open in large chunks after a real context overflow. +- Deterministic prompt assembly: sorted plugin prompts and context-summary + files, stable tool order, no git status inside the cached prefix. +- Cache accounting: `stream_options.include_usage`, OpenAI + `prompt_tokens_details.cached_tokens`, Anthropic + `cache_read_input_tokens` / `cache_creation_input_tokens`, DeepSeek + `prompt_cache_hit_tokens`, and normalized `CachedTokens`. +- Anthropic cache strategy: four breakpoints (tools, system, previous-turn + tail, current tail). +- OpenAI `prompt_cache_key` for routing affinity; Zhipu/GLM providers + (`ZHIPU_API_KEY` / `GLM_API_KEY`). +- TUI: transcript hydration for `--continue`/`--session`, inline colored + unified diffs for edit/write, and diff expand/collapse. - Permission engine (`permission`): modes `plan|build|edit|yolo|auto`, per-tool allow/deny/ask rules from `~/.rakit/settings.json` and `/.rakit/settings.json`, per-call decisions, project/session diff --git a/agent/agent.go b/agent/agent.go index 3332395..1b6f821 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -1711,7 +1711,10 @@ func (m *Manager) runAgentTurn(ctx context.Context, sessionID string) { budgetSpent := 0 consecutiveToolFailures := 0 - transcriptWindowSize := 40 + // Keep the full transcript by default. Sliding the window one message per + // step destroys the provider's cached prefix; we only cut in large chunks + // after a real context-overflow error. + transcriptWindowSize := 1 << 30 isFinalSummaryStep := false for iter := 0; iter < maxReasoningSteps; iter++ { turnNum := iter + 1 @@ -1789,6 +1792,9 @@ func (m *Manager) runAgentTurn(ctx context.Context, sessionID string) { // results, which most providers reject). The transcript is windowed // so long sessions don't bloat the context; the stored session keeps // the full history for the UI. + // Stable prefix: identity/rules, project conventions (without git + // status), skills catalog, memory and plugin prompts. This must stay + // byte-identical across steps for provider prefix caching to hit. history, contextSummary := windowTranscriptWithSummary(sess.Messages, transcriptWindowSize) sysContent := sess.SystemPrompt dialect := sess.Dialect @@ -1800,46 +1806,63 @@ func (m *Manager) runAgentTurn(ctx context.Context, sessionID string) { } } if m.pluginMgr != nil { - for _, pluginPrompt := range m.pluginMgr.ActiveSystemPrompts() { + prompts := m.pluginMgr.ActiveSystemPrompts() + sort.Strings(prompts) + for _, pluginPrompt := range prompts { sysContent += "\n\n" + pluginPrompt } } + + // Volatile runtime context rides in a trailing system-reminder message + // so it never invalidates the cached prefix. + var volatileParts []string + if folder != "" { + if gitBlock := buildGitStatusBlock(folder); gitBlock != "" { + volatileParts = append(volatileParts, gitBlock) + } + } // Completed background tasks (bash run_in_background) announce their // outcome here — the model learns results without polling task_output. if notes := m.toolReg.DrainTaskNotifications(sessionID); len(notes) > 0 { - sysContent += "\n\n" + formatTaskNotifications(notes) + volatileParts = append(volatileParts, formatTaskNotifications(notes)) } // Proactively announced sub-agent completions from child sessions if subNotes := m.drainSubagentNotifications(sessionID); len(subNotes) > 0 { - sysContent += "\n\n" + formatSubagentNotifications(subNotes) + volatileParts = append(volatileParts, formatSubagentNotifications(subNotes)) } if contextSummary != "" { - sysContent += "\n\n" + contextSummary + volatileParts = append(volatileParts, contextSummary) } else if len(history) < len(sess.Messages) { - sysContent += "\n\n(Note: older conversation history was truncated from this request to fit the context window. The original task and recent exchanges are preserved.)" + volatileParts = append(volatileParts, "(Note: older conversation history was truncated from this request to fit the context window. The original task and recent exchanges are preserved.)") } if remaining := maxToolBudget - budgetSpent; remaining > 0 { - sysContent += fmt.Sprintf("\n\nTool budget remaining: %d points (cheap=1, medium=3, high=10). Prefer batched cheap tools (read_multiple, search) over many tiny calls; bash is expensive.", remaining) + volatileParts = append(volatileParts, fmt.Sprintf("Tool budget remaining: %d points (cheap=1, medium=3, high=10). Prefer batched cheap tools (read_multiple, search) over many tiny calls; bash is expensive.", remaining)) } else { - sysContent += "\n\nTool budget exhausted — wrap up with a summary; no more tool calls." + volatileParts = append(volatileParts, "Tool budget exhausted — wrap up with a summary; no more tool calls.") } if p := sess.Progress; p.CurrentGoal != "" || len(p.ActiveTodos) > 0 || len(p.CompletedSteps) > 0 { - sysContent += "\n\n## Current objective & tasks\n" + var pb strings.Builder + pb.WriteString("## Current objective & tasks\n") if p.CurrentGoal != "" { - sysContent += "Goal: " + p.CurrentGoal + "\n" + pb.WriteString("Goal: " + p.CurrentGoal + "\n") } if len(p.CompletedSteps) > 0 { - sysContent += "Completed: " + strings.Join(p.CompletedSteps, ", ") + "\n" + pb.WriteString("Completed: " + strings.Join(p.CompletedSteps, ", ") + "\n") } if len(p.ActiveTodos) > 0 { - sysContent += "Todos: " + strings.Join(p.ActiveTodos, ", ") + "\n" + pb.WriteString("Todos: " + strings.Join(p.ActiveTodos, ", ") + "\n") } if len(p.ActiveTodos) == 0 && len(p.CompletedSteps) > 0 { - sysContent += "(Note: All previous tasks are completed. When the user gives new instructions or next steps, call update_tasks with a fresh task breakdown to replace this checklist with the new plan.)\n" + pb.WriteString("(Note: All previous tasks are completed. When the user gives new instructions or next steps, call update_tasks with a fresh task breakdown to replace this checklist with the new plan.)\n") } if p.CurrentGoal != "" && len(sess.Tasks) == 0 { - sysContent += "(MANDATORY: A goal is active but no tasks exist yet. You MUST call update_tasks as your first tool call to decompose this goal into a sequential task checklist.)\n" + pb.WriteString("(MANDATORY: A goal is active but no tasks exist yet. You MUST call update_tasks as your first tool call to decompose this goal into a sequential task checklist.)\n") } + volatileParts = append(volatileParts, strings.TrimRight(pb.String(), "\n")) + } + volatileContent := "" + if len(volatileParts) > 0 { + volatileContent = "\n" + strings.Join(volatileParts, "\n\n") + "\n" } messages := []llm.LLMMessage{ {Role: llm.RoleSystem, Content: sysContent}, @@ -1883,6 +1906,12 @@ func (m *Manager) runAgentTurn(ctx context.Context, sessionID string) { log.Printf("[agent:%s] turn %d: offering %d tools (jevRouted=%v)", sessionID, turnNum, len(toolDefs), jevRouted) var dialectScanner *xmlDialectScanner + // Apply any late system-prompt additions (e.g. the final-summary + // instruction) before the request is built. + if len(messages) > 0 && messages[0].Role == llm.RoleSystem { + messages[0].Content = sysContent + } + var thoughtScanner *thoughtStreamScanner isPureDialect := dialect == "xml" if isPureDialect { @@ -1900,10 +1929,16 @@ func (m *Manager) runAgentTurn(ctx context.Context, sessionID string) { Role: llm.RoleUser, Content: renderDialectTranscript(history, catalog), }) + if volatileContent != "" { + messages = append(messages, llm.LLMMessage{Role: llm.RoleUser, Content: volatileContent}) + } toolDefs = nil dialectScanner = newXMLDialectScanner() } else { messages = append(messages, m.messagesToLLM(history)...) + if volatileContent != "" { + messages = append(messages, llm.LLMMessage{Role: llm.RoleUser, Content: volatileContent}) + } dialectScanner = nil thoughtScanner = newThoughtStreamScanner() } @@ -2045,9 +2080,21 @@ func (m *Manager) runAgentTurn(ctx context.Context, sessionID string) { continue } if isContextOverflowError(err) && transcriptWindowSize > 6 { + next := transcriptWindowSize / 2 + if transcriptWindowSize >= len(sess.Messages) || transcriptWindowSize > 400 { + // First overflow: cut in one large chunk rather than + // sliding one message per step. + next = len(sess.Messages) / 2 + } + if next < 10 { + next = 10 + } + if next >= transcriptWindowSize { + next = transcriptWindowSize / 2 + } log.Printf("[agent:%s] Context overflow detected (%v); auto-compacting transcript window from %d to %d and retrying step", - sessionID, err, transcriptWindowSize, transcriptWindowSize/2) - transcriptWindowSize = transcriptWindowSize / 2 + sessionID, err, transcriptWindowSize, next) + transcriptWindowSize = next iter-- // Retry the same step with smaller window continue } @@ -2109,14 +2156,14 @@ func (m *Manager) runAgentTurn(ctx context.Context, sessionID string) { sess.TokenUsage.TotalTokens += resp.TokenUsage.TotalTokens sess.ContextTokens = resp.TokenUsage.PromptTokens if sess.TokenUsage.PromptTokens > 0 { - totalCached := sess.TokenUsage.CachedTokens + sess.TokenUsage.PromptCacheHitTokens + totalCached := sess.TokenUsage.CachedTokens sess.CacheHitRate = float64(totalCached) / float64(sess.TokenUsage.PromptTokens) } activeModel := sess.Model if activeModel == "" && m.llmClient != nil { activeModel = m.llmClient.GetActiveModel() } - totalCached := sess.TokenUsage.CachedTokens + sess.TokenUsage.PromptCacheHitTokens + totalCached := sess.TokenUsage.CachedTokens sess.EstimatedCost = llm.EstimateCost(activeModel, sess.TokenUsage.PromptTokens, sess.TokenUsage.CompletionTokens, totalCached) // Sanitize assistant text to ensure no residual in-band markup or unclosed tags leak. cleaned := cleanInbandToolSyntax(assistant.Text()) @@ -4164,14 +4211,20 @@ func buildContextSummary(allMsgs, windowed []AgentMessage) string { sb.WriteString("\n## Relevant Files\n") if len(filesTouched) > 0 { - count := 0 + files := make([]string, 0, len(filesTouched)) for f := range filesTouched { + files = append(files, f) + } + sort.Strings(files) + limit := len(files) + if limit > 15 { + limit = 15 + } + for _, f := range files[:limit] { sb.WriteString(fmt.Sprintf("- `%s`\n", f)) - count++ - if count >= 15 { - sb.WriteString("- ... (additional files omitted)\n") - break - } + } + if len(files) > limit { + sb.WriteString("- ... (additional files omitted)\n") } } else { sb.WriteString("- (none)\n") @@ -4431,7 +4484,7 @@ func (m *Manager) GetUsageSummary() map[string]interface{} { pTokens := sess.TokenUsage.PromptTokens cTokens := sess.TokenUsage.CompletionTokens - cached := sess.TokenUsage.CachedTokens + sess.TokenUsage.PromptCacheHitTokens + cached := sess.TokenUsage.CachedTokens tTokens := sess.TokenUsage.TotalTokens if tTokens == 0 && (pTokens > 0 || cTokens > 0) { tTokens = pTokens + cTokens @@ -5431,16 +5484,25 @@ func buildProjectContext(folder string) string { sb.WriteString("\n") } - // Lightweight git status snapshot (best effort, fast). - if out, err := gitSnapshot(folder); err == nil && out != "" { - sb.WriteString("\nGit status:\n") - sb.WriteString(out) - sb.WriteString("\n") - } + // Git status is deliberately excluded here: it changes constantly and + // would invalidate the prompt-cache prefix. It is injected per step as + // volatile context by buildGitStatusBlock. return sb.String() } +// buildGitStatusBlock renders the per-step volatile git context. +func buildGitStatusBlock(folder string) string { + if folder == "" { + return "" + } + out, err := gitSnapshot(folder) + if err != nil || out == "" { + return "" + } + return "Git status:\n" + out +} + var ( gitSnapMu sync.Mutex gitSnapCache = make(map[string]gitSnapItem) diff --git a/cmd/rakit/tui/diff_test.go b/cmd/rakit/tui/diff_test.go new file mode 100644 index 0000000..3c891dc --- /dev/null +++ b/cmd/rakit/tui/diff_test.go @@ -0,0 +1,29 @@ +package tui + +import ( + "strings" + "testing" +) + +func TestRenderDiff(t *testing.T) { + m := Model{width: 100} + out := m.renderDiff("--- a/file.go\n+++ b/file.go\n@@ -1 +1 @@\n-old line\n+new line\n", false) + if !strings.Contains(out, "new line") || !strings.Contains(out, "old line") { + t.Fatalf("diff body missing: %q", out) + } + if !strings.Contains(out, "+1 −1") { + t.Fatalf("diff summary missing: %q", out) + } +} + +func TestRenderDiffCollapsedLimitsLines(t *testing.T) { + m := Model{width: 100} + var b strings.Builder + for i := 0; i < 30; i++ { + b.WriteString("+line\n") + } + out := m.renderDiff(b.String(), false) + if !strings.Contains(out, "more (ctrl+o expands)") { + t.Fatalf("expected truncation notice: %q", out) + } +} diff --git a/cmd/rakit/tui/transcript.go b/cmd/rakit/tui/transcript.go index 740f0e3..e9d979e 100644 --- a/cmd/rakit/tui/transcript.go +++ b/cmd/rakit/tui/transcript.go @@ -126,6 +126,7 @@ type toolInfo struct { mcp bool // MCP-bridged tool (distinct color) state toolState result string + diff string isErr bool expanded bool started time.Time @@ -452,6 +453,12 @@ func (m *Model) renderToolBlock(i int) string { return line + "\n" + toolErrStyle.Render(" ⎿ ✗ "+m.truncateWidth(oneLine(body))) } + // Edit/write results carry a unified diff; render it with added/removed + // coloring instead of a plain preview. + if t.diff != "" { + return line + "\n" + m.renderDiff(t.diff, t.expanded) + } + if t.expanded { text := strings.TrimRight(t.result, "\n") if len(text) > maxToolVerbose { @@ -506,6 +513,131 @@ func (m *Model) truncateWidth(s string) string { return string(runes[:w]) + "…" } +// renderDiff renders a unified diff with added/removed coloring. Collapsed +// mode shows the first 12 lines plus a summary. +func (m *Model) renderDiff(diff string, expanded bool) string { + lines := strings.Split(strings.TrimRight(diff, "\n"), "\n") + limit := len(lines) + if !expanded && limit > 12 { + limit = 12 + } + var sb strings.Builder + added, removed := 0, 0 + for i, ln := range lines { + switch { + case strings.HasPrefix(ln, "+++"), strings.HasPrefix(ln, "---"): + continue + case strings.HasPrefix(ln, "+"): + added++ + case strings.HasPrefix(ln, "-"): + removed++ + } + if i >= limit { + continue + } + style := resultStyle + switch { + case strings.HasPrefix(ln, "+"): + style = toolDotOKStyle + case strings.HasPrefix(ln, "-"): + style = toolDotErrStyle + case strings.HasPrefix(ln, "@@"): + style = statusStyle + } + mark := " " + if i == 0 { + mark = " ⎿ " + } + sb.WriteString("\n" + resultMarkStyle.Render(mark) + style.Render(m.truncateWidth(ln))) + } + if !expanded && len(lines) > limit { + sb.WriteString(statusStyle.Render(fmt.Sprintf("\n … + %d more (ctrl+o expands)", len(lines)-limit))) + } + sb.WriteString(statusStyle.Render(fmt.Sprintf("\n +%d −%d", added, removed))) + return sb.String() +} + +// setToolDiff attaches a unified diff to a tool block. +func (m *Model) setToolDiff(id, diff string) { + idx, ok := m.toolLines[id] + if !ok || idx >= len(m.transcript) || m.transcript[idx].tool == nil { + return + } + m.transcript[idx].tool.diff = diff + m.refreshViewport() +} + +// hydrateTranscript rebuilds the transcript from the session's stored +// messages so --continue / --session shows the real history. +func (m *Model) hydrateTranscript() { + if m.mgr == nil { + return + } + sess, ok := m.mgr.GetSession(m.sessionID) + if !ok { + return + } + for _, msg := range sess.Messages { + switch msg.Role { + case "user": + for _, b := range msg.Content { + if b.Type == "text" && strings.TrimSpace(b.Text) != "" { + m.appendBlock(userStyle.Render("❯ ") + b.Text) + } + } + case "assistant": + for _, b := range msg.Content { + switch b.Type { + case "text": + if strings.TrimSpace(b.Text) == "" { + continue + } + idx := m.ensureBlock(blockText) + m.transcript[idx].content = b.Text + m.transcript[idx].sealed = true + m.transcript[idx].rendered = renderMarkdown(b.Text, m.markdownWidth()) + m.transcript[idx].renderedWidth = m.markdownWidth() + case "reasoning", "thinking": + if strings.TrimSpace(b.Text) == "" { + continue + } + idx := m.ensureBlock(blockReasoning) + m.transcript[idx].content = b.Text + m.transcript[idx].reasoningOpen = false + m.transcript[idx].sealed = true + case "tool_use", "tool_call": + info := newToolInfo(b.Name, b.Arguments) + info.state = toolOK + idx := m.ensureBlock(blockTool) + m.transcript[idx].tool = info + if b.ToolCallID != "" { + m.toolLines[b.ToolCallID] = idx + } + } + } + case "tool": + for _, b := range msg.Content { + if b.Type != "tool_result" { + continue + } + idx, ok := m.toolLines[b.ToolCallID] + if !ok || idx >= len(m.transcript) || m.transcript[idx].tool == nil { + continue + } + t := m.transcript[idx].tool + t.result = b.Text + t.isErr = b.IsError + if b.IsError { + t.state = toolErr + } else { + t.state = toolOK + } + } + } + } + m.refreshViewport() +} + // toggleReasoning flips the collapsed state of the newest reasoning block. func (m *Model) toggleReasoning() { for i := len(m.transcript) - 1; i >= 0; i-- { diff --git a/cmd/rakit/tui/tui.go b/cmd/rakit/tui/tui.go index 5daaa11..4d5577d 100644 --- a/cmd/rakit/tui/tui.go +++ b/cmd/rakit/tui/tui.go @@ -165,6 +165,7 @@ func New(e *engine.Engine, sessionID string) Model { } m.pullSession() m.appendSessionHeader() + m.hydrateTranscript() return m } @@ -610,6 +611,11 @@ func (m *Model) handleBusEvent(ev events.Event) { result, _ := ev.Data["result"].(string) isErr, _ := ev.Data["is_error"].(bool) m.setToolResult(id, name, result, isErr) + if disp, ok := ev.Data["display"].(map[string]any); ok { + if diff, ok := disp["diff"].(string); ok && diff != "" { + m.setToolDiff(id, diff) + } + } case events.AgentTurnStart: m.closeStream() m.busy = true diff --git a/llm/cache_test.go b/llm/cache_test.go new file mode 100644 index 0000000..15b5ccd --- /dev/null +++ b/llm/cache_test.go @@ -0,0 +1,79 @@ +package llm + +import "testing" + +func TestTokenStatsNormalize(t *testing.T) { + tests := []struct { + name string + in TokenStats + want int + }{ + {"deepseek", TokenStats{PromptCacheHitTokens: 900}, 900}, + {"openai nested", TokenStats{PromptTokensDetails: PromptTokensDetails{CachedTokens: 700}}, 700}, + {"anthropic", TokenStats{CacheReadTokens: 500}, 500}, + {"flat", TokenStats{CachedTokens: 300}, 300}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := tt.in + s.Normalize() + if s.CachedTokens != tt.want { + t.Fatalf("got %d, want %d", s.CachedTokens, tt.want) + } + }) + } +} + +func TestPreparePayloadAnthropicBreakpoints(t *testing.T) { + messages := []LLMMessage{ + {Role: RoleSystem, Content: "system instructions"}, + {Role: RoleUser, Content: "first"}, + {Role: RoleAssistant, Content: "answer"}, + {Role: RoleUser, Content: "second"}, + {Role: RoleAssistant, Content: "answer two"}, + {Role: RoleTool, Content: "tool result"}, + } + tools := []ToolDefinition{{Type: "function", Function: FunctionSpec{Name: "read"}}} + rawMsgs, rawTools := PreparePayloadWithCache("anthropic", "claude-sonnet-4", messages, tools) + + msgMaps, ok := rawMsgs.([]map[string]any) + if !ok { + t.Fatalf("messages type %T", rawMsgs) + } + toolMaps, ok := rawTools.([]map[string]any) + if !ok { + t.Fatalf("tools type %T", rawTools) + } + count := 0 + if _, ok := toolMaps[0]["cache_control"]; ok { + count++ + } + for _, m := range msgMaps { + switch c := m["content"].(type) { + case []map[string]any: + for _, part := range c { + if _, ok := part["cache_control"]; ok { + count++ + } + } + } + } + if count != 4 { + t.Fatalf("cache breakpoints = %d, want 4", count) + } +} + +func TestPreparePayloadDeepSeekHasNoBreakpoints(t *testing.T) { + messages := []LLMMessage{{Role: RoleSystem, Content: "s"}, {Role: RoleUser, Content: "u"}} + rawMsgs, rawTools := PreparePayloadWithCache("deepseek", "deepseek-chat", messages, nil) + msgMaps, ok := rawMsgs.([]LLMMessage) + if !ok { + t.Fatalf("deepseek payload should be unchanged, got %T", rawMsgs) + } + if msgMaps[0].Content != "s" { + t.Fatal("content mutated") + } + if ts, ok := rawTools.([]ToolDefinition); !ok || len(ts) != 0 { + t.Fatalf("tools should stay empty, got %#v", rawTools) + } +} diff --git a/llm/provider.go b/llm/provider.go index 8789a63..2907538 100644 --- a/llm/provider.go +++ b/llm/provider.go @@ -104,6 +104,10 @@ type MCPTool struct { InputSchema map[string]interface{} `json:"input_schema"` } +type PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens,omitempty"` +} + type TokenStats struct { PromptTokens int `json:"prompt_tokens"` CompletionTokens int `json:"completion_tokens"` @@ -113,7 +117,25 @@ type TokenStats struct { // CachedTokens; DeepSeek providers leave it 0. PromptCacheHitTokens int `json:"prompt_cache_hit_tokens"` PromptCacheMissTokens int `json:"prompt_cache_miss_tokens"` - TotalTokens int `json:"total_tokens"` + // CacheReadTokens / CacheWriteTokens: Anthropic native accounting. + CacheReadTokens int `json:"cache_read_input_tokens,omitempty"` + CacheWriteTokens int `json:"cache_creation_input_tokens,omitempty"` + // PromptTokensDetails carries OpenAI's nested cached-token count. + PromptTokensDetails PromptTokensDetails `json:"prompt_tokens_details,omitempty"` + TotalTokens int `json:"total_tokens"` +} + +// Normalize folds the provider-specific cached-token fields into CachedTokens +// so every consumer sees one number. +func (t *TokenStats) Normalize() { + switch { + case t.PromptCacheHitTokens > 0: + t.CachedTokens = t.PromptCacheHitTokens + case t.PromptTokensDetails.CachedTokens > 0: + t.CachedTokens = t.PromptTokensDetails.CachedTokens + case t.CacheReadTokens > 0: + t.CachedTokens = t.CacheReadTokens + } } type LLMResponse struct { @@ -215,6 +237,22 @@ var SupportedProviders = map[string]ProviderConfig{ DefaultModel: "mistral-large-latest", RequiresKey: true, }, + "zhipu": { + ID: "zhipu", + Name: "Zhipu GLM", + BaseURL: "https://open.bigmodel.cn/api/paas/v4", + EnvKey: "ZHIPU_API_KEY", + DefaultModel: "glm-4.6", + RequiresKey: true, + }, + "glm": { + ID: "glm", + Name: "GLM (Zhipu)", + BaseURL: "https://open.bigmodel.cn/api/paas/v4", + EnvKey: "GLM_API_KEY", + DefaultModel: "glm-4.6", + RequiresKey: true, + }, } type Profile struct { @@ -926,6 +964,16 @@ func (c *LLMClient) ChatWithProviderStream( "messages": reqMessages, "stream": isStreaming, } + if isStreaming && supportsIncludeUsage(targetProviderID, baseURL) { + reqBody["stream_options"] = map[string]interface{}{"include_usage": true} + } + // OpenAI routes requests with the same prompt_cache_key to the same + // machine, improving automatic prefix-cache hits. + if targetProviderID == "openai" { + if sid := SessionIDFrom(ctx); sid != "" { + reqBody["prompt_cache_key"] = sid + } + } if len(tools) > 0 { reqBody["tools"] = reqTools @@ -1175,10 +1223,8 @@ func buildLLMResponse(openAIResp openAICompletionResp, body []byte) (*LLMRespons } choice := openAIResp.Choices[0] - cached := openAIResp.Usage.CachedTokens - if openAIResp.Usage.PromptCacheHitTokens > 0 { - cached = openAIResp.Usage.PromptCacheHitTokens - } + usage := openAIResp.Usage + usage.Normalize() reasoning := choice.Message.ReasoningContent if reasoning == "" { @@ -1196,12 +1242,15 @@ func buildLLMResponse(openAIResp openAICompletionResp, body []byte) (*LLMRespons Reasoning: reasoning, ToolCalls: choice.Message.ToolCalls, TokenUsage: TokenStats{ - PromptTokens: openAIResp.Usage.PromptTokens, - CompletionTokens: openAIResp.Usage.CompletionTokens, - CachedTokens: cached, - PromptCacheHitTokens: openAIResp.Usage.PromptCacheHitTokens, - PromptCacheMissTokens: openAIResp.Usage.PromptCacheMissTokens, - TotalTokens: openAIResp.Usage.TotalTokens, + PromptTokens: usage.PromptTokens, + CompletionTokens: usage.CompletionTokens, + CachedTokens: usage.CachedTokens, + PromptCacheHitTokens: usage.PromptCacheHitTokens, + PromptCacheMissTokens: usage.PromptCacheMissTokens, + CacheReadTokens: usage.CacheReadTokens, + CacheWriteTokens: usage.CacheWriteTokens, + PromptTokensDetails: usage.PromptTokensDetails, + TotalTokens: usage.TotalTokens, }, }, nil } @@ -1278,19 +1327,29 @@ func parseSSEStream(ctx context.Context, r io.Reader, onChunk func(deltaContent CachedTokens int `json:"cached_tokens"` // Anthropic-style: cached input subset PromptCacheHitTokens int `json:"prompt_cache_hit_tokens"` // DeepSeek PromptCacheMissTokens int `json:"prompt_cache_miss_tokens"` // DeepSeek - TotalTokens int `json:"total_tokens"` + CacheReadTokens int `json:"cache_read_input_tokens"` // Anthropic native + CacheWriteTokens int `json:"cache_creation_input_tokens"` + PromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` + } `json:"prompt_tokens_details"` + TotalTokens int `json:"total_tokens"` } `json:"usage"` } if err := json.Unmarshal([]byte(data), &chunk); err == nil { if chunk.Usage != nil { - if chunk.Usage.TotalTokens > 0 { - usage.PromptTokens = chunk.Usage.PromptTokens - usage.CompletionTokens = chunk.Usage.CompletionTokens - usage.CachedTokens = chunk.Usage.CachedTokens - usage.PromptCacheHitTokens = chunk.Usage.PromptCacheHitTokens - usage.PromptCacheMissTokens = chunk.Usage.PromptCacheMissTokens - usage.TotalTokens = chunk.Usage.TotalTokens + u := chunk.Usage + if u.TotalTokens > 0 || u.PromptTokens > 0 || u.CachedTokens > 0 || u.PromptCacheHitTokens > 0 || u.CacheReadTokens > 0 { + usage.PromptTokens = u.PromptTokens + usage.CompletionTokens = u.CompletionTokens + usage.CachedTokens = u.CachedTokens + usage.PromptCacheHitTokens = u.PromptCacheHitTokens + usage.PromptCacheMissTokens = u.PromptCacheMissTokens + usage.CacheReadTokens = u.CacheReadTokens + usage.CacheWriteTokens = u.CacheWriteTokens + usage.PromptTokensDetails.CachedTokens = u.PromptTokensDetails.CachedTokens + usage.TotalTokens = u.TotalTokens + usage.Normalize() } } if len(chunk.Choices) > 0 { @@ -1450,49 +1509,37 @@ func PreparePayloadWithCache(providerID, model string, messages []LLMMessage, to msgMaps[i] = mm } - // Breakpoint 2 & 3: Check system message(s) - // If the system message is a single block or multiple, add cache_control to system content blocks - for _, m := range msgMaps { - if m["role"] == "system" && remainingBudget > 0 { - contentStr, isStr := m["content"].(string) - if isStr && contentStr != "" { - // Convert to Anthropic / OpenRouter structured content parts - // e.g. [{"type": "text", "text": "...", "cache_control": {"type": "ephemeral"}}] - m["content"] = []map[string]any{ - { - "type": "text", - "text": contentStr, - "cache_control": map[string]string{"type": "ephemeral"}, - }, - } + // Breakpoint 2: mark the LAST system message. With the stable/volatile + // split this is the cacheable instruction block; the volatile runtime + // context is a trailing user reminder and does not invalidate it. + lastSystem := -1 + for i, m := range msgMaps { + if m["role"] == "system" { + lastSystem = i + } + } + if lastSystem >= 0 && remainingBudget > 0 { + markContent(msgMaps[lastSystem]) + remainingBudget-- + } + + // Breakpoint 3: end of the previous turn (the message before the tail). + // Anchoring here lets the next step reuse the whole completed turn. + if remainingBudget > 0 { + for i := len(msgMaps) - 2; i >= 0; i-- { + if markContent(msgMaps[i]) { remainingBudget-- + break } - break } } - // Breakpoint 4: Conversation tail + // Breakpoint 4: conversation tail // Mark the last user message or tool observation so subsequent turns in the loop hit prefix cache if remainingBudget > 0 && len(msgMaps) > 0 { for i := len(msgMaps) - 1; i >= 0; i-- { - role := msgMaps[i]["role"] - if role == "user" || role == "tool" { - contentVal := msgMaps[i]["content"] - if str, ok := contentVal.(string); ok && str != "" { - msgMaps[i]["content"] = []map[string]any{ - { - "type": "text", - "text": str, - "cache_control": map[string]string{"type": "ephemeral"}, - }, - } - remainingBudget-- - } else if parts, ok := contentVal.([]map[string]any); ok && len(parts) > 0 { - // Mark the last part of this message - lastIdx := len(parts) - 1 - parts[lastIdx]["cache_control"] = map[string]string{"type": "ephemeral"} - remainingBudget-- - } + if markContent(msgMaps[i]) { + remainingBudget-- break } } @@ -1501,6 +1548,30 @@ func PreparePayloadWithCache(providerID, model string, messages []LLMMessage, to return msgMaps, finalTools } +// markContent applies an Anthropic ephemeral cache breakpoint to a message's +// content. Returns false when there is nothing cacheable. +func markContent(m map[string]any) bool { + switch content := m["content"].(type) { + case string: + if content == "" { + return false + } + m["content"] = []map[string]any{{ + "type": "text", + "text": content, + "cache_control": map[string]string{"type": "ephemeral"}, + }} + return true + case []map[string]any: + if len(content) == 0 { + return false + } + content[len(content)-1]["cache_control"] = map[string]string{"type": "ephemeral"} + return true + } + return false +} + func shouldApplyCacheBreakpoints(providerID, model string) bool { p := strings.ToLower(providerID) m := strings.ToLower(model) @@ -1554,6 +1625,23 @@ func isRetryableStatus(code int) bool { return code >= 500 && code <= 599 } +// supportsIncludeUsage reports whether the provider accepts OpenAI's +// stream_options.include_usage. Whitelisted rather than assumed: strict +// gateways reject unknown request fields. +func supportsIncludeUsage(providerID, baseURL string) bool { + p := strings.ToLower(providerID) + u := strings.ToLower(baseURL) + switch p { + case "openai", "deepseek", "zhipu", "glm", "groq", "mistral", "openrouter", "xai", "ollama": + return true + } + // Custom OpenAI-compatible gateways generally accept it. + if strings.Contains(p, "openai") || strings.Contains(u, "openai") || strings.Contains(u, "deepseek") || strings.Contains(u, "bigmodel") { + return true + } + return false +} + func retryAfterDelay(header string) time.Duration { header = strings.TrimSpace(header) if header == "" {