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
360// TestBedrockBuildRequestBtwToolHistory reproduces the /btw HTTP 400 bug:
361// the side-chat sends no tools but the frozen main transcript contains
362// toolUse / toolResult blocks. Bedrock rejects such a request unless
363// toolConfig is present. buildRequest must inject a stub toolConfig.
364func TestBedrockBuildRequestBtwToolHistory(t *testing.T) {
365 client := &bedrockClient{region: "us-east-1"}
366 req, err := client.buildRequest(Request{
367 Model: "amazon.nova-pro-v1:0",
368 // No tools — simulates /btw side-chat.
369 Messages: []Message{
370 // Frozen main transcript: one tool call + result already happened.
371 {Role: RoleAssistant, Content: []Content{
372 ToolCallBlock{ID: "tc-1", Name: "bash", Arguments: json.RawMessage(`{"command":"ls"}`)},
373 }},
374 {Role: RoleTool, Content: []Content{
375 ToolResultBlock{CallID: "tc-1", Content: []Content{TextBlock{Text: "file.go"}}},
376 }},
377 // Side-chat question appended by btwDialog.submit.
378 {Role: RoleUser, Content: []Content{TextBlock{Text: "what does that file do?"}}},
379 },
380 })
381 if err != nil {
382 t.Fatal(err)
383 }
384 if req.ToolConfig == nil {
385 t.Fatal("buildRequest should inject a stub toolConfig when history has tool blocks and req.Tools is empty, but ToolConfig is nil")
386 }
387 if len(req.ToolConfig.Tools) == 0 {
388 t.Fatal("stub toolConfig should have at least one tool")
389 }
390 // The stub must have a valid name and a non-nil schema so Bedrock won't
391 // reject it for a different reason.
392 stub := req.ToolConfig.Tools[0]
393 if stub.ToolSpec.Name == "" {
394 t.Error("stub tool name must not be empty")
395 }
396 if stub.ToolSpec.InputSchema.JSON == nil {
397 t.Error("stub tool schema must not be nil")
398 }
399 // Serialised wire payload must include toolConfig.
400 b, err := json.Marshal(req)
401 if err != nil {
402 t.Fatal(err)
403 }
404 if !strings.Contains(string(b), `"toolConfig"`) {
405 t.Errorf("serialised request missing toolConfig: %s", b)
406 }
407}
408
409// TestBedrockBuildRequestNoToolHistoryNoStub ensures that when there are
410// no tool blocks in the history and no tools in the request (ordinary
411// conversational call), we do NOT inject a toolConfig at all — the extra
412// field is unnecessary and some models may behave differently with it.
413func TestBedrockBuildRequestNoToolHistoryNoStub(t *testing.T) {
414 client := &bedrockClient{region: "us-east-1"}
415 req, err := client.buildRequest(Request{
416 Model: "amazon.nova-pro-v1:0",
417 Messages: []Message{
418 {Role: RoleUser, Content: []Content{TextBlock{Text: "hello"}}},
419 },
420 })
421 if err != nil {
422 t.Fatal(err)
423 }
424 if req.ToolConfig != nil {
425 t.Errorf("expected nil ToolConfig for plain conversation, got: %+v", req.ToolConfig)
426 }
427}
428
429func TestBedrockBuildRequestSkipsEmptyToolMessages(t *testing.T) {
430 client := &bedrockClient{}
431 req, err := client.buildRequest(Request{Messages: []Message{
432 {Role: RoleTool, Content: []Content{ToolResultBlock{CallID: "missing", Content: []Content{TextBlock{Text: "orphan"}}}}},
433 {Role: RoleUser, Content: []Content{TextBlock{Text: "hello"}}},
434 }})
435 if err != nil {
436 t.Fatal(err)
437 }
438 if len(req.Messages) != 1 {
439 t.Fatalf("got %d bedrock messages, want 1: %+v", len(req.Messages), req.Messages)
440 }
441 if req.Messages[0].Role != "user" {
442 t.Fatalf("role = %s, want user", req.Messages[0].Role)
443 }
444}
445
446func TestBedrockEventPayloadHelpers(t *testing.T) {
447 wrapped := []byte(`{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"Hello"}}}`)
448 if got := bedrockEventTypeFromPayload(wrapped); got != "contentBlockDelta" {
449 t.Fatalf("event type = %q, want contentBlockDelta", got)
450 }
451 var delta struct {
452 ContentBlockIndex int `json:"contentBlockIndex"`
453 Delta struct {
454 Text string `json:"text"`
455 } `json:"delta"`
456 }
457 if err := unmarshalBedrockEventPayload(wrapped, "contentBlockDelta", &delta); err != nil {
458 t.Fatal(err)
459 }
460 if delta.ContentBlockIndex != 0 || delta.Delta.Text != "Hello" {
461 t.Fatalf("unexpected wrapped delta: %+v", delta)
462 }
463
464 direct := []byte(`{"contentBlockIndex":1,"delta":{"text":"world"}}`)
465 if err := unmarshalBedrockEventPayload(direct, "contentBlockDelta", &delta); err != nil {
466 t.Fatal(err)
467 }
468 if delta.ContentBlockIndex != 1 || delta.Delta.Text != "world" {
469 t.Fatalf("unexpected direct delta: %+v", delta)
470 }
471}
472
473func TestResolveBedrockInferenceProfileID(t *testing.T) {
474 cases := []struct {
475 model string
476 region string
477 want string
478 }{
479 // Bare Anthropic foundation IDs get the region-matched prefix.
480 {"anthropic.claude-sonnet-4-5-20250929-v1:0", "us-east-1", "us.anthropic.claude-sonnet-4-5-20250929-v1:0"},
481 {"anthropic.claude-sonnet-4-5-20250929-v1:0", "eu-central-1", "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"},
482 {"anthropic.claude-opus-4-6-v1", "ap-southeast-2", "apac.anthropic.claude-opus-4-6-v1"},
483 {"anthropic.claude-opus-4-6-v1", "us-gov-west-1", "us-gov.anthropic.claude-opus-4-6-v1"},
484 {"deepseek.r1-v1:0", "eu-west-1", "eu.deepseek.r1-v1:0"},
485 // Empty region defaults to us.
486 {"anthropic.claude-opus-4-6-v1", "", "us.anthropic.claude-opus-4-6-v1"},
487 // Already-prefixed IDs are left untouched.
488 {"eu.anthropic.claude-sonnet-4-5-20250929-v1:0", "us-east-1", "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"},
489 {"global.anthropic.claude-opus-4-6-v1", "us-east-1", "global.anthropic.claude-opus-4-6-v1"},
490 {"us.anthropic.claude-opus-4-6-v1", "eu-central-1", "us.anthropic.claude-opus-4-6-v1"},
491 // ARNs are passed through verbatim.
492 {"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"},
493 // Families that don't need a profile are untouched.
494 {"amazon.nova-pro-v1:0", "us-east-1", "amazon.nova-pro-v1:0"},
495 {"", "us-east-1", ""},
496 }
497 for _, c := range cases {
498 if got := resolveBedrockInferenceProfileID(c.model, c.region); got != c.want {
499 t.Errorf("resolveBedrockInferenceProfileID(%q, %q) = %q; want %q", c.model, c.region, got, c.want)
500 }
501 }
502}