fix: broaden guardrail/quota coverage for tool calls, cache, oversized tokens & /v1/messages (#448)

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)]

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;
         }

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")]);

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<String> {
                 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<char> = 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]

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<String> {
     if text.is_empty() {
@@ -373,7 +374,12 @@ fn chunk_text(text: &str, max_chars: usize) -> Vec<String> {
                 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<char> = word.chars().collect();
+                for piece in word_chars_vec.chunks(max_chars) {
+                    chunks.push(piece.iter().collect());
+                }
                 continue;
             }
         }

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;

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

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");
+  });
+});