feat: Ochre invocation records to JetStream, delivered to S3 and the log sink (#673)

* feat: record every Bedrock invocation to JetStream, deliver to S3 and logs

Publishes an InvocationRecord for every Converse/InvokeModel call (streaming
or not) to a durable JetStream stream, then fans it out via a delivery
consumer to the account's own predastore bucket (body text, gated per
account by PutModelInvocationLoggingConfiguration) and a metadata-only log
line. Publish failures increment a drop counter instead of failing the
invocation. Streaming pumps record on every exit path via defer, including
client disconnect, and Close() now cancels the upstream request outright so
a disconnected client actually frees backend generation. vLLM streams
estimate output tokens from content-delta counts when real usage isn't
available yet, flagged usage_estimated; Anthropic always reports real usage.

* fix: deliver invocation records to S3 regardless of body, replicate stream

deliver() was gating the S3 write itself on InputText/OutputText being
non-empty, so an account with a destination bucket configured but body
delivery left off got nothing — the delivery flags already gate the body
fields upstream in Record(), and shouldn't gate delivery a second time.
EnsureInvocationStream defaulted to a single replica on a multi-node
cluster, unlike its sibling KV stores, so losing one node lost billing and
audit records without the control plane itself going down. Also renames
the awsgw predastore client from ecrStore to objStore now that it backs
both the ECR registry and Bedrock invocation delivery.
This commit is contained in:
Josh Siviour
2026-08-03 12:44:40 +10:00
committed by GitHub
parent 261d5d505e
commit eef5b0784f
20 changed files with 1703 additions and 115 deletions
+45 -5
View File
@@ -2,6 +2,7 @@ package gateway
import (
"context"
"encoding/json"
"errors"
"io"
"log/slog"
@@ -25,21 +26,39 @@ type bedrockRoute struct {
// bedrockRouteHandler invokes a per-action bedrock (control-plane) gateway
// function. params holds the regex capture groups, PathUnescape'd. resolver
// is gw.bedrockResolver(): the configured credential store, or a no-op
// fallback.
type bedrockRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver) (any, error)
// fallback. loggingStore is gw.bedrockLoggingConfigStore().
type bedrockRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver, loggingStore *gateway_bedrock.LoggingConfigStore) (any, error)
// bedrockRoutes is the dispatch table. More-specific paths must precede
// less-specific ones with the same prefix so the regex matcher picks the
// deeper route first.
var bedrockRoutes = []bedrockRoute{
{"GET", regexp.MustCompile(`^/foundation-models$`), "ListFoundationModels",
func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver) (any, error) {
func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver, _ *gateway_bedrock.LoggingConfigStore) (any, error) {
return gateway_bedrock.ListFoundationModels(ctx, acct, resolver, new(bedrock.ListFoundationModelsInput))
}},
{"GET", regexp.MustCompile(`^/foundation-models/([^/]+)$`), "GetFoundationModel",
func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver) (any, error) {
func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, _ *gateway_bedrock.LoggingConfigStore) (any, error) {
return gateway_bedrock.GetFoundationModel(ctx, acct, p[0])
}},
{"PUT", regexp.MustCompile(`^/logging/modelinvocations$`), "PutModelInvocationLoggingConfiguration",
func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, store *gateway_bedrock.LoggingConfigStore) (any, error) {
input := new(bedrock.PutModelInvocationLoggingConfigurationInput)
if len(b) > 0 {
if err := json.Unmarshal(b, input); err != nil {
return nil, errors.New(awserrors.ErrorValidationException)
}
}
return gateway_bedrock.PutModelInvocationLoggingConfiguration(ctx, acct, store, input)
}},
{"GET", regexp.MustCompile(`^/logging/modelinvocations$`), "GetModelInvocationLoggingConfiguration",
func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, store *gateway_bedrock.LoggingConfigStore) (any, error) {
return gateway_bedrock.GetModelInvocationLoggingConfiguration(ctx, acct, store, new(bedrock.GetModelInvocationLoggingConfigurationInput))
}},
{"DELETE", regexp.MustCompile(`^/logging/modelinvocations$`), "DeleteModelInvocationLoggingConfiguration",
func(ctx context.Context, acct string, p []string, b []byte, _ gateway_bedrock.CredentialResolver, store *gateway_bedrock.LoggingConfigStore) (any, error) {
return gateway_bedrock.DeleteModelInvocationLoggingConfiguration(ctx, acct, store, new(bedrock.DeleteModelInvocationLoggingConfigurationInput))
}},
}
// lookupBedrockAction matches method+path against bedrockRoutes, returning the
@@ -102,7 +121,7 @@ func (gw *GatewayConfig) Bedrock_Request(w http.ResponseWriter, r *http.Request)
return errors.New(awserrors.ErrorInvalidParameterValue)
}
output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver())
output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver(), gw.bedrockLoggingConfigStore())
if err != nil {
return err
}
@@ -120,6 +139,27 @@ func (gw *GatewayConfig) bedrockResolver() gateway_bedrock.CredentialResolver {
return gateway_bedrock.NoopCredentialResolver
}
// bedrockLoggingConfigStore returns gw.BedrockLoggingConfig, or a store
// backed by no JetStream client when unconfigured. Reads/writes then fail
// with an error (no JetStream to open a KV bucket against) rather than
// panicking, which is acceptable for unit tests of unrelated routes that
// never reach a logging-config handler.
func (gw *GatewayConfig) bedrockLoggingConfigStore() *gateway_bedrock.LoggingConfigStore {
if gw.BedrockLoggingConfig != nil {
return gw.BedrockLoggingConfig
}
return gateway_bedrock.NewLoggingConfigStore(nil, 1)
}
// bedrockRecorder returns gw.BedrockRecorder, or the no-op fallback when no
// invocation recorder is configured.
func (gw *GatewayConfig) bedrockRecorder() gateway_bedrock.Recorder {
if gw.BedrockRecorder != nil {
return gw.BedrockRecorder
}
return gateway_bedrock.NoopRecorder
}
// bedrockEndpointResolver returns an EndpointResolver over the configured
// pinned self-host endpoints (gw.BedrockEndpoints). A nil/empty map resolves
// nothing, so self-host models return ModelNotReady until configured.
+23 -1
View File
@@ -272,8 +272,13 @@ func (p *anthropicProvider) ConverseStream(ctx context.Context, modelID string,
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
// reqCtx/cancel let Close() abort the upstream request outright rather
// than merely stopping our own reads — a disconnected client must free
// the Anthropic-side generation, not just this goroutine.
reqCtx, cancel := context.WithCancel(ctx)
httpReq, err := http.NewRequestWithContext(reqCtx, 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 {
cancel()
slog.Error("anthropic stream: failed to build request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
@@ -284,6 +289,7 @@ func (p *anthropicProvider) ConverseStream(ctx context.Context, modelID string,
resp, err := p.httpClient.Do(httpReq) //nolint:gosec // G704: httpReq targets p.baseURL, not user input
if err != nil {
cancel()
slog.Error("anthropic stream: request failed", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
@@ -291,6 +297,7 @@ func (p *anthropicProvider) ConverseStream(ctx context.Context, modelID string,
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
respBody, _ := io.ReadAll(resp.Body)
_ = resp.Body.Close()
cancel()
slog.Error("anthropic stream: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody))
return nil, errors.New(mapUpstreamStatus(resp.StatusCode))
}
@@ -299,6 +306,7 @@ func (p *anthropicProvider) ConverseStream(ctx context.Context, modelID string,
resp: resp,
scanner: newSSEScanner(resp.Body),
start: time.Now(),
cancel: cancel,
}, nil
}
@@ -311,17 +319,31 @@ type anthropicConverseStreamSource struct {
resp *http.Response
scanner *sseScanner
start time.Time
cancel context.CancelFunc
inputTokens int64
outputTokens int64
}
var _ converseStreamSource = (*anthropicConverseStreamSource)(nil)
var _ usageReporter = (*anthropicConverseStreamSource)(nil)
// Close aborts the upstream request via cancel (not just closing the read
// side), so a disconnected client actually frees the Anthropic-side
// generation — real external cost, not just local compute.
func (s *anthropicConverseStreamSource) Close() error {
s.cancel()
return s.resp.Body.Close()
}
// usage always reports real, non-estimated tokens: Anthropic streams
// incremental usage throughout (input at message_start, output at
// message_delta), and unlike self-host generation those tokens carry real
// external cost that must never be guessed at.
func (s *anthropicConverseStreamSource) usage() (inputTokens, outputTokens int64, estimated bool) {
return s.inputTokens, s.outputTokens, false
}
func (s *anthropicConverseStreamSource) Next(_ context.Context) (ConverseStreamEvent, bool, error) {
for {
ev, ok, err := s.scanner.Next()
+3 -3
View File
@@ -56,7 +56,7 @@ func TestContract_ConverseStream_SelfHost_RealSDKConsumerDecodesFrames(t *testin
if !assert.NoError(t, err) {
return
}
err = ConverseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints)
err = ConverseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints, nil)
assert.NoError(t, err)
})
@@ -137,7 +137,7 @@ func TestContract_ConverseStream_Anthropic_RealSDKConsumerDecodesFrames(t *testi
if !assert.NoError(t, err) {
return
}
pumpConverseStream(r.Context(), fw, src, modelID)
pumpConverseStream(r.Context(), fw, src, modelID, streamRecordCtx{})
})
resp, err := svc.ConverseStream(&bedrockruntime.ConverseStreamInput{ModelId: aws.String(modelID)})
@@ -200,7 +200,7 @@ func TestContract_InvokeModelWithResponseStream_SelfHost_RealSDKConsumerDecodesF
if !assert.NoError(t, err) {
return
}
err = InvokeModelWithResponseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints, r.Header.Get("Content-Type"))
err = InvokeModelWithResponseStream(r.Context(), w, "000000000001", m[1], body, nil, endpoints, r.Header.Get("Content-Type"), nil)
assert.NoError(t, err)
})
+77 -5
View File
@@ -7,9 +7,13 @@ import (
"fmt"
"log/slog"
"net/http"
"strings"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/private/protocol/json/jsonutil"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/google/uuid"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
@@ -96,7 +100,13 @@ type converseStreamSource interface {
// 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 {
func ConverseStream(ctx context.Context, w http.ResponseWriter, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder) error {
if recorder == nil {
recorder = NoopRecorder
}
requestID := uuid.NewString()
start := time.Now()
input := new(bedrockruntime.ConverseStreamInput)
if len(body) > 0 {
if err := json.Unmarshal(body, input); err != nil {
@@ -104,7 +114,9 @@ func ConverseStream(ctx context.Context, w http.ResponseWriter, accountID, model
}
}
src, err := NewRouter(resolver, endpointResolver).ConverseStream(ctx, accountID, modelID, input)
entry, _ := lookupCatalogEntry(modelID) // Router.ConverseStream below re-validates; only its Provider tag is needed here.
src, err := NewRouter(resolver, endpointResolver, recorder).ConverseStream(ctx, accountID, modelID, input)
if err != nil {
return err
}
@@ -119,7 +131,19 @@ func ConverseStream(ctx context.Context, w http.ResponseWriter, accountID, model
return err
}
pumpConverseStream(ctx, fw, src, modelID)
inputJSON, jerr := json.Marshal(input)
if jerr != nil {
slog.Error("converse-stream: failed to marshal input for recording", "model", modelID, "err", jerr)
}
pumpConverseStream(ctx, fw, src, modelID, streamRecordCtx{
recorder: recorder,
requestID: requestID,
accountID: accountID,
backend: entry.Provider,
operation: OperationConverseStream,
start: start,
inputText: string(inputJSON),
})
return nil
}
@@ -128,9 +152,48 @@ func ConverseStream(ctx context.Context, w http.ResponseWriter, accountID, model
// 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) {
// HTTP status can no longer change either way. On every exit — clean end,
// client disconnect (ctx.Done), or upstream fault — the deferred closure
// records exactly one InvocationRecord via rc.recorder.
func pumpConverseStream(ctx context.Context, fw *frameWriter, src converseStreamSource, modelID string, rc streamRecordCtx) {
if rc.recorder == nil {
rc.recorder = NoopRecorder
}
partial := true
errCode := ""
var outputText strings.Builder
defer func() {
inputTokens, outputTokens, estimated := int64(0), int64(0), true
if ur, ok := src.(usageReporter); ok {
inputTokens, outputTokens, estimated = ur.usage()
}
rc.recorder.Record(ctx, InvocationRecord{
RequestID: rc.requestID,
AccountID: rc.accountID,
ModelID: modelID,
Operation: rc.operation,
Backend: rc.backend,
LatencyMs: time.Since(rc.start).Milliseconds(),
HTTPStatus: http.StatusOK, // the 200 header is already committed by the time the pump runs
ErrorCode: errCode,
InputTokens: inputTokens,
OutputTokens: outputTokens,
UsageEstimated: estimated,
Partial: partial,
InputText: rc.inputText,
OutputText: outputText.String(),
})
}()
for {
select {
case <-ctx.Done():
errCode = errCodeClientDisconnected
return
default:
}
event, ok, err := src.Next(ctx)
if err != nil {
excType := excInternalServerException
@@ -138,18 +201,26 @@ func pumpConverseStream(ctx context.Context, fw *frameWriter, src converseStream
if errors.As(err, &fault) {
excType = excModelStreamErrorException
}
errCode = awserrors.ValidErrorCodeFromError(err)
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)
errCode = errCodeClientDisconnected
}
return
}
if !ok {
partial = false
return
}
if event.Kind == converseStreamEventContentBlockDelta && event.ContentBlockDelta != nil && event.ContentBlockDelta.Delta != nil {
outputText.WriteString(aws.StringValue(event.ContentBlockDelta.Delta.Text))
}
payload, err := event.payload()
if err != nil {
errCode = awserrors.ErrorInternalError
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)
@@ -158,6 +229,7 @@ func pumpConverseStream(ctx context.Context, fw *frameWriter, src converseStream
}
if werr := fw.writeEvent(event.Kind.eventType(), payload); werr != nil {
errCode = errCodeClientDisconnected
slog.Error("converse-stream: failed to write frame, aborting", "model", modelID, "err", werr)
return
}
@@ -7,7 +7,9 @@ import (
"errors"
"net/http"
"net/http/httptest"
"sync"
"testing"
"time"
esv2 "github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
"github.com/aws/aws-sdk-go/aws"
@@ -17,6 +19,27 @@ import (
"github.com/stretchr/testify/require"
)
// fakeRecorder captures every InvocationRecord passed to Record, for tests
// asserting a record is emitted on a specific pump exit path.
type fakeRecorder struct {
mu sync.Mutex
records []InvocationRecord
}
func (f *fakeRecorder) Record(_ context.Context, rec InvocationRecord) {
f.mu.Lock()
defer f.mu.Unlock()
f.records = append(f.records, rec)
}
func (f *fakeRecorder) all() []InvocationRecord {
f.mu.Lock()
defer f.mu.Unlock()
out := make([]InvocationRecord, len(f.records))
copy(out, f.records)
return out
}
// fakeConverseStreamSource drives pumpConverseStream directly with a scripted
// event/error sequence, independent of any real backend.
type fakeConverseStreamSource struct {
@@ -81,7 +104,7 @@ func TestPumpConverseStream_MidStreamFaultWritesExceptionFrame(t *testing.T) {
err: newStreamFault(errors.New("upstream connection reset")),
}
pumpConverseStream(context.Background(), fw, src, "test-model")
pumpConverseStream(context.Background(), fw, src, "test-model", streamRecordCtx{})
frames := decodeAllFrames(t, rec.Body.Bytes())
require.Len(t, frames, 2)
@@ -101,7 +124,7 @@ func TestPumpConverseStream_NonFaultErrorUsesInternalServerException(t *testing.
require.NoError(t, err)
src := &fakeConverseStreamSource{err: errors.New("plain internal error")}
pumpConverseStream(context.Background(), fw, src, "test-model")
pumpConverseStream(context.Background(), fw, src, "test-model", streamRecordCtx{})
frames := decodeAllFrames(t, rec.Body.Bytes())
require.Len(t, frames, 1)
@@ -119,7 +142,7 @@ func TestPumpConverseStream_CleanEOFWritesNoExceptionFrame(t *testing.T) {
{Kind: converseStreamEventMessageStop, MessageStop: &bedrockruntime.MessageStopEvent{StopReason: aws.String(bedrockruntime.StopReasonEndTurn)}},
},
}
pumpConverseStream(context.Background(), fw, src, "test-model")
pumpConverseStream(context.Background(), fw, src, "test-model", streamRecordCtx{})
frames := decodeAllFrames(t, rec.Body.Bytes())
require.Len(t, frames, 2)
@@ -131,7 +154,7 @@ func TestConverseStream_UnknownModelReturnsResourceNotFoundPreHeader(t *testing.
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)
err := ConverseStream(context.Background(), rec, "000000000001", "does.not-exist-v1:0", body, nil, 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
@@ -141,7 +164,7 @@ func TestConverseStream_UnknownModelReturnsResourceNotFoundPreHeader(t *testing.
func TestConverseStream_MalformedBodyReturnsValidationException(t *testing.T) {
rec := httptest.NewRecorder()
err := ConverseStream(context.Background(), rec, "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte("{not-json"), nil, nil)
err := ConverseStream(context.Background(), rec, "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte("{not-json"), nil, nil, nil)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorValidationException, err.Error())
}
@@ -163,7 +186,7 @@ func TestConverseStream_SelfHostHappyPath_WritesFramedTaxonomy(t *testing.T) {
})
require.NoError(t, err)
err = ConverseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
err = ConverseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil)
require.NoError(t, err)
assert.Equal(t, http.StatusOK, rec.Code)
@@ -180,6 +203,123 @@ func TestConverseStream_SelfHostHappyPath_WritesFramedTaxonomy(t *testing.T) {
assert.Equal(t, "metadata", frames[6].Type)
}
func TestPumpConverseStream_RecordsInvocationOnCleanEnd(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: converseStreamEventContentBlockDelta, ContentBlockDelta: &bedrockruntime.ContentBlockDeltaEvent{Delta: &bedrockruntime.ContentBlockDelta{Text: aws.String("hi")}}},
{Kind: converseStreamEventMessageStop, MessageStop: &bedrockruntime.MessageStopEvent{StopReason: aws.String(bedrockruntime.StopReasonEndTurn)}},
},
}
fr := &fakeRecorder{}
pumpConverseStream(context.Background(), fw, src, "test-model", streamRecordCtx{recorder: fr, requestID: "req-clean", accountID: "acct-a"})
records := fr.all()
require.Len(t, records, 1)
got := records[0]
assert.Equal(t, "req-clean", got.RequestID)
assert.Equal(t, "acct-a", got.AccountID)
assert.False(t, got.Partial)
assert.Empty(t, got.ErrorCode)
assert.Equal(t, "hi", got.OutputText)
}
func TestPumpConverseStream_RecordsInvocationOnClientDisconnect(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
ctx, cancel := context.WithCancel(context.Background())
cancel() // simulate a client that has already disconnected
src := &fakeConverseStreamSource{
events: []ConverseStreamEvent{
{Kind: converseStreamEventMessageStart, MessageStart: &bedrockruntime.MessageStartEvent{Role: aws.String("assistant")}},
},
}
fr := &fakeRecorder{}
pumpConverseStream(ctx, fw, src, "test-model", streamRecordCtx{recorder: fr, requestID: "req-disconnect"})
records := fr.all()
require.Len(t, records, 1)
got := records[0]
assert.True(t, got.Partial)
assert.Equal(t, errCodeClientDisconnected, got.ErrorCode)
// A disconnect must be caught before draining any queued event.
assert.Empty(t, rec.Body.Bytes())
}
func TestPumpConverseStream_RecordsInvocationOnUpstreamFault(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
src := &fakeConverseStreamSource{err: newStreamFault(errors.New("upstream connection reset"))}
fr := &fakeRecorder{}
pumpConverseStream(context.Background(), fw, src, "test-model", streamRecordCtx{recorder: fr, requestID: "req-fault"})
records := fr.all()
require.Len(t, records, 1)
got := records[0]
assert.True(t, got.Partial)
assert.NotEmpty(t, got.ErrorCode)
}
// TestConverseStream_ClientDisconnectAbortsUpstreamRequest proves the
// disconnect path aborts the real upstream request rather than merely
// draining it: the fixture handler blocks on r.Context().Done() and never
// sends EOF, so the only way the handler unblocks is via ConverseStream's
// deferred src.Close() -> cancel() reaching the upstream request context.
func TestConverseStream_ClientDisconnectAbortsUpstreamRequest(t *testing.T) {
serverCtxDone := make(chan struct{})
firstChunkSent := make(chan struct{})
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\":{\"content\":\"hi\"}}]}\n\n"))
w.(http.Flusher).Flush()
close(firstChunkSent)
<-r.Context().Done()
close(serverCtxDone)
}))
defer ts.Close()
modelID := "meta.llama3-2-1b-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)
ctx, cancel := context.WithCancel(context.Background())
done := make(chan struct{})
go func() {
defer close(done)
_ = ConverseStream(ctx, rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil)
}()
select {
case <-firstChunkSent:
case <-time.After(2 * time.Second):
t.Fatal("server never sent its first chunk")
}
cancel()
select {
case <-serverCtxDone:
// Upstream request context was canceled: aborted, not drained.
case <-time.After(2 * time.Second):
t.Fatal("upstream request was never aborted after client disconnect")
}
<-done
}
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")
@@ -197,7 +337,7 @@ func TestConverseStream_NonFlusherWriter_ReturnsPreHeaderError(t *testing.T) {
})
require.NoError(t, err)
err = ConverseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
err = ConverseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorInternalError, err.Error())
assert.False(t, w.wroteHeader)
+55 -18
View File
@@ -4,7 +4,9 @@ import (
"context"
"errors"
"strings"
"time"
"github.com/google/uuid"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
@@ -22,30 +24,59 @@ type InvokeAdapter interface {
type InvokeRouter struct {
resolver CredentialResolver
endpointResolver EndpointResolver
recorder Recorder
}
// NewInvokeRouter constructs an InvokeRouter. A nil resolver or
// endpointResolver falls back to a resolver/resolver that finds nothing, so
// an InvokeRouter is always safe to use even before the real stores are
// wired in.
func NewInvokeRouter(resolver CredentialResolver, endpointResolver EndpointResolver) *InvokeRouter {
// NewInvokeRouter constructs an InvokeRouter. A nil resolver,
// endpointResolver, or recorder falls back to a no-op implementation, so an
// InvokeRouter is always safe to use even before the real stores are wired
// in.
func NewInvokeRouter(resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder) *InvokeRouter {
if resolver == nil {
resolver = NoopCredentialResolver
}
if endpointResolver == nil {
endpointResolver = NewStaticEndpointResolver(nil)
}
return &InvokeRouter{resolver: resolver, endpointResolver: endpointResolver}
if recorder == nil {
recorder = NoopRecorder
}
return &InvokeRouter{resolver: resolver, endpointResolver: endpointResolver, recorder: recorder}
}
// InvokeModel routes modelID to its family adapter via the catalog. Unknown
// modelIds and unresolvable vendors return ResourceNotFoundException; a
// vendor with no resolvable credential returns AccessDeniedException.
func (rt *InvokeRouter) InvokeModel(ctx context.Context, accountID, modelID string, body []byte) ([]byte, string, error) {
// vendor with no resolvable credential returns AccessDeniedException. Every
// exit records an InvocationRecord via the deferred closure.
func (rt *InvokeRouter) InvokeModel(ctx context.Context, accountID, modelID string, body []byte) (respBody []byte, contentType string, err error) {
requestID := uuid.NewString()
start := time.Now()
var backend string
defer func() {
httpStatus, code := recordOutcome(err)
inputTokens, outputTokens, _ := extractTokenUsage(backend, respBody)
rt.recorder.Record(ctx, InvocationRecord{
RequestID: requestID,
AccountID: accountID,
ModelID: modelID,
Operation: OperationInvokeModel,
Backend: backend,
LatencyMs: time.Since(start).Milliseconds(),
HTTPStatus: httpStatus,
ErrorCode: code,
InputTokens: inputTokens,
OutputTokens: outputTokens,
InputText: string(body),
OutputText: string(respBody),
})
}()
entry, ok := lookupCatalogEntry(modelID)
if !ok {
return nil, "", errors.New(awserrors.ErrorResourceNotFoundException)
err = errors.New(awserrors.ErrorResourceNotFoundException)
return nil, "", err
}
backend = entry.Provider
var a InvokeAdapter
switch {
@@ -54,29 +85,35 @@ func (rt *InvokeRouter) InvokeModel(ctx context.Context, accountID, modelID stri
case strings.HasPrefix(entry.Provider, providerPrefix):
switch strings.TrimPrefix(entry.Provider, providerPrefix) {
case vendorAnthropic:
key, ok, err := rt.resolver.Resolve(ctx, accountID, vendorAnthropic)
var key string
var resolvable bool
key, resolvable, err = rt.resolver.Resolve(ctx, accountID, vendorAnthropic)
if err != nil {
return nil, "", err
}
if !ok {
return nil, "", errors.New(awserrors.ErrorAccessDeniedException)
if !resolvable {
err = errors.New(awserrors.ErrorAccessDeniedException)
return nil, "", err
}
a = newAnthropicInvokeAdapter(key)
default:
return nil, "", errors.New(awserrors.ErrorResourceNotFoundException)
err = errors.New(awserrors.ErrorResourceNotFoundException)
return nil, "", err
}
default:
return nil, "", errors.New(awserrors.ErrorResourceNotFoundException)
err = errors.New(awserrors.ErrorResourceNotFoundException)
return nil, "", err
}
return a.InvokeModel(ctx, modelID, body)
respBody, contentType, err = a.InvokeModel(ctx, modelID, body)
return respBody, contentType, err
}
// InvokeModel is the bedrock-runtime InvokeModel entry point used by the
// gateway route table. resolver and endpointResolver may be nil;
// gateway route table. resolver, endpointResolver, and recorder may be nil;
// NewInvokeRouter supplies no-op fallbacks.
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)
func InvokeModel(ctx context.Context, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder) ([]byte, string, error) {
return NewInvokeRouter(resolver, endpointResolver, recorder).InvokeModel(ctx, accountID, modelID, body)
}
// InvokeStreamAdapter is the optional streaming capability an InvokeAdapter
+57 -3
View File
@@ -134,8 +134,13 @@ func (a *anthropicInvokeAdapter) InvokeModelWithResponseStream(ctx context.Conte
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
// reqCtx/cancel let Close() abort the upstream request outright rather
// than merely stopping our own reads — a disconnected client must free
// the Anthropic-side generation, not just this goroutine.
reqCtx, cancel := context.WithCancel(ctx)
httpReq, err := http.NewRequestWithContext(reqCtx, 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 {
cancel()
slog.Error("anthropic invoke-stream: failed to build request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
@@ -146,6 +151,7 @@ func (a *anthropicInvokeAdapter) InvokeModelWithResponseStream(ctx context.Conte
resp, err := a.httpClient.Do(httpReq) //nolint:gosec // G704: httpReq targets a.baseURL, not user input
if err != nil {
cancel()
slog.Error("anthropic invoke-stream: request failed", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
@@ -153,23 +159,34 @@ func (a *anthropicInvokeAdapter) InvokeModelWithResponseStream(ctx context.Conte
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
respBody, _ := io.ReadAll(resp.Body)
_ = resp.Body.Close()
cancel()
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
return &anthropicInvokeStreamSource{resp: resp, scanner: newSSEScanner(resp.Body), cancel: cancel}, nil
}
// anthropicInvokeStreamSource forwards each Anthropic SSE "data:" line
// verbatim as the chunk payload, skipping keepalive "ping" events.
// verbatim as the chunk payload, skipping keepalive "ping" events. It also
// opportunistically taps message_start/message_delta for token usage
// without altering the forwarded bytes.
type anthropicInvokeStreamSource struct {
resp *http.Response
scanner *sseScanner
cancel context.CancelFunc
inputTokens, outputTokens int64
}
var _ invokeStreamSource = (*anthropicInvokeStreamSource)(nil)
var _ usageReporter = (*anthropicInvokeStreamSource)(nil)
// Close aborts the upstream request via cancel (not just closing the read
// side), so a disconnected client actually frees the Anthropic-side
// generation — real external cost, not just local compute.
func (s *anthropicInvokeStreamSource) Close() error {
s.cancel()
return s.resp.Body.Close()
}
@@ -188,6 +205,43 @@ func (s *anthropicInvokeStreamSource) Next(_ context.Context) ([]byte, bool, err
if ev.Event == "ping" || ev.Data == "" {
continue
}
s.tapUsage(ev.Event, ev.Data)
return []byte(ev.Data), true, nil
}
}
// tapUsage opportunistically extracts token usage from message_start/
// message_delta events without altering the bytes InvokeModelWithResponseStream
// forwards verbatim: Bedrock's Claude invoke-stream contract is "forward
// Anthropic's own bytes", not "re-derive them"; a parse miss is silently
// ignored since usage is only ever read back out for the invocation record.
func (s *anthropicInvokeStreamSource) tapUsage(event, data string) {
switch event {
case "message_start":
var payload struct {
Message struct {
Usage struct {
InputTokens int64 `json:"input_tokens"`
} `json:"usage"`
} `json:"message"`
}
if json.Unmarshal([]byte(data), &payload) == nil {
s.inputTokens = payload.Message.Usage.InputTokens
}
case "message_delta":
var payload struct {
Usage struct {
OutputTokens int64 `json:"output_tokens"`
} `json:"usage"`
}
if json.Unmarshal([]byte(data), &payload) == nil {
s.outputTokens = payload.Usage.OutputTokens
}
}
}
// usage always reports real, non-estimated tokens, mirroring
// anthropicConverseStreamSource.usage.
func (s *anthropicInvokeStreamSource) usage() (inputTokens, outputTokens int64, estimated bool) {
return s.inputTokens, s.outputTokens, false
}
+28 -1
View File
@@ -244,8 +244,13 @@ func (a *llamaInvokeAdapter) InvokeModelWithResponseStream(ctx context.Context,
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
// reqCtx/cancel let Close() abort the upstream request outright rather
// than merely stopping our own reads — a disconnected client must free
// the vLLM-side generation, not just this goroutine.
reqCtx, cancel := context.WithCancel(ctx)
httpReq, err := http.NewRequestWithContext(reqCtx, http.MethodPost, baseURL+llamaCompletionsPath, bytes.NewReader(reqBody)) //nolint:gosec // G704: baseURL is a resolved pinned self-host endpoint, not user input
if err != nil {
cancel()
slog.Error("llama invoke-stream: failed to build request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
@@ -254,6 +259,7 @@ func (a *llamaInvokeAdapter) InvokeModelWithResponseStream(ctx context.Context,
resp, err := a.httpClient.Do(httpReq) //nolint:gosec // G704: httpReq targets the resolved pinned self-host endpoint, not user input
if err != nil {
cancel()
slog.Error("llama invoke-stream: request failed", "model", modelID, "endpoint", baseURL, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
@@ -261,6 +267,7 @@ func (a *llamaInvokeAdapter) InvokeModelWithResponseStream(ctx context.Context,
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
respBody, _ := io.ReadAll(resp.Body)
_ = resp.Body.Close()
cancel()
slog.Error("llama invoke-stream: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody))
return nil, errors.New(mapUpstreamStatus(resp.StatusCode))
}
@@ -269,6 +276,7 @@ func (a *llamaInvokeAdapter) InvokeModelWithResponseStream(ctx context.Context,
resp: resp,
scanner: newSSEScanner(resp.Body),
start: time.Now(),
cancel: cancel,
}, nil
}
@@ -284,19 +292,36 @@ type llamaInvokeStreamSource struct {
resp *http.Response
scanner *sseScanner
start time.Time
cancel context.CancelFunc
gotFirstByte bool
firstByte time.Time
promptTokens, completionTokens int
usageKnown bool
deltaChunks int64
pendingFinal *llamaInvokeStreamFinalChunk
}
var _ invokeStreamSource = (*llamaInvokeStreamSource)(nil)
var _ usageReporter = (*llamaInvokeStreamSource)(nil)
// Close aborts the upstream request via cancel (not just closing the read
// side), so a disconnected client actually frees the vLLM-side generation.
func (s *llamaInvokeStreamSource) Close() error {
s.cancel()
return s.resp.Body.Close()
}
// usage returns real usage once vLLM's trailing usage-only chunk has set it
// (usageKnown); until then it falls back to a content-delta chunk count as
// an estimated output-token proxy, flagged accordingly.
func (s *llamaInvokeStreamSource) usage() (inputTokens, outputTokens int64, estimated bool) {
if s.usageKnown {
return int64(s.promptTokens), int64(s.completionTokens), false
}
return 0, s.deltaChunks, true
}
func (s *llamaInvokeStreamSource) Next(_ context.Context) ([]byte, bool, error) {
for {
ev, ok, err := s.scanner.Next()
@@ -329,6 +354,7 @@ func (s *llamaInvokeStreamSource) Next(_ context.Context) ([]byte, bool, error)
if chunk.Usage != nil {
s.promptTokens = chunk.Usage.PromptTokens
s.completionTokens = chunk.Usage.CompletionTokens
s.usageKnown = true
}
if len(chunk.Choices) == 0 {
// The trailing usage-only chunk: emit the deferred final chunk
@@ -348,6 +374,7 @@ func (s *llamaInvokeStreamSource) Next(_ context.Context) ([]byte, bool, error)
continue
}
s.deltaChunks++
out, err := json.Marshal(llamaInvokeStreamChunk{Generation: choice.Text})
if err != nil {
return nil, false, newStreamFault(fmt.Errorf("llama invoke-stream: marshal chunk: %w", err))
+70 -4
View File
@@ -5,6 +5,10 @@ import (
"errors"
"log/slog"
"net/http"
"time"
"github.com/google/uuid"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
// invokeStreamSource is the backend-agnostic InvokeModelWithResponseStream
@@ -24,7 +28,17 @@ type invokeStreamSource interface {
// 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 {
// resolver, endpointResolver, and recorder may be nil; NewInvokeStreamRouter
// and the internal NoopRecorder fallback keep this call safe either way.
func InvokeModelWithResponseStream(ctx context.Context, w http.ResponseWriter, accountID, modelID string, body []byte, resolver CredentialResolver, endpointResolver EndpointResolver, requestContentType string, recorder Recorder) error {
if recorder == nil {
recorder = NoopRecorder
}
requestID := uuid.NewString()
start := time.Now()
entry, _ := lookupCatalogEntry(modelID) // InvokeStreamRouter below re-validates; only its Provider tag is needed here.
src, err := NewInvokeStreamRouter(resolver, endpointResolver).InvokeModelWithResponseStream(ctx, accountID, modelID, body)
if err != nil {
return err
@@ -41,7 +55,15 @@ func InvokeModelWithResponseStream(ctx context.Context, w http.ResponseWriter, a
}
slog.Debug("invoke-with-response-stream: streaming", "model", modelID, "request_content_type", requestContentType)
pumpInvokeStream(ctx, fw, src, modelID)
pumpInvokeStream(ctx, fw, src, modelID, streamRecordCtx{
recorder: recorder,
requestID: requestID,
accountID: accountID,
backend: entry.Provider,
operation: OperationInvokeModelWithResponseStream,
start: start,
inputText: string(body),
})
return nil
}
@@ -49,9 +71,49 @@ func InvokeModelWithResponseStream(ctx context.Context, w http.ResponseWriter, 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) {
// status can no longer change either way. On every exit — clean end, client
// disconnect (ctx.Done), or upstream fault — the deferred closure records
// exactly one InvocationRecord via rc.recorder.
func pumpInvokeStream(ctx context.Context, fw *frameWriter, src invokeStreamSource, modelID string, rc streamRecordCtx) {
if rc.recorder == nil {
rc.recorder = NoopRecorder
}
partial := true
errCode := ""
defer func() {
inputTokens, outputTokens, estimated := int64(0), int64(0), true
if ur, ok := src.(usageReporter); ok {
inputTokens, outputTokens, estimated = ur.usage()
}
rc.recorder.Record(ctx, InvocationRecord{
RequestID: rc.requestID,
AccountID: rc.accountID,
ModelID: modelID,
Operation: rc.operation,
Backend: rc.backend,
LatencyMs: time.Since(rc.start).Milliseconds(),
HTTPStatus: http.StatusOK, // the 200 header is already committed by the time the pump runs
ErrorCode: errCode,
InputTokens: inputTokens,
OutputTokens: outputTokens,
UsageEstimated: estimated,
Partial: partial,
InputText: rc.inputText,
// OutputText is intentionally left empty: chunks are opaque
// provider-native bytes (Llama JSON or raw Anthropic SSE data),
// with no generic way to recover assistant text across both.
})
}()
for {
select {
case <-ctx.Done():
errCode = errCodeClientDisconnected
return
default:
}
chunk, ok, err := src.Next(ctx)
if err != nil {
excType := excInternalServerException
@@ -59,17 +121,21 @@ func pumpInvokeStream(ctx context.Context, fw *frameWriter, src invokeStreamSour
if errors.As(err, &fault) {
excType = excModelStreamErrorException
}
errCode = awserrors.ValidErrorCodeFromError(err)
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)
errCode = errCodeClientDisconnected
}
return
}
if !ok {
partial = false
return
}
if werr := fw.writeChunk(chunk); werr != nil {
errCode = errCodeClientDisconnected
slog.Error("invoke-with-response-stream: failed to write frame, aborting", "model", modelID, "err", werr)
return
}
+54 -5
View File
@@ -46,7 +46,7 @@ func TestPumpInvokeStream_MidStreamFaultWritesExceptionFrame(t *testing.T) {
chunks: [][]byte{[]byte(`{"generation":"hi"}`)},
err: newStreamFault(errors.New("upstream connection reset")),
}
pumpInvokeStream(context.Background(), fw, src, "test-model")
pumpInvokeStream(context.Background(), fw, src, "test-model", streamRecordCtx{})
frames := decodeAllFrames(t, rec.Body.Bytes())
require.Len(t, frames, 2)
@@ -61,16 +61,65 @@ func TestPumpInvokeStream_NonFaultErrorUsesInternalServerException(t *testing.T)
require.NoError(t, err)
src := &fakeInvokeStreamSource{err: errors.New("plain internal error")}
pumpInvokeStream(context.Background(), fw, src, "test-model")
pumpInvokeStream(context.Background(), fw, src, "test-model", streamRecordCtx{})
frames := decodeAllFrames(t, rec.Body.Bytes())
require.Len(t, frames, 1)
assert.Equal(t, excInternalServerException, frames[0].Type)
}
func TestPumpInvokeStream_RecordsInvocationOnCleanEnd(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
src := &fakeInvokeStreamSource{chunks: [][]byte{[]byte(`{"generation":"hi"}`)}}
fr := &fakeRecorder{}
pumpInvokeStream(context.Background(), fw, src, "test-model", streamRecordCtx{recorder: fr, requestID: "req-clean"})
records := fr.all()
require.Len(t, records, 1)
assert.Equal(t, "req-clean", records[0].RequestID)
assert.False(t, records[0].Partial)
assert.Empty(t, records[0].ErrorCode)
}
func TestPumpInvokeStream_RecordsInvocationOnClientDisconnect(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
ctx, cancel := context.WithCancel(context.Background())
cancel()
src := &fakeInvokeStreamSource{chunks: [][]byte{[]byte(`{"generation":"hi"}`)}}
fr := &fakeRecorder{}
pumpInvokeStream(ctx, fw, src, "test-model", streamRecordCtx{recorder: fr, requestID: "req-disconnect"})
records := fr.all()
require.Len(t, records, 1)
assert.True(t, records[0].Partial)
assert.Equal(t, errCodeClientDisconnected, records[0].ErrorCode)
}
func TestPumpInvokeStream_RecordsInvocationOnUpstreamFault(t *testing.T) {
rec := httptest.NewRecorder()
fw, err := newFrameWriter(rec)
require.NoError(t, err)
src := &fakeInvokeStreamSource{err: newStreamFault(errors.New("upstream connection reset"))}
fr := &fakeRecorder{}
pumpInvokeStream(context.Background(), fw, src, "test-model", streamRecordCtx{recorder: fr, requestID: "req-fault"})
records := fr.all()
require.Len(t, records, 1)
assert.True(t, records[0].Partial)
assert.NotEmpty(t, records[0].ErrorCode)
}
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")
err := InvokeModelWithResponseStream(context.Background(), rec, "000000000001", "does.not-exist-v1:0", []byte(`{}`), nil, nil, "application/json", nil)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
assert.Equal(t, 0, rec.Body.Len())
@@ -88,7 +137,7 @@ func TestInvokeModelWithResponseStream_SelfHostHappyPath_WritesChunkFrames(t *te
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")
err := InvokeModelWithResponseStream(context.Background(), rec, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), "application/json", nil)
require.NoError(t, err)
assert.Equal(t, http.StatusOK, rec.Code)
@@ -113,7 +162,7 @@ func TestInvokeModelWithResponseStream_NonFlusherWriter_ReturnsPreHeaderError(t
w := &nonFlusherWriter{}
body := []byte(`{"prompt":"hello"}`)
err := InvokeModelWithResponseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), "application/json")
err := InvokeModelWithResponseStream(context.Background(), w, "000000000001", modelID, body, nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), "application/json", nil)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorInternalError, err.Error())
assert.False(t, w.wroteHeader)
+6 -6
View File
@@ -24,7 +24,7 @@ func TestInvokeRouter_SelfHostSuccess(t *testing.T) {
defer ts.Close()
modelID := "meta.llama3-2-1b-instruct-v1:0"
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil)
respBody, contentType, err := rt.InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`))
require.NoError(t, err)
@@ -36,21 +36,21 @@ func TestInvokeRouter_SelfHostSuccess(t *testing.T) {
}
func TestInvokeRouter_UnknownModelReturnsResourceNotFound(t *testing.T) {
rt := NewInvokeRouter(nil, nil)
rt := NewInvokeRouter(nil, nil, nil)
_, _, err := rt.InvokeModel(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`))
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
}
func TestInvokeRouter_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) {
rt := NewInvokeRouter(stubCredentialResolver{ok: false}, nil)
rt := NewInvokeRouter(stubCredentialResolver{ok: false}, nil, nil)
_, _, err := rt.InvokeModel(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", []byte(`{}`))
require.Error(t, err)
assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error())
}
func TestInvokeRouter_SelfHostNoEndpointReturnsModelNotReady(t *testing.T) {
rt := NewInvokeRouter(nil, nil)
rt := NewInvokeRouter(nil, nil, nil)
_, _, err := rt.InvokeModel(context.Background(), "000000000001", "meta.llama3-2-1b-instruct-v1:0", []byte(`{"prompt":"hello"}`))
require.Error(t, err)
assert.Equal(t, awserrors.ErrorModelNotReadyException, err.Error())
@@ -68,7 +68,7 @@ func TestInvokeModel_PackageEntryPoint_SelfHostSuccess(t *testing.T) {
defer ts.Close()
modelID := "meta.llama3-2-1b-instruct-v1:0"
respBody, contentType, err := InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
respBody, contentType, err := InvokeModel(context.Background(), "000000000001", modelID, []byte(`{"prompt":"hello"}`), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil)
require.NoError(t, err)
assert.Equal(t, "application/json", contentType)
@@ -78,7 +78,7 @@ func TestInvokeModel_PackageEntryPoint_SelfHostSuccess(t *testing.T) {
}
func TestInvokeModel_PackageEntryPoint_UnknownModel(t *testing.T) {
_, _, err := InvokeModel(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`), nil, nil)
_, _, err := InvokeModel(context.Background(), "000000000001", "does.not-exist-v1:0", []byte(`{}`), nil, nil, nil)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
}
+620
View File
@@ -0,0 +1,620 @@
package gateway_bedrock
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"log/slog"
"net/http"
"strings"
"sync"
"sync/atomic"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrock"
"github.com/aws/aws-sdk-go/service/s3"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/kvutil"
"github.com/mulgadc/spinifex/spinifex/objectstore"
"github.com/nats-io/nats.go/jetstream"
)
// Bedrock-runtime operation names carried on every InvocationRecord.
const (
OperationConverse = "Converse"
OperationConverseStream = "ConverseStream"
OperationInvokeModel = "InvokeModel"
OperationInvokeModelWithResponseStream = "InvokeModelWithResponseStream"
)
const (
// InvocationStreamName is the JetStream stream carrying every Bedrock
// invocation record. LimitsPolicy retention (not WorkQueuePolicy) is
// required: two independent durable consumers attach to it — this
// package's own S3+ES delivery consumer (InvocationDeliveryConsumer) and
// a separate usage/cost-aggregation consumer owned elsewhere — and each
// must see every message regardless of the other's ack position.
InvocationStreamName = "OCHRE_INVOCATIONS"
// InvocationStreamSubject is the single subject the stream captures.
InvocationStreamSubject = "bedrock.invocations"
// InvocationDeliveryConsumer is this package's durable pull consumer:
// one record in, one S3 write plus one metadata-only log line out.
InvocationDeliveryConsumer = "ochre-invocation-delivery"
// invocationStreamMaxAge bounds disk growth for a stream nothing besides
// its durable consumers is required to keep forever.
invocationStreamMaxAge = 30 * 24 * time.Hour
)
// EnsureInvocationStream idempotently creates (or updates) the invocation
// stream. Safe to call from every gateway node at boot. replicas must match
// the cluster's node count: at replicas=1, losing the one node holding this
// stream loses billing/audit records even though the control plane survives
// on the others, defeating the point of using JetStream over a lossy buffer.
func EnsureInvocationStream(ctx context.Context, js jetstream.JetStream, replicas int) (jetstream.Stream, error) {
stream, err := js.CreateOrUpdateStream(ctx, jetstream.StreamConfig{
Name: InvocationStreamName,
Subjects: []string{InvocationStreamSubject},
Retention: jetstream.LimitsPolicy,
Storage: jetstream.FileStorage,
MaxAge: invocationStreamMaxAge,
Replicas: replicas,
})
if err != nil {
return nil, fmt.Errorf("ensure invocation stream: %w", err)
}
return stream, nil
}
// EnsureDeliveryConsumer idempotently creates (or updates) this package's
// durable pull consumer on the invocation stream.
func EnsureDeliveryConsumer(ctx context.Context, js jetstream.JetStream) (jetstream.Consumer, error) {
consumer, err := js.CreateOrUpdateConsumer(ctx, InvocationStreamName, jetstream.ConsumerConfig{
Durable: InvocationDeliveryConsumer,
FilterSubject: InvocationStreamSubject,
AckPolicy: jetstream.AckExplicitPolicy,
MaxDeliver: 5,
AckWait: time.Minute,
})
if err != nil {
return nil, fmt.Errorf("ensure invocation delivery consumer: %w", err)
}
return consumer, nil
}
// errCodeClientDisconnected marks an invocation record ended because the
// client went away (context cancelled, or a frame write failed), not because
// of any upstream/model fault. It is this package's own bookkeeping value,
// not a registered AWS error code — InvocationRecord is an internal audit
// trail, not an AWS API response.
const errCodeClientDisconnected = "ClientDisconnected"
// InvocationRecord is one Bedrock invocation's durable audit trail: what was
// called, by whom, at what cost, and how it ended. It is published to the
// invocation stream on every exit path (clean end, client disconnect,
// upstream fault) and fanned out from there to per-account S3 delivery and
// the platform's metadata-only usage/cost sink.
//
// InputText/OutputText are populated by the caller but are only actually
// delivered when StreamRecorder finds AccountID's logging config enables
// text delivery; every other field is metadata safe for the shared platform
// log sink to read.
type InvocationRecord struct {
RequestID string `json:"requestId"`
AccountID string `json:"accountId"`
ModelID string `json:"modelId"`
Operation string `json:"operation"`
Backend string `json:"backend"`
Timestamp time.Time `json:"timestamp"`
LatencyMs int64 `json:"latencyMs"`
HTTPStatus int `json:"httpStatus"`
ErrorCode string `json:"errorCode,omitempty"`
InputTokens int64 `json:"inputTokens"`
OutputTokens int64 `json:"outputTokens"`
UsageEstimated bool `json:"usageEstimated"`
// Partial marks a record whose invocation did not reach a clean end
// (client disconnect, upstream fault, mid-stream cancellation). Partial
// records are still billable — usage reflects whatever tokens were
// actually generated — but flagged so downstream aggregation can treat
// them distinctly if it chooses to.
Partial bool `json:"partial"`
InputText string `json:"inputText,omitempty"`
OutputText string `json:"outputText,omitempty"`
}
// Recorder durably records one invocation. Implementations MUST NOT let a
// recording failure fail the invocation itself — publish errors are logged
// and counted, never returned, so a struggling stream never takes billable
// model traffic down with it.
type Recorder interface {
Record(ctx context.Context, rec InvocationRecord)
}
// usageReporter is an optional capability a streaming source may implement
// to expose token usage collected so far, for the invocation record the
// pump emits on every exit path. vLLM only learns real usage from its
// trailing usage chunk, so it falls back to a content-delta-chunk output
// estimate (estimated=true) when the stream ends before that arrives.
// Anthropic streams real incremental usage throughout and never estimates —
// unlike self-host generation, its tokens carry real external cost.
type usageReporter interface {
usage() (inputTokens, outputTokens int64, estimated bool)
}
// streamRecordCtx carries the fixed per-call metadata a streaming pump needs
// to build its InvocationRecord once the stream ends, on every exit path.
type streamRecordCtx struct {
recorder Recorder
requestID string
accountID string
backend string
operation string
start time.Time
inputText string
}
type noopRecorder struct{}
func (noopRecorder) Record(context.Context, InvocationRecord) {}
// NoopRecorder discards every record. Used as the fallback wherever no
// StreamRecorder is configured (e.g. unit tests of unrelated routes).
var NoopRecorder Recorder = noopRecorder{}
var _ Recorder = noopRecorder{}
// recordOutcome maps a Converse/InvokeModel error to the HTTP status and
// AWS error code an InvocationRecord carries. A nil err is success.
func recordOutcome(err error) (httpStatus int, code string) {
if err == nil {
return http.StatusOK, ""
}
code = awserrors.ValidErrorCodeFromError(err)
return awserrors.LookupErrorMessage("bedrock-runtime", code).HTTPCode, code
}
// extractTokenUsage best-effort parses a non-streaming InvokeModel response
// body for real token counts, reusing each backend's own response struct
// rather than a bespoke shape. ok is false when the body doesn't decode —
// callers must treat that as "unknown", not as a zero-token estimate.
func extractTokenUsage(backend string, respBody []byte) (inputTokens, outputTokens int64, ok bool) {
switch {
case backend == tierSelfHost:
var lr llamaInvokeResponse
if json.Unmarshal(respBody, &lr) != nil {
return 0, 0, false
}
return int64(lr.PromptTokenCount), int64(lr.GenerationTokenCount), true
case strings.HasPrefix(backend, providerPrefix):
var ar anthropicResponse
if json.Unmarshal(respBody, &ar) != nil {
return 0, 0, false
}
return ar.Usage.InputTokens, ar.Usage.OutputTokens, true
default:
return 0, 0, false
}
}
// recordPublishTimeout bounds how long Record waits for the JetStream
// publish once detached from the caller's (possibly already-cancelled)
// context, so a struggling stream can't hang a pump's defer indefinitely.
const recordPublishTimeout = 5 * time.Second
// StreamRecorder publishes InvocationRecords to the invocation stream,
// including body text only when the account's own logging config enables
// delivery — LoggingConfigReader is consulted per record, never cached past
// a single Record call, so a config change takes effect immediately.
//
// A publish failure never fails the caller's invocation: it is logged and
// counted (Dropped), matching the platform's stance that billing/audit
// durability must never come at the cost of the model call itself.
type StreamRecorder struct {
js jetstream.JetStream
configs LoggingConfigReader
dropped atomic.Int64
}
var _ Recorder = (*StreamRecorder)(nil)
// NewStreamRecorder constructs a StreamRecorder. A nil configs falls back to
// a reader that never enables body delivery, so records always carry
// metadata only until a real LoggingConfigStore is wired in.
func NewStreamRecorder(js jetstream.JetStream, configs LoggingConfigReader) *StreamRecorder {
if configs == nil {
configs = noopLoggingConfigReader{}
}
return &StreamRecorder{js: js, configs: configs}
}
// Dropped returns the number of records this recorder has failed to
// publish, for monitoring and tests.
func (r *StreamRecorder) Dropped() int64 {
return r.dropped.Load()
}
// Record fills in Timestamp, clears body text unless AccountID's logging
// config enables it, and publishes to the invocation stream keyed by
// RequestID (jetstream.WithMsgID) so a retried publish dedupes downstream.
// ctx is detached from the caller's cancellation — a client disconnect must
// not prevent recording the very disconnect it caused — but bounded by
// recordPublishTimeout so a stalled stream can't hang the caller's defer.
func (r *StreamRecorder) Record(ctx context.Context, rec InvocationRecord) {
rec.Timestamp = time.Now()
cfg, enabled, err := r.configs.Get(ctx, rec.AccountID)
if err != nil {
slog.Error("bedrock: logging config lookup failed, dropping record body", "account", rec.AccountID, "err", err)
enabled = false
}
if !enabled || !cfg.TextDataDeliveryEnabled {
rec.InputText = ""
rec.OutputText = ""
}
payload, err := json.Marshal(rec)
if err != nil {
slog.Error("bedrock: failed to marshal invocation record", "request_id", rec.RequestID, "err", err)
r.dropped.Add(1)
return
}
pubCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), recordPublishTimeout)
defer cancel()
if _, err := r.js.Publish(pubCtx, InvocationStreamSubject, payload, jetstream.WithMsgID(rec.RequestID)); err != nil {
slog.Error("bedrock: failed to publish invocation record", "request_id", rec.RequestID, "account", rec.AccountID, "err", err)
r.dropped.Add(1)
}
}
// LoggingConfig is the account's own invocation-logging preferences: where
// to deliver, and whether body text may be included at all.
type LoggingConfig struct {
S3BucketName string `json:"s3BucketName"`
S3KeyPrefix string `json:"s3KeyPrefix"`
TextDataDeliveryEnabled bool `json:"textDataDeliveryEnabled"`
ImageDataDeliveryEnabled bool `json:"imageDataDeliveryEnabled"`
}
// LoggingConfigReader resolves accountID's stored logging config. ok is
// false when the account has never configured one (delivery disabled).
type LoggingConfigReader interface {
Get(ctx context.Context, accountID string) (LoggingConfig, bool, error)
}
type noopLoggingConfigReader struct{}
func (noopLoggingConfigReader) Get(context.Context, string) (LoggingConfig, bool, error) {
return LoggingConfig{}, false, nil
}
// bedrockLoggingConfigBucket is the cluster-replicated KV bucket holding
// per-account invocation-logging configuration.
const bedrockLoggingConfigBucket = "bedrock-logging-config"
const bedrockLoggingConfigHistory = 1
// LoggingConfigStore persists per-account LoggingConfig in the
// bedrock-logging-config JetStream KV bucket, and serves as the
// LoggingConfigReader StreamRecorder consults before including body text.
type LoggingConfigStore struct {
js jetstream.JetStream
replicas int
mu sync.Mutex
kv jetstream.KeyValue
}
var _ LoggingConfigReader = (*LoggingConfigStore)(nil)
// NewLoggingConfigStore constructs a LoggingConfigStore.
func NewLoggingConfigStore(js jetstream.JetStream, replicas int) *LoggingConfigStore {
return &LoggingConfigStore{js: js, replicas: replicas}
}
// bucket lazily opens (or creates) the cluster-replicated
// bedrock-logging-config KV bucket, caching the handle for subsequent calls.
func (s *LoggingConfigStore) bucket(ctx context.Context) (jetstream.KeyValue, error) {
if s.js == nil {
return nil, errors.New("bedrock: logging config store has no JetStream client configured")
}
s.mu.Lock()
defer s.mu.Unlock()
if s.kv != nil {
return s.kv, nil
}
kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, s.js, bedrockLoggingConfigBucket, bedrockLoggingConfigHistory, s.replicas)
if err != nil {
return nil, err
}
s.kv = kv
return kv, nil
}
// Get returns accountID's stored logging config, or (zero, false, nil) if
// the account has never configured one.
func (s *LoggingConfigStore) Get(ctx context.Context, accountID string) (LoggingConfig, bool, error) {
kv, err := s.bucket(ctx)
if err != nil {
return LoggingConfig{}, false, err
}
entry, err := kv.Get(ctx, accountID)
if err != nil {
if errors.Is(err, jetstream.ErrKeyNotFound) {
return LoggingConfig{}, false, nil
}
return LoggingConfig{}, false, fmt.Errorf("kv get logging config for %s: %w", accountID, err)
}
var cfg LoggingConfig
if err := json.Unmarshal(entry.Value(), &cfg); err != nil {
return LoggingConfig{}, false, fmt.Errorf("decode logging config for %s: %w", accountID, err)
}
return cfg, true, nil
}
// Put stores accountID's logging config, overwriting any existing one.
func (s *LoggingConfigStore) Put(ctx context.Context, accountID string, cfg LoggingConfig) error {
kv, err := s.bucket(ctx)
if err != nil {
return err
}
data, err := json.Marshal(cfg)
if err != nil {
return fmt.Errorf("encode logging config for %s: %w", accountID, err)
}
if _, err := kv.Put(ctx, accountID, data); err != nil {
return fmt.Errorf("kv put logging config for %s: %w", accountID, err)
}
return nil
}
// Delete removes accountID's logging config, treating an already-absent
// config as success.
func (s *LoggingConfigStore) Delete(ctx context.Context, accountID string) error {
kv, err := s.bucket(ctx)
if err != nil {
return err
}
if err := kv.Delete(ctx, accountID); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
return fmt.Errorf("kv delete logging config for %s: %w", accountID, err)
}
return nil
}
// loggingConfigFromSDK validates and translates a *bedrock.LoggingConfig
// into the package's internal LoggingConfig. CloudWatchConfig is rejected:
// this platform's only delivery target is the account's own S3 bucket —
// body text is never routed to a platform-shared sink.
func loggingConfigFromSDK(in *bedrock.LoggingConfig) (LoggingConfig, error) {
if in == nil {
return LoggingConfig{}, errors.New(awserrors.ErrorValidationException)
}
if in.CloudWatchConfig != nil {
return LoggingConfig{}, errors.New(awserrors.ErrorValidationException)
}
textEnabled := aws.BoolValue(in.TextDataDeliveryEnabled)
imageEnabled := aws.BoolValue(in.ImageDataDeliveryEnabled)
if (textEnabled || imageEnabled) && (in.S3Config == nil || aws.StringValue(in.S3Config.BucketName) == "") {
return LoggingConfig{}, errors.New(awserrors.ErrorValidationException)
}
cfg := LoggingConfig{
TextDataDeliveryEnabled: textEnabled,
ImageDataDeliveryEnabled: imageEnabled,
}
if in.S3Config != nil {
cfg.S3BucketName = aws.StringValue(in.S3Config.BucketName)
cfg.S3KeyPrefix = aws.StringValue(in.S3Config.KeyPrefix)
}
return cfg, nil
}
// loggingConfigToSDK translates the package's internal LoggingConfig back to
// the AWS-parity *bedrock.LoggingConfig shape.
func loggingConfigToSDK(cfg LoggingConfig) *bedrock.LoggingConfig {
out := &bedrock.LoggingConfig{
TextDataDeliveryEnabled: aws.Bool(cfg.TextDataDeliveryEnabled),
ImageDataDeliveryEnabled: aws.Bool(cfg.ImageDataDeliveryEnabled),
}
if cfg.S3BucketName != "" {
out.S3Config = &bedrock.S3Config{
BucketName: aws.String(cfg.S3BucketName),
KeyPrefix: aws.String(cfg.S3KeyPrefix),
}
}
return out
}
// PutModelInvocationLoggingConfiguration stores accountID's invocation
// logging preferences. CloudWatch delivery is not supported (body text is
// delivered only to the account's own S3 bucket, never a platform-shared
// sink); a request naming it is rejected.
func PutModelInvocationLoggingConfiguration(ctx context.Context, accountID string, store *LoggingConfigStore, input *bedrock.PutModelInvocationLoggingConfigurationInput) (*bedrock.PutModelInvocationLoggingConfigurationOutput, error) {
cfg, err := loggingConfigFromSDK(input.LoggingConfig)
if err != nil {
return nil, err
}
if err := store.Put(ctx, accountID, cfg); err != nil {
return nil, err
}
return &bedrock.PutModelInvocationLoggingConfigurationOutput{}, nil
}
// GetModelInvocationLoggingConfiguration returns accountID's stored logging
// config, or an empty LoggingConfig when none has been set.
func GetModelInvocationLoggingConfiguration(ctx context.Context, accountID string, store *LoggingConfigStore, _ *bedrock.GetModelInvocationLoggingConfigurationInput) (*bedrock.GetModelInvocationLoggingConfigurationOutput, error) {
cfg, ok, err := store.Get(ctx, accountID)
if err != nil {
return nil, err
}
if !ok {
return &bedrock.GetModelInvocationLoggingConfigurationOutput{}, nil
}
return &bedrock.GetModelInvocationLoggingConfigurationOutput{LoggingConfig: loggingConfigToSDK(cfg)}, nil
}
// DeleteModelInvocationLoggingConfiguration removes accountID's stored
// logging config.
func DeleteModelInvocationLoggingConfiguration(ctx context.Context, accountID string, store *LoggingConfigStore, _ *bedrock.DeleteModelInvocationLoggingConfigurationInput) (*bedrock.DeleteModelInvocationLoggingConfigurationOutput, error) {
if err := store.Delete(ctx, accountID); err != nil {
return nil, err
}
return &bedrock.DeleteModelInvocationLoggingConfigurationOutput{}, nil
}
// DeliveryConsumer drains InvocationDeliveryConsumer, writing each record's
// body (if present) to the account's own S3 bucket and always emitting one
// structured slog line carrying metadata only — the platform's log-shipping
// pipeline (filebeat/Alloy -> Elasticsearch) is the "ES delivery" path,
// there being no bespoke Elasticsearch client in this codebase. The slog
// call reads named metadata fields only and never references
// rec.InputText/rec.OutputText, so the no-body-reaches-ES invariant holds
// structurally even if StreamRecorder's own gating were ever bypassed.
type DeliveryConsumer struct {
store objectstore.ObjectStore
configs LoggingConfigReader
}
// NewDeliveryConsumer constructs a DeliveryConsumer. store is the
// platform-trusted object store client (the same one the ECR registry
// writes through) — the account's own bucket name comes from its logging
// config, not a per-account credential, mirroring how Bedrock's own service
// role writes into a customer bucket, simplified for this platform's trust
// model.
func NewDeliveryConsumer(store objectstore.ObjectStore, configs LoggingConfigReader) *DeliveryConsumer {
return &DeliveryConsumer{store: store, configs: configs}
}
// Run pulls records from consumer until ctx is cancelled, acking each after
// successful delivery and nak'ing on failure so JetStream redelivers it.
func (c *DeliveryConsumer) Run(ctx context.Context, consumer jetstream.Consumer) {
iter, err := consumer.Messages()
if err != nil {
slog.Error("bedrock: failed to start invocation delivery consumer", "err", err)
return
}
defer iter.Stop()
// iter.Next() blocks indefinitely when no message is pending, so ctx
// cancellation alone would never unblock the loop below — Stop() is what
// actually wakes it, letting Next() return ErrMsgIteratorClosed.
stopped := make(chan struct{})
defer close(stopped)
go func() {
select {
case <-ctx.Done():
iter.Stop()
case <-stopped:
}
}()
for {
msg, err := iter.Next()
if err != nil {
if errors.Is(err, jetstream.ErrMsgIteratorClosed) || ctx.Err() != nil {
return
}
slog.Error("bedrock: invocation delivery consumer fetch failed", "err", err)
continue
}
c.deliver(ctx, msg)
}
}
// deliver decodes, writes, and acks (or naks) one message. A decode failure
// acks and drops the message outright: it can never succeed on redelivery.
func (c *DeliveryConsumer) deliver(ctx context.Context, msg jetstream.Msg) {
var rec InvocationRecord
if err := json.Unmarshal(msg.Data(), &rec); err != nil {
slog.Error("bedrock: invocation delivery consumer: undecodable record, dropping", "err", err)
if ackErr := msg.Ack(); ackErr != nil {
slog.Error("bedrock: invocation delivery consumer: ack failed", "err", ackErr)
}
return
}
// deliverS3 is unconditional: any account with a resolvable S3 bucket
// gets its invocation records, body or not. The delivery *flags*
// (TextDataDeliveryEnabled et al) already gated InputText/OutputText
// back in Record — this is not a second gate on top of that one.
if err := c.deliverS3(ctx, rec); err != nil {
slog.Error("bedrock: invocation delivery consumer: S3 delivery failed", "request_id", rec.RequestID, "account", rec.AccountID, "err", err)
if nakErr := msg.Nak(); nakErr != nil {
slog.Error("bedrock: invocation delivery consumer: nak failed", "err", nakErr)
}
return
}
// Metadata only: never rec.InputText/rec.OutputText.
slog.Info("bedrock invocation",
"request_id", rec.RequestID,
"account_id", rec.AccountID,
"model_id", rec.ModelID,
"operation", rec.Operation,
"backend", rec.Backend,
"timestamp", rec.Timestamp,
"latency_ms", rec.LatencyMs,
"http_status", rec.HTTPStatus,
"error_code", rec.ErrorCode,
"input_tokens", rec.InputTokens,
"output_tokens", rec.OutputTokens,
"usage_estimated", rec.UsageEstimated,
"partial", rec.Partial,
)
if err := msg.Ack(); err != nil {
slog.Error("bedrock: invocation delivery consumer: ack failed", "request_id", rec.RequestID, "err", err)
}
}
// deliverS3 writes rec's body to the account's own bucket, or does nothing
// if the account has since disabled/removed delivery.
func (c *DeliveryConsumer) deliverS3(ctx context.Context, rec InvocationRecord) error {
cfg, ok, err := c.configs.Get(ctx, rec.AccountID)
if err != nil {
return fmt.Errorf("resolve logging config: %w", err)
}
if !ok || cfg.S3BucketName == "" {
return nil
}
if err := c.store.EnsureBucket(ctx, cfg.S3BucketName); err != nil {
return fmt.Errorf("ensure bucket %s: %w", cfg.S3BucketName, err)
}
body, err := json.Marshal(rec)
if err != nil {
return fmt.Errorf("marshal record: %w", err)
}
key := s3ObjectKey(cfg.S3KeyPrefix, rec)
if _, err := c.store.PutObject(ctx, &s3.PutObjectInput{
Bucket: aws.String(cfg.S3BucketName),
Key: aws.String(key),
Body: bytes.NewReader(body),
}); err != nil {
return fmt.Errorf("put object %s/%s: %w", cfg.S3BucketName, key, err)
}
return nil
}
// s3ObjectKey lays records out as
// <prefix>/<modelId>/<year>/<month>/<day>/<requestId>.json, so an account's
// bucket stays browsable by model and date without a manifest.
func s3ObjectKey(prefix string, rec InvocationRecord) string {
var parts []string
if prefix != "" {
parts = append(parts, strings.Trim(prefix, "/"))
}
parts = append(parts,
rec.ModelID,
rec.Timestamp.Format("2006/01/02"),
rec.RequestID+".json",
)
return strings.Join(parts, "/")
}
+272
View File
@@ -0,0 +1,272 @@
package gateway_bedrock
import (
"bytes"
"context"
"encoding/json"
"io"
"log/slog"
"testing"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrock"
"github.com/aws/aws-sdk-go/service/s3"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/objectstore"
"github.com/mulgadc/spinifex/spinifex/testutil"
"github.com/nats-io/nats.go/jetstream"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// stubLoggingConfigReader returns a fixed (cfg, ok, err) for every Get call.
type stubLoggingConfigReader struct {
cfg LoggingConfig
ok bool
err error
}
func (s stubLoggingConfigReader) Get(context.Context, string) (LoggingConfig, bool, error) {
return s.cfg, s.ok, s.err
}
func TestLoggingConfigStore_PutGetDelete(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewLoggingConfigStore(js, 1)
ctx := context.Background()
_, ok, err := store.Get(ctx, "000000000001")
require.NoError(t, err)
assert.False(t, ok)
cfg := LoggingConfig{S3BucketName: "acct-bucket", S3KeyPrefix: "bedrock", TextDataDeliveryEnabled: true}
require.NoError(t, store.Put(ctx, "000000000001", cfg))
got, ok, err := store.Get(ctx, "000000000001")
require.NoError(t, err)
require.True(t, ok)
assert.Equal(t, cfg, got)
require.NoError(t, store.Delete(ctx, "000000000001"))
_, ok, err = store.Get(ctx, "000000000001")
require.NoError(t, err)
assert.False(t, ok)
// Deleting an already-absent config is idempotent, not an error.
require.NoError(t, store.Delete(ctx, "000000000001"))
}
func TestPutModelInvocationLoggingConfiguration_RejectsCloudWatch(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewLoggingConfigStore(js, 1)
input := &bedrock.PutModelInvocationLoggingConfigurationInput{
LoggingConfig: &bedrock.LoggingConfig{
CloudWatchConfig: &bedrock.CloudWatchConfig{LogGroupName: aws.String("g"), RoleArn: aws.String("arn:aws:iam::000000000001:role/r")},
TextDataDeliveryEnabled: aws.Bool(true),
},
}
_, err := PutModelInvocationLoggingConfiguration(context.Background(), "000000000001", store, input)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorValidationException, err.Error())
}
func TestPutModelInvocationLoggingConfiguration_RequiresBucketWhenDeliveryEnabled(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewLoggingConfigStore(js, 1)
input := &bedrock.PutModelInvocationLoggingConfigurationInput{
LoggingConfig: &bedrock.LoggingConfig{TextDataDeliveryEnabled: aws.Bool(true)},
}
_, err := PutModelInvocationLoggingConfiguration(context.Background(), "000000000001", store, input)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorValidationException, err.Error())
}
// TestModelInvocationLoggingConfiguration_RoundTripsViaSDKStructs drives the
// three control-plane handlers with real aws-sdk-go v1 bedrock structs, the
// same shapes a genuine SDK client marshals, rather than this package's own
// internal LoggingConfig.
func TestModelInvocationLoggingConfiguration_RoundTripsViaSDKStructs(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
store := NewLoggingConfigStore(js, 1)
ctx := context.Background()
putInput := &bedrock.PutModelInvocationLoggingConfigurationInput{
LoggingConfig: &bedrock.LoggingConfig{
S3Config: &bedrock.S3Config{BucketName: aws.String("acct-bucket"), KeyPrefix: aws.String("bedrock")},
TextDataDeliveryEnabled: aws.Bool(true),
ImageDataDeliveryEnabled: aws.Bool(false),
},
}
_, err := PutModelInvocationLoggingConfiguration(ctx, "000000000001", store, putInput)
require.NoError(t, err)
getOut, err := GetModelInvocationLoggingConfiguration(ctx, "000000000001", store, new(bedrock.GetModelInvocationLoggingConfigurationInput))
require.NoError(t, err)
require.NotNil(t, getOut.LoggingConfig)
require.NotNil(t, getOut.LoggingConfig.S3Config)
assert.Equal(t, "acct-bucket", aws.StringValue(getOut.LoggingConfig.S3Config.BucketName))
assert.Equal(t, "bedrock", aws.StringValue(getOut.LoggingConfig.S3Config.KeyPrefix))
assert.True(t, aws.BoolValue(getOut.LoggingConfig.TextDataDeliveryEnabled))
assert.False(t, aws.BoolValue(getOut.LoggingConfig.ImageDataDeliveryEnabled))
_, err = DeleteModelInvocationLoggingConfiguration(ctx, "000000000001", store, new(bedrock.DeleteModelInvocationLoggingConfigurationInput))
require.NoError(t, err)
getOut, err = GetModelInvocationLoggingConfiguration(ctx, "000000000001", store, new(bedrock.GetModelInvocationLoggingConfigurationInput))
require.NoError(t, err)
assert.Nil(t, getOut.LoggingConfig)
}
// TestStreamRecorder_DropsAndCountsWhenStreamMissing proves the "publish
// failure must never fail the invocation" contract: Record has no error
// return at all, so the only way to observe a failed publish is Dropped().
func TestStreamRecorder_DropsAndCountsWhenStreamMissing(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
// Deliberately skip EnsureInvocationStream: publishing to a subject no
// stream captures must fail cleanly (no responders), not panic or block.
recorder := NewStreamRecorder(js, nil)
recorder.Record(context.Background(), InvocationRecord{RequestID: "req-1", AccountID: "000000000001"})
assert.Equal(t, int64(1), recorder.Dropped())
}
// TestStreamRecorder_Record_BodyOnlyWhenLoggingEnabled is the D11 gate at the
// publisher: body text only survives onto the stream when the account's own
// logging config enables delivery.
func TestStreamRecorder_Record_BodyOnlyWhenLoggingEnabled(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
ctx := context.Background()
_, err := EnsureInvocationStream(ctx, js, 1)
require.NoError(t, err)
consumer, err := EnsureDeliveryConsumer(ctx, js)
require.NoError(t, err)
enabled := NewStreamRecorder(js, stubLoggingConfigReader{cfg: LoggingConfig{TextDataDeliveryEnabled: true}, ok: true})
enabled.Record(ctx, InvocationRecord{RequestID: "req-enabled", AccountID: "acct-a", InputText: "prompt", OutputText: "completion"})
disabled := NewStreamRecorder(js, stubLoggingConfigReader{})
disabled.Record(ctx, InvocationRecord{RequestID: "req-disabled", AccountID: "acct-b", InputText: "prompt", OutputText: "completion"})
batch, err := consumer.Fetch(2, jetstream.FetchMaxWait(2*time.Second))
require.NoError(t, err)
byID := map[string]InvocationRecord{}
for msg := range batch.Messages() {
var r InvocationRecord
require.NoError(t, json.Unmarshal(msg.Data(), &r))
byID[r.RequestID] = r
require.NoError(t, msg.Ack())
}
require.NoError(t, batch.Error())
require.Len(t, byID, 2)
assert.Equal(t, "prompt", byID["req-enabled"].InputText)
assert.Equal(t, "completion", byID["req-enabled"].OutputText)
assert.Empty(t, byID["req-disabled"].InputText)
assert.Empty(t, byID["req-disabled"].OutputText)
}
// TestDeliveryConsumer_Run_WritesS3AndNeverLogsBodyText is the D11 invariant
// at the delivery side: the account's own body text reaches only its S3
// bucket; the metadata log line — this package's stand-in for "ES delivery"
// (see DeliveryConsumer's doc comment) — must never contain it, even when a
// record legitimately carries a populated body.
func TestDeliveryConsumer_Run_WritesS3AndNeverLogsBodyText(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
_, err := EnsureInvocationStream(ctx, js, 1)
require.NoError(t, err)
consumer, err := EnsureDeliveryConsumer(ctx, js)
require.NoError(t, err)
store := objectstore.NewMemoryObjectStore()
configs := stubLoggingConfigReader{cfg: LoggingConfig{S3BucketName: "acct-bucket", TextDataDeliveryEnabled: true}, ok: true}
const secretPrompt = "the secret prompt text"
const secretCompletion = "the secret completion text"
recorder := NewStreamRecorder(js, configs)
recorder.Record(ctx, InvocationRecord{
RequestID: "req-body-test", AccountID: "acct-a", ModelID: "meta.llama3-70b-instruct-v1:0",
Operation: OperationConverse, InputText: secretPrompt, OutputText: secretCompletion,
})
var logBuf bytes.Buffer
prevLogger := slog.Default()
slog.SetDefault(slog.New(slog.NewJSONHandler(&logBuf, nil)))
defer slog.SetDefault(prevLogger)
dc := NewDeliveryConsumer(store, configs)
done := make(chan struct{})
go func() {
dc.Run(ctx, consumer)
close(done)
}()
require.Eventually(t, func() bool { return store.Count() == 1 }, 2*time.Second, 10*time.Millisecond)
cancel()
<-done
logOutput := logBuf.String()
assert.NotContains(t, logOutput, secretPrompt)
assert.NotContains(t, logOutput, secretCompletion)
assert.Contains(t, logOutput, "req-body-test")
}
// TestDeliveryConsumer_Run_DeliversMetadataOnlyRecordToS3 covers the common
// case: an account configured a destination bucket but left body delivery
// disabled. It must still receive its invocation record in S3 (metadata
// only, InputText/OutputText empty) — S3 delivery is not itself gated on
// body presence, only the body fields are.
func TestDeliveryConsumer_Run_DeliversMetadataOnlyRecordToS3(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
_, err := EnsureInvocationStream(ctx, js, 1)
require.NoError(t, err)
consumer, err := EnsureDeliveryConsumer(ctx, js)
require.NoError(t, err)
store := objectstore.NewMemoryObjectStore()
// TextDataDeliveryEnabled is deliberately false: the account asked for a
// destination bucket but not body delivery.
configs := stubLoggingConfigReader{cfg: LoggingConfig{S3BucketName: "acct-bucket"}, ok: true}
recorder := NewStreamRecorder(js, configs)
recorder.Record(ctx, InvocationRecord{
RequestID: "req-metadata-only", AccountID: "acct-a", ModelID: "meta.llama3-70b-instruct-v1:0",
Operation: OperationConverse, InputText: "prompt", OutputText: "completion",
})
dc := NewDeliveryConsumer(store, configs)
done := make(chan struct{})
go func() {
dc.Run(ctx, consumer)
close(done)
}()
require.Eventually(t, func() bool { return store.Count() == 1 }, 2*time.Second, 10*time.Millisecond)
cancel()
<-done
listing, err := store.ListObjectsV2(context.Background(), &s3.ListObjectsV2Input{Bucket: aws.String("acct-bucket")})
require.NoError(t, err)
require.Len(t, listing.Contents, 1, "expected the record to land in the account's bucket despite body delivery being disabled")
obj, err := store.GetObject(context.Background(), &s3.GetObjectInput{Bucket: aws.String("acct-bucket"), Key: listing.Contents[0].Key})
require.NoError(t, err)
body, err := io.ReadAll(obj.Body)
require.NoError(t, err)
var stored InvocationRecord
require.NoError(t, json.Unmarshal(body, &stored))
assert.Equal(t, "req-metadata-only", stored.RequestID)
assert.Empty(t, stored.InputText)
assert.Empty(t, stored.OutputText)
}
+73 -17
View File
@@ -2,11 +2,15 @@ package gateway_bedrock
import (
"context"
"encoding/json"
"errors"
"log/slog"
"strings"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
"github.com/google/uuid"
"github.com/mulgadc/spinifex/spinifex/awserrors"
)
@@ -29,29 +33,75 @@ type Provider interface {
type Router struct {
resolver CredentialResolver
endpointResolver EndpointResolver
recorder Recorder
}
// NewRouter constructs a Router. A nil resolver or endpointResolver falls
// back to a resolver/resolver that finds nothing, so a Router is always safe
// NewRouter constructs a Router. A nil resolver, endpointResolver, or
// recorder falls back to a no-op implementation, so a Router is always safe
// to use even before the real stores are wired in.
func NewRouter(resolver CredentialResolver, endpointResolver EndpointResolver) *Router {
func NewRouter(resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder) *Router {
if resolver == nil {
resolver = NoopCredentialResolver
}
if endpointResolver == nil {
endpointResolver = NewStaticEndpointResolver(nil)
}
return &Router{resolver: resolver, endpointResolver: endpointResolver}
if recorder == nil {
recorder = NoopRecorder
}
return &Router{resolver: resolver, endpointResolver: endpointResolver, recorder: recorder}
}
// Converse routes modelID to its provider via the catalog. Unknown modelIds
// and unresolvable vendors return ResourceNotFoundException; a vendor with no
// resolvable credential returns AccessDeniedException.
func (rt *Router) Converse(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseInput) (*bedrockruntime.ConverseOutput, error) {
// resolvable credential returns AccessDeniedException. Every exit records an
// InvocationRecord via the deferred closure, matching pumpConverseStream's
// treatment of the streaming path.
func (rt *Router) Converse(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseInput) (out *bedrockruntime.ConverseOutput, err error) {
requestID := uuid.NewString()
start := time.Now()
var backend string
defer func() {
httpStatus, code := recordOutcome(err)
var inputTokens, outputTokens int64
if out != nil && out.Usage != nil {
inputTokens = aws.Int64Value(out.Usage.InputTokens)
outputTokens = aws.Int64Value(out.Usage.OutputTokens)
}
inputText, ierr := json.Marshal(input)
if ierr != nil {
slog.Error("bedrock converse: failed to marshal input for recording", "model", modelID, "err", ierr)
}
var outputText []byte
if out != nil {
var oerr error
outputText, oerr = json.Marshal(out)
if oerr != nil {
slog.Error("bedrock converse: failed to marshal output for recording", "model", modelID, "err", oerr)
}
}
rt.recorder.Record(ctx, InvocationRecord{
RequestID: requestID,
AccountID: accountID,
ModelID: modelID,
Operation: OperationConverse,
Backend: backend,
LatencyMs: time.Since(start).Milliseconds(),
HTTPStatus: httpStatus,
ErrorCode: code,
InputTokens: inputTokens,
OutputTokens: outputTokens,
InputText: string(inputText),
OutputText: string(outputText),
})
}()
entry, ok := lookupCatalogEntry(modelID)
if !ok {
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
err = errors.New(awserrors.ErrorResourceNotFoundException)
return nil, err
}
backend = entry.Provider
var p Provider
switch {
@@ -60,29 +110,35 @@ func (rt *Router) Converse(ctx context.Context, accountID, modelID string, input
case strings.HasPrefix(entry.Provider, providerPrefix):
switch strings.TrimPrefix(entry.Provider, providerPrefix) {
case vendorAnthropic:
key, ok, err := rt.resolver.Resolve(ctx, accountID, vendorAnthropic)
var key string
var resolvable bool
key, resolvable, err = rt.resolver.Resolve(ctx, accountID, vendorAnthropic)
if err != nil {
return nil, err
}
if !ok {
return nil, errors.New(awserrors.ErrorAccessDeniedException)
if !resolvable {
err = errors.New(awserrors.ErrorAccessDeniedException)
return nil, err
}
p = newAnthropicProvider(key)
default:
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
err = errors.New(awserrors.ErrorResourceNotFoundException)
return nil, err
}
default:
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
err = errors.New(awserrors.ErrorResourceNotFoundException)
return nil, err
}
return p.Converse(ctx, modelID, input)
out, err = p.Converse(ctx, modelID, input)
return out, err
}
// Converse is the bedrock-runtime Converse entry point used by the gateway
// route table. resolver and endpointResolver may be nil; NewRouter supplies
// no-op fallbacks.
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)
// route table. resolver, endpointResolver, and recorder may be nil; NewRouter
// supplies no-op fallbacks.
func Converse(ctx context.Context, accountID, modelID string, input *bedrockruntime.ConverseInput, resolver CredentialResolver, endpointResolver EndpointResolver, recorder Recorder) (*bedrockruntime.ConverseOutput, error) {
return NewRouter(resolver, endpointResolver, recorder).Converse(ctx, accountID, modelID, input)
}
// ConverseStreamProvider is the optional streaming capability a Provider may
+10 -10
View File
@@ -46,7 +46,7 @@ func TestRouter_Converse_SelfHostSuccess(t *testing.T) {
defer ts.Close()
modelID := "meta.llama3-2-1b-instruct-v1:0"
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil)
out, err := rt.Converse(context.Background(), "000000000001", modelID, converseInput())
require.NoError(t, err)
@@ -55,14 +55,14 @@ func TestRouter_Converse_SelfHostSuccess(t *testing.T) {
}
func TestRouter_Converse_UnknownModelReturnsResourceNotFound(t *testing.T) {
rt := NewRouter(nil, nil)
rt := NewRouter(nil, nil, nil)
_, err := rt.Converse(context.Background(), "000000000001", "does.not-exist-v1:0", converseInput())
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
}
func TestRouter_Converse_AnthropicNoCredentialReturnsAccessDenied(t *testing.T) {
rt := NewRouter(stubCredentialResolver{ok: false}, nil)
rt := NewRouter(stubCredentialResolver{ok: false}, nil, nil)
_, err := rt.Converse(context.Background(), "000000000001", "anthropic.claude-3-5-sonnet-20240620-v1:0", converseInput())
require.Error(t, err)
assert.Equal(t, awserrors.ErrorAccessDeniedException, err.Error())
@@ -80,14 +80,14 @@ func TestConverse_PackageEntryPoint_SelfHostSuccess(t *testing.T) {
defer ts.Close()
modelID := "meta.llama3-2-1b-instruct-v1:0"
out, err := Converse(context.Background(), "000000000001", modelID, converseInput(), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
out, err := Converse(context.Background(), "000000000001", modelID, converseInput(), nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil)
require.NoError(t, err)
require.NotNil(t, out.Output.Message)
assert.Equal(t, "via package Converse", *out.Output.Message.Content[0].Text)
}
func TestConverse_PackageEntryPoint_UnknownModel(t *testing.T) {
_, err := Converse(context.Background(), "000000000001", "does.not-exist-v1:0", converseInput(), nil, nil)
_, err := Converse(context.Background(), "000000000001", "does.not-exist-v1:0", converseInput(), nil, nil, nil)
require.Error(t, err)
assert.Equal(t, awserrors.ErrorResourceNotFoundException, err.Error())
}
@@ -109,7 +109,7 @@ func TestRouter_ConverseStream_SelfHostSuccess(t *testing.T) {
defer ts.Close()
modelID := "meta.llama3-2-1b-instruct-v1:0"
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}))
rt := NewRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil)
src, err := rt.ConverseStream(context.Background(), "000000000001", modelID, converseStreamInput())
require.NoError(t, err)
@@ -120,28 +120,28 @@ func TestRouter_ConverseStream_SelfHostSuccess(t *testing.T) {
}
func TestRouter_ConverseStream_UnknownModelReturnsResourceNotFound(t *testing.T) {
rt := NewRouter(nil, nil)
rt := NewRouter(nil, 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)
rt := NewRouter(stubCredentialResolver{ok: false}, nil, 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)
rt := NewRouter(nil, nil, nil)
_, err := rt.ConverseStream(context.Background(), "000000000001", "meta.llama3-2-1b-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)
rt := NewRouter(nil, nil, nil)
require.NotNil(t, rt)
// Self-host model with no endpoint resolver configured resolves nothing,
+26 -1
View File
@@ -269,8 +269,13 @@ func (p *vllmProvider) ConverseStream(ctx context.Context, modelID string, input
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
// reqCtx/cancel let Close() abort the upstream request outright rather
// than merely stopping our own reads — a disconnected client must free
// the vLLM-side generation, not just this goroutine.
reqCtx, cancel := context.WithCancel(ctx)
httpReq, err := http.NewRequestWithContext(reqCtx, http.MethodPost, baseURL+vllmChatCompletionsPath, bytes.NewReader(reqBody)) //nolint:gosec // G704: baseURL is a resolved pinned self-host endpoint, not user input
if err != nil {
cancel()
slog.Error("vllm stream: failed to build request", "model", modelID, "err", err)
return nil, errors.New(awserrors.ErrorInternalError)
}
@@ -279,6 +284,7 @@ func (p *vllmProvider) ConverseStream(ctx context.Context, modelID string, input
resp, err := p.httpClient.Do(httpReq) //nolint:gosec // G704: httpReq targets the resolved pinned self-host endpoint, not user input
if err != nil {
cancel()
slog.Error("vllm stream: request failed", "model", modelID, "endpoint", baseURL, "err", err)
return nil, errors.New(awserrors.ErrorServiceUnavailableException)
}
@@ -286,6 +292,7 @@ func (p *vllmProvider) ConverseStream(ctx context.Context, modelID string, input
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
respBody, _ := io.ReadAll(resp.Body)
_ = resp.Body.Close()
cancel()
slog.Error("vllm stream: upstream error", "model", modelID, "status", resp.StatusCode, "body", string(respBody))
return nil, errors.New(mapUpstreamStatus(resp.StatusCode))
}
@@ -294,6 +301,7 @@ func (p *vllmProvider) ConverseStream(ctx context.Context, modelID string, input
resp: resp,
scanner: newSSEScanner(resp.Body),
start: time.Now(),
cancel: cancel,
}, nil
}
@@ -308,6 +316,7 @@ type vllmConverseStreamSource struct {
resp *http.Response
scanner *sseScanner
start time.Time
cancel context.CancelFunc
queue []ConverseStreamEvent
started bool
@@ -315,14 +324,29 @@ type vllmConverseStreamSource struct {
metadataEmitted bool
inputTokens int64
outputTokens int64
deltaChunks int64
}
var _ converseStreamSource = (*vllmConverseStreamSource)(nil)
var _ usageReporter = (*vllmConverseStreamSource)(nil)
// Close aborts the upstream request via cancel (not just closing the read
// side), so a disconnected client actually frees the vLLM-side generation.
func (s *vllmConverseStreamSource) Close() error {
s.cancel()
return s.resp.Body.Close()
}
// usage returns real usage once the trailing usage-only chunk has been
// consumed (metadataEmitted); until then it falls back to a content-delta
// chunk count as an estimated output-token proxy, flagged accordingly.
func (s *vllmConverseStreamSource) usage() (inputTokens, outputTokens int64, estimated bool) {
if s.metadataEmitted {
return s.inputTokens, s.outputTokens, false
}
return 0, s.deltaChunks, true
}
func (s *vllmConverseStreamSource) Next(_ context.Context) (ConverseStreamEvent, bool, error) {
for len(s.queue) == 0 {
ev, ok, err := s.scanner.Next()
@@ -381,6 +405,7 @@ func (s *vllmConverseStreamSource) consume(chunk vllmStreamChunk) {
s.ensureStarted()
if choice.Delta.Content != "" {
s.deltaChunks++
s.queue = append(s.queue, ConverseStreamEvent{
Kind: converseStreamEventContentBlockDelta,
ContentBlockDelta: &bedrockruntime.ContentBlockDeltaEvent{
+81
View File
@@ -6,6 +6,7 @@ import (
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/bedrockruntime"
@@ -271,6 +272,86 @@ func TestVLLMProvider_ConverseStream_EmptyStreamIsWellFormed(t *testing.T) {
}, kinds)
}
// TestVLLMConverseStreamSource_EstimatesOutputTokensBeforeUsageChunk proves
// vLLM's estimation flag: while the trailing usage-only chunk has not yet
// been consumed (e.g. the stream was cut short by a fault or disconnect),
// usage() falls back to a content-delta chunk count and reports estimated.
func TestVLLMConverseStreamSource_EstimatesOutputTokensBeforeUsageChunk(t *testing.T) {
const sse = "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n" +
"data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\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))
w.(http.Flusher).Flush()
// Deliberately never sends the trailing usage-only chunk or EOF,
// simulating a stream cut short before real usage is known.
<-r.Context().Done()
}))
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() }()
// Drain exactly the 4 events both chunks queue (messageStart,
// contentBlockStart, and two content deltas) without reaching EOF.
for range 4 {
_, ok, nerr := src.Next(context.Background())
require.NoError(t, nerr)
require.True(t, ok)
}
ur, ok := src.(usageReporter)
require.True(t, ok)
inputTokens, outputTokens, estimated := ur.usage()
assert.True(t, estimated)
assert.Equal(t, int64(0), inputTokens)
assert.Equal(t, int64(2), outputTokens)
}
// TestVLLMConverseStreamSource_CloseAbortsUpstreamRequest proves Close()
// truly aborts the upstream request rather than merely stopping local reads:
// the handler blocks on r.Context().Done() and never sends EOF itself, so it
// can only unblock via the cancel() invoked from Close().
func TestVLLMConverseStreamSource_CloseAbortsUpstreamRequest(t *testing.T) {
serverCtxDone := make(chan struct{})
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\":{\"content\":\"hi\"}}]}\n\n"))
w.(http.Flusher).Flush()
<-r.Context().Done()
close(serverCtxDone)
}))
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)
// Drain one event so the handler's write/flush has definitely happened.
_, ok, err := src.Next(context.Background())
require.NoError(t, err)
require.True(t, ok)
require.NoError(t, src.Close())
select {
case <-serverCtxDone:
// aborted, as expected
case <-time.After(2 * time.Second):
t.Fatal("server request context was never canceled: Close() did not abort the upstream request")
}
}
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")
+9 -8
View File
@@ -26,22 +26,23 @@ type bedrockRuntimeRoute struct {
// bedrockRuntimeRouteHandler invokes a per-action bedrock-runtime (data-plane)
// gateway function. params holds the regex capture groups, PathUnescape'd.
// resolver is gw.bedrockResolver() (credential store or no-op); endpoints is
// gw.bedrockEndpointResolver() over the configured pinned self-host endpoints.
type bedrockRuntimeRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver) (any, error)
// gw.bedrockEndpointResolver() over the configured pinned self-host
// endpoints; recorder is gw.bedrockRecorder() (invocation recorder or no-op).
type bedrockRuntimeRouteHandler func(ctx context.Context, accountID string, params []string, body []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver, recorder gateway_bedrock.Recorder) (any, error)
// bedrockRuntimeRoutes is the dispatch table. InvokeModel has no handler
// function here: BedrockRuntime_Request special-cases its action to bypass
// the JSON-marshaling dispatch below, since its response is raw bytes.
var bedrockRuntimeRoutes = []bedrockRuntimeRoute{
{"POST", regexp.MustCompile(`^/model/([^/]+)/converse$`), "Converse",
func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver) (any, error) {
func(ctx context.Context, acct string, p []string, b []byte, resolver gateway_bedrock.CredentialResolver, endpoints gateway_bedrock.EndpointResolver, recorder gateway_bedrock.Recorder) (any, error) {
input := new(bedrockruntime.ConverseInput)
if len(b) > 0 {
if err := json.Unmarshal(b, input); err != nil {
return nil, errors.New(awserrors.ErrorValidationException)
}
}
return gateway_bedrock.Converse(ctx, acct, p[0], input, resolver, endpoints)
return gateway_bedrock.Converse(ctx, acct, p[0], input, resolver, endpoints, recorder)
}},
{"POST", regexp.MustCompile(`^/model/([^/]+)/invoke$`), "InvokeModel", nil},
{"POST", regexp.MustCompile(`^/model/([^/]+)/converse-stream$`), "ConverseStream", nil},
@@ -111,7 +112,7 @@ func (gw *GatewayConfig) BedrockRuntime_Request(w http.ResponseWriter, r *http.R
// InvokeModel returns provider-native bytes, not a struct WriteJSONResponse
// could marshal, so it writes its own response body directly.
if action == "InvokeModel" {
respBody, contentType, err := gateway_bedrock.InvokeModel(r.Context(), accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver())
respBody, contentType, err := gateway_bedrock.InvokeModel(r.Context(), accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder())
if err != nil {
return err
}
@@ -126,13 +127,13 @@ func (gw *GatewayConfig) BedrockRuntime_Request(w http.ResponseWriter, r *http.R
// (-> 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())
return gateway_bedrock.ConverseStream(r.Context(), w, accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder())
}
if action == "InvokeModelWithResponseStream" {
return gateway_bedrock.InvokeModelWithResponseStream(r.Context(), w, accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), r.Header.Get("Content-Type"))
return gateway_bedrock.InvokeModelWithResponseStream(r.Context(), w, accountID, params[0], body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), r.Header.Get("Content-Type"), gw.bedrockRecorder())
}
output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver(), gw.bedrockEndpointResolver())
output, err := handler(r.Context(), accountID, params, body, gw.bedrockResolver(), gw.bedrockEndpointResolver(), gw.bedrockRecorder())
if err != nil {
return err
}
+8
View File
@@ -124,6 +124,14 @@ type GatewayConfig struct {
// config; a later phase backs it with the daemon's dynamic endpoint
// registry. Nil/empty means no self-hosted models are reachable.
BedrockEndpoints map[string]string
// BedrockLoggingConfig persists per-account invocation-logging preferences
// (PutModelInvocationLoggingConfiguration and friends). Nil falls back to
// an unconfigured store, under which reads/writes error rather than panic.
BedrockLoggingConfig *gateway_bedrock.LoggingConfigStore
// BedrockRecorder durably records every Bedrock invocation. Nil falls back
// to a no-op recorder, so routes stay safe to exercise before the
// invocation stream is wired in (e.g. unit tests of unrelated routes).
BedrockRecorder gateway_bedrock.Recorder
}
var supportedServices = map[string]bool{
+39 -21
View File
@@ -245,13 +245,13 @@ func launchService(config *config.ClusterConfig) error {
// predastore from the gateway; repo/tag/manifest metadata and in-progress
// uploads are owned by the daemon and reached over NATS request/reply. The
// /v2 auth bridge resolves the per-request account from a verified token.
ecrStore := objectstore.NewS3ObjectStoreFromConfig(
objStore := objectstore.NewS3ObjectStoreFromConfig(
admin.DialTarget(nodeConfig.Predastore.Host),
nodeConfig.Predastore.Region,
nodeConfig.Predastore.AccessKey,
nodeConfig.Predastore.SecretKey,
)
ecrRegistry := gateway_ecr.NewRegistry(ecrStore, ecr.NewNATSMetaStore(natsConn), config.Bootstrap.AccountID)
ecrRegistry := gateway_ecr.NewRegistry(objStore, ecr.NewNATSMetaStore(natsConn), config.Bootstrap.AccountID)
// Lifecycle expiry sweep applies each repo's stored lifecycle policy and
// deletes the expired set via the registry GC path. It runs here (not the
@@ -313,26 +313,44 @@ func launchService(config *config.ClusterConfig) error {
// is a comma-separated list of modelId=baseURL pairs.
bedrockEndpoints := parseBedrockEndpoints(os.Getenv("OCHRE_VLLM_ENDPOINTS"))
// Bedrock invocation records: every Converse/InvokeModel call (streaming
// or not) is published to the invocation stream, then fanned out by
// deliveryConsumer to any account with a configured destination bucket
// and a metadata-only log line. bedrockLoggingConfig separately gates
// whether the record written to that bucket includes body text.
bedrockLoggingConfig := gateway_bedrock.NewLoggingConfigStore(js, len(config.Nodes))
if _, err := gateway_bedrock.EnsureInvocationStream(janitorCtx, js, len(config.Nodes)); err != nil {
return fmt.Errorf("bedrock: ensure invocation stream: %w", err)
}
deliveryConsumer, err := gateway_bedrock.EnsureDeliveryConsumer(janitorCtx, js)
if err != nil {
return fmt.Errorf("bedrock: ensure invocation delivery consumer: %w", err)
}
bedrockRecorder := gateway_bedrock.NewStreamRecorder(js, bedrockLoggingConfig)
go gateway_bedrock.NewDeliveryConsumer(objStore, bedrockLoggingConfig).Run(janitorCtx, deliveryConsumer)
gw := gateway.GatewayConfig{
Debug: nodeConfig.AWSGW.Debug,
DisableLogging: false,
NATSConn: natsConn,
Config: nodeConfig.AWSGW.Config,
ExpectedNodes: len(config.Nodes),
Region: nodeConfig.Region,
InternalSuffix: config.AWS.InternalSuffix,
RegistryPort: registryPort,
RegistryHost: registryHost,
AZ: nodeConfig.AZ,
IAMService: iamService,
STSService: stsService,
Version: version,
Commit: commit,
ECRRegistry: ecrRegistry,
ECRTokenIssuer: gateway_ecrauth.NewIssuer(signingKey, ecrAudience),
ECRTokenVerifier: gateway_ecrauth.NewVerifier(verifyKeys, ecrAudience),
BedrockCredentials: bedrockCredentials,
BedrockEndpoints: bedrockEndpoints,
Debug: nodeConfig.AWSGW.Debug,
DisableLogging: false,
NATSConn: natsConn,
Config: nodeConfig.AWSGW.Config,
ExpectedNodes: len(config.Nodes),
Region: nodeConfig.Region,
InternalSuffix: config.AWS.InternalSuffix,
RegistryPort: registryPort,
RegistryHost: registryHost,
AZ: nodeConfig.AZ,
IAMService: iamService,
STSService: stsService,
Version: version,
Commit: commit,
ECRRegistry: ecrRegistry,
ECRTokenIssuer: gateway_ecrauth.NewIssuer(signingKey, ecrAudience),
ECRTokenVerifier: gateway_ecrauth.NewVerifier(verifyKeys, ecrAudience),
BedrockCredentials: bedrockCredentials,
BedrockEndpoints: bedrockEndpoints,
BedrockLoggingConfig: bedrockLoggingConfig,
BedrockRecorder: bedrockRecorder,
}
// Rotate the ECR signing key on a 30-day cadence, retaining the previous keys