1package provider
2
3// Amazon Bedrock Converse-Stream client.
4//
5// Endpoint: POST https://bedrock-runtime.{region}.amazonaws.com/model/{id}/converse-stream
6//
7// Auth: two paths supported.
8//
9// - **Bearer token** via AWS_BEARER_TOKEN_BEDROCK. Modern, simple,
10// `bedrock:CallWithBearerToken` IAM permission required.
11// - **SigV4 signing** via AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY
12// (+ optional AWS_SESSION_TOKEN for temporary credentials, e.g. SSO
13// or IRSA-issued). Hand-rolled here so we don't pull in aws-sdk-go-v2.
14//
15// If both are present, bearer wins (simpler / fewer moving parts).
16//
17// Wire format: AWS event-stream binary framing. Each message is:
18//
19// [4-byte total length BE]
20// [4-byte headers length BE]
21// [4-byte CRC32 of the 8-byte prelude]
22// [headers ...] each = [1B name-length][name][1B value-type][2B value-length][value]
23// [payload ...] payload-length = total - headers - 4 (trailer) - 12 (prelude+prelude-crc)
24// [4-byte CRC32 of message bytes minus trailer]
25//
26// We only consume the stream — never produce it — so we don't validate
27// CRCs (the TLS+TCP transport already protects in-flight bytes). We do
28// validate length-bounds to avoid OOB reads on a hostile / corrupted
29// stream.
30//
31// Event types we care about (from `:event-type` header):
32// - messageStart { role: "assistant" }
33// - contentBlockStart { contentBlockIndex, start: { toolUse: { toolUseId, name } } }
34// - contentBlockDelta { contentBlockIndex, delta: { text?, toolUse?: { input } } }
35// - contentBlockStop { contentBlockIndex }
36// - messageStop { stopReason }
37// - metadata { usage: { inputTokens, outputTokens } }
38
39import (
40 "bytes"
41 "context"
42 "crypto/hmac"
43 "crypto/sha256"
44 "encoding/base64"
45 "encoding/binary"
46 "encoding/hex"
47 "encoding/json"
48 "fmt"
49 "io"
50 "net/http"
51 "net/url"
52 "os"
53 "os/exec"
54 "sort"
55 "strings"
56 "time"
57)
58
59type bedrockClient struct {
60 // Auth mode is determined at construction time. Exactly one of
61 // bearerToken or sigv4 is populated.
62 bearerToken string
63 sigv4 *bedrockSigV4Creds
64
65 region string
66 baseURL string
67 http *http.Client
68}
69
70// bedrockSigV4Creds holds the IAM credentials used to sign each request.
71type bedrockSigV4Creds struct {
72 accessKeyID string
73 secretAccessKey string
74 sessionToken string // optional; STS / SSO / IRSA temp creds
75}
76
77// NewBedrockClient returns a Bedrock client.
78//
79// Auth resolution (first match wins):
80//
81// 1. apiKey == real bearer-ish string (not "<aws>") -> bearer route.
82// 2. AWS_BEARER_TOKEN_BEDROCK env var -> bearer route.
83// 3. AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY (+ optional
84// AWS_SESSION_TOKEN) -> SigV4 route.
85// 4. AWS_PROFILE -> read ~/.aws/credentials, take that profile's keys.
86// 5. AWS_PROFILE -> `aws configure export-credentials` (SSO, assume-role,
87// credential_process) when the profile has no static keys.
88//
89// region defaults to us-east-1 unless AWS_REGION / AWS_DEFAULT_REGION is
90// set or the baseURL embeds a region.
91func NewBedrockClient(apiKey, baseURL string) Client {
92 region := bedrockResolveRegion()
93 if baseURL == "" {
94 baseURL = "https://bedrock-runtime." + region + ".amazonaws.com"
95 }
96 c := &bedrockClient{
97 region: region,
98 baseURL: strings.TrimRight(baseURL, "/"),
99 http: &http.Client{Timeout: 0},
100 }
101
102 bearer, sigv4 := resolveBedrockAuth(apiKey)
103 if bearer != "" {
104 c.bearerToken = bearer
105 return c
106 }
107 if sigv4 != nil {
108 c.sigv4 = sigv4
109 return c
110 }
111
112 return &unimplementedClient{
113 name: "amazon-bedrock",
114 hint: "no Bedrock credentials found (set AWS_BEARER_TOKEN_BEDROCK, AWS_ACCESS_KEY_ID+AWS_SECRET_ACCESS_KEY, or AWS_PROFILE)",
115 }
116}
117
118// bedrockResolveRegion returns the AWS region for Bedrock calls,
119// honoring AWS_REGION then AWS_DEFAULT_REGION, defaulting to us-east-1.
120func bedrockResolveRegion() string {
121 region := os.Getenv("AWS_REGION")
122 if region == "" {
123 region = os.Getenv("AWS_DEFAULT_REGION")
124 }
125 if region == "" {
126 region = "us-east-1"
127 }
128 return region
129}
130
131// resolveBedrockAuth determines Bedrock credentials from the same
132// sources NewBedrockClient uses, so discovery and inference share one
133// resolution path. Exactly one of (bearer, sigv4) is non-empty on
134// success; both are zero when no credentials are found.
135//
136// Resolution order (first match wins): explicit apiKey bearer, then
137// AWS_BEARER_TOKEN_BEDROCK, then AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY
138// env vars, then ~/.aws/credentials via AWS_PROFILE.
139func resolveBedrockAuth(apiKey string) (bearer string, sigv4 *bedrockSigV4Creds) {
140 token := apiKey
141 if token == "" || token == "<aws>" {
142 token = os.Getenv("AWS_BEARER_TOKEN_BEDROCK")
143 }
144 if token != "" && token != "<aws>" {
145 return token, nil
146 }
147
148 return "", resolveBedrockDiscoveryCreds()
149}
150
151// resolveBedrockDiscoveryCreds ignores inference bearer tokens: Bedrock's
152// control-plane listing endpoints require SigV4 even when inference uses a bearer.
153func resolveBedrockDiscoveryCreds() *bedrockSigV4Creds {
154 ak := os.Getenv("AWS_ACCESS_KEY_ID")
155 sk := os.Getenv("AWS_SECRET_ACCESS_KEY")
156 st := os.Getenv("AWS_SESSION_TOKEN")
157 if ak != "" && sk != "" {
158 return &bedrockSigV4Creds{accessKeyID: ak, secretAccessKey: sk, sessionToken: st}
159 }
160
161 if profile := os.Getenv("AWS_PROFILE"); profile != "" {
162 if creds, err := readAWSCredentialsFile(profile); err == nil {
163 return creds
164 }
165 }
166
167 // Last resort: ask the AWS CLI to resolve credentials for the active
168 // profile. This handles SSO, assumed roles, credential_process, and
169 // web-identity setups whose credentials are not stored as static keys
170 // in ~/.aws/credentials. Best-effort: absent CLI or an expired login
171 // simply yields no credentials.
172 if creds := resolveBedrockCredsViaCLI(); creds != nil {
173 return creds
174 }
175
176 return nil
177}
178
179// resolveBedrockCredsViaCLI shells out to `aws configure
180// export-credentials --format process`, which returns fully-resolved
181// credentials for the active profile regardless of how that profile is
182// configured (static keys, SSO, assume-role, credential_process). The
183// AWS CLI must be installed and the profile's login (if any) still
184// valid. Returns nil on any failure so the caller can fall through.
185func resolveBedrockCredsViaCLI() *bedrockSigV4Creds {
186 if _, err := exec.LookPath("aws"); err != nil {
187 return nil
188 }
189 ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
190 defer cancel()
191 // AWS_PROFILE / AWS_REGION already in the environment are inherited,
192 // so the CLI resolves the same profile zot was pointed at.
193 cmd := exec.CommandContext(ctx, "aws", "configure", "export-credentials", "--format", "process")
194 out, err := cmd.Output()
195 if err != nil {
196 return nil
197 }
198 var exported struct {
199 AccessKeyID string `json:"AccessKeyId"`
200 SecretAccessKey string `json:"SecretAccessKey"`
201 SessionToken string `json:"SessionToken"`
202 }
203 if err := json.Unmarshal(out, &exported); err != nil {
204 return nil
205 }
206 if exported.AccessKeyID == "" || exported.SecretAccessKey == "" {
207 return nil
208 }
209 return &bedrockSigV4Creds{
210 accessKeyID: exported.AccessKeyID,
211 secretAccessKey: exported.SecretAccessKey,
212 sessionToken: exported.SessionToken,
213 }
214}
215
216// readAWSCredentialsFile parses ~/.aws/credentials and returns the
217// access-key/secret-key (and optional session-token) for the named
218// profile. The file is an INI-like format with `[profile]` headers.
219func readAWSCredentialsFile(profile string) (*bedrockSigV4Creds, error) {
220 home, err := os.UserHomeDir()
221 if err != nil {
222 return nil, err
223 }
224 b, err := os.ReadFile(home + "/.aws/credentials")
225 if err != nil {
226 return nil, err
227 }
228 var current string
229 creds := map[string]map[string]string{}
230 for _, line := range strings.Split(string(b), "\n") {
231 line = strings.TrimSpace(line)
232 if line == "" || strings.HasPrefix(line, "#") {
233 continue
234 }
235 if strings.HasPrefix(line, "[") && strings.HasSuffix(line, "]") {
236 current = strings.TrimSpace(line[1 : len(line)-1])
237 creds[current] = map[string]string{}
238 continue
239 }
240 if current == "" {
241 continue
242 }
243 eq := strings.IndexByte(line, '=')
244 if eq < 0 {
245 continue
246 }
247 k := strings.TrimSpace(line[:eq])
248 v := strings.TrimSpace(line[eq+1:])
249 creds[current][k] = v
250 }
251 p, ok := creds[profile]
252 if !ok {
253 return nil, fmt.Errorf("aws profile %q not found in ~/.aws/credentials", profile)
254 }
255 ak := p["aws_access_key_id"]
256 sk := p["aws_secret_access_key"]
257 if ak == "" || sk == "" {
258 return nil, fmt.Errorf("aws profile %q missing aws_access_key_id or aws_secret_access_key", profile)
259 }
260 return &bedrockSigV4Creds{
261 accessKeyID: ak,
262 secretAccessKey: sk,
263 sessionToken: p["aws_session_token"],
264 }, nil
265}
266
267func (c *bedrockClient) Name() string { return "amazon-bedrock" }
268
269// ---- request building ----
270
271type bedrockMessage struct {
272 Role string `json:"role"`
273 Content []map[string]interface{} `json:"content"`
274}
275
276// bedrockCachePoint is the Converse API cache checkpoint content block.
277// Appending this to a system or message content array tells Bedrock to
278// create a cache boundary at that position in the prompt prefix.
279var bedrockCachePoint = map[string]interface{}{
280 "cachePoint": map[string]interface{}{"type": "default"},
281}
282
283type bedrockToolSpec struct {
284 ToolSpec struct {
285 Name string `json:"name"`
286 Description string `json:"description,omitempty"`
287 InputSchema struct {
288 JSON json.RawMessage `json:"json"`
289 } `json:"inputSchema"`
290 } `json:"toolSpec"`
291}
292
293type bedrockRequest struct {
294 Messages []bedrockMessage `json:"messages"`
295 System []map[string]interface{} `json:"system,omitempty"`
296 AdditionalModelRequestFields map[string]interface{} `json:"additionalModelRequestFields,omitempty"`
297 InferenceConfig struct {
298 MaxTokens int `json:"maxTokens,omitempty"`
299 Temperature *float32 `json:"temperature,omitempty"`
300 } `json:"inferenceConfig,omitempty"`
301 ToolConfig *struct {
302 Tools []bedrockToolSpec `json:"tools"`
303 } `json:"toolConfig,omitempty"`
304}
305
306func normalizeBedrockToolResults(msgs []Message) []Message {
307 resultByID := map[string]ToolResultBlock{}
308 for _, m := range msgs {
309 for _, c := range m.Content {
310 if tr, ok := c.(ToolResultBlock); ok {
311 if _, exists := resultByID[tr.CallID]; !exists {
312 resultByID[tr.CallID] = tr
313 }
314 }
315 }
316 }
317
318 out := make([]Message, 0, len(msgs))
319 for _, m := range msgs {
320 copy := m
321 copy.Content = nil
322 var toolCalls []ToolCallBlock
323 for _, c := range m.Content {
324 switch v := c.(type) {
325 case ToolResultBlock:
326 // Bedrock requires toolResult blocks immediately after the
327 // assistant toolUse they answer. Reinsert them from the
328 // assistant pass below instead of preserving their original
329 // location, which may be separated by user text in active
330 // sessions.
331 continue
332 case ToolCallBlock:
333 toolCalls = append(toolCalls, v)
334 }
335 copy.Content = append(copy.Content, c)
336 }
337 if len(copy.Content) > 0 {
338 out = append(out, copy)
339 }
340 if m.Role != RoleAssistant || len(toolCalls) == 0 {
341 continue
342 }
343
344 results := make([]Content, 0, len(toolCalls))
345 for _, tc := range toolCalls {
346 if tr, ok := resultByID[tc.ID]; ok {
347 results = append(results, tr)
348 continue
349 }
350 results = append(results, ToolResultBlock{
351 CallID: tc.ID,
352 IsError: true,
353 Content: []Content{TextBlock{Text: "tool call did not complete before the next user message"}},
354 })
355 }
356 out = append(out, Message{Role: RoleTool, Content: results, Time: m.Time})
357 }
358 return out
359}
360
361// bedrockModelSupportsCaching reports whether the resolved model ID
362// supports explicit prompt caching via cachePoint markers on Bedrock.
363// We use PriceCacheWrite > 0 as a proxy: every Bedrock-hosted Claude
364// model with a write price in the catalog supports cachePoint markers.
365// Nova models use automatic caching and don't need explicit markers.
366func bedrockModelSupportsCaching(modelID string) bool {
367 // Strip geo prefix (us./eu./apac./au./global.) before catalog lookup.
368 for _, p := range bedrockGeoPrefixes {
369 if strings.HasPrefix(modelID, p+".") {
370 modelID = modelID[len(p)+1:]
371 break
372 }
373 }
374 if m, err := FindModel("amazon-bedrock", modelID); err == nil {
375 return m.PriceCacheWrite > 0
376 }
377 // Unknown model: enable for Anthropic Claude families — cachePoint is
378 // silently ignored by the API if the model doesn't support it.
379 return strings.HasPrefix(modelID, "anthropic.claude-")
380}
381
382// bedrockMessagesHaveToolBlocks reports whether any message in msgs
383// contains a toolUse or toolResult content block. Bedrock's Converse API
384// rejects the request with HTTP 400 if those block types are present but
385// toolConfig is absent.
386func bedrockMessagesHaveToolBlocks(msgs []bedrockMessage) bool {
387 for _, m := range msgs {
388 for _, c := range m.Content {
389 if _, ok := c["toolUse"]; ok {
390 return true
391 }
392 if _, ok := c["toolResult"]; ok {
393 return true
394 }
395 }
396 }
397 return false
398}
399
400func (c *bedrockClient) buildRequest(req Request) (*bedrockRequest, error) {
401 out := &bedrockRequest{}
402
403 // Resolve the model ID as it will appear on the wire so the caching
404 // check operates on the same ID used for FindModel.
405 resolvedModel := resolveBedrockInferenceProfileID(req.Model, c.region)
406 caching := bedrockModelSupportsCaching(resolvedModel)
407 model := Model{ID: resolvedModel}
408 if catalogModel, err := FindModel("amazon-bedrock", resolvedModel); err == nil {
409 model = catalogModel
410 }
411 reasoning := ClampReasoningForModel(model, req.Reasoning)
412 adaptive := usesAdaptiveThinking(model)
413
414 if req.System != "" {
415 sysBlock := map[string]interface{}{"text": req.System}
416 if caching {
417 // Append cachePoint after the system text so the stable system
418 // prompt is cached as the first breakpoint.
419 out.System = []map[string]interface{}{sysBlock, bedrockCachePoint}
420 } else {
421 out.System = []map[string]interface{}{sysBlock}
422 }
423 }
424 if !adaptive {
425 out.InferenceConfig.Temperature = req.Temperature
426 }
427 if adaptive && reasoning != "" {
428 out.AdditionalModelRequestFields = map[string]interface{}{
429 "thinking": map[string]interface{}{"type": "adaptive"},
430 "output_config": map[string]interface{}{"effort": AnthropicAdaptiveEffort(reasoning)},
431 }
432 }
433 out.InferenceConfig.MaxTokens = req.MaxTokens
434 if out.InferenceConfig.MaxTokens == 0 {
435 out.InferenceConfig.MaxTokens = 4096
436 }
437 for _, m := range normalizeBedrockToolResults(req.Messages) {
438 role := string(m.Role)
439 if role == "tool" {
440 role = "user"
441 }
442 bm := bedrockMessage{Role: role}
443 for _, c := range m.Content {
444 switch v := c.(type) {
445 case TextBlock:
446 bm.Content = append(bm.Content, map[string]interface{}{"text": v.Text})
447 case ToolCallBlock:
448 var input map[string]interface{}
449 _ = json.Unmarshal(v.Arguments, &input)
450 if input == nil {
451 input = map[string]interface{}{}
452 }
453 bm.Content = append(bm.Content, map[string]interface{}{
454 "toolUse": map[string]interface{}{
455 "toolUseId": v.ID, "name": v.Name, "input": input,
456 },
457 })
458 case ImageBlock:
459 data, mime := anthShrinkImageBytesIfTooBig(v.Data, v.MimeType)
460 // Bedrock Converse uses the bare format name (jpeg/png/gif/webp).
461 imgFmt := strings.TrimPrefix(mime, "image/")
462 bm.Content = append(bm.Content, map[string]interface{}{
463 "image": map[string]interface{}{
464 "format": imgFmt,
465 "source": map[string]interface{}{
466 "bytes": base64.StdEncoding.EncodeToString(data),
467 },
468 },
469 })
470 case ToolResultBlock:
471 var resultContent []map[string]interface{}
472 for _, inner := range v.Content {
473 switch iv := inner.(type) {
474 case TextBlock:
475 resultContent = append(resultContent, map[string]interface{}{"text": iv.Text})
476 case ImageBlock:
477 data, mime := anthShrinkImageBytesIfTooBig(iv.Data, iv.MimeType)
478 imgFmt := strings.TrimPrefix(mime, "image/")
479 resultContent = append(resultContent, map[string]interface{}{
480 "image": map[string]interface{}{
481 "format": imgFmt,
482 "source": map[string]interface{}{
483 "bytes": base64.StdEncoding.EncodeToString(data),
484 },
485 },
486 })
487 }
488 }
489 status := "success"
490 if v.IsError {
491 status = "error"
492 }
493 bm.Content = append(bm.Content, map[string]interface{}{
494 "toolResult": map[string]interface{}{
495 "toolUseId": v.CallID,
496 "content": resultContent,
497 "status": status,
498 },
499 })
500 }
501 }
502 if len(bm.Content) == 0 {
503 continue
504 }
505 out.Messages = append(out.Messages, bm)
506 }
507 activeTools := activeToolDefinitions(req.Tools, req.Messages)
508 if len(activeTools) > 0 {
509 tc := struct {
510 Tools []bedrockToolSpec `json:"tools"`
511 }{}
512 for _, t := range activeTools {
513 var ts bedrockToolSpec
514 ts.ToolSpec.Name = t.Name
515 ts.ToolSpec.Description = t.Description
516 ts.ToolSpec.InputSchema.JSON = t.Schema
517 tc.Tools = append(tc.Tools, ts)
518 }
519 out.ToolConfig = &tc
520 } else if bedrockMessagesHaveToolBlocks(out.Messages) {
521 // Bedrock requires toolConfig to be present whenever the message
522 // history contains toolUse or toolResult content blocks, even if
523 // the current request does not intend to call any tools (e.g. the
524 // /btw side-chat sends the frozen main transcript which may include
525 // prior tool turns). Inject a stub tool so the API accepts the
526 // request; it will never be invoked since no tool_use stop reason
527 // can be triggered when the real tool list is empty.
528 var stub bedrockToolSpec
529 stub.ToolSpec.Name = "_placeholder"
530 stub.ToolSpec.Description = "placeholder required by Bedrock when tool history is present"
531 stub.ToolSpec.InputSchema.JSON = json.RawMessage(`{"type":"object","properties":{}}`)
532 out.ToolConfig = &struct {
533 Tools []bedrockToolSpec `json:"tools"`
534 }{Tools: []bedrockToolSpec{stub}}
535 }
536
537 if caching {
538 // Tag the last user message with a cachePoint. This extends the
539 // cached prefix to cover the full conversation history up to the
540 // current turn, so the next turn reads that history cheaply.
541 bedrockTagLastUserCache(out.Messages)
542 }
543
544 return out, nil
545}
546
547// bedrockTagLastUserCache appends a cachePoint block to the last user
548// message in the Bedrock message list. It is the Bedrock equivalent of
549// Anthropic's cache_control:{type:"ephemeral"} on the last user message.
550func bedrockTagLastUserCache(msgs []bedrockMessage) {
551 for i := len(msgs) - 1; i >= 0; i-- {
552 if msgs[i].Role == "user" {
553 msgs[i].Content = append(msgs[i].Content, bedrockCachePoint)
554 return
555 }
556 }
557}
558
559// resolveBedrockInferenceProfileID maps a bare foundation-model ID to
560// its region-matched cross-region inference-profile ID.
561//
562// Several newer Bedrock models (Anthropic Claude 4.x, DeepSeek, etc.)
563// cannot be invoked with on-demand throughput by their plain
564// foundation-model ID; Bedrock returns HTTP 400 demanding "the ID or
565// ARN of an inference profile that contains this model". The profile
566// ID is the same model ID with a geographic prefix (us/eu/apac/...).
567//
568// We only rewrite IDs that (a) lack an existing geo prefix and (b)
569// belong to a model family that requires a profile. IDs that already
570// carry a prefix (e.g. "eu.anthropic...", "global.anthropic...") or
571// fully-qualified ARNs are returned unchanged, so explicit user
572// choices and custom application inference profiles still work.
573func resolveBedrockInferenceProfileID(modelID, region string) string {
574 if modelID == "" {
575 return modelID
576 }
577 // ARNs are already inference-profile references; leave untouched.
578 if strings.HasPrefix(modelID, "arn:") {
579 return modelID
580 }
581 // Already geo-prefixed (us. / eu. / apac. / ap. / us-gov. / global.)?
582 if bedrockHasGeoPrefix(modelID) {
583 return modelID
584 }
585 // Opus 5 is exposed for on-demand throughput through the global
586 // inference profile rather than a region-specific profile.
587 if modelID == "anthropic.claude-opus-5" {
588 return "global." + modelID
589 }
590 if !bedrockRequiresInferenceProfile(modelID) {
591 return modelID
592 }
593 prefix := bedrockGeoPrefixForRegion(region)
594 if prefix == "" {
595 return modelID
596 }
597 return prefix + "." + modelID
598}
599
600// bedrockGeoPrefixes are the cross-region inference-profile geo
601// prefixes Bedrock uses. A model ID that starts with one of these
602// (followed by a dot) is already a profile reference.
603var bedrockGeoPrefixes = []string{"us-gov", "us", "eu", "apac", "ap", "global", "au"}
604
605func bedrockHasGeoPrefix(modelID string) bool {
606 for _, p := range bedrockGeoPrefixes {
607 if strings.HasPrefix(modelID, p+".") {
608 return true
609 }
610 }
611 return false
612}
613
614// bedrockRequiresInferenceProfile reports whether a bare
615// foundation-model ID is one of the families AWS only exposes through
616// a cross-region inference profile for on-demand throughput.
617func bedrockRequiresInferenceProfile(modelID string) bool {
618 switch {
619 case strings.HasPrefix(modelID, "anthropic.claude-"):
620 return true
621 case strings.HasPrefix(modelID, "deepseek."):
622 return true
623 default:
624 return false
625 }
626}
627
628// bedrockGeoPrefixForRegion maps an AWS region to the geo prefix used
629// by its cross-region inference profiles. Returns "" when the region
630// has no known mapping, in which case the model ID is left unchanged.
631func bedrockGeoPrefixForRegion(region string) string {
632 switch {
633 case region == "":
634 return "us"
635 case strings.HasPrefix(region, "us-gov-"):
636 return "us-gov"
637 case strings.HasPrefix(region, "us-"):
638 return "us"
639 case strings.HasPrefix(region, "eu-"):
640 return "eu"
641 case strings.HasPrefix(region, "ap-"):
642 return "apac"
643 case strings.HasPrefix(region, "ca-"):
644 return "us"
645 case strings.HasPrefix(region, "sa-"):
646 return "us"
647 default:
648 return "us"
649 }
650}
651
652func (c *bedrockClient) Stream(ctx context.Context, req Request) (<-chan Event, error) {
653 wire, err := c.buildRequest(req)
654 if err != nil {
655 return nil, err
656 }
657 body, err := json.Marshal(wire)
658 if err != nil {
659 return nil, err
660 }
661 modelID := resolveBedrockInferenceProfileID(req.Model, c.region)
662 url := c.baseURL + "/model/" + modelID + "/converse-stream"
663 httpReq, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
664 if err != nil {
665 return nil, err
666 }
667 httpReq.Header.Set("content-type", "application/json")
668 httpReq.Header.Set("accept", "application/vnd.amazon.eventstream")
669 if c.bearerToken != "" {
670 httpReq.Header.Set("authorization", "Bearer "+c.bearerToken)
671 } else if c.sigv4 != nil {
672 if err := signSigV4(httpReq, body, "bedrock", c.region, c.sigv4, time.Now().UTC()); err != nil {
673 return nil, fmt.Errorf("bedrock: sigv4 sign: %w", err)
674 }
675 } else {
676 return nil, fmt.Errorf("bedrock: no auth configured")
677 }
678 resp, err := c.http.Do(httpReq)
679 if err != nil {
680 return nil, fmt.Errorf("bedrock: %w", err)
681 }
682 if resp.StatusCode != http.StatusOK {
683 b, _ := io.ReadAll(resp.Body)
684 resp.Body.Close()
685 msg := strings.TrimSpace(string(b))
686 // A 403 on the bearer route is almost always a region mismatch:
687 // short-term Bedrock API keys are scoped to the region of the
688 // console session that minted them, but zot defaults to
689 // us-east-1. Surface the resolved region and the fix so the user
690 // is not left guessing why a freshly-copied key is "invalid".
691 if resp.StatusCode == http.StatusForbidden && c.bearerToken != "" {
692 return nil, fmt.Errorf(
693 "bedrock: http 403 (region=%s): %s\nhint: Bedrock API keys are region-scoped. If your key was created in another region, set AWS_REGION (e.g. AWS_REGION=eu-central-1) or pass --base-url https://bedrock-runtime.<region>.amazonaws.com",
694 c.region, msg)
695 }
696 return nil, fmt.Errorf("bedrock: http %d: %s", resp.StatusCode, msg)
697 }
698 out := make(chan Event, 16)
699 go c.runStream(ctx, resp, req, out)
700 return out, nil
701}
702
703func (c *bedrockClient) runStream(ctx context.Context, resp *http.Response, req Request, out chan<- Event) {
704 defer close(out)
705 defer resp.Body.Close()
706
707 out <- EventStart{Provider: "amazon-bedrock", Model: req.Model}
708
709 // State accumulated across deltas.
710 contentBlocks := map[int]*bedrockBlockState{}
711 stop := StopEnd
712 finalMsg := Message{Role: RoleAssistant, Time: time.Now()}
713 var usage Usage
714
715 for {
716 if ctx.Err() != nil {
717 stop = StopAborted
718 break
719 }
720 evt, err := readEventStreamMessage(resp.Body)
721 if err != nil {
722 if err == io.EOF {
723 break
724 }
725 out <- EventDone{Stop: StopError, Err: err, Message: finalMsg}
726 return
727 }
728 eventType := evt.headerString(":event-type")
729 if eventType == "" {
730 eventType = bedrockEventTypeFromPayload(evt.payload)
731 }
732 messageType := evt.headerString(":message-type")
733 if messageType == "exception" {
734 out <- EventDone{Stop: StopError, Err: fmt.Errorf("bedrock exception (%s): %s", evt.headerString(":exception-type"), string(evt.payload)), Message: finalMsg}
735 return
736 }
737 switch eventType {
738 case "messageStart":
739 // nothing to do; role is always "assistant"
740 case "contentBlockStart":
741 var d struct {
742 ContentBlockIndex int `json:"contentBlockIndex"`
743 Start struct {
744 ToolUse *struct {
745 ToolUseID string `json:"toolUseId"`
746 Name string `json:"name"`
747 } `json:"toolUse"`
748 } `json:"start"`
749 }
750 if err := unmarshalBedrockEventPayload(evt.payload, "contentBlockStart", &d); err != nil {
751 continue
752 }
753 st := &bedrockBlockState{}
754 contentBlocks[d.ContentBlockIndex] = st
755 if d.Start.ToolUse != nil {
756 st.isToolUse = true
757 st.toolID = d.Start.ToolUse.ToolUseID
758 st.toolName = d.Start.ToolUse.Name
759 out <- EventToolStart{ID: st.toolID, Name: st.toolName}
760 }
761 case "contentBlockDelta":
762 var d struct {
763 ContentBlockIndex int `json:"contentBlockIndex"`
764 Delta struct {
765 Text string `json:"text"`
766 ToolUse *struct {
767 Input string `json:"input"`
768 } `json:"toolUse"`
769 } `json:"delta"`
770 }
771 if err := unmarshalBedrockEventPayload(evt.payload, "contentBlockDelta", &d); err != nil {
772 continue
773 }
774 st := contentBlocks[d.ContentBlockIndex]
775 if st == nil {
776 st = &bedrockBlockState{}
777 contentBlocks[d.ContentBlockIndex] = st
778 }
779 if d.Delta.Text != "" {
780 st.text.WriteString(d.Delta.Text)
781 out <- EventTextDelta{Delta: d.Delta.Text}
782 }
783 if d.Delta.ToolUse != nil && d.Delta.ToolUse.Input != "" {
784 st.toolArgs.WriteString(d.Delta.ToolUse.Input)
785 out <- EventToolArgs{ID: st.toolID, Delta: d.Delta.ToolUse.Input}
786 }
787 case "contentBlockStop":
788 var d struct {
789 ContentBlockIndex int `json:"contentBlockIndex"`
790 }
791 if err := unmarshalBedrockEventPayload(evt.payload, "contentBlockStop", &d); err != nil {
792 continue
793 }
794 st := contentBlocks[d.ContentBlockIndex]
795 if st == nil {
796 continue
797 }
798 if st.isToolUse {
799 args := st.toolArgs.String()
800 if args == "" {
801 args = "{}"
802 }
803 finalMsg.Content = append(finalMsg.Content, ToolCallBlock{
804 ID: st.toolID, Name: st.toolName, Arguments: json.RawMessage(args),
805 })
806 out <- EventToolEnd{ID: st.toolID}
807 } else if st.text.Len() > 0 {
808 finalMsg.Content = append(finalMsg.Content, TextBlock{Text: st.text.String()})
809 }
810 case "messageStop":
811 var d struct {
812 StopReason string `json:"stopReason"`
813 }
814 _ = unmarshalBedrockEventPayload(evt.payload, "messageStop", &d)
815 switch d.StopReason {
816 case "tool_use":
817 stop = StopToolUse
818 case "end_turn":
819 stop = StopEnd
820 case "max_tokens":
821 stop = StopLength
822 default:
823 stop = StopEnd
824 }
825 case "metadata":
826 var d struct {
827 Usage struct {
828 InputTokens int `json:"inputTokens"`
829 OutputTokens int `json:"outputTokens"`
830 CacheReadInputTokens int `json:"cacheReadInputTokens"`
831 CacheWriteInputTokens int `json:"cacheWriteInputTokens"`
832 } `json:"usage"`
833 }
834 if err := unmarshalBedrockEventPayload(evt.payload, "metadata", &d); err == nil {
835 usage.InputTokens = d.Usage.InputTokens
836 usage.OutputTokens = d.Usage.OutputTokens
837 usage.CacheReadTokens = d.Usage.CacheReadInputTokens
838 usage.CacheWriteTokens = d.Usage.CacheWriteInputTokens
839 if m, err := FindModel("amazon-bedrock", req.Model); err == nil {
840 usage.CostUSD = ComputeCost(m, usage)
841 }
842 out <- EventUsage{Usage: usage}
843 }
844 }
845 }
846 out <- EventDone{Stop: stop, Message: finalMsg}
847}
848
849type bedrockBlockState struct {
850 isToolUse bool
851 toolID string
852 toolName string
853 toolArgs strings.Builder
854 text strings.Builder
855}
856
857func bedrockEventTypeFromPayload(payload []byte) string {
858 var outer map[string]json.RawMessage
859 if err := json.Unmarshal(payload, &outer); err != nil {
860 return ""
861 }
862 for _, name := range []string{"messageStart", "contentBlockStart", "contentBlockDelta", "contentBlockStop", "messageStop", "metadata"} {
863 if _, ok := outer[name]; ok {
864 return name
865 }
866 }
867 return ""
868}
869
870func unmarshalBedrockEventPayload(payload []byte, eventType string, dst any) error {
871 var outer map[string]json.RawMessage
872 if err := json.Unmarshal(payload, &outer); err == nil {
873 if wrapped, ok := outer[eventType]; ok {
874 return json.Unmarshal(wrapped, dst)
875 }
876 }
877 if err := json.Unmarshal(payload, dst); err != nil {
878 return fmt.Errorf("bedrock: parse %s payload: %w", eventType, err)
879 }
880 return nil
881}
882
883// ---- event-stream binary framing parser ----
884
885type eventStreamMessage struct {
886 headers map[string]string
887 payload []byte
888}
889
890func (m *eventStreamMessage) headerString(name string) string { return m.headers[name] }
891
892// readEventStreamMessage reads one event-stream message from r. Returns
893// io.EOF when the stream is finished cleanly.
894func readEventStreamMessage(r io.Reader) (*eventStreamMessage, error) {
895 var prelude [8]byte
896 if _, err := io.ReadFull(r, prelude[:]); err != nil {
897 return nil, err
898 }
899 totalLen := binary.BigEndian.Uint32(prelude[0:4])
900 headersLen := binary.BigEndian.Uint32(prelude[4:8])
901 if totalLen < 16 || headersLen > totalLen-16 {
902 return nil, fmt.Errorf("bedrock: malformed event-stream prelude (total=%d headers=%d)", totalLen, headersLen)
903 }
904 // Skip prelude CRC (4 bytes); we don't validate.
905 var preludeCRC [4]byte
906 if _, err := io.ReadFull(r, preludeCRC[:]); err != nil {
907 return nil, err
908 }
909 headersBuf := make([]byte, headersLen)
910 if _, err := io.ReadFull(r, headersBuf); err != nil {
911 return nil, err
912 }
913 payloadLen := int(totalLen) - 16 - int(headersLen)
914 payload := make([]byte, payloadLen)
915 if _, err := io.ReadFull(r, payload); err != nil {
916 return nil, err
917 }
918 var trailer [4]byte
919 if _, err := io.ReadFull(r, trailer[:]); err != nil {
920 return nil, err
921 }
922 // Parse headers.
923 hdrs := map[string]string{}
924 i := 0
925 for i < len(headersBuf) {
926 if i+1 > len(headersBuf) {
927 break
928 }
929 nameLen := int(headersBuf[i])
930 i++
931 if i+nameLen > len(headersBuf) {
932 return nil, fmt.Errorf("bedrock: header name overflow")
933 }
934 name := string(headersBuf[i : i+nameLen])
935 i += nameLen
936 if i+1 > len(headersBuf) {
937 return nil, fmt.Errorf("bedrock: header value-type missing")
938 }
939 valueType := headersBuf[i]
940 i++
941 // Type 7 = string with 2-byte length. Type 6 = byte-array similar.
942 // Other types (bool/int/long/timestamp/uuid) we treat as opaque
943 // strings since bedrock's event-type headers are always type 7.
944 switch valueType {
945 case 7, 6:
946 if i+2 > len(headersBuf) {
947 return nil, fmt.Errorf("bedrock: header value length missing")
948 }
949 vLen := int(binary.BigEndian.Uint16(headersBuf[i : i+2]))
950 i += 2
951 if i+vLen > len(headersBuf) {
952 return nil, fmt.Errorf("bedrock: header value overflow")
953 }
954 hdrs[name] = string(headersBuf[i : i+vLen])
955 i += vLen
956 case 0: // bool true
957 hdrs[name] = "true"
958 case 1: // bool false
959 hdrs[name] = "false"
960 case 2: // byte
961 if i+1 > len(headersBuf) {
962 return nil, fmt.Errorf("bedrock: header byte missing")
963 }
964 i++
965 case 3: // int16
966 if i+2 > len(headersBuf) {
967 return nil, fmt.Errorf("bedrock: header int16 missing")
968 }
969 i += 2
970 case 4: // int32
971 if i+4 > len(headersBuf) {
972 return nil, fmt.Errorf("bedrock: header int32 missing")
973 }
974 i += 4
975 case 5: // int64
976 if i+8 > len(headersBuf) {
977 return nil, fmt.Errorf("bedrock: header int64 missing")
978 }
979 i += 8
980 case 8: // timestamp (int64)
981 if i+8 > len(headersBuf) {
982 return nil, fmt.Errorf("bedrock: header timestamp missing")
983 }
984 i += 8
985 case 9: // uuid (16 bytes)
986 if i+16 > len(headersBuf) {
987 return nil, fmt.Errorf("bedrock: header uuid missing")
988 }
989 i += 16
990 default:
991 return nil, fmt.Errorf("bedrock: unknown header value-type %d", valueType)
992 }
993 }
994 return &eventStreamMessage{headers: hdrs, payload: payload}, nil
995}
996
997// ---- SigV4 signing ----
998//
999// AWS Signature Version 4. Implemented from scratch so we don't pull in
1000// aws-sdk-go-v2 (large transitive dep tree just for one signer). The
1001// algorithm is documented at:
1002// https://docs.aws.amazon.com/general/latest/gr/sigv4-signed-request-examples.html
1003//
1004// Steps:
1005//
1006// 1. Build the canonical request: method + path + canonical query +
1007// canonical headers + signed headers + payload-hash.
1008// 2. Build the string-to-sign: algorithm + amz-date + credential scope +
1009// sha256(canonical request).
1010// 3. Derive the signing key by chaining HMAC-SHA256:
1011// kSecret -> kDate -> kRegion -> kService -> kSigning
1012// 4. signature = HMAC-SHA256(kSigning, stringToSign)
1013// 5. Authorization header: AWS4-HMAC-SHA256 Credential=..., SignedHeaders=..., Signature=...
1014//
1015// Payload hash is sha256 of the JSON body. For Bedrock's streaming
1016// endpoint this works because the request body is fully buffered before
1017// the call — we don't need the streaming-payload variant.
1018
1019func sigHMAC(key, msg []byte) []byte {
1020 m := hmac.New(sha256.New, key)
1021 m.Write(msg)
1022 return m.Sum(nil)
1023}
1024
1025func sigSHA256Hex(b []byte) string {
1026 h := sha256.Sum256(b)
1027 return hex.EncodeToString(h[:])
1028}
1029
1030func signSigV4(req *http.Request, payload []byte, service, region string, creds *bedrockSigV4Creds, now time.Time) error {
1031 amzDate := now.Format("20060102T150405Z")
1032 shortDate := now.Format("20060102")
1033 credentialScope := shortDate + "/" + region + "/" + service + "/aws4_request"
1034
1035 // Required headers.
1036 req.Header.Set("host", req.URL.Host)
1037 req.Header.Set("x-amz-date", amzDate)
1038 if creds.sessionToken != "" {
1039 req.Header.Set("x-amz-security-token", creds.sessionToken)
1040 }
1041 payloadHash := sigSHA256Hex(payload)
1042 req.Header.Set("x-amz-content-sha256", payloadHash)
1043
1044 // Canonical headers: lowercase name, trimmed value, sorted by name,
1045 // each on its own line ending with \n.
1046 var names []string
1047 values := map[string]string{}
1048 for k, vs := range req.Header {
1049 lower := strings.ToLower(k)
1050 // Skip authorization (we're computing it) and any
1051 // connection-specific headers.
1052 if lower == "authorization" {
1053 continue
1054 }
1055 names = append(names, lower)
1056 values[lower] = strings.Join(vs, ",")
1057 }
1058 // host is special — it's not in req.Header by default; set above.
1059 sort.Strings(names)
1060 var canonHdrs strings.Builder
1061 for _, n := range names {
1062 canonHdrs.WriteString(n)
1063 canonHdrs.WriteByte(':')
1064 canonHdrs.WriteString(strings.TrimSpace(values[n]))
1065 canonHdrs.WriteByte('\n')
1066 }
1067 signedHeaders := strings.Join(names, ";")
1068
1069 // Canonical query string: keys sorted by name, each key and value
1070 // URI-escaped (RFC 3986).
1071 q := req.URL.Query()
1072 var qKeys []string
1073 for k := range q {
1074 qKeys = append(qKeys, k)
1075 }
1076 sort.Strings(qKeys)
1077 var canonQuery strings.Builder
1078 for i, k := range qKeys {
1079 if i > 0 {
1080 canonQuery.WriteByte('&')
1081 }
1082 canonQuery.WriteString(url.QueryEscape(k))
1083 canonQuery.WriteByte('=')
1084 canonQuery.WriteString(url.QueryEscape(q.Get(k)))
1085 }
1086
1087 // Canonical request.
1088 canonReq := req.Method + "\n" +
1089 canonicalSigV4Path(req.URL.Path) + "\n" +
1090 canonQuery.String() + "\n" +
1091 canonHdrs.String() + "\n" +
1092 signedHeaders + "\n" +
1093 payloadHash
1094
1095 // String to sign.
1096 stringToSign := "AWS4-HMAC-SHA256\n" + amzDate + "\n" + credentialScope + "\n" + sigSHA256Hex([]byte(canonReq))
1097
1098 // Signing key.
1099 kDate := sigHMAC([]byte("AWS4"+creds.secretAccessKey), []byte(shortDate))
1100 kRegion := sigHMAC(kDate, []byte(region))
1101 kService := sigHMAC(kRegion, []byte(service))
1102 kSigning := sigHMAC(kService, []byte("aws4_request"))
1103
1104 // Signature.
1105 signature := hex.EncodeToString(sigHMAC(kSigning, []byte(stringToSign)))
1106
1107 // Authorization header.
1108 auth := "AWS4-HMAC-SHA256 " +
1109 "Credential=" + creds.accessKeyID + "/" + credentialScope + ", " +
1110 "SignedHeaders=" + signedHeaders + ", " +
1111 "Signature=" + signature
1112 req.Header.Set("Authorization", auth)
1113 return nil
1114}
1115
1116// canonicalSigV4Path returns the path with each segment URI-encoded
1117// according to SigV4 rules (everything except unreserved chars
1118// A-Za-z0-9-_.~ and / between segments). Bedrock model IDs contain
1119// dots, colons, and slashes; the colons must be percent-encoded.
1120func canonicalSigV4Path(p string) string {
1121 if p == "" {
1122 return "/"
1123 }
1124 segments := strings.Split(p, "/")
1125 for i, s := range segments {
1126 segments[i] = sigV4EncodeSegment(s)
1127 }
1128 return strings.Join(segments, "/")
1129}
1130
1131func sigV4EncodeSegment(s string) string {
1132 var b strings.Builder
1133 for i := 0; i < len(s); i++ {
1134 c := s[i]
1135 // Unreserved per RFC 3986.
1136 if (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') ||
1137 c == '-' || c == '_' || c == '.' || c == '~' {
1138 b.WriteByte(c)
1139 continue
1140 }
1141 fmt.Fprintf(&b, "%%%02X", c)
1142 }
1143 return b.String()
1144}