From 23452ba71adb96bcec75f10f1cae4ea9534c7d1c Mon Sep 17 00:00:00 2001 From: Josh Siviour Date: Sat, 15 Aug 2026 18:49:21 +1000 Subject: [PATCH] fix: fail closed when guardrail OUTPUT extraction fails (#797) The InvokeModel OUTPUT guardrail hook logged and forwarded the model response unguarded when the per-family completion extractor could not parse the body, so output enforcement failed open while the INPUT path already failed closed. Both the non-stream OUTPUT hook and the buffered InvokeModelWithResponseStream source now fail closed: on an unparseable completion (non-stream) or an undecodable assistant chunk (stream) they load the guardrail view and emit a blocked response carrying BlockedOutputsMessaging. A genuinely empty completion (valid chunks, no assistant text) still passes through. Shared resolve/load prologue extracted into loadGuardrailView. --- .../gateway/bedrock/guardrail_inference.go | 31 ++- .../bedrock/guardrail_inference_test.go | 193 ++++++++++++++++++ spinifex/gateway/bedrock/guardrail_invoke.go | 47 +++-- spinifex/gateway/bedrock/invoke.go | 12 +- 4 files changed, 259 insertions(+), 24 deletions(-) diff --git a/spinifex/gateway/bedrock/guardrail_inference.go b/spinifex/gateway/bedrock/guardrail_inference.go index f73347ae7..da082c635 100644 --- a/spinifex/gateway/bedrock/guardrail_inference.go +++ b/spinifex/gateway/bedrock/guardrail_inference.go @@ -11,38 +11,49 @@ import ( "github.com/aws/aws-sdk-go/service/bedrockruntime" ) -// enforceGuardrail is the single resolve->load->filter path Converse and -// ConverseStream both call, mirroring ApplyGuardrail's own sequence. An -// unresolvable or foreign guardrail fails closed: ResourceNotFoundException. -func enforceGuardrail(ctx context.Context, store *GuardrailStore, accountID, ident, version, source string, texts []string) (blocked bool, blockedMessage string, redactedTexts []string, assessments []*bedrockruntime.GuardrailAssessment, err error) { +// loadGuardrailView resolves ident to the stored record's configured view for +// version (DRAFT when empty), the resolve->load->version-select prologue shared +// with enforceGuardrail. An unresolvable or foreign guardrail fails closed. +func loadGuardrailView(ctx context.Context, store *GuardrailStore, accountID, ident, version string) (guardrailView, error) { if store == nil || ident == "" { - return false, "", texts, nil, errGuardrailNotFound(ident, "") + return guardrailView{}, errGuardrailNotFound(ident, "") } id, err := resolveGuardrailID(ident, store.region, accountID) if err != nil { - return false, "", texts, nil, err + return guardrailView{}, err } kv, err := store.bucket(ctx) if err != nil { - return false, "", texts, nil, err + return guardrailView{}, err } rec, found, err := getGuardrailRecord(ctx, kv, accountID, id) if err != nil { - return false, "", texts, nil, err + return guardrailView{}, err } if !found { - return false, "", texts, nil, errGuardrailNotFound(id, "") + return guardrailView{}, errGuardrailNotFound(id, "") } view := rec.guardrailView if version != "" && version != guardrailDraftVersion { snap, ok := rec.Versions[version] if !ok { - return false, "", texts, nil, errGuardrailNotFound(id, version) + return guardrailView{}, errGuardrailNotFound(id, version) } view = snap.guardrailView } + return view, nil +} + +// enforceGuardrail is the single resolve->load->filter path Converse and +// ConverseStream both call, mirroring ApplyGuardrail's own sequence. An +// unresolvable or foreign guardrail fails closed: ResourceNotFoundException. +func enforceGuardrail(ctx context.Context, store *GuardrailStore, accountID, ident, version, source string, texts []string) (blocked bool, blockedMessage string, redactedTexts []string, assessments []*bedrockruntime.GuardrailAssessment, err error) { + view, err := loadGuardrailView(ctx, store, accountID, ident, version) + if err != nil { + return false, "", texts, nil, err + } action, gassessments, outputs, _ := applyGuardrailPolicies(view, texts, source) if action == bedrockruntime.GuardrailActionGuardrailIntervened { diff --git a/spinifex/gateway/bedrock/guardrail_inference_test.go b/spinifex/gateway/bedrock/guardrail_inference_test.go index b8de4a36b..cefb32185 100644 --- a/spinifex/gateway/bedrock/guardrail_inference_test.go +++ b/spinifex/gateway/bedrock/guardrail_inference_test.go @@ -9,6 +9,7 @@ import ( "testing" "github.com/aws/aws-sdk-go/aws" + "github.com/aws/aws-sdk-go/service/bedrock" "github.com/aws/aws-sdk-go/service/bedrockruntime" "github.com/mulgadc/spinifex/spinifex/awserrors" "github.com/stretchr/testify/assert" @@ -721,3 +722,195 @@ func TestGuardrailStreamSource_TraceDisabledOmitsAssessment(t *testing.T) { require.Equal(t, converseStreamEventMetadata, metadata.Kind) assert.Nil(t, metadata.Metadata.Trace) } + +// TestAnthropicInvokeAdapter_UndecodableOutputBlocksWithMessaging proves the +// OUTPUT hook's fail-closed building blocks for the Anthropic family end to +// end at the adapter boundary: anthropicInvokeAdapter.InvokeModel forwards +// the upstream body verbatim (unlike Llama's, which always re-marshals a +// validated Go struct before returning), so a malformed upstream response +// really does reach extractAnthropicCompletionTexts undecodable in +// production -- exactly the extraction failure invoke.go's OUTPUT hook now +// fails closed on, using invokeGuardrailBlockedResponse to build the guarded +// reply. +func TestAnthropicInvokeAdapter_UndecodableOutputBlocksWithMessaging(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("not json")) + })) + defer ts.Close() + + a := &anthropicInvokeAdapter{httpClient: ts.Client(), baseURL: ts.URL} + reqBody := []byte(`{"anthropic_version":"bedrock-2023-05-31","max_tokens":100,"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`) + + respBody, _, err := a.InvokeModel(context.Background(), "anthropic.claude-3-5-sonnet-20240620-v1:0", reqBody, "sk-test") + require.NoError(t, err, "the adapter forwards a non-2xx-but-still-200 malformed body verbatim rather than erroring itself") + assert.Equal(t, "not json", string(respBody)) + + _, ok := extractAnthropicCompletionTexts(respBody) + require.False(t, ok, "a non-JSON upstream body must fail to decode as the expected shape") + + blocked, err := invokeGuardrailBlockedResponse(providerPrefix+vendorAnthropic, "anthropic.claude-3-5-sonnet-20240620-v1:0", "The model response violates our policy.") + require.NoError(t, err) + + texts, ok := extractAnthropicCompletionTexts(blocked) + require.True(t, ok) + assert.Equal(t, []string{"The model response violates our policy."}, texts) + assert.NotContains(t, string(blocked), "not json") +} + +// TestLlamaCompletionGuardrailHelpers_UndecodableBodyBlocksWithMessaging +// exercises the same extractLlamaCompletionTexts/invokeGuardrailBlockedResponse +// pair the OUTPUT hook's fail-closed branch calls for the Llama family. +// Unlike Anthropic, llamaInvokeAdapter.InvokeModel always re-marshals a +// validated llamaInvokeResponse before returning, so an undecodable body can +// never actually reach the hook through the real self-host adapter chain -- +// this proves the building blocks fail closed regardless. +func TestLlamaCompletionGuardrailHelpers_UndecodableBodyBlocksWithMessaging(t *testing.T) { + badBody := []byte(`{"generation": 12345}`) + + _, ok := extractLlamaCompletionTexts(badBody) + require.False(t, ok, "a numeric generation field must fail to decode as the expected string shape") + + blocked, err := invokeGuardrailBlockedResponse(tierSelfHost, "meta.llama3-2-1b-instruct-v1:0", "The model response violates our policy.") + require.NoError(t, err) + + var out llamaInvokeResponse + require.NoError(t, json.Unmarshal(blocked, &out)) + assert.Equal(t, "The model response violates our policy.", out.Generation) + assert.Equal(t, bedrockruntime.StopReasonGuardrailIntervened, out.StopReason) + assert.NotContains(t, string(blocked), "12345") +} + +// TestInvokeRouter_GuardrailOutputPath_DeletedMidRequest_FailsClosed proves +// loadGuardrailView's reload on the OUTPUT path fails closed: the guardrail +// is deleted from inside the backend handler, after the INPUT check already +// passed (the backend is only ever reached once INPUT allowed the request +// through) but before the OUTPUT check runs its own reload. +func TestInvokeRouter_GuardrailOutputPath_DeletedMidRequest_FailsClosed(t *testing.T) { + store := newGuardrailTestStore(t) + ctx := context.Background() + createOut, err := CreateGuardrail(ctx, grCallerAccount, store, createGuardrailInput("invoke-output-guardrail-deleted-mid-request")) + require.NoError(t, err) + + modelID := "meta.llama3-2-1b-instruct-v1:0" + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, derr := DeleteGuardrail(context.Background(), grCallerAccount, store, &bedrock.DeleteGuardrailInput{GuardrailIdentifier: createOut.GuardrailId}) + assert.NoError(t, derr) + + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "choices": [{"text": "hi", "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1} + }`)) + })) + defer ts.Close() + + rt := NewInvokeRouter(nil, NewStaticEndpointResolver(map[string]string{modelID: ts.URL}), nil, grantAll{}, nil, store) + + respBody, _, err := rt.InvokeModel(ctx, grCallerAccount, modelID, []byte(`{"prompt":"hello"}`), aws.StringValue(createOut.GuardrailId), guardrailDraftVersion) + require.Error(t, err) + assert.True(t, awserrors.IsErrorCode(err, awserrors.ErrorResourceNotFoundException)) + assert.Nil(t, respBody, "a fail-closed OUTPUT-path guardrail-load failure must not leak a body") +} + +// TestGuardrailInvokeStreamSource_OutputBlock_UndecodableChunk_Llama drives +// guardrailInvokeStreamSource.Next directly with a scripted chunk sequence +// (independent of any real backend) whose second chunk fails to decode -- +// the buffered/SYNC invoke-stream analog of TestInvokeRouter_ +// GuardrailOutputBlock_UndecodableBody_Anthropic. +func TestGuardrailInvokeStreamSource_OutputBlock_UndecodableChunk_Llama(t *testing.T) { + store := newGuardrailTestStore(t) + ctx := context.Background() + createOut, err := CreateGuardrail(ctx, grCallerAccount, store, createGuardrailInput("invoke-stream-output-undecodable-llama")) + require.NoError(t, err) + + inner := &fakeInvokeStreamSource{chunks: [][]byte{ + []byte(`{"generation":"partial safe text "}`), + []byte(`{"generation": 123}`), + }} + g := newGuardrailInvokeStreamSource(inner, store, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, tierSelfHost) + + chunks := drainInvokeStream(t, g) + require.Len(t, chunks, 2, "the guarded blocked sequence replaces the raw buffer wholesale") + + var delta llamaInvokeStreamChunk + require.NoError(t, json.Unmarshal(chunks[0], &delta)) + assert.Equal(t, "The model response violates our policy.", delta.Generation) + + var final llamaInvokeStreamFinalChunk + require.NoError(t, json.Unmarshal(chunks[1], &final)) + assert.Equal(t, bedrockruntime.StopReasonGuardrailIntervened, final.StopReason) + + for _, c := range chunks { + assert.NotContains(t, string(c), "partial safe text") + } +} + +// TestGuardrailInvokeStreamSource_OutputBlock_UndecodableChunk_Anthropic is +// TestGuardrailInvokeStreamSource_OutputBlock_UndecodableChunk_Llama's +// Anthropic-family sibling: a content_block_delta chunk whose text field is +// the wrong JSON type fails to decode outright, distinct from a +// structurally valid control chunk (content_block_start/stop) carrying no +// text at all. +func TestGuardrailInvokeStreamSource_OutputBlock_UndecodableChunk_Anthropic(t *testing.T) { + store := newGuardrailTestStore(t) + ctx := context.Background() + createOut, err := CreateGuardrail(ctx, grCallerAccount, store, createGuardrailInput("invoke-stream-output-undecodable-anthropic")) + require.NoError(t, err) + + backend := providerPrefix + vendorAnthropic + inner := &fakeInvokeStreamSource{chunks: [][]byte{ + []byte(`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial safe text "}}`), + []byte(`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":123}}`), + }} + g := newGuardrailInvokeStreamSource(inner, store, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, backend) + + chunks := drainInvokeStream(t, g) + require.Len(t, chunks, 5, "anthropicGuardedStreamChunks' minimal event sequence") + + var contentDelta struct { + Delta struct { + Text string `json:"text"` + } `json:"delta"` + } + require.NoError(t, json.Unmarshal(chunks[1], &contentDelta)) + assert.Equal(t, "The model response violates our policy.", contentDelta.Delta.Text) + + var msgDelta struct { + Delta struct { + StopReason string `json:"stop_reason"` + } `json:"delta"` + } + require.NoError(t, json.Unmarshal(chunks[3], &msgDelta)) + assert.Equal(t, bedrockruntime.StopReasonGuardrailIntervened, msgDelta.Delta.StopReason) + + for _, c := range chunks { + assert.NotContains(t, string(c), "partial safe text") + } +} + +// TestGuardrailInvokeStreamSource_EmptyCompletion_PassesThroughUnblocked_Regression +// proves a genuinely empty completion (every chunk decodes fine, none carry +// assistant text) is forwarded unchanged rather than blocked -- the decode- +// failure signal must never fire on a benign control/empty chunk. +func TestGuardrailInvokeStreamSource_EmptyCompletion_PassesThroughUnblocked_Regression(t *testing.T) { + store := newGuardrailTestStore(t) + ctx := context.Background() + createOut, err := CreateGuardrail(ctx, grCallerAccount, store, createGuardrailInput("invoke-stream-output-empty-regression")) + require.NoError(t, err) + + rawChunks := [][]byte{ + []byte(`{"generation":""}`), + []byte(`{"generation":"","prompt_token_count":1,"generation_token_count":0,"stop_reason":"stop"}`), + } + inner := &fakeInvokeStreamSource{chunks: append([][]byte{}, rawChunks...)} + g := newGuardrailInvokeStreamSource(inner, store, grCallerAccount, aws.StringValue(createOut.GuardrailId), guardrailDraftVersion, tierSelfHost) + + chunks := drainInvokeStream(t, g) + require.Len(t, chunks, len(rawChunks)) + for i, c := range chunks { + assert.JSONEq(t, string(rawChunks[i]), string(c)) + } +} diff --git a/spinifex/gateway/bedrock/guardrail_invoke.go b/spinifex/gateway/bedrock/guardrail_invoke.go index 7f79772d3..3d3b6278f 100644 --- a/spinifex/gateway/bedrock/guardrail_invoke.go +++ b/spinifex/gateway/bedrock/guardrail_invoke.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "errors" + "log/slog" "strings" "github.com/aws/aws-sdk-go/service/bedrockruntime" @@ -266,20 +267,19 @@ func invokeGuardrailRedactedCompletion(backend string, respBody []byte, texts [] return respBody, nil } -// extractInvokeStreamChunkText extracts the assistant-text delta (if any) -// carried in one raw provider-native invoke-stream chunk, for buffered -// OUTPUT guardrail accumulation. ok is false when the chunk carries no -// assistant text (e.g. a control/metadata event), not a decode error. -func extractInvokeStreamChunkText(backend string, chunk []byte) (text string, ok bool) { +// extractInvokeStreamChunkText pulls the assistant-text delta from one raw +// invoke-stream chunk for buffered OUTPUT accumulation. ok=false is a benign +// no-text chunk; decodeErr=true is a real JSON-decode failure (fail closed). +func extractInvokeStreamChunkText(backend string, chunk []byte) (text string, ok bool, decodeErr bool) { switch { case backend == tierSelfHost: var c struct { Generation string `json:"generation"` } if json.Unmarshal(chunk, &c) != nil { - return "", false + return "", false, true } - return c.Generation, c.Generation != "" + return c.Generation, c.Generation != "", false case strings.HasPrefix(backend, providerPrefix): switch strings.TrimPrefix(backend, providerPrefix) { case vendorAnthropic: @@ -289,13 +289,16 @@ func extractInvokeStreamChunkText(backend string, chunk []byte) (text string, ok Text string `json:"text"` } `json:"delta"` } - if json.Unmarshal(chunk, &c) != nil || c.Type != "content_block_delta" { - return "", false + if json.Unmarshal(chunk, &c) != nil { + return "", false, true } - return c.Delta.Text, c.Delta.Text != "" + if c.Type != "content_block_delta" { + return "", false, false + } + return c.Delta.Text, c.Delta.Text != "", false } } - return "", false + return "", false, false } // llamaGuardedStreamChunks builds the two-chunk (delta + final) Llama @@ -443,6 +446,7 @@ func (g *guardrailInvokeStreamSource) Next(ctx context.Context) ([]byte, bool, e return nil, false, nil } + var extractionFailed bool for { chunk, ok, err := g.inner.Next(ctx) if err != nil { @@ -451,13 +455,32 @@ func (g *guardrailInvokeStreamSource) Next(ctx context.Context) ([]byte, bool, e if !ok { break } + text, hasText, decodeErr := extractInvokeStreamChunkText(g.backend, chunk) + if decodeErr { + extractionFailed = true + break + } g.buffered = append(g.buffered, chunk) - if text, hasText := extractInvokeStreamChunkText(g.backend, chunk); hasText { + if hasText { g.text.WriteString(text) } } g.assessed = true + if extractionFailed { + slog.Error("invoke-stream: failed to decode assistant chunk for guardrail OUTPUT check, blocking", "backend", g.backend) + view, verr := loadGuardrailView(ctx, g.store, g.accountID, g.ident, g.version) + if verr != nil { + return nil, false, verr + } + var berr error + g.queue, berr = buildGuardedInvokeStreamChunks(g.backend, view.BlockedOutputsMessaging, true) + if berr != nil { + return nil, false, berr + } + return g.Next(ctx) + } + original := g.text.String() blocked, message, redacted, _, err := enforceGuardrail(ctx, g.store, g.accountID, g.ident, g.version, bedrockruntime.GuardrailContentSourceOutput, []string{original}) diff --git a/spinifex/gateway/bedrock/invoke.go b/spinifex/gateway/bedrock/invoke.go index 141dc49af..ce61a82aa 100644 --- a/spinifex/gateway/bedrock/invoke.go +++ b/spinifex/gateway/bedrock/invoke.go @@ -194,8 +194,16 @@ func (rt *InvokeRouter) InvokeModel(ctx context.Context, accountID, modelID stri texts, extractOK := extractInvokeCompletionTexts(backend, respBody) if !extractOK { - slog.Error("invoke: failed to extract completion text for guardrail OUTPUT check, forwarding unguarded", "model", modelID, "backend", backend) - return respBody, contentType, err + slog.Error("invoke: failed to extract completion text for guardrail OUTPUT check, blocking", "model", modelID, "backend", backend) + view, verr := loadGuardrailView(ctx, rt.guardrails, accountID, guardrailIdent, guardrailVersion) + if verr != nil { + return nil, "", verr + } + respBody, err = invokeGuardrailBlockedResponse(backend, modelID, view.BlockedOutputsMessaging) + if err != nil { + return nil, "", err + } + return respBody, "application/json", nil } var blockedOut bool var messageOut string