1package provider
2
3import (
4 "bytes"
5 "context"
6 "encoding/binary"
7 "encoding/json"
8 "io"
9 "net/http"
10 "strings"
11 "testing"
12)
13
14func TestNormalizeToolArgs(t *testing.T) {
15 tab, nl, bell := string(rune(9)), string(rune(10)), string(rune(7))
16 tests := []struct {
17 name string
18 in string
19 want string
20 }{
21 {name: "empty", in: "", want: `{}`},
22 {name: "valid kept verbatim", in: `{"a":"b\tc"}`, want: `{"a":"b\tc"}`},
23 {name: "raw tab in string", in: `{"old":"` + tab + `x"}`, want: `{"old":"\tx"}`},
24 {name: "raw newline in string", in: `{"s":"a` + nl + `b"}`, want: `{"s":"a\nb"}`},
25 {name: "other control char", in: `{"s":"` + bell + `"}`, want: `{"s":"\u0007"}`},
26 {name: "whitespace outside strings untouched", in: "{" + nl + tab + `"s":"` + tab + `"}`, want: "{" + nl + tab + `"s":"\t"}`},
27 {name: "escaped quote keeps string state", in: `{"s":"\"` + tab + `"}`, want: `{"s":"\"\t"}`},
28 {name: "escaped backslash closes string", in: `{"a":"x\\"` + tab + `}`, want: `{"a":"x\\"` + tab + `}`},
29 {name: "multibyte utf-8 kept", in: `{"a":"é` + tab + `"}`, want: `{"a":"é\t"}`},
30 {name: "unrepairable", in: `{"s":`, want: `{}`},
31 }
32 for _, tt := range tests {
33 t.Run(tt.name, func(t *testing.T) {
34 got := normalizeToolArgs(tt.in)
35 if string(got) != tt.want {
36 t.Fatalf("normalizeToolArgs(%q) = %q, want %q", tt.in, got, tt.want)
37 }
38 if !json.Valid(got) {
39 t.Fatalf("result %q is not valid JSON", got)
40 }
41 })
42 }
43}
44
45func TestAnthropicStreamRepairsRawTabInToolArgs(t *testing.T) {
46 // partial_json is a JSON string carrying the model's raw argument text;
47 // "\\t" below decodes to a real tab byte inside the argument string.
48 body := "event: content_block_start\ndata: {\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call-1\",\"name\":\"edit\",\"input\":{}}}\n\n" +
49 "event: content_block_delta\ndata: {\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"old\\\":\\\"\\tx\\\"}\"}}\n\n" +
50 "event: content_block_stop\ndata: {\"index\":0}\n\n" +
51 "event: message_delta\ndata: {\"delta\":{\"stop_reason\":\"tool_use\"}}\n\n"
52 resp := &http.Response{Body: io.NopCloser(strings.NewReader(body))}
53 out := make(chan Event, 16)
54 go (&anthropicClient{}).runStream(context.Background(), resp, Request{Model: "test"}, out)
55 assertToolArgOld(t, out, "\tx")
56}
57
58func TestBedrockStreamRepairsRawTabInToolArgs(t *testing.T) {
59 // Header-less event-stream frames; the reader takes the event type from
60 // the wrapped payload and does not validate CRCs. The `\t` escape in the input
61 // field decodes to a real tab byte inside the argument string.
62 var body bytes.Buffer
63 for _, payload := range []string{
64 `{"contentBlockStart":{"contentBlockIndex":0,"start":{"toolUse":{"toolUseId":"call-1","name":"edit"}}}}`,
65 `{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"toolUse":{"input":"{\"old\":\"\tx\"}"}}}}`,
66 `{"contentBlockStop":{"contentBlockIndex":0}}`,
67 `{"messageStop":{"stopReason":"tool_use"}}`,
68 } {
69 var prelude [12]byte
70 binary.BigEndian.PutUint32(prelude[0:4], uint32(16+len(payload)))
71 body.Write(prelude[:])
72 body.WriteString(payload)
73 body.Write(make([]byte, 4))
74 }
75 resp := &http.Response{Body: io.NopCloser(&body)}
76 out := make(chan Event, 16)
77 go (&bedrockClient{}).runStream(context.Background(), resp, Request{Model: "test"}, out)
78 assertToolArgOld(t, out, "\tx")
79}
80
81func assertToolArgOld(t *testing.T, out <-chan Event, want string) {
82 t.Helper()
83 var done EventDone
84 for ev := range out {
85 if d, ok := ev.(EventDone); ok {
86 done = d
87 }
88 }
89 if done.Err != nil {
90 t.Fatalf("stream error: %v", done.Err)
91 }
92 if len(done.Message.Content) != 1 {
93 t.Fatalf("content = %#v, want one tool call", done.Message.Content)
94 }
95 tc, ok := done.Message.Content[0].(ToolCallBlock)
96 if !ok {
97 t.Fatalf("content[0] = %T, want ToolCallBlock", done.Message.Content[0])
98 }
99 var args struct{ Old string }
100 if err := json.Unmarshal(tc.Arguments, &args); err != nil {
101 t.Fatalf("arguments %q not valid JSON: %v", tc.Arguments, err)
102 }
103 if args.Old != want {
104 t.Fatalf("old = %q, want %q", args.Old, want)
105 }
106}