feat(bedrock): Ochre Phase 2 — ConverseStream + InvokeModelWithResponseStream (#526)

* feat(bedrock): Ochre streaming — ConverseStream + InvokeModelWithResponseStream

Adds the two streaming bedrock-runtime operations behind a shared
eventstream.Encoder frame writer: ConverseStream reframes vLLM/Anthropic SSE
into the normalized Bedrock event taxonomy (messageStart -> contentBlock*
-> messageStop -> metadata), InvokeModelWithResponseStream wraps each
provider-native chunk in a PayloadPart frame. Fixes otelsetup's
statusRecorder to implement Flusher/Unwrap so flushes actually reach the
socket through the middleware chain, which streaming depends on.

* fix(bedrock): keep vLLM ConverseStream taxonomy ordered on edge cases

Route contentBlockStop/messageStop through an idempotent ensureStarted so a
usage-only chunk arriving before any finish_reason no longer emits metadata
ahead of messageStop, and a zero-content stream no longer emits stop events
without a preceding messageStart. Regression tests for both.

* fix(bedrock): route /model/ data-plane calls to bedrock-runtime

AWS SDKs sign bedrock-runtime with the SigV4 signing name "bedrock", the
same as the control plane; the two services are separated only by endpoint
hostname. The gateway serves one endpoint and routed purely on the signing
name, so every /model/{id}/converse*|invoke* call landed in Bedrock_Request
(control plane) and returned InvalidAction — no standard SDK/CLI could reach
any bedrock-runtime endpoint.

Disambiguate by path in GetService: a "bedrock"-scoped request to a /model/
path (data-plane-exclusive) resolves to bedrock-runtime. Control-plane paths
(/foundation-models, /custom-models, /model-invocation-job) are unaffected.

* test(bedrock): satisfy testifylint go-require/empty in streaming tests

Use assert (guarded where a value is dereferenced afterward) instead of
require inside httptest handler closures, and assert.Empty for empty-string
checks, per dev's stricter testifylint config.
This commit is contained in:
Josh Siviour
2026-07-19 21:11:03 +10:00
committed by GitHub
parent 8205f88b04
commit 055fa26f37
28 changed files with 3082 additions and 14 deletions
+1 -1
View File
@@ -5,6 +5,7 @@ go 1.26.5
require (
github.com/aws/aws-sdk-go v1.55.8
github.com/aws/aws-sdk-go-v2 v1.42.1
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14
github.com/aws/aws-sdk-go-v2/config v1.32.29
github.com/aws/aws-sdk-go-v2/credentials v1.19.28
github.com/cenkalti/backoff/v4 v4.3.0
@@ -61,7 +62,6 @@ require (
github.com/antithesishq/antithesis-sdk-go v0.7.0 // indirect
github.com/armon/go-metrics v0.4.1 // indirect
github.com/atotto/clipboard v0.1.4 // indirect
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.14 // indirect
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.30 // indirect
github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.30 // indirect
github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.30 // indirect
+232
View File
@@ -5,6 +5,7 @@ import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
@@ -58,6 +59,7 @@ type anthropicRequest struct {
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
StopSequences []string `json:"stop_sequences,omitempty"`
Stream bool `json:"stream,omitempty"`
}
type anthropicUsage struct {
@@ -115,6 +117,12 @@ func (b *anthropicBoundProvider) Converse(ctx context.Context, modelID string, i
return b.inner.Converse(ctx, modelID, input, b.apiKey)
}
var _ ConverseStreamProvider = (*anthropicBoundProvider)(nil)
func (b *anthropicBoundProvider) ConverseStream(ctx context.Context, modelID string, input *bedrockruntime.ConverseStreamInput) (converseStreamSource, error) {
return b.inner.ConverseStream(ctx, modelID, input, b.apiKey)
}
// newAnthropicProvider returns a Provider that calls the Anthropic Messages
// API with apiKey.
func newAnthropicProvider(apiKey string) Provider {
@@ -249,3 +257,227 @@ func anthropicToConverseOutput(ar anthropicResponse, latency time.Duration) *bed
},
}
}
// ConverseStream translates input to an Anthropic Messages request with
// stream:true, opens the SSE connection, and returns a converseStreamSource
// that reframes it into the normalized Bedrock ConverseStream taxonomy.
// Anthropic's own event taxonomy already matches Bedrock's nearly 1:1.
func (p *anthropicProvider) ConverseStream(ctx context.Context, modelID string, input *bedrockruntime.ConverseStreamInput, apiKey string) (converseStreamSource, error) {
req := buildAnthropicRequest(modelID, converseStreamToConverseInput(input))
req.Stream = true
reqBody, err := json.Marshal(req)
if err != nil {
slog.Error("anthropic stream: failed to marshal request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+anthropicMessagesPath, bytes.NewReader(reqBody)) //nolint:gosec // G704: p.baseURL is the hardcoded Anthropic API endpoint (or an httptest stub in tests), never user input
if err != nil {
slog.Error("anthropic stream: failed to build request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq.Header.Set("X-Api-Key", apiKey)
httpReq.Header.Set("Anthropic-Version", anthropicAPIVersion)
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Accept", "text/event-stream")
resp, err := p.httpClient.Do(httpReq) //nolint:gosec // G704: httpReq targets p.baseURL, not user input
if err != nil {
slog.Error("anthropic stream: request failed", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
respBody, _ := io.ReadAll(resp.Body)
_ = resp.Body.Close()
slog.Error("anthropic stream: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody))
return nil, errors.New(mapUpstreamStatus(resp.StatusCode))
}
return &anthropicConverseStreamSource{
resp: resp,
scanner: newSSEScanner(resp.Body),
start: time.Now(),
}, nil
}
// anthropicConverseStreamSource reframes an Anthropic Messages SSE stream
// into the normalized Bedrock ConverseStream taxonomy. Anthropic emits one
// SSE event per normalized Bedrock event, except message_delta (captures
// stop_reason/output tokens, emitted as messageStop) and message_stop
// (triggers the trailing metadata event, carrying accumulated usage).
type anthropicConverseStreamSource struct {
resp *http.Response
scanner *sseScanner
start time.Time
inputTokens int64
outputTokens int64
}
var _ converseStreamSource = (*anthropicConverseStreamSource)(nil)
func (s *anthropicConverseStreamSource) Close() error {
return s.resp.Body.Close()
}
func (s *anthropicConverseStreamSource) Next(_ context.Context) (ConverseStreamEvent, bool, error) {
for {
ev, ok, err := s.scanner.Next()
if err != nil {
if errors.Is(err, io.EOF) {
return ConverseStreamEvent{}, false, nil
}
return ConverseStreamEvent{}, false, newStreamFault(err)
}
if !ok {
continue
}
switch ev.Event {
case "", "ping":
continue
case "message_start":
return s.handleMessageStart(ev.Data)
case "content_block_start":
return s.handleContentBlockStart(ev.Data)
case "content_block_delta":
return s.handleContentBlockDelta(ev.Data)
case "content_block_stop":
return s.handleContentBlockStop(ev.Data)
case "message_delta":
return s.handleMessageDelta(ev.Data)
case "message_stop":
return s.handleMessageStop()
case "error":
return ConverseStreamEvent{}, false, newStreamFault(fmt.Errorf("anthropic stream: upstream error event: %s", ev.Data))
default:
continue // forward-compat: skip unrecognised event types
}
}
}
func (s *anthropicConverseStreamSource) handleMessageStart(data string) (ConverseStreamEvent, bool, error) {
var payload struct {
Message struct {
Role string `json:"role"`
Usage struct {
InputTokens int64 `json:"input_tokens"`
} `json:"usage"`
} `json:"message"`
}
if err := json.Unmarshal([]byte(data), &payload); err != nil {
return ConverseStreamEvent{}, false, newStreamFault(fmt.Errorf("anthropic stream: decode message_start: %w", err))
}
s.inputTokens = payload.Message.Usage.InputTokens
role := payload.Message.Role
if role == "" {
role = bedrockruntime.ConversationRoleAssistant
}
return ConverseStreamEvent{
Kind: converseStreamEventMessageStart,
MessageStart: &bedrockruntime.MessageStartEvent{Role: aws.String(role)},
}, true, nil
}
func (s *anthropicConverseStreamSource) handleContentBlockStart(data string) (ConverseStreamEvent, bool, error) {
var payload struct {
Index int64 `json:"index"`
ContentBlock struct {
Type string `json:"type"`
ID string `json:"id"`
Name string `json:"name"`
} `json:"content_block"`
}
if err := json.Unmarshal([]byte(data), &payload); err != nil {
return ConverseStreamEvent{}, false, newStreamFault(fmt.Errorf("anthropic stream: decode content_block_start: %w", err))
}
start := &bedrockruntime.ContentBlockStart{}
if payload.ContentBlock.Type == "tool_use" {
start.ToolUse = &bedrockruntime.ToolUseBlockStart{
ToolUseId: aws.String(payload.ContentBlock.ID),
Name: aws.String(payload.ContentBlock.Name),
}
}
return ConverseStreamEvent{
Kind: converseStreamEventContentBlockStart,
ContentBlockStart: &bedrockruntime.ContentBlockStartEvent{
ContentBlockIndex: aws.Int64(payload.Index),
Start: start,
},
}, true, nil
}
func (s *anthropicConverseStreamSource) handleContentBlockDelta(data string) (ConverseStreamEvent, bool, error) {
var payload struct {
Index int64 `json:"index"`
Delta struct {
Type string `json:"type"`
Text string `json:"text"`
PartialJSON string `json:"partial_json"`
} `json:"delta"`
}
if err := json.Unmarshal([]byte(data), &payload); err != nil {
return ConverseStreamEvent{}, false, newStreamFault(fmt.Errorf("anthropic stream: decode content_block_delta: %w", err))
}
delta := &bedrockruntime.ContentBlockDelta{}
if payload.Delta.Type == "input_json_delta" {
delta.ToolUse = &bedrockruntime.ToolUseBlockDelta{Input: aws.String(payload.Delta.PartialJSON)}
} else {
delta.Text = aws.String(payload.Delta.Text)
}
return ConverseStreamEvent{
Kind: converseStreamEventContentBlockDelta,
ContentBlockDelta: &bedrockruntime.ContentBlockDeltaEvent{
ContentBlockIndex: aws.Int64(payload.Index),
Delta: delta,
},
}, true, nil
}
func (s *anthropicConverseStreamSource) handleContentBlockStop(data string) (ConverseStreamEvent, bool, error) {
var payload struct {
Index int64 `json:"index"`
}
if err := json.Unmarshal([]byte(data), &payload); err != nil {
return ConverseStreamEvent{}, false, newStreamFault(fmt.Errorf("anthropic stream: decode content_block_stop: %w", err))
}
return ConverseStreamEvent{
Kind: converseStreamEventContentBlockStop,
ContentBlockStop: &bedrockruntime.ContentBlockStopEvent{ContentBlockIndex: aws.Int64(payload.Index)},
}, true, nil
}
func (s *anthropicConverseStreamSource) handleMessageDelta(data string) (ConverseStreamEvent, bool, error) {
var payload struct {
Delta struct {
StopReason string `json:"stop_reason"`
} `json:"delta"`
Usage struct {
OutputTokens int64 `json:"output_tokens"`
} `json:"usage"`
}
if err := json.Unmarshal([]byte(data), &payload); err != nil {
return ConverseStreamEvent{}, false, newStreamFault(fmt.Errorf("anthropic stream: decode message_delta: %w", err))
}
s.outputTokens = payload.Usage.OutputTokens
return ConverseStreamEvent{
Kind: converseStreamEventMessageStop,
MessageStop: &bedrockruntime.MessageStopEvent{StopReason: aws.String(mapAnthropicStopReason(payload.Delta.StopReason))},
}, true, nil
}
func (s *anthropicConverseStreamSource) handleMessageStop() (ConverseStreamEvent, bool, error) {
return ConverseStreamEvent{
Kind: converseStreamEventMetadata,
Metadata: &bedrockruntime.ConverseStreamMetadataEvent{
Usage: &bedrockruntime.TokenUsage{
InputTokens: aws.Int64(s.inputTokens),
OutputTokens: aws.Int64(s.outputTokens),
TotalTokens: aws.Int64(s.inputTokens + s.outputTokens),
},
Metrics: &bedrockruntime.ConverseStreamMetrics{LatencyMs: aws.Int64(time.Since(s.start).Milliseconds())},
},
}, true, nil
}
+208
View File
@@ -129,6 +129,214 @@ func TestAnthropicBoundProvider_Converse(t *testing.T) {
assert.Equal(t, "bound response", *out.Output.Message.Content[0].Text)
}
// anthropicStreamFixture is a canned Anthropic Messages streaming SSE body
// covering the full taxonomy, including a keepalive ping and a
// tool_use content block to exercise structural passthrough.
const anthropicStreamFixture = `event: message_start
data: {"type":"message_start","message":{"id":"msg_1","role":"assistant","usage":{"input_tokens":12}}}
event: content_block_start
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
event: ping
data: {"type":"ping"}
event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" there"}}
event: content_block_stop
data: {"type":"content_block_stop","index":0}
event: content_block_start
data: {"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"toolu_1","name":"get_weather"}}
event: content_block_delta
data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"city\":"}}
event: content_block_stop
data: {"type":"content_block_stop","index":1}
event: message_delta
data: {"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":5}}
event: message_stop
data: {"type":"message_stop"}
`
func TestAnthropicProvider_ConverseStream_MapsSSEToTaxonomy(t *testing.T) {
var capturedBody anthropicRequest
var capturedAPIKey string
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
capturedAPIKey = r.Header.Get("X-Api-Key")
assert.Equal(t, "text/event-stream", r.Header.Get("Accept"))
assert.NoError(t, json.NewDecoder(r.Body).Decode(&capturedBody))
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(anthropicStreamFixture))
}))
defer ts.Close()
p := &anthropicProvider{httpClient: ts.Client(), baseURL: ts.URL}
input := &bedrockruntime.ConverseStreamInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hi")}}},
},
}
src, err := p.ConverseStream(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", input, "sk-test-key")
require.NoError(t, err)
defer func() { _ = src.Close() }()
assert.Equal(t, "sk-test-key", capturedAPIKey)
assert.True(t, capturedBody.Stream)
events := drainConverseStream(t, src)
require.Len(t, events, 10)
kinds := make([]converseStreamEventKind, len(events))
for i, ev := range events {
kinds[i] = ev.Kind
}
assert.Equal(t, []converseStreamEventKind{
converseStreamEventMessageStart,
converseStreamEventContentBlockStart,
converseStreamEventContentBlockDelta,
converseStreamEventContentBlockDelta,
converseStreamEventContentBlockStop,
converseStreamEventContentBlockStart,
converseStreamEventContentBlockDelta,
converseStreamEventContentBlockStop,
converseStreamEventMessageStop,
converseStreamEventMetadata,
}, kinds)
assert.Equal(t, bedrockruntime.ConversationRoleAssistant, *events[0].MessageStart.Role)
assert.Equal(t, "Hello", *events[2].ContentBlockDelta.Delta.Text)
assert.Equal(t, " there", *events[3].ContentBlockDelta.Delta.Text)
toolStart := events[5].ContentBlockStart
require.NotNil(t, toolStart.Start.ToolUse)
assert.Equal(t, "toolu_1", *toolStart.Start.ToolUse.ToolUseId)
assert.Equal(t, "get_weather", *toolStart.Start.ToolUse.Name)
toolDelta := events[6].ContentBlockDelta
require.NotNil(t, toolDelta.Delta.ToolUse)
assert.Equal(t, `{"city":`, *toolDelta.Delta.ToolUse.Input)
assert.Equal(t, bedrockruntime.StopReasonToolUse, *events[8].MessageStop.StopReason)
assert.Equal(t, int64(12), *events[9].Metadata.Usage.InputTokens)
assert.Equal(t, int64(5), *events[9].Metadata.Usage.OutputTokens)
assert.Equal(t, int64(17), *events[9].Metadata.Usage.TotalTokens)
}
func TestAnthropicConverseStreamSource_MessageStopEmitsMetadata(t *testing.T) {
raw := `event: message_start
data: {"type":"message_start","message":{"role":"assistant","usage":{"input_tokens":4}}}
event: message_delta
data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":2}}
event: message_stop
data: {"type":"message_stop"}
`
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(raw))
}))
defer ts.Close()
p := &anthropicProvider{httpClient: ts.Client(), baseURL: ts.URL}
input := &bedrockruntime.ConverseStreamInput{
Messages: []*bedrockruntime.Message{{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hi")}}}},
}
src, err := p.ConverseStream(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", input, "k")
require.NoError(t, err)
defer func() { _ = src.Close() }()
events := drainConverseStream(t, src)
require.Len(t, events, 3)
assert.Equal(t, converseStreamEventMetadata, events[2].Kind)
assert.Equal(t, int64(4), *events[2].Metadata.Usage.InputTokens)
assert.Equal(t, int64(2), *events[2].Metadata.Usage.OutputTokens)
assert.Equal(t, int64(6), *events[2].Metadata.Usage.TotalTokens)
}
func TestAnthropicConverseStreamSource_ErrorEventSurfacesAsStreamFault(t *testing.T) {
raw := "event: error\ndata: {\"error\":{\"message\":\"overloaded\"}}\n\n"
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(raw))
}))
defer ts.Close()
p := &anthropicProvider{httpClient: ts.Client(), baseURL: ts.URL}
input := &bedrockruntime.ConverseStreamInput{
Messages: []*bedrockruntime.Message{{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hi")}}}},
}
src, err := p.ConverseStream(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", input, "k")
require.NoError(t, err)
defer func() { _ = src.Close() }()
_, ok, err := src.Next(context.Background())
assert.False(t, ok)
require.Error(t, err)
var fault *streamFaultError
assert.ErrorAs(t, err, &fault)
assert.Contains(t, err.Error(), "overloaded")
}
func TestAnthropicProvider_ConverseStream_401ReturnsAccessDenied(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"error": {"message": "invalid x-api-key"}}`))
}))
defer ts.Close()
p := &anthropicProvider{httpClient: ts.Client(), baseURL: ts.URL}
input := &bedrockruntime.ConverseStreamInput{
Messages: []*bedrockruntime.Message{{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hi")}}}},
}
_, err := p.ConverseStream(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", input, "bad-key")
require.Error(t, err)
assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error())
}
func TestAnthropicBoundProvider_ConverseStream(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "sk-bound-test", r.Header.Get("X-Api-Key"))
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(anthropicStreamFixture))
}))
defer ts.Close()
p := newAnthropicProvider("sk-bound-test")
bound, ok := p.(*anthropicBoundProvider)
require.True(t, ok)
bound.inner.httpClient = ts.Client()
bound.inner.baseURL = ts.URL
sp, ok := p.(ConverseStreamProvider)
require.True(t, ok)
input := &bedrockruntime.ConverseStreamInput{
Messages: []*bedrockruntime.Message{{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hi")}}}},
}
src, err := sp.ConverseStream(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", input)
require.NoError(t, err)
defer func() { _ = src.Close() }()
events := drainConverseStream(t, src)
require.NotEmpty(t, events)
assert.Equal(t, converseStreamEventMessageStart, events[0].Kind)
}
func TestAnthropicModelID(t *testing.T) {
cases := []struct {
name, in, want string
+240
View File
@@ -0,0 +1,240 @@
package gateway_bedrock
import (
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"regexp"
"testing"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/awstesting/unit"
"github.com/aws/aws-sdk-go/private/protocol/eventstream/eventstreamtest"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// This file is the contract test the plan's acceptance criterion calls for:
// feed our production reframers' output through the *real* aws-sdk-go v1
// bedrockruntime client (not just decode the raw bytes ourselves), proving a
// genuine SDK consumer can read what we frame — byte framing, headers, and
// JSON payload shape all have to be right for this to work at all.
// contractSession stands up an httptest server (TLS+h2, matching a real
// event-stream transport) running handler, and returns a session pointed at
// it plus a bedrockruntime client built on that session.
func contractSession(t *testing.T, handler http.HandlerFunc) *bedrockruntime.BedrockRuntime {
t.Helper()
sess, cleanup, err := eventstreamtest.SetupEventStreamSession(t, handler, true)
require.NoError(t, err)
t.Cleanup(cleanup)
return bedrockruntime.New(sess, &aws.Config{Credentials: unit.Session.Config.Credentials})
}
var converseStreamPathPattern = regexp.MustCompile(`^/model/([^/]+)/converse-stream$`)
var invokeStreamPathPattern = regexp.MustCompile(`^/model/([^/]+)/invoke-with-response-stream$`)
func TestContract_ConverseStream_SelfHost_RealSDKConsumerDecodesFrames(t *testing.T) {
vllmStub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(vllmStreamFixture))
}))
defer vllmStub.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
endpoints := NewStaticEndpointResolver(map[string]string{modelID: vllmStub.URL})
svc := contractSession(t, func(w http.ResponseWriter, r *http.Request) {
m := converseStreamPathPattern.FindStringSubmatch(r.URL.Path)
if !assert.NotNil(t, m, "unexpected path %s", r.URL.Path) {
return
}
body, err := readAll(r)
if !assert.NoError(t, err) {
return
}
err = ConverseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints)
assert.NoError(t, err)
})
resp, err := svc.ConverseStream(&bedrockruntime.ConverseStreamInput{
ModelId: aws.String(modelID),
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}},
},
})
require.NoError(t, err)
defer func() { _ = resp.GetStream().Close() }()
var kinds []string
var deltas []string
for event := range resp.GetStream().Events() {
switch e := event.(type) {
case *bedrockruntime.MessageStartEvent:
kinds = append(kinds, "messageStart")
assert.Equal(t, bedrockruntime.ConversationRoleAssistant, *e.Role)
case *bedrockruntime.ContentBlockStartEvent:
kinds = append(kinds, "contentBlockStart")
case *bedrockruntime.ContentBlockDeltaEvent:
kinds = append(kinds, "contentBlockDelta")
deltas = append(deltas, *e.Delta.Text)
case *bedrockruntime.ContentBlockStopEvent:
kinds = append(kinds, "contentBlockStop")
case *bedrockruntime.MessageStopEvent:
kinds = append(kinds, "messageStop")
assert.Equal(t, bedrockruntime.StopReasonEndTurn, *e.StopReason)
case *bedrockruntime.ConverseStreamMetadataEvent:
kinds = append(kinds, "metadata")
assert.Equal(t, int64(8), *e.Usage.InputTokens)
assert.Equal(t, int64(3), *e.Usage.OutputTokens)
default:
t.Fatalf("unexpected event type %T", event)
}
}
require.NoError(t, resp.GetStream().Err())
assert.Equal(t, []string{
"messageStart", "contentBlockStart", "contentBlockDelta", "contentBlockDelta",
"contentBlockStop", "messageStop", "metadata",
}, kinds)
assert.Equal(t, []string{"Hello", " world"}, deltas)
}
func TestContract_ConverseStream_Anthropic_RealSDKConsumerDecodesFrames(t *testing.T) {
anthropicStub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(anthropicStreamFixture))
}))
defer anthropicStub.Close()
modelID := "anthropic.claude-3-5-sonnet-20240620-v1:0"
// The top-level ConverseStream entrypoint routes provider-direct traffic
// through Router, which hardcodes the real Anthropic base URL — there is
// no injection seam for it (by design: it's a live vendor endpoint, not
// a pinned self-host one). So this contract test drives the same
// production pieces (anthropicProvider.ConverseStream + pumpConverseStream
// + frameWriter) directly against the stub, rather than through
// ConverseStream(ctx, w, ...): the framing/encoding under test is
// identical either way, since both operations share pumpConverseStream.
svc := contractSession(t, func(w http.ResponseWriter, r *http.Request) {
p := &anthropicProvider{httpClient: anthropicStub.Client(), baseURL: anthropicStub.URL}
src, err := p.ConverseStream(r.Context(), modelID, &bedrockruntime.ConverseStreamInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hi")}}},
},
}, "sk-test")
if !assert.NoError(t, err) {
return
}
defer func() { _ = src.Close() }()
fw, err := newFrameWriter(w)
if !assert.NoError(t, err) {
return
}
pumpConverseStream(r.Context(), fw, src, modelID)
})
resp, err := svc.ConverseStream(&bedrockruntime.ConverseStreamInput{ModelId: aws.String(modelID)})
require.NoError(t, err)
defer func() { _ = resp.GetStream().Close() }()
var kinds []string
var sawToolUse bool
for event := range resp.GetStream().Events() {
switch e := event.(type) {
case *bedrockruntime.MessageStartEvent:
kinds = append(kinds, "messageStart")
case *bedrockruntime.ContentBlockStartEvent:
kinds = append(kinds, "contentBlockStart")
if e.Start != nil && e.Start.ToolUse != nil {
sawToolUse = true
assert.Equal(t, "get_weather", *e.Start.ToolUse.Name)
}
case *bedrockruntime.ContentBlockDeltaEvent:
kinds = append(kinds, "contentBlockDelta")
case *bedrockruntime.ContentBlockStopEvent:
kinds = append(kinds, "contentBlockStop")
case *bedrockruntime.MessageStopEvent:
kinds = append(kinds, "messageStop")
assert.Equal(t, bedrockruntime.StopReasonToolUse, *e.StopReason)
case *bedrockruntime.ConverseStreamMetadataEvent:
kinds = append(kinds, "metadata")
assert.Equal(t, int64(12), *e.Usage.InputTokens)
assert.Equal(t, int64(5), *e.Usage.OutputTokens)
default:
t.Fatalf("unexpected event type %T", event)
}
}
require.NoError(t, resp.GetStream().Err())
assert.Equal(t, []string{
"messageStart", "contentBlockStart", "contentBlockDelta", "contentBlockDelta", "contentBlockStop",
"contentBlockStart", "contentBlockDelta", "contentBlockStop", "messageStop", "metadata",
}, kinds)
assert.True(t, sawToolUse)
}
func TestContract_InvokeModelWithResponseStream_SelfHost_RealSDKConsumerDecodesFrames(t *testing.T) {
llamaStub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(llamaCompletionsStreamFixture))
}))
defer llamaStub.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
endpoints := NewStaticEndpointResolver(map[string]string{modelID: llamaStub.URL})
svc := contractSession(t, func(w http.ResponseWriter, r *http.Request) {
m := invokeStreamPathPattern.FindStringSubmatch(r.URL.Path)
if !assert.NotNil(t, m, "unexpected path %s", r.URL.Path) {
return
}
body, err := readAll(r)
if !assert.NoError(t, err) {
return
}
err = InvokeModelWithResponseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints, r.Header.Get("Content-Type"))
assert.NoError(t, err)
})
resp, err := svc.InvokeModelWithResponseStream(&bedrockruntime.InvokeModelWithResponseStreamInput{
ModelId: aws.String(modelID),
Body: []byte(`{"prompt":"hello","max_gen_len":128}`),
})
require.NoError(t, err)
defer func() { _ = resp.GetStream().Close() }()
var generations []string
var sawFinal bool
for event := range resp.GetStream().Events() {
part, ok := event.(*bedrockruntime.PayloadPart)
require.True(t, ok, "unexpected event type %T", event)
var probe struct {
Generation string `json:"generation"`
StopReason string `json:"stop_reason"`
}
require.NoError(t, json.Unmarshal(part.Bytes, &probe))
generations = append(generations, probe.Generation)
if probe.StopReason != "" {
sawFinal = true
assert.Equal(t, "stop", probe.StopReason)
}
}
require.NoError(t, resp.GetStream().Err())
assert.Equal(t, []string{"Hello", " world", ""}, generations)
assert.True(t, sawFinal)
}
func readAll(r *http.Request) ([]byte, error) {
defer func() { _ = r.Body.Close() }()
return io.ReadAll(r.Body)
}
+165
View File
@@ -0,0 +1,165 @@
package gateway_bedrock
import (
"context"
"encoding/json"
"errors"
"fmt"
"log/slog"
"net/http"
"github.com/aws/aws-sdk-go/private/protocol/json/jsonutil"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
// converseStreamEventKind tags which aws-sdk-go v1 event struct a
// ConverseStreamEvent carries, driving both the frame's ":event-type" header
// and which field payload() marshals.
type converseStreamEventKind int
const (
converseStreamEventMessageStart converseStreamEventKind = iota
converseStreamEventContentBlockStart
converseStreamEventContentBlockDelta
converseStreamEventContentBlockStop
converseStreamEventMessageStop
converseStreamEventMetadata
)
// eventType is the Bedrock wire ":event-type" for kind.
func (k converseStreamEventKind) eventType() string {
switch k {
case converseStreamEventMessageStart:
return "messageStart"
case converseStreamEventContentBlockStart:
return "contentBlockStart"
case converseStreamEventContentBlockDelta:
return "contentBlockDelta"
case converseStreamEventContentBlockStop:
return "contentBlockStop"
case converseStreamEventMessageStop:
return "messageStop"
case converseStreamEventMetadata:
return "metadata"
default:
return "unknown"
}
}
// ConverseStreamEvent is one normalized Bedrock ConverseStream event. Exactly
// one payload field is set, matching Kind; reframers (vLLM, Anthropic) build
// these from their own SSE so the writer/pump stays backend-agnostic.
type ConverseStreamEvent struct {
Kind converseStreamEventKind
MessageStart *bedrockruntime.MessageStartEvent
ContentBlockStart *bedrockruntime.ContentBlockStartEvent
ContentBlockDelta *bedrockruntime.ContentBlockDeltaEvent
ContentBlockStop *bedrockruntime.ContentBlockStopEvent
MessageStop *bedrockruntime.MessageStopEvent
Metadata *bedrockruntime.ConverseStreamMetadataEvent
}
// payload JSON-marshals the event struct matching Kind via jsonutil.BuildJSON
// (the SDK's own restjson marshaler), honouring the locationName tags a real
// SDK consumer's UnmarshalEvent expects — mirroring WriteJSONResponse.
func (e ConverseStreamEvent) payload() ([]byte, error) {
switch e.Kind {
case converseStreamEventMessageStart:
return jsonutil.BuildJSON(e.MessageStart)
case converseStreamEventContentBlockStart:
return jsonutil.BuildJSON(e.ContentBlockStart)
case converseStreamEventContentBlockDelta:
return jsonutil.BuildJSON(e.ContentBlockDelta)
case converseStreamEventContentBlockStop:
return jsonutil.BuildJSON(e.ContentBlockStop)
case converseStreamEventMessageStop:
return jsonutil.BuildJSON(e.MessageStop)
case converseStreamEventMetadata:
return jsonutil.BuildJSON(e.Metadata)
default:
return nil, fmt.Errorf("converse stream: unknown event kind %d", e.Kind)
}
}
// converseStreamSource is the backend-agnostic reframer contract: Next
// returns the next normalized event, or (zero,false,nil) at clean EOF.
type converseStreamSource interface {
Next(ctx context.Context) (ConverseStreamEvent, bool, error)
Close() error
}
// ConverseStream is the bedrock-runtime ConverseStream entry point used by
// the gateway route table. Unlike the JSON-dispatch handlers it owns w
// directly: a pre-stream failure (unknown model, unresolved credential,
// upstream connect error) returns an awserrors code for the normal
// ErrorHandler envelope. Once the first frame is written it always returns
// nil — any later failure is an in-band exception event, since the HTTP
// status can no longer change.
func ConverseStream(ctx context.Context, w http.ResponseWriter, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver) error {
input := new(bedrockruntime.ConverseStreamInput)
if len(body) > 0 {
if err := json.Unmarshal(body, input); err != nil {
return errors.New(awserrors.ErrorValidationException)
}
}
src, err := NewRouter(resolver, endpointResolver).ConverseStream(ctx, accountID, modelID, input)
if err != nil {
return err
}
defer func() {
if cerr := src.Close(); cerr != nil {
slog.Error("converse-stream: failed to close upstream source", "model", modelID, "err", cerr)
}
}()
fw, err := newFrameWriter(w)
if err != nil {
return err
}
pumpConverseStream(ctx, fw, src, modelID)
return nil
}
// pumpConverseStream drains src and writes each normalized event as a frame,
// in AWS order (messageStart -> per block contentBlockStart/Delta*/Stop ->
// messageStop -> metadata). A mid-stream Next error surfaces as an in-band
// exception event and stops the pump; a write/flush error also stops the
// pump silently, since the client connection is already broken and the
// HTTP status can no longer change either way.
func pumpConverseStream(ctx context.Context, fw *frameWriter, src converseStreamSource, modelID string) {
for {
event, ok, err := src.Next(ctx)
if err != nil {
excType := excInternalServerException
var fault *streamFaultError
if errors.As(err, &fault) {
excType = excModelStreamErrorException
}
slog.Error("converse-stream: upstream fault", "model", modelID, "err", err)
if werr := fw.writeException(excType, exceptionPayload(err)); werr != nil {
slog.Error("converse-stream: failed to write exception frame", "model", modelID, "err", werr)
}
return
}
if !ok {
return
}
payload, err := event.payload()
if err != nil {
slog.Error("converse-stream: failed to marshal event", "model", modelID, "kind", event.Kind, "err", err)
if werr := fw.writeException(excInternalServerException, exceptionPayload(err)); werr != nil {
slog.Error("converse-stream: failed to write exception frame", "model", modelID, "err", werr)
}
return
}
if werr := fw.writeEvent(event.Kind.eventType(), payload); werr != nil {
slog.Error("converse-stream: failed to write frame, aborting", "model", modelID, "err", werr)
return
}
}
}
@@ -0,0 +1,204 @@
package gateway_bedrock
import (
"bytes"
"context"
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"testing"
esv2 "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// fakeConverseStreamSource drives pumpConverseStream directly with a scripted
// event/error sequence, independent of any real backend.
type fakeConverseStreamSource struct {
events []ConverseStreamEvent
err error
closed bool
}
func (f *fakeConverseStreamSource) Next(_ context.Context) (ConverseStreamEvent, bool, error) {
if len(f.events) > 0 {
ev := f.events[0]
f.events = f.events[1:]
return ev, true, nil
}
if f.err != nil {
return ConverseStreamEvent{}, false, f.err
}
return ConverseStreamEvent{}, false, nil
}
func (f *fakeConverseStreamSource) Close() error {
f.closed = true
return nil
}
// decodeAllFrames decodes every event-stream message in raw via the v2
// package Decoder, returning each frame's ":message-type"/":event-type" (or
// ":exception-type") header and payload.
type decodedFrame struct {
MessageType string
Type string
Payload []byte
}
func decodeAllFrames(t *testing.T, raw []byte) []decodedFrame {
t.Helper()
dec := esv2.NewDecoder()
r := bytes.NewReader(raw)
var out []decodedFrame
for r.Len() > 0 {
msg, err := dec.Decode(r, nil)
require.NoError(t, err)
mt := headerString(t, msg, ":message-type")
typeHeader := ":event-type"
if mt == "exception" {
typeHeader = ":exception-type"
}
out = append(out, decodedFrame{MessageType: mt, Type: headerString(t, msg, typeHeader), Payload: msg.Payload})
}
return out
}
func TestPumpConverseStream_MidStreamFaultWritesExceptionFrame(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
src := &fakeConverseStreamSource{
events: []ConverseStreamEvent{
{Kind: converseStreamEventMessageStart, MessageStart: &bedrockruntime.MessageStartEvent{Role: aws.String("assistant")}},
},
err: newStreamFault(errors.New("upstream connection reset")),
}
pumpConverseStream(context.Background(), fw, src, "test-model")
frames := decodeAllFrames(t, rec.Body.Bytes())
require.Len(t, frames, 2)
assert.Equal(t, "event", frames[0].MessageType)
assert.Equal(t, "messageStart", frames[0].Type)
assert.Equal(t, "exception", frames[1].MessageType)
assert.Equal(t, excModelStreamErrorException, frames[1].Type)
assert.JSONEq(t, `{"message":"upstream connection reset"}`, string(frames[1].Payload))
// The status must still be 200: post-header, a fault cannot change it.
assert.Equal(t, http.StatusOK, rec.Code)
}
func TestPumpConverseStream_NonFaultErrorUsesInternalServerException(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
src := &fakeConverseStreamSource{err: errors.New("plain internal error")}
pumpConverseStream(context.Background(), fw, src, "test-model")
frames := decodeAllFrames(t, rec.Body.Bytes())
require.Len(t, frames, 1)
assert.Equal(t, excInternalServerException, frames[0].Type)
}
func TestPumpConverseStream_CleanEOFWritesNoExceptionFrame(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
src := &fakeConverseStreamSource{
events: []ConverseStreamEvent{
{Kind: converseStreamEventMessageStart, MessageStart: &bedrockruntime.MessageStartEvent{Role: aws.String("assistant")}},
{Kind: converseStreamEventMessageStop, MessageStop: &bedrockruntime.MessageStopEvent{StopReason: aws.String(bedrockruntime.StopReasonEndTurn)}},
},
}
pumpConverseStream(context.Background(), fw, src, "test-model")
frames := decodeAllFrames(t, rec.Body.Bytes())
require.Len(t, frames, 2)
assert.Equal(t, "messageStart", frames[0].Type)
assert.Equal(t, "messageStop", frames[1].Type)
}
func TestConverseStream_UnknownModelReturnsResourceNotFoundPreHeader(t *testing.T) {
rec := httptest.NewRecorder()
body := []byte(`{"messages":[{"role":"user","content":[{"text":"hi"}]}]}`)
err := ConverseStream(context.Background(), rec, "000000000001", "does.not-exist-v1:0", body, nil, nil)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
// A pre-stream failure must not have written anything: the gateway's
// ErrorHandler owns turning this into the JSON error envelope.
assert.Equal(t, 0, rec.Body.Len())
}
func TestConverseStream_MalformedBodyReturnsValidationException(t *testing.T) {
rec := httptest.NewRecorder()
err := ConverseStream(context.Background(), rec, "000000000001", "meta.llama3-70b-instruct-v1:0", []byte("{not-json"), nil, nil)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorValidationException, err.Error())
}
func TestConverseStream_SelfHostHappyPath_WritesFramedTaxonomy(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(vllmStreamFixture))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
rec := httptest.NewRecorder()
body, err := json.Marshal(&bedrockruntime.ConverseStreamInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}},
},
})
require.NoError(t, err)
err = ConverseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
require.NoError(t, err)
assert.Equal(t, http.StatusOK, rec.Code)
assert.Equal(t, eventStreamContentType, rec.Header().Get("Content-Type"))
frames := decodeAllFrames(t, rec.Body.Bytes())
require.Len(t, frames, 7)
assert.Equal(t, "messageStart", frames[0].Type)
assert.Equal(t, "contentBlockStart", frames[1].Type)
assert.Equal(t, "contentBlockDelta", frames[2].Type)
assert.Equal(t, "contentBlockDelta", frames[3].Type)
assert.Equal(t, "contentBlockStop", frames[4].Type)
assert.Equal(t, "messageStop", frames[5].Type)
assert.Equal(t, "metadata", frames[6].Type)
}
func TestConverseStream_NonFlusherWriter_ReturnsPreHeaderError(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(vllmStreamFixture))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
w := &nonFlusherWriter{}
body, err := json.Marshal(&bedrockruntime.ConverseStreamInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}},
},
})
require.NoError(t, err)
err = ConverseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
require.Error(t, err)
assert.Equal(t, awserrors.ErrorInternalError, err.Error())
assert.False(t, w.wroteHeader)
}
+146
View File
@@ -0,0 +1,146 @@
package gateway_bedrock
import (
"encoding/json"
"errors"
"net/http"
esv2 "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
"github.com/aws/aws-sdk-go/private/protocol/json/jsonutil"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
// eventStreamContentType is the wire content type for both ConverseStream and
// InvokeModelWithResponseStream responses.
const eventStreamContentType = "application/vnd.amazon.eventstream"
// Bedrock in-band exception types, used as the frame ":exception-type"
// header once WriteHeader(200) has already been sent and the HTTP status can
// no longer change. Every pre-header failure instead returns a normal
// awserrors code for the ordinary JSON error envelope.
const (
excInternalServerException = "internalServerException"
excModelStreamErrorException = "modelStreamErrorException"
)
// streamFaultError marks a reframer error as originating from the upstream
// connection (a broken/malformed provider stream) rather than our own
// internal encoding, so the pump surfaces it as modelStreamErrorException
// instead of the default internalServerException.
type streamFaultError struct{ err error }
// newStreamFault wraps err as a streamFaultError, or returns nil for a nil err.
func newStreamFault(err error) error {
if err == nil {
return nil
}
return &streamFaultError{err: err}
}
func (f *streamFaultError) Error() string { return f.err.Error() }
func (f *streamFaultError) Unwrap() error { return f.err }
// exceptionMessage is the payload shape for an in-band exception frame,
// matching the "message" field every generated Bedrock exception struct
// (InternalServerException, ModelStreamErrorException, ...) carries.
type exceptionMessage struct {
Message string `json:"message"`
}
// exceptionPayload marshals err as an exception frame payload. Marshaling a
// single-string struct cannot fail; the fallback exists only to satisfy the
// type checker.
func exceptionPayload(err error) []byte {
body, marshalErr := json.Marshal(exceptionMessage{Message: err.Error()})
if marshalErr != nil {
return []byte(`{"message":"internal error"}`)
}
return body
}
// flushSupported reports whether w (or anything it wraps via Unwrap, per the
// otelsetup statusRecorder fix) implements http.Flusher, without invoking
// Flush — a pure capability probe, so a non-Flusher writer can be rejected
// pre-header via the normal awserrors path rather than after WriteHeader(200)
// has already committed the response to 200.
func flushSupported(w http.ResponseWriter) bool {
for {
if _, ok := w.(http.Flusher); ok {
return true
}
u, ok := w.(interface{ Unwrap() http.ResponseWriter })
if !ok {
return false
}
w = u.Unwrap()
}
}
// frameWriter writes AWS event-stream frames to an HTTP response, flushing
// after each one so a streaming SDK consumer sees bytes incrementally rather
// than buffered until the handler returns. One frameWriter serves both
// ConverseStream and InvokeModelWithResponseStream.
type frameWriter struct {
w http.ResponseWriter
enc *esv2.Encoder
ctrl *http.ResponseController
}
// newFrameWriter probes w for flush support (pre-header safe: flushSupported
// has no side effects on failure), then sets the event-stream content type
// and writes the 200 header. Once this returns successfully, no further
// failure may change the HTTP status — only in-band exception frames.
func newFrameWriter(w http.ResponseWriter) (*frameWriter, error) {
if !flushSupported(w) {
return nil, errors.New(awserrors.ErrorInternalError)
}
w.Header().Set("Content-Type", eventStreamContentType)
w.WriteHeader(http.StatusOK)
return &frameWriter{
w: w,
enc: esv2.NewEncoder(),
ctrl: http.NewResponseController(w),
}, nil
}
// writeMessage encodes headers+payload as one event-stream message and
// flushes it to the client.
func (fw *frameWriter) writeMessage(headers esv2.Headers, payload []byte) error {
headers.Set(":content-type", esv2.StringValue("application/json"))
if err := fw.enc.Encode(fw.w, esv2.Message{Headers: headers, Payload: payload}); err != nil {
return err
}
return fw.ctrl.Flush()
}
// writeEvent frames payload as a normal ":message-type: event" frame of
// eventType (e.g. "messageStart", "chunk").
func (fw *frameWriter) writeEvent(eventType string, payload []byte) error {
var headers esv2.Headers
headers.Set(":message-type", esv2.StringValue("event"))
headers.Set(":event-type", esv2.StringValue(eventType))
return fw.writeMessage(headers, payload)
}
// writeException frames payload as an in-band ":message-type: exception"
// frame of excType. Used only after the 200 header has already been sent —
// the post-header error contract forbids changing the HTTP status.
func (fw *frameWriter) writeException(excType string, payload []byte) error {
var headers esv2.Headers
headers.Set(":message-type", esv2.StringValue("exception"))
headers.Set(":exception-type", esv2.StringValue(excType))
return fw.writeMessage(headers, payload)
}
// writeChunk frames a provider-native chunk as the
// InvokeModelWithResponseStream "chunk" event, base64-encoding nativeChunk
// into Bedrock's PayloadPart{Bytes} shape via the SDK's own restjson
// marshaler (jsonutil.BuildJSON), matching WriteJSONResponse's convention.
func (fw *frameWriter) writeChunk(nativeChunk []byte) error {
payload, err := jsonutil.BuildJSON(&bedrockruntime.PayloadPart{Bytes: nativeChunk})
if err != nil {
return err
}
return fw.writeEvent("chunk", payload)
}
@@ -0,0 +1,146 @@
package gateway_bedrock
import (
"bytes"
"errors"
"net/http"
"net/http/httptest"
"testing"
esv2 "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// decodeFrame decodes one event-stream message off raw using the
// aws-sdk-go-v2 eventstream package's own Decoder, proving our encoded bytes
// round-trip through prelude/headers/CRC framing exactly as a real consumer
// would decode them.
func decodeFrame(t *testing.T, raw []byte) esv2.Message {
t.Helper()
dec := esv2.NewDecoder()
msg, err := dec.Decode(bytes.NewReader(raw), nil)
require.NoError(t, err)
return msg
}
func headerString(t *testing.T, msg esv2.Message, name string) string {
t.Helper()
v := msg.Headers.Get(name)
require.NotNilf(t, v, "header %q not present", name)
return v.String()
}
func TestFrameWriter_WriteEvent_RoundTripsThroughDecoder(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
require.NoError(t, fw.writeEvent("messageStart", []byte(`{"role":"assistant"}`)))
assert.Equal(t, http.StatusOK, rec.Code)
assert.Equal(t, eventStreamContentType, rec.Header().Get("Content-Type"))
msg := decodeFrame(t, rec.Body.Bytes())
assert.Equal(t, "event", headerString(t, msg, ":message-type"))
assert.Equal(t, "messageStart", headerString(t, msg, ":event-type"))
assert.Equal(t, "application/json", headerString(t, msg, ":content-type"))
assert.JSONEq(t, `{"role":"assistant"}`, string(msg.Payload))
}
func TestFrameWriter_WriteException_RoundTripsThroughDecoder(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
require.NoError(t, fw.writeException(excModelStreamErrorException, exceptionPayload(errors.New("upstream broke"))))
msg := decodeFrame(t, rec.Body.Bytes())
assert.Equal(t, "exception", headerString(t, msg, ":message-type"))
assert.Equal(t, excModelStreamErrorException, headerString(t, msg, ":exception-type"))
assert.JSONEq(t, `{"message":"upstream broke"}`, string(msg.Payload))
}
func TestFrameWriter_WriteChunk_Base64EncodesPayload(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
require.NoError(t, fw.writeChunk([]byte(`{"type":"content_block_delta"}`)))
msg := decodeFrame(t, rec.Body.Bytes())
assert.Equal(t, "chunk", headerString(t, msg, ":event-type"))
assert.JSONEq(t, `{"bytes":"eyJ0eXBlIjoiY29udGVudF9ibG9ja19kZWx0YSJ9"}`, string(msg.Payload))
}
func TestFrameWriter_MultipleFrames_EachDecodesIndependently(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
require.NoError(t, fw.writeEvent("messageStart", []byte(`{"role":"assistant"}`)))
require.NoError(t, fw.writeEvent("contentBlockDelta", []byte(`{"contentBlockIndex":0}`)))
body := rec.Body.Bytes()
dec := esv2.NewDecoder()
r := bytes.NewReader(body)
msg1, err := dec.Decode(r, nil)
require.NoError(t, err)
assert.Equal(t, "messageStart", headerString(t, msg1, ":event-type"))
msg2, err := dec.Decode(r, nil)
require.NoError(t, err)
assert.Equal(t, "contentBlockDelta", headerString(t, msg2, ":event-type"))
assert.Equal(t, 0, r.Len(), "decoder should have consumed the entire body across both frames")
}
// nonFlusherWriter implements only the bare http.ResponseWriter methods: no
// Flush, no Unwrap. It stands in for a hand-rolled ResponseWriter that
// cannot support streaming, proving newFrameWriter rejects it pre-header
// (Header/WriteHeader/Write are never called) rather than writing a 200 and
// then failing.
type nonFlusherWriter struct {
header http.Header
wroteHeader bool
}
func (w *nonFlusherWriter) Header() http.Header {
if w.header == nil {
w.header = http.Header{}
}
return w.header
}
func (w *nonFlusherWriter) Write(b []byte) (int, error) { return len(b), nil }
func (w *nonFlusherWriter) WriteHeader(int) { w.wroteHeader = true }
func TestNewFrameWriter_NonFlusherWriter_CleanPreHeaderError(t *testing.T) {
w := &nonFlusherWriter{}
fw, err := newFrameWriter(w)
require.Error(t, err)
assert.Nil(t, fw)
assert.Equal(t, awserrors.ErrorInternalError, err.Error())
assert.False(t, w.wroteHeader, "newFrameWriter must not write a header when the writer can't support streaming")
}
func TestFlushSupported_WalksUnwrapChain(t *testing.T) {
rec := httptest.NewRecorder()
wrapped := &unwrappingWriter{inner: rec}
assert.True(t, flushSupported(wrapped))
}
// unwrappingWriter has no Flush of its own but exposes Unwrap, mirroring
// otelsetup's fixed statusRecorder: flushSupported must walk through it to
// find rec's Flush.
type unwrappingWriter struct {
inner http.ResponseWriter
}
func (w *unwrappingWriter) Header() http.Header { return w.inner.Header() }
func (w *unwrappingWriter) Write(b []byte) (int, error) { return w.inner.Write(b) }
func (w *unwrappingWriter) WriteHeader(code int) { w.inner.WriteHeader(code) }
func (w *unwrappingWriter) Unwrap() http.ResponseWriter { return w.inner }
+68
View File
@@ -78,3 +78,71 @@ func (rt *InvokeRouter) InvokeModel(ctx context.Context, accountID, modelID stri
func InvokeModel(ctx context.Context, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver) ([]byte, string, error) {
return NewInvokeRouter(resolver, endpointResolver).InvokeModel(ctx, accountID, modelID, body)
}
// InvokeStreamAdapter is the optional streaming capability an InvokeAdapter
// may implement. Both shipped families (Llama/vLLM, Anthropic) do;
// InvokeStreamRouter type-asserts the resolved InvokeAdapter to it rather
// than widening InvokeAdapter itself, mirroring ConverseStreamProvider.
type InvokeStreamAdapter interface {
InvokeModelWithResponseStream(ctx context.Context, modelID string, body []byte) (invokeStreamSource, error)
}
// InvokeStreamRouter resolves a modelId to its catalog entry and dispatches
// to the matching InvokeStreamAdapter, resolving self-host endpoints and
// provider credentials as needed.
type InvokeStreamRouter struct {
resolver CredentialResolver
endpointResolver EndpointResolver
}
// NewInvokeStreamRouter constructs an InvokeStreamRouter. A nil resolver or
// endpointResolver falls back to a resolver/resolver that finds nothing, so
// an InvokeStreamRouter is always safe to use even before the real stores
// are wired in.
func NewInvokeStreamRouter(resolver CredentialResolver, endpointResolver EndpointResolver) *InvokeStreamRouter {
if resolver == nil {
resolver = NoopCredentialResolver
}
if endpointResolver == nil {
endpointResolver = NewStaticEndpointResolver(nil)
}
return &InvokeStreamRouter{resolver: resolver, endpointResolver: endpointResolver}
}
// InvokeModelWithResponseStream routes modelID to its family adapter via the
// catalog, exactly like InvokeRouter.InvokeModel, then requires the resolved
// adapter to also implement InvokeStreamAdapter.
func (rt *InvokeStreamRouter) InvokeModelWithResponseStream(ctx context.Context, accountID, modelID string, body []byte) (invokeStreamSource, error) {
entry, ok := lookupCatalogEntry(modelID)
if !ok {
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
}
var a InvokeAdapter
switch {
case entry.Provider == tierSelfHost:
a = newLlamaInvokeAdapter(rt.endpointResolver)
case strings.HasPrefix(entry.Provider, providerPrefix):
switch strings.TrimPrefix(entry.Provider, providerPrefix) {
case vendorAnthropic:
key, ok, err := rt.resolver.Resolve(ctx, accountID, vendorAnthropic)
if err != nil {
return nil, err
}
if !ok {
return nil, errors.New(awserrors.ErrorAccessDeniedException)
}
a = newAnthropicInvokeAdapter(key)
default:
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
}
default:
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
}
sa, ok := a.(InvokeStreamAdapter)
if !ok {
return nil, errors.New(awserrors.ErrorValidationException)
}
return sa.InvokeModelWithResponseStream(ctx, modelID, body)
}
@@ -40,6 +40,12 @@ func (b *boundAnthropicInvokeAdapter) InvokeModel(ctx context.Context, modelID s
return b.inner.InvokeModel(ctx, modelID, body, b.apiKey)
}
var _ InvokeStreamAdapter = (*boundAnthropicInvokeAdapter)(nil)
func (b *boundAnthropicInvokeAdapter) InvokeModelWithResponseStream(ctx context.Context, modelID string, body []byte) (invokeStreamSource, error) {
return b.inner.InvokeModelWithResponseStream(ctx, modelID, body, b.apiKey)
}
// newAnthropicInvokeAdapter returns an InvokeAdapter that forwards to the
// Anthropic Messages API with apiKey.
func newAnthropicInvokeAdapter(apiKey string) InvokeAdapter {
@@ -100,3 +106,88 @@ func (a *anthropicInvokeAdapter) InvokeModel(ctx context.Context, modelID string
return respBody, "application/json", nil
}
// InvokeModelWithResponseStream rewrites the request body exactly like
// InvokeModel (drops anthropic_version, injects model), sets stream:true,
// and returns a source that forwards each Anthropic SSE "data:" line
// verbatim as the chunk payload — Bedrock's Claude invoke-stream bytes *are*
// the Anthropic event JSON, so no shape translation is needed.
func (a *anthropicInvokeAdapter) InvokeModelWithResponseStream(ctx context.Context, modelID string, body []byte, apiKey string) (invokeStreamSource, error) {
var fields map[string]json.RawMessage
if err := json.Unmarshal(body, &fields); err != nil {
slog.Error("anthropic invoke-stream: failed to parse request body", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorValidationException)
}
delete(fields, "anthropic_version")
modelJSON, err := json.Marshal(anthropicModelID(modelID))
if err != nil {
slog.Error("anthropic invoke-stream: failed to marshal model id", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
fields["model"] = modelJSON
fields["stream"] = json.RawMessage("true")
reqBody, err := json.Marshal(fields)
if err != nil {
slog.Error("anthropic invoke-stream: failed to marshal request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, a.baseURL+anthropicMessagesPath, bytes.NewReader(reqBody)) //nolint:gosec // G704: a.baseURL is the hardcoded Anthropic API endpoint (or an httptest stub in tests), never user input
if err != nil {
slog.Error("anthropic invoke-stream: failed to build request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq.Header.Set("X-Api-Key", apiKey)
httpReq.Header.Set("Anthropic-Version", anthropicAPIVersion)
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Accept", "text/event-stream")
resp, err := a.httpClient.Do(httpReq) //nolint:gosec // G704: httpReq targets a.baseURL, not user input
if err != nil {
slog.Error("anthropic invoke-stream: request failed", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
respBody, _ := io.ReadAll(resp.Body)
_ = resp.Body.Close()
slog.Error("anthropic invoke-stream: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody))
return nil, errors.New(mapUpstreamStatus(resp.StatusCode))
}
return &anthropicInvokeStreamSource{resp: resp, scanner: newSSEScanner(resp.Body)}, nil
}
// anthropicInvokeStreamSource forwards each Anthropic SSE "data:" line
// verbatim as the chunk payload, skipping keepalive "ping" events.
type anthropicInvokeStreamSource struct {
resp *http.Response
scanner *sseScanner
}
var _ invokeStreamSource = (*anthropicInvokeStreamSource)(nil)
func (s *anthropicInvokeStreamSource) Close() error {
return s.resp.Body.Close()
}
func (s *anthropicInvokeStreamSource) Next(_ context.Context) ([]byte, bool, error) {
for {
ev, ok, err := s.scanner.Next()
if err != nil {
if errors.Is(err, io.EOF) {
return nil, false, nil
}
return nil, false, newStreamFault(err)
}
if !ok {
continue
}
if ev.Event == "ping" || ev.Data == "" {
continue
}
return []byte(ev.Data), true, nil
}
}
@@ -77,6 +77,90 @@ func TestAnthropicInvokeAdapter_InvokeModel_401ReturnsAccessDenied(t *testing.T)
// hardcodes the real Anthropic base URL, so — same package as production code
// — the test reaches into the unexported inner client to redirect it at an
// httptest stub instead of the network.
func TestAnthropicInvokeAdapter_InvokeModelWithResponseStream_ForwardsSSEVerbatim(t *testing.T) {
const raw = `event: message_start
data: {"type":"message_start","message":{"id":"msg_1"}}
event: ping
data: {"type":"ping"}
event: content_block_delta
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}
event: message_stop
data: {"type":"message_stop"}
`
var captured map[string]json.RawMessage
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "text/event-stream", r.Header.Get("Accept"))
assert.NoError(t, json.NewDecoder(r.Body).Decode(&captured))
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(raw))
}))
defer ts.Close()
a := &anthropicInvokeAdapter{httpClient: ts.Client(), baseURL: ts.URL}
reqBody := []byte(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":256,"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`)
src, err := a.InvokeModelWithResponseStream(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", reqBody, "sk-test-key")
require.NoError(t, err)
defer func() { _ = src.Close() }()
var streamField bool
require.NoError(t, json.Unmarshal(captured["stream"], &streamField))
assert.True(t, streamField)
_, hasVersion := captured["anthropic_version"]
assert.False(t, hasVersion)
chunks := drainInvokeStream(t, src)
require.Len(t, chunks, 3, "ping is skipped; message_start, content_block_delta, message_stop forward verbatim")
assert.JSONEq(t, `{"type":"message_start","message":{"id":"msg_1"}}`, string(chunks[0]))
assert.JSONEq(t, `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}`, string(chunks[1]))
assert.JSONEq(t, `{"type":"message_stop"}`, string(chunks[2]))
}
func TestAnthropicInvokeAdapter_InvokeModelWithResponseStream_401ReturnsAccessDenied(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusUnauthorized)
_, _ = w.Write([]byte(`{"error": {"message": "invalid x-api-key"}}`))
}))
defer ts.Close()
a := &anthropicInvokeAdapter{httpClient: ts.Client(), baseURL: ts.URL}
_, err := a.InvokeModelWithResponseStream(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", []byte(`{"anthropic_version":"bedrock-2023-05-31","messages":[]}`), "bad-key")
require.Error(t, err)
assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error())
}
func TestBoundAnthropicInvokeAdapter_InvokeModelWithResponseStream(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "sk-bound-test", r.Header.Get("X-Api-Key"))
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("data: {\"type\":\"message_stop\"}\n\n"))
}))
defer ts.Close()
adapter := newAnthropicInvokeAdapter("sk-bound-test")
bound, ok := adapter.(*boundAnthropicInvokeAdapter)
require.True(t, ok)
bound.inner.httpClient = ts.Client()
bound.inner.baseURL = ts.URL
sa, ok := adapter.(InvokeStreamAdapter)
require.True(t, ok)
src, err := sa.InvokeModelWithResponseStream(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", []byte(`{"anthropic_version":"bedrock-2023-05-31","messages":[]}`))
require.NoError(t, err)
defer func() { _ = src.Close() }()
chunks := drainInvokeStream(t, src)
require.Len(t, chunks, 1)
assert.JSONEq(t, `{"type":"message_stop"}`, string(chunks[0]))
}
func TestBoundAnthropicInvokeAdapter_InvokeModel(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "sk-bound-test", r.Header.Get("X-Api-Key"))
+216 -5
View File
@@ -5,9 +5,12 @@ import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"strings"
"time"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
@@ -32,11 +35,50 @@ type llamaInvokeResponse struct {
// llamaCompletionsRequest is the OpenAI /v1/completions request vLLM serves.
type llamaCompletionsRequest struct {
Model string `json:"model"`
Prompt string `json:"prompt"`
MaxTokens int `json:"max_tokens,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
Model string `json:"model"`
Prompt string `json:"prompt"`
MaxTokens int `json:"max_tokens,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
Stream bool `json:"stream,omitempty"`
StreamOptions *vllmStreamOptions `json:"stream_options,omitempty"`
}
// llamaCompletionsStreamChunk is one OpenAI /v1/completions streaming SSE
// "data:" chunk.
type llamaCompletionsStreamChunk struct {
Choices []llamaCompletionsStreamChoice `json:"choices"`
Usage *llamaCompletionsUsage `json:"usage"`
}
type llamaCompletionsStreamChoice struct {
Text string `json:"text"`
FinishReason *string `json:"finish_reason"`
}
// llamaInvokeStreamChunk is a non-final Bedrock-native Llama
// invoke-with-response-stream chunk.
type llamaInvokeStreamChunk struct {
Generation string `json:"generation"`
}
// llamaInvocationMetrics is Bedrock's per-invocation latency/token summary,
// carried on the final Llama invoke-with-response-stream chunk.
type llamaInvocationMetrics struct {
InputTokenCount int `json:"inputTokenCount"`
OutputTokenCount int `json:"outputTokenCount"`
InvocationLatency int64 `json:"invocationLatency"`
FirstByteLatency int64 `json:"firstByteLatency"`
}
// llamaInvokeStreamFinalChunk is the final Bedrock-native Llama
// invoke-with-response-stream chunk, carrying stop/usage metadata.
type llamaInvokeStreamFinalChunk struct {
Generation string `json:"generation"`
PromptTokenCount int `json:"prompt_token_count"`
GenerationTokenCount int `json:"generation_token_count"`
StopReason string `json:"stop_reason"`
InvocationMetrics llamaInvocationMetrics `json:"amazon-bedrock-invocationMetrics"`
}
type llamaCompletionsChoice struct {
@@ -164,3 +206,172 @@ func (a *llamaInvokeAdapter) InvokeModel(ctx context.Context, modelID string, bo
return out, "application/json", nil
}
var _ InvokeStreamAdapter = (*llamaInvokeAdapter)(nil)
// InvokeModelWithResponseStream resolves modelID's endpoint, opens a
// streaming OpenAI /v1/completions request (stream:true,
// stream_options.include_usage for the trailing usage chunk), and returns a
// source that translates each chunk to the Bedrock-native Llama streaming
// shape.
func (a *llamaInvokeAdapter) InvokeModelWithResponseStream(ctx context.Context, modelID string, body []byte) (invokeStreamSource, error) {
baseURL, ok, err := a.endpointResolver.Endpoint(ctx, modelID)
if err != nil {
slog.Error("llama invoke-stream: endpoint resolution failed", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
if !ok {
return nil, errors.New(awserrors.ErrorModelNotReadyException)
}
var lr llamaInvokeRequest
if err := json.Unmarshal(body, &lr); err != nil {
slog.Error("llama invoke-stream: failed to parse request body", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorValidationException)
}
reqBody, err := json.Marshal(llamaCompletionsRequest{
Model: modelID,
Prompt: lr.Prompt,
MaxTokens: lr.MaxGenLen,
Temperature: lr.Temperature,
TopP: lr.TopP,
Stream: true,
StreamOptions: &vllmStreamOptions{IncludeUsage: true},
})
if err != nil {
slog.Error("llama invoke-stream: failed to marshal request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+llamaCompletionsPath, bytes.NewReader(reqBody)) //nolint:gosec // G704: baseURL is a resolved pinned self-host endpoint, not user input
if err != nil {
slog.Error("llama invoke-stream: failed to build request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Accept", "text/event-stream")
resp, err := a.httpClient.Do(httpReq) //nolint:gosec // G704: httpReq targets the resolved pinned self-host endpoint, not user input
if err != nil {
slog.Error("llama invoke-stream: request failed", "model", modelID, "endpoint", baseURL, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
respBody, _ := io.ReadAll(resp.Body)
_ = resp.Body.Close()
slog.Error("llama invoke-stream: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody))
return nil, errors.New(mapUpstreamStatus(resp.StatusCode))
}
return &llamaInvokeStreamSource{
resp: resp,
scanner: newSSEScanner(resp.Body),
start: time.Now(),
}, nil
}
// llamaInvokeStreamSource translates an OpenAI /v1/completions SSE stream
// into the Bedrock-native Llama invoke-with-response-stream chunk shape: a
// {"generation": "..."} chunk per text delta, and a final chunk carrying
// stop_reason, token counts, and amazon-bedrock-invocationMetrics. The final
// chunk's token counts arrive on vLLM's trailing usage-only chunk (via
// stream_options.include_usage), which follows the finish_reason chunk, so
// emitting the final chunk is deferred until that usage chunk (or EOF) is
// seen — otherwise the token counts would always be zero.
type llamaInvokeStreamSource struct {
resp *http.Response
scanner *sseScanner
start time.Time
gotFirstByte bool
firstByte time.Time
promptTokens, completionTokens int
pendingFinal *llamaInvokeStreamFinalChunk
}
var _ invokeStreamSource = (*llamaInvokeStreamSource)(nil)
func (s *llamaInvokeStreamSource) Close() error {
return s.resp.Body.Close()
}
func (s *llamaInvokeStreamSource) Next(_ context.Context) ([]byte, bool, error) {
for {
ev, ok, err := s.scanner.Next()
if err != nil {
if errors.Is(err, io.EOF) {
if s.pendingFinal != nil {
return s.emitFinal()
}
return nil, false, nil
}
return nil, false, newStreamFault(err)
}
if !ok {
continue
}
data := strings.TrimSpace(ev.Data)
if data == "" || data == "[DONE]" {
continue
}
var chunk llamaCompletionsStreamChunk
if err := json.Unmarshal([]byte(ev.Data), &chunk); err != nil {
return nil, false, newStreamFault(fmt.Errorf("llama invoke-stream: decode chunk: %w", err))
}
if !s.gotFirstByte {
s.gotFirstByte = true
s.firstByte = time.Now()
}
if chunk.Usage != nil {
s.promptTokens = chunk.Usage.PromptTokens
s.completionTokens = chunk.Usage.CompletionTokens
}
if len(chunk.Choices) == 0 {
// The trailing usage-only chunk: emit the deferred final chunk
// now that token counts are known.
if s.pendingFinal != nil {
return s.emitFinal()
}
continue
}
choice := chunk.Choices[0]
if choice.FinishReason != nil {
s.pendingFinal = &llamaInvokeStreamFinalChunk{
Generation: choice.Text,
StopReason: mapLlamaStopReason(*choice.FinishReason),
}
continue
}
out, err := json.Marshal(llamaInvokeStreamChunk{Generation: choice.Text})
if err != nil {
return nil, false, newStreamFault(fmt.Errorf("llama invoke-stream: marshal chunk: %w", err))
}
return out, true, nil
}
}
// emitFinal fills in token counts/metrics on the deferred final chunk and
// marshals it.
func (s *llamaInvokeStreamSource) emitFinal() ([]byte, bool, error) {
final := *s.pendingFinal
s.pendingFinal = nil
final.PromptTokenCount = s.promptTokens
final.GenerationTokenCount = s.completionTokens
final.InvocationMetrics = llamaInvocationMetrics{
InputTokenCount: s.promptTokens,
OutputTokenCount: s.completionTokens,
InvocationLatency: time.Since(s.start).Milliseconds(),
FirstByteLatency: s.firstByte.Sub(s.start).Milliseconds(),
}
out, err := json.Marshal(final)
if err != nil {
return nil, false, newStreamFault(fmt.Errorf("llama invoke-stream: marshal final chunk: %w", err))
}
return out, true, nil
}
@@ -102,3 +102,106 @@ func TestLlamaInvokeAdapter_InvokeModel_UnresolvedEndpointReturnsModelNotReady(t
require.Error(t, err)
assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error())
}
// llamaCompletionsStreamFixture is a canned OpenAI /v1/completions streaming
// SSE body: two text deltas, a finish_reason chunk, a trailing usage-only
// chunk, then [DONE].
const llamaCompletionsStreamFixture = `data: {"choices":[{"text":"Hello","finish_reason":null}]}
data: {"choices":[{"text":" world","finish_reason":null}]}
data: {"choices":[{"text":"","finish_reason":"stop"}]}
data: {"choices":[],"usage":{"prompt_tokens":6,"completion_tokens":2}}
data: [DONE]
`
func drainInvokeStream(t *testing.T, src invokeStreamSource) [][]byte {
t.Helper()
var chunks [][]byte
for {
chunk, ok, err := src.Next(context.Background())
require.NoError(t, err)
if !ok {
return chunks
}
chunks = append(chunks, chunk)
}
}
func TestLlamaInvokeAdapter_InvokeModelWithResponseStream_TranslatesChunks(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var req llamaCompletionsRequest
assert.NoError(t, json.NewDecoder(r.Body).Decode(&req))
assert.True(t, req.Stream)
if !assert.NotNil(t, req.StreamOptions) {
return
}
assert.True(t, req.StreamOptions.IncludeUsage)
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(llamaCompletionsStreamFixture))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
a := newLlamaInvokeAdapter(NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
a.httpClient = ts.Client()
src, err := a.InvokeModelWithResponseStream(context.Background(), modelID, []byte(`{"prompt":"hello","max_gen_len":128}`))
require.NoError(t, err)
defer func() { _ = src.Close() }()
chunks := drainInvokeStream(t, src)
require.Len(t, chunks, 3)
var c1, c2 llamaInvokeStreamChunk
require.NoError(t, json.Unmarshal(chunks[0], &c1))
assert.Equal(t, "Hello", c1.Generation)
require.NoError(t, json.Unmarshal(chunks[1], &c2))
assert.Equal(t, " world", c2.Generation)
var final llamaInvokeStreamFinalChunk
require.NoError(t, json.Unmarshal(chunks[2], &final))
assert.Empty(t, final.Generation)
assert.Equal(t, "stop", final.StopReason)
assert.Equal(t, 6, final.PromptTokenCount)
assert.Equal(t, 2, final.GenerationTokenCount)
assert.Equal(t, 6, final.InvocationMetrics.InputTokenCount)
assert.Equal(t, 2, final.InvocationMetrics.OutputTokenCount)
assert.GreaterOrEqual(t, final.InvocationMetrics.InvocationLatency, int64(0))
}
func TestLlamaInvokeAdapter_InvokeModelWithResponseStream_UnresolvedEndpointReturnsModelNotReady(t *testing.T) {
a := newLlamaInvokeAdapter(NewStaticEndpointResolver(nil))
_, err := a.InvokeModelWithResponseStream(context.Background(), "meta.llama3-70b-instruct-v1:0", []byte(`{"prompt":"hello"}`))
require.Error(t, err)
assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error())
}
func TestLlamaInvokeStreamSource_MidStreamDecodeErrorSurfacesAsStreamFault(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("data: {not-json\n\n"))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
a := newLlamaInvokeAdapter(NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
a.httpClient = ts.Client()
src, err := a.InvokeModelWithResponseStream(context.Background(), modelID, []byte(`{"prompt":"hello"}`))
require.NoError(t, err)
defer func() { _ = src.Close() }()
_, ok, err := src.Next(context.Background())
assert.False(t, ok)
require.Error(t, err)
var fault *streamFaultError
assert.ErrorAs(t, err, &fault)
}
+77
View File
@@ -0,0 +1,77 @@
package gateway_bedrock
import (
"context"
"errors"
"log/slog"
"net/http"
)
// invokeStreamSource is the backend-agnostic InvokeModelWithResponseStream
// reframer contract: Next returns the next provider-native chunk (verbatim
// for Anthropic passthrough, translated to the Bedrock-native shape for
// Llama), or (nil,false,nil) at clean EOF.
type invokeStreamSource interface {
Next(ctx context.Context) (chunk []byte, ok bool, err error)
Close() error
}
// InvokeModelWithResponseStream is the bedrock-runtime
// InvokeModelWithResponseStream entry point used by the gateway route table.
// Unlike the JSON-dispatch handlers it owns w directly: a pre-stream failure
// (unknown model, unresolved credential, upstream connect error) returns an
// awserrors code for the normal ErrorHandler envelope. Once the first frame
// is written it always returns nil — any later failure is an in-band
// exception event, since the HTTP status can no longer change.
// requestContentType is the client's declared Content-Type, logged only.
func InvokeModelWithResponseStream(ctx context.Context, w http.ResponseWriter, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, requestContentType string) error {
src, err := NewInvokeStreamRouter(resolver, endpointResolver).InvokeModelWithResponseStream(ctx, accountID, modelID, body)
if err != nil {
return err
}
defer func() {
if cerr := src.Close(); cerr != nil {
slog.Error("invoke-with-response-stream: failed to close upstream source", "model", modelID, "err", cerr)
}
}()
fw, err := newFrameWriter(w)
if err != nil {
return err
}
slog.Debug("invoke-with-response-stream: streaming", "model", modelID, "request_content_type", requestContentType)
pumpInvokeStream(ctx, fw, src, modelID)
return nil
}
// pumpInvokeStream drains src and writes each provider-native chunk as a
// "chunk" frame. A mid-stream Next error surfaces as an in-band exception
// event and stops the pump; a write/flush error also stops the pump
// silently, since the client connection is already broken and the HTTP
// status can no longer change either way.
func pumpInvokeStream(ctx context.Context, fw *frameWriter, src invokeStreamSource, modelID string) {
for {
chunk, ok, err := src.Next(ctx)
if err != nil {
excType := excInternalServerException
var fault *streamFaultError
if errors.As(err, &fault) {
excType = excModelStreamErrorException
}
slog.Error("invoke-with-response-stream: upstream fault", "model", modelID, "err", err)
if werr := fw.writeException(excType, exceptionPayload(err)); werr != nil {
slog.Error("invoke-with-response-stream: failed to write exception frame", "model", modelID, "err", werr)
}
return
}
if !ok {
return
}
if werr := fw.writeChunk(chunk); werr != nil {
slog.Error("invoke-with-response-stream: failed to write frame, aborting", "model", modelID, "err", werr)
return
}
}
}
@@ -0,0 +1,120 @@
package gateway_bedrock
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// fakeInvokeStreamSource drives pumpInvokeStream directly with a scripted
// chunk/error sequence, independent of any real backend.
type fakeInvokeStreamSource struct {
chunks [][]byte
err error
closed bool
}
func (f *fakeInvokeStreamSource) Next(_ context.Context) ([]byte, bool, error) {
if len(f.chunks) > 0 {
c := f.chunks[0]
f.chunks = f.chunks[1:]
return c, true, nil
}
if f.err != nil {
return nil, false, f.err
}
return nil, false, nil
}
func (f *fakeInvokeStreamSource) Close() error {
f.closed = true
return nil
}
func TestPumpInvokeStream_MidStreamFaultWritesExceptionFrame(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
src := &fakeInvokeStreamSource{
chunks: [][]byte{[]byte(`{"generation":"hi"}`)},
err: newStreamFault(errors.New("upstream connection reset")),
}
pumpInvokeStream(context.Background(), fw, src, "test-model")
frames := decodeAllFrames(t, rec.Body.Bytes())
require.Len(t, frames, 2)
assert.Equal(t, "chunk", frames[0].Type)
assert.Equal(t, "exception", frames[1].MessageType)
assert.Equal(t, excModelStreamErrorException, frames[1].Type)
}
func TestPumpInvokeStream_NonFaultErrorUsesInternalServerException(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
src := &fakeInvokeStreamSource{err: errors.New("plain internal error")}
pumpInvokeStream(context.Background(), fw, src, "test-model")
frames := decodeAllFrames(t, rec.Body.Bytes())
require.Len(t, frames, 1)
assert.Equal(t, excInternalServerException, frames[0].Type)
}
func TestInvokeModelWithResponseStream_UnknownModelReturnsResourceNotFoundPreHeader(t *testing.T) {
rec := httptest.NewRecorder()
err := InvokeModelWithResponseStream(context.Background(), rec, "000000000001", "does.not-exist-v1:0", []byte(`{}`), nil, nil, "application/json")
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
assert.Equal(t, 0, rec.Body.Len())
}
func TestInvokeModelWithResponseStream_SelfHostHappyPath_WritesChunkFrames(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(llamaCompletionsStreamFixture))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
rec := httptest.NewRecorder()
body := []byte(`{"prompt":"hello","max_gen_len":128}`)
err := InvokeModelWithResponseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), "application/json")
require.NoError(t, err)
assert.Equal(t, http.StatusOK, rec.Code)
assert.Equal(t, eventStreamContentType, rec.Header().Get("Content-Type"))
frames := decodeAllFrames(t, rec.Body.Bytes())
require.Len(t, frames, 3)
for _, f := range frames {
assert.Equal(t, "chunk", f.Type)
}
}
func TestInvokeModelWithResponseStream_NonFlusherWriter_ReturnsPreHeaderError(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(llamaCompletionsStreamFixture))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
w := &nonFlusherWriter{}
body := []byte(`{"prompt":"hello"}`)
err := InvokeModelWithResponseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), "application/json")
require.Error(t, err)
assert.Equal(t, awserrors.ErrorInternalError, err.Error())
assert.False(t, w.wroteHeader)
}
+40
View File
@@ -82,3 +82,43 @@ func TestInvokeModel_PackageEntryPoint_UnknownModel(t *testing.T) {
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
}
func TestInvokeStreamRouter_SelfHostSuccess(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(llamaCompletionsStreamFixture))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
rt := NewInvokeStreamRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
src, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`))
require.NoError(t, err)
defer func() { _ = src.Close() }()
chunks := drainInvokeStream(t, src)
assert.Len(t, chunks, 3)
}
func TestInvokeStreamRouter_UnknownModelReturnsResourceNotFound(t *testing.T) {
rt := NewInvokeStreamRouter(nil, nil)
_, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`))
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
}
func TestInvokeStreamRouter_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) {
rt := NewInvokeStreamRouter(stubCredentialResolver{ok: false}, nil)
_, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", []byte(`{}`))
require.Error(t, err)
assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error())
}
func TestInvokeStreamRouter_SelfHostNoEndpointReturnsModelNotReady(t *testing.T) {
rt := NewInvokeStreamRouter(nil, nil)
_, err := rt.InvokeModelWithResponseStream(context.Background(), "000000000001", "meta.llama3-70b-instruct-v1:0", []byte(`{"prompt":"hello"}`))
require.Error(t, err)
assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error())
}
+59
View File
@@ -84,3 +84,62 @@ func (rt *Router) Converse(ctx context.Context, accountID, modelID string, input
func Converse(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseInput, resolver CredentialResolver, endpointResolver EndpointResolver) (*bedrockruntime.ConverseOutput, error) {
return NewRouter(resolver, endpointResolver).Converse(ctx, accountID, modelID, input)
}
// ConverseStreamProvider is the optional streaming capability a Provider may
// implement. Both shipped backends (vLLM, Anthropic) do; Router.ConverseStream
// type-asserts the resolved Provider to it rather than widening Provider
// itself, so a future non-streaming-capable backend still compiles.
type ConverseStreamProvider interface {
ConverseStream(ctx context.Context, modelID string, input *bedrockruntime.ConverseStreamInput) (converseStreamSource, error)
}
// converseStreamToConverseInput adapts a ConverseStreamInput to a
// ConverseInput carrying only the fields buildVLLMRequest/buildAnthropicRequest
// read (Messages, System, InferenceConfig). The two generated types are
// otherwise structurally identical field-for-field; this lets the streaming
// path reuse the exact same request builders as Converse.
func converseStreamToConverseInput(input *bedrockruntime.ConverseStreamInput) *bedrockruntime.ConverseInput {
return &bedrockruntime.ConverseInput{
Messages: input.Messages,
System: input.System,
InferenceConfig: input.InferenceConfig,
}
}
// ConverseStream routes modelID to its provider via the catalog, exactly like
// Converse, then requires the resolved provider to also implement
// ConverseStreamProvider.
func (rt *Router) ConverseStream(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseStreamInput) (converseStreamSource, error) {
entry, ok := lookupCatalogEntry(modelID)
if !ok {
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
}
var p Provider
switch {
case entry.Provider == tierSelfHost:
p = newVLLMProvider(rt.endpointResolver)
case strings.HasPrefix(entry.Provider, providerPrefix):
switch strings.TrimPrefix(entry.Provider, providerPrefix) {
case vendorAnthropic:
key, ok, err := rt.resolver.Resolve(ctx, accountID, vendorAnthropic)
if err != nil {
return nil, err
}
if !ok {
return nil, errors.New(awserrors.ErrorAccessDeniedException)
}
p = newAnthropicProvider(key)
default:
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
}
default:
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
}
sp, ok := p.(ConverseStreamProvider)
if !ok {
return nil, errors.New(awserrors.ErrorValidationException)
}
return sp.ConverseStream(ctx, modelID, input)
}
+48
View File
@@ -92,6 +92,54 @@ func TestConverse_PackageEntryPoint_UnknownModel(t *testing.T) {
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
}
func converseStreamInput() *bedrockruntime.ConverseStreamInput {
return &bedrockruntime.ConverseStreamInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}},
},
}
}
func TestRouter_ConverseStream_SelfHostSuccess(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(vllmStreamFixture))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
src, err := rt.ConverseStream(context.Background(), "000000000001", modelID, converseStreamInput())
require.NoError(t, err)
defer func() { _ = src.Close() }()
events := drainConverseStream(t, src)
assert.Equal(t, converseStreamEventMessageStart, events[0].Kind)
}
func TestRouter_ConverseStream_UnknownModelReturnsResourceNotFound(t *testing.T) {
rt := NewRouter(nil, nil)
_, err := rt.ConverseStream(context.Background(), "000000000001", "does.not-exist-v1:0", converseStreamInput())
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
}
func TestRouter_ConverseStream_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) {
rt := NewRouter(stubCredentialResolver{ok: false}, nil)
_, err := rt.ConverseStream(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", converseStreamInput())
require.Error(t, err)
assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error())
}
func TestRouter_ConverseStream_SelfHostNoEndpointReturnsModelNotReady(t *testing.T) {
rt := NewRouter(nil, nil)
_, err := rt.ConverseStream(context.Background(), "000000000001", "meta.llama3-70b-instruct-v1:0", converseStreamInput())
require.Error(t, err)
assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error())
}
func TestNewRouter_NilArgumentsFallBackToNoops(t *testing.T) {
rt := NewRouter(nil, nil)
require.NotNil(t, rt)
+71
View File
@@ -0,0 +1,71 @@
package gateway_bedrock
import (
"bufio"
"io"
"strings"
)
// sseMaxLineBuf is the maximum single-line size the scanner accepts. Data
// lines carry one JSON chunk each; 1MiB comfortably covers any single
// streamed delta.
const sseMaxLineBuf = 1024 * 1024
// sseEvent is one parsed text/event-stream record: the optional "event:"
// type (empty when absent, per the SSE spec's implicit "message" default)
// and its "data:" lines joined with "\n".
type sseEvent struct {
Event string
Data string
}
// sseScanner reads text/event-stream framing shared by vLLM/OpenAI (bare
// "data:" lines) and Anthropic ("event:"+"data:" pairs): blank-line
// terminated records, skipping comment lines (leading ':') and blank
// keepalive lines between records.
type sseScanner struct {
sc *bufio.Scanner
}
// newSSEScanner wraps r for SSE line scanning.
func newSSEScanner(r io.Reader) *sseScanner {
sc := bufio.NewScanner(r)
sc.Buffer(make([]byte, 0, 64*1024), sseMaxLineBuf)
return &sseScanner{sc: sc}
}
// Next returns the next SSE record, or (zero, false, err) at EOF/read error.
// err is io.EOF at a clean stream end.
func (s *sseScanner) Next() (sseEvent, bool, error) {
var event string
var dataLines []string
sawAny := false
for s.sc.Scan() {
line := s.sc.Text()
if line == "" {
if sawAny {
return sseEvent{Event: event, Data: strings.Join(dataLines, "\n")}, true, nil
}
continue // keepalive blank line between records
}
switch {
case strings.HasPrefix(line, ":"):
// comment/keepalive, ignore — does not start a record
case strings.HasPrefix(line, "event:"):
sawAny = true
event = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
case strings.HasPrefix(line, "data:"):
sawAny = true
dataLines = append(dataLines, strings.TrimPrefix(strings.TrimPrefix(line, "data:"), " "))
}
}
if err := s.sc.Err(); err != nil {
return sseEvent{}, false, err
}
if sawAny {
return sseEvent{Event: event, Data: strings.Join(dataLines, "\n")}, true, nil
}
return sseEvent{}, false, io.EOF
}
+104
View File
@@ -0,0 +1,104 @@
package gateway_bedrock
import (
"io"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestSSEScanner_VLLMStyleBareDataLines(t *testing.T) {
raw := "data: {\"a\":1}\n\ndata: {\"a\":2}\n\ndata: [DONE]\n\n"
sc := newSSEScanner(strings.NewReader(raw))
ev1, ok, err := sc.Next()
require.NoError(t, err)
require.True(t, ok)
assert.Empty(t, ev1.Event)
assert.Equal(t, `{"a":1}`, ev1.Data)
ev2, ok, err := sc.Next()
require.NoError(t, err)
require.True(t, ok)
assert.Equal(t, `{"a":2}`, ev2.Data)
ev3, ok, err := sc.Next()
require.NoError(t, err)
require.True(t, ok)
assert.Equal(t, "[DONE]", ev3.Data)
_, ok, err = sc.Next()
assert.False(t, ok)
assert.ErrorIs(t, err, io.EOF)
}
func TestSSEScanner_AnthropicStyleEventAndData(t *testing.T) {
raw := "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: ping\ndata: {}\n\n"
sc := newSSEScanner(strings.NewReader(raw))
ev1, ok, err := sc.Next()
require.NoError(t, err)
require.True(t, ok)
assert.Equal(t, "message_start", ev1.Event)
assert.Equal(t, `{"type":"message_start"}`, ev1.Data)
ev2, ok, err := sc.Next()
require.NoError(t, err)
require.True(t, ok)
assert.Equal(t, "ping", ev2.Event)
}
func TestSSEScanner_MultiLineData(t *testing.T) {
raw := "data: line one\ndata: line two\n\n"
sc := newSSEScanner(strings.NewReader(raw))
ev, ok, err := sc.Next()
require.NoError(t, err)
require.True(t, ok)
assert.Equal(t, "line one\nline two", ev.Data)
}
func TestSSEScanner_SkipsCommentsAndKeepaliveBlankLines(t *testing.T) {
raw := ":keepalive\n\n:another comment\ndata: {\"a\":1}\n\n"
sc := newSSEScanner(strings.NewReader(raw))
ev, ok, err := sc.Next()
require.NoError(t, err)
require.True(t, ok)
assert.Equal(t, `{"a":1}`, ev.Data)
}
func TestSSEScanner_TrailingRecordWithoutFinalBlankLine(t *testing.T) {
raw := "data: {\"a\":1}\n" // no trailing blank line before EOF
sc := newSSEScanner(strings.NewReader(raw))
ev, ok, err := sc.Next()
require.NoError(t, err)
require.True(t, ok)
assert.Equal(t, `{"a":1}`, ev.Data)
_, ok, err = sc.Next()
assert.False(t, ok)
assert.ErrorIs(t, err, io.EOF)
}
func TestSSEScanner_EmptyStreamIsCleanEOF(t *testing.T) {
sc := newSSEScanner(strings.NewReader(""))
_, ok, err := sc.Next()
assert.False(t, ok)
assert.ErrorIs(t, err, io.EOF)
}
type errReader struct{ err error }
func (r errReader) Read(_ []byte) (int, error) { return 0, r.err }
func TestSSEScanner_ReadErrorPropagates(t *testing.T) {
boom := assert.AnError
sc := newSSEScanner(errReader{err: boom})
_, ok, err := sc.Next()
assert.False(t, ok)
assert.ErrorIs(t, err, boom)
}
+235 -6
View File
@@ -5,6 +5,7 @@ import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
@@ -47,12 +48,31 @@ type vllmMessage struct {
}
type vllmChatRequest struct {
Model string `json:"model"`
Messages []vllmMessage `json:"messages"`
MaxTokens *int64 `json:"max_tokens,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
Stop []string `json:"stop,omitempty"`
Model string `json:"model"`
Messages []vllmMessage `json:"messages"`
MaxTokens *int64 `json:"max_tokens,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
Stop []string `json:"stop,omitempty"`
Stream bool `json:"stream,omitempty"`
StreamOptions *vllmStreamOptions `json:"stream_options,omitempty"`
}
// vllmStreamOptions requests the trailing usage-only chunk vLLM/OpenAI
// otherwise omits from a streaming response.
type vllmStreamOptions struct {
IncludeUsage bool `json:"include_usage"`
}
// vllmStreamChunk is one OpenAI chat-completions streaming SSE "data:" chunk.
type vllmStreamChunk struct {
Choices []vllmStreamChoice `json:"choices"`
Usage *vllmUsage `json:"usage"`
}
type vllmStreamChoice struct {
Delta vllmMessage `json:"delta"`
FinishReason *string `json:"finish_reason"`
}
type vllmChoice struct {
@@ -222,3 +242,212 @@ func vllmToConverseOutput(cr vllmChatResponse, latency time.Duration) (*bedrockr
Metrics: &bedrockruntime.ConverseMetrics{LatencyMs: aws.Int64(latency.Milliseconds())},
}, nil
}
var _ ConverseStreamProvider = (*vllmProvider)(nil)
// ConverseStream resolves modelID's endpoint, opens a streaming OpenAI
// chat-completions request (stream:true, stream_options.include_usage for
// the trailing usage chunk), and returns a converseStreamSource that
// reframes the SSE into the normalized Bedrock ConverseStream taxonomy.
func (p *vllmProvider) ConverseStream(ctx context.Context, modelID string, input *bedrockruntime.ConverseStreamInput) (converseStreamSource, error) {
baseURL, ok, err := p.endpointResolver.Endpoint(ctx, modelID)
if err != nil {
slog.Error("vllm stream: endpoint resolution failed", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
if !ok {
return nil, errors.New(awserrors.ErrorModelNotReadyException)
}
req := buildVLLMRequest(modelID, converseStreamToConverseInput(input))
req.Stream = true
req.StreamOptions = &vllmStreamOptions{IncludeUsage: true}
reqBody, err := json.Marshal(req)
if err != nil {
slog.Error("vllm stream: failed to marshal request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+vllmChatCompletionsPath, bytes.NewReader(reqBody)) //nolint:gosec // G704: baseURL is a resolved pinned self-host endpoint, not user input
if err != nil {
slog.Error("vllm stream: failed to build request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Accept", "text/event-stream")
resp, err := p.httpClient.Do(httpReq) //nolint:gosec // G704: httpReq targets the resolved pinned self-host endpoint, not user input
if err != nil {
slog.Error("vllm stream: request failed", "model", modelID, "endpoint", baseURL, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
respBody, _ := io.ReadAll(resp.Body)
_ = resp.Body.Close()
slog.Error("vllm stream: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody))
return nil, errors.New(mapUpstreamStatus(resp.StatusCode))
}
return &vllmConverseStreamSource{
resp: resp,
scanner: newSSEScanner(resp.Body),
start: time.Now(),
}, nil
}
// vllmConverseStreamSource reframes a vLLM (OpenAI chat-completions) SSE
// stream into the normalized Bedrock ConverseStream event taxonomy:
// messageStart+contentBlockStart on the first chunk, contentBlockDelta per
// content delta, contentBlockStop+messageStop on finish_reason, metadata on
// the trailing usage-only chunk (or defensively at EOF if the upstream never
// sends one). Events are queued because a single upstream chunk can fan out
// to more than one normalized event.
type vllmConverseStreamSource struct {
resp *http.Response
scanner *sseScanner
start time.Time
queue []ConverseStreamEvent
started bool
stopped bool
metadataEmitted bool
inputTokens int64
outputTokens int64
}
var _ converseStreamSource = (*vllmConverseStreamSource)(nil)
func (s *vllmConverseStreamSource) Close() error {
return s.resp.Body.Close()
}
func (s *vllmConverseStreamSource) Next(_ context.Context) (ConverseStreamEvent, bool, error) {
for len(s.queue) == 0 {
ev, ok, err := s.scanner.Next()
if err != nil {
if errors.Is(err, io.EOF) {
s.closeOut(bedrockruntime.StopReasonEndTurn)
if !s.metadataEmitted {
s.emitMetadata()
}
break
}
return ConverseStreamEvent{}, false, newStreamFault(err)
}
if !ok {
continue
}
data := strings.TrimSpace(ev.Data)
if data == "" || data == "[DONE]" {
continue
}
var chunk vllmStreamChunk
if err := json.Unmarshal([]byte(ev.Data), &chunk); err != nil {
return ConverseStreamEvent{}, false, newStreamFault(fmt.Errorf("vllm stream: decode chunk: %w", err))
}
s.consume(chunk)
}
if len(s.queue) == 0 {
return ConverseStreamEvent{}, false, nil
}
next := s.queue[0]
s.queue = s.queue[1:]
return next, true, nil
}
// consume folds one decoded chat-completion chunk into the queued normalized
// events and running usage totals.
func (s *vllmConverseStreamSource) consume(chunk vllmStreamChunk) {
if chunk.Usage != nil {
s.inputTokens = chunk.Usage.PromptTokens
s.outputTokens = chunk.Usage.CompletionTokens
}
if len(chunk.Choices) == 0 {
// The trailing usage-only chunk from stream_options.include_usage.
// Close the block/message before metadata so the taxonomy stays ordered
// even if the upstream sends usage without a prior finish_reason chunk.
if chunk.Usage != nil {
s.closeOut(bedrockruntime.StopReasonEndTurn)
s.emitMetadata()
}
return
}
choice := chunk.Choices[0]
s.ensureStarted()
if choice.Delta.Content != "" {
s.queue = append(s.queue, ConverseStreamEvent{
Kind: converseStreamEventContentBlockDelta,
ContentBlockDelta: &bedrockruntime.ContentBlockDeltaEvent{
ContentBlockIndex: aws.Int64(0),
Delta: &bedrockruntime.ContentBlockDelta{Text: aws.String(choice.Delta.Content)},
},
})
}
if choice.FinishReason != nil {
s.closeOut(mapVLLMFinishReason(*choice.FinishReason))
}
}
// ensureStarted queues messageStart+contentBlockStart exactly once, so the stop
// events can never appear without a preceding start on a zero-content stream.
func (s *vllmConverseStreamSource) ensureStarted() {
if s.started {
return
}
s.started = true
s.queue = append(s.queue,
ConverseStreamEvent{
Kind: converseStreamEventMessageStart,
MessageStart: &bedrockruntime.MessageStartEvent{Role: aws.String(bedrockruntime.ConversationRoleAssistant)},
},
ConverseStreamEvent{
Kind: converseStreamEventContentBlockStart,
ContentBlockStart: &bedrockruntime.ContentBlockStartEvent{ContentBlockIndex: aws.Int64(0), Start: &bedrockruntime.ContentBlockStart{}},
},
)
}
// closeOut queues contentBlockStop+messageStop exactly once.
func (s *vllmConverseStreamSource) closeOut(stopReason string) {
if s.stopped {
return
}
s.ensureStarted()
s.stopped = true
s.queue = append(s.queue,
ConverseStreamEvent{
Kind: converseStreamEventContentBlockStop,
ContentBlockStop: &bedrockruntime.ContentBlockStopEvent{ContentBlockIndex: aws.Int64(0)},
},
ConverseStreamEvent{
Kind: converseStreamEventMessageStop,
MessageStop: &bedrockruntime.MessageStopEvent{StopReason: aws.String(stopReason)},
},
)
}
// emitMetadata queues the trailing metadata event exactly once.
func (s *vllmConverseStreamSource) emitMetadata() {
if s.metadataEmitted {
return
}
s.metadataEmitted = true
s.queue = append(s.queue, ConverseStreamEvent{
Kind: converseStreamEventMetadata,
Metadata: &bedrockruntime.ConverseStreamMetadataEvent{
Usage: &bedrockruntime.TokenUsage{
InputTokens: aws.Int64(s.inputTokens),
OutputTokens: aws.Int64(s.outputTokens),
TotalTokens: aws.Int64(s.inputTokens + s.outputTokens),
},
Metrics: &bedrockruntime.ConverseStreamMetrics{LatencyMs: aws.Int64(time.Since(s.start).Milliseconds())},
},
})
}
+225
View File
@@ -74,3 +74,228 @@ func TestVLLMProvider_Converse_UnresolvedEndpointReturnsModelNotReady(t *testing
require.Error(t, err)
assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error())
}
// vllmStreamFixture is a canned OpenAI chat-completions streaming SSE body:
// a role-only first chunk, two content deltas, a finish_reason chunk, a
// trailing usage-only chunk (stream_options.include_usage), then [DONE].
const vllmStreamFixture = `data: {"choices":[{"delta":{"role":"assistant"},"finish_reason":null}]}
data: {"choices":[{"delta":{"content":"Hello"},"finish_reason":null}]}
data: {"choices":[{"delta":{"content":" world"},"finish_reason":null}]}
data: {"choices":[{"delta":{},"finish_reason":"stop"}]}
data: {"choices":[],"usage":{"prompt_tokens":8,"completion_tokens":3}}
data: [DONE]
`
// drainConverseStream reads src to clean EOF, returning the ordered events.
func drainConverseStream(t *testing.T, src converseStreamSource) []ConverseStreamEvent {
t.Helper()
var events []ConverseStreamEvent
for {
ev, ok, err := src.Next(context.Background())
require.NoError(t, err)
if !ok {
return events
}
events = append(events, ev)
}
}
func TestVLLMProvider_ConverseStream_MapsSSEToTaxonomy(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "text/event-stream", r.Header.Get("Accept"))
var req vllmChatRequest
assert.NoError(t, json.NewDecoder(r.Body).Decode(&req))
assert.True(t, req.Stream)
if !assert.NotNil(t, req.StreamOptions) {
return
}
assert.True(t, req.StreamOptions.IncludeUsage)
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(vllmStreamFixture))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
p := newVLLMProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
p.httpClient = ts.Client()
input := &bedrockruntime.ConverseStreamInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}},
},
}
src, err := p.ConverseStream(context.Background(), modelID, input)
require.NoError(t, err)
defer func() { _ = src.Close() }()
events := drainConverseStream(t, src)
require.Len(t, events, 7)
kinds := make([]converseStreamEventKind, len(events))
for i, ev := range events {
kinds[i] = ev.Kind
}
assert.Equal(t, []converseStreamEventKind{
converseStreamEventMessageStart,
converseStreamEventContentBlockStart,
converseStreamEventContentBlockDelta,
converseStreamEventContentBlockDelta,
converseStreamEventContentBlockStop,
converseStreamEventMessageStop,
converseStreamEventMetadata,
}, kinds)
assert.Equal(t, bedrockruntime.ConversationRoleAssistant, *events[0].MessageStart.Role)
assert.Equal(t, int64(0), *events[1].ContentBlockStart.ContentBlockIndex)
assert.Equal(t, "Hello", *events[2].ContentBlockDelta.Delta.Text)
assert.Equal(t, " world", *events[3].ContentBlockDelta.Delta.Text)
assert.Equal(t, int64(0), *events[4].ContentBlockStop.ContentBlockIndex)
assert.Equal(t, bedrockruntime.StopReasonEndTurn, *events[5].MessageStop.StopReason)
assert.Equal(t, int64(8), *events[6].Metadata.Usage.InputTokens)
assert.Equal(t, int64(3), *events[6].Metadata.Usage.OutputTokens)
assert.Equal(t, int64(11), *events[6].Metadata.Usage.TotalTokens)
assert.GreaterOrEqual(t, *events[6].Metadata.Metrics.LatencyMs, int64(0))
}
func TestVLLMProvider_ConverseStream_UnresolvedEndpointReturnsModelNotReady(t *testing.T) {
p := newVLLMProvider(NewStaticEndpointResolver(nil))
input := &bedrockruntime.ConverseStreamInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}},
},
}
_, err := p.ConverseStream(context.Background(), "meta.llama3-70b-instruct-v1:0", input)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error())
}
func TestVLLMProvider_ConverseStream_UpstreamNon2xxMapsStatus(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusTooManyRequests)
_, _ = w.Write([]byte(`{"error":"rate limited"}`))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
p := newVLLMProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
p.httpClient = ts.Client()
input := &bedrockruntime.ConverseStreamInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}},
},
}
_, err := p.ConverseStream(context.Background(), modelID, input)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorThrottlingException, err.Error())
}
// A usage-only chunk that arrives without a preceding finish_reason chunk must
// still yield contentBlockStop+messageStop before metadata, not metadata first.
func TestVLLMProvider_ConverseStream_UsageWithoutFinishReasonKeepsOrder(t *testing.T) {
sse := "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\",\"content\":\"Hi\"}}]}\n\n" +
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":8,\"completion_tokens\":3}}\n\n" +
"data: [DONE]\n\n"
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(sse))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
p := newVLLMProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
p.httpClient = ts.Client()
src, err := p.ConverseStream(context.Background(), modelID, &bedrockruntime.ConverseStreamInput{})
require.NoError(t, err)
defer func() { _ = src.Close() }()
events := drainConverseStream(t, src)
kinds := make([]converseStreamEventKind, len(events))
for i, ev := range events {
kinds[i] = ev.Kind
}
assert.Equal(t, []converseStreamEventKind{
converseStreamEventMessageStart,
converseStreamEventContentBlockStart,
converseStreamEventContentBlockDelta,
converseStreamEventContentBlockStop,
converseStreamEventMessageStop,
converseStreamEventMetadata,
}, kinds)
}
// A zero-content stream (upstream closes before any choice chunk) must still be
// well-formed: messageStart before the synthesized stop events, never orphaned.
func TestVLLMProvider_ConverseStream_EmptyStreamIsWellFormed(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
p := newVLLMProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
p.httpClient = ts.Client()
src, err := p.ConverseStream(context.Background(), modelID, &bedrockruntime.ConverseStreamInput{})
require.NoError(t, err)
defer func() { _ = src.Close() }()
events := drainConverseStream(t, src)
require.NotEmpty(t, events)
assert.Equal(t, converseStreamEventMessageStart, events[0].Kind)
kinds := make([]converseStreamEventKind, len(events))
for i, ev := range events {
kinds[i] = ev.Kind
}
assert.Equal(t, []converseStreamEventKind{
converseStreamEventMessageStart,
converseStreamEventContentBlockStart,
converseStreamEventContentBlockStop,
converseStreamEventMessageStop,
converseStreamEventMetadata,
}, kinds)
}
func TestVLLMConverseStreamSource_MidStreamDecodeErrorSurfacesAsStreamFault(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("data: {not-json\n\n"))
}))
defer ts.Close()
modelID := "meta.llama3-70b-instruct-v1:0"
p := newVLLMProvider(NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
p.httpClient = ts.Client()
input := &bedrockruntime.ConverseStreamInput{
Messages: []*bedrockruntime.Message{
{Role: aws.String(bedrockruntime.ConversationRoleUser), Content: []*bedrockruntime.ContentBlock{{Text: aws.String("hello")}}},
},
}
src, err := p.ConverseStream(context.Background(), modelID, input)
require.NoError(t, err)
defer func() { _ = src.Close() }()
_, ok, err := src.Next(context.Background())
assert.False(t, ok)
require.Error(t, err)
var fault *streamFaultError
assert.ErrorAs(t, err, &fault)
}
+99
View File
@@ -200,3 +200,102 @@ func TestBedrockRuntimeRequest_MissingAccountIDReturnsServerInternal(t *testing.
require.Error(t, err)
assert.Equal(t, awserrors.ErrorServerInternal, err.Error())
}
// newVLLMStreamStub stands up an httptest server answering the OpenAI
// chat-completions streaming (SSE) wire, for the ConverseStream self-host path.
func newVLLMStreamStub(t *testing.T) *httptest.Server {
t.Helper()
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"role":"assistant"},"finish_reason":null}]}
data: {"choices":[{"delta":{"content":"hi from vllm stream"},"finish_reason":null}]}
data: {"choices":[{"delta":{},"finish_reason":"stop"}]}
data: {"choices":[],"usage":{"prompt_tokens":4,"completion_tokens":2}}
data: [DONE]
`))
}))
t.Cleanup(ts.Close)
return ts
}
// newLlamaCompletionsStreamStub stands up an httptest server answering the
// OpenAI completions streaming (SSE) wire, for the
// InvokeModelWithResponseStream self-host path.
func newLlamaCompletionsStreamStub(t *testing.T) *httptest.Server {
t.Helper()
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`data: {"choices":[{"text":"hi from llama stream","finish_reason":null}]}
data: {"choices":[{"text":"","finish_reason":"stop"}]}
data: {"choices":[],"usage":{"prompt_tokens":5,"completion_tokens":3}}
data: [DONE]
`))
}))
t.Cleanup(ts.Close)
return ts
}
func TestBedrockRuntimeRequest_ConverseStream(t *testing.T) {
ts := newVLLMStreamStub(t)
gw := newBedrockRequestGateway(t, ts.URL)
body := `{"messages":[{"role":"user","content":[{"text":"hello"}]}]}`
req := bedrockRequestWithAccount(http.MethodPost, "/model/"+bedrockTestLlamaModelID+"/converse-stream", body)
w := httptest.NewRecorder()
require.NoError(t, gw.BedrockRuntime_Request(w, req))
require.Equal(t, http.StatusOK, w.Code)
assert.Equal(t, "application/vnd.amazon.eventstream", w.Header().Get("Content-Type"))
assert.NotZero(t, w.Body.Len())
}
func TestBedrockRuntimeRequest_ConverseStream_UnknownModelReturnsResourceNotFound(t *testing.T) {
ts := newVLLMStreamStub(t)
gw := newBedrockRequestGateway(t, ts.URL)
body := `{"messages":[{"role":"user","content":[{"text":"hello"}]}]}`
req := bedrockRequestWithAccount(http.MethodPost, "/model/does.not-exist-v1:0/converse-stream", body)
w := httptest.NewRecorder()
err := gw.BedrockRuntime_Request(w, req)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
// Pre-stream failure: BedrockRuntime_Request must not have written
// anything, leaving ErrorHandler free to write the JSON envelope.
assert.Zero(t, w.Body.Len())
}
func TestBedrockRuntimeRequest_InvokeModelWithResponseStream(t *testing.T) {
ts := newLlamaCompletionsStreamStub(t)
gw := newBedrockRequestGateway(t, ts.URL)
body := `{"prompt":"hello","max_gen_len":128}`
req := bedrockRequestWithAccount(http.MethodPost, "/model/"+bedrockTestLlamaModelID+"/invoke-with-response-stream", body)
w := httptest.NewRecorder()
require.NoError(t, gw.BedrockRuntime_Request(w, req))
require.Equal(t, http.StatusOK, w.Code)
assert.Equal(t, "application/vnd.amazon.eventstream", w.Header().Get("Content-Type"))
assert.NotZero(t, w.Body.Len())
}
func TestBedrockRuntimeRequest_InvokeModelWithResponseStream_UnknownModelReturnsResourceNotFound(t *testing.T) {
ts := newLlamaCompletionsStreamStub(t)
gw := newBedrockRequestGateway(t, ts.URL)
body := `{"prompt":"hello"}`
req := bedrockRequestWithAccount(http.MethodPost, "/model/does.not-exist-v1:0/invoke-with-response-stream", body)
w := httptest.NewRecorder()
err := gw.BedrockRuntime_Request(w, req)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
assert.Zero(t, w.Body.Len())
}
+15
View File
@@ -44,6 +44,8 @@ var bedrockRuntimeRoutes = []bedrockRuntimeRoute{
return gateway_bedrock.Converse(ctx, acct, p[0], input, resolver, endpoints)
}},
{"POST", regexp.MustCompile(`^/model/([^/]+)/invoke$`), "InvokeModel", nil},
{"POST", regexp.MustCompile(`^/model/([^/]+)/converse-stream$`), "ConverseStream", nil},
{"POST", regexp.MustCompile(`^/model/([^/]+)/invoke-with-response-stream$`), "InvokeModelWithResponseStream", nil},
}
// lookupBedrockRuntimeAction matches method+path against bedrockRuntimeRoutes,
@@ -117,6 +119,19 @@ func (gw *GatewayConfig) BedrockRuntime_Request(w http.ResponseWriter, r *http.R
return nil
}
// ConverseStream and InvokeModelWithResponseStream own w directly and
// write framed event-stream bytes as they arrive, rather than one
// buffered struct/body WriteJSONResponse/WriteRawResponse could send in
// one shot. Each returns an error ONLY for a pre-first-frame failure
// (-> ErrorHandler); once streaming starts they always return nil,
// surfacing any further failure as an in-band exception event.
if action == "ConverseStream" {
return gateway_bedrock.ConverseStream(r.Context(), w, accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver())
}
if action == "InvokeModelWithResponseStream" {
return gateway_bedrock.InvokeModelWithResponseStream(r.Context(), w, accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), r.Header.Get("Content-Type"))
}
output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver(), gw.bedrockEndpointResolver())
if err != nil {
return err
+9
View File
@@ -387,6 +387,15 @@ func (gw *GatewayConfig) GetService(r *http.Request) (string, error) {
if !ok {
return "", errors.New(awserrors.ErrorAuthFailure)
}
// bedrock and bedrock-runtime share the SigV4 signing name "bedrock", so the
// credential scope alone cannot tell the control plane from the data plane —
// AWS separates them by endpoint hostname, but the gateway serves one
// endpoint. The request path is the discriminator: /model/... is exclusive to
// the data plane (Converse/InvokeModel and their streaming variants), so a
// "bedrock"-scoped call to it is really bedrock-runtime.
if svc == "bedrock" && strings.HasPrefix(r.URL.Path, "/model/") {
svc = "bedrock-runtime"
}
if !supportedServices[svc] {
slog.Debug("Unsupported service", "service", svc)
return "", errors.New(awserrors.ErrorUnsupportedOperation)
+35 -2
View File
@@ -459,7 +459,8 @@ func TestGetService(t *testing.T) {
tests := []struct {
name string
ctxVal any // value to set in ctxService, nil means no value
ctxVal any // value to set in ctxService, nil means no value
path string // request path, empty means "/"
wantSvc string
wantError string
}{
@@ -493,11 +494,43 @@ func TestGetService(t *testing.T) {
ctxVal: "tagging",
wantSvc: "tagging",
},
{
// bedrock and bedrock-runtime share the SigV4 signing name
// "bedrock"; a /model/ path is data-plane-exclusive, so it must
// resolve to bedrock-runtime even though the scope reads "bedrock".
name: "bedrock scope on data-plane path resolves to bedrock-runtime",
ctxVal: "bedrock",
path: "/model/meta.llama3-70b-instruct-v1:0/converse",
wantSvc: "bedrock-runtime",
},
{
name: "bedrock scope on streaming data-plane path resolves to bedrock-runtime",
ctxVal: "bedrock",
path: "/model/meta.llama3-70b-instruct-v1:0/converse-stream",
wantSvc: "bedrock-runtime",
},
{
name: "bedrock scope on control-plane path stays bedrock",
ctxVal: "bedrock",
path: "/foundation-models",
wantSvc: "bedrock",
},
{
// A native bedrock-runtime scope is left untouched by the path check.
name: "bedrock-runtime scope on data-plane path stays bedrock-runtime",
ctxVal: "bedrock-runtime",
path: "/model/meta.llama3-70b-instruct-v1:0/invoke",
wantSvc: "bedrock-runtime",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
req := httptest.NewRequest(http.MethodPost, "/", nil)
path := tc.path
if path == "" {
path = "/"
}
req := httptest.NewRequest(http.MethodPost, path, nil)
if tc.ctxVal != nil {
ctx := context.WithValue(req.Context(), ctxService, tc.ctxVal)
req = req.WithContext(ctx)
+14
View File
@@ -28,6 +28,20 @@ func (w *statusRecorder) WriteHeader(code int) {
w.ResponseWriter.WriteHeader(code)
}
// Unwrap returns the wrapped ResponseWriter, letting http.ResponseController
// (and any other Unwrap-aware caller) walk past statusRecorder to reach the
// real flusher underneath.
func (w *statusRecorder) Unwrap() http.ResponseWriter { return w.ResponseWriter }
// Flush implements http.Flusher by delegating to the wrapped ResponseWriter.
// Without this, statusRecorder blocks every downstream Flush call (including
// chi's WrapResponseWriter, which wraps statusRecorder) from ever reaching
// the real writer: the call succeeds as a no-op but no bytes reach the
// socket, silently breaking streaming handlers.
func (w *statusRecorder) Flush() {
_ = http.NewResponseController(w.ResponseWriter).Flush()
}
// HTTPMiddleware opens a server span per request, honoring an inbound W3C
// traceparent header, and records request count/duration metrics. Handlers
// rename the span (and SetRequestAction) once they resolve a logical
+27
View File
@@ -23,6 +23,33 @@ func withRecorder(t *testing.T) *tracetest.SpanRecorder {
return sr
}
// flusherSpy records whether Flush was called, satisfying http.Flusher so a
// statusRecorder wrapping it can prove the flush chain reaches it.
type flusherSpy struct {
http.ResponseWriter
flushed bool
}
func (f *flusherSpy) Flush() { f.flushed = true }
// TestStatusRecorderFlushReachesUnderlyingFlusher guards the §0 fix: before
// Unwrap()/Flush() were added, http.NewResponseController(rec).Flush() would
// fail with ErrNotSupported (or, via a Flusher-only outer wrapper such as
// chi's WrapResponseWriter, silently no-op) because statusRecorder itself was
// not a Flusher and exposed no way to reach the one it wraps.
func TestStatusRecorderFlushReachesUnderlyingFlusher(t *testing.T) {
spy := &flusherSpy{ResponseWriter: httptest.NewRecorder()}
rec := &statusRecorder{ResponseWriter: spy, status: http.StatusOK}
if err := http.NewResponseController(rec).Flush(); err != nil {
t.Fatalf("Flush: %v", err)
}
if !spy.flushed {
t.Error("flush did not reach the underlying Flusher spy")
}
}
func TestHTTPMiddlewareSpanPerRequest(t *testing.T) {
sr := withRecorder(t)