""" sidecar/gate.py — Concurrent Guard + LLM Fan-Out Gate. The Gate is the heart of the sidecar. It starts the Prompt Guard check and the LLM stream simultaneously, then: • If Guard BLOCKS → cancel the LLM stream, yield a BlockEvent, stop. • If Guard PASSES → open the gate; yield all buffered + live tokens. This means the guard's latency (~50–150ms) is hidden inside the LLM's time-to-first-token (~300–600ms). On the happy path the user effectively sees ZERO added latency from the guard. ┌────────────────────────────────┐ Client ──prompt──►│ Gate.run() │ │ │ │ Task A: guard.screen(prompt) │ ~100ms │ Task B: adapter.stream_chat() │ ~500ms TTFT │ │ │ Guard finishes first? │ │ blocked=True → cancel B │──► BlockEvent │ blocked=False → open gate │──► yield buffered tokens │ │──► yield live tokens… └────────────────────────────────┘ """ import asyncio import logging import time from dataclasses import dataclass, field from typing import AsyncGenerator, List, Optional, Union log = logging.getLogger(__name__) # ── Gate events ─────────────────────────────────────────────────────────────── @dataclass class TokenEvent: """A raw token chunk from the LLM — pass through to client.""" text: str type: str = field(default="token", init=False) @dataclass class BlockEvent: """Guard blocked the prompt — LLM never delivered output to client.""" reason: str threat_score: int flags: List[str] pg_score: float guard_ms: float type: str = field(default="blocked", init=False) GateEvent = Union[TokenEvent, BlockEvent] # ── The Gate ────────────────────────────────────────────────────────────────── class ShieldGate: """ Orchestrates concurrent Guard check + LLM stream with a buffering gate. Args: guard: PromptGuardTextGuard instance. adapter: LLMAdapter with stream_chat(). guard_timeout_sec: Max seconds to wait for guard before passing anyway. (Safety valve — guard should always be < 300ms.) """ def __init__(self, guard, adapter, guard_timeout_sec: float = 3.0): self._guard = guard self._adapter = adapter self._timeout = guard_timeout_sec async def run( self, prompt: str, system_prompt: str = "", trace=None, # Optional[RequestTrace] — for pipeline telemetry ) -> AsyncGenerator[GateEvent, None]: """ Async generator — yields GateEvents to the SSE layer. Yields either a single BlockEvent (if guard fires) or a stream of TokenEvents followed by nothing (caller handles done/close). """ t_start = time.perf_counter() # ── Token buffer — accumulates LLM tokens while guard decides ───── token_buffer: List[str] = [] buffer_lock = asyncio.Lock() first_token_seen = False # ── Task A: Guard (runs in thread pool — it's synchronous) ──────── async def run_guard() -> dict: if trace: trace.on_guard_start() loop = asyncio.get_event_loop() result = await loop.run_in_executor(None, self._guard.screen, prompt) return result # ── Task B: LLM stream — feeds token_buffer until gate opens ────── async def run_stream(): nonlocal first_token_seen if trace: trace.on_llm_start() try: token_count = 0 total_chars = 0 async for token in self._adapter.stream_chat(prompt, system_prompt or None): if not first_token_seen: first_token_seen = True if trace: trace.on_first_token() async with buffer_lock: token_buffer.append(token) token_count += 1 total_chars += len(token) # Emit a token-flow tick every 10 tokens for the dashboard if trace and token_count % 10 == 0: trace.on_token_tick(token_count, total_chars) except asyncio.CancelledError: pass # Guard blocked — stream cancelled cleanly # Start both tasks concurrently guard_task = asyncio.create_task(run_guard(), name="shield-guard") stream_task = asyncio.create_task(run_stream(), name="llm-stream") # ── Wait for guard verdict ───────────────────────────────────────── try: guard_result = await asyncio.wait_for( asyncio.shield(guard_task), timeout=self._timeout, ) except asyncio.TimeoutError: log.error("[Gate] Guard timed out after %.1fs — defaulting to PASS", self._timeout) guard_result = {"blocked": False, "threat_score": 0, "flags": [], "checks": {}} guard_ms = round((time.perf_counter() - t_start) * 1000, 2) pg_score = guard_result.get("checks", {}).get("prompt_guard", {}).get("malicious_score", 0.0) # ── Guard says BLOCK ─────────────────────────────────────────────── if guard_result.get("blocked"): stream_task.cancel() try: await stream_task except (asyncio.CancelledError, Exception): pass log.warning( "[Gate] Prompt BLOCKED — score=%.3f flags=%s latency=%.0fms", pg_score, guard_result.get("flags", []), guard_ms, ) if trace: trace.on_guard_block( pg_score = pg_score, threat_score = guard_result.get("threat_score", 100), reason = guard_result.get("reason", "GUARD_BLOCKED"), flags = guard_result.get("flags", []), ) yield BlockEvent( reason = guard_result.get("reason", "GUARD_BLOCKED"), threat_score = guard_result.get("threat_score", 100), flags = guard_result.get("flags", []), pg_score = pg_score, guard_ms = guard_ms, ) return # ── Guard says PASS — drain buffer then yield live tokens ────────── log.info( "[Gate] Prompt PASSED — score=%.3f latency=%.0fms — opening gate", pg_score, guard_ms, ) if trace: trace.on_guard_pass( pg_score = pg_score, threat_score = guard_result.get("threat_score", 0), flags = guard_result.get("flags", []), ) # Drain anything buffered while guard was deciding async with buffer_lock: buffered = token_buffer.copy() token_buffer.clear() for token in buffered: yield TokenEvent(text=token) # Wait for stream task to complete, yielding live tokens as they arrive # We poll the buffer so we can yield from the async generator context while not stream_task.done(): await asyncio.sleep(0.005) # 5ms poll — tight enough for streaming async with buffer_lock: live = token_buffer.copy() token_buffer.clear() for token in live: yield TokenEvent(text=token) # Final drain after stream task completes async with buffer_lock: for token in token_buffer: yield TokenEvent(text=token) # Propagate any stream errors if not stream_task.cancelled() and stream_task.exception(): log.error("[Gate] LLM stream error: %s", stream_task.exception())