From 055fa26f3797ea75b021732aa98a197a44b6d204 Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Sun, 19 Jul 2026 21:11:03 +1000 Subject: [PATCH] =?UTF-8?q?feat(bedrock):=20Ochre=20Phase=202=20=E2=80=94?= =?UTF-8?q?=20ConverseStream=20+=20InvokeModelWithResponseStream=20(#526)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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. --- go.mod | 2 +- spinifex/gateway/bedrock/anthropic.go | 232 +++++++++++++++++ spinifex/gateway/bedrock/anthropic_test.go | 208 +++++++++++++++ spinifex/gateway/bedrock/contract_test.go | 240 +++++++++++++++++ spinifex/gateway/bedrock/converse_stream.go | 165 ++++++++++++ .../gateway/bedrock/converse_stream_test.go | 204 +++++++++++++++ spinifex/gateway/bedrock/eventstream.go | 146 +++++++++++ spinifex/gateway/bedrock/eventstream_test.go | 146 +++++++++++ spinifex/gateway/bedrock/invoke.go | 68 +++++ spinifex/gateway/bedrock/invoke_anthropic.go | 91 +++++++ .../gateway/bedrock/invoke_anthropic_test.go | 84 ++++++ spinifex/gateway/bedrock/invoke_llama.go | 221 +++++++++++++++- spinifex/gateway/bedrock/invoke_llama_test.go | 103 ++++++++ spinifex/gateway/bedrock/invoke_stream.go | 77 ++++++ .../gateway/bedrock/invoke_stream_test.go | 120 +++++++++ spinifex/gateway/bedrock/invoke_test.go | 40 +++ spinifex/gateway/bedrock/provider.go | 59 +++++ spinifex/gateway/bedrock/router_test.go | 48 ++++ spinifex/gateway/bedrock/sse.go | 71 ++++++ spinifex/gateway/bedrock/sse_test.go | 104 ++++++++ spinifex/gateway/bedrock/vllm.go | 241 +++++++++++++++++- spinifex/gateway/bedrock/vllm_test.go | 225 ++++++++++++++++ spinifex/gateway/bedrock_request_test.go | 99 +++++++ spinifex/gateway/bedrockruntime.go | 15 ++ spinifex/gateway/gateway.go | 9 + spinifex/gateway/gateway_test.go | 37 ++- spinifex/otelsetup/http.go | 14 + spinifex/otelsetup/http_test.go | 27 ++ 28 files changed, 3082 insertions(+), 14 deletions(-) create mode 100644 spinifex/gateway/bedrock/contract_test.go create mode 100644 spinifex/gateway/bedrock/converse_stream.go create mode 100644 spinifex/gateway/bedrock/converse_stream_test.go create mode 100644 spinifex/gateway/bedrock/eventstream.go create mode 100644 spinifex/gateway/bedrock/eventstream_test.go create mode 100644 spinifex/gateway/bedrock/invoke_stream.go create mode 100644 spinifex/gateway/bedrock/invoke_stream_test.go create mode 100644 spinifex/gateway/bedrock/sse.go create mode 100644 spinifex/gateway/bedrock/sse_test.go diff --git a/go.mod b/go.mod index 04d654db9..2c6611199 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/spinifex/gateway/bedrock/anthropic.go b/spinifex/gateway/bedrock/anthropic.go index f246dee97..9bfc59cb8 100644 --- a/spinifex/gateway/bedrock/anthropic.go +++ b/spinifex/gateway/bedrock/anthropic.go @@ -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 +} diff --git a/spinifex/gateway/bedrock/anthropic_test.go b/spinifex/gateway/bedrock/anthropic_test.go index cf0b4318b..427c22ad1 100644 --- a/spinifex/gateway/bedrock/anthropic_test.go +++ b/spinifex/gateway/bedrock/anthropic_test.go @@ -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 diff --git a/spinifex/gateway/bedrock/contract_test.go b/spinifex/gateway/bedrock/contract_test.go new file mode 100644 index 000000000..d3abd903e --- /dev/null +++ b/spinifex/gateway/bedrock/contract_test.go @@ -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) +} diff --git a/spinifex/gateway/bedrock/converse_stream.go b/spinifex/gateway/bedrock/converse_stream.go new file mode 100644 index 000000000..d05096d25 --- /dev/null +++ b/spinifex/gateway/bedrock/converse_stream.go @@ -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 + } + } +} diff --git a/spinifex/gateway/bedrock/converse_stream_test.go b/spinifex/gateway/bedrock/converse_stream_test.go new file mode 100644 index 000000000..9468555fa --- /dev/null +++ b/spinifex/gateway/bedrock/converse_stream_test.go @@ -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) +} diff --git a/spinifex/gateway/bedrock/eventstream.go b/spinifex/gateway/bedrock/eventstream.go new file mode 100644 index 000000000..7801e90f7 --- /dev/null +++ b/spinifex/gateway/bedrock/eventstream.go @@ -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) +} diff --git a/spinifex/gateway/bedrock/eventstream_test.go b/spinifex/gateway/bedrock/eventstream_test.go new file mode 100644 index 000000000..220d5caed --- /dev/null +++ b/spinifex/gateway/bedrock/eventstream_test.go @@ -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 } diff --git a/spinifex/gateway/bedrock/invoke.go b/spinifex/gateway/bedrock/invoke.go index d1872cdbb..5b232ef7e 100644 --- a/spinifex/gateway/bedrock/invoke.go +++ b/spinifex/gateway/bedrock/invoke.go @@ -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) +} diff --git a/spinifex/gateway/bedrock/invoke_anthropic.go b/spinifex/gateway/bedrock/invoke_anthropic.go index 62ebcc6e6..c2f221a26 100644 --- a/spinifex/gateway/bedrock/invoke_anthropic.go +++ b/spinifex/gateway/bedrock/invoke_anthropic.go @@ -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 + } +} diff --git a/spinifex/gateway/bedrock/invoke_anthropic_test.go b/spinifex/gateway/bedrock/invoke_anthropic_test.go index 32564f9ac..20f9131b2 100644 --- a/spinifex/gateway/bedrock/invoke_anthropic_test.go +++ b/spinifex/gateway/bedrock/invoke_anthropic_test.go @@ -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")) diff --git a/spinifex/gateway/bedrock/invoke_llama.go b/spinifex/gateway/bedrock/invoke_llama.go index ba4313151..bc3b84088 100644 --- a/spinifex/gateway/bedrock/invoke_llama.go +++ b/spinifex/gateway/bedrock/invoke_llama.go @@ -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 +} diff --git a/spinifex/gateway/bedrock/invoke_llama_test.go b/spinifex/gateway/bedrock/invoke_llama_test.go index 605f3efc5..c8a6e5394 100644 --- a/spinifex/gateway/bedrock/invoke_llama_test.go +++ b/spinifex/gateway/bedrock/invoke_llama_test.go @@ -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) +} diff --git a/spinifex/gateway/bedrock/invoke_stream.go b/spinifex/gateway/bedrock/invoke_stream.go new file mode 100644 index 000000000..f3a13dc19 --- /dev/null +++ b/spinifex/gateway/bedrock/invoke_stream.go @@ -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 + } + } +} diff --git a/spinifex/gateway/bedrock/invoke_stream_test.go b/spinifex/gateway/bedrock/invoke_stream_test.go new file mode 100644 index 000000000..182e85cbb --- /dev/null +++ b/spinifex/gateway/bedrock/invoke_stream_test.go @@ -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) +} diff --git a/spinifex/gateway/bedrock/invoke_test.go b/spinifex/gateway/bedrock/invoke_test.go index 31b4b5142..b279fff97 100644 --- a/spinifex/gateway/bedrock/invoke_test.go +++ b/spinifex/gateway/bedrock/invoke_test.go @@ -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()) +} diff --git a/spinifex/gateway/bedrock/provider.go b/spinifex/gateway/bedrock/provider.go index e684566bd..ab98a667e 100644 --- a/spinifex/gateway/bedrock/provider.go +++ b/spinifex/gateway/bedrock/provider.go @@ -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) +} diff --git a/spinifex/gateway/bedrock/router_test.go b/spinifex/gateway/bedrock/router_test.go index 91e3c4730..6cee3c9e2 100644 --- a/spinifex/gateway/bedrock/router_test.go +++ b/spinifex/gateway/bedrock/router_test.go @@ -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) diff --git a/spinifex/gateway/bedrock/sse.go b/spinifex/gateway/bedrock/sse.go new file mode 100644 index 000000000..1e1061250 --- /dev/null +++ b/spinifex/gateway/bedrock/sse.go @@ -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 +} diff --git a/spinifex/gateway/bedrock/sse_test.go b/spinifex/gateway/bedrock/sse_test.go new file mode 100644 index 000000000..6c6915f3a --- /dev/null +++ b/spinifex/gateway/bedrock/sse_test.go @@ -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) +} diff --git a/spinifex/gateway/bedrock/vllm.go b/spinifex/gateway/bedrock/vllm.go index 853e706f2..3d1a1fa56 100644 --- a/spinifex/gateway/bedrock/vllm.go +++ b/spinifex/gateway/bedrock/vllm.go @@ -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())}, + }, + }) +} diff --git a/spinifex/gateway/bedrock/vllm_test.go b/spinifex/gateway/bedrock/vllm_test.go index 015039cd1..7d3ca90a9 100644 --- a/spinifex/gateway/bedrock/vllm_test.go +++ b/spinifex/gateway/bedrock/vllm_test.go @@ -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) +} diff --git a/spinifex/gateway/bedrock_request_test.go b/spinifex/gateway/bedrock_request_test.go index af445691c..581fc4380 100644 --- a/spinifex/gateway/bedrock_request_test.go +++ b/spinifex/gateway/bedrock_request_test.go @@ -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()) +} diff --git a/spinifex/gateway/bedrockruntime.go b/spinifex/gateway/bedrockruntime.go index 95143fb55..96641fc43 100644 --- a/spinifex/gateway/bedrockruntime.go +++ b/spinifex/gateway/bedrockruntime.go @@ -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 diff --git a/spinifex/gateway/gateway.go b/spinifex/gateway/gateway.go index 404ea755f..7df705e8f 100644 --- a/spinifex/gateway/gateway.go +++ b/spinifex/gateway/gateway.go @@ -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) diff --git a/spinifex/gateway/gateway_test.go b/spinifex/gateway/gateway_test.go index de8952f3d..7e731e4bd 100644 --- a/spinifex/gateway/gateway_test.go +++ b/spinifex/gateway/gateway_test.go @@ -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) diff --git a/spinifex/otelsetup/http.go b/spinifex/otelsetup/http.go index 495510acd..867e91897 100644 --- a/spinifex/otelsetup/http.go +++ b/spinifex/otelsetup/http.go @@ -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 diff --git a/spinifex/otelsetup/http_test.go b/spinifex/otelsetup/http_test.go index d0d1ca424..b1bf960c7 100644 --- a/spinifex/otelsetup/http_test.go +++ b/spinifex/otelsetup/http_test.go @@ -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)