AutoZero/llm/system_cache_test.go

package llm

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
)

// TestSystemCacheSplit：稳定 system 与动态 system 必须分开下发，且收尾回合
// 只禁用工具调用、不撤掉工具定义——两者都是为了保住 prompt cache 前缀。
func TestSystemCacheSplit(t *testing.T) {
	tools := []ToolDef{{Name: "read", Description: "d", Parameters: map[string]any{"type": "object"}}}
	opts := StreamOptions{SystemDynamic: "DYNAMIC", DisableTools: true}

	t.Run("anthropic：动态段是第二个 system 块且不带 cache_control", func(t *testing.T) {
		var body map[string]any
		srv := captureBodyServer(t, &body,
			"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"ok\"}}\n\n"+
				"data: {\"type\":\"message_stop\"}\n\n")
		defer srv.Close()
		p := &AnthropicProvider{BaseURL: srv.URL, APIKey: "k", Model: "m", MaxTokens: 128, HTTP: srv.Client()}
		if _, err := p.Stream(context.Background(), "STABLE",
			[]Message{{Role: RoleUser, Content: "hi"}}, tools, opts, nil); err != nil {
			t.Fatal(err)
		}
		blocks := body["system"].([]any)
		if len(blocks) != 2 {
			t.Fatalf("want 2 system blocks, got %d: %v", len(blocks), blocks)
		}
		first := blocks[0].(map[string]any)
		second := blocks[1].(map[string]any)
		if first["text"] != "STABLE" || first["cache_control"].(map[string]any)["type"] != "ephemeral" {
			t.Fatalf("stable block must be cached: %v", first)
		}
		if second["text"] != "DYNAMIC" {
			t.Fatalf("dynamic block text: %v", second)
		}
		if _, ok := second["cache_control"]; ok {
			t.Fatalf("dynamic block must not carry cache_control: %v", second)
		}
		if body["tool_choice"].(map[string]any)["type"] != "none" {
			t.Fatalf("tool_choice: %v", body["tool_choice"])
		}
		if len(body["tools"].([]any)) != 1 {
			t.Fatalf("tools must stay in the request: %v", body["tools"])
		}
	})

	t.Run("chat_completions：动态段是独立的第二条 system 消息", func(t *testing.T) {
		var body map[string]any
		srv := captureBodyServer(t, &body,
			"data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n")
		defer srv.Close()
		p := &OpenAIProvider{BaseURL: srv.URL, APIKey: "k", Model: "m", HTTP: srv.Client()}
		if _, err := p.Stream(context.Background(), "STABLE",
			[]Message{{Role: RoleUser, Content: "hi"}}, tools, opts, nil); err != nil {
			t.Fatal(err)
		}
		msgs := body["messages"].([]any)
		if len(msgs) != 2 {
			t.Fatalf("want one system message + user, got %v", msgs)
		}
		if msgs[0].(map[string]any)["content"] != "STABLE\n\nDYNAMIC" {
			t.Fatalf("system must keep the stable prefix and append the dynamic tail: %v", msgs[0])
		}
		if body["tool_choice"] != "none" {
			t.Fatalf("tool_choice: %v", body["tool_choice"])
		}
		if len(body["tools"].([]any)) != 1 {
			t.Fatalf("tools must stay in the request: %v", body["tools"])
		}
	})

	t.Run("responses：动态段是独立的第二条 system 消息", func(t *testing.T) {
		var body map[string]any
		srv := captureBodyServer(t, &body,
			"data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n"+
				"data: {\"type\":\"response.completed\",\"response\":{\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n\n")
		defer srv.Close()
		p := &ResponsesProvider{BaseURL: srv.URL, APIKey: "k", Model: "m", HTTP: srv.Client()}
		if _, err := p.Stream(context.Background(), "STABLE",
			[]Message{{Role: RoleUser, Content: "hi"}}, tools, opts, nil); err != nil {
			t.Fatal(err)
		}
		input := body["input"].([]any)
		if len(input) != 2 {
			t.Fatalf("want one system message + user, got %v", input)
		}
		if input[0].(map[string]any)["content"] != "STABLE\n\nDYNAMIC" {
			t.Fatalf("system must keep the stable prefix and append the dynamic tail: %v", input[0])
		}
		if body["tool_choice"] != "none" {
			t.Fatalf("tool_choice: %v", body["tool_choice"])
		}
	})
}

func captureBodyServer(t *testing.T, body *map[string]any, sse string) *httptest.Server {
	t.Helper()
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewDecoder(r.Body).Decode(body)
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = w.Write([]byte(sse))
	}))
}