mirror of
https://github.com/mulgadc/spinifex.git
synced 2026-10-20 11:21:49 +00:00
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:
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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 }
|
||||
@@ -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"))
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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())},
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user