sidecar/gate.py

"""
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())