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,164 @@
+package protocol
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"strings"
+	"testing"
+)
+
+func limitEvent(s string) string { return "data: " + s + "\n\n" }
+func TestSSEReadLimits(t *testing.T) {
+	for _, to := range []Kind{Chat, Anthropic} {
+		var out bytes.Buffer
+		_, err := streamWithLimits(context.Background(), &out, strings.NewReader(strings.Repeat("x", 10000)), Chat, to, "", 64, 8, 64)
+		if err == nil || !strings.Contains(err.Error(), "size limit") || out.Len() != 0 {
+			t.Fatalf("%s: %v", to, err)
+		}
+	}
+	for _, input := range []string{"12345678", "1234567\n"} {
+		line, err := readSSELine(context.Background(), bufio.NewReader(strings.NewReader(input)), 8)
+		if line != input || (err != nil && err != io.EOF) {
+			t.Fatalf("%q %v", line, err)
+		}
+	}
+	var out bytes.Buffer
+	_, err := streamWithLimits(context.Background(), &out, strings.NewReader(":123456\n:123456\n"), Chat, Chat, "", 12, 8, 16)
+	if err == nil || out.String() != ":123456\n" {
+		t.Fatalf("%q %v", out.String(), err)
+	}
+}
+func TestStreamPassthroughTotalLimit(t *testing.T) {
+	frame := "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":null}]}\n\n"
+	var out bytes.Buffer
+	_, err := streamWithLimits(context.Background(), &out, strings.NewReader(strings.Repeat(frame, 4)), Chat, Chat, "", 512, 8, 2*len(frame))
+	if err == nil || !strings.Contains(err.Error(), "passthrough") || out.Len() != 2*len(frame) {
+		t.Fatalf("out=%d %q %v", out.Len(), out.String(), err)
+	}
+}
+func TestStreamCumulativeLimits(t *testing.T) {
+	for _, field := range []string{"text", "arguments", "identity"} {
+		delta := object{"content": strings.Repeat("x", 100)}
+		if field != "text" {
+			call := object{"index": 0}
+			if field == "arguments" {
+				call["function"] = object{"arguments": strings.Repeat("x", 100)}
+			} else {
+				call["id"] = strings.Repeat("x", 100)
+			}
+			delta = object{"tool_calls": []any{call}}
+		}
+		event := limitEvent(fmt.Sprintf("{\"choices\":[{\"index\":0,\"delta\":%s}]}", mustLimitJSON(t, delta)))
+		_, err := streamWithLimits(context.Background(), io.Discard, strings.NewReader(strings.Repeat(event, 6)), Chat, Responses, "", 512, 8, 512)
+		if err == nil || !strings.Contains(err.Error(), "retained content") {
+			t.Fatalf("%s: %v", field, err)
+		}
+	}
+}
+func TestStreamBlockLimits(t *testing.T) {
+	for _, kind := range []string{"text", "thinking", "reasoning", "summary", "tool"} {
+		from := Anthropic
+		in := limitEvent("{\"type\":\"message_start\",\"message\":{\"content\":[]}}")
+		if kind == "reasoning" || kind == "summary" || kind == "tool" {
+			from = Responses
+			in = ""
+		}
+		if kind == "summary" {
+			in = limitEvent("{\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"reasoning\",\"id\":\"r\"}}")
+		}
+		for i := 0; i < 4; i++ {
+			m := object{"type": "content_block_start", "index": i, "content_block": object{"type": kind, "text": ""}}
+			switch kind {
+			case "thinking":
+				m["content_block"] = object{"type": "thinking", "thinking": "", "signature": ""}
+			case "reasoning", "tool":
+				item := object{"type": "reasoning", "id": "r"}
+				if kind == "tool" {
+					item = object{"type": "function_call", "call_id": "c", "name": "f", "arguments": ""}
+				}
+				m = object{"type": "response.output_item.added", "output_index": i, "item": item}
+			case "summary":
+				m = object{"type": "response.reasoning_summary_part.added", "output_index": 0, "item_id": "r", "summary_index": i, "part": object{"type": "summary_text", "text": ""}}
+			}
+			in += limitEvent(mustLimitJSON(t, m))
+		}
+		_, err := streamWithLimits(context.Background(), io.Discard, strings.NewReader(in), from, Chat, "", 1024, 3, 1024)
+		if err == nil || !strings.Contains(err.Error(), "3 blocks") {
+			t.Fatalf("%s: %v", kind, err)
+		}
+	}
+}
+
+func TestReasoningIdentityCumulativeLimit(t *testing.T) {
+	var in strings.Builder
+	for i := 0; i < 6; i++ {
+		in.WriteString(limitEvent(mustLimitJSON(t, object{"type": "response.output_item.added", "output_index": i, "item": object{"type": "reasoning", "id": strings.Repeat("r", 100)}})))
+	}
+	_, err := streamWithLimits(context.Background(), io.Discard, strings.NewReader(in.String()), Responses, Chat, "", 512, 8, 512)
+	if err == nil || !strings.Contains(err.Error(), "retained content") {
+		t.Fatal(err)
+	}
+}
+func TestStreamBufferSemantics(t *testing.T) {
+	var out bytes.Buffer
+	s := &streamState{ctx: context.Background(), w: &out, from: Chat, to: Responses, source: map[int]*streamBlock{}, byteLimit: 64, blockLimit: 4}
+	b, err := s.add(0, part{typ: "call"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err = s.delta(b, "{\"a\":"); err != nil {
+		t.Fatal(err)
+	}
+	if out.Len() != 0 {
+		t.Fatal("incomplete identity emitted")
+	}
+	b.id, b.name = "c", "f"
+	if err = s.delta(b, "1}"); err != nil {
+		t.Fatal(err)
+	}
+	if b.args != "{\"a\":1}" || b.sent != 7 {
+		t.Fatal(b.part)
+	}
+	if err = s.close(b); err != nil {
+		t.Fatal(err)
+	}
+	text, err := s.add(1, part{typ: "text"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, v := range []string{"ab", "cd"} {
+		if err = s.delta(text, v); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if text.text != "abcd" {
+		t.Fatal(text.text)
+	}
+	if err = s.finish("stop"); err != nil {
+		t.Fatal(err)
+	}
+	if err = s.end(); err != nil {
+		t.Fatal(err)
+	}
+	if !strings.Contains(out.String(), "abcd") || !strings.Contains(out.String(), "response.completed") {
+		t.Fatal(out.String())
+	}
+	s.finished = false
+	text.closed = false
+	s.byteLimit = s.retained
+	if err = s.delta(text, "x"); err == nil || text.text != "abcd" {
+		t.Fatal("overflow changed content")
+	}
+}
+func mustLimitJSON(t *testing.T, m object) string {
+	t.Helper()
+	b, err := json.Marshal(m)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return string(b)
+}