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