From 02aae2bbb49b941290a01e7ded8a1ca051603775 Mon Sep 17 00:00:00 2001 From: Aska Date: Sun, 6 Sep 2026 05:07:13 +0800 Subject: [PATCH] =?UTF-8?q?feat(guardrail):=20=E8=BE=93=E5=85=A5/=E8=BE=93?= =?UTF-8?q?=E5=87=BA=E6=8A=A4=E6=A0=8F=EF=BC=88=E5=85=B3=E9=94=AE=E8=AF=8D?= =?UTF-8?q?/=E6=95=8F=E6=84=9F=E4=BF=A1=E6=81=AF/=E9=95=BF=E5=BA=A6?= =?UTF-8?q?=EF=BC=8Cpipeline=20fail-fast=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/runtime/guardrail/builtin.go | 247 ++++++++++ internal/runtime/guardrail/guardrail.go | 166 +++++++ internal/runtime/guardrail/guardrail_test.go | 450 +++++++++++++++++++ 3 files changed, 863 insertions(+) create mode 100644 internal/runtime/guardrail/builtin.go create mode 100644 internal/runtime/guardrail/guardrail.go create mode 100644 internal/runtime/guardrail/guardrail_test.go diff --git a/internal/runtime/guardrail/builtin.go b/internal/runtime/guardrail/builtin.go new file mode 100644 index 0000000..314daab --- /dev/null +++ b/internal/runtime/guardrail/builtin.go @@ -0,0 +1,247 @@ +package guardrail + +import ( + "context" + "errors" + "regexp" + "strings" + "unicode/utf8" + + "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model" +) + +// allow returns a passing verdict. +func allow() *Verdict { return &Verdict{Allow: true} } + +// deny returns a blocking verdict with the given reason. +func deny(reason string) *Verdict { return &Verdict{Allow: false, Reason: reason} } + +// requestText returns the concatenation of every text-bearing content part in a +// request's messages. Image parts are ignored. +func requestText(input *model.Request) string { + var b strings.Builder + for _, m := range input.Messages { + b.WriteString(messageText(m)) + } + return b.String() +} + +// responseText returns the text-bearing output of a response. It prefers the +// convenience Content field and falls back to the normalized Message text parts +// when Content is empty, so a response that duplicates its content in both is +// measured once rather than twice. +func responseText(output *model.Response) string { + if output.Content != "" { + return output.Content + } + var b strings.Builder + for _, c := range output.Message.Content { + if c.Type == model.ContentTypeText { + b.WriteString(c.Text) + } + } + return b.String() +} + +// messageText returns the text-bearing content parts of a single message. +func messageText(msg model.Message) string { + var b strings.Builder + for _, c := range msg.Content { + if c.Type == model.ContentTypeText { + b.WriteString(c.Text) + } + } + return b.String() +} + +// KeywordGuard is a content-policy guardrail that blocks any input or output +// containing one of its configured keywords. Matching is a case-insensitive +// substring search, so each keyword is matched as a plain fragment rather than +// a regexp. +// +// NewKeywordGuard builds a KeywordGuard from a word list. With no keywords the +// guard allows everything. +type KeywordGuard struct { + keywords []string // lowercased +} + +// NewKeywordGuard builds a KeywordGuard that blocks payloads containing any of +// the given keywords. +func NewKeywordGuard(keywords ...string) *KeywordGuard { + normalized := make([]string, 0, len(keywords)) + for _, k := range keywords { + if k = strings.ToLower(strings.TrimSpace(k)); k != "" { + normalized = append(normalized, k) + } + } + return &KeywordGuard{keywords: normalized} +} + +// CheckInput blocks the request when any keyword is present in its text. +func (g *KeywordGuard) CheckInput(ctx context.Context, input *model.Request) (*Verdict, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if input == nil { + return nil, errors.New("guardrail: nil request") + } + return g.check(requestText(input)) +} + +// CheckOutput blocks the response when any keyword is present in its text. +func (g *KeywordGuard) CheckOutput(ctx context.Context, output *model.Response) (*Verdict, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if output == nil { + return nil, errors.New("guardrail: nil response") + } + return g.check(responseText(output)) +} + +func (g *KeywordGuard) check(text string) (*Verdict, error) { + if len(g.keywords) == 0 { + return allow(), nil + } + lower := strings.ToLower(text) + for _, kw := range g.keywords { + if strings.Contains(lower, kw) { + return deny("content_policy_blocked: " + kw), nil + } + } + return allow(), nil +} + +// SensitiveMode selects how a SensitiveGuard reacts to a match. +type SensitiveMode int + +const ( + // SensitiveBlock denies any payload that matches a pattern. + SensitiveBlock SensitiveMode = iota + // SensitiveMark allows the payload but flags it in Verdict.Reason. + SensitiveMark +) + +// SensitivePattern pairs a compiled regular expression with a human label. The +// label is what appears in Verdict.Reason; the pattern's captured secret is +// never echoed anywhere. +type SensitivePattern struct { + Label string + Regexp *regexp.Regexp +} + +// DefaultSensitivePatterns returns the built-in set of patterns used to detect +// common API keys, bearer tokens, cloud credentials, passwords and private-key +// material. Callers may extend or replace these with their own patterns. +func DefaultSensitivePatterns() []SensitivePattern { + return []SensitivePattern{ + {Label: "api_key", Regexp: regexp.MustCompile(`(?i)\bapi[_-]?key\s*[:=]\s*\S+`)}, + {Label: "openai_key", Regexp: regexp.MustCompile(`\bsk-[A-Za-z0-9_-]{20,}\b`)}, + {Label: "aws_access_key", Regexp: regexp.MustCompile(`\bAKIA[0-9A-Z]{16}\b`)}, + {Label: "bearer_token", Regexp: regexp.MustCompile(`(?i)\bbearer\s+[A-Za-z0-9._~-]+`)}, + {Label: "password", Regexp: regexp.MustCompile(`(?i)\b(?:bind_?)?password\s*[:=]\s*\S+`)}, + {Label: "private_key", Regexp: regexp.MustCompile(`-----BEGIN (?:RSA |EC |DSA |OPENSSH |PGP )?PRIVATE KEY-----`)}, + } +} + +// SensitiveGuard is a guardrail that detects sensitive information such as API +// keys, tokens and passwords using regular expressions. Matching is +// case-insensitive where the pattern says so; it never writes the matched secret +// into a Reason, log or return value. +// +// NewSensitiveGuard builds a SensitiveGuard from patterns and a mode. +type SensitiveGuard struct { + patterns []SensitivePattern + mode SensitiveMode +} + +// NewSensitiveGuard builds a SensitiveGuard. In SensitiveBlock mode a match +// denies the payload; in SensitiveMark mode a match allows it but records the +// flagged label in Verdict.Reason. +func NewSensitiveGuard(patterns []SensitivePattern, mode SensitiveMode) *SensitiveGuard { + return &SensitiveGuard{patterns: patterns, mode: mode} +} + +// CheckInput evaluates the request text against every pattern. +func (g *SensitiveGuard) CheckInput(ctx context.Context, input *model.Request) (*Verdict, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if input == nil { + return nil, errors.New("guardrail: nil request") + } + return g.check(requestText(input)) +} + +// CheckOutput evaluates the response text against every pattern. +func (g *SensitiveGuard) CheckOutput(ctx context.Context, output *model.Response) (*Verdict, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if output == nil { + return nil, errors.New("guardrail: nil response") + } + return g.check(responseText(output)) +} + +func (g *SensitiveGuard) check(text string) (*Verdict, error) { + for _, p := range g.patterns { + if p.Regexp == nil || !p.Regexp.MatchString(text) { + continue + } + if g.mode == SensitiveBlock { + return deny("sensitive_information_blocked: " + p.Label), nil + } + return &Verdict{Allow: true, Reason: "sensitive_information_detected: " + p.Label}, nil + } + return allow(), nil +} + +// LengthGuard is a guardrail that blocks input or output exceeding a configured +// character limit. A limit of zero disables that side. +// +// NewLengthGuard builds a LengthGuard with per-side limits. +type LengthGuard struct { + maxInputChars int + maxOutputChars int +} + +// NewLengthGuard builds a LengthGuard. Limits are measured in Unicode code +// points; a limit of zero means unlimited. +func NewLengthGuard(maxInputChars, maxOutputChars int) *LengthGuard { + return &LengthGuard{maxInputChars: maxInputChars, maxOutputChars: maxOutputChars} +} + +// CheckInput blocks the request when its text exceeds the input character limit. +func (g *LengthGuard) CheckInput(ctx context.Context, input *model.Request) (*Verdict, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if input == nil { + return nil, errors.New("guardrail: nil request") + } + if g.maxInputChars <= 0 { + return allow(), nil + } + if n := utf8.RuneCountInString(requestText(input)); n > g.maxInputChars { + return deny("input_length_exceeded"), nil + } + return allow(), nil +} + +// CheckOutput blocks the response when its text exceeds the output limit. +func (g *LengthGuard) CheckOutput(ctx context.Context, output *model.Response) (*Verdict, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if output == nil { + return nil, errors.New("guardrail: nil response") + } + if g.maxOutputChars <= 0 { + return allow(), nil + } + if n := utf8.RuneCountInString(responseText(output)); n > g.maxOutputChars { + return deny("output_length_exceeded"), nil + } + return allow(), nil +} diff --git a/internal/runtime/guardrail/guardrail.go b/internal/runtime/guardrail/guardrail.go new file mode 100644 index 0000000..8731a28 --- /dev/null +++ b/internal/runtime/guardrail/guardrail.go @@ -0,0 +1,166 @@ +// Package guardrail provides input/output validation and safety interception +// for the Cyrene agent runtime. +// +// It implements the Guardrail design in DevDocs/design/04-运行时.md §5: one or +// more checks run in parallel with agent execution and fail fast on the first +// denial. A guardrail evaluates a model.Request (before it reaches the LLM) or +// a model.Response (before it is returned to the user) and returns a Verdict +// that either allows the payload or blocks it with a human-readable Reason. +// +// The package ships three built-in, configurable checks: +// +// - KeywordGuard: content-policy keyword/category matching. +// - SensitiveGuard: regex detection of API keys, tokens, passwords and other +// secrets, in either block or mark mode. It never echoes the matched secret +// into a Reason, log or return value. +// - LengthGuard: input/output character-length limits. +// +// Checks are composed into a Pipeline (see New). All errors surfaced by a +// Pipeline wrap the sentinel ErrGuardrail, so callers can test with errors.Is. +// +// Guardrail is a leaf package: it depends only on the standard library and +// internal/runtime/model types, and it never imports the agent loop or any other +// runtime package, so it can sit alongside the loop without an import cycle. +package guardrail + +import ( + "context" + "errors" + "fmt" + + "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model" +) + +// Verdict is the outcome of a single guardrail evaluation. +// +// When Allow is true the payload is accepted; when it is false the payload is +// blocked and Reason explains why. Reason is also set on an allow verdict to +// flag context (for example a sensitive-information mark) without carrying the +// offending secret. +type Verdict struct { + Allow bool + Reason string +} + +// Guardrail evaluates a request or a response and returns a Verdict. +// +// Implementations must be safe for concurrent use and honor ctx cancellation: +// a canceled context returns immediately with a non-nil error. A nil Verdict +// on a nil error is treated as allow by the Pipeline. +type Guardrail interface { + // CheckInput evaluates a completion request before it reaches the LLM. + CheckInput(ctx context.Context, input *model.Request) (*Verdict, error) + // CheckOutput evaluates a completion response before it is returned. + CheckOutput(ctx context.Context, output *model.Response) (*Verdict, error) +} + +// ErrGuardrail is the sentinel wrapped around any error returned by a +// Pipeline's CheckInput or CheckOutput. Use errors.Is to detect it. +var ErrGuardrail = errors.New("guardrail: evaluation failed") + +// wrapErr attaches the ErrGuardrail sentinel to a guard's internal error so the +// caller can tell a guardrail failure apart from any other runtime error. +func wrapErr(err error) error { + return fmt.Errorf("%w: %w", ErrGuardrail, err) +} + +// Pipeline runs a sequence of Guardrails in the order they were registered and +// fails fast: the first deny verdict short-circuits the rest and is returned +// as-is. Because the sentinel is non-fatal, allow verdicts let the remaining +// guards run, and the reason of the last flagged verdict is preserved. +type Pipeline struct { + guards []Guardrail +} + +// New builds a Pipeline from the given options. Options are applied in order, +// so the guards run in the same order they are passed. With no options the +// pipeline allows every payload. +func New(opts ...Option) *Pipeline { + p := &Pipeline{} + for _, o := range opts { + o(p) + } + return p +} + +// Option configures a Pipeline. Options are applied in the order passed to New. +type Option func(*Pipeline) + +// WithGuard appends a custom Guardrail to the pipeline. +func WithGuard(g Guardrail) Option { + return func(p *Pipeline) { p.guards = append(p.guards, g) } +} + +// WithKeywords appends a KeywordGuard that blocks any payload containing one of +// the given keywords (case-insensitive substring match). +func WithKeywords(keywords ...string) Option { + return func(p *Pipeline) { p.guards = append(p.guards, NewKeywordGuard(keywords...)) } +} + +// WithSensitiveInfo appends a SensitiveGuard using the built-in default patterns +// in the given mode. Use WithSensitivePatterns to supply your own patterns. +func WithSensitiveInfo(mode SensitiveMode) Option { + return WithSensitivePatterns(mode, DefaultSensitivePatterns()...) +} + +// WithSensitivePatterns appends a SensitiveGuard that matches the given patterns +// in the given mode. +func WithSensitivePatterns(mode SensitiveMode, patterns ...SensitivePattern) Option { + return func(p *Pipeline) { p.guards = append(p.guards, NewSensitiveGuard(patterns, mode)) } +} + +// WithLengthLimits appends a LengthGuard capping input and output in characters. +// A zero limit disables that side. +func WithLengthLimits(maxInputChars, maxOutputChars int) Option { + return func(p *Pipeline) { p.guards = append(p.guards, NewLengthGuard(maxInputChars, maxOutputChars)) } +} + +// CheckInput runs every registered guard against the request in order and +// returns the first deny verdict. An internal guard error is wrapped with +// ErrGuardrail. +func (p *Pipeline) CheckInput(ctx context.Context, input *model.Request) (*Verdict, error) { + return p.run(ctx, input) +} + +// CheckOutput runs every registered guard against the response in order and +// returns the first deny verdict. An internal guard error is wrapped with +// ErrGuardrail. +func (p *Pipeline) CheckOutput(ctx context.Context, output *model.Response) (*Verdict, error) { + return p.run(ctx, output) +} + +// run is the shared fail-fast loop. eval dispatches the payload to the guard +// and returns a verdict, an error, or both (an error wins). The first deny +// verdict stops the loop, while a flagged-allow verdict is remembered as the +// final reason. +func (p *Pipeline) run(ctx context.Context, payload any) (*Verdict, error) { + var reason string + for _, g := range p.guards { + v, err := p.eval(ctx, g, payload) + if err != nil { + return nil, wrapErr(err) + } + if v == nil { + continue + } + if !v.Allow { + return v, nil // fail-fast: first denial + } + if v.Reason != "" { + reason = v.Reason + } + } + return &Verdict{Allow: true, Reason: reason}, nil +} + +// eval dispatches a payload to the matching guard method. +func (p *Pipeline) eval(ctx context.Context, g Guardrail, payload any) (*Verdict, error) { + switch t := payload.(type) { + case *model.Request: + return g.CheckInput(ctx, t) + case *model.Response: + return g.CheckOutput(ctx, t) + default: + return nil, fmt.Errorf("guardrail: unsupported payload %T", payload) + } +} diff --git a/internal/runtime/guardrail/guardrail_test.go b/internal/runtime/guardrail/guardrail_test.go new file mode 100644 index 0000000..79b1cb4 --- /dev/null +++ b/internal/runtime/guardrail/guardrail_test.go @@ -0,0 +1,450 @@ +package guardrail + +import ( + "context" + "errors" + "regexp" + "strings" + "testing" + + "git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model" +) + +// textRequest builds a text-only request from the given message bodies. +func textRequest(t *testing.T, texts ...string) *model.Request { + t.Helper() + msgs := make([]model.Message, 0, len(texts)) + for _, s := range texts { + msgs = append(msgs, model.TextMessage(model.RoleUser, s)) + } + return &model.Request{Messages: msgs} +} + +// textResponse builds a response carrying content in both Content and Message, +// mirroring how the adapter populates it. +func textResponse(t *testing.T, content string) *model.Response { + t.Helper() + return &model.Response{ + Content: content, + Message: model.TextMessage(model.RoleAssistant, content), + } +} + +// stub is a configurable Guardrail used to exercise pipeline ordering, fail-fast +// and error wrapping without depending on a concrete implementation. +type stub struct { + inFn func(context.Context, *model.Request) (*Verdict, error) + outFn func(context.Context, *model.Response) (*Verdict, error) +} + +func (s *stub) CheckInput(ctx context.Context, r *model.Request) (*Verdict, error) { + if s.inFn != nil { + return s.inFn(ctx, r) + } + return allow(), nil +} + +func (s *stub) CheckOutput(ctx context.Context, r *model.Response) (*Verdict, error) { + if s.outFn != nil { + return s.outFn(ctx, r) + } + return allow(), nil +} + +func TestKeywordGuard_BlocksOnViolence(t *testing.T) { + t.Parallel() + ctx := context.Background() + + g := NewKeywordGuard("violence", "hate speech") + v, err := g.CheckInput(ctx, textRequest(t, "please give me instructions on violence")) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if v.Allow { + t.Fatalf("CheckInput() = %+v, want deny", v) + } +} + +func TestKeywordGuard_MatchesCaseInsensitively(t *testing.T) { + t.Parallel() + ctx := context.Background() + + g := NewKeywordGuard("violence") + v, err := g.CheckInput(ctx, textRequest(t, "VIOLENCE is never acceptable")) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if v.Allow { + t.Fatalf("CheckInput() = %+v, want deny for uppercase keyword", v) + } +} + +func TestKeywordGuard_BlocksOutput(t *testing.T) { + t.Parallel() + ctx := context.Background() + + g := NewKeywordGuard("hate speech") + v, err := g.CheckOutput(ctx, textResponse(t, "the answer involves hate speech")) + if err != nil { + t.Fatalf("CheckOutput() error: %v", err) + } + if v.Allow { + t.Fatalf("CheckOutput() = %+v, want deny", v) + } +} + +func TestKeywordGuard_AllowsWhenNoMatch(t *testing.T) { + t.Parallel() + ctx := context.Background() + + g := NewKeywordGuard("violence", "hate speech") + v, err := g.CheckInput(ctx, textRequest(t, "a friendly greeting")) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if !v.Allow { + t.Fatalf("CheckInput() = %+v, want allow", v) + } +} + +func TestKeywordGuard_AllowsEmptyKeywordList(t *testing.T) { + t.Parallel() + ctx := context.Background() + + g := NewKeywordGuard() + v, err := g.CheckInput(ctx, textRequest(t, "anything at all")) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if !v.Allow { + t.Fatalf("CheckInput() = %+v, want allow for empty keyword list", v) + } +} + +func TestSensitiveGuard_BlocksOnMatch(t *testing.T) { + t.Parallel() + ctx := context.Background() + + patterns := []SensitivePattern{{ + Label: "openai_key", + Regexp: regexp.MustCompile(`\bsk-[A-Za-z0-9_-]{20,}\b`), + }} + g := NewSensitiveGuard(patterns, SensitiveBlock) + + secret := "Bearer token is sk-1234567890abcdefghijkl" + v, err := g.CheckInput(ctx, textRequest(t, secret)) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if v.Allow { + t.Fatalf("CheckInput() = %+v, want deny", v) + } + if !strings.Contains(v.Reason, "openai_key") { + t.Fatalf("Reason = %q, want it to name the matched pattern", v.Reason) + } + if strings.Contains(v.Reason, "sk-1234567890abcdefghijkl") { + t.Fatalf("Reason = %q, must not echo the secret", v.Reason) + } +} + +func TestSensitiveGuard_MarksOnMatch(t *testing.T) { + t.Parallel() + ctx := context.Background() + + patterns := []SensitivePattern{{ + Label: "mysql_password", + Regexp: regexp.MustCompile(`(?i)password\s*[:=]\s*\S+`), + }} + g := NewSensitiveGuard(patterns, SensitiveMark) + + v, err := g.CheckInput(ctx, textRequest(t, "connect using password=hunter2")) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if !v.Allow { + t.Fatalf("CheckInput() = %+v, want allow in mark mode", v) + } + if !strings.Contains(v.Reason, "mysql_password") { + t.Fatalf("Reason = %q, want the flagged label", v.Reason) + } + if strings.Contains(v.Reason, "hunter2") { + t.Fatalf("Reason = %q, must not echo the secret", v.Reason) + } +} + +func TestSensitiveGuard_AllowsWhenNoMatch(t *testing.T) { + t.Parallel() + ctx := context.Background() + + patterns := []SensitivePattern{{ + Label: "openai_key", + Regexp: regexp.MustCompile(`\bsk-[A-Za-z0-9_-]{20,}\b`), + }} + g := NewSensitiveGuard(patterns, SensitiveBlock) + + v, err := g.CheckInput(ctx, textRequest(t, "a harmless sentence")) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if !v.Allow { + t.Fatalf("CheckInput() = %+v, want allow", v) + } +} + +func TestSensitiveGuard_SkipsNilPattern(t *testing.T) { + t.Parallel() + ctx := context.Background() + + g := NewSensitiveGuard([]SensitivePattern{{Label: "broken", Regexp: nil}}, SensitiveBlock) + v, err := g.CheckInput(ctx, textRequest(t, "any text")) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if !v.Allow { + t.Fatalf("CheckInput() = %+v, want allow when pattern regexp is nil", v) + } +} + +func TestSensitiveGuard_DefaultPatternsCatchBearerToken(t *testing.T) { + t.Parallel() + ctx := context.Background() + + g := NewSensitiveGuard(DefaultSensitivePatterns(), SensitiveBlock) + v, err := g.CheckInput(ctx, textRequest(t, "Authorization: Bearer abc.def-ghi_jkl")) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if v.Allow { + t.Fatalf("CheckInput() = %+v, want deny for a bearer token", v) + } +} + +func TestSensitiveGuard_ChecksOutput(t *testing.T) { + t.Parallel() + ctx := context.Background() + + patterns := []SensitivePattern{{ + Label: "api_key", + Regexp: regexp.MustCompile(`(?i)api[_-]?key\s*[:=]\s*\S+`), + }} + g := NewSensitiveGuard(patterns, SensitiveBlock) + v, err := g.CheckOutput(ctx, textResponse(t, "API_KEY=sk-secret-value")) + if err != nil { + t.Fatalf("CheckOutput() error: %v", err) + } + if v.Allow { + t.Fatalf("CheckOutput() = %+v, want deny", v) + } +} + +func TestLengthGuard_BlocksOverLimit(t *testing.T) { + t.Parallel() + ctx := context.Background() + + tests := []struct { + name string + limit int + text string + want string + }{ + {name: "input_over", limit: 5, text: "hello world", want: "input_length_exceeded"}, + {name: "output_over", limit: 5, text: "a very long answer", want: "output_length_exceeded"}, + {name: "unicode_runes_not_bytes", limit: 2, text: "你好啊", want: "input_length_exceeded"}, + } + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + g := NewLengthGuard(tt.limit, tt.limit) + var v *Verdict + var err error + if tt.name == "output_over" { + v, err = g.CheckOutput(ctx, textResponse(t, tt.text)) + } else { + v, err = g.CheckInput(ctx, textRequest(t, tt.text)) + } + if err != nil { + t.Fatalf("check error: %v", err) + } + if v.Allow { + t.Fatalf("got %+v, want deny", v) + } + if !strings.Contains(v.Reason, tt.want) { + t.Fatalf("Reason = %q, want %q", v.Reason, tt.want) + } + }) + } +} + +func TestLengthGuard_AllowsWithinLimit(t *testing.T) { + t.Parallel() + ctx := context.Background() + + g := NewLengthGuard(20, 20) + v, err := g.CheckInput(ctx, textRequest(t, "ok")) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if !v.Allow { + t.Fatalf("CheckInput() = %+v, want allow under limit", v) + } +} + +func TestLengthGuard_ZeroLimitIsUnlimited(t *testing.T) { + t.Parallel() + ctx := context.Background() + + g := NewLengthGuard(0, 0) + v, err := g.CheckInput(ctx, textRequest(t, strings.Repeat("x", 100))) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if !v.Allow { + t.Fatalf("CheckInput() = %+v, want allow for unlimited", v) + } +} + +func TestPipeline_AllowsCleanPayload(t *testing.T) { + t.Parallel() + ctx := context.Background() + + p := New( + WithKeywords("violence", "hate"), + WithSensitiveInfo(SensitiveBlock), + WithLengthLimits(100, 100), + ) + res, err := p.CheckInput(ctx, textRequest(t, "a friendly greeting")) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if !res.Allow { + t.Fatalf("CheckInput() = %+v, want allow", res) + } + + res, err = p.CheckOutput(ctx, textResponse(t, "a helpful reply")) + if err != nil { + t.Fatalf("CheckOutput() error: %v", err) + } + if !res.Allow { + t.Fatalf("CheckOutput() = %+v, want allow", res) + } +} + +func TestPipeline_EmptyPipelineAllows(t *testing.T) { + t.Parallel() + ctx := context.Background() + + p := New() + v, err := p.CheckInput(ctx, textRequest(t, "anything")) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if !v.Allow { + t.Fatalf("CheckInput() = %+v, want allow with no guards", v) + } +} + +func TestPipeline_FailFastOnInput(t *testing.T) { + t.Parallel() + ctx := context.Background() + + denyFirst := &stub{inFn: func(context.Context, *model.Request) (*Verdict, error) { + return deny("first"), nil + }} + keywordCalled := false + second := &stub{inFn: func(context.Context, *model.Request) (*Verdict, error) { + keywordCalled = true + return deny("second"), nil + }} + + p := New(WithGuard(denyFirst), WithGuard(second)) + v, err := p.CheckInput(ctx, textRequest(t, "text")) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if v.Allow || v.Reason != "first" { + t.Fatalf("verdict = %+v, want deny reason %q", v, "first") + } + if keywordCalled { + t.Fatal("second guard ran after a denial; want fail-fast") + } +} + +func TestPipeline_FailFastOnOutput(t *testing.T) { + t.Parallel() + ctx := context.Background() + + denyFirst := &stub{outFn: func(context.Context, *model.Response) (*Verdict, error) { + return deny("blocked-output"), nil + }} + second := &stub{outFn: func(context.Context, *model.Response) (*Verdict, error) { + t.Fatal("second guard ran; want fail-fast") + return allow(), nil + }} + p := New(WithGuard(denyFirst), WithGuard(second)) + v, err := p.CheckOutput(ctx, textResponse(t, "text")) + if err != nil { + t.Fatalf("CheckOutput() error: %v", err) + } + if v.Allow || v.Reason != "blocked-output" { + t.Fatalf("verdict = %+v, want deny reason %q", v, "blocked-output") + } +} + +func TestPipeline_ErrorWrapsErrGuardrail(t *testing.T) { + t.Parallel() + ctx := context.Background() + + boom := errors.New("boom") + g := &stub{inFn: func(context.Context, *model.Request) (*Verdict, error) { + return nil, boom + }} + p := New(WithGuard(g)) + _, err := p.CheckInput(ctx, textRequest(t, "text")) + if err == nil { + t.Fatal("CheckInput() returned nil error, want ErrGuardrail") + } + if !errors.Is(err, ErrGuardrail) { + t.Fatalf("err = %v, want it to wrap ErrGuardrail", err) + } + if !errors.Is(err, boom) { + t.Fatalf("err = %v, want it to preserve the cause", err) + } +} + +func TestPipeline_CanceledContextWraps(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + g := NewKeywordGuard("x") + p := New(WithGuard(g)) + _, err := p.CheckInput(ctx, textRequest(t, "text")) + if err == nil { + t.Fatal("CheckInput() returned nil error, want wrapped ctx error") + } + if !errors.Is(err, ErrGuardrail) { + t.Fatalf("err = %v, want it to wrap ErrGuardrail", err) + } +} + +func TestPipeline_PreservesMarkReason(t *testing.T) { + t.Parallel() + ctx := context.Background() + + patterns := []SensitivePattern{{ + Label: "api_key", + Regexp: regexp.MustCompile(`(?i)api[_-]?key\s*[:=]\s*\S+`), + }} + p := New(WithSensitivePatterns(SensitiveMark, patterns...)) + v, err := p.CheckInput(ctx, textRequest(t, "API_KEY=secret")) + if err != nil { + t.Fatalf("CheckInput() error: %v", err) + } + if !v.Allow { + t.Fatalf("CheckInput() = %+v, want allow in mark mode", v) + } + if !strings.Contains(v.Reason, "api_key") { + t.Fatalf("Reason = %q, want the flagged label preserved", v.Reason) + } +}