| # Phase 2: Streaming Stabilization & Observability - PRD & MCP |
|
|
| ## 1. Product Requirements Document (PRD) |
|
|
| ### 1.1 Objective |
| The goal of Phase 2 is to harden the "Zero-Buffer Streaming Proxy" architecture implemented in Phase 1. This phase focuses on **reliability**, **observability**, and **error resilience**. We aim to ensure the system can handle high concurrency, network instability, and malformed upstream responses without crashing or leaking resources, while providing deep visibility into the streaming pipeline. |
|
|
| ### 1.2 Functional Requirements |
|
|
| #### 1.2.1 Advanced Request Logging |
| * **Streaming Support:** The logger must support true streaming for *both* request and response bodies. The current `[]byte` buffer for request bodies must be replaced with a stream-aware interface (`io.Reader`). |
| * **Sanitization:** Sensitive data (Thinking blocks, Tool arguments) must be redacted in real-time with zero-latency overhead. |
| * **Format:** Logs should optionally support **NDJSON** (Newline Delimited JSON) to facilitate machine parsing and ingestion into observability platforms. |
|
|
| #### 1.2.2 AMP Response Rewriter Resilience |
| * **Edge Case Handling:** The rewriter must gracefully handle: |
| * Split JSON tokens across chunk boundaries (already implemented, needs verification). |
| * Invalid or malformed JSON from upstream. |
| * Mixed content types (e.g., error responses sent as plain text instead of SSE). |
| * **Fallback Strategy:** If rewriting fails (e.g., parsing error), the proxy must fallback to passing the raw chunk through to avoid disrupting the client, logging the error asynchronously. |
|
|
| #### 1.2.3 Proxy Resilience & Timeouts |
| * **Context Management:** Request context must be propagated correctly. Client disconnection must immediately cancel the upstream request to save costs. |
| * **Timeouts:** The proxy must enforce explicit timeouts: |
| * **Connect Timeout:** Max 10s. |
| * **Header Timeout:** Max 30s (TTFB). |
| * **Idle Timeout:** Max 60s (for SSE streams). |
|
|
| ### 1.3 Non-Functional Requirements |
|
|
| * **Performance:** |
| * **Latency Overhead:** < 5ms added by the proxy layer (decompression + rewriting + logging). |
| * **Memory Usage:** Constant memory usage per request (O(1)), independent of response size. Target: < 64KB overhead per active stream. |
| * **Concurrency:** Support 1000+ concurrent streaming connections on a standard instance without OOM. |
| * **Observability:** Expose Prometheus metrics for: |
| * Active streams. |
| * Log queue depth. |
| * Sanitization hit rate. |
| * Upstream latency histograms. |
|
|
| --- |
|
|
| ## 2. Master Control Plan (MCP) |
|
|
| ### 2.1 Architecture Review |
| The Phase 1 refactor successfully removed full-body buffering from the Response path. However, the **Request path** still buffers the full body in memory (`RequestLogger` interface takes `body []byte`). Additionally, the current implementation lacks comprehensive error handling for network interruptions during streaming and doesn't expose internal metrics. |
|
|
| **Refinement Areas:** |
| * **Logging Interface:** Refactor `RequestLogger` to accept `io.Reader` for the request body. |
| * **Async Safety:** Ensure the `eventLoop` in the logger handles channel overflows gracefully (drop strategy vs. block strategy). |
| * **Transport Configuration:** The `httputil.ReverseProxy` needs a custom `Transport` with tuned timeouts. |
|
|
| ### 2.2 Implementation Roadmap |
|
|
| #### Step 1: Logging Interface Refactor (True Zero-Buffer) |
| * **Task:** Modify `RequestLogger.LogStreamingRequest` signature. |
| * **From:** `LogStreamingRequest(..., body []byte, ...)` |
| * **To:** `LogStreamingRequest(..., body io.Reader, ...)` |
| * **Action:** Update `internal/logging/request_logger.go`. |
| * **Action:** Update `internal/api/middleware/request_logging.go` to pass the request body stream directly (using `io.TeeReader` if necessary to log *and* process, though usually we log what we read). |
|
|
| #### Step 2: Proxy Timeout & Transport Hardening |
| * **Task:** Configure `httputil.ReverseProxy` with a custom `http.Transport`. |
| * **Action:** In `internal/api/modules/amp/proxy.go`, define: |
| ```go |
| Transport: &http.Transport{ |
| DialContext: (&net.Dialer{ |
| Timeout: 10 * time.Second, |
| KeepAlive: 30 * time.Second, |
| }).DialContext, |
| ResponseHeaderTimeout: 30 * time.Second, |
| IdleConnTimeout: 90 * time.Second, |
| } |
| ``` |
| * **Task:** Ensure `req.Context()` cancellation propagates upstream. |
| |
| #### Step 3: Observability Integration |
| * **Task:** Instrument the `FileStreamingLogWriter` and `ResponseRewriter`. |
| * **Action:** Add counters for `log_events_dropped`, `chunks_processed`, `thinking_blocks_redacted`. |
|
|
| #### Step 4: Testing & Validation |
| * **Task:** Add Unit/Integration Tests. |
| * Test `StreamingSanitizer` with random chunk splits (Fuzzing). |
| * Test `ResponseRewriter` with invalid JSON. |
| * Test Memory usage under load (using `pprof`). |
|
|
| ### 2.3 Technical Specifications |
|
|
| #### New Logger Interface |
| ```go |
| type RequestLogger interface { |
| // ... existing LogRequest ... |
| |
| // LogStreamingRequest now accepts a reader for the request body |
| LogStreamingRequest(ctx context.Context, url, method string, headers map[string][]string, bodyStream io.Reader, requestID string) (StreamingLogWriter, error) |
| } |
| ``` |
|
|
| #### Struct Modifications |
| **FileStreamingLogWriter:** |
| * Add `dropCount atomic.Uint64` to track buffer overflows. |
| * Add `metrics MetricsCollector` interface dependency. |
|
|
| **ResponseRewriter:** |
| * Add `FallbackMode bool` flag. If `true`, parser errors disable rewriting for the rest of the stream to ensure delivery. |
|
|
| ### 2.4 Execution Strategy |
| 1. **Refactor Logger Interface:** High impact, touches middleware. Do this first. |
| 2. **Harden Proxy:** Low risk, high value for reliability. |
| 3. **Add Tests:** Critical for verifying the stability of the complex streaming logic. |
| 4. **Add Metrics:** Final polish for operations. |
|
|