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