[Idea]: RFC: Ray Serve LLM Governance Middleware Integration

[nagasatish007]

Your Idea

A community contributor (@Divyam19) volunteered on Ray issue #65259 to implement a governance middleware layer for Ray Serve LLM. They proposed TealTiger as the reference integration behind a backend-agnostic LLMMiddleware interface in Ray.

This discussion is to collaborate on the interface design before submitting an RFC/PR to Ray.

Proposed architecture:

1. LLMMiddleware interface (lives in Ray) — typed hooks at the semantic LLM request layer:

   • before_inference(request, context) → request | BlockedResponse
   • after_inference(request, response, context) → response
   • on_stream_chunk(chunk, context) → chunk | None (optional)

2. ray-tealtiger package (lives here) — reference implementation wrapping TealTiger's engine behind the interface. PII scanning, cost budgets, model access policies, audit receipts. Zero-LLM governance path.

Open design questions for the RFC:

• Where should hooks run — ingress/router level, per-model server level, or both?
• How to handle streaming responses (chunk-by-chunk vs full-response audit)?
• Configuration scope — app-level global policies vs per-model LLMConfig.middleware?

Category

Integration

What problem does this solve?

Ray Serve LLM has no built-in governance/guardrails layer. Teams in regulated industries (healthcare, finance, government) currently have to either build custom deployments from scratch or add a separate proxy in front of Ray — losing all benefits of LLMRouter, autoscaling, and multi-model management. A middleware layer lets them use the standard LLMConfig pipeline with governance built in.

Potential Impact

• Any team deploying LLMs on Ray in production (regulated or multi-tenant)
• Ray Serve LLM gets a governance story without building one in-house
• TealTiger gets a high-visibility integration (Ray has 43K+ stars)

Implementation Thoughts

• TealTiger already works as ASGI middleware — adapting to typed LLMMiddleware hooks is straightforward
• We provide: PII patterns, governance policy schema, TEEC receipt format, ENFORCE/MONITOR/OBSERVE modes
• Ray provides: the hook interface, pipeline integration, streaming contract
• Keeps Ray from taking a hard dependency on any vendor

Examples or References

