chore: initial commit

2026-10-23 18:35:44 +08:00
commit f6921b7b23
69 changed files with 12148 additions and 0 deletions
@@ -0,0 +1,89 @@
+package protocol
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"reflect"
+	"strings"
+	"testing"
+)
+
+func TestMergeUsageSizeBoundary(t *testing.T) {
+	for _, kind := range []Kind{Chat, Responses, Anthropic} {
+		t.Run(string(kind), func(t *testing.T) {
+			m := object{"padding": strings.Repeat("x", maxStreamUsageBytes-len(`{"padding":""}`))}
+			u, err := mergeUsage(kind, Usage{}, object{"usage": m})
+			if err != nil || len(u.Raw) != maxStreamUsageBytes {
+				t.Fatalf("exact limit: size=%d err=%v", len(u.Raw), err)
+			}
+			for _, update := range []object{
+				{"usage": object{"padding": str(m["padding"]) + "x"}},
+				{"usage": object{"new_key": 1}},
+			} {
+				got, err := mergeUsage(kind, u, update)
+				if err == nil || !strings.Contains(err.Error(), "usage exceeds 16384 bytes") || !reflect.DeepEqual(got, u) {
+					t.Fatalf("overflow: retained=%d err=%v", len(got.Raw), err)
+				}
+			}
+			for i := 0; i < 3; i++ {
+				got, err := mergeUsage(kind, u, object{"usage": m})
+				if err != nil || !reflect.DeepEqual(got, u) {
+					t.Fatalf("replacement at limit: %v", err)
+				}
+			}
+		})
+	}
+}
+
+func TestStreamUsageSizeLimits(t *testing.T) {
+	for _, from := range []Kind{Chat, Responses, Anthropic} {
+		for _, to := range []Kind{Chat, Responses, Anthropic} {
+			for _, cumulative := range []bool{false, true} {
+				t.Run(fmt.Sprintf("%s/%s/cumulative=%t", from, to, cumulative), func(t *testing.T) {
+					var input strings.Builder
+					writeUsage := func(u object) {
+						m := object{"usage": u}
+						switch from {
+						case Anthropic:
+							m["type"] = "ping"
+						case Responses:
+							m = object{"type": "response.in_progress", "response": m}
+						}
+						input.WriteString(limitEvent(mustLimitJSON(t, m)))
+					}
+					writeUsage(object{"prompt_tokens": 7, "input_tokens": 7})
+					if cumulative {
+						for i := 0; i < 32; i++ {
+							writeUsage(object{fmt.Sprintf("key_%d", i): strings.Repeat("x", 1024)})
+						}
+					} else {
+						writeUsage(object{"padding": strings.Repeat("x", maxStreamUsageBytes)})
+					}
+					u, err := Stream(context.Background(), io.Discard, strings.NewReader(input.String()), from, to, "")
+					if err == nil || !strings.Contains(err.Error(), "usage exceeds 16384 bytes") {
+						t.Fatalf("expected usage limit error, got %v", err)
+					}
+					if !u.Known || u.Input != 7 || len(u.Raw) > maxStreamUsageBytes {
+						t.Fatalf("lost valid usage: known=%t input=%d size=%d", u.Known, u.Input, len(u.Raw))
+					}
+				})
+			}
+		}
+	}
+}
+
+func TestMergeUsagePartialUpdates(t *testing.T) {
+	u, err := mergeUsage(Anthropic, Usage{}, object{"message": object{"usage": object{"input_tokens": 5, "cache_read_input_tokens": 2}}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	u, err = mergeUsage(Anthropic, u, object{"usage": object{"output_tokens": 3}})
+	if err != nil || !u.Known || u.Input != 7 || u.Output != 3 || u.CacheRead != 2 {
+		t.Fatalf("partial update: %+v err=%v", u, err)
+	}
+	got, err := mergeUsage(Anthropic, u, object{})
+	if err != nil || !reflect.DeepEqual(got, u) {
+		t.Fatalf("missing usage changed totals: %+v err=%v", got, err)
+	}
+}