ddapm-test-agent 1.36.0__py3-none-any.whl → 1.38.0__py3-none-any.whl

ddapm_test_agent/web.py

@@ -0,0 +1,1523 @@
+import asyncio
+import datetime
+import json
+import logging
+from pathlib import Path
+import time
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Protocol
+from typing import Tuple
+import urllib.parse
+import weakref
+
+from aiohttp import web
+from aiohttp.web import StreamResponse
+from jinja2 import Environment
+from jinja2 import FileSystemLoader
+
+
+log = logging.getLogger(__name__)
+
+
+class BodyProcessor:
+    """Utility class for processing request/response bodies"""
+
+    @staticmethod
+    def process_body(body_data: bytes, content_type: str) -> Tuple[str, bool]:
+        """
+        Process body data for display
+        Returns: (processed_body, is_binary)
+        """
+        if not body_data:
+            return "", False
+
+        content_type = content_type or ""
+
+        # Check if this is binary content that should be base64 encoded
+        if (
+            "msgpack" in content_type.lower()
+            or "application/octet-stream" in content_type.lower()
+            or "multipart" in content_type.lower()
+        ):
+            import base64
+
+            return base64.b64encode(body_data).decode("ascii"), True
+        else:
+            # Try to decode as text
+            try:
+                return body_data.decode("utf-8", errors="replace"), False
+            except Exception:
+                return "[Binary data]", False
+
+
+class TraceProcessor:
+    """Utility class for processing trace data"""
+
+    @staticmethod
+    def process_traces(raw_data: bytes, content_type: str, path: str, suppress_errors: bool = False) -> Dict[str, Any]:
+        """Process trace data and return standardized trace information"""
+        if not raw_data:
+            return TraceProcessor._empty_trace_result(path)
+
+        try:
+            # Import decode functions
+            from .trace import decode_v1 as trace_decode_v1
+            from .trace import decode_v04 as trace_decode_v04
+            from .trace import decode_v05 as trace_decode_v05
+            from .trace import decode_v07 as trace_decode_v07
+
+            # Decode based on path
+            traces = []
+            log.info(f"Decoding trace data for path: {path}, data length: {len(raw_data)}")
+            if path == "/v0.4/traces":
+                traces = trace_decode_v04(content_type, raw_data, suppress_errors)
+            elif path == "/v0.5/traces":
+                traces = trace_decode_v05(raw_data)
+            elif path == "/v0.7/traces":
+                traces = trace_decode_v07(raw_data)
+            elif path == "/v1.0/traces":
+                traces = trace_decode_v1(raw_data)
+            log.info(f"Decoded {len(traces)} traces for {path}")
+
+            return TraceProcessor._process_decoded_traces(traces, path)
+
+        except Exception as e:
+            log.error(f"Failed to process trace data: {e}")
+            return TraceProcessor._empty_trace_result(path)
+
+    @staticmethod
+    def _process_decoded_traces(traces: List[Any], path: str) -> Dict[str, Any]:
+        """Convert decoded traces to standardized format"""
+        # Count traces and spans
+        trace_count = len(traces)
+        span_count = sum(len(trace) for trace in traces)
+
+        # Convert spans to dict format for JSON serialization
+        trace_data = []
+        for trace in traces:
+            trace_spans = []
+            for span in trace:
+                # Handle both dict and object formats
+                if hasattr(span, "trace_id"):
+                    # Span object with attributes
+                    span_dict = {
+                        "trace_id": span.trace_id,
+                        "span_id": span.span_id,
+                        "parent_id": span.parent_id,
+                        "service": span.service,
+                        "name": span.name,
+                        "resource": span.resource,
+                        "type": span.type,
+                        "start": span.start,
+                        "duration": span.duration,
+                        "error": span.error,
+                        "meta": span.meta,
+                        "metrics": span.metrics,
+                    }
+                else:
+                    # Already a dictionary
+                    span_dict = {
+                        "trace_id": span.get("trace_id"),
+                        "span_id": span.get("span_id"),
+                        "parent_id": span.get("parent_id"),
+                        "service": span.get("service"),
+                        "name": span.get("name"),
+                        "resource": span.get("resource"),
+                        "type": span.get("type"),
+                        "start": span.get("start"),
+                        "duration": span.get("duration"),
+                        "error": span.get("error"),
+                        "meta": span.get("meta", {}),
+                        "metrics": span.get("metrics", {}),
+                    }
+                trace_spans.append(span_dict)
+            trace_data.append(trace_spans)
+
+        # Encode trace data as base64 JSON for JavaScript
+        import base64
+        import json
+
+        trace_json = json.dumps(trace_data)
+        trace_data_b64 = base64.b64encode(trace_json.encode("utf-8")).decode("ascii")
+
+        return {
+            "is_trace_request": True,
+            "path": path,
+            "trace_count": trace_count,
+            "span_count": span_count,
+            "traces": trace_data,
+            "trace_data_b64": trace_data_b64,
+        }
+
+    @staticmethod
+    def _empty_trace_result(path: str) -> Dict[str, Any]:
+        """Return empty trace result for error cases"""
+        return {
+            "is_trace_request": True,
+            "path": path,
+            "trace_count": 0,
+            "span_count": 0,
+            "traces": [],
+            "trace_data_b64": "",
+        }
+
+
+class RequestObserver(Protocol):
+    """Observer interface for request notifications"""
+
+    async def notify_request(self, request_data: Dict[str, Any]) -> None:
+        pass
+
+
+class RequestStorage:
+    """Centralized storage for all request/response data"""
+
+    def __init__(self, max_requests: int = 200):
+        self._requests: List[Dict[str, Any]] = []
+        self._max_requests = max_requests
+        self._observers: List[RequestObserver] = []
+
+    def add_request(self, request_data: Dict[str, Any]) -> None:
+        """Add a new request and notify observers"""
+        self._requests.append(request_data)
+
+        # Maintain size limit
+        if len(self._requests) > self._max_requests:
+            self._requests.pop(0)
+
+        # Notify observers asynchronously
+        import asyncio
+
+        for observer in self._observers:
+            try:
+                asyncio.create_task(observer.notify_request(request_data))
+            except Exception:
+                # Don't let observer failures break request processing
+                pass
+
+    def get_all_requests(self) -> List[Dict[str, Any]]:
+        """Get all stored requests (most recent first)"""
+        return list(reversed(self._requests))
+
+    def clear_requests(self) -> None:
+        """Clear all stored requests"""
+        self._requests.clear()
+
+    def add_observer(self, observer: RequestObserver) -> None:
+        """Add an observer for request notifications"""
+        self._observers.append(observer)
+
+    def remove_observer(self, observer: RequestObserver) -> None:
+        """Remove an observer"""
+        if observer in self._observers:
+            self._observers.remove(observer)
+
+    def __len__(self) -> int:
+        return len(self._requests)
+
+
+# Global request storage instance (initialized with default, can be reconfigured)
+request_storage = None
+
+
+@web.middleware
+async def request_response_capture_middleware(request: web.Request, handler: Any) -> Any:
+    """Middleware to capture all request/response data for WebUI"""
+    request_start_time = time.time()
+
+    # Capture request data
+    request_body = b""
+    if request.has_body and request.can_read_body:
+        try:
+            request_body = await request.read()
+            # Store the body back on the request for handlers to use
+            request._payload.feed_eof()
+            request._payload.feed_data(request_body)
+        except Exception as e:
+            log.debug(f"Failed to read request body: {e}")
+
+    # Process the request
+    response_status = 500
+    response_headers = {}
+    response_body = b""
+
+    try:
+        response = await handler(request)
+        response_status = response.status
+        response_headers = dict(response.headers)
+
+        # Capture response body if it's a regular response
+        if hasattr(response, "body") and response.body:
+            response_body = response.body
+        elif hasattr(response, "text") and response.text:
+            response_body = response.text.encode("utf-8")
+
+    except Exception as e:
+        # Handle errors
+        response_status = 500
+        response_headers = {}
+        response_body = str(e).encode("utf-8")
+        # Store data even on error, then re-raise
+        request_data = {
+            "timestamp": request_start_time,
+            "method": request.method,
+            "path": request.path_qs,
+            "headers": dict(request.headers),
+            "content_type": request.headers.get("Content-Type", ""),
+            "remote_addr": request.remote or "",
+            "request_body": request_body,
+            "response": {
+                "status": response_status,
+                "headers": response_headers,
+                "body": response_body,
+            },
+            "duration_ms": (time.time() - request_start_time) * 1000,
+        }
+        if request_storage is not None:
+            request_storage.add_request(request_data)
+        raise
+
+    # Store captured data for successful requests
+    request_data = {
+        "timestamp": request_start_time,
+        "method": request.method,
+        "path": request.path_qs,
+        "headers": dict(request.headers),
+        "content_type": request.headers.get("Content-Type", ""),
+        "remote_addr": request.remote or "",
+        "request_body": request_body,
+        "response": {
+            "status": response_status,
+            "headers": response_headers,
+            "body": response_body,
+        },
+        "duration_ms": (time.time() - request_start_time) * 1000,
+    }
+
+    # Store in unified request storage (with automatic size limit and observer notifications)
+    if request_storage is not None:
+        request_storage.add_request(request_data)
+
+    return response
+
+
+MAX_STORED_REQUESTS = 200
+
+
+class WebUI:
+    """Web UI module for the dd-apm-test-agent"""
+
+    def __init__(self, agent: Any, config: Optional[Dict[str, Any]] = None) -> None:
+        self.agent = agent
+        self.config = config or {}
+
+        # Initialize or reconfigure global request storage with max requests from config
+        global request_storage
+        max_requests = self.config.get("max_requests", 200)
+        if request_storage is None:
+            request_storage = RequestStorage(max_requests=max_requests)
+        else:
+            # Update existing storage limit if needed
+            request_storage._max_requests = max_requests
+
+        # Track SSE connections for real-time updates
+        self._sse_connections: weakref.WeakSet[Any] = weakref.WeakSet()
+
+        # Register as observer for request notifications
+        request_storage.add_observer(self)
+
+        # Set up Jinja2 template environment
+        templates_dir = Path(__file__).parent / "templates"
+        self.jinja_env = Environment(loader=FileSystemLoader(str(templates_dir)), autoescape=True)
+
+        # Add custom filters
+        def timestamp_format(timestamp):
+            """Format timestamp for display"""
+            dt = datetime.datetime.fromtimestamp(timestamp)
+            return dt.strftime("%Y-%m-%d %H:%M:%S")
+
+        self.jinja_env.filters["timestamp_format"] = timestamp_format
+
+    async def notify_request(self, request_data: Dict[str, Any]) -> None:
+        """Request Observer implementation - notify SSE connections of new requests"""
+        if not self._sse_connections:
+            return
+
+        # Process the request data for WebUI display
+        processed_request = self._process_single_request(request_data)
+        if not processed_request:
+            return
+
+        message = json.dumps(
+            {
+                "type": "new_request",
+                "request": processed_request,
+                "total_count": len(request_storage) if request_storage is not None else 0,
+            }
+        )
+
+        # Send to all connected SSE clients
+        dead_connections = []
+        for connection in self._sse_connections:
+            try:
+                await connection.write(f"data: {message}\n\n".encode())
+            except Exception:
+                dead_connections.append(connection)
+
+        # Clean up dead connections
+        for connection in dead_connections:
+            self._sse_connections.discard(connection)
+
+    def get_requests_from_agent(self) -> List[Dict[str, Any]]:
+        """Get processed request data from unified request storage"""
+        processed_requests = []
+
+        # Get all requests from unified storage (already in most recent first order)
+        all_requests = request_storage.get_all_requests()[:MAX_STORED_REQUESTS] if request_storage is not None else []
+
+        for req_data in all_requests:
+            processed_request = self._process_single_request(req_data)
+            if processed_request:
+                processed_requests.append(processed_request)
+
+        return processed_requests
+
+    def _process_single_request(self, req_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """Process a single request data dict into WebUI format"""
+        try:
+            # Process request and response bodies using utility
+            request_body, request_body_is_binary = BodyProcessor.process_body(
+                req_data["request_body"], req_data["content_type"]
+            )
+            response_body, response_body_is_binary = BodyProcessor.process_body(
+                req_data["response"]["body"],
+                req_data["response"]["headers"].get("Content-Type", ""),
+            )
+
+            # Extract query string from path
+            path_parts = req_data["path"].split("?", 1)
+            path = path_parts[0]
+            query_string = path_parts[1] if len(path_parts) > 1 else ""
+
+            # Parse query parameters into a dict for template rendering
+            query_params = {}
+            if query_string:
+                query_params = urllib.parse.parse_qs(query_string, keep_blank_values=True)
+
+            # Check if this is a trace request and process trace data
+            trace_data = None
+            if path in ["/v0.4/traces", "/v0.5/traces", "/v0.7/traces", "/v1.0/traces"]:
+                trace_data = self._process_middleware_trace_data(req_data, path)
+
+            return {
+                "method": req_data["method"],
+                "path": path,
+                "query_string": query_string,
+                "query_params": query_params,
+                "headers": req_data["headers"],
+                "content_type": req_data["content_type"],
+                "content_length": len(req_data["request_body"]) if req_data["request_body"] else 0,
+                "remote_addr": req_data["remote_addr"],
+                "timestamp": req_data["timestamp"],
+                "session_token": req_data["headers"].get("X-Datadog-Test-Session-Token"),
+                "body": request_body,
+                "body_is_binary": request_body_is_binary,
+                "trace_data": trace_data,
+                "response": {
+                    "status": req_data["response"]["status"],
+                    "headers": req_data["response"]["headers"],
+                    "content_type": req_data["response"]["headers"].get("Content-Type", ""),
+                    "body": response_body,
+                    "body_is_binary": response_body_is_binary,
+                },
+                "duration_ms": req_data["duration_ms"],
+            }
+        except Exception as e:
+            log.debug(f"Failed to process captured request: {e}")
+            return None
+
+    def _get_request_body(self, req: Any) -> Any:
+        """Extract and format request body for display"""
+        if "_testagent_data" not in req:
+            return ""
+
+        data = req["_testagent_data"]
+        content_type = req.content_type or ""
+
+        if "msgpack" in content_type.lower():
+            # For binary data, show as base64
+            import base64
+
+            return base64.b64encode(data).decode("ascii")
+        else:
+            # For text data, decode as UTF-8
+            try:
+                return data.decode("utf-8", errors="replace")
+            except (UnicodeDecodeError, AttributeError):
+                return str(data)
+
+    def _get_basic_trace_info(self, req):
+        """Get basic trace info without full decoding for request list"""
+        if not self._is_trace_request(req):
+            return None
+
+        # Just return basic info indicating this is a trace request
+        # The actual decoding will happen when viewing the trace details
+        return {
+            "is_trace_request": True,
+            "path": req.path,
+            "trace_count": "?",  # Unknown until decoded
+            "span_count": "?",  # Unknown until decoded
+            "traces": None,
+            "trace_data_b64": "",  # Will be filled when actually needed
+        }
+
+    def _get_trace_data(self, req):
+        """Extract and process trace data for trace requests"""
+        if not self._is_trace_request(req):
+            return None
+
+        # Get raw data
+        if "_testagent_data" not in req:
+            return None
+
+        raw_data = req["_testagent_data"]
+        content_type = req.content_type or ""
+
+        return TraceProcessor.process_traces(raw_data, content_type, req.path, suppress_errors=False)
+
+    def _process_middleware_trace_data(self, req_data, path):
+        """Process trace data from middleware-captured request data"""
+        raw_data = req_data.get("request_body", b"")
+        content_type = req_data.get("content_type", "")
+
+        return TraceProcessor.process_traces(raw_data, content_type, path, suppress_errors=True)
+
+    def _is_trace_request(self, req: Any) -> bool:
+        """Check if request is a trace request based on path"""
+        trace_paths = ["/v0.4/traces", "/v0.5/traces", "/v0.7/traces", "/v1.0/traces"]
+        return req.path in trace_paths
+
+    def _clean_trace_data_for_json(self, traces):
+        """Clean trace data to ensure JSON serializability"""
+        import json
+
+        def clean_value(obj):
+            if isinstance(obj, (int, float, str, bool)) or obj is None:
+                return obj
+            elif isinstance(obj, bytes):
+                # Convert bytes to base64 string
+                import base64
+
+                try:
+                    return base64.b64encode(obj).decode("ascii")
+                except Exception:
+                    return "[Binary data]"
+            elif isinstance(obj, dict):
+                return {str(k): clean_value(v) for k, v in obj.items()}
+            elif isinstance(obj, (list, tuple)):
+                return [clean_value(item) for item in obj]
+            else:
+                # For any other type, convert to string
+                try:
+                    return str(obj)
+                except Exception:
+                    return "[Unserializable object]"
+
+        try:
+            cleaned = clean_value(traces)
+            # Test JSON serializability
+            json.dumps(cleaned)
+            return cleaned
+        except Exception:
+            # Failed to clean trace data
+            return []
+
+    def make_app(self) -> web.Application:
+        """Create the web UI application"""
+        app = web.Application()
+
+        # Set up routes
+        app.add_routes(
+            [
+                web.get("/", self.handle_dashboard),
+                web.get("/requests", self.handle_requests),
+                web.get("/requests/stream", self.handle_requests_sse),
+                web.post("/requests/clear", self.handle_clear_requests),
+                web.get("/requests/download", self.handle_download_requests),
+                web.get("/traces", self.handle_requests),  # Redirect old traces URL
+                web.get("/traces/{trace_id}", self.handle_trace_detail),
+                web.get("/config", self.handle_config),
+                web.post("/config/create", self.handle_config_create),
+                web.post("/config/update", self.handle_config_update),
+                web.post("/config/create_path", self.handle_config_create_path),
+                web.post("/config/clear", self.handle_config_clear),
+                web.get("/tracerflares", self.handle_tracer_flares),
+                web.post("/tracerflares/start", self.handle_start_flare),
+                web.post("/tracerflares/stop", self.handle_stop_flare),
+                web.get("/tracerflares/download", self.handle_download_tracer_flare),
+                web.get("/snapshots", self.handle_snapshots),
+                web.get("/snapshots/{filename}", self.handle_snapshot_detail),
+                # HTMX endpoints for server-side processing
+                web.post("/api/render-waterfall", self.handle_render_waterfall),
+                web.post("/api/render-json", self.handle_render_json),
+                # Static files
+                web.static("/static", Path(__file__).parent / "static"),
+            ]
+        )
+
+        return app
+
+    async def handle_dashboard(self, request: web.Request) -> web.Response:
+        """Handle dashboard page"""
+        template = self.jinja_env.get_template("dashboard.html")
+
+        # Get server configuration information
+
+        # Get server configuration (not runtime status)
+        enabled_servers = {
+            "web_ui": True,  # If we're serving this page, Web UI is enabled
+            "apm_server": True,  # Always enabled in the main app
+            "otlp_http": True,  # Always enabled
+            "otlp_grpc": True,  # Always enabled
+        }
+
+        # Get configuration from main app config stored in agent
+        main_config = self.config
+        agent_url = main_config.get("agent_url", "")
+        is_proxying = bool(agent_url)
+
+        # Get actual port configuration from main app
+        actual_apm_port = main_config.get("port", 8126)
+        actual_otlp_http_port = main_config.get("otlp_http_port", 4318)
+        actual_otlp_grpc_port = main_config.get("otlp_grpc_port", 4317)
+
+        content = template.render(
+            title="Dashboard",
+            total_requests=len(request_storage) if request_storage is not None else 0,
+            # Server configuration
+            web_ui_port=self.config.get("web_ui_port", 8080),
+            apm_port=actual_apm_port,
+            otlp_http_port=actual_otlp_http_port,
+            otlp_grpc_port=actual_otlp_grpc_port,
+            enabled_servers=enabled_servers,
+            is_proxying=is_proxying,
+            agent_url=agent_url,
+            snapshot_dir=main_config.get("snapshot_dir", "snapshots"),
+            vcr_enabled=bool(main_config.get("vcr_cassettes_directory")),
+            error_responses_disabled=main_config.get("disable_error_responses", False),
+            max_requests=main_config.get("max_requests", 200),
+        )
+        return web.Response(text=content, content_type="text/html")
+
+    async def handle_requests(self, request: web.Request) -> web.Response:
+        """Handle requests page - live view of all requests"""
+        template = self.jinja_env.get_template("requests.html")
+
+        # Get processed requests from agent
+        request_data = self.get_requests_from_agent()
+
+        content = template.render(
+            requests=request_data,
+            total_requests=len(request_storage) if request_storage is not None else 0,
+        )
+        return web.Response(text=content, content_type="text/html")
+
+    async def handle_trace_detail(self, request: web.Request) -> web.Response:
+        """Handle individual trace detail page"""
+        trace_id = int(request.match_info["trace_id"])
+        template = self.jinja_env.get_template("trace_detail.html")
+        try:
+            trace = await self.agent._trace_by_trace_id(trace_id)
+            content = template.render(
+                title=f"Trace {trace_id}",
+                trace_id=trace_id,
+                trace=trace,
+            )
+        except KeyError:
+            content = template.render(
+                title=f"Trace {trace_id}",
+                trace_id=trace_id,
+                trace=None,
+                error="Trace not found",
+            )
+        return web.Response(text=content, content_type="text/html")
+
+    async def handle_config(self, request: web.Request) -> web.Response:
+        """Handle config page"""
+        template = self.jinja_env.get_template("config.html")
+
+        # Get selected token from query parameter
+        selected_token: Optional[str] = request.query.get("token", "")
+        if selected_token == "null" or selected_token == "":
+            selected_token = None
+
+        # Get available session tokens from stored requests
+        session_tokens = set()
+        for req_info in self.get_requests_from_agent():
+            token = req_info.get("session_token")
+            if token:
+                session_tokens.add(token)
+
+        # Always include None (default) option
+        all_tokens = [None] + sorted(session_tokens)
+
+        # Get current remote config data for selected token
+        current_config = {}
+        try:
+            current_config = await self.agent._rc_server.get_config_response(selected_token)
+            if not current_config:
+                current_config = {}
+        except Exception:
+            current_config = {}
+
+        # Get all configs for display
+        all_configs = {}
+        for token in all_tokens:
+            try:
+                token_config = await self.agent._rc_server.get_config_response(token)
+                if token_config:
+                    all_configs[str(token)] = token_config
+            except Exception:
+                pass
+
+        content = template.render(
+            title="Configuration",
+            session_tokens=all_tokens,
+            selected_token=selected_token,
+            current_config_json=json.dumps(current_config, indent=2) if current_config else "{}",
+            config_data=json.dumps(all_configs, indent=2) if all_configs else "{}",
+        )
+        return web.Response(text=content, content_type="text/html")
+
+    async def handle_config_create(self, request: web.Request) -> web.Response:
+        """Handle creating a new remote config response"""
+        try:
+            data = await request.post()
+            token = data.get("token") or None
+            config_data = data.get("config_data", "{}")
+
+            # Parse and validate JSON
+            try:
+                if isinstance(config_data, bytes):
+                    config_data = config_data.decode("utf-8")
+                elif not isinstance(config_data, str):
+                    config_data = str(config_data)
+                parsed_config = json.loads(config_data)
+            except json.JSONDecodeError as e:
+                return web.json_response({"error": f"Invalid JSON: {e}"}, status=400)
+
+            self.agent._rc_server.create_config_response(token, parsed_config)
+            return web.json_response({"status": "success", "message": "Config created successfully"})
+
+        except Exception as e:
+            return web.json_response({"error": str(e)}, status=500)
+
+    async def handle_config_update(self, request: web.Request) -> web.Response:
+        """Handle updating a remote config response"""
+        try:
+            data = await request.post()
+            token = data.get("token") or None
+            config_data = data.get("config_data", "{}")
+
+            # Parse and validate JSON
+            try:
+                if isinstance(config_data, bytes):
+                    config_data = config_data.decode("utf-8")
+                elif not isinstance(config_data, str):
+                    config_data = str(config_data)
+                parsed_config = json.loads(config_data)
+            except json.JSONDecodeError as e:
+                return web.json_response({"error": f"Invalid JSON: {e}"}, status=400)
+
+            self.agent._rc_server.update_config_response(token, parsed_config)
+            return web.json_response({"status": "success", "message": "Config updated successfully"})
+
+        except Exception as e:
+            return web.json_response({"error": str(e)}, status=500)
+
+    async def handle_config_create_path(self, request: web.Request) -> web.Response:
+        """Handle creating a remote config path response"""
+        try:
+            data = await request.post()
+            token = data.get("token") or None
+            path = data.get("path", "")
+            message = data.get("message", "{}")
+
+            if not path:
+                return web.json_response({"error": "Path is required"}, status=400)
+
+            # Validate and parse message JSON
+            try:
+                if isinstance(message, bytes):
+                    message = message.decode("utf-8")
+                elif not isinstance(message, str):
+                    message = str(message)
+                parsed_message = json.loads(message)
+            except json.JSONDecodeError as e:
+                return web.json_response({"error": f"Invalid message JSON: {e}"}, status=400)
+
+            self.agent._rc_server.create_config_path_response(token, path, parsed_message)
+            return web.json_response({"status": "success", "message": "Config path created successfully"})
+
+        except Exception as e:
+            return web.json_response({"error": str(e)}, status=500)
+
+    async def handle_config_clear(self, request: web.Request) -> web.Response:
+        """Handle clearing remote config for a token"""
+        try:
+            data = await request.post()
+            token = data.get("token") or None
+
+            # Clear by creating empty response
+            self.agent._rc_server.create_config_response(token, {})
+            return web.json_response({"status": "success", "message": "Config cleared successfully"})
+
+        except Exception as e:
+            return web.json_response({"error": str(e)}, status=500)
+
+    async def handle_tracer_flares(self, request: web.Request) -> web.Response:
+        """Handle tracer flares page"""
+        template = self.jinja_env.get_template("tracer_flares.html")
+
+        # Get ALL tracer flares regardless of session token
+        tracer_flares = []
+        try:
+            # Get all flares by getting flares for None (default session) and then all other sessions
+            all_flares = []
+
+            # Get flares for default session (None)
+            default_flares = await self.agent._tracerflares_by_session(None)
+            for flare in default_flares:
+                flare["session_token"] = None  # Mark session token for display
+            all_flares.extend(default_flares)
+
+            # Get flares for all other session tokens
+            session_tokens = set()
+            for req_data in request_storage.get_all_requests() if request_storage is not None else []:
+                token = req_data.get("headers", {}).get("X-Datadog-Test-Session-Token")
+                if token:
+                    session_tokens.add(token)
+
+            for token in session_tokens:
+                token_flares = await self.agent._tracerflares_by_session(token)
+                for flare in token_flares:
+                    flare["session_token"] = token  # Mark session token for display
+                all_flares.extend(token_flares)
+
+            tracer_flares = all_flares
+        except Exception as e:
+            print(f"Error getting tracer flares: {e}")
+
+        # Check for active flare - check all active flares
+        active_flares = getattr(self.agent, "_active_flares", {})
+        active_flare = None
+        # Find any active flare (there should only be one at a time)
+        for flare_info in active_flares.values():
+            active_flare = flare_info
+            break
+
+        content = template.render(
+            title="Tracer Flares",
+            tracer_flares=tracer_flares,
+            total_flares=len(tracer_flares),
+            active_flare=active_flare,
+        )
+        return web.Response(text=content, content_type="text/html")
+
+    async def handle_start_flare(self, request: web.Request) -> web.Response:
+        """Handle starting a tracer flare collection"""
+        try:
+            data = await request.post()
+            session_token = data.get("token") or None
+
+            # Generate unique UUID for this flare request
+            import uuid
+
+            flare_uuid = str(uuid.uuid4())
+
+            # Step 1: Send AGENT_CONFIG to enable debug logging
+            debug_config = {"config": {"log_level": "debug"}}
+            self.agent._rc_server.create_config_path_response(
+                session_token, "datadog/2/AGENT_CONFIG/flare_debug/config", debug_config
+            )
+
+            # Store the active flare info
+            if not hasattr(self.agent, "_active_flares"):
+                self.agent._active_flares = {}
+
+            import time
+
+            self.agent._active_flares[session_token or "default"] = {
+                "uuid": flare_uuid,
+                "start_time": time.time(),
+                "case_id": str(int(time.time())),  # Pure numeric timestamp
+                "hostname": "test-agent",
+                "user_handle": "test@example.com",
+            }
+
+            return web.json_response(
+                {
+                    "status": "success",
+                    "message": "Tracer flare collection started",
+                    "uuid": flare_uuid,
+                }
+            )
+
+        except Exception as e:
+            return web.json_response({"error": str(e)}, status=500)
+
+    async def handle_stop_flare(self, request: web.Request) -> web.Response:
+        """Handle stopping a tracer flare collection and triggering upload"""
+        try:
+            data = await request.post()
+            session_token = data.get("token") or None
+
+            # Check if there's an active flare
+            if not hasattr(self.agent, "_active_flares"):
+                return web.json_response({"error": "No active flare collection"}, status=400)
+
+            active_flares = getattr(self.agent, "_active_flares", {})
+            flare_key = session_token or "default"
+
+            if flare_key not in active_flares:
+                return web.json_response({"error": "No active flare for this session"}, status=400)
+
+            flare_info = active_flares[flare_key]
+
+            # Step 2: Send AGENT_TASK to trigger flare upload
+            task_config = {
+                "task_type": "tracer_flare",
+                "uuid": flare_info["uuid"],
+                "args": {
+                    "case_id": flare_info["case_id"],
+                    "hostname": flare_info["hostname"],
+                    "user_handle": flare_info["user_handle"],
+                },
+            }
+
+            self.agent._rc_server.create_config_path_response(
+                session_token, "datadog/2/AGENT_TASK/flare_upload/config", task_config
+            )
+
+            # Remove from active flares
+            del active_flares[flare_key]
+
+            return web.json_response(
+                {
+                    "status": "success",
+                    "message": "Tracer flare upload triggered",
+                    "case_id": flare_info["case_id"],
+                }
+            )
+
+        except Exception as e:
+            return web.json_response({"error": str(e)}, status=500)
+
+    async def handle_download_tracer_flare(self, request: web.Request) -> web.Response:
+        """Handle downloading a tracer flare ZIP file"""
+        try:
+            case_id = request.query.get("case_id")
+
+            if not case_id:
+                return web.json_response({"error": "case_id parameter required"}, status=400)
+
+            # Get ALL tracer flares from all sessions
+            all_flares = []
+
+            # Get flares for default session (None)
+            default_flares = await self.agent._tracerflares_by_session(None)
+            all_flares.extend(default_flares)
+
+            # Get flares for all other session tokens
+            session_tokens = set()
+            for req_data in request_storage.get_all_requests() if request_storage is not None else []:
+                token = req_data.get("headers", {}).get("X-Datadog-Test-Session-Token")
+                if token:
+                    session_tokens.add(token)
+
+            for token in session_tokens:
+                token_flares = await self.agent._tracerflares_by_session(token)
+                all_flares.extend(token_flares)
+
+            # Debug logging
+            # Looking for case_id in flares
+            # Found flares
+            # Find the flare with matching case_id
+            target_flare = None
+            for flare in all_flares:
+                if flare.get("case_id") == case_id:
+                    target_flare = flare
+                    break
+
+            if not target_flare:
+                # Create detailed error message with available case_ids
+                available_cases = [f"'{flare.get('case_id')}'" for flare in all_flares]
+                error_msg = f"Flare not found. Looking for '{case_id}', available: {available_cases}"
+                # Error processing flare
+                return web.json_response({"error": error_msg}, status=404)
+
+            if "flare_file" not in target_flare:
+                return web.json_response({"error": "No flare file in this flare"}, status=404)
+
+            try:
+                # Decode base64 flare file data
+                import base64
+
+                flare_data = base64.b64decode(target_flare["flare_file"])
+
+                # Return as ZIP file download
+                return web.Response(
+                    body=flare_data,
+                    content_type="application/zip",
+                    headers={"Content-Disposition": f'attachment; filename="tracer_flare_{case_id}.zip"'},
+                )
+            except Exception as decode_error:
+                return web.json_response({"error": f"Error decoding flare file: {decode_error}"}, status=500)
+
+        except Exception as e:
+            return web.json_response({"error": str(e)}, status=500)
+
+    async def handle_snapshots(self, request: web.Request) -> web.Response:
+        """Handle snapshots page"""
+        template = self.jinja_env.get_template("snapshots.html")
+
+        # Get snapshot directory from main app config
+        main_config = self.config
+        snapshot_dir = main_config.get("snapshot_dir", "snapshots")
+        snapshots = []
+
+        try:
+            snapshot_path = Path(snapshot_dir)
+            if snapshot_path.exists() and snapshot_path.is_dir():
+                # Get all JSON files in the snapshot directory
+                json_files = list(snapshot_path.glob("*.json"))
+                json_files.sort()  # Sort alphabetically
+
+                for file_path in json_files:
+                    try:
+                        stat = file_path.stat()
+                        snapshots.append(
+                            {
+                                "filename": file_path.name,
+                                "size": stat.st_size,
+                                "modified": stat.st_mtime,
+                            }
+                        )
+                    except (OSError, IOError):
+                        # Skip files we can't stat
+                        continue
+        except Exception as e:
+            # Handle directory access errors
+            error_msg: Optional[str] = f"Error accessing snapshot directory: {e}"
+        else:
+            error_msg = None
+
+        content = template.render(
+            title="Snapshots",
+            snapshots=snapshots,
+            snapshot_dir=snapshot_dir,
+            error=error_msg,
+        )
+        return web.Response(text=content, content_type="text/html")
+
+    async def handle_snapshot_detail(self, request: web.Request) -> web.Response:
+        """Handle individual snapshot detail page"""
+        filename = request.match_info["filename"]
+        template = self.jinja_env.get_template("snapshot_detail.html")
+
+        # Get snapshot directory from main app config
+        main_config = self.config
+        snapshot_dir = main_config.get("snapshot_dir", "snapshots")
+
+        try:
+            snapshot_path = Path(snapshot_dir) / filename
+
+            # Security check: ensure the file is within the snapshot directory
+            snapshot_path = snapshot_path.resolve()
+            allowed_dir = Path(snapshot_dir).resolve()
+            if not str(snapshot_path).startswith(str(allowed_dir)):
+                raise ValueError("File path not allowed")
+
+            if not snapshot_path.exists() or not snapshot_path.is_file():
+                raise FileNotFoundError("Snapshot file not found")
+
+            # Read and parse JSON content
+            with open(snapshot_path, "r", encoding="utf-8") as f:
+                raw_content = f.read()
+
+            try:
+                parsed_data = json.loads(raw_content)
+                # Pretty-format the JSON
+                formatted_content = json.dumps(parsed_data, indent=2, ensure_ascii=False)
+                is_valid_json = True
+                parse_error: Optional[str] = None
+
+                # Count traces and spans
+                trace_count = 0
+                span_count = 0
+                if isinstance(parsed_data, list):
+                    trace_count = len(parsed_data)
+                    for trace in parsed_data:
+                        if isinstance(trace, list):
+                            span_count += len(trace)
+                        elif isinstance(trace, dict):
+                            span_count += 1
+                elif isinstance(parsed_data, dict):
+                    # Single trace case
+                    trace_count = 1
+                    span_count = 1
+
+            except json.JSONDecodeError as e:
+                formatted_content = raw_content
+                is_valid_json = False
+                parse_error = str(e)
+                trace_count = 0
+                span_count = 0
+            else:
+                parse_error = None
+
+            # Get file stats
+            stat = snapshot_path.stat()
+            file_info = {
+                "filename": filename,
+                "size": stat.st_size,
+                "modified": stat.st_mtime,
+            }
+
+            content = template.render(
+                title=f"Snapshot: {filename}",
+                filename=filename,
+                file_info=file_info,
+                raw_content=formatted_content,
+                is_valid_json=is_valid_json,
+                parse_error=parse_error,
+                trace_count=trace_count,
+                span_count=span_count,
+                trace_data=json.dumps(parsed_data) if is_valid_json else None,
+                error=None,
+            )
+
+        except (FileNotFoundError, ValueError, OSError, IOError) as e:
+            content = template.render(
+                title=f"Snapshot: {filename}",
+                filename=filename,
+                file_info=None,
+                raw_content=None,
+                is_valid_json=False,
+                parse_error=None,
+                error=f"Error loading snapshot: {e}",
+            )
+
+        return web.Response(text=content, content_type="text/html")
+
+    async def handle_requests_sse(self, request: web.Request) -> StreamResponse:
+        """Handle Server-Sent Events for real-time request updates"""
+        response = StreamResponse(
+            status=200,
+            reason="OK",
+            headers={
+                "Content-Type": "text/event-stream",
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "Access-Control-Allow-Origin": "*",
+            },
+        )
+
+        await response.prepare(request)
+
+        # Add this connection to our set
+        self._sse_connections.add(response)
+
+        try:
+            # Send initial request count from unified storage
+            last_count = len(request_storage) if request_storage is not None else 0
+            await response.write(f"data: {json.dumps({'type': 'count', 'count': last_count})}\n\n".encode())
+
+            # Keep connection alive with heartbeats
+            # New requests are now sent via notify_new_request from middleware
+            while True:
+                await asyncio.sleep(5)  # Heartbeat every 5 seconds
+                # Send heartbeat (will raise exception if connection is closed)
+                await response.write(f"data: {json.dumps({'type': 'heartbeat'})}\n\n".encode())
+
+        except (ConnectionResetError, asyncio.CancelledError, Exception):
+            pass
+        finally:
+            # Clean up connection
+            self._sse_connections.discard(response)
+
+        return response
+
+    async def notify_new_request(self, request_info: Dict[str, Any]) -> None:
+        """Notify all SSE connections about a new request"""
+        if not self._sse_connections:
+            return
+
+        # Add trace data processing for streaming requests if this is a trace request
+        trace_data = None
+        if request_info.get("path") in [
+            "/v0.4/traces",
+            "/v0.5/traces",
+            "/v0.7/traces",
+            "/v1.0/traces",
+        ]:
+            try:
+                import base64
+                import json
+
+                # Parse trace data using the proper decoder based on path
+                from .trace import decode_v1
+                from .trace import decode_v04
+                from .trace import decode_v05
+                from .trace import decode_v07
+
+                raw_body = request_info.get("body", "")
+                content_type = request_info.get("content_type", "")
+                path = request_info.get("path", "")
+
+                if raw_body:
+                    # Decode from base64 if needed (msgpack requests are base64 encoded in agent.py)
+                    if "msgpack" in content_type.lower():
+                        binary_data = base64.b64decode(raw_body)
+                    else:
+                        binary_data = raw_body.encode("utf-8")
+
+                    # Use the appropriate decoder based on path
+                    if path == "/v0.4/traces":
+                        traces = decode_v04(content_type, binary_data, False)
+                    elif path == "/v0.5/traces":
+                        traces = decode_v05(binary_data)
+                    elif path == "/v0.7/traces":
+                        traces = decode_v07(binary_data)
+                    elif path == "/v1.0/traces":
+                        traces = decode_v1(binary_data)
+                    else:
+                        traces = []
+
+                    # Decoded trace chunks
+                else:
+                    traces = None
+
+                if traces:
+                    # Count traces and spans using the same logic as the static version
+                    trace_count = len(traces)
+                    span_count = sum(len(trace) for trace in traces)
+
+                    # Clean the data for JSON serialization
+                    clean_traces = self._clean_trace_data_for_json(traces)
+
+                    # Cleaned trace data
+
+                    # Create base64-encoded version for safe HTML transport
+                    trace_data_json = json.dumps(
+                        {
+                            "traces": clean_traces,
+                            "trace_count": trace_count,
+                            "span_count": span_count,
+                        }
+                    )
+                    trace_data_b64 = base64.b64encode(trace_data_json.encode("utf-8")).decode("ascii")
+
+                    trace_data = {
+                        "traces": clean_traces,
+                        "trace_count": trace_count,
+                        "span_count": span_count,
+                        "trace_data_b64": trace_data_b64,
+                    }
+            except Exception:
+                # Failed to process streaming trace data
+                trace_data = None
+
+        # Add trace_data to request_info if present
+        if trace_data:
+            request_info = {**request_info, "trace_data": trace_data}
+
+        message = json.dumps(
+            {
+                "type": "new_request",
+                "request": request_info,
+                "total_count": len(request_storage) if request_storage is not None else 0,
+            }
+        )
+
+        # Send to all connected clients
+        dead_connections = []
+        for connection in self._sse_connections:
+            try:
+                await connection.write(f"data: {message}\n\n".encode())
+            except Exception:
+                # Connection is dead, mark for removal
+                dead_connections.append(connection)
+
+        # Clean up dead connections
+        for connection in dead_connections:
+            self._sse_connections.discard(connection)
+
+    async def notify_latest_request(self) -> None:
+        """Notify all SSE connections about the latest request"""
+        if not self._sse_connections:
+            return
+
+        try:
+            # Get the latest processed request
+            processed_requests = self.get_requests_from_agent()
+            if not processed_requests:
+                return
+
+            latest_request = processed_requests[0]  # Most recent request
+
+            message = json.dumps(
+                {
+                    "type": "new_request",
+                    "request": latest_request,
+                    "total_count": len(request_storage) if request_storage is not None else 0,
+                }
+            )
+
+            # Send to all connected clients
+            dead_connections = []
+            for connection in self._sse_connections:
+                try:
+                    await connection.write(f"data: {message}\n\n".encode())
+                except Exception:
+                    # Connection is dead, mark for removal
+                    dead_connections.append(connection)
+
+            # Clean up dead connections
+            for connection in dead_connections:
+                self._sse_connections.discard(connection)
+
+        except Exception as e:
+            log.error(f"Error in notify_latest_request: {e}")
+
+    async def handle_clear_requests(self, request: web.Request) -> web.Response:
+        """Handle clearing all stored requests"""
+        # Clear unified request storage
+        if request_storage is not None:
+            request_storage.clear_requests()
+        return web.json_response({"status": "success", "message": "All requests cleared"})
+
+    async def handle_download_requests(self, request: web.Request) -> web.Response:
+        """Handle downloading all current requests as JSON"""
+        # Use the already processed requests from unified storage
+        requests_data = self.get_requests_from_agent()
+
+        # Create JSON response with proper headers for download
+        json_content = json.dumps(requests_data, indent=2)
+        return web.Response(
+            text=json_content,
+            content_type="application/json",
+            headers={
+                "Content-Disposition": "attachment; filename=requests.json",
+                "Content-Length": str(len(json_content.encode("utf-8"))),
+            },
+        )
+
+    async def handle_render_waterfall(self, request):
+        """HTMX endpoint: Render waterfall view for trace data"""
+        try:
+            # Get request data from POST body
+            req_data = await request.json()
+
+            # Extract trace data
+            trace_data = req_data.get("trace_data")
+            if not trace_data or not isinstance(trace_data, dict):
+                return web.Response(
+                    text='<div class="empty-state">No trace data available</div>',
+                    content_type="text/html",
+                )
+
+            # Use existing trace processing
+            traces = trace_data.get("traces", [])
+            if not traces:
+                return web.Response(
+                    text='<div class="empty-state">No trace data available for waterfall view</div>',
+                    content_type="text/html",
+                )
+
+            # Render waterfall HTML server-side
+            html = self._render_waterfall_html(traces)
+            return web.Response(text=html, content_type="text/html")
+
+        except Exception as e:
+            log.error(f"Error rendering waterfall: {e}")
+            return web.Response(
+                text=f'<div class="error-message">Error generating waterfall view: {str(e)}</div>',
+                content_type="text/html",
+            )
+
+    async def handle_render_json(self, request):
+        """HTMX endpoint: Render pretty-printed JSON for request body"""
+        try:
+            # Get request data from POST body
+            req_data = await request.json()
+
+            # Extract body data and content type
+            body_data = req_data.get("body_data", "")
+            content_type = req_data.get("content_type", "")
+
+            if not body_data:
+                return web.Response(
+                    text='<div class="empty-state">No body data</div>',
+                    content_type="text/html",
+                )
+
+            # Process the body data
+            processed_body, is_binary = BodyProcessor.process_body(
+                body_data.encode("utf-8") if isinstance(body_data, str) else body_data,
+                content_type,
+            )
+
+            if is_binary:
+                html = f'<pre class="json-display binary">{processed_body}</pre>'
+            else:
+                # Try to pretty-print as JSON
+                try:
+                    parsed_json = json.loads(processed_body)
+                    pretty_json = json.dumps(parsed_json, indent=2)
+                    html = f'<pre class="json-display">{pretty_json}</pre>'
+                except json.JSONDecodeError:
+                    # Not JSON, display as text
+                    html = f'<pre class="json-display text">{processed_body}</pre>'
+
+            return web.Response(text=html, content_type="text/html")
+
+        except Exception as e:
+            log.error(f"Error rendering JSON: {e}")
+            return web.Response(
+                text=f'<div class="error-message">Error processing body data: {str(e)}</div>',
+                content_type="text/html",
+            )
+
+    def _render_waterfall_html(self, traces: List[Any]) -> str:
+        """Render waterfall HTML from trace data"""
+        if not traces:
+            return '<div class="empty-state">No trace data available</div>'
+
+        html = '<div class="waterfall-traces">'
+
+        for trace_index, trace in enumerate(traces):
+            if not trace or not isinstance(trace, list):
+                continue
+
+            # Process spans
+            valid_spans = [span for span in trace if span and isinstance(span, dict)]
+            if not valid_spans:
+                continue
+
+            # Calculate timing
+            min_start = min(span.get("start", 0) for span in valid_spans)
+            max_end = max(span.get("start", 0) + span.get("duration", 0) for span in valid_spans)
+            total_duration = max_end - min_start
+
+            if total_duration == 0:
+                total_duration = max(span.get("duration", 0) for span in valid_spans)
+
+            trace_id = valid_spans[0].get("trace_id", trace_index)
+
+            html += f"""
+            <div class="waterfall-trace">
+                <div class="trace-header">
+                    <div class="trace-header-left">
+                        <h4>Trace {trace_id}</h4>
+                    </div>
+                    <span class="trace-duration">{self._format_duration(total_duration)}</span>
+                </div>
+                <div class="spans-timeline">
+            """
+
+            # Sort spans and render
+            sorted_spans = sorted(valid_spans, key=lambda s: s.get("start", 0))
+            span_hierarchy = self._build_span_hierarchy(sorted_spans)
+
+            for span_info in span_hierarchy:
+                html += self._render_span_html(span_info, min_start, total_duration, 0)
+
+            html += "</div></div>"
+
+        html += "</div>"
+        return html
+
+    def _build_span_hierarchy(self, spans: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Build hierarchical span structure"""
+        # Create span lookup
+        span_lookup = {span.get("span_id"): span for span in spans}
+        root_spans = []
+
+        for span in spans:
+            parent_id = span.get("parent_id")
+            if not parent_id or parent_id not in span_lookup:
+                # Root span
+                root_spans.append({"span": span, "children": self._get_children(span, spans)})
+
+        return root_spans
+
+    def _get_children(self, parent_span: Dict[str, Any], all_spans: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Get child spans recursively"""
+        parent_id = parent_span.get("span_id")
+        children = []
+
+        for span in all_spans:
+            if span.get("parent_id") == parent_id and span.get("span_id") != parent_id:
+                children.append({"span": span, "children": self._get_children(span, all_spans)})
+
+        return sorted(children, key=lambda c: c["span"].get("start", 0))
+
+    def _render_span_html(self, span_info: Dict[str, Any], min_start: int, total_duration: int, depth: int) -> str:
+        """Render HTML for a single span and its children"""
+        span = span_info["span"]
+        children = span_info.get("children", [])
+
+        # Calculate positioning
+        start_time = span.get("start", 0)
+        duration = span.get("duration", 0)
+
+        if total_duration > 0:
+            left_percent = ((start_time - min_start) / total_duration) * 100
+            width_percent = (duration / total_duration) * 100
+        else:
+            left_percent = 0
+            width_percent = 100
+
+        # Generate span classes
+        span_class = "waterfall-span"
+        if span.get("error"):
+            span_class += " error"
+
+        service = span.get("service", "unknown")
+        operation = span.get("name", "unknown")
+
+        indent_style = f"margin-left: {depth * 20}px;"
+        bar_style = f"left: {left_percent}%; width: {max(width_percent, 0.5)}%;"
+
+        html = f"""
+        <div class="{span_class}" style="{indent_style}">
+            <div class="span-info">
+                <span class="service-name">{service}</span>
+                <span class="operation-name">{operation}</span>
+                <span class="span-duration">{self._format_duration(duration)}</span>
+            </div>
+            <div class="span-bar" style="{bar_style}"></div>
+        </div>
+        """
+
+        # Add children
+        for child in children:
+            html += self._render_span_html(child, min_start, total_duration, depth + 1)
+
+        return html
+
+    def _format_duration(self, duration_ns: int) -> str:
+        """Format duration from nanoseconds to human readable"""
+        if duration_ns < 1000:
+            return f"{duration_ns}ns"
+        elif duration_ns < 1000000:
+            return f"{duration_ns / 1000:.1f}μs"
+        elif duration_ns < 1000000000:
+            return f"{duration_ns / 1000000:.1f}ms"
+        else:
+            return f"{duration_ns / 1000000000:.2f}s"