diff --git a/pkg/agent/cost_test.go b/pkg/agent/cost_test.go new file mode 100644 index 0000000..a3c3540 --- /dev/null +++ b/pkg/agent/cost_test.go @@ -0,0 +1,100 @@ +package agent + +import ( + "math" + "testing" + + "ollie/pkg/backend" +) + +func approxEqual(a, b, tol float64) bool { + return math.Abs(a-b) <= tol +} + +// TestComputeCostUSD_BaseTokens verifies input+output pricing with no cache fields. +func TestComputeCostUSD_BaseTokens(t *testing.T) { + // claude-sonnet-4: $3/M input, $15/M output + got := computeCostUSD("claude-sonnet-4-5", backend.Usage{ + InputTokens: 1_000_000, + OutputTokens: 1_000_000, + }) + want := 3.00 + 15.00 + if !approxEqual(got, want, 0.001) { + t.Errorf("got %.4f; want %.4f", got, want) + } +} + +// TestComputeCostUSD_ClaudeCacheReadDiscount verifies that cached input tokens +// are charged at 10% of the normal input rate for Claude models. +func TestComputeCostUSD_ClaudeCacheReadDiscount(t *testing.T) { + // claude-sonnet-4: $3/M input → cache read = $0.30/M + got := computeCostUSD("claude-sonnet-4-5", backend.Usage{ + InputTokens: 0, + CachedInputTokens: 1_000_000, + OutputTokens: 0, + }) + want := 0.30 // 10% of $3.00 + if !approxEqual(got, want, 0.001) { + t.Errorf("claude cache read: got %.4f; want %.4f", got, want) + } +} + +// TestComputeCostUSD_NonClaudeCacheReadDiscount verifies that cached input tokens +// are charged at 50% of the normal input rate for non-Claude models (OpenAI). +func TestComputeCostUSD_NonClaudeCacheReadDiscount(t *testing.T) { + // gpt-4o: $2.50/M input → cache read = $1.25/M + got := computeCostUSD("gpt-4o", backend.Usage{ + InputTokens: 0, + CachedInputTokens: 1_000_000, + OutputTokens: 0, + }) + want := 1.25 // 50% of $2.50 + if !approxEqual(got, want, 0.001) { + t.Errorf("gpt-4o cache read: got %.4f; want %.4f", got, want) + } +} + +// TestComputeCostUSD_CacheCreationSurcharge verifies cache creation tokens are +// charged at 125% of the normal input rate. +func TestComputeCostUSD_CacheCreationSurcharge(t *testing.T) { + // claude-sonnet-4: $3/M input → cache creation = $3.75/M + got := computeCostUSD("claude-sonnet-4-5", backend.Usage{ + InputTokens: 0, + CacheCreationTokens: 1_000_000, + OutputTokens: 0, + }) + want := 3.75 // 125% of $3.00 + if !approxEqual(got, want, 0.001) { + t.Errorf("cache creation surcharge: got %.4f; want %.4f", got, want) + } +} + +// TestComputeCostUSD_AllFieldsCombined verifies that all four token categories +// are summed correctly in a single call. +func TestComputeCostUSD_AllFieldsCombined(t *testing.T) { + // claude-sonnet-4: $3/M in, $15/M out + // 100k normal input = $0.30 + // 200k cached input = $0.06 (10% of $3/M) + // 50k cache create = $0.1875 (125% of $3/M) + // 100k output = $1.50 + got := computeCostUSD("claude-sonnet-4-5", backend.Usage{ + InputTokens: 100_000, + CachedInputTokens: 200_000, + CacheCreationTokens: 50_000, + OutputTokens: 100_000, + }) + want := 0.30 + 0.06 + 0.1875 + 1.50 + if !approxEqual(got, want, 0.0001) { + t.Errorf("combined: got %.6f; want %.6f", got, want) + } +} + +// TestComputeCostUSD_UnknownModelZero verifies that unknown/local models return 0. +func TestComputeCostUSD_UnknownModelZero(t *testing.T) { + got := computeCostUSD("llama-3-local", backend.Usage{ + InputTokens: 1_000_000, OutputTokens: 1_000_000, + }) + if got != 0 { + t.Errorf("unknown model: got %.4f; want 0", got) + } +} diff --git a/pkg/agent/loop_cache_test.go b/pkg/agent/loop_cache_test.go index 4262fab..0ff83e9 100644 --- a/pkg/agent/loop_cache_test.go +++ b/pkg/agent/loop_cache_test.go @@ -143,6 +143,43 @@ func TestResultCache_ErrorNotCached(t *testing.T) { } } +// TestResultCache_IdenticalParallelReads verifies that two identical read-safe +// tool calls issued in the same parallel batch (both goroutines miss the cache +// simultaneously) both complete without error and return consistent results. +func TestResultCache_IdenticalParallelReads(t *testing.T) { + var execCount int32 + be := defaultBE() + be.respond = toolsStream([]backend.ToolCall{ + {ID: "1", Name: "file_read", Arguments: json.RawMessage(`{"path":"a.txt"}`)}, + {ID: "2", Name: "file_read", Arguments: json.RawMessage(`{"path":"a.txt"}`)}, + }) + + c := newCore(t, be, nil) + c.cfg.ClassifyTool = func(string) bool { return true } // both parallel-safe → same batch + c.cfg.Exec = func(_ context.Context, _ string, _ json.RawMessage) (string, error) { + atomic.AddInt32(&execCount, 1) + return "file contents", nil + } + + evs := collectEvents(context.Background(), c, "identical parallel") + + toolEvs := byRole(evs, "tool") + if len(toolEvs) != 2 { + t.Fatalf("tool events = %d; want 2", len(toolEvs)) + } + for i, ev := range toolEvs { + if ev != "file contents" { + t.Errorf("tool event[%d] = %q; want %q", i, ev, "file contents") + } + } + // Both may have executed (cache miss race) or one may have hit cache — + // either is correct. What must not happen: panic, empty result, or wrong value. + n := atomic.LoadInt32(&execCount) + if n < 1 || n > 2 { + t.Errorf("Exec called %d times; want 1 or 2", n) + } +} + type mockErr struct{ msg string } func (e *mockErr) Error() string { return e.msg } diff --git a/pkg/agent/loop_prune_test.go b/pkg/agent/loop_prune_test.go index 9b184d1..e9a7f2f 100644 --- a/pkg/agent/loop_prune_test.go +++ b/pkg/agent/loop_prune_test.go @@ -168,3 +168,86 @@ func TestExtractExecCodePaths_Basic(t *testing.T) { t.Errorf("writes = %v; want [/b.go /c.go]", writes) } } + +// TestPruneStaleReads_ReadWriteReadCycle is the primary correctness invariant: +// read before a write is pruned; read after the same write is kept. +func TestPruneStaleReads_ReadWriteReadCycle(t *testing.T) { + history := []backend.Message{ + {Role: "user", Content: "update foo"}, + makeExecCodeMsg("c1", []map[string]any{{"tool": "file_read", "args": []string{"/foo.go"}}}), + makeToolResult("c1", "old content"), + makeExecCodeMsg("c2", []map[string]any{{"tool": "file_write", "args": []string{"/foo.go", "new content"}}}), + makeToolResult("c2", "ok"), + makeExecCodeMsg("c3", []map[string]any{{"tool": "file_read", "args": []string{"/foo.go"}}}), + makeToolResult("c3", "new content"), + } + + out := pruneStaleReads(history) + + // c1 (pre-write read) must be gone; c3 (post-write read) must remain. + for _, m := range out { + if m.Role == "tool" && m.ToolCallID == "c1" { + t.Error("pre-write read result c1 should have been pruned") + } + } + found := false + for _, m := range out { + if m.Role == "tool" && m.ToolCallID == "c3" { + found = true + } + } + if !found { + t.Error("post-write read result c3 should be kept") + } + if len(out) != len(history)-1 { + t.Errorf("len(out) = %d; want %d", len(out), len(history)-1) + } +} + +// TestPruneStaleReads_AssistantMessagesUnaffected verifies that the assistant +// messages that issued the tool calls are never removed — only the tool result +// messages (role:"tool") for stale reads are dropped. +func TestPruneStaleReads_AssistantMessagesUnaffected(t *testing.T) { + history := []backend.Message{ + {Role: "user", Content: "go"}, + makeExecCodeMsg("c1", []map[string]any{{"tool": "file_read", "args": []string{"/foo.go"}}}), + makeToolResult("c1", "old"), + makeExecCodeMsg("c2", []map[string]any{{"tool": "file_write", "args": []string{"/foo.go", "new"}}}), + makeToolResult("c2", "ok"), + } + + out := pruneStaleReads(history) + + var assistantCount int + for _, m := range out { + if m.Role == "assistant" { + assistantCount++ + } + } + if assistantCount != 2 { + t.Errorf("assistant messages = %d; want 2 (pruning must not remove assistant messages)", assistantCount) + } +} + +// TestPruneStaleReads_NonExecuteCodeUnaffected verifies that tool calls with +// names other than "execute_code" are ignored by the path tracker and never +// cause unrelated tool results to be pruned. +func TestPruneStaleReads_NonExecuteCodeUnaffected(t *testing.T) { + history := []backend.Message{ + {Role: "user", Content: "go"}, + // A non-execute_code call that happens to reference a path. + {Role: "assistant", ToolCalls: []backend.ToolCall{ + {ID: "c1", Name: "some_other_tool", Arguments: json.RawMessage(`{"path":"/foo.go"}`)}, + }}, + makeToolResult("c1", "other tool result"), + // execute_code write — should not prune the unrelated result above. + makeExecCodeMsg("c2", []map[string]any{{"tool": "file_write", "args": []string{"/foo.go", "new"}}}), + makeToolResult("c2", "ok"), + } + + out := pruneStaleReads(history) + + if len(out) != len(history) { + t.Errorf("len(out) = %d; want %d (non-execute_code results must not be pruned)", len(out), len(history)) + } +} diff --git a/pkg/backend/openai_encode_test.go b/pkg/backend/openai_encode_test.go index 2e2e9ff..5db196a 100644 --- a/pkg/backend/openai_encode_test.go +++ b/pkg/backend/openai_encode_test.go @@ -1,7 +1,12 @@ package backend import ( + "context" "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" "strings" "testing" ) @@ -228,6 +233,135 @@ func TestEncodeMessages_OpenRouterClaudeCacheControl(t *testing.T) { func strPtr(s string) *string { return &s } +// --- OpenRouter cache injection: positive and negative cases --- + +// captureRawRequest spins up an httptest server that captures the raw request +// body JSON, calls the backend, and drains the stream. +func captureRawRequest(t *testing.T, name, model string, msgs []Message, tools []Tool) map[string]json.RawMessage { + t.Helper() + var rawBody []byte + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + rawBody, _ = io.ReadAll(r.Body) + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint(w, "data: [DONE]\n\n") + })) + defer srv.Close() + + b, err := NewOpenAI(name, srv.URL, "key") + if err != nil { + t.Fatal(err) + } + b.SetModel(model) + ch, err := b.ChatStream(context.Background(), msgs, tools, GenerationParams{}) + if err != nil { + t.Fatal(err) + } + for range ch { + } + + var top map[string]json.RawMessage + json.Unmarshal(rawBody, &top) //nolint:errcheck + return top +} + +// TestOpenRouterClaude_SystemCacheControl verifies that when the backend is +// "openrouter" and the model contains "claude-", the system message content is +// encoded as a JSON array with a cache_control:ephemeral block. +func TestOpenRouterClaude_SystemCacheControl(t *testing.T) { + top := captureRawRequest(t, "openrouter", "claude-sonnet-4-5", + []Message{{Role: "system", Content: "be helpful"}, {Role: "user", Content: "hi"}}, + nil, + ) + + var messages []map[string]json.RawMessage + json.Unmarshal(top["messages"], &messages) //nolint:errcheck + if len(messages) == 0 { + t.Fatal("no messages in captured request") + } + sys := messages[0] + + // content must be an array (first byte '['), not a plain string. + content := sys["content"] + if len(content) == 0 || content[0] != '[' { + t.Errorf("system content should be array; got %s", content) + } + if !strings.Contains(string(content), "cache_control") { + t.Errorf("system content missing cache_control: %s", content) + } + if !strings.Contains(string(content), "ephemeral") { + t.Errorf("system content missing ephemeral: %s", content) + } +} + +// TestOpenRouterClaude_LastToolCacheControl verifies that the last tool in the +// array gets cache_control:ephemeral on openrouter+claude requests. +func TestOpenRouterClaude_LastToolCacheControl(t *testing.T) { + top := captureRawRequest(t, "openrouter", "claude-sonnet-4-5", + []Message{{Role: "user", Content: "hi"}}, + []Tool{ + {Name: "alpha", Parameters: json.RawMessage(`{}`)}, + {Name: "beta", Parameters: json.RawMessage(`{}`)}, + }, + ) + + var tools []map[string]json.RawMessage + json.Unmarshal(top["tools"], &tools) //nolint:errcheck + if len(tools) != 2 { + t.Fatalf("tools len = %d; want 2", len(tools)) + } + if strings.Contains(string(tools[0]["cache_control"]), "ephemeral") { + t.Errorf("first tool should not have cache_control:ephemeral") + } + if !strings.Contains(string(tools[1]["cache_control"]), "ephemeral") { + t.Errorf("last tool missing cache_control:ephemeral; got %s", tools[1]["cache_control"]) + } +} + +// TestOpenRouterNonClaude_NoCacheControl verifies that a non-Claude model on +// openrouter does NOT get cache_control injected into system messages or tools. +func TestOpenRouterNonClaude_NoCacheControl(t *testing.T) { + top := captureRawRequest(t, "openrouter", "mistral-7b", + []Message{{Role: "system", Content: "be helpful"}, {Role: "user", Content: "hi"}}, + []Tool{{Name: "run", Parameters: json.RawMessage(`{}`)}}, + ) + + var messages []map[string]json.RawMessage + json.Unmarshal(top["messages"], &messages) //nolint:errcheck + content := messages[0]["content"] + // Must be a plain JSON string, not an array. + if len(content) == 0 || content[0] != '"' { + t.Errorf("non-claude system content should be plain string; got %s", content) + } + if strings.Contains(string(content), "cache_control") { + t.Errorf("non-claude system content must not contain cache_control: %s", content) + } + + var tools []map[string]json.RawMessage + json.Unmarshal(top["tools"], &tools) //nolint:errcheck + if len(tools) > 0 && strings.Contains(string(tools[0]["cache_control"]), "ephemeral") { + t.Errorf("non-claude tool should not have cache_control:ephemeral") + } +} + +// TestOpenAI_NoCacheControl verifies that a plain "openai" backend (not +// openrouter) never injects cache_control, even for a claude-named model. +func TestOpenAI_NoCacheControl(t *testing.T) { + top := captureRawRequest(t, "openai", "claude-sonnet-4-5", + []Message{{Role: "system", Content: "sys"}, {Role: "user", Content: "hi"}}, + []Tool{{Name: "run", Parameters: json.RawMessage(`{}`)}}, + ) + + var messages []map[string]json.RawMessage + json.Unmarshal(top["messages"], &messages) //nolint:errcheck + content := messages[0]["content"] + if len(content) == 0 || content[0] != '"' { + t.Errorf("openai system content should be plain string; got %s", content) + } + if strings.Contains(string(content), "cache_control") { + t.Errorf("openai backend must not inject cache_control: %s", content) + } +} + func TestEncodeRequest_ToolsOmittedWhenEmpty(t *testing.T) { req := openAIChatRequest{ Model: "m",