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 TestBedrockBuildRequestSkipsEmptyToolMessages(t *testing.T) {
135 client := &bedrockClient{}
136 req, err := client.buildRequest(Request{Messages: []Message{
137 {Role: RoleTool, Content: []Content{ToolResultBlock{CallID: "missing", Content: []Content{TextBlock{Text: "orphan"}}}}},
138 {Role: RoleUser, Content: []Content{TextBlock{Text: "hello"}}},
139 }})
140 if err != nil {
141 t.Fatal(err)
142 }
143 if len(req.Messages) != 1 {
144 t.Fatalf("got %d bedrock messages, want 1: %+v", len(req.Messages), req.Messages)
145 }
146 if req.Messages[0].Role != "user" {
147 t.Fatalf("role = %s, want user", req.Messages[0].Role)
148 }
149}
150
151func TestBedrockEventPayloadHelpers(t *testing.T) {
152 wrapped := []byte(`{"contentBlockDelta":{"contentBlockIndex":0,"delta":{"text":"Hello"}}}`)
153 if got := bedrockEventTypeFromPayload(wrapped); got != "contentBlockDelta" {
154 t.Fatalf("event type = %q, want contentBlockDelta", got)
155 }
156 var delta struct {
157 ContentBlockIndex int `json:"contentBlockIndex"`
158 Delta struct {
159 Text string `json:"text"`
160 } `json:"delta"`
161 }
162 if err := unmarshalBedrockEventPayload(wrapped, "contentBlockDelta", &delta); err != nil {
163 t.Fatal(err)
164 }
165 if delta.ContentBlockIndex != 0 || delta.Delta.Text != "Hello" {
166 t.Fatalf("unexpected wrapped delta: %+v", delta)
167 }
168
169 direct := []byte(`{"contentBlockIndex":1,"delta":{"text":"world"}}`)
170 if err := unmarshalBedrockEventPayload(direct, "contentBlockDelta", &delta); err != nil {
171 t.Fatal(err)
172 }
173 if delta.ContentBlockIndex != 1 || delta.Delta.Text != "world" {
174 t.Fatalf("unexpected direct delta: %+v", delta)
175 }
176}
177
178func TestResolveBedrockInferenceProfileID(t *testing.T) {
179 cases := []struct {
180 model string
181 region string
182 want string
183 }{
184 // Bare Anthropic foundation IDs get the region-matched prefix.
185 {"anthropic.claude-sonnet-4-5-20250929-v1:0", "us-east-1", "us.anthropic.claude-sonnet-4-5-20250929-v1:0"},
186 {"anthropic.claude-sonnet-4-5-20250929-v1:0", "eu-central-1", "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"},
187 {"anthropic.claude-opus-4-6-v1", "ap-southeast-2", "apac.anthropic.claude-opus-4-6-v1"},
188 {"anthropic.claude-opus-4-6-v1", "us-gov-west-1", "us-gov.anthropic.claude-opus-4-6-v1"},
189 {"deepseek.r1-v1:0", "eu-west-1", "eu.deepseek.r1-v1:0"},
190 // Empty region defaults to us.
191 {"anthropic.claude-opus-4-6-v1", "", "us.anthropic.claude-opus-4-6-v1"},
192 // Already-prefixed IDs are left untouched.
193 {"eu.anthropic.claude-sonnet-4-5-20250929-v1:0", "us-east-1", "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"},
194 {"global.anthropic.claude-opus-4-6-v1", "us-east-1", "global.anthropic.claude-opus-4-6-v1"},
195 {"us.anthropic.claude-opus-4-6-v1", "eu-central-1", "us.anthropic.claude-opus-4-6-v1"},
196 // ARNs are passed through verbatim.
197 {"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"},
198 // Families that don't need a profile are untouched.
199 {"amazon.nova-pro-v1:0", "us-east-1", "amazon.nova-pro-v1:0"},
200 {"", "us-east-1", ""},
201 }
202 for _, c := range cases {
203 if got := resolveBedrockInferenceProfileID(c.model, c.region); got != c.want {
204 t.Errorf("resolveBedrockInferenceProfileID(%q, %q) = %q; want %q", c.model, c.region, got, c.want)
205 }
206 }
207}