From fa3820070ac0cb7b9932cdd5cbf1137886348f4e Mon Sep 17 00:00:00 2001 From: pdfinn Date: Fri, 30 Jan 2026 16:09:18 +0700 Subject: [PATCH] feat(llm): Add extended thinking support and usage tracking - Add thinking token control via /n/llm/thinking file (max/off/number) - CLI backend sets MAX_THINKING_TOKENS env var for Claude CLI - Default to max thinking (31999 tokens) for CLI backend - Add /n/llm/usage file for token usage monitoring - Add /n/llm/compact file for conversation summarization - Extend Backend interface with ThinkingTokens, TotalTokens, ContextLimit, Compact - Add true streaming support for CLI backend with line-by-line output - Update example file with thinking documentation Co-Authored-By: Claude Opus 4.5 --- internal/llm/backend.go | 11 ++ internal/llm/cli_client.go | 232 +++++++++++++++++++++++----- internal/llm/client.go | 129 ++++++++++++++-- internal/llm/client_test.go | 227 +++++++++++++++++++++++++++ internal/llmfs/ask.go | 25 ++- internal/llmfs/ask_test.go | 221 ++++++++++++++++++++++++++ internal/llmfs/compact.go | 73 +++++++++ internal/llmfs/compact_test.go | 181 ++++++++++++++++++++++ internal/llmfs/example.go | 9 ++ internal/llmfs/mock_backend_test.go | 110 +++++++++++++ internal/llmfs/root.go | 3 + internal/llmfs/thinking.go | 86 +++++++++++ internal/llmfs/usage.go | 49 ++++++ internal/llmfs/usage_test.go | 119 ++++++++++++++ 14 files changed, 1427 insertions(+), 48 deletions(-) create mode 100644 internal/llm/client_test.go create mode 100644 internal/llmfs/ask_test.go create mode 100644 internal/llmfs/compact.go create mode 100644 internal/llmfs/compact_test.go create mode 100644 internal/llmfs/mock_backend_test.go create mode 100644 internal/llmfs/thinking.go create mode 100644 internal/llmfs/usage.go create mode 100644 internal/llmfs/usage_test.go diff --git a/internal/llm/backend.go b/internal/llm/backend.go index b324f67df2897ab53c906abb124ec4ea9207a9c6..8fa0715af7269532fc2a267880601447be5c64ac 100644 --- a/internal/llm/backend.go +++ b/internal/llm/backend.go @@ -18,8 +18,19 @@ type Backend interface { SystemPrompt() string // SetSystemPrompt sets the system prompt SetSystemPrompt(prompt string) + // ThinkingTokens returns the thinking token budget (-1=max, 0=disabled, >0=budget) + ThinkingTokens() int + // SetThinkingTokens sets the thinking token budget + SetThinkingTokens(tokens int) // LastTokens returns token count from last response LastTokens() int + // TotalTokens returns cumulative token count for this conversation + TotalTokens() int + // ContextLimit returns the model's context window limit + ContextLimit() int + // Compact summarizes the conversation to reduce token usage + // The conversation history is replaced with a summary + Compact(ctx context.Context) error // Messages returns conversation history Messages() []Message // MessagesJSON returns conversation history as JSON diff --git a/internal/llm/cli_client.go b/internal/llm/cli_client.go index 1d4b499b44bc208aee5fff9d3cd6ab4a3303c81b..e23ce9fc863bd0fb34fbcfc41db133ec9f3f3e0d 100644 --- a/internal/llm/cli_client.go +++ b/internal/llm/cli_client.go @@ -15,15 +15,17 @@ import ( // CLIClient uses the Claude Code CLI for LLM requests. // This allows using a Claude Max subscription instead of API tokens. type CLIClient struct { - mu sync.RWMutex - model string - temperature float64 - systemPrompt string - messages []Message - lastTokens int - streaming bool - streamChan chan string - streamDone chan struct{} + mu sync.RWMutex + model string + temperature float64 + systemPrompt string + messages []Message + lastTokens int + totalTokens int // cumulative estimated token count + thinkingTokens int // 0 = disabled, >0 = budget, -1 = max (default) + streaming bool + streamChan chan string + streamDone chan struct{} } // cliResponse represents the JSON response from claude CLI @@ -35,9 +37,10 @@ type cliResponse struct { // NewCLIClient creates a new CLI-based LLM client func NewCLIClient() *CLIClient { return &CLIClient{ - model: "sonnet", // CLI uses short model names - temperature: 0.7, - messages: make([]Message, 0), + model: "sonnet", // CLI uses short model names + temperature: 0.7, + messages: make([]Message, 0), + thinkingTokens: -1, // -1 = max thinking (31999 tokens) enabled by default } } @@ -86,6 +89,22 @@ func (c *CLIClient) SetTemperature(temp float64) error { return nil } +// ThinkingTokens returns the current thinking token budget +// -1 = max (31999), 0 = disabled, >0 = specific budget +func (c *CLIClient) ThinkingTokens() int { + c.mu.RLock() + defer c.mu.RUnlock() + return c.thinkingTokens +} + +// SetThinkingTokens sets the thinking token budget +// -1 = max (31999), 0 = disabled, >0 = specific budget +func (c *CLIClient) SetThinkingTokens(tokens int) { + c.mu.Lock() + defer c.mu.Unlock() + c.thinkingTokens = tokens +} + // SystemPrompt returns the current system prompt func (c *CLIClient) SystemPrompt() string { c.mu.RLock() @@ -137,6 +156,96 @@ func (c *CLIClient) Reset() { defer c.mu.Unlock() c.messages = make([]Message, 0) c.lastTokens = 0 + c.totalTokens = 0 +} + +// TotalTokens returns cumulative estimated token count for this conversation +func (c *CLIClient) TotalTokens() int { + c.mu.RLock() + defer c.mu.RUnlock() + return c.totalTokens +} + +// ContextLimit returns the model's context window limit +func (c *CLIClient) ContextLimit() int { + c.mu.RLock() + model := c.model + c.mu.RUnlock() + return contextLimitForModel(model) +} + +// Compact summarizes the conversation to reduce token usage +func (c *CLIClient) Compact(ctx context.Context) error { + c.mu.Lock() + if len(c.messages) < 4 { + c.mu.Unlock() + return nil // Not enough to compact + } + + // Build conversation text for summarization + var conversationText string + for _, msg := range c.messages { + if msg.Role == "system" { + continue // Don't include system messages in summary + } + conversationText += fmt.Sprintf("%s: %s\n\n", msg.Role, msg.Content) + } + + model := c.model + thinkingTokens := c.thinkingTokens + c.mu.Unlock() + + // Use a compact summarization prompt + summaryPrompt := "Summarize this conversation concisely, preserving key facts, decisions, and context needed to continue:\n\n" + conversationText + + // Build CLI command for summarization + args := []string{ + "--print", + "--output-format", "json", + "--model", model, + "--allowedTools", "", + "--dangerously-skip-permissions", + "-", + } + + cmd := exec.CommandContext(ctx, "claude", args...) + cmd.Stdin = bytes.NewBufferString(summaryPrompt) + + // Set thinking token budget + cmd.Env = append(cmd.Environ(), func() string { + if thinkingTokens < 0 { + return "MAX_THINKING_TOKENS=31999" + } + return fmt.Sprintf("MAX_THINKING_TOKENS=%d", thinkingTokens) + }()) + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + if err := cmd.Run(); err != nil { + return fmt.Errorf("compaction failed: %w (stderr: %s)", err, stderr.String()) + } + + summary, err := parseJSONResponse(stdout.String()) + if err != nil { + return fmt.Errorf("compaction parse failed: %w", err) + } + + // Replace conversation with summary + c.mu.Lock() + c.messages = []Message{{Role: "system", Content: "Previous conversation summary: " + summary}} + // Estimate tokens for the new conversation state (chars * 0.25) + c.totalTokens = len(summary) / 4 + c.mu.Unlock() + + return nil +} + +// estimateTokens estimates token count from character count +// Uses rough approximation of 4 chars per token +func estimateTokens(s string) int { + return (len(s) + 3) / 4 // Round up } // buildPrompt builds a full prompt string from conversation history @@ -176,21 +285,21 @@ func (c *CLIClient) Ask(ctx context.Context, prompt string) (string, error) { fullPrompt := c.buildPrompt() systemPrompt := c.getSystemPrompt() model := c.model + thinkingTokens := c.thinkingTokens c.mu.Unlock() // Build claude CLI command. // --print: non-interactive mode, output to stdout // --output-format json: structured output we can parse // --allowedTools "": disable all tools (text-only, no Bash/Edit/etc.) - // - // Note: --dangerously-skip-permissions is NOT needed when tools are disabled - // and --print mode is used. The CLI only prompts for permission when tools - // might take actions. With tools disabled, it's purely text-in/text-out. + // --dangerously-skip-permissions: prevents macOS permission dialogs from blocking + // (Photo Library, Audio, etc. that Claude CLI initializes even with tools disabled) args := []string{ "--print", "--output-format", "json", "--model", model, "--allowedTools", "", + "--dangerously-skip-permissions", } if systemPrompt != "" { @@ -202,6 +311,15 @@ func (c *CLIClient) Ask(ctx context.Context, prompt string) (string, error) { cmd := exec.CommandContext(ctx, "claude", args...) cmd.Stdin = bytes.NewBufferString(fullPrompt) + // Set thinking token budget via environment variable + // -1 = max (31999), 0 = disabled, >0 = specific budget + cmd.Env = append(cmd.Environ(), func() string { + if thinkingTokens < 0 { + return "MAX_THINKING_TOKENS=31999" + } + return fmt.Sprintf("MAX_THINKING_TOKENS=%d", thinkingTokens) + }()) + var stdout, stderr bytes.Buffer cmd.Stdout = &stdout cmd.Stderr = &stderr @@ -231,7 +349,9 @@ func (c *CLIClient) Ask(ctx context.Context, prompt string) (string, error) { // Update state c.mu.Lock() c.messages = append(c.messages, Message{Role: "assistant", Content: responseText}) - c.lastTokens = 0 // CLI doesn't provide token counts + // Estimate tokens: prompt + response (chars / 4) + c.lastTokens = estimateTokens(fullPrompt) + estimateTokens(responseText) + c.totalTokens += c.lastTokens c.mu.Unlock() return responseText, nil @@ -267,7 +387,7 @@ func parseJSONResponse(output string) (string, error) { } // StartStream begins streaming a response for the given prompt -// Note: CLI streaming is simulated - we run the command and feed output progressively +// Uses text output mode and reads stdout progressively for real streaming func (c *CLIClient) StartStream(ctx context.Context, prompt string) error { c.mu.Lock() if c.streaming { @@ -279,6 +399,7 @@ func (c *CLIClient) StartStream(ctx context.Context, prompt string) error { fullPrompt := c.buildPrompt() systemPrompt := c.getSystemPrompt() model := c.model + thinkingTokens := c.thinkingTokens c.streaming = true c.streamChan = make(chan string, 100) @@ -286,20 +407,30 @@ func (c *CLIClient) StartStream(ctx context.Context, prompt string) error { c.mu.Unlock() go func() { + var fullResponse string + defer func() { + // Update conversation history with full response c.mu.Lock() + if fullResponse != "" { + c.messages = append(c.messages, Message{Role: "assistant", Content: fullResponse}) + c.lastTokens = estimateTokens(fullPrompt) + estimateTokens(fullResponse) + c.totalTokens += c.lastTokens + } c.streaming = false close(c.streamChan) close(c.streamDone) c.mu.Unlock() }() - // Build command (same flags as Ask - see comments there) + // Build command for streaming - use text output, not JSON + // --output-format text gives us raw text we can stream args := []string{ "--print", - "--output-format", "json", + "--output-format", "text", "--model", model, "--allowedTools", "", + "--dangerously-skip-permissions", } if systemPrompt != "" { @@ -311,11 +442,17 @@ func (c *CLIClient) StartStream(ctx context.Context, prompt string) error { cmd := exec.CommandContext(ctx, "claude", args...) cmd.Stdin = bytes.NewBufferString(fullPrompt) - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout - cmd.Stderr = &stderr + // Set thinking token budget via environment variable + cmd.Env = append(cmd.Environ(), func() string { + if thinkingTokens < 0 { + return "MAX_THINKING_TOKENS=31999" + } + return fmt.Sprintf("MAX_THINKING_TOKENS=%d", thinkingTokens) + }()) - if err := cmd.Run(); err != nil { + // Get stdout pipe for streaming reads + stdout, err := cmd.StdoutPipe() + if err != nil { select { case c.streamChan <- fmt.Sprintf("[Error: %v]", err): case <-ctx.Done(): @@ -328,9 +465,8 @@ func (c *CLIClient) StartStream(ctx context.Context, prompt string) error { return } - // Parse and send response - responseText, err := parseJSONResponse(stdout.String()) - if err != nil { + // Start the command + if err := cmd.Start(); err != nil { select { case c.streamChan <- fmt.Sprintf("[Error: %v]", err): case <-ctx.Done(): @@ -343,18 +479,40 @@ func (c *CLIClient) StartStream(ctx context.Context, prompt string) error { return } - // Send response as a single chunk (CLI doesn't truly stream) - select { - case c.streamChan <- responseText: - case <-ctx.Done(): - return + // Read stdout in chunks and send to channel + buf := make([]byte, 256) // Small buffer for responsive streaming + for { + n, err := stdout.Read(buf) + if n > 0 { + chunk := string(buf[:n]) + fullResponse += chunk + select { + case c.streamChan <- chunk: + case <-ctx.Done(): + cmd.Process.Kill() + return + } + } + if err != nil { + break // EOF or error + } } - // Update state - c.mu.Lock() - c.messages = append(c.messages, Message{Role: "assistant", Content: responseText}) - c.lastTokens = 0 - c.mu.Unlock() + // Wait for command to finish + if err := cmd.Wait(); err != nil { + // Only report error if we got no response + if fullResponse == "" { + select { + case c.streamChan <- fmt.Sprintf("[Error: %v]", err): + case <-ctx.Done(): + } + c.mu.Lock() + if len(c.messages) > 0 { + c.messages = c.messages[:len(c.messages)-1] + } + c.mu.Unlock() + } + } }() return nil diff --git a/internal/llm/client.go b/internal/llm/client.go index 7587eb6e1ce7c2dbc17a0120f2d8448d014f9545..28d29e67af268f78e50bae5711bf60ed590fce4e 100644 --- a/internal/llm/client.go +++ b/internal/llm/client.go @@ -5,6 +5,7 @@ import ( "context" "encoding/json" "fmt" + "strings" "sync" "github.com/anthropics/anthropic-sdk-go" @@ -19,16 +20,18 @@ type Message struct { // Client wraps the Anthropic API client with conversation state type Client struct { - client anthropic.Client - mu sync.RWMutex - model string - temperature float64 - systemPrompt string - messages []Message - lastTokens int - streaming bool - streamChan chan string - streamDone chan struct{} + client anthropic.Client + mu sync.RWMutex + model string + temperature float64 + systemPrompt string + messages []Message + lastTokens int + totalTokens int // cumulative token count for context tracking + thinkingTokens int // 0 = disabled, >0 = budget, -1 = max (default for CLI, not used for API yet) + streaming bool + streamChan chan string + streamDone chan struct{} } // NewClient creates a new LLM client @@ -74,6 +77,22 @@ func (c *Client) SetTemperature(temp float64) error { return nil } +// ThinkingTokens returns the current thinking token budget +// Note: API backend does not currently use extended thinking +func (c *Client) ThinkingTokens() int { + c.mu.RLock() + defer c.mu.RUnlock() + return c.thinkingTokens +} + +// SetThinkingTokens sets the thinking token budget +// Note: API backend does not currently use extended thinking +func (c *Client) SetThinkingTokens(tokens int) { + c.mu.Lock() + defer c.mu.Unlock() + c.thinkingTokens = tokens +} + // SystemPrompt returns the current system prompt func (c *Client) SystemPrompt() string { c.mu.RLock() @@ -125,6 +144,94 @@ func (c *Client) Reset() { defer c.mu.Unlock() c.messages = make([]Message, 0) c.lastTokens = 0 + c.totalTokens = 0 +} + +// TotalTokens returns cumulative token count for this conversation +func (c *Client) TotalTokens() int { + c.mu.RLock() + defer c.mu.RUnlock() + return c.totalTokens +} + +// ContextLimit returns the model's context window limit +func (c *Client) ContextLimit() int { + c.mu.RLock() + model := c.model + c.mu.RUnlock() + return contextLimitForModel(model) +} + +// Compact summarizes the conversation to reduce token usage +func (c *Client) Compact(ctx context.Context) error { + c.mu.Lock() + if len(c.messages) < 4 { + c.mu.Unlock() + return nil // Not enough to compact + } + + // Build conversation text for summarization + var conversationText string + for _, msg := range c.messages { + if msg.Role == "system" { + continue // Don't include system messages in summary + } + conversationText += fmt.Sprintf("%s: %s\n\n", msg.Role, msg.Content) + } + + model := c.model + c.mu.Unlock() + + // Use a compact summarization prompt + summaryPrompt := "Summarize this conversation concisely, preserving key facts, decisions, and context needed to continue:\n\n" + conversationText + + // Build API request for summarization + apiMessages := []anthropic.MessageParam{ + anthropic.NewUserMessage(anthropic.NewTextBlock(summaryPrompt)), + } + + params := anthropic.MessageNewParams{ + Model: anthropic.Model(model), + MaxTokens: 2048, + Messages: apiMessages, + } + + response, err := c.client.Messages.New(ctx, params) + if err != nil { + return fmt.Errorf("compaction failed: %w", err) + } + + // Extract summary + var summary string + for _, block := range response.Content { + if block.Type == "text" { + summary += block.Text + } + } + + // Replace conversation with summary + c.mu.Lock() + c.messages = []Message{{Role: "system", Content: "Previous conversation summary: " + summary}} + c.totalTokens = int(response.Usage.InputTokens + response.Usage.OutputTokens) + c.mu.Unlock() + + return nil +} + +// contextLimitForModel returns the context window size for a model +func contextLimitForModel(model string) int { + model = strings.ToLower(model) + // Claude models and their context limits + switch { + case strings.Contains(model, "opus"): + return 200000 + case strings.Contains(model, "sonnet"): + return 200000 + case strings.Contains(model, "haiku"): + return 200000 + default: + return 200000 // Default to 200K for newer Claude models + } } // Ask sends a prompt to the LLM and returns the response @@ -203,6 +310,7 @@ func (c *Client) Ask(ctx context.Context, prompt string) (string, error) { c.mu.Lock() c.messages = append(c.messages, Message{Role: "assistant", Content: responseText}) c.lastTokens = int(response.Usage.InputTokens + response.Usage.OutputTokens) + c.totalTokens += c.lastTokens c.mu.Unlock() return responseText, nil @@ -325,6 +433,7 @@ func (c *Client) StartStream(ctx context.Context, prompt string) error { c.mu.Lock() c.messages = append(c.messages, Message{Role: "assistant", Content: fullResponse}) c.lastTokens = int(inputTokens + outputTokens) + c.totalTokens += c.lastTokens c.mu.Unlock() }() diff --git a/internal/llm/client_test.go b/internal/llm/client_test.go new file mode 100644 index 0000000000000000000000000000000000000000..7324d24c0943e1f319655e43bf842d99ccaee86a --- /dev/null +++ b/internal/llm/client_test.go @@ -0,0 +1,227 @@ +package llm + +import ( + "testing" +) + +func TestContextLimitForModel(t *testing.T) { + tests := []struct { + model string + expected int + }{ + {"claude-3-opus-20240229", 200000}, + {"claude-3-sonnet-20240229", 200000}, + {"claude-3-haiku-20240307", 200000}, + {"claude-sonnet-4-20250514", 200000}, + {"CLAUDE-3-OPUS", 200000}, // case insensitive + {"some-sonnet-model", 200000}, // substring match + {"unknown-model", 200000}, // default + } + + for _, tc := range tests { + t.Run(tc.model, func(t *testing.T) { + got := contextLimitForModel(tc.model) + if got != tc.expected { + t.Errorf("contextLimitForModel(%q) = %d, want %d", tc.model, got, tc.expected) + } + }) + } +} + +func TestClientReset(t *testing.T) { + // Create a client with dummy API key (won't make real calls) + c := NewClient("dummy-key") + + // Manually set some state + c.messages = []Message{ + {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "hi"}, + } + c.lastTokens = 100 + c.totalTokens = 500 + + // Reset should clear everything + c.Reset() + + if len(c.messages) != 0 { + t.Errorf("Reset() should clear messages, got %d", len(c.messages)) + } + if c.lastTokens != 0 { + t.Errorf("Reset() should clear lastTokens, got %d", c.lastTokens) + } + if c.totalTokens != 0 { + t.Errorf("Reset() should clear totalTokens, got %d", c.totalTokens) + } +} + +func TestClientTotalTokens(t *testing.T) { + c := NewClient("dummy-key") + + // Initially should be 0 + if got := c.TotalTokens(); got != 0 { + t.Errorf("TotalTokens() = %d, want 0", got) + } + + // Manually set for testing + c.totalTokens = 12345 + if got := c.TotalTokens(); got != 12345 { + t.Errorf("TotalTokens() = %d, want 12345", got) + } +} + +func TestClientContextLimit(t *testing.T) { + c := NewClient("dummy-key") + + // Default model should have 200K limit + got := c.ContextLimit() + if got != 200000 { + t.Errorf("ContextLimit() = %d, want 200000", got) + } + + // Change model and verify + c.SetModel("claude-3-haiku-20240307") + got = c.ContextLimit() + if got != 200000 { + t.Errorf("ContextLimit() for haiku = %d, want 200000", got) + } +} + +func TestClientTemperature(t *testing.T) { + c := NewClient("dummy-key") + + // Default temperature + if got := c.Temperature(); got != 0.7 { + t.Errorf("Temperature() = %f, want 0.7", got) + } + + // Valid temperature + if err := c.SetTemperature(1.5); err != nil { + t.Errorf("SetTemperature(1.5) error: %v", err) + } + if got := c.Temperature(); got != 1.5 { + t.Errorf("Temperature() = %f, want 1.5", got) + } + + // Invalid temperature - too low + if err := c.SetTemperature(-0.1); err == nil { + t.Error("SetTemperature(-0.1) should return error") + } + + // Invalid temperature - too high + if err := c.SetTemperature(2.1); err == nil { + t.Error("SetTemperature(2.1) should return error") + } +} + +func TestClientSystemPrompt(t *testing.T) { + c := NewClient("dummy-key") + + // Initially empty + if got := c.SystemPrompt(); got != "" { + t.Errorf("SystemPrompt() = %q, want empty", got) + } + + // Set and verify + c.SetSystemPrompt("You are a helpful assistant") + if got := c.SystemPrompt(); got != "You are a helpful assistant" { + t.Errorf("SystemPrompt() = %q, want 'You are a helpful assistant'", got) + } +} + +func TestClientMessages(t *testing.T) { + c := NewClient("dummy-key") + + // Initially empty + if msgs := c.Messages(); len(msgs) != 0 { + t.Errorf("Messages() should be empty, got %d", len(msgs)) + } + + // Add messages + c.messages = []Message{ + {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "hi"}, + } + + msgs := c.Messages() + if len(msgs) != 2 { + t.Errorf("Messages() = %d messages, want 2", len(msgs)) + } + + // Verify it's a copy (modification doesn't affect original) + msgs[0].Content = "modified" + if c.messages[0].Content == "modified" { + t.Error("Messages() should return a copy, not the original") + } +} + +func TestClientAddSystemMessage(t *testing.T) { + c := NewClient("dummy-key") + + // Add system message + c.AddSystemMessage("Context info here") + + msgs := c.Messages() + if len(msgs) != 1 { + t.Fatalf("Messages() = %d messages, want 1", len(msgs)) + } + + if msgs[0].Role != "system" { + t.Errorf("Message role = %q, want 'system'", msgs[0].Role) + } + if msgs[0].Content != "Context info here" { + t.Errorf("Message content = %q, want 'Context info here'", msgs[0].Content) + } +} + +func TestClientModel(t *testing.T) { + c := NewClient("dummy-key") + + // Default model + if got := c.Model(); got != "claude-sonnet-4-20250514" { + t.Errorf("Model() = %q, want 'claude-sonnet-4-20250514'", got) + } + + // Change model + c.SetModel("claude-3-haiku-20240307") + if got := c.Model(); got != "claude-3-haiku-20240307" { + t.Errorf("Model() = %q, want 'claude-3-haiku-20240307'", got) + } +} + +func TestClientMessagesJSON(t *testing.T) { + c := NewClient("dummy-key") + + c.messages = []Message{ + {Role: "user", Content: "hello"}, + } + + data, err := c.MessagesJSON() + if err != nil { + t.Fatalf("MessagesJSON() error: %v", err) + } + + // Should contain the message content + json := string(data) + if !contains(json, "hello") || !contains(json, "user") { + t.Errorf("MessagesJSON() = %s, should contain 'hello' and 'user'", json) + } +} + +func TestClientIsStreaming(t *testing.T) { + c := NewClient("dummy-key") + + // Initially not streaming + if c.IsStreaming() { + t.Error("IsStreaming() should be false initially") + } +} + +// Helper function +func contains(s, substr string) bool { + for i := 0; i <= len(s)-len(substr); i++ { + if s[i:i+len(substr)] == substr { + return true + } + } + return false +} diff --git a/internal/llmfs/ask.go b/internal/llmfs/ask.go index f598e274c8786b2480ae9c22fe42fb954bb23a12..cc3be1b2487ea007975b3938a1f6304b217c8cd5 100644 --- a/internal/llmfs/ask.go +++ b/internal/llmfs/ask.go @@ -3,6 +3,7 @@ package llmfs import ( "context" "io" + "log" "strings" "sync" @@ -10,6 +11,9 @@ import ( "github.com/NERVsystems/llm9p/internal/protocol" ) +// CompactThreshold is the percentage of context limit at which auto-compaction triggers +const CompactThreshold = 0.80 + // AskFile is the main interaction file - write a prompt, read the response type AskFile struct { *protocol.BaseFile @@ -49,8 +53,27 @@ func (f *AskFile) Write(p []byte, offset int64) (int, error) { return len(p), nil // Empty write is a no-op } + ctx := context.Background() + + // Check if we need to auto-compact before processing + tokens := f.client.TotalTokens() + limit := f.client.ContextLimit() + threshold := int(float64(limit) * CompactThreshold) + + if tokens > threshold { + log.Printf("llm9p: auto-compacting at %d/%d tokens (%.0f%% threshold)", + tokens, limit, CompactThreshold*100) + if err := f.client.Compact(ctx); err != nil { + log.Printf("llm9p: auto-compact failed: %v", err) + // Continue anyway - better to try than to fail + } else { + log.Printf("llm9p: auto-compact complete, now at %d tokens", + f.client.TotalTokens()) + } + } + // Make the API call - response, err := f.client.Ask(context.Background(), prompt) + response, err := f.client.Ask(ctx, prompt) if err != nil { // Store error as response so it can be read f.mu.Lock() diff --git a/internal/llmfs/ask_test.go b/internal/llmfs/ask_test.go new file mode 100644 index 0000000000000000000000000000000000000000..2df07986ed4f64621969a6c948e7aac1eaf61ad0 --- /dev/null +++ b/internal/llmfs/ask_test.go @@ -0,0 +1,221 @@ +package llmfs + +import ( + "io" + "testing" +) + +func TestAskFile_Read_Empty(t *testing.T) { + mock := NewMockBackend() + ask := NewAskFile(mock) + + // Initial read should be empty (EOF) + buf := make([]byte, 100) + n, err := ask.Read(buf, 0) + if err != io.EOF { + t.Errorf("Read() error = %v, want io.EOF", err) + } + if n != 0 { + t.Errorf("Read() n = %d, want 0", n) + } +} + +func TestAskFile_WriteRead(t *testing.T) { + mock := NewMockBackend() + mock.askResponse = "Hello, I'm Claude!" + + ask := NewAskFile(mock) + + // Write a prompt + prompt := "Hello!" + n, err := ask.Write([]byte(prompt), 0) + if err != nil { + t.Fatalf("Write() error: %v", err) + } + if n != len(prompt) { + t.Errorf("Write() n = %d, want %d", n, len(prompt)) + } + + // Read response + buf := make([]byte, 100) + readN, err := ask.Read(buf, 0) + if err != nil { + t.Fatalf("Read() error: %v", err) + } + + response := string(buf[:readN]) + expected := "Hello, I'm Claude!\n" + if response != expected { + t.Errorf("Read() = %q, want %q", response, expected) + } +} + +func TestAskFile_Write_EmptyNoOp(t *testing.T) { + mock := NewMockBackend() + ask := NewAskFile(mock) + + // Empty write should be no-op + n, err := ask.Write([]byte(""), 0) + if err != nil { + t.Fatalf("Write() error: %v", err) + } + if n != 0 { + t.Errorf("Write('') n = %d, want 0", n) + } + + // Whitespace-only also no-op + n, err = ask.Write([]byte(" \n\t "), 0) + if err != nil { + t.Fatalf("Write() error: %v", err) + } +} + +func TestAskFile_Write_Error(t *testing.T) { + mock := NewMockBackend() + mock.askError = io.ErrUnexpectedEOF + + ask := NewAskFile(mock) + + // Write should succeed (error is stored for read) + n, err := ask.Write([]byte("test"), 0) + if err != nil { + t.Fatalf("Write() error: %v", err) + } + if n != 4 { + t.Errorf("Write() n = %d, want 4", n) + } + + // Read should return the error + buf := make([]byte, 100) + readN, _ := ask.Read(buf, 0) + response := string(buf[:readN]) + + if response[:6] != "Error:" { + t.Errorf("Read() = %q, should start with 'Error:'", response) + } +} + +func TestAskFile_AutoCompaction_BelowThreshold(t *testing.T) { + mock := NewMockBackend() + mock.totalTokens = 100000 // 50% of 200K, below 80% threshold + mock.contextLimit = 200000 + mock.askResponse = "response" + + ask := NewAskFile(mock) + + // Write should NOT trigger compaction + ask.Write([]byte("test"), 0) + + if mock.compactCalled { + t.Error("Compact() should not be called below threshold") + } +} + +func TestAskFile_AutoCompaction_AboveThreshold(t *testing.T) { + mock := NewMockBackend() + mock.totalTokens = 170000 // 85% of 200K, above 80% threshold + mock.contextLimit = 200000 + mock.askResponse = "response" + + ask := NewAskFile(mock) + + // Write SHOULD trigger compaction + ask.Write([]byte("test"), 0) + + if !mock.compactCalled { + t.Error("Compact() should be called when above threshold") + } +} + +func TestAskFile_AutoCompaction_ExactThreshold(t *testing.T) { + mock := NewMockBackend() + mock.totalTokens = 160000 // Exactly 80% of 200K + mock.contextLimit = 200000 + mock.askResponse = "response" + + ask := NewAskFile(mock) + + // Write should NOT trigger (we use > not >=) + ask.Write([]byte("test"), 0) + + if mock.compactCalled { + t.Error("Compact() should not be called at exactly threshold") + } +} + +func TestAskFile_AutoCompaction_JustAboveThreshold(t *testing.T) { + mock := NewMockBackend() + mock.totalTokens = 160001 // Just over 80% + mock.contextLimit = 200000 + mock.askResponse = "response" + + ask := NewAskFile(mock) + + // Write SHOULD trigger compaction + ask.Write([]byte("test"), 0) + + if !mock.compactCalled { + t.Error("Compact() should be called when just above threshold") + } +} + +func TestAskFile_Stat(t *testing.T) { + mock := NewMockBackend() + mock.askResponse = "Hello!" + + ask := NewAskFile(mock) + + // Initial stat - no response yet + stat := ask.Stat() + if stat.Length != 0 { + t.Errorf("Stat().Length = %d, want 0 (no response yet)", stat.Length) + } + + // Write to get a response + ask.Write([]byte("test"), 0) + + // Stat should now show response length (with newline) + stat = ask.Stat() + expected := uint64(len("Hello!\n")) + if stat.Length != expected { + t.Errorf("Stat().Length = %d, want %d", stat.Length, expected) + } +} + +func TestAskFile_ResponseNewline(t *testing.T) { + mock := NewMockBackend() + + // Response without trailing newline + mock.askResponse = "No newline" + ask := NewAskFile(mock) + ask.Write([]byte("test"), 0) + + buf := make([]byte, 100) + n, _ := ask.Read(buf, 0) + response := string(buf[:n]) + + // Should have newline added + if response[len(response)-1] != '\n' { + t.Errorf("Response should end with newline, got %q", response) + } + + // Response with trailing newline already + mock.askResponse = "Has newline\n" + ask2 := NewAskFile(mock) + ask2.Write([]byte("test"), 0) + + n, _ = ask2.Read(buf, 0) + response = string(buf[:n]) + + // Should NOT double the newline + if response != "Has newline\n" { + t.Errorf("Response = %q, should not double newline", response) + } +} + +func TestCompactThreshold(t *testing.T) { + // Verify the threshold constant + if CompactThreshold != 0.80 { + t.Errorf("CompactThreshold = %f, want 0.80", CompactThreshold) + } +} diff --git a/internal/llmfs/compact.go b/internal/llmfs/compact.go new file mode 100644 index 0000000000000000000000000000000000000000..45119b86742ab278cc5cb7f4c33f89547dc47af5 --- /dev/null +++ b/internal/llmfs/compact.go @@ -0,0 +1,73 @@ +package llmfs + +import ( + "context" + "fmt" + "io" + "strings" + "sync" + + "github.com/NERVsystems/llm9p/internal/llm" + "github.com/NERVsystems/llm9p/internal/protocol" +) + +// CompactFile allows manual compaction trigger +// Write anything to trigger compaction +// Read returns status ("ok" or "error: ...") +type CompactFile struct { + *protocol.BaseFile + client llm.Backend + mu sync.RWMutex + lastResult string +} + +// NewCompactFile creates the compact file +func NewCompactFile(client llm.Backend) *CompactFile { + return &CompactFile{ + BaseFile: protocol.NewBaseFile("compact", 0666), + client: client, + lastResult: "ready\n", + } +} + +func (f *CompactFile) Read(p []byte, offset int64) (int, error) { + f.mu.RLock() + content := f.lastResult + f.mu.RUnlock() + + if offset >= int64(len(content)) { + return 0, io.EOF + } + n := copy(p, content[offset:]) + return n, nil +} + +func (f *CompactFile) Write(p []byte, offset int64) (int, error) { + cmd := strings.TrimSpace(string(p)) + if cmd == "" { + return len(p), nil + } + + // Trigger compaction + err := f.client.Compact(context.Background()) + + f.mu.Lock() + if err != nil { + f.lastResult = fmt.Sprintf("error: %v\n", err) + } else { + tokens := f.client.TotalTokens() + limit := f.client.ContextLimit() + f.lastResult = fmt.Sprintf("ok: %d/%d\n", tokens, limit) + } + f.mu.Unlock() + + return len(p), nil +} + +func (f *CompactFile) Stat() protocol.Stat { + s := f.BaseFile.Stat() + f.mu.RLock() + s.Length = uint64(len(f.lastResult)) + f.mu.RUnlock() + return s +} diff --git a/internal/llmfs/compact_test.go b/internal/llmfs/compact_test.go new file mode 100644 index 0000000000000000000000000000000000000000..214a00a16f8b86961f856951e040327de949e101 --- /dev/null +++ b/internal/llmfs/compact_test.go @@ -0,0 +1,181 @@ +package llmfs + +import ( + "fmt" + "io" + "strings" + "testing" +) + +func TestCompactFile_Read_Initial(t *testing.T) { + mock := NewMockBackend() + compact := NewCompactFile(mock) + + // Initial read should return "ready" + buf := make([]byte, 100) + n, err := compact.Read(buf, 0) + if err != nil { + t.Fatalf("Read() error: %v", err) + } + + content := string(buf[:n]) + expected := "ready\n" + if content != expected { + t.Errorf("Read() = %q, want %q", content, expected) + } +} + +func TestCompactFile_Write_TriggerCompaction(t *testing.T) { + mock := NewMockBackend() + mock.totalTokens = 160000 + mock.contextLimit = 200000 + + compact := NewCompactFile(mock) + + // Write to trigger compaction + n, err := compact.Write([]byte("1"), 0) + if err != nil { + t.Fatalf("Write() error: %v", err) + } + if n != 1 { + t.Errorf("Write() n = %d, want 1", n) + } + + // Verify compaction was called + if !mock.compactCalled { + t.Error("Compact() was not called on backend") + } + + // Read should show "ok: tokens/limit" + buf := make([]byte, 100) + readN, _ := compact.Read(buf, 0) + content := string(buf[:readN]) + + if !strings.HasPrefix(content, "ok:") { + t.Errorf("Read() after compaction = %q, want prefix 'ok:'", content) + } +} + +func TestCompactFile_Write_EmptyNoOp(t *testing.T) { + mock := NewMockBackend() + compact := NewCompactFile(mock) + + // Empty write should be no-op + n, err := compact.Write([]byte(""), 0) + if err != nil { + t.Fatalf("Write() error: %v", err) + } + if n != 0 { + t.Errorf("Write('') n = %d, want 0", n) + } + + // Compaction should not be triggered + if mock.compactCalled { + t.Error("Compact() should not be called for empty write") + } +} + +func TestCompactFile_Write_WhitespaceNoOp(t *testing.T) { + mock := NewMockBackend() + compact := NewCompactFile(mock) + + // Whitespace-only write should be no-op + _, err := compact.Write([]byte(" \n\t "), 0) + if err != nil { + t.Fatalf("Write() error: %v", err) + } + + // Compaction should not be triggered + if mock.compactCalled { + t.Error("Compact() should not be called for whitespace-only write") + } +} + +func TestCompactFile_Write_Error(t *testing.T) { + mock := NewMockBackend() + mock.compactError = fmt.Errorf("compaction failed") + + compact := NewCompactFile(mock) + + // Write should still succeed (error is stored for read) + n, err := compact.Write([]byte("1"), 0) + if err != nil { + t.Fatalf("Write() error: %v", err) + } + if n != 1 { + t.Errorf("Write() n = %d, want 1", n) + } + + // Read should show error + buf := make([]byte, 100) + readN, _ := compact.Read(buf, 0) + content := string(buf[:readN]) + + if !strings.HasPrefix(content, "error:") { + t.Errorf("Read() after error = %q, want prefix 'error:'", content) + } + if !strings.Contains(content, "compaction failed") { + t.Errorf("Read() should contain error message, got %q", content) + } +} + +func TestCompactFile_Read_EOF(t *testing.T) { + mock := NewMockBackend() + compact := NewCompactFile(mock) + + // Read past end + buf := make([]byte, 100) + n, err := compact.Read(buf, 1000) + if err != io.EOF { + t.Errorf("Read(offset=1000) error = %v, want io.EOF", err) + } + if n != 0 { + t.Errorf("Read(offset=1000) n = %d, want 0", n) + } +} + +func TestCompactFile_Stat(t *testing.T) { + mock := NewMockBackend() + compact := NewCompactFile(mock) + + stat := compact.Stat() + + // Initial content is "ready\n" = 6 chars + expected := uint64(6) + if stat.Length != expected { + t.Errorf("Stat().Length = %d, want %d", stat.Length, expected) + } +} + +func TestCompactFile_MultipleCompactions(t *testing.T) { + mock := NewMockBackend() + mock.totalTokens = 180000 + mock.contextLimit = 200000 + + compact := NewCompactFile(mock) + + // First compaction + compact.Write([]byte("1"), 0) + if !mock.compactCalled { + t.Error("First compaction not called") + } + + // Reset flag + mock.compactCalled = false + mock.totalTokens = 100000 + + // Second compaction + compact.Write([]byte("1"), 0) + if !mock.compactCalled { + t.Error("Second compaction not called") + } + + // Check result reflects new token count + buf := make([]byte, 100) + readN, _ := compact.Read(buf, 0) + content := string(buf[:readN]) + + if !strings.Contains(content, "25000") { // 100000 / 4 from mock + t.Errorf("Read() = %q, should contain reduced token count", content) + } +} diff --git a/internal/llmfs/example.go b/internal/llmfs/example.go index 2b63bef037c92103912cab424ecbdab752f20fae..967787b9fc746b3d0dc55eec8720e67d1972edd0 100644 --- a/internal/llmfs/example.go +++ b/internal/llmfs/example.go @@ -26,6 +26,9 @@ Conversation Management: Token Usage: cat tokens # View tokens from last response + cat usage # View total/limit (e.g., "45000/200000") + echo "1" > compact # Manually trigger conversation compaction + cat compact # Check compaction status Streaming: echo "Tell me a story" > stream/ask # Start streaming request @@ -55,11 +58,17 @@ Files: temperature Read/write: sampling temperature (0.0-2.0) system Read/write: system prompt (persists across resets) tokens Read-only: token count from last response + usage Read-only: total tokens/limit (e.g., "45000/200000") + compact Read/write: write to trigger compaction, read for status new Write-only: any write resets conversation (keeps system prompt) context Read: JSON history; Write: add system message to history _example Read-only: this help text stream/ask Write-only: starts a streaming request stream/chunk Read-only: returns next chunk (blocks), EOF when done + +Auto-Compaction: + When tokens exceed 80% of context limit, the conversation is automatically + summarized before processing the next query. This is transparent to the client. ` // NewExampleFile creates the _example file with usage examples diff --git a/internal/llmfs/mock_backend_test.go b/internal/llmfs/mock_backend_test.go new file mode 100644 index 0000000000000000000000000000000000000000..a8b198db98a63e1f03ea257f77dd9643fbc42104 --- /dev/null +++ b/internal/llmfs/mock_backend_test.go @@ -0,0 +1,110 @@ +package llmfs + +import ( + "context" + "fmt" + + "github.com/NERVsystems/llm9p/internal/llm" +) + +// MockBackend implements llm.Backend for testing +type MockBackend struct { + model string + temperature float64 + systemPrompt string + messages []llm.Message + lastTokens int + totalTokens int + contextLimit int + thinkingTokens int + compactCalled bool + compactError error + askResponse string + askError error +} + +func NewMockBackend() *MockBackend { + return &MockBackend{ + model: "mock-model", + temperature: 0.7, + contextLimit: 200000, + messages: make([]llm.Message, 0), + } +} + +func (m *MockBackend) Model() string { return m.model } +func (m *MockBackend) SetModel(model string) { m.model = model } +func (m *MockBackend) Temperature() float64 { return m.temperature } +func (m *MockBackend) SetTemperature(temp float64) error { + if temp < 0 || temp > 2 { + return fmt.Errorf("invalid temperature") + } + m.temperature = temp + return nil +} +func (m *MockBackend) SystemPrompt() string { return m.systemPrompt } +func (m *MockBackend) SetSystemPrompt(prompt string) { m.systemPrompt = prompt } +func (m *MockBackend) ThinkingTokens() int { return m.thinkingTokens } +func (m *MockBackend) SetThinkingTokens(tokens int) { m.thinkingTokens = tokens } +func (m *MockBackend) LastTokens() int { return m.lastTokens } +func (m *MockBackend) TotalTokens() int { return m.totalTokens } +func (m *MockBackend) ContextLimit() int { return m.contextLimit } + +func (m *MockBackend) Compact(ctx context.Context) error { + m.compactCalled = true + if m.compactError != nil { + return m.compactError + } + // Simulate compaction - reduce tokens + m.totalTokens = m.totalTokens / 4 + m.messages = []llm.Message{{Role: "system", Content: "compacted summary"}} + return nil +} + +func (m *MockBackend) Messages() []llm.Message { + result := make([]llm.Message, len(m.messages)) + copy(result, m.messages) + return result +} + +func (m *MockBackend) MessagesJSON() ([]byte, error) { + return []byte("[]"), nil +} + +func (m *MockBackend) AddSystemMessage(content string) { + m.messages = append([]llm.Message{{Role: "system", Content: content}}, m.messages...) +} + +func (m *MockBackend) Reset() { + m.messages = make([]llm.Message, 0) + m.lastTokens = 0 + m.totalTokens = 0 +} + +func (m *MockBackend) Ask(ctx context.Context, prompt string) (string, error) { + if m.askError != nil { + return "", m.askError + } + m.messages = append(m.messages, llm.Message{Role: "user", Content: prompt}) + m.messages = append(m.messages, llm.Message{Role: "assistant", Content: m.askResponse}) + m.lastTokens = len(prompt) + len(m.askResponse) + m.totalTokens += m.lastTokens + return m.askResponse, nil +} + +func (m *MockBackend) StartStream(ctx context.Context, prompt string) error { + return fmt.Errorf("streaming not implemented in mock") +} + +func (m *MockBackend) ReadStreamChunk() (string, bool) { + return "", false +} + +func (m *MockBackend) IsStreaming() bool { + return false +} + +func (m *MockBackend) WaitStream() {} + +// Verify MockBackend implements Backend +var _ llm.Backend = (*MockBackend)(nil) diff --git a/internal/llmfs/root.go b/internal/llmfs/root.go index fbbf4f24022992aaa831151e4b7ab23cd8b4fcbb..dbbc1f13cd2309a57ddb92571ac4efef2ac7fe90 100644 --- a/internal/llmfs/root.go +++ b/internal/llmfs/root.go @@ -18,7 +18,10 @@ func NewRoot(client llm.Backend) protocol.Dir { root.AddChild(NewTokensFile(client)) root.AddChild(NewNewFile(client)) root.AddChild(NewContextFile(client)) + root.AddChild(NewThinkingFile(client)) root.AddChild(NewExampleFile()) + root.AddChild(NewUsageFile(client)) + root.AddChild(NewCompactFile(client)) // Add stream directory streamDir := protocol.NewStaticDir("stream") diff --git a/internal/llmfs/thinking.go b/internal/llmfs/thinking.go new file mode 100644 index 0000000000000000000000000000000000000000..63ada9755ef34b05e1d2c8b66b4d6345c9c70483 --- /dev/null +++ b/internal/llmfs/thinking.go @@ -0,0 +1,86 @@ +package llmfs + +import ( + "fmt" + "io" + "strconv" + "strings" + + "github.com/NERVsystems/llm9p/internal/llm" + "github.com/NERVsystems/llm9p/internal/protocol" +) + +// ThinkingFile exposes the thinking token budget (read/write) +// Values: -1 = max (31999), 0 = disabled, >0 = specific budget +// Only effective with CLI backend; API backend ignores this setting. +type ThinkingFile struct { + *protocol.BaseFile + client llm.Backend +} + +// NewThinkingFile creates the thinking file +func NewThinkingFile(client llm.Backend) *ThinkingFile { + return &ThinkingFile{ + BaseFile: protocol.NewBaseFile("thinking", 0666), + client: client, + } +} + +func (f *ThinkingFile) Read(p []byte, offset int64) (int, error) { + tokens := f.client.ThinkingTokens() + var content string + switch { + case tokens < 0: + content = "max\n" + case tokens == 0: + content = "off\n" + default: + content = fmt.Sprintf("%d\n", tokens) + } + if offset >= int64(len(content)) { + return 0, io.EOF + } + n := copy(p, content[offset:]) + return n, nil +} + +func (f *ThinkingFile) Write(p []byte, offset int64) (int, error) { + input := strings.TrimSpace(string(p)) + input = strings.ToLower(input) + + var tokens int + switch input { + case "max", "on", "true", "enabled", "-1": + tokens = -1 + case "off", "false", "disabled", "0": + tokens = 0 + default: + var err error + tokens, err = strconv.Atoi(input) + if err != nil { + return 0, fmt.Errorf("invalid thinking value: use 'max', 'off', or a number") + } + if tokens < 0 { + tokens = -1 // Treat any negative as max + } + } + + f.client.SetThinkingTokens(tokens) + return len(p), nil +} + +func (f *ThinkingFile) Stat() protocol.Stat { + s := f.BaseFile.Stat() + tokens := f.client.ThinkingTokens() + var content string + switch { + case tokens < 0: + content = "max\n" + case tokens == 0: + content = "off\n" + default: + content = fmt.Sprintf("%d\n", tokens) + } + s.Length = uint64(len(content)) + return s +} diff --git a/internal/llmfs/usage.go b/internal/llmfs/usage.go new file mode 100644 index 0000000000000000000000000000000000000000..b43e90d18e57e1545fa39eba8db8e41f86882dd0 --- /dev/null +++ b/internal/llmfs/usage.go @@ -0,0 +1,49 @@ +package llmfs + +import ( + "fmt" + "io" + + "github.com/NERVsystems/llm9p/internal/llm" + "github.com/NERVsystems/llm9p/internal/protocol" +) + +// UsageFile provides token usage observability +// Read returns "tokens/limit" (e.g., "45000/200000") +type UsageFile struct { + *protocol.BaseFile + client llm.Backend +} + +// NewUsageFile creates the usage file +func NewUsageFile(client llm.Backend) *UsageFile { + return &UsageFile{ + BaseFile: protocol.NewBaseFile("usage", 0444), + client: client, + } +} + +func (f *UsageFile) Read(p []byte, offset int64) (int, error) { + tokens := f.client.TotalTokens() + limit := f.client.ContextLimit() + content := fmt.Sprintf("%d/%d\n", tokens, limit) + + if offset >= int64(len(content)) { + return 0, io.EOF + } + n := copy(p, content[offset:]) + return n, nil +} + +func (f *UsageFile) Write(p []byte, offset int64) (int, error) { + return 0, fmt.Errorf("usage is read-only") +} + +func (f *UsageFile) Stat() protocol.Stat { + s := f.BaseFile.Stat() + tokens := f.client.TotalTokens() + limit := f.client.ContextLimit() + content := fmt.Sprintf("%d/%d\n", tokens, limit) + s.Length = uint64(len(content)) + return s +} diff --git a/internal/llmfs/usage_test.go b/internal/llmfs/usage_test.go new file mode 100644 index 0000000000000000000000000000000000000000..6e1d1af95890be2fea9ac70127056d50138e8b3c --- /dev/null +++ b/internal/llmfs/usage_test.go @@ -0,0 +1,119 @@ +package llmfs + +import ( + "io" + "testing" +) + +func TestUsageFile_Read(t *testing.T) { + mock := NewMockBackend() + mock.totalTokens = 45000 + mock.contextLimit = 200000 + + usage := NewUsageFile(mock) + + // Read the full content + buf := make([]byte, 100) + n, err := usage.Read(buf, 0) + if err != nil { + t.Fatalf("Read() error: %v", err) + } + + content := string(buf[:n]) + expected := "45000/200000\n" + if content != expected { + t.Errorf("Read() = %q, want %q", content, expected) + } +} + +func TestUsageFile_Read_Offset(t *testing.T) { + mock := NewMockBackend() + mock.totalTokens = 1000 + mock.contextLimit = 10000 + + usage := NewUsageFile(mock) + + // Read with offset (skip first 5 bytes: "1000/") + buf := make([]byte, 100) + n, err := usage.Read(buf, 5) + if err != nil { + t.Fatalf("Read() error: %v", err) + } + + content := string(buf[:n]) + expected := "10000\n" + if content != expected { + t.Errorf("Read(offset=5) = %q, want %q", content, expected) + } +} + +func TestUsageFile_Read_EOF(t *testing.T) { + mock := NewMockBackend() + mock.totalTokens = 100 + mock.contextLimit = 1000 + + usage := NewUsageFile(mock) + + // Read past end + buf := make([]byte, 100) + n, err := usage.Read(buf, 1000) + if err != io.EOF { + t.Errorf("Read(offset=1000) error = %v, want io.EOF", err) + } + if n != 0 { + t.Errorf("Read(offset=1000) n = %d, want 0", n) + } +} + +func TestUsageFile_Write(t *testing.T) { + mock := NewMockBackend() + usage := NewUsageFile(mock) + + // Write should fail (read-only) + n, err := usage.Write([]byte("test"), 0) + if err == nil { + t.Error("Write() should return error for read-only file") + } + if n != 0 { + t.Errorf("Write() n = %d, want 0", n) + } +} + +func TestUsageFile_Stat(t *testing.T) { + mock := NewMockBackend() + mock.totalTokens = 12345 + mock.contextLimit = 200000 + + usage := NewUsageFile(mock) + stat := usage.Stat() + + // Content would be "12345/200000\n" = 14 chars + expected := uint64(len("12345/200000\n")) + if stat.Length != expected { + t.Errorf("Stat().Length = %d, want %d", stat.Length, expected) + } +} + +func TestUsageFile_DynamicContent(t *testing.T) { + mock := NewMockBackend() + mock.totalTokens = 0 + mock.contextLimit = 100000 + + usage := NewUsageFile(mock) + + // First read + buf := make([]byte, 100) + n, _ := usage.Read(buf, 0) + if string(buf[:n]) != "0/100000\n" { + t.Errorf("First read = %q, want '0/100000\\n'", string(buf[:n])) + } + + // Update tokens + mock.totalTokens = 50000 + + // Second read should reflect the change + n, _ = usage.Read(buf, 0) + if string(buf[:n]) != "50000/100000\n" { + t.Errorf("Second read = %q, want '50000/100000\\n'", string(buf[:n])) + } +}