diff --git a/spinifex/gateway/bedrock.go b/spinifex/gateway/bedrock.go index 46ef4acbf..1c29d105c 100644 --- a/spinifex/gateway/bedrock.go +++ b/spinifex/gateway/bedrock.go @@ -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. diff --git a/spinifex/gateway/bedrock/anthropic.go b/spinifex/gateway/bedrock/anthropic.go index 9bfc59cb8..90d2e1d1e 100644 --- a/spinifex/gateway/bedrock/anthropic.go +++ b/spinifex/gateway/bedrock/anthropic.go @@ -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() diff --git a/spinifex/gateway/bedrock/contract_test.go b/spinifex/gateway/bedrock/contract_test.go index dc079de61..ed582ff7f 100644 --- a/spinifex/gateway/bedrock/contract_test.go +++ b/spinifex/gateway/bedrock/contract_test.go @@ -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) }) diff --git a/spinifex/gateway/bedrock/converse_stream.go b/spinifex/gateway/bedrock/converse_stream.go index d05096d25..81224761d 100644 --- a/spinifex/gateway/bedrock/converse_stream.go +++ b/spinifex/gateway/bedrock/converse_stream.go @@ -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 } diff --git a/spinifex/gateway/bedrock/converse_stream_test.go b/spinifex/gateway/bedrock/converse_stream_test.go index 0afb8d927..2a4a3a6ec 100644 --- a/spinifex/gateway/bedrock/converse_stream_test.go +++ b/spinifex/gateway/bedrock/converse_stream_test.go @@ -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) diff --git a/spinifex/gateway/bedrock/invoke.go b/spinifex/gateway/bedrock/invoke.go index 5b232ef7e..6ba425930 100644 --- a/spinifex/gateway/bedrock/invoke.go +++ b/spinifex/gateway/bedrock/invoke.go @@ -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 diff --git a/spinifex/gateway/bedrock/invoke_anthropic.go b/spinifex/gateway/bedrock/invoke_anthropic.go index c2f221a26..12a74d6b0 100644 --- a/spinifex/gateway/bedrock/invoke_anthropic.go +++ b/spinifex/gateway/bedrock/invoke_anthropic.go @@ -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 +} diff --git a/spinifex/gateway/bedrock/invoke_llama.go b/spinifex/gateway/bedrock/invoke_llama.go index bc3b84088..369c85d69 100644 --- a/spinifex/gateway/bedrock/invoke_llama.go +++ b/spinifex/gateway/bedrock/invoke_llama.go @@ -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)) diff --git a/spinifex/gateway/bedrock/invoke_stream.go b/spinifex/gateway/bedrock/invoke_stream.go index f3a13dc19..84c44dc38 100644 --- a/spinifex/gateway/bedrock/invoke_stream.go +++ b/spinifex/gateway/bedrock/invoke_stream.go @@ -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 } diff --git a/spinifex/gateway/bedrock/invoke_stream_test.go b/spinifex/gateway/bedrock/invoke_stream_test.go index 1e23094f8..c1052157e 100644 --- a/spinifex/gateway/bedrock/invoke_stream_test.go +++ b/spinifex/gateway/bedrock/invoke_stream_test.go @@ -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) diff --git a/spinifex/gateway/bedrock/invoke_test.go b/spinifex/gateway/bedrock/invoke_test.go index 756c3e8d4..ae19a50a6 100644 --- a/spinifex/gateway/bedrock/invoke_test.go +++ b/spinifex/gateway/bedrock/invoke_test.go @@ -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()) } diff --git a/spinifex/gateway/bedrock/logging.go b/spinifex/gateway/bedrock/logging.go new file mode 100644 index 000000000..cedbe3160 --- /dev/null +++ b/spinifex/gateway/bedrock/logging.go @@ -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 +// /////.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, "/") +} diff --git a/spinifex/gateway/bedrock/logging_test.go b/spinifex/gateway/bedrock/logging_test.go new file mode 100644 index 000000000..19736ae56 --- /dev/null +++ b/spinifex/gateway/bedrock/logging_test.go @@ -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) +} diff --git a/spinifex/gateway/bedrock/provider.go b/spinifex/gateway/bedrock/provider.go index ab98a667e..07575959f 100644 --- a/spinifex/gateway/bedrock/provider.go +++ b/spinifex/gateway/bedrock/provider.go @@ -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 diff --git a/spinifex/gateway/bedrock/router_test.go b/spinifex/gateway/bedrock/router_test.go index 9d6343033..f1009ab4b 100644 --- a/spinifex/gateway/bedrock/router_test.go +++ b/spinifex/gateway/bedrock/router_test.go @@ -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, diff --git a/spinifex/gateway/bedrock/vllm.go b/spinifex/gateway/bedrock/vllm.go index 3d1a1fa56..97f9739bc 100644 --- a/spinifex/gateway/bedrock/vllm.go +++ b/spinifex/gateway/bedrock/vllm.go @@ -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{ diff --git a/spinifex/gateway/bedrock/vllm_test.go b/spinifex/gateway/bedrock/vllm_test.go index 22bba34ed..49dd4d310 100644 --- a/spinifex/gateway/bedrock/vllm_test.go +++ b/spinifex/gateway/bedrock/vllm_test.go @@ -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") diff --git a/spinifex/gateway/bedrockruntime.go b/spinifex/gateway/bedrockruntime.go index 96641fc43..da934a1c9 100644 --- a/spinifex/gateway/bedrockruntime.go +++ b/spinifex/gateway/bedrockruntime.go @@ -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 } diff --git a/spinifex/gateway/gateway.go b/spinifex/gateway/gateway.go index a5f91ad0b..cecc8ed9f 100644 --- a/spinifex/gateway/gateway.go +++ b/spinifex/gateway/gateway.go @@ -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{ diff --git a/spinifex/services/awsgw/awsgw.go b/spinifex/services/awsgw/awsgw.go index 67915091f..d5ff2bd89 100644 --- a/spinifex/services/awsgw/awsgw.go +++ b/spinifex/services/awsgw/awsgw.go @@ -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