1package provider
2
3import (
4 "bytes"
5 "encoding/json"
6 "net/http"
7 "strings"
8 "testing"
9 "time"
10)
11
12// TestSigV4Known checks our signer against AWS's published test vector
13// from the SigV4 documentation. The expected Authorization header below
14// is taken from "get-vanilla" in the suite (adapted: we use POST with
15// an empty body since Bedrock is always POST).
16//
17// We don't run the official suite (it lives in S3 as a tarball and the
18// signing rules vary slightly between examples); instead this test
19// pins a fixed input/output so future refactors can't silently break
20// the signer.
21func TestSigV4Deterministic(t *testing.T) {
22 creds := &bedrockSigV4Creds{
23 accessKeyID: "AKIDEXAMPLE",
24 secretAccessKey: "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
25 }
26 body := []byte(`{"hello":"world"}`)
27 req, err := http.NewRequest("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/model/foo/converse-stream", bytes.NewReader(body))
28 if err != nil {
29 t.Fatal(err)
30 }
31 req.Header.Set("content-type", "application/json")
32 when := time.Date(2024, 1, 15, 12, 0, 0, 0, time.UTC)
33 if err := signSigV4(req, body, "bedrock", "us-east-1", creds, when); err != nil {
34 t.Fatal(err)
35 }
36 auth := req.Header.Get("Authorization")
37 if !strings.HasPrefix(auth, "AWS4-HMAC-SHA256 ") {
38 t.Fatalf("missing AWS4 prefix: %q", auth)
39 }
40 if !strings.Contains(auth, "Credential=AKIDEXAMPLE/20240115/us-east-1/bedrock/aws4_request") {
41 t.Errorf("wrong credential scope: %q", auth)
42 }
43 if !strings.Contains(auth, "SignedHeaders=") {
44 t.Errorf("missing SignedHeaders: %q", auth)
45 }
46 if !strings.Contains(auth, "Signature=") {
47 t.Errorf("missing Signature: %q", auth)
48 }
49 // Sign again with identical inputs -> identical signature (determinism).
50 req2, _ := http.NewRequest("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/model/foo/converse-stream", bytes.NewReader(body))
51 req2.Header.Set("content-type", "application/json")
52 if err := signSigV4(req2, body, "bedrock", "us-east-1", creds, when); err != nil {
53 t.Fatal(err)
54 }
55 if req2.Header.Get("Authorization") != auth {
56 t.Errorf("signature not deterministic:\n a=%q\n b=%q", auth, req2.Header.Get("Authorization"))
57 }
58}
59
60func TestSigV4PathEncoding(t *testing.T) {
61 // Bedrock model IDs include colons and dots. The colon must encode
62 // to %3A in the canonical path (per SigV4 rules).
63 cases := map[string]string{
64 "/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse-stream": "/model/anthropic.claude-sonnet-4-5-20250929-v1%3A0/converse-stream",
65 "/": "/",
66 "": "/",
67 "/a": "/a",
68 }
69 for in, want := range cases {
70 got := canonicalSigV4Path(in)
71 if got != want {
72 t.Errorf("canonicalSigV4Path(%q) = %q, want %q", in, got, want)
73 }
74 }
75}
76
77func TestReadAWSCredentialsFile(t *testing.T) {
78 // Smoke test: parses an INI-ish file when present. We don't require
79 // the file to exist on the test host.
80 if _, err := readAWSCredentialsFile("default"); err != nil {
81 // missing file or missing profile is fine — both are non-panics
82 t.Logf("no aws creds available (expected on CI): %v", err)
83 }
84}
85
86func TestNormalizeBedrockToolResultsMovesResultsAdjacentToToolUse(t *testing.T) {
87 msgs := []Message{
88 {Role: RoleAssistant, Content: []Content{ToolCallBlock{ID: "tool-1", Name: "edit", Arguments: json.RawMessage(`{"path":"a"}`)}}},
89 {Role: RoleUser, Content: []Content{TextBlock{Text: "what's wrong"}}},
90 {Role: RoleTool, Content: []Content{ToolResultBlock{CallID: "tool-1", Content: []Content{TextBlock{Text: "edited"}}}}},
91 }
92
93 out := normalizeBedrockToolResults(msgs)
94 if len(out) != 3 {
95 t.Fatalf("got %d messages, want 3: %+v", len(out), out)
96 }
97 if out[0].Role != RoleAssistant {
98 t.Fatalf("first role = %s, want assistant", out[0].Role)
99 }
100 if out[1].Role != RoleTool {
101 t.Fatalf("second role = %s, want tool", out[1].Role)
102 }
103 tr, ok := out[1].Content[0].(ToolResultBlock)
104 if !ok {
105 t.Fatalf("second message content = %T, want ToolResultBlock", out[1].Content[0])
106 }
107 if tr.CallID != "tool-1" || tr.IsError {
108 t.Fatalf("unexpected tool result: %+v", tr)
109 }
110 if out[2].Role != RoleUser {
111 t.Fatalf("third role = %s, want user", out[2].Role)
112 }
113}
114
115func TestNormalizeBedrockToolResultsInjectsMissingResult(t *testing.T) {
116 msgs := []Message{
117 {Role: RoleAssistant, Content: []Content{ToolCallBlock{ID: "tool-1", Name: "edit", Arguments: json.RawMessage(`{}`)}}},
118 {Role: RoleUser, Content: []Content{TextBlock{Text: "what's wrong"}}},
119 }
120
121 out := normalizeBedrockToolResults(msgs)
122 if len(out) != 3 {
123 t.Fatalf("got %d messages, want 3: %+v", len(out), out)
124 }
125 tr, ok := out[1].Content[0].(ToolResultBlock)
126 if !ok {
127 t.Fatalf("second message content = %T, want ToolResultBlock", out[1].Content[0])
128 }
129 if tr.CallID != "tool-1" || !tr.IsError {
130 t.Fatalf("unexpected synthetic result: %+v", tr)
131 }
132}
133
134func TestBedrockModelSupportsCaching(t *testing.T) {
135 cases := []struct {
136 model string
137 want bool
138 }{
139 // Bare Claude IDs (as they come from catalog_builtin)
140 {"anthropic.claude-sonnet-4-5-20250929-v1:0", true},
141 {"anthropic.claude-haiku-4-5-20251001-v1:0", true},
142 {"anthropic.claude-opus-4-5-20251101-v1:0", true},
143 // Geo-prefixed (resolved form that arrives at bedrockModelSupportsCaching)
144 {"us.anthropic.claude-sonnet-4-5-20250929-v1:0", true},
145 {"eu.anthropic.claude-haiku-4-5-20251001-v1:0", true},
146 {"global.anthropic.claude-opus-4-5-20251101-v1:0", true},
147 // Nova models have PriceCacheWrite==0 — they use automatic caching
148 {"amazon.nova-pro-v1:0", false},
149 {"amazon.nova-lite-v1:0", false},
150 {"amazon.nova-micro-v1:0", false},
151 // DeepSeek — no cache write price
152 {"deepseek.r1-v1:0", false},
153 // Unknown model with claude prefix
154 {"anthropic.claude-future-v99:0", true},
155 // Unknown non-claude model
156 {"some.unknown-model-v1:0", false},
157 }
158 for _, c := range cases {
159 got := bedrockModelSupportsCaching(c.model)
160 if got != c.want {
161 t.Errorf("bedrockModelSupportsCaching(%q) = %v, want %v", c.model, got, c.want)
162 }
163 }
164}
165
166func TestBedrockBuildRequestCachingClaudeModel(t *testing.T) {
167 // A Claude model (PriceCacheWrite > 0) should get cachePoint markers
168 // in the system array and on the last user message.
169 client := &bedrockClient{region: "us-east-1"}
170 req, err := client.buildRequest(Request{
171 Model: "anthropic.claude-sonnet-4-5-20250929-v1:0",
172 System: "You are a helpful assistant.",
173 Messages: []Message{
174 {Role: RoleUser, Content: []Content{TextBlock{Text: "hello"}}},
175 },
176 })
177 if err != nil {
178 t.Fatal(err)
179 }
180
181 // system should be: [{text: ...}, {cachePoint: {type: default}}]
182 if len(req.System) != 2 {
183 t.Fatalf("system len = %d, want 2 (text + cachePoint)", len(req.System))
184 }
185 if _, ok := req.System[0]["text"]; !ok {
186 t.Errorf("system[0] missing text key: %v", req.System[0])
187 }
188 if _, ok := req.System[1]["cachePoint"]; !ok {
189 t.Errorf("system[1] missing cachePoint key: %v", req.System[1])
190 }
191
192 // last user message should end with a cachePoint block
193 if len(req.Messages) == 0 {
194 t.Fatal("no messages")
195 }
196 lastMsg := req.Messages[len(req.Messages)-1]
197 if lastMsg.Role != "user" {
198 t.Fatalf("last message role = %q, want user", lastMsg.Role)
199 }
200 lastBlock := lastMsg.Content[len(lastMsg.Content)-1]
201 if _, ok := lastBlock["cachePoint"]; !ok {
202 t.Errorf("last user message final block missing cachePoint: %v", lastBlock)
203 }
204}
205
206func TestBedrockBuildRequestNoCachingNovaModel(t *testing.T) {
207 // A Nova model (PriceCacheWrite == 0) should NOT get cachePoint markers.
208 client := &bedrockClient{region: "us-east-1"}
209 req, err := client.buildRequest(Request{
210 Model: "amazon.nova-pro-v1:0",
211 System: "You are helpful.",
212 Messages: []Message{
213 {Role: RoleUser, Content: []Content{TextBlock{Text: "hi"}}},
214 },
215 })
216 if err != nil {
217 t.Fatal(err)
218 }
219
220 // system should be plain: [{text: ...}] with no cachePoint
221 if len(req.System) != 1 {
222 t.Fatalf("system len = %d, want 1 (text only)", len(req.System))
223 }
224 if _, ok := req.System[0]["cachePoint"]; ok {
225 t.Errorf("Nova system unexpectedly contains cachePoint")
226 }
227
228 // last user message should NOT end with a cachePoint block
229 if len(req.Messages) == 0 {
230 t.Fatal("no messages")
231 }
232 lastMsg := req.Messages[len(req.Messages)-1]
233 for _, block := range lastMsg.Content {
234 if _, ok := block["cachePoint"]; ok {
235 t.Errorf("Nova user message unexpectedly contains cachePoint: %v", block)
236 }
237 }
238}
239
240func TestBedrockBuildRequestCachingMultiTurn(t *testing.T) {
241 // In a multi-turn conversation the cachePoint should be on the LAST
242 // user message only, not on earlier ones.
243 client := &bedrockClient{region: "us-east-1"}
244 req, err := client.buildRequest(Request{
245 Model: "anthropic.claude-sonnet-4-5-20250929-v1:0",
246 Messages: []Message{
247 {Role: RoleUser, Content: []Content{TextBlock{Text: "turn 1"}}},
248 {Role: RoleAssistant, Content: []Content{TextBlock{Text: "response 1"}}},
249 {Role: RoleUser, Content: []Content{TextBlock{Text: "turn 2"}}},
250 },
251 })
252 if err != nil {
253 t.Fatal(err)
254 }
255
256 // Find all user messages and verify only the last has a cachePoint.
257 var userMsgs []bedrockMessage
258 for _, m := range req.Messages {
259 if m.Role == "user" {
260 userMsgs = append(userMsgs, m)
261 }
262 }
263 if len(userMsgs) != 2 {
264 t.Fatalf("expected 2 user messages, got %d", len(userMsgs))
265 }
266
267 // First user message: no cachePoint
268 for _, block := range userMsgs[0].Content {
269 if _, ok := block["cachePoint"]; ok {
270 t.Errorf("first user message should not have cachePoint: %v", block)
271 }
272 }
273
274 // Last user message: ends with cachePoint
275 lastContent := userMsgs[len(userMsgs)-1].Content
276 lastBlock := lastContent[len(lastContent)-1]
277 if _, ok := lastBlock["cachePoint"]; !ok {
278 t.Errorf("last user message should end with cachePoint, got: %v", lastBlock)
279 }
280}
281
282func TestBedrockBuildRequestCachingNoSystemNoCrash(t *testing.T) {
283 // No system prompt: system array should be nil, not a bare cachePoint.
284 client := &bedrockClient{region: "us-east-1"}
285 req, err := client.buildRequest(Request{
286 Model: "anthropic.claude-sonnet-4-5-20250929-v1:0",
287 Messages: []Message{{Role: RoleUser, Content: []Content{TextBlock{Text: "hi"}}}},
288 })
289 if err != nil {
290 t.Fatal(err)
291 }
292 if len(req.System) != 0 {
293 t.Errorf("empty system prompt should produce nil system, got: %v", req.System)
294 }
295 // Message cachePoint still present
296 lastBlock := req.Messages[0].Content[len(req.Messages[0].Content)-1]
297 if _, ok := lastBlock["cachePoint"]; !ok {
298 t.Errorf("user message should still have cachePoint when system is empty")
299 }
300}
301
302func TestBedrockTagLastUserCache(t *testing.T) {
303 msgs := []bedrockMessage{
304 {Role: "user", Content: []map[string]interface{}{{"text": "hi"}}},
305 {Role: "assistant", Content: []map[string]interface{}{{"text": "hello"}}},
306 {Role: "user", Content: []map[string]interface{}{{"text": "followup"}}},
307 }
308 bedrockTagLastUserCache(msgs)
309
310 // Only the last user message (index 2) should have a cachePoint appended.
311 last := msgs[2].Content
312 if _, ok := last[len(last)-1]["cachePoint"]; !ok {
313 t.Errorf("last user message should end with cachePoint")
314 }
315 // The first user message should be untouched.
316 if len(msgs[0].Content) != 1 {
317 t.Errorf("first user message content len = %d, want 1", len(msgs[0].Content))
318 }
319 if _, ok := msgs[0].Content[0]["cachePoint"]; ok {
320 t.Errorf("first user message should not have cachePoint")
321 }
322}
323
324func TestBedrockTagLastUserCacheEmpty(t *testing.T) {
325 // Should not panic on empty or assistant-only history.
326 bedrockTagLastUserCache(nil)
327 bedrockTagLastUserCache([]bedrockMessage{})
328 msgs := []bedrockMessage{
329 {Role: "assistant", Content: []map[string]interface{}{{"text": "hi"}}},
330 }
331 bedrockTagLastUserCache(msgs) // should not panic, no user message to tag
332}
333
334func TestBedrockBuildRequestCachingWireJSON(t *testing.T) {
335 // Verify the JSON shape Bedrock actually receives has the right keys.
336 client := &bedrockClient{region: "us-east-1"}
337 breq, err := client.buildRequest(Request{
338 Model: "anthropic.claude-sonnet-4-5-20250929-v1:0",
339 System: "be helpful",
340 Messages: []Message{
341 {Role: RoleUser, Content: []Content{TextBlock{Text: "hello"}}},
342 },
343 })
344 if err != nil {
345 t.Fatal(err)
346 }
347 b, err := json.Marshal(breq)
348 if err != nil {
349 t.Fatal(err)
350 }
351 s := string(b)
352 if !strings.Contains(s, `"cachePoint"`) {
353 t.Errorf("serialised request missing cachePoint key: %s", s)
354 }
355 if !strings.Contains(s, `"type":"default"`) {
356 t.Errorf("serialised request missing cachePoint type:default: %s", s)
357 }
358}
359
360func TestBedrockBuildRequestSkipsEmptyToolMessages(t *testing.T) {
361 client := &bedrockClient{}
362 req, err := client.buildRequest(Request{Messages: []Message{
363 {Role: RoleTool, Content: []Content{ToolResultBlock{CallID: "missing", Content: []Content{TextBlock{Text: "orphan"}}}}},
364 {Role: RoleUser, Content: []Content{TextBlock{Text: "hello"}}},
365 }})
366 if err != nil {
367 t.Fatal(err)
368 }
369 if len(req.Messages) != 1 {
370 t.Fatalf("got %d bedrock messages, want 1: %+v", len(req.Messages), req.Messages)
371 }
372 if req.Messages[0].Role != "user" {
373 t.Fatalf("role = %s, want user", req.Messages[0].Role)
374 }
375}
376
377func TestBedrockEventPayloadHelpers(t *testing.T) {
378 wrapped := []byte(`{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"Hello"}}}`)
379 if got := bedrockEventTypeFromPayload(wrapped); got != "contentBlockDelta" {
380 t.Fatalf("event type = %q, want contentBlockDelta", got)
381 }
382 var delta struct {
383 ContentBlockIndex int `json:"contentBlockIndex"`
384 Delta struct {
385 Text string `json:"text"`
386 } `json:"delta"`
387 }
388 if err := unmarshalBedrockEventPayload(wrapped, "contentBlockDelta", &delta); err != nil {
389 t.Fatal(err)
390 }
391 if delta.ContentBlockIndex != 0 || delta.Delta.Text != "Hello" {
392 t.Fatalf("unexpected wrapped delta: %+v", delta)
393 }
394
395 direct := []byte(`{"contentBlockIndex":1,"delta":{"text":"world"}}`)
396 if err := unmarshalBedrockEventPayload(direct, "contentBlockDelta", &delta); err != nil {
397 t.Fatal(err)
398 }
399 if delta.ContentBlockIndex != 1 || delta.Delta.Text != "world" {
400 t.Fatalf("unexpected direct delta: %+v", delta)
401 }
402}
403
404func TestResolveBedrockInferenceProfileID(t *testing.T) {
405 cases := []struct {
406 model string
407 region string
408 want string
409 }{
410 // Bare Anthropic foundation IDs get the region-matched prefix.
411 {"anthropic.claude-sonnet-4-5-20250929-v1:0", "us-east-1", "us.anthropic.claude-sonnet-4-5-20250929-v1:0"},
412 {"anthropic.claude-sonnet-4-5-20250929-v1:0", "eu-central-1", "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"},
413 {"anthropic.claude-opus-4-6-v1", "ap-southeast-2", "apac.anthropic.claude-opus-4-6-v1"},
414 {"anthropic.claude-opus-4-6-v1", "us-gov-west-1", "us-gov.anthropic.claude-opus-4-6-v1"},
415 {"deepseek.r1-v1:0", "eu-west-1", "eu.deepseek.r1-v1:0"},
416 // Empty region defaults to us.
417 {"anthropic.claude-opus-4-6-v1", "", "us.anthropic.claude-opus-4-6-v1"},
418 // Already-prefixed IDs are left untouched.
419 {"eu.anthropic.claude-sonnet-4-5-20250929-v1:0", "us-east-1", "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"},
420 {"global.anthropic.claude-opus-4-6-v1", "us-east-1", "global.anthropic.claude-opus-4-6-v1"},
421 {"us.anthropic.claude-opus-4-6-v1", "eu-central-1", "us.anthropic.claude-opus-4-6-v1"},
422 // ARNs are passed through verbatim.
423 {"arn:aws:bedrock:us-east-1:123:inference-profile/us.anthropic.claude-opus-4-6-v1", "eu-west-1", "arn:aws:bedrock:us-east-1:123:inference-profile/us.anthropic.claude-opus-4-6-v1"},
424 // Families that don't need a profile are untouched.
425 {"amazon.nova-pro-v1:0", "us-east-1", "amazon.nova-pro-v1:0"},
426 {"", "us-east-1", ""},
427 }
428 for _, c := range cases {
429 if got := resolveBedrockInferenceProfileID(c.model, c.region); got != c.want {
430 t.Errorf("resolveBedrockInferenceProfileID(%q, %q) = %q; want %q", c.model, c.region, got, c.want)
431 }
432 }
433}