• Ray issue #65259 — original proposal
• Ray Discuss: preprocessing in Ray Serve LLM — community demand
• AG2 integration (PR #3128) — similar pattern: typed middleware interface in the framework, TealTiger as a pluggable extension
• LangChain middleware listing — same approach, framework-native hook + TealTiger adapter

[kindrat86]

Thanks for the invite @nagasatish007 — joining from the cost-enforcement angle. My goal here is to help keep LLMMiddleware generic enough that any governance engine (TealTiger, AgentShield, custom in-house policies) can implement it. Two independent implementations is also a stronger argument for Ray maintainers than one — it proves the abstraction isn't shaped around a single vendor.

Four concrete interface suggestions:

1. RequestContext needs cost-estimation inputs for before_inference

Pre-dispatch budget enforcement only works if the hook can see estimated cost before the model runs. Suggest the context carry:

• model_id (resolved routing target — also catches silent substitution to a pricier variant)
• estimated_input_tokens (Ray already parses/tokenizes at this layer)
• max_tokens from the request (bounds worst-case output cost)
• tenant / session / agent identifiers (for cumulative budget scoping)

With those, an engine can compute worst-case cost = input_tokens × input_price + max_tokens × output_price and block before dispatch. Without them, engines are forced back to reactive post-hoc tracking — which defeats the "cost budgets" item in the original proposal.

2. BlockedResponse should be structured, not just an HTTP status

Suggest: {decision, rule_triggered, reason, severity, retry_after?}. This (a) feeds the audit trail directly, (b) lets clients distinguish 402-style budget exhaustion (don't retry) from 429-style velocity blocks (retry after window), and (c) keeps engines swappable because the decision schema belongs to Ray, not the vendor.

3. Streaming: pre-dispatch is unaffected, but usage reconciliation matters

Budget/velocity checks run in before_inference regardless of streaming. The subtle gap is afterwards: cumulative budgets need actual usage, and estimated ≠ actual. Suggest an on_inference_complete(usage, context) hook — fires after the stream ends with final token counts — distinct from after_inference content scanning. Engines tracking cumulative spend reconcile estimates against actuals there; otherwise long streamed responses under-count against daily/session caps.

4. Config scope + middleware ordering

Agree with app-level + per-model. One addition: chain order should be explicit and deterministic (list order = execution order), so cheap checks (budget rules, sub-ms) can run before expensive ones (PII scanning) — a blocked request should never pay the scanning cost.

Happy to build an AgentShield adapter against whatever interface comes out of this as a second reference engine — the pre-dispatch spend rules (transaction_limit, daily_total, velocity, cascade_cost, merchant_allowlist, category_block) map 1:1 onto before_inference + on_inference_complete. That gives the Ray RFC two independent engines to validate the abstraction against.

[Divyam19]

@kindrat86 — Great feedback. I'm taking all four points as design requirements, not optional.

Summary of what i think needs implementing:

RequestContext with cost-estimation inputs — model_id, estimated_input_tokens, max_tokens, tenant/session identifiers
Structured BlockedResponse — {decision, rule_triggered, reason, severity, retry_after?}
on_inference_complete(usage, context) hook for post-stream usage reconciliation (distinct from after_inference content scanning)
Explicit middleware ordering (list order = execution order) so cheap checks run first

Next steps:

Finalize the interface definition here (you, @nagasatish007, and any Ray maintainers watching)
Once we agree on the hooks, I'll PR the base LLMMiddleware class to Ray + TealTiger adapter
Your AgentShield implementation validates the abstraction — that's exactly what we want

[Divyam19]

@nagasatish007 — Ray just assigned this to @richardliaw at Anyscale. I've posted on the main issue asking about implementation expectations. Will sync here as we move toward the RFC PR.

[kindrat86]

@Divyam19 — appreciate you treating all four as requirements. That summary is exactly the interface I had in mind.

On the plan: finalize interface here → PR base LLMMiddleware to Ray + TealTiger adapter → AgentShield adapter as second reference engine. That sequencing works. Two independent implementations hitting the same hooks is the strongest argument Ray maintainers can see.

What I'll bring when the interface draft lands: AgentShield's 7 rule types all map to before_inference + on_inference_complete without any impedance mismatch:

• transaction_limit, daily_total, session_budget, cascade_cost → before_inference (pre-dispatch cost evaluation against the RequestContext fields you listed)
• velocity → before_inference (rolling-window rate limiting from the context's session/tenant identifier)
• merchant_allowlist, category_block → before_inference (model routing validation via model_id)
• on_inference_complete → reconciliation pass (actual vs. estimated token counts against cumulative budgets)

I'll sketch the agentshield-ray adapter as a concrete code snippet against whatever hook signature you post — not an abstract "yeah this works," but actual class AgentShieldMiddleware(LLMMiddleware) with the engine wired in. That way Ray maintainers reviewing the RFC can see the abstraction exercised from both the PII/content side (TealTiger) and the cost-enforcement side (AgentShield) without either vendor shaping the interface.

Also worth noting: Ray assigned the issue to @richardliaw at Anyscale — that's about as good an assignment signal as we could hope for. The interface needs to be tight when it reaches his desk.

I'll watch this thread. Tag me when the draft hooks are up.

[Divyam19]

cool need to do some R&D on both of your repos. will get back to you guys

[kindrat86]

@Divyam19 — this summary is exactly right. All four points as stated.

On the AgentShield side, the mapping is straightforward — all 7 rule types slot cleanly into two hooks:

before_inference: transaction_limit, daily_total, velocity, merchant_allowlist, category_block — all pre-dispatch spend rules. They're sub-ms checks (pure arithmetic against SQLite counters), so they belong at the front of the chain per the ordering design.

on_inference_complete: cascade_cost reconciliation — adjusts cumulative budgets from estimates to actual token counts after the stream ends.

Happy to sketch the agentshield-ray adapter as a concrete code snippet once the hook signatures are final. The two-engine validation is the strongest argument for Ray maintainers.

@nagasatish007 — your call on when the interface draft is ready for the next review round.

[nagasatish007]

Great progress here. @kindrat86's four design requirements are exactly right — adopting all of them.

On the interface draft, here's where TealTiger maps to the proposed hooks:

Hook	TealTiger capability
before_inference	PII scan (40+ patterns), tool allowlist, model access policy, pre-dispatch cost estimation
after_inference	Output PII/secrets scan, content moderation
on_inference_complete	Actual token usage → reconcile against cumulative budgets, emit TEEC audit receipt
BlockedResponse	Already emit structured decisions: {action, reason_codes, risk_score, decision_id, policy_digest} — maps directly

On ordering: agree completely. TealTiger's policy evaluation is sub-ms (fnmatch + regex), so it belongs early in the chain. PII scanning is the most expensive check (~1ms for 40 patterns) and should run last among our policies. The explicit list-order-equals-execution-order design handles this naturally.

On RequestContext fields: TealTiger already consumes agent_id/session_id/model_id for scoping — mapping these to the context fields @kindrat86 proposed is 1:1.

Next step from my side: once @Divyam19 posts the draft hook signatures, I'll write the ray-tealtiger adapter as a concrete class TealTigerMiddleware(LLMMiddleware) implementation — not pseudocode, actual working code that can be tested against the interface.

The two-engine validation (TealTiger for PII/content, AgentShield for cost enforcement) is the strongest possible signal to @richardliaw that the abstraction is vendor-neutral. Let's make sure both adapters exist before the RFC reaches his desk.

[Divyam19]

@nagasatish007 @kindrat86 — did the R&D. Went through ray/llm/_internal/serve/core/ end to end. Here's where things land.

How it works in practice

Ray already supports custom ingress classes through ingress_cls_config + ingress_extra_kwargs — documented in doc/source/serve/llm/architecture/overview.md as the intended extension point for auth/tracing. So the approach is: subclass OpenAiIngress, override _process_llm_request, inject the middleware chain.

All three hooks live inside _process_llm_request:

• before_inference → before _get_response() (this is where DeploymentHandle.chat.remote fires). Can block the request before it ever touches the model.
• after_inference → after the generator peek, before building JSONResponse / StreamingResponse. For non-streaming, the full response is available to modify. For streaming, this fires on the first chunk.
• on_inference_complete → async generator wrapper, fires after the stream exhausts. This is where @kindrat86's usage reconciliation happens — actual token counts from the final chunk (when stream_options.include_usage=True), not estimates.

Context availability

Mapped the RequestContext fields to actual Ray data sources:

• model_id → body.model — available directly
• request_id → x-request-id header / SetRequestIdMiddleware — available directly
• session_id → x-session-id header — available directly
• max_tokens → body.max_tokens / body.max_completion_tokens — available directly
• user_id → request.state.user_id (auth MW must set it), or body.user, or X-User-Id header — needs convention
• tenant_id → not first-class in Ray, comes from headers/JWT — user-provided
• estimated_input_tokens → not on the request, need to pre-tokenize at ingress (~0.5ms) — estimate only, actuals come in on_inference_complete

Streaming approach

v1 is audit-only. Generator gets wrapped, chunks pass through to the client untouched, we accumulate text + usage as they flow, on_inference_complete fires after exhaustion. No buffering, no TTFT impact.

If someone needs to redact secrets from the response before the client sees them — that requires buffering the entire response, which kills streaming. That's a v2 opt-in with explicit latency tradeoffs documented.

Scope clarification

Want to make sure we're not stepping on each other — I'm handling the Ray-side work: LLMMiddleware base class, GovernanceIngress subclass, hooks, tests, PR to Ray. @nagasatish007 you're building the TealTigerMiddleware adapter. @kindrat86 you're building the AgentShieldMiddleware adapter. That right?

What's next from me

Draft hook signatures this week, will post here for review. Once we're aligned, I'll open a draft PR on Ray.

@kindrat86 — on your question about estimated_input_tokens: at the ingress level we'd have to pre-tokenize to get accurate counts. Ray doesn't expose token counts at that layer — the engine tokenizes internally after the request crosses the DeploymentHandle boundary. Options are (a) estimate from character count, (b) load a tokenizer at the ingress, or (c) accept the estimate gap and rely on on_inference_complete for accurate reconciliation. What's your preference for pre-dispatch checks — rough estimate good enough, or do you need precision?

---

[kindrat86]

@Divyam19 — solid breakdown, thanks for mapping it all out.

On the token estimate question: rough estimate is fine for AgentShield's pre-dispatch. We're a firewall, not an auditor — the job is catching obvious over-spend before it happens, not precise billing. Character-count ÷ 4 gets us close enough for the blocking threshold. Accurate reconciliation in on_inference_complete handles the bookkeeping side.

That said, if someone else (TealTiger?) needs precision at the ingress, option (b) — loading the tokenizer — is the right call. Happy to keep it optional/configurable so each middleware adapter picks its own strategy.

Re: next steps — your plan sounds right. I'll sketch the AgentShield adapter once the hook signatures are up for review. No rush on my side.

[nagasatish007]

@Divyam19 Sounds good. I will work on Tealtiger's end. Go ahead with your plan.

[kindrat86]

@nagasatish007 Done — starred the repo and joined the Discord. Happy to keep contributing as the interface draft takes shape.

[nagasatish007]

@kindrat86 @Divyam19 Please join TealTiger Discord at : https://discord.gg/ZnCEd5Mk2 and award github star too, if that is not a big ask :)

[nagasatish007]

@kindrat86 @Divyam19 Please join TealTiger Discord at : https://discord.gg/ZnCEd5Mk2 and award github star too, if that is not a big ask :)

Thanks @kindrat86 , your discord name is Maryan?

[kindrat86]

Yes — Discord name is Maryan (maryan_33221). Joined and starred 👍

[Divyam19]

@kindrat86 @nagasatish007 I have raised the pr for hooks please go through it once to understand what I have done. Ur review will help. Also I'll try to get it merged, facing a micro check issue rn. On the CI. Working on that.

[kindrat86]

Took a proper look at the PR (ray-project/ray#65514), good work @Divyam19. All four points from this thread landed in the interface: structured BlockedResponse with the THROTTLED → 429 + Retry-After mapping, on_inference_complete kept separate from content scanning, explicit list ordering in the chain, and the cost-estimation fields on RequestContext. Reviewing as promised; three things from reading the diff:

1. Error handling is asymmetric across the chain. MiddlewareChain.before_inference and after_inference log and re-raise, so an engine bug surfaces to the client as an unstructured 500. on_inference_complete logs and swallows. Neither is what a governance deployment wants: in ENFORCE an engine failure should fail closed with a structured governance error, in MONITOR / REPORT_ONLY it should fail open (log, continue). A per-middleware on_error: "fail_open" | "fail_closed" constructor arg on LLMMiddleware, honored by the chain, fixes it cheaply. Sensible defaults: fail_closed for before_inference, fail_open for the other two hooks.

2. estimated_input_tokens is declared but never populated. build_request_context always leaves it None, so today an engine cannot compute worst-case pre-dispatch cost: max_tokens bounds the output side but the input side is missing. Even a cheap chars/4 heuristic would unblock budget engines; exact tokenization clearly does not belong on the hot path. Is populating it in scope for this PR or planned as a follow-up?

3. context.headers carries the full inbound header set, including Authorization and API keys. Exposing it is fine, but governance.md should warn engine authors not to write headers wholesale into audit receipts, or every governance log becomes a secret store.

One more doc suggestion: state explicitly in the streaming section that after_inference sees only the first stream chunk today (buffering/redaction being out of scope), so output PII scanning on streamed responses is best-effort until that lands. Better for engine authors to read that than infer it.

Planning to point the AgentShield adapter at this branch once the interface settles. Anything it can't express will come back here as concrete issues.

[Divyam19]

let me address these issues by tonight. Meanwhile @nagasatish007 kindly give your inputs as well. Also i guess with the current structure you guys can begin planning the adapter infra.

[kindrat86]

Sounds good @Divyam19 — no rush, the review comments are ordered by impact so tackle them in whatever sequence works.

On planning the adapter infra: starting now. First pass at class AgentShieldMiddleware(LLMMiddleware) will run against your branch, not main, so the fail-open/fail-closed semantics from point 1 get exercised against real behavior instead of assumptions. Two things I expect to shake out of it for this thread:

• how a purely pre-dispatch engine (all 7 AgentShield rule types are before_inference checks) composes with an engine that also does content scanning, when both sit in the same ordered chain
• whether BlockedResponse carries enough for a cost engine to emit a useful retry_after without knowing anything about the model backends

Anything the adapter cannot express comes back here as a concrete issue against the PR, which is the whole point of two independent engines hitting the same hooks.

@nagasatish007 — looking forward to your read on the PR. Worth a look on your side: point 3 (headers in audit receipts) interacts directly with the TEEC receipt format, since TealTiger is the engine most likely to be writing receipts in production.

[nagasatish007]

@Divyam19 — went through ray-project/ray#65514. This is solid work. All four design requirements from this thread are represented in the interface, the commit structure is clean (one concern per commit), and the test coverage is thorough — unit tests, mocked-handle ingress, and HTTP e2e. Labels (serve, llm, docs, community-contribution) landed automatically from CODEOWNERS, which means the right people will see it.

Feedback on @kindrat86's three points, from the TealTiger adapter perspective:

Point 1: Error handling — on_error: "fail_open" | "fail_closed"

+1. This is exactly how TealTiger's governance modes work internally:

• ENFORCE = fail closed. If the engine errors, the request is blocked. You don't get to bypass governance because the scanner crashed.
• MONITOR = fail open. Engine errors are logged but the request proceeds. You get observability without breaking production.
• OBSERVE = fail open + silent. Same as MONITOR but no client-visible headers.

The proposed on_error constructor arg maps cleanly to this. For the TealTiger adapter, we'd wire it to the governance mode:

class TealTigerMiddleware(LLMMiddleware):
    def __init__(self, engine, mode="ENFORCE"):
        # ENFORCE → fail_closed, MONITOR/OBSERVE → fail_open
        super().__init__(on_error="fail_closed" if mode == "ENFORCE" else "fail_open")

Agree with kindrat86's defaults: fail_closed for before_inference (safety-critical path), fail_open for after_inference and on_inference_complete (observability paths where blocking after-the-fact adds no safety).

Point 2: estimated_input_tokens — populate with chars/4 heuristic

For TealTiger's pre-dispatch cost checks, chars/4 is good enough. Our cost policy evaluates worst-case: (estimated_input × input_price) + (max_tokens × output_price) against the session budget. The estimate doesn't need to be precise — it needs to be conservative enough to catch budget exhaustion before dispatch. If chars/4 over-estimates by 10-15%, that's fine; on_inference_complete reconciles with actuals afterward.

Suggest: populate estimated_input_tokens with len(request_text) // 4 as the default. Engines that need precision can load a tokenizer; the interface shouldn't mandate that cost.

Point 3: Headers in audit receipts — TEEC format interaction

This is directly relevant to TealTiger. Our TEEC receipt format captures governance evidence for compliance audits (SOC2, HIPAA). The receipt includes:

{
  "decision_id": "uuid",
  "action": "DENY",
  "risk_score": 85,
  "reason_codes": ["PII_DETECTED:ssn"],
  "policy_id": "pii_block_v1",
  "evaluation_time_ms": 0.8,
  "agent_id": "...",
  "session_id": "...",
  "correlation_id": "..."
}

What we do NOT include: raw request headers, auth tokens, API keys. The receipt captures the governance decision, not the request context. We pull identifiers (agent_id, session_id, model_id) from RequestContext fields — not from raw headers.

The risk kindrat86 flagged is real: if an engine naively does receipt.headers = context.headers, every governance log becomes a credential store. TealTiger's adapter will explicitly select only the governance-relevant fields from RequestContext and never persist raw headers.

Recommendation for the governance.md docs: add a section under "Writing a Custom Middleware" that states:

⚠️ Security: context.headers is provided for authentication and routing decisions. Do not persist raw headers in audit logs or external stores — they may contain Authorization tokens, API keys, or session secrets. Extract only the identifiers your engine needs (user_id, tenant_id, session_id) and discard the rest.

This is cheap to add and prevents the inevitable "governance engine leaks API keys to log aggregator" incident.

Additional observation: after_inference streaming semantics

The PR description says after_inference sees "the first stream chunk" for streaming responses. That's correct for v1 (audit-only, no buffering). Worth documenting explicitly in the user guide that output scanning on streamed responses is best-effort in v1 — engines see only the first chunk, not the full response. Full-response scanning requires buffering, which kills TTFT and is a v2 opt-in.

For TealTiger, this means our output PII scanner only works on non-streaming responses in v1. For streaming, we'd rely on on_inference_complete to log a post-hoc audit (the accumulated text is available there after the stream exhausts), but we can't redact — the client already received it. That's an acceptable tradeoff for v1 and should be stated in docs.

---

On the TealTiger adapter — next steps

Interface looks stable enough to start coding against. I'll begin the ray-tealtiger adapter this week:

class TealTigerMiddleware(LLMMiddleware):
    """TealTiger governance engine for Ray Serve LLM."""
    
    async def before_inference(self, request, context):
        # PII scan, tool/model allowlist, cost budget check
        decision = self.engine.evaluate(request, context)
        if decision.action == "DENY":
            return BlockedResponse(
                rule_triggered=decision.reason_codes[0],
                reason=decision.policy_id,
                severity=decision.risk_score,
            )
        return request
    
    async def on_inference_complete(self, usage, context):
        # Reconcile actual tokens against budget, emit TEEC receipt
        self.engine.reconcile(usage, context)
        self.engine.emit_receipt(context)

Will build against Divyam19's branch. Once kindrat86's AgentShield adapter also hits the hooks, we'll have the two-engine validation ready for @richardliaw's review.

@Divyam19 — good momentum. The CI micro check issue — is that the lint step or an actual test failure? Happy to help debug if needed.

---

Summary of action items

Who	Action
@Divyam19	Fix kindrat86's 3 points (error handling, estimated_input_tokens, headers warning in docs)
@nagasatish007	Begin ray-tealtiger adapter against Divyam19's branch
@kindrat86	Begin agentshield-ray adapter against same branch
All	Reconvene once both adapters exercise the hooks — file concrete issues for anything that doesn't fit

[kindrat86]

@nagasatish007: all three points are settled from my side, no changes:

• on_error: agree with the mapping. From AgentShield's cost-enforcement angle the same split holds: budget/velocity rules that BLOCK before dispatch are fail_closed (a request must not bypass a spend cap because the counter crashed), and reconciliation/observability paths are fail_open. The ENFORCE/MONITOR/OBSERVE wiring you sketched is exactly what I'll do in AgentShieldMiddleware.

• estimated_input_tokens = chars/4: good default. For cost enforcement the estimate's job is worst-case budget exhaustion, and over-estimating 10-15% is the safe direction. on_inference_complete reconciles to actuals, which is the part AgentShield actually depends on.

• Headers warning in governance.md: +1, add it as written. AgentShield pulls identifiers from RequestContext only and never persists raw headers, same rationale.

One addition from the cost-enforcement side, on BlockedResponse. TealTiger's decisions are ALLOW/DENY, but AgentShield emits four verdicts: APPROVED, BLOCKED, FLAGGED (medium-severity signal, proceed but flag), and REVIEW (the hitl_threshold rule escalates to a human instead of hard-blocking). So the interface should carry a verdict field (allow / block / flag / review) rather than only a block path; otherwise engines with HITL semantics can't express themselves. Cheap to add now, painful to retrofit after the RFC lands.

On streaming: agree v1 is audit-only, no buffering. For AgentShield this is a non-issue at dispatch (spend checks run in before_inference regardless of stream vs non-stream), and the on_inference_complete reconciliation with accumulated actuals is what keeps cumulative budgets honest on long streams. So the design already covers the cost-enforcement case without needing redaction.

Action item confirmed: I'll start the agentshield-ray adapter against @Divyam19's branch this week, same cadence as your ray-tealtiger adapter, and file concrete issues for anything that doesn't fit the hooks.