diff --git a/crates/aisix-gateway/src/chat.rs b/crates/aisix-gateway/src/chat.rs index 7be97f65..97777332 100644 --- a/crates/aisix-gateway/src/chat.rs +++ b/crates/aisix-gateway/src/chat.rs @@ -346,6 +346,33 @@ pub struct ChatResponse { pub usage: UsageStats, } +impl ChatResponse { + /// The client-visible output text that content/DLP output guardrails + /// must inspect: the assistant `content` plus any `tool_calls` + /// material (function names + arguments, and Anthropic `tool_use` + /// normalized into the same `extra["tool_calls"]` slot). Tool-call + /// output is rendered to clients but would otherwise bypass output + /// guardrails that only read `message.content` (#448). + /// + /// Reasoning/thinking content is intentionally NOT included — it is + /// left out of output-guardrail scope by design. + pub fn guardrail_output_text(&self) -> String { + let mut out = self.message.content.clone(); + if let Some(tool_calls) = self.message.extra.get("tool_calls") { + if !tool_calls.is_null() { + if !out.is_empty() { + out.push('\n'); + } + // Serialize the whole tool-call payload so no function + // name or argument can escape inspection regardless of the + // provider-specific shape. + out.push_str(&tool_calls.to_string()); + } + } + out + } +} + /// One streamed delta event. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] diff --git a/crates/aisix-guardrails/src/bedrock.rs b/crates/aisix-guardrails/src/bedrock.rs index edc22972..5891b136 100644 --- a/crates/aisix-guardrails/src/bedrock.rs +++ b/crates/aisix-guardrails/src/bedrock.rs @@ -290,7 +290,7 @@ impl Guardrail for BedrockGuardrail { ) { return GuardrailVerdict::Allow; } - let text = resp.message.content.clone(); + let text = resp.guardrail_output_text(); if text.is_empty() { return GuardrailVerdict::Allow; } diff --git a/crates/aisix-guardrails/src/keyword.rs b/crates/aisix-guardrails/src/keyword.rs index ab42f741..d252a14f 100644 --- a/crates/aisix-guardrails/src/keyword.rs +++ b/crates/aisix-guardrails/src/keyword.rs @@ -120,7 +120,9 @@ impl Guardrail for KeywordBlocklist { if !self.check_output_enabled { return GuardrailVerdict::Allow; } - match self.first_match(&resp.message.content) { + // Inspect content + tool-call output (#448), not just content. + let text = resp.guardrail_output_text(); + match self.first_match(&text) { Some(rule) => GuardrailVerdict::Block { reason: format!("output blocked by {}", rule.description()), }, @@ -197,6 +199,34 @@ mod tests { assert!(v.is_block()); } + #[tokio::test] + async fn output_check_inspects_tool_call_arguments() { + // A forbidden word that appears ONLY inside tool_call arguments + // (message.content is empty) must still be blocked — tool-call + // output is client-visible and must not bypass guardrails (#448). + let g = KeywordBlocklist::new(vec![KeywordRule::literal("dangerous")]); + let mut msg = ChatMessage::assistant(""); + msg.extra.insert( + "tool_calls".into(), + serde_json::json!([{ + "id": "call_1", + "type": "function", + "function": { + "name": "run", + "arguments": "{\"cmd\":\"do something dangerous\"}" + } + }]), + ); + let r = ChatResponse { + id: "r".into(), + model: "m".into(), + message: msg, + finish_reason: FinishReason::Stop, + usage: UsageStats::new(0, 0), + }; + assert!(g.check_output(&r).await.is_block()); + } + #[tokio::test] async fn input_only_skips_output_checks() { let g = KeywordBlocklist::input_only(vec![KeywordRule::literal("zeta")]); diff --git a/crates/aisix-guardrails/src/prompt_shield.rs b/crates/aisix-guardrails/src/prompt_shield.rs index 99d6b81a..f7ecd37e 100644 --- a/crates/aisix-guardrails/src/prompt_shield.rs +++ b/crates/aisix-guardrails/src/prompt_shield.rs @@ -287,7 +287,7 @@ impl Guardrail for PromptShieldGuardrail { ) { return GuardrailVerdict::Allow; } - let text = resp.message.content.clone(); + let text = resp.guardrail_output_text(); if text.is_empty() { return GuardrailVerdict::Allow; } @@ -328,8 +328,15 @@ fn chunk_text(text: &str, max_chars: usize) -> Vec { chunks.push(std::mem::take(&mut current)); } if word_chars > max_chars { - // Single word longer than the limit — truncate and flush. - chunks.push(word.chars().take(max_chars).collect()); + // Single word longer than the limit — split it into + // max_chars-sized pieces so the ENTIRE token is evaluated. + // Truncating to the first max_chars (the previous behavior) + // let the trailing part of an oversized whitespace-free + // input reach the model unscanned (#448). + let word_chars_vec: Vec = word.chars().collect(); + for piece in word_chars_vec.chunks(max_chars) { + chunks.push(piece.iter().collect()); + } continue; } } @@ -455,11 +462,19 @@ mod tests { } #[test] - fn chunk_text_single_oversized_word_is_truncated() { + fn chunk_text_single_oversized_word_is_fully_covered() { + // A whitespace-free token longer than the limit must be split so + // the ENTIRE token is evaluated — not truncated to the prefix, + // which let the trailing part reach the model unscanned (#448). let word: String = "x".repeat(15_000); let chunks = chunk_text(&word, MAX_PROMPT_CHARS); - assert_eq!(chunks.len(), 1); - assert_eq!(chunks[0].chars().count(), MAX_PROMPT_CHARS); + assert_eq!(chunks.len(), 2, "15k chars over a 10k limit → 2 chunks"); + for c in &chunks { + assert!(c.chars().count() <= MAX_PROMPT_CHARS); + } + let total: usize = chunks.iter().map(|c| c.chars().count()).sum(); + assert_eq!(total, 15_000, "no characters may be dropped"); + assert_eq!(chunks.concat(), word, "chunks must reconstruct the input"); } #[test] diff --git a/crates/aisix-guardrails/src/text_moderation.rs b/crates/aisix-guardrails/src/text_moderation.rs index f2a7b9b5..9a360050 100644 --- a/crates/aisix-guardrails/src/text_moderation.rs +++ b/crates/aisix-guardrails/src/text_moderation.rs @@ -343,7 +343,7 @@ impl Guardrail for TextModerationGuardrail { ) { return GuardrailVerdict::Allow; } - let text = resp.message.content.clone(); + let text = resp.guardrail_output_text(); if text.is_empty() { return GuardrailVerdict::Allow; } @@ -354,7 +354,8 @@ impl Guardrail for TextModerationGuardrail { } /// Split `text` into chunks of at most `max_chars` characters on -/// whitespace boundaries. A single word over the limit is hard-truncated. +/// whitespace boundaries. A single word over the limit is split into +/// max_chars-sized pieces so the entire token is evaluated (#448). /// (Forked from `prompt_shield::chunk_text`; see the module note.) fn chunk_text(text: &str, max_chars: usize) -> Vec { if text.is_empty() { @@ -373,7 +374,12 @@ fn chunk_text(text: &str, max_chars: usize) -> Vec { chunks.push(std::mem::take(&mut current)); } if word_chars > max_chars { - chunks.push(word.chars().take(max_chars).collect()); + // Split the oversized token fully instead of truncating to + // the prefix, which let the trailing part bypass scanning. + let word_chars_vec: Vec = word.chars().collect(); + for piece in word_chars_vec.chunks(max_chars) { + chunks.push(piece.iter().collect()); + } continue; } } diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index a015cc4d..e3517fa0 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -1053,6 +1053,28 @@ async fn dispatch( match cache.get(key).await { Ok(Some(cached)) => { reservation.commit_tokens(0); + // #448: a cache hit is client-visible output just like a + // fresh upstream response, so it must run output guardrails + // before being returned — not bypass them. + match resolved_chain.check_output(&cached).await { + GuardrailVerdict::Block { reason } => { + tracing::warn!( + guardrail_hook = "output", + model = %req.model, + reason = %reason, + "guardrail blocked cached response", + ); + return Err(with_model(ProxyError::ContentFiltered( + "response blocked by content policy".into(), + ))); + } + GuardrailVerdict::Bypass { reason } => { + if bypass_reason.is_none() { + bypass_reason = Some(reason); + } + } + _ => {} + } let prompt = cached.usage.prompt_tokens as u64; let completion = cached.usage.completion_tokens as u64; let total = cached.usage.total_tokens as u64; diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 5d4455c4..184b18e1 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -243,6 +243,47 @@ async fn dispatch( crate::quota::ModelRateLimit::from_model(&model_name, &model_entry.id, &model_entry.value); let _reservation = crate::quota::enforce(state, auth, Some(&model_rl)).await?; + // #448 (#22): /v1/messages must run input guardrails + the budget + // pre-check like /v1/chat/completions — previously prompts reached the + // upstream without any content/DLP check. Translate the Anthropic- + // shaped body into the internal ChatFormat and run the resolved input + // guardrail chain; a Block short-circuits before dispatch. (Input + // Rewrite/Bypass on this endpoint is not yet applied to the outgoing + // Anthropic body — only Block is enforced here.) + let guardrail_ctx = aisix_guardrails::RequestContext { + model_id: &model_entry.id, + api_key_id: &auth.entry.id, + team_id: auth.key().team_id.as_deref(), + }; + let resolved_chain = state.guardrail_index.resolve(&guardrail_ctx); + if !resolved_chain.is_empty() { + if let Ok(chat) = aisix_provider_anthropic::parse_inbound_request(body) { + if let aisix_guardrails::GuardrailVerdict::Block { reason } = + aisix_guardrails::Guardrail::check_input(&resolved_chain, &chat).await + { + tracing::warn!( + guardrail_hook = "input", + model = %model_name, + reason = %reason, + "guardrail blocked /v1/messages request", + ); + return Err(ProxyError::ContentFiltered( + "request blocked by content policy".into(), + )); + } + } + } + + // Budget pre-check via cp-api (mirrors /v1/chat/completions). + let budget_decision = state.budgets.check(&auth.entry.id).await; + if !budget_decision.allowed { + return Err(ProxyError::BudgetExceeded(Box::new( + budget_decision.reason.unwrap_or_else(|| { + crate::budget::BudgetReason::message_only(auth.entry.id.clone()) + }), + ))); + } + // Resolve the attempt list. For a Model Group (routing model) this // walks `routing.targets` and health-filters them; for a direct // model it's just the model itself. Shared with /v1/chat/completions diff --git a/tests/e2e/src/cases/cache-hit-output-guardrail-e2e.test.ts b/tests/e2e/src/cases/cache-hit-output-guardrail-e2e.test.ts new file mode 100644 index 00000000..fd44b6a0 --- /dev/null +++ b/tests/e2e/src/cases/cache-hit-output-guardrail-e2e.test.ts @@ -0,0 +1,107 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: a non-streaming cache HIT must still run output guardrails (#448). +// Pre-fix, cache hits returned the stored body without any output check, +// so a response cached before an output guardrail existed (or under a key +// without one) could be replayed past a guardrail that should block it. +// +// The mock upstream's canned reply is "mock reply"; we attach an output +// guardrail blocking the literal "reply" AFTER the response is cached, +// then re-issue the identical (cached) request and require it to be +// blocked rather than served from cache. + +const CALLER = "sk-cache-gr-caller"; +const HASH = createHash("sha256").update(CALLER).digest("hex"); +const CACHED_PROMPT = "cache-and-guard-me"; + +describe("cache hit runs output guardrails (#448)", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + upstream = await startOpenAiUpstream(); + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + const pk = await admin.createProviderKey({ + display_name: "cache-gr-pk", + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + await admin.createModel({ + display_name: "cache-gr", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + }); + await admin.createApiKey({ key_hash: HASH, allowed_models: ["cache-gr"] }); + await admin.json("POST", "/admin/v1/cache_policies", { + name: "cache-gr-policy", + enabled: true, + applies_to: "all", + }); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + }); + + const chat = (content: string) => + fetch(`${app!.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { "content-type": "application/json", authorization: `Bearer ${CALLER}` }, + body: JSON.stringify({ model: "cache-gr", messages: [{ role: "user", content }] }), + }); + + test("a response cached before the guardrail is blocked on the cache hit", async (ctx) => { + if (!etcdReachable || !app || !upstream) { + ctx.skip(); + return; + } + + // Wait until model+key+pk+cache are live (clean prompt → 200 + cache miss). + await waitConfigPropagation(async () => (await chat("ready-probe")).ok); + + // 1) Cache the response BEFORE any output guardrail exists. + const first = await chat(CACHED_PROMPT); + expect(first.status, "first request should succeed and populate the cache").toBe(200); + expect(first.headers.get("x-aisix-cache")).toBe("miss"); + + // 2) Attach an output guardrail blocking the canned reply text. + await admin!.json("POST", "/admin/v1/guardrails", { + name: "cache-gr-output-keyword", + enabled: true, + hook_point: "output", + kind: "keyword", + patterns: [{ kind: "literal", value: "reply" }], + }); + + // Gate on guardrail propagation: a FRESH prompt (cache miss) returns + // the canned "mock reply", which the output guardrail must now block. + await waitConfigPropagation(async () => (await chat(`probe-${Math.random()}`)).status === 422); + + // 3) Re-issue the cached request: it is a cache hit, and must now be + // blocked by the output guardrail rather than replayed from cache. + const hit = await chat(CACHED_PROMPT); + expect( + hit.status, + "cache hit must run output guardrails and block the stored reply", + ).toBe(422); + const body = await hit.json(); + expect(JSON.stringify(body)).toContain("content_filter"); + }); +}); diff --git a/tests/e2e/src/cases/messages-input-guardrail-e2e.test.ts b/tests/e2e/src/cases/messages-input-guardrail-e2e.test.ts new file mode 100644 index 00000000..bcd33b0c --- /dev/null +++ b/tests/e2e/src/cases/messages-input-guardrail-e2e.test.ts @@ -0,0 +1,92 @@ +import { createHash } from "node:crypto"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: /v1/messages runs input guardrails (#448 #22). Pre-fix the +// Anthropic /v1/messages path dispatched without any guardrail check, so +// prompts reached the upstream unscanned. The handler now translates the +// body to the internal ChatFormat and runs the resolved input guardrail +// chain before dispatch — a blocked prompt must never hit the upstream. + +const CALLER = "sk-msg-gr-caller"; +const HASH = createHash("sha256").update(CALLER).digest("hex"); +const FORBIDDEN = "forbiddenmsgword"; + +describe("/v1/messages input guardrail (#448)", () => { + let app: SpawnedApp | undefined; + let upstream: OpenAiUpstream | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + upstream = await startOpenAiUpstream(); + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + const pk = await admin.createProviderKey({ + display_name: "msg-gr-pk", + secret: "sk-anth-mock", + api_base: upstream.baseUrl, + }); + await admin.createModel({ + display_name: "msg-gr", + provider: "anthropic", + model_name: "claude-3-5-haiku-20241022", + provider_key_id: pk.id, + }); + await admin.createApiKey({ key_hash: HASH, allowed_models: ["msg-gr"] }); + await admin.json("POST", "/admin/v1/guardrails", { + name: "msg-gr-input-keyword", + enabled: true, + hook_point: "input", + kind: "keyword", + patterns: [{ kind: "literal", value: FORBIDDEN }], + }); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + }); + + const messages = (content: string) => + fetch(`${app!.proxyUrl}/v1/messages`, { + method: "POST", + headers: { "content-type": "application/json", "x-api-key": CALLER }, + body: JSON.stringify({ + model: "msg-gr", + max_tokens: 64, + messages: [{ role: "user", content }], + }), + }); + + test("a forbidden /v1/messages prompt is blocked before hitting upstream", async (ctx) => { + if (!etcdReachable || !app || !upstream) { + ctx.skip(); + return; + } + // Gate on guardrail propagation: the forbidden prompt must be rejected. + await waitConfigPropagation(async () => (await messages(`probe ${FORBIDDEN}`)).status >= 400); + + const hitsBefore = upstream.receivedRequests.length; + const blocked = await messages(`please do ${FORBIDDEN} now`); + expect(blocked.status, "forbidden prompt must be rejected").toBeGreaterThanOrEqual(400); + expect( + upstream.receivedRequests.length, + "blocked prompt must not reach the upstream", + ).toBe(hitsBefore); + + // A benign prompt is not blocked by the input guardrail. + const ok = await messages("hello there"); + expect(ok.status, "benign prompt should not be content-blocked").toBeLessThan(400); + }); +});