From df885c5a8efe0ceedbbd648217570306e60938f2 Mon Sep 17 00:00:00 2001 From: Hannes Hapke Date: Fri, 12 Jun 2026 19:26:22 -0700 Subject: [PATCH 01/13] fix: streaming support to work with Claude Code --- src/backend/providers/anthropic.go | 60 +++++-- src/backend/proxy/handler.go | 35 ++++ src/backend/proxy/streaming.go | 260 +++++++++++++++++++++++++++++ src/backend/proxy/transparent.go | 31 +++- 4 files changed, 371 insertions(+), 15 deletions(-) create mode 100644 src/backend/proxy/streaming.go diff --git a/src/backend/providers/anthropic.go b/src/backend/providers/anthropic.go index e5cd7168..ce268cf1 100644 --- a/src/backend/providers/anthropic.go +++ b/src/backend/providers/anthropic.go @@ -55,6 +55,20 @@ func (p *AnthropicProvider) ExtractRequestText(data map[string]interface{}) (str } if content, ok := msgMap["content"].(string); ok { result.WriteString(content + "\n") + } else if blocks, ok := msgMap["content"].([]interface{}); ok { + // Messages API content-block array: collect text from text blocks. + for _, blk := range blocks { + blkMap, ok := blk.(map[string]interface{}) + if !ok { + continue + } + if t, _ := blkMap["type"].(string); t != "text" { + continue + } + if text, ok := blkMap["text"].(string); ok { + result.WriteString(text + "\n") + } + } } } return result.String(), nil @@ -91,24 +105,46 @@ func (p *AnthropicProvider) CreateMaskedRequest(maskedRequest map[string]interfa return maskedToOriginal, &entities, fmt.Errorf("no messages field in request") } + // mask runs PII detection over a single piece of text and merges the + // resulting entities and mappings into the accumulators above. + mask := func(text string) string { + maskedText, _maskedToOriginal, _entities := maskPIIInText(text, "[MaskedRequest]") + entities = append(entities, _entities...) + for k, v := range _maskedToOriginal { + maskedToOriginal[k] = v + } + return maskedText + } + for _, msg := range messages { msgMap, ok := msg.(map[string]interface{}) if !ok { continue } - content, ok := msgMap["content"].(string) - if !ok { - continue - } - - // Mask PII in this message's content and update message content with masked text - maskedText, _maskedToOriginal, _entities := maskPIIInText(content, "[MaskedRequest]") - msgMap["content"] = maskedText - // Collect entities and mappings - entities = append(entities, _entities...) - for k, v := range _maskedToOriginal { - maskedToOriginal[k] = v + // The Messages API allows `content` to be either a plain string or an + // array of typed content blocks (Claude Code always uses the latter). + // Handle both so PII is masked in either shape. + switch content := msgMap["content"].(type) { + case string: + msgMap["content"] = mask(content) + case []interface{}: + for _, blk := range content { + blkMap, ok := blk.(map[string]interface{}) + if !ok { + continue + } + // Only text blocks carry free text; skip image / tool_use / + // tool_result blocks (different/nested shapes). + if t, _ := blkMap["type"].(string); t != "text" { + continue + } + text, ok := blkMap["text"].(string) + if !ok { + continue + } + blkMap["text"] = mask(text) + } } } diff --git a/src/backend/proxy/handler.go b/src/backend/proxy/handler.go index 5f5163e5..75238427 100644 --- a/src/backend/proxy/handler.go +++ b/src/backend/proxy/handler.go @@ -671,6 +671,41 @@ func (h *Handler) ProcessResponseBody(ctx context.Context, body []byte, contentT return modifiedBody } +// LogStreamedResponse records a streamed (SSE) response in the logging DB. The +// usual ProcessResponseBody path is skipped for streams (the body is pumped +// straight to the client), so this mirrors its response-logging half: a +// response_masked row (the text the model actually returned, before PII +// restoration) and a response_original row (the text delivered to the client, +// after restoration), correlated by transactionID. Only assistant text is +// captured; tool_use / tool_result blocks are not included. +func (h *Handler) LogStreamedResponse(ctx context.Context, transactionID, maskedText, restoredText string) { + if h.loggingDB == nil { + return + } + logCtx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + + // Wrap the assistant text in a minimal Anthropic-shaped envelope so the + // transaction ID can be attached and the log UI can parse it consistently. + wrap := func(text string) string { + b, err := json.Marshal(map[string]interface{}{ + "streamed": true, + "content": []map[string]interface{}{{"type": "text", "text": text}}, + }) + if err != nil { + return text + } + return h.addTransactionID(string(b), transactionID) + } + + if err := h.loggingDB.InsertLog(logCtx, wrap(maskedText), "response_masked", []pii.Entity{}, false); err != nil { + log.Printf("[Proxy] ⚠️ Failed to log masked streamed response: %v", err) + } + if err := h.loggingDB.InsertLog(logCtx, wrap(restoredText), "response_original", []pii.Entity{}, false); err != nil { + log.Printf("[Proxy] ⚠️ Failed to log restored streamed response: %v", err) + } +} + // addTransactionID adds transaction ID to JSON message for log correlation func (h *Handler) addTransactionID(message string, transactionID string) string { // Try to parse as JSON and add transaction_id field diff --git a/src/backend/proxy/streaming.go b/src/backend/proxy/streaming.go new file mode 100644 index 00000000..3a9516f4 --- /dev/null +++ b/src/backend/proxy/streaming.go @@ -0,0 +1,260 @@ +package proxy + +import ( + "bufio" + "bytes" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "strings" + "time" +) + +// streamingClient is used for requests whose responses are streamed (SSE). +// Unlike the shared handler client it has no overall timeout, so long-lived +// token streams are not cut off after 30s. A response-header timeout still +// guards against a dead upstream. Proxy is nil to avoid looping back through +// ourselves. +var streamingClient = &http.Client{ + Transport: &http.Transport{ + Proxy: nil, + ResponseHeaderTimeout: 60 * time.Second, + }, +} + +// requestWantsStream reports whether the (masked) request body asks for a +// streaming response, i.e. it contains "stream": true. +func requestWantsStream(body []byte) bool { + var m map[string]interface{} + if err := json.Unmarshal(body, &m); err != nil { + return false + } + v, ok := m["stream"].(bool) + return ok && v +} + +// isEventStream reports whether the response is a Server-Sent Events stream. +func isEventStream(resp *http.Response) bool { + return strings.HasPrefix( + strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") +} + +// sseRestorer restores masked PII inside an SSE token stream. Placeholder +// (dummy) values can be split across consecutive content_block_delta events, so +// it keeps a per-content-block carry buffer and only emits text that is far +// enough from the tail that any placeholder starting in it is guaranteed +// complete. The held-back tail is flushed when the content block stops. +type sseRestorer struct { + mapping map[string]string + keep int // bytes to hold back = longest dummy length - 1 + carry map[int]string // un-emitted tail per content-block index + masked strings.Builder // raw model text (pre-restore), accumulated for logging +} + +func newSSERestorer(mapping map[string]string) *sseRestorer { + keep := 0 + for masked := range mapping { + if len(masked) > keep { + keep = len(masked) + } + } + if keep > 0 { + keep-- + } + return &sseRestorer{mapping: mapping, keep: keep, carry: map[int]string{}} +} + +// restore replaces every masked (dummy) value with its original. Replacing +// already-restored text is idempotent provided an original value does not +// contain a dummy placeholder as a substring. +func (s *sseRestorer) restore(text string) string { + for masked, original := range s.mapping { + text = strings.ReplaceAll(text, masked, original) + } + return text +} + +// splitSafe returns the prefix that is safe to emit now and the tail that must +// be held back so a placeholder straddling the boundary can still complete. +func splitSafe(s string, keep int) (emit, hold string) { + if len(s) <= keep { + return "", s + } + return s[:len(s)-keep], s[len(s)-keep:] +} + +// transformEvent rewrites a single, complete SSE event (the raw lines up to and +// including the blank terminator). For a content_block_delta text delta it +// restores PII and rewrites only the JSON on the data: line, leaving the +// surrounding event:/blank framing intact so the event is always well formed — +// even when the entire delta is held back (an empty text delta is emitted). Any +// held-back tail is flushed as a synthetic delta immediately before the +// matching content_block_stop. Every other event passes through byte-for-byte. +func (s *sseRestorer) transformEvent(lines [][]byte) []byte { + dataIdx := -1 + var evt struct { + Type string `json:"type"` + Index int `json:"index"` + Delta struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"delta"` + } + for i, ln := range lines { + t := bytes.TrimRight(ln, "\r\n") + if bytes.HasPrefix(t, []byte("data: ")) { + if json.Unmarshal(t[len("data: "):], &evt) == nil { + dataIdx = i + } + break + } + } + if dataIdx == -1 { + return concatLines(lines) // no recognisable data line: pass through + } + + switch { + case evt.Type == "content_block_delta" && evt.Delta.Type == "text_delta": + s.masked.WriteString(evt.Delta.Text) // raw model text, for audit logging + buf := s.carry[evt.Index] + evt.Delta.Text + emit, hold := splitSafe(s.restore(buf), s.keep) + s.carry[evt.Index] = hold + // Rewrite only the data: line; keep the original event:/blank lines. + out := make([][]byte, len(lines)) + copy(out, lines) + out[dataIdx] = textDeltaDataLine(evt.Index, emit) + return concatLines(out) + + case evt.Type == "content_block_stop": + var out []byte + if tail := s.carry[evt.Index]; tail != "" { + out = append(out, textDeltaEvent(evt.Index, tail)...) + delete(s.carry, evt.Index) + } + return append(out, concatLines(lines)...) + + default: + return concatLines(lines) + } +} + +func concatLines(lines [][]byte) []byte { + var b []byte + for _, ln := range lines { + b = append(b, ln...) + } + return b +} + +// textDeltaDataLine builds just the `data: {...}\n` line for a text delta. +func textDeltaDataLine(index int, text string) []byte { + payload, _ := json.Marshal(map[string]interface{}{ + "type": "content_block_delta", + "index": index, + "delta": map[string]interface{}{"type": "text_delta", "text": text}, + }) + out := append([]byte("data: "), payload...) + return append(out, '\n') +} + +// textDeltaEvent builds a complete content_block_delta SSE event, including its +// trailing blank line, used to flush a held-back tail. +func textDeltaEvent(index int, text string) []byte { + out := []byte("event: content_block_delta\n") + out = append(out, textDeltaDataLine(index, text)...) + return append(out, '\n') +} + +// isBlankLine reports whether a raw line is an SSE event terminator. +func isBlankLine(line []byte) bool { + return len(bytes.TrimRight(line, "\r\n")) == 0 +} + +// streamSSEResponse writes an SSE response to the (HTTP/1.1) client connection, +// restoring masked PII incrementally and flushing after every event so the +// client receives tokens as they arrive. Events are buffered only until their +// blank-line terminator, never the whole body, and chunked transfer encoding is +// used instead of Content-Length. +func streamSSEResponse(conn net.Conn, resp *http.Response, restorer *sseRestorer) error { + bw := bufio.NewWriter(conn) + + // Status line. + if _, err := fmt.Fprintf(bw, "HTTP/1.1 %d %s\r\n", resp.StatusCode, http.StatusText(resp.StatusCode)); err != nil { + return err + } + + // Copy headers, but take control of framing: stream as chunked, no length. + for key, values := range resp.Header { + switch strings.ToLower(key) { + case "content-length", "transfer-encoding", "connection": + continue + } + for _, v := range values { + if _, err := fmt.Fprintf(bw, "%s: %s\r\n", key, v); err != nil { + return err + } + } + } + if _, err := bw.WriteString("Transfer-Encoding: chunked\r\nConnection: keep-alive\r\n\r\n"); err != nil { + return err + } + if err := bw.Flush(); err != nil { + return err + } + + writeChunk := func(b []byte) error { + if len(b) == 0 { + return nil + } + if _, err := fmt.Fprintf(bw, "%x\r\n", len(b)); err != nil { + return err + } + if _, err := bw.Write(b); err != nil { + return err + } + if _, err := bw.WriteString("\r\n"); err != nil { + return err + } + return bw.Flush() // flush each event so the client streams in real time + } + + reader := bufio.NewReader(resp.Body) + var event [][]byte + flush := func() error { + if len(event) == 0 { + return nil + } + out := restorer.transformEvent(event) + event = event[:0] + return writeChunk(out) + } + + for { + line, err := reader.ReadBytes('\n') + if len(line) > 0 { + event = append(event, line) + if isBlankLine(line) { // blank line terminates the SSE event + if werr := flush(); werr != nil { + return werr + } + } + } + if err != nil { + if err == io.EOF { + if ferr := flush(); ferr != nil { // trailing event w/o blank line + return ferr + } + break + } + return err + } + } + + // Terminating zero-length chunk. + if _, err := bw.WriteString("0\r\n\r\n"); err != nil { + return err + } + return bw.Flush() +} diff --git a/src/backend/proxy/transparent.go b/src/backend/proxy/transparent.go index 80a2bc92..cabd3852 100644 --- a/src/backend/proxy/transparent.go +++ b/src/backend/proxy/transparent.go @@ -457,9 +457,16 @@ func (tp *TransparentProxy) interceptHTTPOverTLS(conn net.Conn, r *http.Request, // Explicitly set Accept-Encoding to identity to avoid compressed responses proxyReq.Header.Set("Accept-Encoding", "identity") - // Forward request using handler's HTTP client (bypasses proxy to prevent infinite loop) - log.Printf("[TransparentProxy] Forwarding TLS request directly to %s (bypassing proxy)", targetURL) - resp, err := tp.handler.GetHTTPClient().Do(proxyReq) + // Forward request using handler's HTTP client (bypasses proxy to prevent infinite loop). + // Streaming requests use a client without an overall timeout so long-lived + // SSE token streams are not cut off mid-response. + wantStream := requestWantsStream(processed.RedactedBody) + httpClient := tp.handler.GetHTTPClient() + if wantStream { + httpClient = streamingClient + } + log.Printf("[TransparentProxy] Forwarding TLS request directly to %s (bypassing proxy, stream=%t)", targetURL, wantStream) + resp, err := httpClient.Do(proxyReq) if err != nil { log.Printf("[TransparentProxy] ❌ Failed to forward request: %v", err) tp.writeErrorResponse(conn, http.StatusBadGateway, fmt.Sprintf("Failed to forward request: %v", err)) @@ -467,6 +474,24 @@ func (tp *TransparentProxy) interceptHTTPOverTLS(conn net.Conn, r *http.Request, } defer resp.Body.Close() + // Stream Server-Sent Events straight through to the client, restoring PII + // incrementally. Buffering an SSE stream (as the non-streaming path below + // does) breaks streaming clients like Claude Code and hangs until the + // upstream timeout fires. + if wantStream && isEventStream(resp) { + log.Printf("[TransparentProxy] Streaming SSE response for %s", r.URL.Path) + restorer := newSSERestorer(processed.MaskedToOriginal) + if streamErr := streamSSEResponse(conn, resp, restorer); streamErr != nil { + log.Printf("[TransparentProxy] ❌ Failed to stream SSE response: %v", streamErr) + } + // Record the streamed response for audit: masked = the text the model + // actually returned (pre-restore), restored = what the client received. + maskedText := restorer.masked.String() + tp.handler.LogStreamedResponse(ctx, processed.TransactionID, maskedText, restorer.restore(maskedText)) + log.Printf("[TransparentProxy] Streamed %s %s - Status: %d", r.Method, r.URL.Path, resp.StatusCode) + return + } + // Read response body respBody, err := io.ReadAll(resp.Body) if err != nil { From b7caf5abe744241a22bf5b90d55485f51dd9a74f Mon Sep 17 00:00:00 2001 From: Hannes Hapke Date: Sat, 13 Jun 2026 10:13:31 -0700 Subject: [PATCH 02/13] codex streaming --- src/backend/config/config.go | 8 +- src/backend/providers/openai.go | 31 ++++- src/backend/providers/provider.go | 3 +- src/backend/proxy/codec_anthropic.go | 122 +++++++++++++++++++ src/backend/proxy/codec_openai.go | 133 +++++++++++++++++++++ src/backend/proxy/handler.go | 5 +- src/backend/proxy/streaming.go | 170 +++++++++++++-------------- src/backend/proxy/transparent.go | 8 +- 8 files changed, 377 insertions(+), 103 deletions(-) create mode 100644 src/backend/proxy/codec_anthropic.go create mode 100644 src/backend/proxy/codec_openai.go diff --git a/src/backend/config/config.go b/src/backend/config/config.go index f92fa02f..f12c40cb 100644 --- a/src/backend/config/config.go +++ b/src/backend/config/config.go @@ -343,13 +343,19 @@ func DefaultConfig() *Config { // GetInterceptDomains returns the list of intercept domains (as a union of all provider domains) func (pc ProvidersConfig) GetInterceptDomains() []string { - return []string{ + domains := []string{ interceptDomain(pc.AnthropicProviderConfig.APIDomain), interceptDomain(pc.OpenAIProviderConfig.APIDomain), interceptDomain(pc.GeminiProviderConfig.APIDomain), interceptDomain(pc.MistralProviderConfig.APIDomain), interceptDomain(pc.CustomProviderConfig.APIDomain), } + // ChatGPT-login Codex talks to chatgpt.com instead of the configured OpenAI + // API domain, so it must be intercepted explicitly whenever OpenAI is enabled. + if pc.OpenAIProviderConfig.APIDomain != "" { + domains = append(domains, "chatgpt.com") + } + return domains } func interceptDomain(apiDomain string) string { diff --git a/src/backend/providers/openai.go b/src/backend/providers/openai.go index 0042cf3c..10418a3e 100644 --- a/src/backend/providers/openai.go +++ b/src/backend/providers/openai.go @@ -11,11 +11,16 @@ import ( ) const ( - ProviderTypeOpenAI ProviderType = "openai" - ProviderSubpathOpenAI string = "/v1/chat/completions" - ProviderSubpathOpenAIResp string = "/v1/responses" - ProviderAPIDomainOpenAI string = "api.openai.com" - ProviderNameOpenAI string = "OpenAI" + ProviderTypeOpenAI ProviderType = "openai" + ProviderSubpathOpenAI string = "/v1/chat/completions" + ProviderSubpathOpenAIResp string = "/v1/responses" + ProviderAPIDomainOpenAI string = "api.openai.com" + ProviderNameOpenAI string = "OpenAI" + // ProviderAPIDomainCodex is the host used by ChatGPT-login Codex (the OpenAI + // CLI). It hits chatgpt.com/backend-api/codex/responses with an OAuth bearer + // token instead of api.openai.com, so it must be routed to and intercepted by + // the OpenAI provider alongside the API-key host. + ProviderAPIDomainCodex string = "chatgpt.com" ) // reasoningModelFamilies lists OpenAI model family prefixes that require the @@ -548,6 +553,22 @@ func restoreResponsesAPIResponse(maskedResponse map[string]interface{}, maskedTo if !ok { continue } + + // function_call output items carry model-generated tool arguments as a + // JSON string, which can echo masked PII just like assistant text. Restore + // it here so buffered responses match the streaming codec's behavior. + if args, ok := itemMap["arguments"].(string); ok { + restoredArgs := restorePII(args, maskedToOriginal) + if restoredArgs != args && getLogResponses() { + log.Printf("PII restored in response tool-call arguments") + if getLogVerbose() { + log.Printf("Original tool-call arguments: %s", args) + log.Printf("Restored tool-call arguments: %s", restoredArgs) + } + } + itemMap["arguments"] = restoredArgs + } + contents, ok := itemMap["content"].([]interface{}) if !ok { continue diff --git a/src/backend/providers/provider.go b/src/backend/providers/provider.go index fb5890d2..8304afa7 100644 --- a/src/backend/providers/provider.go +++ b/src/backend/providers/provider.go @@ -175,7 +175,8 @@ func (p *Providers) GetProviderFromHost(host string, logPrefix string) (*Provide } switch { - case p.OpenAIProvider != nil && providerHostMatches(host, p.OpenAIProvider.apiDomain): + case p.OpenAIProvider != nil && (providerHostMatches(host, p.OpenAIProvider.apiDomain) || host == ProviderAPIDomainCodex): + // chatgpt.com is the ChatGPT-login Codex host; route it to OpenAI too. provider = p.OpenAIProvider case p.AnthropicProvider != nil && providerHostMatches(host, p.AnthropicProvider.apiDomain): provider = p.AnthropicProvider diff --git a/src/backend/proxy/codec_anthropic.go b/src/backend/proxy/codec_anthropic.go new file mode 100644 index 00000000..8336fa75 --- /dev/null +++ b/src/backend/proxy/codec_anthropic.go @@ -0,0 +1,122 @@ +package proxy + +import ( + "bytes" + "encoding/json" +) + +// anthropicCodec restores masked PII inside an Anthropic Messages-API SSE +// stream. Placeholder (dummy) values can be split across consecutive +// content_block_delta events, so it keeps a per-content-block carry buffer and +// only emits text that is far enough from the tail that any placeholder starting +// in it is guaranteed complete. The held-back tail is flushed when the content +// block stops. +type anthropicCodec struct { + restoreCore + carry map[int]string // un-emitted tail per content-block index + kind map[int]string // delta type per index ("text_delta" / "input_json_delta") +} + +func newAnthropicCodec(mapping map[string]string) *anthropicCodec { + return &anthropicCodec{ + restoreCore: newRestoreCore(mapping), + carry: map[int]string{}, + kind: map[int]string{}, + } +} + +// transformEvent rewrites a single, complete SSE event (the raw lines up to and +// including the blank terminator). For a content_block_delta text delta it +// restores PII and rewrites only the JSON on the data: line, leaving the +// surrounding event:/blank framing intact so the event is always well formed — +// even when the entire delta is held back (an empty text delta is emitted). Any +// held-back tail is flushed as a synthetic delta immediately before the +// matching content_block_stop. Every other event passes through byte-for-byte. +func (s *anthropicCodec) transformEvent(lines [][]byte) []byte { + dataIdx := -1 + var evt struct { + Type string `json:"type"` + Index int `json:"index"` + Delta struct { + Type string `json:"type"` + Text string `json:"text"` + PartialJSON string `json:"partial_json"` + } `json:"delta"` + } + for i, ln := range lines { + t := bytes.TrimRight(ln, "\r\n") + if bytes.HasPrefix(t, []byte("data: ")) { + if json.Unmarshal(t[len("data: "):], &evt) == nil { + dataIdx = i + } + break + } + } + if dataIdx == -1 { + return concatLines(lines) // no recognisable data line: pass through + } + + switch { + case evt.Type == "content_block_delta" && (evt.Delta.Type == "text_delta" || evt.Delta.Type == "input_json_delta"): + // text_delta carries assistant text; input_json_delta carries a fragment + // of a tool call's JSON arguments. PII can hide in either, and can be + // split across consecutive deltas, so both go through the carry buffer. + raw := evt.Delta.Text + if evt.Delta.Type == "input_json_delta" { + raw = evt.Delta.PartialJSON + } + s.masked.WriteString(raw) // raw model output (text or tool args), for audit + s.kind[evt.Index] = evt.Delta.Type + buf := s.carry[evt.Index] + raw + emit, hold := splitSafe(s.restore(buf), s.keep) + s.carry[evt.Index] = hold + // Rewrite only the data: line; keep the original event:/blank lines. + out := make([][]byte, len(lines)) + copy(out, lines) + out[dataIdx] = deltaDataLine(evt.Index, evt.Delta.Type, emit) + return concatLines(out) + + case evt.Type == "content_block_stop": + var out []byte + if tail := s.carry[evt.Index]; tail != "" { + out = append(out, deltaEvent(evt.Index, s.kind[evt.Index], tail)...) + delete(s.carry, evt.Index) + delete(s.kind, evt.Index) + } + return append(out, concatLines(lines)...) + + default: + return concatLines(lines) + } +} + +// deltaField returns the JSON field that carries the payload for a delta type. +func deltaField(deltaType string) string { + if deltaType == "input_json_delta" { + return "partial_json" + } + return "text" +} + +// deltaDataLine builds the `data: {...}\n` line for a text_delta or +// input_json_delta carrying the given (restored) payload. +func deltaDataLine(index int, deltaType, payload string) []byte { + b, _ := json.Marshal(map[string]interface{}{ + "type": "content_block_delta", + "index": index, + "delta": map[string]interface{}{"type": deltaType, deltaField(deltaType): payload}, + }) + out := append([]byte("data: "), b...) + return append(out, '\n') +} + +// deltaEvent builds a complete content_block_delta SSE event (with trailing +// blank line) used to flush a held-back tail. Defaults to a text delta. +func deltaEvent(index int, deltaType, payload string) []byte { + if deltaType == "" { + deltaType = "text_delta" + } + out := []byte("event: content_block_delta\n") + out = append(out, deltaDataLine(index, deltaType, payload)...) + return append(out, '\n') +} diff --git a/src/backend/proxy/codec_openai.go b/src/backend/proxy/codec_openai.go new file mode 100644 index 00000000..a85d0275 --- /dev/null +++ b/src/backend/proxy/codec_openai.go @@ -0,0 +1,133 @@ +package proxy + +import ( + "bytes" + "encoding/json" + "fmt" + "strings" +) + +// openaiCodec restores masked PII inside an OpenAI Responses-API SSE stream +// (used by both api.openai.com and ChatGPT-login Codex on chatgpt.com). +// +// Incremental text arrives in `*.delta` events whose payload is the string field +// "delta" (e.g. response.output_text.delta, response.function_call_arguments. +// delta). A matching `*.done` event then repeats the complete value in a field +// named for its kind ("text", "arguments", or "refusal"). Placeholders can be +// split across consecutive deltas, so delta text goes through a per-channel +// carry buffer keyed by output_index/content_index; the held-back tail is +// flushed as a synthetic delta when the channel's `.done` arrives, and the +// `.done` payload itself is restored whole. +// +// The codec matches on the stable `.delta`/`.done` type suffix and the payload +// field names rather than enumerating exact event types, so it tolerates new +// Responses event kinds (text, tool args, refusals) without changes. Events it +// does not recognise pass through byte-for-byte. +type openaiCodec struct { + restoreCore + carry map[string]string // un-emitted (already-restored) tail per channel +} + +func newOpenAICodec(mapping map[string]string) *openaiCodec { + return &openaiCodec{ + restoreCore: newRestoreCore(mapping), + carry: map[string]string{}, + } +} + +// doneFields are the fields a Responses `*.done` event uses to carry the full +// value, by event kind: output text, tool-call arguments, and refusals. +var doneFields = []string{"text", "arguments", "refusal"} + +func (c *openaiCodec) transformEvent(lines [][]byte) []byte { + dataIdx := -1 + var raw []byte + for i, ln := range lines { + t := bytes.TrimRight(ln, "\r\n") + if bytes.HasPrefix(t, []byte("data: ")) { + raw = t[len("data: "):] + dataIdx = i + break + } + } + if dataIdx == -1 { + return concatLines(lines) // no data line (comment/keep-alive): pass through + } + + var obj map[string]interface{} + if err := json.Unmarshal(raw, &obj); err != nil { + return concatLines(lines) // non-JSON payload (e.g. "[DONE]"): pass through + } + typ, _ := obj["type"].(string) + key := channelKey(obj) + + switch { + case strings.HasSuffix(typ, ".delta"): + delta, ok := obj["delta"].(string) + if !ok { + return concatLines(lines) // non-text delta (e.g. audio): pass through + } + c.masked.WriteString(delta) // raw model output, for audit + emit, hold := splitSafe(c.restore(c.carry[key]+delta), c.keep) + c.carry[key] = hold + obj["delta"] = emit + return rewriteDataLine(lines, dataIdx, obj) + + case strings.HasSuffix(typ, ".done"): + var out []byte + // Flush any held-back delta tail as a synthetic delta so incremental + // renderers see the complete text before the terminating .done event. + if tail := c.carry[key]; tail != "" { + delete(c.carry, key) + out = append(out, c.tailDeltaEvent(typ, obj, tail)...) + } + // Restore the full value the .done event repeats. (Not added to the audit + // accumulator: the deltas above already captured this text.) + for _, field := range doneFields { + if full, ok := obj[field].(string); ok { + obj[field] = c.restore(full) + } + } + return append(out, rewriteDataLine(lines, dataIdx, obj)...) + + default: + return concatLines(lines) + } +} + +// tailDeltaEvent builds a synthetic `*.delta` event mirroring a `*.done` event's +// channel, carrying the already-restored held-back tail. The payload is emitted +// as-is (the carry buffer stores post-restore text). +func (c *openaiCodec) tailDeltaEvent(doneType string, obj map[string]interface{}, tail string) []byte { + deltaType := strings.TrimSuffix(doneType, ".done") + ".delta" + d := map[string]interface{}{"type": deltaType, "delta": tail} + for _, f := range []string{"item_id", "output_index", "content_index"} { + if v, ok := obj[f]; ok { + d[f] = v + } + } + b, _ := json.Marshal(d) + out := []byte("event: " + deltaType + "\n") + out = append(out, append(append([]byte("data: "), b...), '\n')...) + return append(out, '\n') +} + +// channelKey identifies an independent text channel within a Responses stream so +// each gets its own carry buffer. Deltas and their .done share output_index and +// content_index. +func channelKey(obj map[string]interface{}) string { + oi, _ := obj["output_index"].(float64) + ci, _ := obj["content_index"].(float64) + return fmt.Sprintf("%d:%d", int(oi), int(ci)) +} + +// rewriteDataLine re-marshals obj into the event's data: line, preserving every +// other line (event:, blank terminator) so framing stays intact. +func rewriteDataLine(lines [][]byte, dataIdx int, obj map[string]interface{}) []byte { + b, _ := json.Marshal(obj) + out := make([][]byte, len(lines)) + copy(out, lines) + line := append([]byte("data: "), b...) + out[dataIdx] = append(line, '\n') + return concatLines(out) +} diff --git a/src/backend/proxy/handler.go b/src/backend/proxy/handler.go index 75238427..def7de5c 100644 --- a/src/backend/proxy/handler.go +++ b/src/backend/proxy/handler.go @@ -676,8 +676,9 @@ func (h *Handler) ProcessResponseBody(ctx context.Context, body []byte, contentT // straight to the client), so this mirrors its response-logging half: a // response_masked row (the text the model actually returned, before PII // restoration) and a response_original row (the text delivered to the client, -// after restoration), correlated by transactionID. Only assistant text is -// captured; tool_use / tool_result blocks are not included. +// after restoration), correlated by transactionID. Both assistant text and +// tool-call argument JSON are captured (concatenated in the logged text); +// tool_result blocks on the request side are not. func (h *Handler) LogStreamedResponse(ctx context.Context, transactionID, maskedText, restoredText string) { if h.loggingDB == nil { return diff --git a/src/backend/proxy/streaming.go b/src/backend/proxy/streaming.go index 3a9516f4..88d3c6d8 100644 --- a/src/backend/proxy/streaming.go +++ b/src/backend/proxy/streaming.go @@ -6,12 +6,42 @@ import ( "encoding/json" "fmt" "io" + "log" "net" "net/http" + "os" + "path/filepath" "strings" "time" + + "github.com/dataiku/kiji-proxy/src/backend/providers" ) +// sseCaptureDirEnv names a directory to mirror raw upstream SSE streams into, +// one file per stream, for debugging provider event grammars (e.g. confirming +// the real ChatGPT-login Codex Responses event names before trusting the codec). +// Capture is off unless this env var is set to a writable directory. +const sseCaptureDirEnv = "KIJI_SSE_CAPTURE_DIR" + +// newSSECapture opens a per-stream capture file when sseCaptureDirEnv is set, +// or returns nil to disable capture. The returned writer receives the upstream +// bytes byte-for-byte (pre-restore), so the captured file is the provider's raw +// SSE as sent. +func newSSECapture() io.WriteCloser { + dir := os.Getenv(sseCaptureDirEnv) + if dir == "" { + return nil + } + name := fmt.Sprintf("sse-%d.log", time.Now().UnixNano()) + f, err := os.Create(filepath.Join(dir, name)) + if err != nil { + log.Printf("[stream] SSE capture disabled, cannot create file in %s: %v", dir, err) + return nil + } + log.Printf("[stream] Capturing raw upstream SSE to %s", f.Name()) + return f +} + // streamingClient is used for requests whose responses are streamed (SSE). // Unlike the shared handler client it has no overall timeout, so long-lived // token streams are not cut off after 30s. A response-header timeout still @@ -41,19 +71,32 @@ func isEventStream(resp *http.Response) bool { strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") } -// sseRestorer restores masked PII inside an SSE token stream. Placeholder -// (dummy) values can be split across consecutive content_block_delta events, so -// it keeps a per-content-block carry buffer and only emits text that is far -// enough from the tail that any placeholder starting in it is guaranteed -// complete. The held-back tail is flushed when the content block stops. -type sseRestorer struct { +// streamCodec restores masked PII inside one provider's SSE token stream. The +// engine below (streamSSEResponse) owns the transport-level framing and event +// splitting; a codec only knows how to rewrite a single, complete SSE event for +// its provider's grammar. transformEvent receives the raw lines of one event +// (up to and including the blank terminator) and returns the bytes to write. +type streamCodec interface { + transformEvent(lines [][]byte) []byte + // restore replaces every masked (dummy) value with its original. Exposed so + // the caller can restore the accumulated audit text after the stream ends. + restore(text string) string + // maskedOutput returns the raw model output (pre-restore) accumulated across + // the stream, for audit logging. + maskedOutput() string +} + +// restoreCore holds the provider-agnostic restore state shared by every codec: +// the dummy→original mapping, the hold-back length, and the pre-restore model +// output accumulated for audit logging. Codecs embed it so they only implement +// their own event grammar. +type restoreCore struct { mapping map[string]string keep int // bytes to hold back = longest dummy length - 1 - carry map[int]string // un-emitted tail per content-block index - masked strings.Builder // raw model text (pre-restore), accumulated for logging + masked strings.Builder // raw model output (pre-restore), accumulated for logging } -func newSSERestorer(mapping map[string]string) *sseRestorer { +func newRestoreCore(mapping map[string]string) restoreCore { keep := 0 for masked := range mapping { if len(masked) > keep { @@ -63,19 +106,23 @@ func newSSERestorer(mapping map[string]string) *sseRestorer { if keep > 0 { keep-- } - return &sseRestorer{mapping: mapping, keep: keep, carry: map[int]string{}} + return restoreCore{mapping: mapping, keep: keep} } // restore replaces every masked (dummy) value with its original. Replacing // already-restored text is idempotent provided an original value does not // contain a dummy placeholder as a substring. -func (s *sseRestorer) restore(text string) string { - for masked, original := range s.mapping { +func (c *restoreCore) restore(text string) string { + for masked, original := range c.mapping { text = strings.ReplaceAll(text, masked, original) } return text } +func (c *restoreCore) maskedOutput() string { + return c.masked.String() +} + // splitSafe returns the prefix that is safe to emit now and the tail that must // be held back so a placeholder straddling the boundary can still complete. func splitSafe(s string, keep int) (emit, hold string) { @@ -85,61 +132,6 @@ func splitSafe(s string, keep int) (emit, hold string) { return s[:len(s)-keep], s[len(s)-keep:] } -// transformEvent rewrites a single, complete SSE event (the raw lines up to and -// including the blank terminator). For a content_block_delta text delta it -// restores PII and rewrites only the JSON on the data: line, leaving the -// surrounding event:/blank framing intact so the event is always well formed — -// even when the entire delta is held back (an empty text delta is emitted). Any -// held-back tail is flushed as a synthetic delta immediately before the -// matching content_block_stop. Every other event passes through byte-for-byte. -func (s *sseRestorer) transformEvent(lines [][]byte) []byte { - dataIdx := -1 - var evt struct { - Type string `json:"type"` - Index int `json:"index"` - Delta struct { - Type string `json:"type"` - Text string `json:"text"` - } `json:"delta"` - } - for i, ln := range lines { - t := bytes.TrimRight(ln, "\r\n") - if bytes.HasPrefix(t, []byte("data: ")) { - if json.Unmarshal(t[len("data: "):], &evt) == nil { - dataIdx = i - } - break - } - } - if dataIdx == -1 { - return concatLines(lines) // no recognisable data line: pass through - } - - switch { - case evt.Type == "content_block_delta" && evt.Delta.Type == "text_delta": - s.masked.WriteString(evt.Delta.Text) // raw model text, for audit logging - buf := s.carry[evt.Index] + evt.Delta.Text - emit, hold := splitSafe(s.restore(buf), s.keep) - s.carry[evt.Index] = hold - // Rewrite only the data: line; keep the original event:/blank lines. - out := make([][]byte, len(lines)) - copy(out, lines) - out[dataIdx] = textDeltaDataLine(evt.Index, emit) - return concatLines(out) - - case evt.Type == "content_block_stop": - var out []byte - if tail := s.carry[evt.Index]; tail != "" { - out = append(out, textDeltaEvent(evt.Index, tail)...) - delete(s.carry, evt.Index) - } - return append(out, concatLines(lines)...) - - default: - return concatLines(lines) - } -} - func concatLines(lines [][]byte) []byte { var b []byte for _, ln := range lines { @@ -148,25 +140,6 @@ func concatLines(lines [][]byte) []byte { return b } -// textDeltaDataLine builds just the `data: {...}\n` line for a text delta. -func textDeltaDataLine(index int, text string) []byte { - payload, _ := json.Marshal(map[string]interface{}{ - "type": "content_block_delta", - "index": index, - "delta": map[string]interface{}{"type": "text_delta", "text": text}, - }) - out := append([]byte("data: "), payload...) - return append(out, '\n') -} - -// textDeltaEvent builds a complete content_block_delta SSE event, including its -// trailing blank line, used to flush a held-back tail. -func textDeltaEvent(index int, text string) []byte { - out := []byte("event: content_block_delta\n") - out = append(out, textDeltaDataLine(index, text)...) - return append(out, '\n') -} - // isBlankLine reports whether a raw line is an SSE event terminator. func isBlankLine(line []byte) bool { return len(bytes.TrimRight(line, "\r\n")) == 0 @@ -176,8 +149,9 @@ func isBlankLine(line []byte) bool { // restoring masked PII incrementally and flushing after every event so the // client receives tokens as they arrive. Events are buffered only until their // blank-line terminator, never the whole body, and chunked transfer encoding is -// used instead of Content-Length. -func streamSSEResponse(conn net.Conn, resp *http.Response, restorer *sseRestorer) error { +// used instead of Content-Length. The provider-specific rewriting is delegated +// to codec. +func streamSSEResponse(conn net.Conn, resp *http.Response, codec streamCodec) error { bw := bufio.NewWriter(conn) // Status line. @@ -220,13 +194,18 @@ func streamSSEResponse(conn net.Conn, resp *http.Response, restorer *sseRestorer return bw.Flush() // flush each event so the client streams in real time } - reader := bufio.NewReader(resp.Body) + var src io.Reader = resp.Body + if cap := newSSECapture(); cap != nil { + defer cap.Close() + src = io.TeeReader(resp.Body, cap) // mirror raw upstream bytes for debugging + } + reader := bufio.NewReader(src) var event [][]byte flush := func() error { if len(event) == 0 { return nil } - out := restorer.transformEvent(event) + out := codec.transformEvent(event) event = event[:0] return writeChunk(out) } @@ -258,3 +237,14 @@ func streamSSEResponse(conn net.Conn, resp *http.Response, restorer *sseRestorer } return bw.Flush() } + +// codecForProvider selects the SSE codec matching the upstream provider's stream +// grammar. OpenAI (incl. ChatGPT-login Codex on chatgpt.com) speaks the +// Responses-API event shape; everything else uses the Anthropic grammar, which +// is also the safe default. +func codecForProvider(provider *providers.Provider, mapping map[string]string) streamCodec { + if provider != nil && (*provider).GetType() == providers.ProviderTypeOpenAI { + return newOpenAICodec(mapping) + } + return newAnthropicCodec(mapping) +} diff --git a/src/backend/proxy/transparent.go b/src/backend/proxy/transparent.go index cabd3852..9e49d986 100644 --- a/src/backend/proxy/transparent.go +++ b/src/backend/proxy/transparent.go @@ -480,14 +480,14 @@ func (tp *TransparentProxy) interceptHTTPOverTLS(conn net.Conn, r *http.Request, // upstream timeout fires. if wantStream && isEventStream(resp) { log.Printf("[TransparentProxy] Streaming SSE response for %s", r.URL.Path) - restorer := newSSERestorer(processed.MaskedToOriginal) - if streamErr := streamSSEResponse(conn, resp, restorer); streamErr != nil { + codec := codecForProvider(provider, processed.MaskedToOriginal) + if streamErr := streamSSEResponse(conn, resp, codec); streamErr != nil { log.Printf("[TransparentProxy] ❌ Failed to stream SSE response: %v", streamErr) } // Record the streamed response for audit: masked = the text the model // actually returned (pre-restore), restored = what the client received. - maskedText := restorer.masked.String() - tp.handler.LogStreamedResponse(ctx, processed.TransactionID, maskedText, restorer.restore(maskedText)) + maskedText := codec.maskedOutput() + tp.handler.LogStreamedResponse(ctx, processed.TransactionID, maskedText, codec.restore(maskedText)) log.Printf("[TransparentProxy] Streamed %s %s - Status: %d", r.Method, r.URL.Path, resp.StatusCode) return } From 8426f288baa5275c9baf9439efe3eb583c01e928 Mon Sep 17 00:00:00 2001 From: Hannes Hapke Date: Sat, 13 Jun 2026 10:33:26 -0700 Subject: [PATCH 03/13] documentation --- README.md | 2 + docs/09-coding-agents.md | 173 +++++++++++++++++++++++++++++++++++++++ docs/README.md | 22 +++++ 3 files changed, 197 insertions(+) create mode 100644 docs/09-coding-agents.md diff --git a/README.md b/README.md index 73513e4b..287e5e2b 100644 --- a/README.md +++ b/README.md @@ -217,6 +217,7 @@ Complete documentation is available in [docs/README.md](docs/README.md): - **[Chrome Extension](docs/06-chrome-extension.md)** - Building, configuring, and publishing the PII Guard extension - **[Customizing the PII Model](docs/07-customizing-pii-model.md)** - Training a model with your own entity types - **[Masking Controls & Review](docs/08-masking-controls.md)** - Disable entity types, custom regex, mapping review +- **[Coding Agents (Codex & Claude Code)](docs/09-coding-agents.md)** - Route terminal coding agents through the proxy **Quick Links:** - [Installation Guide](docs/01-getting-started.md#quick-installation) @@ -225,6 +226,7 @@ Complete documentation is available in [docs/README.md](docs/README.md): - [Build for macOS](docs/03-building-deployment.md#building-for-macos) - [Build for Linux](docs/03-building-deployment.md#building-for-linux) - [Masking Controls](docs/08-masking-controls.md) - disable entities, custom regex, review mappings +- [Coding Agents Setup](docs/09-coding-agents.md) - Codex & Claude Code via the proxy --- diff --git a/docs/09-coding-agents.md b/docs/09-coding-agents.md new file mode 100644 index 00000000..5ce2e9ab --- /dev/null +++ b/docs/09-coding-agents.md @@ -0,0 +1,173 @@ +# Chapter 9: Coding Agents (Codex & Claude Code) + +Route terminal coding agents — OpenAI **Codex** and Anthropic **Claude Code** — through Kiji Privacy Proxy so that PII in your prompts, files, and tool calls is masked before it reaches the model, and restored in the model's replies. Streaming responses are restored token-by-token, so the agent still feels live. + +This chapter focuses on **what to set up on the client side**. For how the proxy itself works, see [Advanced Topics](05-advanced-topics.md#transparent-proxy--mitm). + +## How it works + +Coding agents talk to their provider over HTTPS. The proxy runs in **transparent (MITM) mode** (port `8081`): it intercepts traffic to known provider hosts, masks PII in the outgoing request, forwards it, then restores PII in the response — buffered or streamed (SSE). + +Hosts the proxy intercepts for coding agents: + +| Agent | Host(s) | Notes | +|-------|---------|-------| +| Claude Code | `api.anthropic.com` | | +| Codex (API key) | `api.openai.com` | `/v1/responses` and `/v1/chat/completions` | +| Codex (ChatGPT login) | `chatgpt.com` | `/backend-api/codex/responses` | + +For **any** agent, two things must be true: + +1. **The agent sends its HTTPS traffic through the proxy** — via `HTTP_PROXY` / `HTTPS_PROXY`. The macOS PAC auto-configuration only routes **browsers**; command-line agents must be pointed at the proxy explicitly. +2. **The agent trusts the proxy's CA** — so the MITM TLS handshake is accepted. Each agent reads its trusted CA from a different place (see below). + +## Prerequisites + +- Kiji Privacy Proxy is **running** (desktop app on macOS, or the standalone backend on Linux). See [Getting Started](01-getting-started.md). +- You know the path to the proxy CA certificate: + + | Platform | CA certificate path | + |----------|---------------------| + | macOS | `$HOME/Library/Application Support/Kiji Privacy Proxy/certs/ca.crt` | + | Linux | `~/.kiji-proxy/certs/ca.crt` | + +Throughout this chapter the proxy endpoint is `http://127.0.0.1:8081` (the transparent proxy port). Adjust if you changed `proxy_port`. + +## Claude Code + +Claude Code is a Node.js application, so it uses the standard Node proxy and CA variables. + +### Environment variables + +```bash +export HTTP_PROXY=http://127.0.0.1:8081 +export HTTPS_PROXY=http://127.0.0.1:8081 + +# macOS +export NODE_EXTRA_CA_CERTS="$HOME/Library/Application Support/Kiji Privacy Proxy/certs/ca.crt" +# Linux +export NODE_EXTRA_CA_CERTS="$HOME/.kiji-proxy/certs/ca.crt" +``` + +Then run `claude` in the same shell. Requests to `api.anthropic.com` now flow through the proxy. + +### Making it persistent + +Instead of exporting in every shell, set the variables in Claude Code's settings file so they apply to every session. Add an `env` block to `~/.claude/settings.json`: + +```json +{ + "env": { + "HTTP_PROXY": "http://127.0.0.1:8081", + "HTTPS_PROXY": "http://127.0.0.1:8081", + "NODE_EXTRA_CA_CERTS": "/Users/you/Library/Application Support/Kiji Privacy Proxy/certs/ca.crt" + } +} +``` + +(The path may contain spaces — that's fine inside the JSON string. Use the absolute path; `~`/`$HOME` are not expanded here.) + +## Codex + +Codex (`codex-cli`) is a **native Rust binary that uses rustls** for TLS, not Node and not the macOS keychain. Two consequences: + +- `NODE_EXTRA_CA_CERTS` is a Node concept — but Codex's CA loader happens to honor it as a fallback, so it still works (see below). +- Adding the CA to the macOS **System keychain alone is not enough**, because rustls uses its own root store. You must point Codex at the CA **file** via an environment variable. + +### Environment variables + +```bash +export HTTP_PROXY=http://127.0.0.1:8081 +export HTTPS_PROXY=http://127.0.0.1:8081 + +# macOS +export CODEX_CA_CERTIFICATE="$HOME/Library/Application Support/Kiji Privacy Proxy/certs/ca.crt" +# Linux +export CODEX_CA_CERTIFICATE="$HOME/.kiji-proxy/certs/ca.crt" +``` + +Then run `codex`. `CODEX_CA_CERTIFICATE` is Codex's native variable. If it is unset, Codex falls back — in order — to these standard CA-bundle variables, so any of them works too: + +``` +CODEX_CA_CERTIFICATE → SSL_CERT_FILE → REQUESTS_CA_BUNDLE → CURL_CA_BUNDLE + → NODE_EXTRA_CA_CERTS → GIT_SSL_CAINFO → BUNDLE_SSL_CA_CERT +``` + +This means if you already export `NODE_EXTRA_CA_CERTS` globally for Claude Code, Codex will pick up the same CA automatically — but setting `CODEX_CA_CERTIFICATE` explicitly is clearest. + +### API-key vs ChatGPT-login Codex + +- **API-key Codex** (`OPENAI_API_KEY` set) talks to `api.openai.com`. Your API key is forwarded untouched. +- **ChatGPT-login Codex** (signed in with `codex login`) talks to `chatgpt.com/backend-api/codex/responses` with an OAuth bearer token. The proxy leaves the `Authorization` header untouched and only masks/restores content, so your session keeps working. + +Both are intercepted with the same setup above; no extra configuration is needed to switch between them. + +## A shared snippet for both agents + +Drop this in your shell profile (`~/.zshrc` / `~/.bashrc`) to cover both agents at once on macOS: + +```bash +KIJI_CA="$HOME/Library/Application Support/Kiji Privacy Proxy/certs/ca.crt" +export HTTP_PROXY=http://127.0.0.1:8081 +export HTTPS_PROXY=http://127.0.0.1:8081 +export NODE_EXTRA_CA_CERTS="$KIJI_CA" # Claude Code (and Codex fallback) +export CODEX_CA_CERTIFICATE="$KIJI_CA" # Codex (explicit) +``` + +## What gets masked and restored + +| Direction | Covered | +|-----------|---------| +| Request → model | Chat `messages`; Responses-API `input` (string or message/part arrays), `instructions` (system prompt), tool-result `output`, and tool-call `arguments` | +| Model → response | Assistant text and tool-call `arguments`, for both **streaming** (SSE) and **buffered** replies | + +Every interception is recorded in the proxy's request log (visible in the desktop app), with the masked text the model actually saw and the restored text the agent received. See [Masking Controls & Review](08-masking-controls.md) to tune what gets masked and to review or delete recorded mappings. + +## Verifying interception + +1. **Check the proxy log.** Run a prompt in the agent, then open the desktop app's request log (or the standalone audit log). You should see an entry for the provider host with masked/restored content. +2. **Capture the raw stream (debugging).** Set `KIJI_SSE_CAPTURE_DIR` before starting the proxy to mirror each upstream SSE stream to a file — useful for confirming exactly what the agent received: + + ```bash + mkdir -p /tmp/agent-sse + KIJI_SSE_CAPTURE_DIR=/tmp/agent-sse + # run one agent request, then inspect: + grep -o '"type":"[^"]*"' /tmp/agent-sse/sse-*.log | sort -u + ``` + + Leave `KIJI_SSE_CAPTURE_DIR` unset in normal use; capture is off by default. + +## Troubleshooting + +**TLS / certificate error (`unable to get local issuer`, `invalid peer certificate`, handshake refused)** +- *Cause:* the agent doesn't trust the proxy CA. +- *Fix:* confirm the CA variable points at a file that exists and is readable. For Codex, use `CODEX_CA_CERTIFICATE` (a **file path**, not a directory); the macOS keychain alone won't satisfy rustls. For Claude Code, use `NODE_EXTRA_CA_CERTS`. Quote paths that contain spaces. + +**Traffic isn't being intercepted (no log entries)** +- *Cause:* the agent isn't using the proxy. +- *Fix:* ensure `HTTP_PROXY` and `HTTPS_PROXY` are exported in the **same shell/process** that runs the agent. Check that `NO_PROXY`/`no_proxy` doesn't list `openai.com`, `chatgpt.com`, or `anthropic.com`. Confirm the proxy is listening on the port you set. + +**ChatGPT-login Codex still fails after setup** +- *Cause:* the proxy build doesn't intercept `chatgpt.com`, or the CA isn't trusted on that host. +- *Fix:* verify `chatgpt.com` is in the proxy's intercept domains (it is added automatically when the OpenAI provider is configured) and that `CODEX_CA_CERTIFICATE` is set. As a last resort for diagnosis you can test with API-key Codex against `api.openai.com` to isolate whether the issue is host-specific. + +**Streaming feels stuck or arrives all at once** +- *Cause:* a buffering layer between the agent and the proxy. +- *Fix:* the proxy streams SSE through chunked and flushes per event. Make sure no additional proxy sits between the agent and Kiji, and that you point the agent directly at `127.0.0.1`. + +## Alternative: forward proxy without CA trust + +If you'd rather not install/trust the CA, agents that let you override the base URL can use the **forward proxy** (port `8080`) instead. The client talks plain HTTP to the proxy, which makes the upstream TLS connection itself — so no client-side CA trust is needed. + +```bash +# Claude Code → forward proxy (no CA needed) +export ANTHROPIC_BASE_URL=http://127.0.0.1:8080 +``` + +This works for API-key clients with a configurable endpoint. It does **not** work for **ChatGPT-login Codex**, whose `chatgpt.com` endpoint is fixed — that path requires the transparent proxy + CA trust described above. + +## See also + +- [Getting Started](01-getting-started.md) — installing the proxy and CA certificate +- [Advanced Topics](05-advanced-topics.md#transparent-proxy--mitm) — MITM architecture, CA management, CORS +- [Masking Controls & Review](08-masking-controls.md) — what gets masked, reviewing mappings diff --git a/docs/README.md b/docs/README.md index 88e35338..64a7c2eb 100644 --- a/docs/README.md +++ b/docs/README.md @@ -146,6 +146,22 @@ Control what gets masked and review what already has been, from the desktop app. --- +### [Chapter 9: Coding Agents (Codex & Claude Code)](09-coding-agents.md) + +Route terminal coding agents through the proxy so PII in prompts, code, and tool calls is masked before it reaches the model and restored in replies. + +**Topics:** +- How agents are intercepted (hosts, masking, streaming restore) +- Claude Code setup (`HTTP_PROXY`/`HTTPS_PROXY`, `NODE_EXTRA_CA_CERTS`, `settings.json`) +- Codex setup (rustls CA trust via `CODEX_CA_CERTIFICATE`, API-key vs ChatGPT-login) +- A shared shell snippet for both agents +- Verifying interception and troubleshooting TLS/proxy issues +- Forward-proxy alternative without CA trust + +**Start here if you're:** Using OpenAI Codex or Claude Code and want their traffic masked by Kiji. + +--- + ## Quick Links ### Getting Started @@ -187,6 +203,12 @@ Control what gets masked and review what already has been, from the desktop app. - [Custom Regex Patterns](08-masking-controls.md#custom-regex-patterns) - [Review & Delete Mappings](08-masking-controls.md#reviewing-and-deleting-masked-entities) +### Coding Agents +- [Claude Code Setup](09-coding-agents.md#claude-code) +- [Codex Setup](09-coding-agents.md#codex) +- [Shared Shell Snippet](09-coding-agents.md#a-shared-snippet-for-both-agents) +- [Troubleshooting](09-coding-agents.md#troubleshooting) + ## Document Status These documents consolidate and supersede the following original files: From 354cf5dbe0e6eaf1449dcf9001de9e3218096ed4 Mon Sep 17 00:00:00 2001 From: Hannes Hapke Date: Sat, 13 Jun 2026 10:59:17 -0700 Subject: [PATCH 04/13] linter fixes --- src/backend/proxy/codec_anthropic.go | 22 ++++++++++++++-------- src/backend/proxy/codec_openai.go | 6 +++--- src/backend/proxy/handler.go | 10 ++++++++-- 3 files changed, 25 insertions(+), 13 deletions(-) diff --git a/src/backend/proxy/codec_anthropic.go b/src/backend/proxy/codec_anthropic.go index 8336fa75..2181c827 100644 --- a/src/backend/proxy/codec_anthropic.go +++ b/src/backend/proxy/codec_anthropic.go @@ -5,6 +5,12 @@ import ( "encoding/json" ) +// Anthropic content_block_delta payload types. +const ( + deltaTypeText = "text_delta" + deltaTypeInputJSON = "input_json_delta" +) + // anthropicCodec restores masked PII inside an Anthropic Messages-API SSE // stream. Placeholder (dummy) values can be split across consecutive // content_block_delta events, so it keeps a per-content-block carry buffer and @@ -57,12 +63,12 @@ func (s *anthropicCodec) transformEvent(lines [][]byte) []byte { } switch { - case evt.Type == "content_block_delta" && (evt.Delta.Type == "text_delta" || evt.Delta.Type == "input_json_delta"): + case evt.Type == "content_block_delta" && (evt.Delta.Type == deltaTypeText || evt.Delta.Type == deltaTypeInputJSON): // text_delta carries assistant text; input_json_delta carries a fragment // of a tool call's JSON arguments. PII can hide in either, and can be // split across consecutive deltas, so both go through the carry buffer. raw := evt.Delta.Text - if evt.Delta.Type == "input_json_delta" { + if evt.Delta.Type == deltaTypeInputJSON { raw = evt.Delta.PartialJSON } s.masked.WriteString(raw) // raw model output (text or tool args), for audit @@ -92,19 +98,19 @@ func (s *anthropicCodec) transformEvent(lines [][]byte) []byte { // deltaField returns the JSON field that carries the payload for a delta type. func deltaField(deltaType string) string { - if deltaType == "input_json_delta" { + if deltaType == deltaTypeInputJSON { return "partial_json" } - return "text" + return jsonKeyText } // deltaDataLine builds the `data: {...}\n` line for a text_delta or // input_json_delta carrying the given (restored) payload. func deltaDataLine(index int, deltaType, payload string) []byte { b, _ := json.Marshal(map[string]interface{}{ - "type": "content_block_delta", - "index": index, - "delta": map[string]interface{}{"type": deltaType, deltaField(deltaType): payload}, + jsonKeyType: "content_block_delta", + "index": index, + "delta": map[string]interface{}{jsonKeyType: deltaType, deltaField(deltaType): payload}, }) out := append([]byte("data: "), b...) return append(out, '\n') @@ -114,7 +120,7 @@ func deltaDataLine(index int, deltaType, payload string) []byte { // blank line) used to flush a held-back tail. Defaults to a text delta. func deltaEvent(index int, deltaType, payload string) []byte { if deltaType == "" { - deltaType = "text_delta" + deltaType = deltaTypeText } out := []byte("event: content_block_delta\n") out = append(out, deltaDataLine(index, deltaType, payload)...) diff --git a/src/backend/proxy/codec_openai.go b/src/backend/proxy/codec_openai.go index a85d0275..7696ebb3 100644 --- a/src/backend/proxy/codec_openai.go +++ b/src/backend/proxy/codec_openai.go @@ -37,7 +37,7 @@ func newOpenAICodec(mapping map[string]string) *openaiCodec { // doneFields are the fields a Responses `*.done` event uses to carry the full // value, by event kind: output text, tool-call arguments, and refusals. -var doneFields = []string{"text", "arguments", "refusal"} +var doneFields = []string{jsonKeyText, "arguments", "refusal"} func (c *openaiCodec) transformEvent(lines [][]byte) []byte { dataIdx := -1 @@ -58,7 +58,7 @@ func (c *openaiCodec) transformEvent(lines [][]byte) []byte { if err := json.Unmarshal(raw, &obj); err != nil { return concatLines(lines) // non-JSON payload (e.g. "[DONE]"): pass through } - typ, _ := obj["type"].(string) + typ, _ := obj[jsonKeyType].(string) key := channelKey(obj) switch { @@ -100,7 +100,7 @@ func (c *openaiCodec) transformEvent(lines [][]byte) []byte { // as-is (the carry buffer stores post-restore text). func (c *openaiCodec) tailDeltaEvent(doneType string, obj map[string]interface{}, tail string) []byte { deltaType := strings.TrimSuffix(doneType, ".done") + ".delta" - d := map[string]interface{}{"type": deltaType, "delta": tail} + d := map[string]interface{}{jsonKeyType: deltaType, "delta": tail} for _, f := range []string{"item_id", "output_index", "content_index"} { if v, ok := obj[f]; ok { d[f] = v diff --git a/src/backend/proxy/handler.go b/src/backend/proxy/handler.go index def7de5c..dcb03110 100644 --- a/src/backend/proxy/handler.go +++ b/src/backend/proxy/handler.go @@ -27,6 +27,12 @@ import ( const paramLimit = "limit" +// JSON field names reused across the response wrappers and the SSE codecs. +const ( + jsonKeyType = "type" + jsonKeyText = "text" +) + // Handler handles HTTP requests and proxies them to LLM provider type Handler struct { client *http.Client @@ -280,7 +286,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) { } piiEntities = append(piiEntities, map[string]interface{}{ - "text": entity.Text, + jsonKeyText: entity.Text, "masked_text": maskedText, "label": entity.Label, "confidence": entity.Confidence, @@ -691,7 +697,7 @@ func (h *Handler) LogStreamedResponse(ctx context.Context, transactionID, masked wrap := func(text string) string { b, err := json.Marshal(map[string]interface{}{ "streamed": true, - "content": []map[string]interface{}{{"type": "text", "text": text}}, + "content": []map[string]interface{}{{jsonKeyType: jsonKeyText, jsonKeyText: text}}, }) if err != nil { return text From 427df29dc3dde103dc010e46d6b80e4c07d425f7 Mon Sep 17 00:00:00 2001 From: Hannes Hapke Date: Sun, 28 Jun 2026 16:42:41 -0700 Subject: [PATCH 05/13] updated package-lock.json --- package-lock.json | 52 ++++++++++++++++++++++++++++++++++++++++++++++- package.json | 3 +++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/package-lock.json b/package-lock.json index d8c22e68..33ef9e2b 100644 --- a/package-lock.json +++ b/package-lock.json @@ -11,6 +11,9 @@ "workspaces": [ "src/frontend" ], + "dependencies": { + "shepherd.js": "^15.2.2" + }, "devDependencies": { "@changesets/cli": "^2.27.1", "@electron/notarize": "^3.1.1" @@ -2606,6 +2609,31 @@ "module-details-from-path": "^1.0.4" } }, + "node_modules/@floating-ui/core": { + "version": "1.7.5", + "resolved": "https://registry.npmjs.org/@floating-ui/core/-/core-1.7.5.tgz", + "integrity": "sha512-1Ih4WTWyw0+lKyFMcBHGbb5U5FtuHJuujoyyr5zTaWS5EYMeT6Jb2AuDeftsCsEuchO+mM2ij5+q9crhydzLhQ==", + "license": "MIT", + "dependencies": { + "@floating-ui/utils": "^0.2.11" + } + }, + "node_modules/@floating-ui/dom": { + "version": "1.7.6", + "resolved": "https://registry.npmjs.org/@floating-ui/dom/-/dom-1.7.6.tgz", + "integrity": "sha512-9gZSAI5XM36880PPMm//9dfiEngYoC6Am2izES1FF406YFsjvyBMmeJ2g4SAju3xWwtuynNRFL2s9hgxpLI5SQ==", + "license": "MIT", + "dependencies": { + "@floating-ui/core": "^1.7.5", + "@floating-ui/utils": "^0.2.11" + } + }, + "node_modules/@floating-ui/utils": { + "version": "0.2.11", + "resolved": "https://registry.npmjs.org/@floating-ui/utils/-/utils-0.2.11.tgz", + "integrity": "sha512-RiB/yIh78pcIxl6lLMG0CgBXAZ2Y0eVHqMPYugu+9U0AeT6YBeiJpf7lbdJNIugFP5SIjwNRgo4DhR1Qxi26Gg==", + "license": "MIT" + }, "node_modules/@humanfs/core": { "version": "0.19.1", "dev": true, @@ -6845,6 +6873,15 @@ "dev": true, "license": "MIT" }, + "node_modules/deepmerge-ts": { + "version": "7.1.5", + "resolved": "https://registry.npmjs.org/deepmerge-ts/-/deepmerge-ts-7.1.5.tgz", + "integrity": "sha512-HOJkrhaYsweh+W+e74Yn7YStZOilkoPb6fycpwNLKzSPtruFs48nYis0zy5yJz1+ktUhHxoRDJ27RQAWLIJVJw==", + "license": "BSD-3-Clause", + "engines": { + "node": ">=16.0.0" + } + }, "node_modules/default-browser": { "version": "5.4.0", "dev": true, @@ -12568,6 +12605,19 @@ "url": "https://github.com/sponsors/ljharb" } }, + "node_modules/shepherd.js": { + "version": "15.2.2", + "resolved": "https://registry.npmjs.org/shepherd.js/-/shepherd.js-15.2.2.tgz", + "integrity": "sha512-WDJDjZ8tUG2HF93C9mJCXsdGJ01bpZKjonyMABMkU6I4KjSD9pZ0zFKaeSducdFIx9inKHbDmcFE82uJ65HmgA==", + "license": "AGPL-3.0", + "dependencies": { + "@floating-ui/dom": "^1.7.5", + "deepmerge-ts": "^7.1.5" + }, + "engines": { + "node": ">= 20" + } + }, "node_modules/side-channel": { "version": "1.1.0", "license": "MIT", @@ -14190,7 +14240,7 @@ }, "src/frontend": { "name": "kiji-privacy-proxy", - "version": "1.4.0", + "version": "1.4.1", "dependencies": { "@sentry/electron": "^7.13.0", "@types/dompurify": "^3.2.0", diff --git a/package.json b/package.json index 993b7310..c2691644 100644 --- a/package.json +++ b/package.json @@ -43,5 +43,8 @@ "engines": { "node": ">=18.0.0", "npm": ">=9.0.0" + }, + "dependencies": { + "shepherd.js": "^15.2.2" } } From 5bfbc69e72a6bf7b84aec9b83e47a01d005a8ca2 Mon Sep 17 00:00:00 2001 From: Davidnet Date: Thu, 2 Jul 2026 16:16:29 -0400 Subject: [PATCH 06/13] fix: record dashboard metrics for streamed responses and add streaming tests The rebase onto main brought in the dashboard metrics added in #547, which are recorded after the response is written. The SSE streaming path returns early, so streamed requests (Claude Code, Codex) were invisible to the dashboard; record them there too, with the request-masking time as the proxy-overhead latency (response restoration is interleaved with the upstream stream). Also adds the test coverage the streaming feature landed without: the SSE engine end-to-end over a wire (chunked framing, header rewrite, PII restore), both provider codecs (placeholder split across deltas, tool-call argument restore, per-channel carry buffers, passthrough events), stream sniffing, Codex host routing, and chatgpt.com interception gating. --- src/backend/config/config_test.go | 22 ++ src/backend/providers/providers_test.go | 1 + src/backend/proxy/streaming_test.go | 356 ++++++++++++++++++++++++ src/backend/proxy/transparent.go | 4 + 4 files changed, 383 insertions(+) create mode 100644 src/backend/proxy/streaming_test.go diff --git a/src/backend/config/config_test.go b/src/backend/config/config_test.go index a120f986..5b8e1cc2 100644 --- a/src/backend/config/config_test.go +++ b/src/backend/config/config_test.go @@ -354,6 +354,28 @@ func stringContains(s, substr string) bool { return strings.Contains(s, substr) } +// chatgpt.com is the ChatGPT-login Codex host; it must be intercepted whenever +// the OpenAI provider is configured, and only then. +func TestGetInterceptDomains_CodexHost(t *testing.T) { + pc := DefaultConfig().Providers + found := false + for _, d := range pc.GetInterceptDomains() { + if d == "chatgpt.com" { + found = true + } + } + if !found { + t.Errorf("GetInterceptDomains() = %v, want it to include chatgpt.com when OpenAI is configured", pc.GetInterceptDomains()) + } + + pc.OpenAIProviderConfig.APIDomain = "" + for _, d := range pc.GetInterceptDomains() { + if d == "chatgpt.com" { + t.Error("chatgpt.com intercepted even though OpenAI provider is disabled") + } + } +} + func TestDefaultConfig_Detectors(t *testing.T) { cfg := DefaultConfig() want := []string{DetectorTypeONNX, DetectorTypeRegex} diff --git a/src/backend/providers/providers_test.go b/src/backend/providers/providers_test.go index 47567af6..2d789179 100644 --- a/src/backend/providers/providers_test.go +++ b/src/backend/providers/providers_test.go @@ -1511,6 +1511,7 @@ func TestProviders_GetProviderFromHost(t *testing.T) { }{ {"OpenAI host", "api.openai.com", "OpenAI", false}, {"OpenAI host with port", "api.openai.com:443", "OpenAI", false}, + {"Codex (ChatGPT-login) host routes to OpenAI", "chatgpt.com", "OpenAI", false}, {"Anthropic host", "api.anthropic.com", "Anthropic", false}, {"Gemini host", "generativelanguage.googleapis.com", "Gemini", false}, {"Mistral host", "api.mistral.ai", "Mistral", false}, diff --git a/src/backend/proxy/streaming_test.go b/src/backend/proxy/streaming_test.go new file mode 100644 index 00000000..2360ca39 --- /dev/null +++ b/src/backend/proxy/streaming_test.go @@ -0,0 +1,356 @@ +package proxy + +import ( + "bufio" + "encoding/json" + "io" + "net" + "net/http" + "strings" + "testing" +) + +// --- restoreCore / splitSafe --- + +func TestSplitSafe(t *testing.T) { + tests := []struct { + name string + s string + keep int + emit, hold string + }{ + {"shorter than keep is all held", "ab", 5, "", "ab"}, + {"exactly keep is all held", "abcde", 5, "", "abcde"}, + {"longer than keep splits at tail", "abcdefgh", 3, "abcde", "fgh"}, + {"keep zero emits everything", "abc", 0, "abc", ""}, + {"empty string", "", 3, "", ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + emit, hold := splitSafe(tt.s, tt.keep) + if emit != tt.emit || hold != tt.hold { + t.Errorf("splitSafe(%q, %d) = (%q, %q), want (%q, %q)", + tt.s, tt.keep, emit, hold, tt.emit, tt.hold) + } + }) + } +} + +func TestRestoreCore(t *testing.T) { + core := newRestoreCore(map[string]string{ + "John Smith": "Jane Doe", + "Bob": "Alice", + }) + // keep = longest dummy - 1 + if core.keep != len("John Smith")-1 { + t.Errorf("keep = %d, want %d", core.keep, len("John Smith")-1) + } + got := core.restore("Hi John Smith, meet Bob.") + want := "Hi Jane Doe, meet Alice." + if got != want { + t.Errorf("restore = %q, want %q", got, want) + } + + empty := newRestoreCore(nil) + if empty.keep != 0 { + t.Errorf("empty mapping keep = %d, want 0", empty.keep) + } + if got := empty.restore("unchanged"); got != "unchanged" { + t.Errorf("restore with empty mapping = %q", got) + } +} + +// --- helpers --- + +// event turns raw SSE text into the line slices transformEvent receives. +func sseLines(raw string) [][]byte { + var lines [][]byte + for _, ln := range strings.SplitAfter(raw, "\n") { + if ln != "" { + lines = append(lines, []byte(ln)) + } + } + return lines +} + +// deltaPayloads concatenates the text/partial_json/delta payloads of every +// data: line in emitted SSE bytes — i.e. the text a streaming client would +// assemble. Restored values may straddle event boundaries, so assertions on +// restored text must run against this, not the raw frames. +func deltaPayloads(t *testing.T, raw string) string { + t.Helper() + var out strings.Builder + for _, ln := range strings.Split(raw, "\n") { + if !strings.HasPrefix(ln, "data: ") { + continue + } + var evt struct { + Delta json.RawMessage `json:"delta"` + } + if json.Unmarshal([]byte(ln[len("data: "):]), &evt) != nil || evt.Delta == nil { + continue + } + // Anthropic: delta is an object with text/partial_json. OpenAI: delta is a string. + var s string + if json.Unmarshal(evt.Delta, &s) == nil { + out.WriteString(s) + continue + } + var obj struct { + Text string `json:"text"` + PartialJSON string `json:"partial_json"` + } + if json.Unmarshal(evt.Delta, &obj) == nil { + out.WriteString(obj.Text) + out.WriteString(obj.PartialJSON) + } + } + return out.String() +} + +// --- anthropicCodec --- + +func TestAnthropicCodec_PlaceholderSplitAcrossDeltas(t *testing.T) { + codec := newAnthropicCodec(map[string]string{"John Smith": "Jane Doe"}) + + // The placeholder arrives split across two deltas. + ev1 := codec.transformEvent(sseLines( + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello John S"}}` + "\n\n")) + ev2 := codec.transformEvent(sseLines( + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"mith, bye"}}` + "\n\n")) + stop := codec.transformEvent(sseLines( + "event: content_block_stop\n" + + `data: {"type":"content_block_stop","index":0}` + "\n\n")) + + combined := string(ev1) + string(ev2) + string(stop) + assembled := deltaPayloads(t, combined) + if strings.Contains(assembled, "John") { + t.Errorf("masked name leaked to client: %s", assembled) + } + if assembled != "Hello Jane Doe, bye" { + t.Errorf("client-assembled text = %q, want %q", assembled, "Hello Jane Doe, bye") + } + // Every emitted delta must remain a well-formed SSE event (framing intact). + if !strings.HasPrefix(string(ev1), "event: content_block_delta\ndata: ") { + t.Errorf("event framing broken: %q", ev1) + } + // Audit accumulator holds the raw (masked) model output. + if got := codec.maskedOutput(); got != "Hello John Smith, bye" { + t.Errorf("maskedOutput = %q", got) + } +} + +func TestAnthropicCodec_ToolArgsRestoredInInputJSONDelta(t *testing.T) { + codec := newAnthropicCodec(map[string]string{"MASKED_EMAIL": "real@example.com"}) + + ev := codec.transformEvent(sseLines( + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":1,"delta":{"type":"input_json_delta","partial_json":"{\"to\":\"MASKED_EMAIL\"}"}}` + "\n\n")) + stop := codec.transformEvent(sseLines( + "event: content_block_stop\n" + + `data: {"type":"content_block_stop","index":1}` + "\n\n")) + + combined := string(ev) + string(stop) + assembled := deltaPayloads(t, combined) + if strings.Contains(assembled, "MASKED_EMAIL") { + t.Errorf("masked value leaked in tool args: %s", assembled) + } + if assembled != `{"to":"real@example.com"}` { + t.Errorf("client-assembled tool args = %q", assembled) + } + // The flushed tail must keep the input_json_delta type, not fall back to text. + if !strings.Contains(combined, "input_json_delta") { + t.Errorf("flushed tail lost its delta type: %s", combined) + } +} + +func TestAnthropicCodec_PassthroughEvents(t *testing.T) { + codec := newAnthropicCodec(map[string]string{"X": "Y"}) + raw := "event: message_start\n" + + `data: {"type":"message_start","message":{"id":"msg_1"}}` + "\n\n" + got := codec.transformEvent(sseLines(raw)) + if string(got) != raw { + t.Errorf("non-delta event modified:\ngot %q\nwant %q", got, raw) + } +} + +func TestAnthropicCodec_IndependentContentBlocks(t *testing.T) { + codec := newAnthropicCodec(map[string]string{"LONGPLACEHOLDER": "short"}) + // Block 0 and block 1 interleave; carry buffers must not mix. + codec.transformEvent(sseLines( + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"LONGPLACE"}}` + "\n\n")) + codec.transformEvent(sseLines( + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"other text entirely"}}` + "\n\n")) + ev := codec.transformEvent(sseLines( + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"HOLDER done"}}` + "\n\n")) + stop := codec.transformEvent(sseLines( + "event: content_block_stop\n" + + `data: {"type":"content_block_stop","index":0}` + "\n\n")) + + combined := string(ev) + string(stop) + if strings.Contains(combined, "LONGPLACEHOLDER") { + t.Errorf("placeholder leaked when split across block-0 deltas: %s", combined) + } + if !strings.Contains(combined, "short") { + t.Errorf("restored value missing: %s", combined) + } +} + +// --- openaiCodec --- + +func TestOpenAICodec_DeltaAndDone(t *testing.T) { + codec := newOpenAICodec(map[string]string{"John Smith": "Jane Doe"}) + + ev1 := codec.transformEvent(sseLines( + "event: response.output_text.delta\n" + + `data: {"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"Hi John S"}` + "\n\n")) + ev2 := codec.transformEvent(sseLines( + "event: response.output_text.delta\n" + + `data: {"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"mith!"}` + "\n\n")) + done := codec.transformEvent(sseLines( + "event: response.output_text.done\n" + + `data: {"type":"response.output_text.done","output_index":0,"content_index":0,"text":"Hi John Smith!"}` + "\n\n")) + + combined := string(ev1) + string(ev2) + string(done) + if strings.Contains(combined, "John") { + t.Errorf("masked name leaked: %s", combined) + } + // Restored text must appear both in the flushed deltas and in the .done payload. + if strings.Count(combined, "Jane Doe") < 2 { + t.Errorf("expected restored name in deltas and .done: %s", combined) + } + if got := codec.maskedOutput(); got != "Hi John Smith!" { + t.Errorf("maskedOutput = %q", got) + } +} + +func TestOpenAICodec_FunctionCallArguments(t *testing.T) { + codec := newOpenAICodec(map[string]string{"MASKED_EMAIL": "real@example.com"}) + + ev := codec.transformEvent(sseLines( + "event: response.function_call_arguments.delta\n" + + `data: {"type":"response.function_call_arguments.delta","output_index":0,"delta":"{\"to\":\"MASKED_EMAIL\"}"}` + "\n\n")) + done := codec.transformEvent(sseLines( + "event: response.function_call_arguments.done\n" + + `data: {"type":"response.function_call_arguments.done","output_index":0,"arguments":"{\"to\":\"MASKED_EMAIL\"}"}` + "\n\n")) + + combined := string(ev) + string(done) + if strings.Contains(combined, "MASKED_EMAIL") { + t.Errorf("masked value leaked in tool args: %s", combined) + } + if !strings.Contains(combined, "real@example.com") { + t.Errorf("restored value missing: %s", combined) + } +} + +func TestOpenAICodec_PassthroughDoneSentinelAndUnknown(t *testing.T) { + codec := newOpenAICodec(map[string]string{"X": "Y"}) + for _, raw := range []string{ + "data: [DONE]\n\n", + "event: response.created\n" + `data: {"type":"response.created","response":{"id":"resp_1"}}` + "\n\n", + ": keep-alive comment\n\n", + } { + got := codec.transformEvent(sseLines(raw)) + if string(got) != raw { + t.Errorf("passthrough event modified:\ngot %q\nwant %q", got, raw) + } + } +} + +// --- request/response sniffing --- + +func TestRequestWantsStream(t *testing.T) { + tests := []struct { + body string + want bool + }{ + {`{"stream":true,"model":"claude"}`, true}, + {`{"stream":false}`, false}, + {`{"model":"claude"}`, false}, + {`not json`, false}, + {`{"stream":"true"}`, false}, // wrong type + } + for _, tt := range tests { + if got := requestWantsStream([]byte(tt.body)); got != tt.want { + t.Errorf("requestWantsStream(%q) = %v, want %v", tt.body, got, tt.want) + } + } +} + +func TestIsEventStream(t *testing.T) { + resp := &http.Response{Header: http.Header{"Content-Type": []string{"text/event-stream; charset=utf-8"}}} + if !isEventStream(resp) { + t.Error("text/event-stream not detected") + } + resp.Header.Set("Content-Type", "application/json") + if isEventStream(resp) { + t.Error("application/json misdetected as event stream") + } +} + +// --- streamSSEResponse end-to-end --- + +func TestStreamSSEResponse_EndToEnd(t *testing.T) { + upstream := "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello John Smith, welcome"}}` + "\n\n" + + "event: content_block_stop\n" + + `data: {"type":"content_block_stop","index":0}` + "\n\n" + + "event: message_stop\n" + + `data: {"type":"message_stop"}` + "\n\n" + + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "Content-Length": []string{"9999"}, // must be dropped in favor of chunked + }, + Body: io.NopCloser(strings.NewReader(upstream)), + } + + client, server := net.Pipe() + codec := newAnthropicCodec(map[string]string{"John Smith": "Jane Doe"}) + + errCh := make(chan error, 1) + go func() { + errCh <- streamSSEResponse(server, resp, codec) + server.Close() + }() + + // Parse what a real HTTP client would see on the wire. + parsed, err := http.ReadResponse(bufio.NewReader(client), nil) + if err != nil { + t.Fatalf("client failed to parse response: %v", err) + } + body, err := io.ReadAll(parsed.Body) + if err != nil { + t.Fatalf("client failed to read body: %v", err) + } + if err := <-errCh; err != nil { + t.Fatalf("streamSSEResponse returned error: %v", err) + } + + if parsed.StatusCode != http.StatusOK { + t.Errorf("status = %d", parsed.StatusCode) + } + if got := parsed.Header.Get("Content-Type"); got != "text/event-stream" { + t.Errorf("Content-Type = %q", got) + } + if parsed.ContentLength >= 0 { + t.Errorf("Content-Length should be dropped for chunked streaming, got %d", parsed.ContentLength) + } + text := string(body) + if strings.Contains(text, "John Smith") { + t.Errorf("masked name leaked to client:\n%s", text) + } + if !strings.Contains(text, "Jane Doe") { + t.Errorf("restored name missing:\n%s", text) + } + if !strings.Contains(text, "message_stop") { + t.Errorf("terminal event missing:\n%s", text) + } +} diff --git a/src/backend/proxy/transparent.go b/src/backend/proxy/transparent.go index 9e49d986..87895e9e 100644 --- a/src/backend/proxy/transparent.go +++ b/src/backend/proxy/transparent.go @@ -489,6 +489,10 @@ func (tp *TransparentProxy) interceptHTTPOverTLS(conn net.Conn, r *http.Request, maskedText := codec.maskedOutput() tp.handler.LogStreamedResponse(ctx, processed.TransactionID, maskedText, codec.restore(maskedText)) log.Printf("[TransparentProxy] Streamed %s %s - Status: %d", r.Method, r.URL.Path, resp.StatusCode) + // Count streamed requests in the dashboard metrics too. Response + // restoration is interleaved with the upstream stream here, so the + // recorded overhead is the request-masking time only. + tp.handler.recordMetrics(provider, processed, sourceFromRequest(r), requestMaskTime, resp.StatusCode) return } From 6235d4a0752ac1fa4832d39ecc65705a0a583ebc Mon Sep 17 00:00:00 2001 From: Davidnet Date: Thu, 2 Jul 2026 16:23:18 -0400 Subject: [PATCH 07/13] chore: drop stray shepherd.js root dependency shepherd.js is a frontend dependency already declared in src/frontend/package.json; the root entry was added accidentally. Regenerating the lockfile also syncs the recorded src/frontend version with main (1.4.0 -> 1.4.1). --- package-lock.json | 50 ----------------------------------------------- package.json | 3 --- 2 files changed, 53 deletions(-) diff --git a/package-lock.json b/package-lock.json index 33ef9e2b..c0802b61 100644 --- a/package-lock.json +++ b/package-lock.json @@ -11,9 +11,6 @@ "workspaces": [ "src/frontend" ], - "dependencies": { - "shepherd.js": "^15.2.2" - }, "devDependencies": { "@changesets/cli": "^2.27.1", "@electron/notarize": "^3.1.1" @@ -2609,31 +2606,6 @@ "module-details-from-path": "^1.0.4" } }, - "node_modules/@floating-ui/core": { - "version": "1.7.5", - "resolved": "https://registry.npmjs.org/@floating-ui/core/-/core-1.7.5.tgz", - "integrity": "sha512-1Ih4WTWyw0+lKyFMcBHGbb5U5FtuHJuujoyyr5zTaWS5EYMeT6Jb2AuDeftsCsEuchO+mM2ij5+q9crhydzLhQ==", - "license": "MIT", - "dependencies": { - "@floating-ui/utils": "^0.2.11" - } - }, - "node_modules/@floating-ui/dom": { - "version": "1.7.6", - "resolved": "https://registry.npmjs.org/@floating-ui/dom/-/dom-1.7.6.tgz", - "integrity": "sha512-9gZSAI5XM36880PPMm//9dfiEngYoC6Am2izES1FF406YFsjvyBMmeJ2g4SAju3xWwtuynNRFL2s9hgxpLI5SQ==", - "license": "MIT", - "dependencies": { - "@floating-ui/core": "^1.7.5", - "@floating-ui/utils": "^0.2.11" - } - }, - "node_modules/@floating-ui/utils": { - "version": "0.2.11", - "resolved": "https://registry.npmjs.org/@floating-ui/utils/-/utils-0.2.11.tgz", - "integrity": "sha512-RiB/yIh78pcIxl6lLMG0CgBXAZ2Y0eVHqMPYugu+9U0AeT6YBeiJpf7lbdJNIugFP5SIjwNRgo4DhR1Qxi26Gg==", - "license": "MIT" - }, "node_modules/@humanfs/core": { "version": "0.19.1", "dev": true, @@ -6873,15 +6845,6 @@ "dev": true, "license": "MIT" }, - "node_modules/deepmerge-ts": { - "version": "7.1.5", - "resolved": "https://registry.npmjs.org/deepmerge-ts/-/deepmerge-ts-7.1.5.tgz", - "integrity": "sha512-HOJkrhaYsweh+W+e74Yn7YStZOilkoPb6fycpwNLKzSPtruFs48nYis0zy5yJz1+ktUhHxoRDJ27RQAWLIJVJw==", - "license": "BSD-3-Clause", - "engines": { - "node": ">=16.0.0" - } - }, "node_modules/default-browser": { "version": "5.4.0", "dev": true, @@ -12605,19 +12568,6 @@ "url": "https://github.com/sponsors/ljharb" } }, - "node_modules/shepherd.js": { - "version": "15.2.2", - "resolved": "https://registry.npmjs.org/shepherd.js/-/shepherd.js-15.2.2.tgz", - "integrity": "sha512-WDJDjZ8tUG2HF93C9mJCXsdGJ01bpZKjonyMABMkU6I4KjSD9pZ0zFKaeSducdFIx9inKHbDmcFE82uJ65HmgA==", - "license": "AGPL-3.0", - "dependencies": { - "@floating-ui/dom": "^1.7.5", - "deepmerge-ts": "^7.1.5" - }, - "engines": { - "node": ">= 20" - } - }, "node_modules/side-channel": { "version": "1.1.0", "license": "MIT", diff --git a/package.json b/package.json index c2691644..993b7310 100644 --- a/package.json +++ b/package.json @@ -43,8 +43,5 @@ "engines": { "node": ">=18.0.0", "npm": ">=9.0.0" - }, - "dependencies": { - "shepherd.js": "^15.2.2" } } From 3f89060b2df876f766ceb9b42c8e00319e15e282 Mon Sep 17 00:00:00 2001 From: Davidnet Date: Fri, 3 Jul 2026 14:37:18 -0400 Subject: [PATCH 08/13] fix: support ChatGPT-login Codex through the transparent proxy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex CLI startup and completions both broke when routed through the MITM proxy. Five fixes, verified end-to-end against a live Codex run: - Gate chatgpt.com interception by path: only /backend-api/codex/responses is masked; the streaming MCP transport, model list, and telemetry pass through verbatim via a new in-tunnel passthrough that streams incrementally and preserves framing across keep-alive requests. - Pin relayed responses to HTTP/1.1: upstream legs negotiate HTTP/2, and relaying resp.Proto verbatim produced "HTTP/2.0" status lines that strict clients (Codex's hyper) reject as malformed. - Create the MITM loop's bufio.Reader once per connection so buffered bytes of a pipelined next request are not dropped between iterations. - Sniff SSE when Content-Type is absent: the Codex backend streams SSE with no Content-Type header, which sent responses down the buffered path and skipped PII restoration. - Restore PII in nested .done/.completed Responses-API payloads (content_part.done, output_item.done, response.completed) — Codex renders its final message from these, not from the flat deltas. - Reject WebSocket upgrades on masked endpoints (fail closed): PII in switched protocols cannot be masked. Codex is steered to SSE via a provider config with supports_websockets=false, documented in docs/09-coding-agents.md. --- docs/09-coding-agents.md | 65 +++++- docs/README.md | 2 +- src/backend/config/config.go | 13 ++ src/backend/config/config_test.go | 23 +++ src/backend/providers/openai.go | 5 + src/backend/proxy/codec_openai.go | 40 +++- src/backend/proxy/cors.go | 2 +- src/backend/proxy/passthrough.go | 202 +++++++++++++++++++ src/backend/proxy/passthrough_test.go | 245 +++++++++++++++++++++++ src/backend/proxy/router.go | 40 +++- src/backend/proxy/router_test.go | 81 ++++++++ src/backend/proxy/streaming.go | 37 ++++ src/backend/proxy/streaming_test.go | 78 ++++++++ src/backend/proxy/transparent.go | 58 ++++-- src/backend/proxy/transparent_factory.go | 2 +- 15 files changed, 862 insertions(+), 31 deletions(-) create mode 100644 src/backend/proxy/passthrough.go create mode 100644 src/backend/proxy/passthrough_test.go create mode 100644 src/backend/proxy/router_test.go diff --git a/docs/09-coding-agents.md b/docs/09-coding-agents.md index 5ce2e9ab..5da7c592 100644 --- a/docs/09-coding-agents.md +++ b/docs/09-coding-agents.md @@ -14,7 +14,7 @@ Hosts the proxy intercepts for coding agents: |-------|---------|-------| | Claude Code | `api.anthropic.com` | | | Codex (API key) | `api.openai.com` | `/v1/responses` and `/v1/chat/completions` | -| Codex (ChatGPT login) | `chatgpt.com` | `/backend-api/codex/responses` | +| Codex (ChatGPT login) | `chatgpt.com` | Only `/backend-api/codex/responses` is masked; other paths (MCP, model list, telemetry) pass through verbatim. Requires forcing SSE — see [below](#chatgpt-login-codex-force-sse-websockets-cannot-be-masked) | For **any** agent, two things must be true: @@ -100,20 +100,69 @@ This means if you already export `NODE_EXTRA_CA_CERTS` globally for Claude Code, - **API-key Codex** (`OPENAI_API_KEY` set) talks to `api.openai.com`. Your API key is forwarded untouched. - **ChatGPT-login Codex** (signed in with `codex login`) talks to `chatgpt.com/backend-api/codex/responses` with an OAuth bearer token. The proxy leaves the `Authorization` header untouched and only masks/restores content, so your session keeps working. -Both are intercepted with the same setup above; no extra configuration is needed to switch between them. +Both use the environment variables above, but ChatGPT-login Codex needs **one extra piece of configuration** — see the next section. -## A shared snippet for both agents +### ChatGPT-login Codex: force SSE (WebSockets cannot be masked) -Drop this in your shell profile (`~/.zshrc` / `~/.bashrc`) to cover both agents at once on macOS: +Recent Codex versions stream completions over a **WebSocket** by default when signed in with ChatGPT. The prompt then travels inside WebSocket frames, which the proxy cannot inspect or mask — so the proxy **rejects protocol upgrades on masked endpoints** (fail-closed; you'll see Codex print `ERROR: Reconnecting... n/5` and the proxy log `Rejecting "websocket" upgrade on intercepted path`). + +The fix is to steer Codex back to SSE. Its built-in `openai` provider cannot be overridden, so define a custom provider that points at the same ChatGPT backend with WebSockets disabled. Add to `~/.codex/config.toml`: + +```toml +model_provider = "kiji" + +[model_providers.kiji] +name = "Kiji" +base_url = "https://chatgpt.com/backend-api/codex" +wire_api = "responses" +requires_openai_auth = true +supports_websockets = false +``` + +Your ChatGPT login keeps working (`requires_openai_auth = true` reuses the existing OAuth session). For a one-off run without touching `config.toml`: + +```bash +codex exec \ + -c 'model_provider="kiji"' \ + -c 'model_providers.kiji={ name = "Kiji", base_url = "https://chatgpt.com/backend-api/codex", wire_api = "responses", requires_openai_auth = true, supports_websockets = false }' \ + "your prompt" +``` + +The proxy/CA environment variables must still be set **in the same shell that runs `codex`** — the provider config alone does not route traffic through the proxy: + +| Variable | Value | Why | +|----------|-------|-----| +| `HTTP_PROXY` | `http://127.0.0.1:8081` | Routes Codex's HTTP traffic through the proxy | +| `HTTPS_PROXY` | `http://127.0.0.1:8081` | Routes Codex's HTTPS traffic (the completions call) through the proxy | +| `CODEX_CA_CERTIFICATE` | path to `ca.crt` (see [Prerequisites](#prerequisites)) | Makes Codex trust the proxy's MITM certificate | + +> **Warning — bypass is silent.** If the proxy variables are missing, Codex talks to `chatgpt.com` directly: everything works and looks identical, but nothing is masked. Verify interception after setup (see [Verifying interception](#verifying-interception)) — a correct-looking reply with an **empty request log** means your prompt went out unmasked. + +Only the completions path (`/backend-api/codex/responses`) is masked. Codex's other `chatgpt.com` traffic — its MCP transport, model list, telemetry — is passed through the proxy verbatim so the CLI works normally; none of it carries your prompt. + +## The environment variables at a glance + +Four variables cover both agents. Set them in the shell that runs the agent (or persist them in your shell profile — `~/.zshrc`, `~/.bashrc`, or `~/.config/fish/config.fish`): + +| Variable | Value | Needed by | Why | +|----------|-------|-----------|-----| +| `HTTP_PROXY` | `http://127.0.0.1:8081` | both | Routes the agent's HTTP traffic through the proxy | +| `HTTPS_PROXY` | `http://127.0.0.1:8081` | both | Routes the agent's HTTPS traffic (the actual API calls) through the proxy | +| `NODE_EXTRA_CA_CERTS` | path to `ca.crt` | Claude Code | Node.js reads extra trusted CAs from this variable; without it Claude Code rejects the proxy's MITM certificate. (Codex also honors it, but only as a fallback.) | +| `CODEX_CA_CERTIFICATE` | path to `ca.crt` | Codex | Codex's native CA variable for its rustls TLS stack | + +Example for a bash/zsh profile on macOS: ```bash KIJI_CA="$HOME/Library/Application Support/Kiji Privacy Proxy/certs/ca.crt" export HTTP_PROXY=http://127.0.0.1:8081 export HTTPS_PROXY=http://127.0.0.1:8081 -export NODE_EXTRA_CA_CERTS="$KIJI_CA" # Claude Code (and Codex fallback) -export CODEX_CA_CERTIFICATE="$KIJI_CA" # Codex (explicit) +export NODE_EXTRA_CA_CERTS="$KIJI_CA" # Claude Code +export CODEX_CA_CERTIFICATE="$KIJI_CA" # Codex ``` +If you only use one agent, you only need its CA variable (plus the two proxy variables). Note that with these set globally, any tool honoring `HTTP(S)_PROXY` fails when the proxy isn't running — unset them (or comment them out) if you stop Kiji. + ## What gets masked and restored | Direction | Covered | @@ -147,6 +196,10 @@ Every interception is recorded in the proxy's request log (visible in the deskto - *Cause:* the agent isn't using the proxy. - *Fix:* ensure `HTTP_PROXY` and `HTTPS_PROXY` are exported in the **same shell/process** that runs the agent. Check that `NO_PROXY`/`no_proxy` doesn't list `openai.com`, `chatgpt.com`, or `anthropic.com`. Confirm the proxy is listening on the port you set. +**Codex prints `ERROR: Reconnecting... n/5` and the turn never completes** +- *Cause:* Codex is trying to stream the completion over a WebSocket, which the proxy rejects because WebSocket frames cannot be masked. The proxy log shows `Rejecting "websocket" upgrade on intercepted path`. +- *Fix:* force SSE with the custom provider config — see [ChatGPT-login Codex: force SSE](#chatgpt-login-codex-force-sse-websockets-cannot-be-masked). + **ChatGPT-login Codex still fails after setup** - *Cause:* the proxy build doesn't intercept `chatgpt.com`, or the CA isn't trusted on that host. - *Fix:* verify `chatgpt.com` is in the proxy's intercept domains (it is added automatically when the OpenAI provider is configured) and that `CODEX_CA_CERTIFICATE` is set. As a last resort for diagnosis you can test with API-key Codex against `api.openai.com` to isolate whether the issue is host-specific. diff --git a/docs/README.md b/docs/README.md index 64a7c2eb..d0fc73d2 100644 --- a/docs/README.md +++ b/docs/README.md @@ -206,7 +206,7 @@ Route terminal coding agents through the proxy so PII in prompts, code, and tool ### Coding Agents - [Claude Code Setup](09-coding-agents.md#claude-code) - [Codex Setup](09-coding-agents.md#codex) -- [Shared Shell Snippet](09-coding-agents.md#a-shared-snippet-for-both-agents) +- [Environment Variables](09-coding-agents.md#the-environment-variables-at-a-glance) - [Troubleshooting](09-coding-agents.md#troubleshooting) ## Document Status diff --git a/src/backend/config/config.go b/src/backend/config/config.go index f12c40cb..37783592 100644 --- a/src/backend/config/config.go +++ b/src/backend/config/config.go @@ -358,6 +358,19 @@ func (pc ProvidersConfig) GetInterceptDomains() []string { return domains } +// GetInterceptPathPrefixes returns, per intercept host, the path prefixes that +// should be intercepted/masked. Hosts absent from the map intercept all paths. +// chatgpt.com serves far more than the Codex completions endpoint (a streaming +// MCP transport, model refresh, telemetry); only the completions path is +// masked — the rest is passed through verbatim so Codex startup isn't broken. +func (pc ProvidersConfig) GetInterceptPathPrefixes() map[string][]string { + prefixes := map[string][]string{} + if pc.OpenAIProviderConfig.APIDomain != "" { + prefixes[providers.ProviderAPIDomainCodex] = []string{providers.ProviderSubpathCodexResponses} + } + return prefixes +} + func interceptDomain(apiDomain string) string { if apiDomain == "" { return "" diff --git a/src/backend/config/config_test.go b/src/backend/config/config_test.go index 5b8e1cc2..0b4f208d 100644 --- a/src/backend/config/config_test.go +++ b/src/backend/config/config_test.go @@ -5,6 +5,8 @@ import ( "reflect" "strings" "testing" + + "github.com/dataiku/kiji-proxy/src/backend/providers" ) func TestValidatePort(t *testing.T) { @@ -376,6 +378,27 @@ func TestGetInterceptDomains_CodexHost(t *testing.T) { } } +// Only the Codex completions path on chatgpt.com should be masked; the host's +// other endpoints (MCP transport, telemetry) must pass through. Hosts other +// than chatgpt.com carry no allowlist (all paths intercepted). +func TestGetInterceptPathPrefixes(t *testing.T) { + pc := DefaultConfig().Providers + + prefixes := pc.GetInterceptPathPrefixes() + want := []string{providers.ProviderSubpathCodexResponses} + if got := prefixes[providers.ProviderAPIDomainCodex]; !reflect.DeepEqual(got, want) { + t.Errorf("GetInterceptPathPrefixes()[chatgpt.com] = %v, want %v", got, want) + } + if len(prefixes) != 1 { + t.Errorf("GetInterceptPathPrefixes() has %d entries, want only chatgpt.com: %v", len(prefixes), prefixes) + } + + pc.OpenAIProviderConfig.APIDomain = "" + if prefixes := pc.GetInterceptPathPrefixes(); len(prefixes) != 0 { + t.Errorf("GetInterceptPathPrefixes() = %v with OpenAI disabled, want empty", prefixes) + } +} + func TestDefaultConfig_Detectors(t *testing.T) { cfg := DefaultConfig() want := []string{DetectorTypeONNX, DetectorTypeRegex} diff --git a/src/backend/providers/openai.go b/src/backend/providers/openai.go index 10418a3e..d7be205a 100644 --- a/src/backend/providers/openai.go +++ b/src/backend/providers/openai.go @@ -21,6 +21,11 @@ const ( // token instead of api.openai.com, so it must be routed to and intercepted by // the OpenAI provider alongside the API-key host. ProviderAPIDomainCodex string = "chatgpt.com" + // ProviderSubpathCodexResponses is the completions endpoint ChatGPT-login + // Codex posts to on chatgpt.com. It is the only chatgpt.com path that should + // be intercepted/masked — everything else on that host (the streaming MCP + // transport, model refresh, telemetry) must be passed through verbatim. + ProviderSubpathCodexResponses string = "/backend-api/codex/responses" ) // reasoningModelFamilies lists OpenAI model family prefixes that require the diff --git a/src/backend/proxy/codec_openai.go b/src/backend/proxy/codec_openai.go index 7696ebb3..cfe8ba1a 100644 --- a/src/backend/proxy/codec_openai.go +++ b/src/backend/proxy/codec_openai.go @@ -73,7 +73,7 @@ func (c *openaiCodec) transformEvent(lines [][]byte) []byte { obj["delta"] = emit return rewriteDataLine(lines, dataIdx, obj) - case strings.HasSuffix(typ, ".done"): + case strings.HasSuffix(typ, ".done"), strings.HasSuffix(typ, ".completed"): var out []byte // Flush any held-back delta tail as a synthetic delta so incremental // renderers see the complete text before the terminating .done event. @@ -81,13 +81,13 @@ func (c *openaiCodec) transformEvent(lines [][]byte) []byte { delete(c.carry, key) out = append(out, c.tailDeltaEvent(typ, obj, tail)...) } - // Restore the full value the .done event repeats. (Not added to the audit + // Restore the full values the event repeats, including nested copies: + // content_part.done nests part.text, output_item.done nests + // item.content[].text, and response.completed nests the whole response + // with output[].content[].text — clients (e.g. Codex) render their final + // message from these, not from the deltas. (Not added to the audit // accumulator: the deltas above already captured this text.) - for _, field := range doneFields { - if full, ok := obj[field].(string); ok { - obj[field] = c.restore(full) - } - } + restoreTextFields(obj, c.restore) return append(out, rewriteDataLine(lines, dataIdx, obj)...) default: @@ -95,6 +95,32 @@ func (c *openaiCodec) transformEvent(lines [][]byte) []byte { } } +// restoreTextFields recursively restores PII in every string value held by a +// known text-carrying field (doneFields) anywhere in the event payload. Only +// those fields are rewritten — ids, signatures, and other opaque strings that +// could contain a dummy value as a substring pass through untouched. +func restoreTextFields(v interface{}, restore func(string) string) { + switch t := v.(type) { + case map[string]interface{}: + for k, val := range t { + if s, ok := val.(string); ok { + for _, field := range doneFields { + if k == field { + t[k] = restore(s) + break + } + } + continue + } + restoreTextFields(val, restore) + } + case []interface{}: + for _, item := range t { + restoreTextFields(item, restore) + } + } +} + // tailDeltaEvent builds a synthetic `*.delta` event mirroring a `*.done` event's // channel, carrying the already-restored held-back tail. The payload is emitted // as-is (the carry buffer stores post-restore text). diff --git a/src/backend/proxy/cors.go b/src/backend/proxy/cors.go index 6105e1c1..232fa0de 100644 --- a/src/backend/proxy/cors.go +++ b/src/backend/proxy/cors.go @@ -51,7 +51,7 @@ func writeCORSPreflightOverTLS(conn net.Conn, r *http.Request) { resp := &http.Response{ StatusCode: http.StatusNoContent, Status: http.StatusText(http.StatusNoContent), - Proto: "HTTP/1.1", + Proto: protoHTTP11, ProtoMajor: 1, ProtoMinor: 1, Header: make(http.Header), diff --git a/src/backend/proxy/passthrough.go b/src/backend/proxy/passthrough.go new file mode 100644 index 00000000..f07fae76 --- /dev/null +++ b/src/backend/proxy/passthrough.go @@ -0,0 +1,202 @@ +package proxy + +import ( + "io" + "log" + "net" + "net/http" + "time" +) + +// passthroughClient forwards non-intercepted in-tunnel requests. Like +// streamingClient it has no overall timeout so long-lived streams (e.g. the +// Codex MCP transport) are not cut off; a response-header timeout still guards +// against a dead upstream. Compression is disabled so the response bytes are +// relayed exactly as the upstream sent them, and redirects are returned to the +// client rather than followed. Proxy is nil to avoid looping back through +// ourselves. +var passthroughClient = &http.Client{ + Transport: &http.Transport{ + Proxy: nil, + DisableCompression: true, + ResponseHeaderTimeout: 60 * time.Second, + }, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, +} + +// protoHTTP11 is the protocol written on every response relayed to the client +// side of the tunnel, which always speaks HTTP/1.1 regardless of what the +// upstream leg negotiated. +const protoHTTP11 = "HTTP/1.1" + +// hopByHopHeaders are connection-level headers that must not be forwarded +// upstream (RFC 9110 §7.6.1). +var hopByHopHeaders = []string{ + "Connection", + "Keep-Alive", + "Proxy-Authenticate", + "Proxy-Authorization", + "Proxy-Connection", + "Te", + "Trailer", + "Transfer-Encoding", + "Upgrade", +} + +// passthroughHTTPOverTLS forwards a decrypted in-tunnel request verbatim to +// the real upstream and streams the response back over the TLS connection — +// no PII masking, no logging to the request log. Used for requests to +// intercept hosts whose path is not in the intercept allowlist (e.g. the Codex +// MCP transport and telemetry endpoints on chatgpt.com). +// +// Returns whether the connection can be reused for the next request in the +// MITM loop (false when the response framing cannot keep client and loop in +// sync, e.g. close-delimited responses or protocol upgrades). +func (tp *TransparentProxy) passthroughHTTPOverTLS(conn net.Conn, r *http.Request, targetHost string) bool { + log.Printf("[TransparentProxy] Passing through %s %s%s (path not intercepted)", r.Method, targetHost, r.URL.Path) + + // Disarm the MITM loop's 30s read deadline: the upstream response may be a + // long-lived stream. The loop re-arms the deadline on the next iteration. + if err := conn.SetReadDeadline(time.Time{}); err != nil { + log.Printf("[TransparentProxy] ❌ Failed to clear read deadline: %v", err) + } + + targetURL := tp.buildTargetURL(r, targetHost, "https") + return tp.forwardVerbatim(conn, r, targetURL, passthroughClient) +} + +// forwardVerbatim re-issues r against targetURL with client and relays the +// response to conn, preserving framing so the caller's next http.ReadRequest +// on the connection stays in sync. The request body is streamed, not +// buffered, and the response body is copied per-Read so streaming upstreams +// flush incrementally. Returns whether conn can be reused for another request. +func (tp *TransparentProxy) forwardVerbatim(conn net.Conn, r *http.Request, targetURL string, client *http.Client) bool { + proxyReq, err := http.NewRequestWithContext(r.Context(), r.Method, targetURL, r.Body) + if err != nil { + log.Printf("[TransparentProxy] ❌ Failed to create passthrough request: %v", err) + drainAndClose(r.Body) + tp.writeErrorResponse(conn, http.StatusBadGateway, "Failed to create passthrough request") + return false + } + proxyReq.ContentLength = r.ContentLength + + // Copy headers verbatim (including the client's own Authorization — Codex + // sends an OAuth bearer that must reach chatgpt.com untouched), minus + // hop-by-hop headers, which describe the client↔proxy connection. + proxyReq.Header = r.Header.Clone() + for _, h := range hopByHopHeaders { + proxyReq.Header.Del(h) + } + + // Re-add upgrade negotiation for protocol-switch requests (e.g. WebSocket) + // so the upstream can accept the upgrade; the 101 response is handled below. + if upgrade := r.Header.Get("Upgrade"); upgrade != "" { + proxyReq.Header.Set("Upgrade", upgrade) + proxyReq.Header.Set("Connection", "Upgrade") + } + + resp, err := client.Do(proxyReq) + if err != nil { + log.Printf("[TransparentProxy] ❌ Passthrough request failed: %v", err) + drainAndClose(r.Body) + tp.writeErrorResponse(conn, http.StatusBadGateway, "Failed to forward request") + return false + } + defer resp.Body.Close() + + // Protocol switch (e.g. WebSocket): relay the 101 response headers, then + // pipe bytes in both directions until either side closes. The connection + // cannot be reused for HTTP afterwards. + if resp.StatusCode == http.StatusSwitchingProtocols { + tp.relayProtocolSwitch(conn, resp) + return false + } + + // client.Do consumed the request body; drain any leftover bytes so they + // cannot desync the next http.ReadRequest on this connection. + drainAndClose(r.Body) + + // Pin the relayed proto to HTTP/1.1: the client side of the tunnel speaks + // HTTP/1.1 even when the upstream leg negotiated HTTP/2, and strict clients + // (e.g. Codex's hyper) drop the connection on an "HTTP/2.0" status line. + resp.Proto = protoHTTP11 + resp.ProtoMajor = 1 + resp.ProtoMinor = 1 + + // A close-delimited body (no Content-Length, no chunking) has no framing + // the client can detect on a reused connection; re-frame it as chunked. + if resp.ContentLength < 0 && len(resp.TransferEncoding) == 0 && bodyAllowedForStatus(r.Method, resp.StatusCode) { + resp.TransferEncoding = []string{"chunked"} + } + + // Write the response directly to the un-buffered connection: Response.Write + // emits status line + headers, then copies the body per-Read, so a + // long-lived stream (SSE, MCP) reaches the client chunk by chunk. + if err := resp.Write(conn); err != nil { + log.Printf("[TransparentProxy] ❌ Failed to write passthrough response: %v", err) + return false + } + + log.Printf("[TransparentProxy] Passed through %s %s - Status: %d", r.Method, r.URL.Path, resp.StatusCode) + + // Response.Write emits "Connection: close" when resp.Close is set, and a + // client that asked to close won't send another request. + return !resp.Close && !r.Close +} + +// relayProtocolSwitch writes a 101 response's status line and headers to conn, +// then copies bytes in both directions between the client connection and the +// upstream (resp.Body is an io.ReadWriteCloser for 101 responses). Afterwards +// the connection no longer speaks HTTP and must not be reused. +func (tp *TransparentProxy) relayProtocolSwitch(conn net.Conn, resp *http.Response) { + upstream, ok := resp.Body.(io.ReadWriteCloser) + if !ok { + log.Printf("[TransparentProxy] ❌ 101 response body is not writable; dropping connection") + return + } + + // Write the status line and headers manually: Response.Write would try to + // serialize the (bidirectional, unbounded) body as well. + if _, err := io.WriteString(conn, "HTTP/1.1 101 Switching Protocols\r\n"); err != nil { + log.Printf("[TransparentProxy] ❌ Failed to write 101 status line: %v", err) + return + } + if err := resp.Header.Write(conn); err != nil { + log.Printf("[TransparentProxy] ❌ Failed to write 101 headers: %v", err) + return + } + if _, err := io.WriteString(conn, "\r\n"); err != nil { + log.Printf("[TransparentProxy] ❌ Failed to finish 101 headers: %v", err) + return + } + + done := make(chan struct{}, 2) + go func() { + _, _ = io.Copy(upstream, conn) + upstream.Close() + done <- struct{}{} + }() + go func() { + _, _ = io.Copy(conn, upstream) + done <- struct{}{} + }() + <-done +} + +// bodyAllowedForStatus reports whether a response to the given request method +// and status code carries a body (mirrors net/http rules: no body for HEAD, +// 1xx, 204, 304). +func bodyAllowedForStatus(method string, status int) bool { + if method == http.MethodHead { + return false + } + switch { + case status >= 100 && status <= 199: + return false + case status == http.StatusNoContent, status == http.StatusNotModified: + return false + } + return true +} diff --git a/src/backend/proxy/passthrough_test.go b/src/backend/proxy/passthrough_test.go new file mode 100644 index 00000000..dc985ccc --- /dev/null +++ b/src/backend/proxy/passthrough_test.go @@ -0,0 +1,245 @@ +package proxy + +import ( + "bufio" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// runForwardVerbatim runs forwardVerbatim against the server end of a pipe and +// returns the client end plus a channel carrying the reuse result. The pipe is +// synchronous, so forwardVerbatim must run concurrently with the reader. +func runForwardVerbatim(t *testing.T, req *http.Request, targetURL string, client *http.Client) (net.Conn, <-chan bool) { + t.Helper() + clientEnd, serverEnd := net.Pipe() + t.Cleanup(func() { + clientEnd.Close() + serverEnd.Close() + }) + + tp := &TransparentProxy{} + reuse := make(chan bool, 1) + go func() { + reuse <- tp.forwardVerbatim(serverEnd, req, targetURL, client) + }() + return clientEnd, reuse +} + +// forwardVerbatim must relay method, path, headers (including Authorization) +// and body to the upstream unchanged, and relay the response back. +func TestForwardVerbatim_RoundTrip(t *testing.T) { + var gotMethod, gotPath, gotAuth, gotBody string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotMethod = r.Method + gotPath = r.URL.Path + gotAuth = r.Header.Get("Authorization") + body, _ := io.ReadAll(r.Body) + gotBody = string(body) + w.Header().Set("X-Upstream", "yes") + w.WriteHeader(http.StatusCreated) + fmt.Fprint(w, `{"ok":true}`) + })) + defer upstream.Close() + + req := httptest.NewRequest(http.MethodPost, "https://chatgpt.com/backend-api/ps/mcp", strings.NewReader(`{"jsonrpc":"2.0"}`)) + req.Header.Set("Authorization", "Bearer codex-oauth-token") + req.Header.Set("Content-Type", "application/json") + + clientEnd, reuse := runForwardVerbatim(t, req, upstream.URL+"/backend-api/ps/mcp", upstream.Client()) + + resp, err := http.ReadResponse(bufio.NewReader(clientEnd), req) + if err != nil { + t.Fatalf("failed to read relayed response: %v", err) + } + defer resp.Body.Close() + + if gotMethod != http.MethodPost || gotPath != "/backend-api/ps/mcp" { + t.Errorf("upstream saw %s %s, want POST /backend-api/ps/mcp", gotMethod, gotPath) + } + if gotAuth != "Bearer codex-oauth-token" { + t.Errorf("upstream saw Authorization %q, want the client's own OAuth bearer untouched", gotAuth) + } + if gotBody != `{"jsonrpc":"2.0"}` { + t.Errorf("upstream saw body %q, want it verbatim", gotBody) + } + + if resp.StatusCode != http.StatusCreated { + t.Errorf("relayed status = %d, want %d", resp.StatusCode, http.StatusCreated) + } + if resp.Header.Get("X-Upstream") != "yes" { + t.Error("upstream response header not relayed") + } + body, _ := io.ReadAll(resp.Body) + if string(body) != `{"ok":true}` { + t.Errorf("relayed body = %q, want %q", body, `{"ok":true}`) + } + + if !<-reuse { + t.Error("forwardVerbatim returned false, want connection reusable after framed response") + } +} + +// A streaming upstream must reach the client incrementally: the first chunk +// must be readable while the upstream is still blocked before its second +// chunk (proves the response is not buffered to completion). +func TestForwardVerbatim_StreamsIncrementally(t *testing.T) { + release := make(chan struct{}) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + flusher := w.(http.Flusher) + fmt.Fprint(w, "chunk-one") + flusher.Flush() + <-release // block until the test has read chunk-one + fmt.Fprint(w, "chunk-two") + })) + defer upstream.Close() + + req := httptest.NewRequest(http.MethodGet, "https://chatgpt.com/backend-api/ps/mcp", nil) + clientEnd, reuse := runForwardVerbatim(t, req, upstream.URL+"/backend-api/ps/mcp", upstream.Client()) + + resp, err := http.ReadResponse(bufio.NewReader(clientEnd), req) + if err != nil { + t.Fatalf("failed to read relayed response: %v", err) + } + defer resp.Body.Close() + + buf := make([]byte, len("chunk-one")) + if err := clientEnd.SetReadDeadline(time.Now().Add(5 * time.Second)); err != nil { + t.Fatalf("failed to set read deadline: %v", err) + } + if _, err := io.ReadFull(resp.Body, buf); err != nil { + t.Fatalf("failed to read first chunk while upstream still streaming: %v", err) + } + if string(buf) != "chunk-one" { + t.Errorf("first chunk = %q, want %q", buf, "chunk-one") + } + + close(release) // let the upstream finish + rest, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("failed to read remainder: %v", err) + } + if string(rest) != "chunk-two" { + t.Errorf("remainder = %q, want %q", rest, "chunk-two") + } + <-reuse +} + +// Two sequential passthrough requests over the same connection must stay in +// sync (framing preserved so the next read starts at a request boundary). +func TestForwardVerbatim_KeepAliveSequential(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fmt.Fprintf(w, "echo:%s", r.URL.Path) + })) + defer upstream.Close() + + clientEnd, serverEnd := net.Pipe() + defer clientEnd.Close() + defer serverEnd.Close() + + tp := &TransparentProxy{} + done := make(chan bool, 2) + go func() { + for _, path := range []string{"/first", "/second"} { + req := httptest.NewRequest(http.MethodGet, "https://chatgpt.com"+path, nil) + done <- tp.forwardVerbatim(serverEnd, req, upstream.URL+path, upstream.Client()) + } + }() + + reader := bufio.NewReader(clientEnd) + for _, want := range []string{"echo:/first", "echo:/second"} { + resp, err := http.ReadResponse(reader, nil) + if err != nil { + t.Fatalf("failed to read relayed response: %v", err) + } + body, _ := io.ReadAll(resp.Body) + resp.Body.Close() + if string(body) != want { + t.Errorf("relayed body = %q, want %q", body, want) + } + if !<-done { + t.Errorf("forwardVerbatim for %q returned false, want reusable", want) + } + } +} + +// roundTripperFunc adapts a function to http.RoundTripper. +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +// When the upstream leg negotiated HTTP/2, the status line relayed to the +// (HTTP/1.1) client must still say HTTP/1.1 — strict clients like Codex's +// hyper drop the connection on an "HTTP/2.0" status line. +func TestForwardVerbatim_NormalizesHTTP2ProtoToHTTP11(t *testing.T) { + client := &http.Client{Transport: roundTripperFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Status: "200 OK", + Proto: "HTTP/2.0", + ProtoMajor: 2, + ProtoMinor: 0, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"models":[]}`)), + ContentLength: int64(len(`{"models":[]}`)), + }, nil + })} + + req := httptest.NewRequest(http.MethodGet, "https://chatgpt.com/backend-api/codex/models", nil) + clientEnd, reuse := runForwardVerbatim(t, req, "https://chatgpt.com/backend-api/codex/models", client) + + reader := bufio.NewReader(clientEnd) + statusLine, err := reader.ReadString('\n') + if err != nil { + t.Fatalf("failed to read status line: %v", err) + } + if !strings.HasPrefix(statusLine, "HTTP/1.1 200") { + t.Errorf("status line = %q, want it to start with %q", statusLine, "HTTP/1.1 200") + } + + resp, err := http.ReadResponse(bufio.NewReader(io.MultiReader(strings.NewReader(statusLine), reader)), req) + if err != nil { + t.Fatalf("relayed response unparseable as HTTP/1.1: %v", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + if string(body) != `{"models":[]}` { + t.Errorf("relayed body = %q, want %q", body, `{"models":[]}`) + } + if !<-reuse { + t.Error("forwardVerbatim returned false, want connection reusable") + } +} + +// An unreachable upstream must produce a 502 on the client connection instead +// of silently dropping it. +func TestForwardVerbatim_UpstreamError(t *testing.T) { + // Grab a port with nothing listening on it. + l, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("failed to reserve port: %v", err) + } + deadURL := "http://" + l.Addr().String() + l.Close() + + req := httptest.NewRequest(http.MethodGet, "https://chatgpt.com/backend-api/ps/mcp", nil) + clientEnd, reuse := runForwardVerbatim(t, req, deadURL+"/backend-api/ps/mcp", &http.Client{Timeout: 5 * time.Second}) + + resp, err := http.ReadResponse(bufio.NewReader(clientEnd), req) + if err != nil { + t.Fatalf("failed to read error response: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusBadGateway { + t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusBadGateway) + } + if <-reuse { + t.Error("forwardVerbatim returned true after upstream failure, want false") + } +} diff --git a/src/backend/proxy/router.go b/src/backend/proxy/router.go index b1110b30..415265a7 100644 --- a/src/backend/proxy/router.go +++ b/src/backend/proxy/router.go @@ -9,12 +9,18 @@ import ( // Router handles request routing based on target domains type Router struct { interceptDomains []string + // pathPrefixes optionally restricts interception per host: a host with an + // entry here only intercepts requests whose path matches one of the + // prefixes; hosts without an entry intercept all paths. + pathPrefixes map[string][]string } -// NewRouter creates a new router with the given intercept domains -func NewRouter(interceptDomains []string) *Router { +// NewRouter creates a new router with the given intercept domains and +// per-host path-prefix allowlists (see Router.pathPrefixes). +func NewRouter(interceptDomains []string, pathPrefixes map[string][]string) *Router { return &Router{ interceptDomains: interceptDomains, + pathPrefixes: pathPrefixes, } } @@ -41,6 +47,36 @@ func (r *Router) ShouldIntercept(host string) bool { return false } +// ShouldInterceptRequest checks if a request to the given host and path should +// be intercepted (masked). It refines ShouldIntercept: for hosts with a +// path-prefix allowlist only matching paths are intercepted; other requests to +// such hosts should be passed through verbatim. +func (r *Router) ShouldInterceptRequest(host string, path string) bool { + if !r.ShouldIntercept(host) { + return false + } + + hostname, _, err := net.SplitHostPort(host) + if err != nil { + hostname = host + } + hostname = strings.ToLower(hostname) + + prefixes, ok := r.pathPrefixes[hostname] + if !ok { + return true // no allowlist: intercept all paths on this host + } + + for _, prefix := range prefixes { + if strings.HasPrefix(path, prefix) { + return true + } + } + + log.Printf("[Router] Passing through request to %s%s (path not in intercept allowlist)", hostname, path) + return false +} + // IsTargetDomain is an alias for ShouldIntercept for consistency func (r *Router) IsTargetDomain(host string) bool { return r.ShouldIntercept(host) diff --git a/src/backend/proxy/router_test.go b/src/backend/proxy/router_test.go new file mode 100644 index 00000000..7418c321 --- /dev/null +++ b/src/backend/proxy/router_test.go @@ -0,0 +1,81 @@ +package proxy + +import ( + "testing" + + "github.com/dataiku/kiji-proxy/src/backend/providers" +) + +func newTestRouter() *Router { + return NewRouter( + []string{"api.openai.com", "api.anthropic.com", providers.ProviderAPIDomainCodex}, + map[string][]string{ + providers.ProviderAPIDomainCodex: {providers.ProviderSubpathCodexResponses}, + }, + ) +} + +func TestRouter_ShouldIntercept(t *testing.T) { + router := newTestRouter() + + tests := []struct { + name string + host string + want bool + }{ + {"intercept domain", "api.openai.com", true}, + {"intercept domain with port", "api.anthropic.com:443", true}, + {"subdomain of intercept domain", "eu.api.openai.com", true}, + {"codex host", "chatgpt.com", true}, + {"non-intercepted host", "example.com", false}, + {"suffix but not subdomain", "notchatgpt.com", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := router.ShouldIntercept(tt.host); got != tt.want { + t.Errorf("ShouldIntercept(%q) = %v, want %v", tt.host, got, tt.want) + } + }) + } +} + +func TestRouter_ShouldInterceptRequest(t *testing.T) { + router := newTestRouter() + + tests := []struct { + name string + host string + path string + want bool + }{ + {"codex completions path", "chatgpt.com", "/backend-api/codex/responses", true}, + {"codex completions subpath", "chatgpt.com", "/backend-api/codex/responses/123", true}, + {"codex MCP transport passes through", "chatgpt.com", "/backend-api/ps/mcp", false}, + {"codex telemetry passes through", "chatgpt.com", "/backend-api/codex/events", false}, + {"codex root passes through", "chatgpt.com", "/", false}, + {"codex host with port normalizes", "chatgpt.com:443", "/backend-api/codex/responses", true}, + {"codex host mixed case normalizes", "ChatGPT.com", "/backend-api/ps/mcp", false}, + {"host without allowlist intercepts all paths", "api.openai.com", "/v1/chat/completions", true}, + {"host without allowlist intercepts any path", "api.openai.com", "/some/other/path", true}, + {"anthropic host intercepts all paths", "api.anthropic.com", "/v1/messages", true}, + {"non-intercepted host is never intercepted", "example.com", "/backend-api/codex/responses", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := router.ShouldInterceptRequest(tt.host, tt.path); got != tt.want { + t.Errorf("ShouldInterceptRequest(%q, %q) = %v, want %v", tt.host, tt.path, got, tt.want) + } + }) + } +} + +// A router constructed with a nil allowlist (legacy behavior) must intercept +// every path on intercept hosts. +func TestRouter_ShouldInterceptRequest_NilPrefixes(t *testing.T) { + router := NewRouter([]string{"chatgpt.com"}, nil) + if !router.ShouldInterceptRequest("chatgpt.com", "/backend-api/ps/mcp") { + t.Error("ShouldInterceptRequest with nil allowlist should intercept all paths") + } +} diff --git a/src/backend/proxy/streaming.go b/src/backend/proxy/streaming.go index 88d3c6d8..afc10fd7 100644 --- a/src/backend/proxy/streaming.go +++ b/src/backend/proxy/streaming.go @@ -71,6 +71,43 @@ func isEventStream(resp *http.Response) bool { strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") } +// responseLooksLikeSSE reports whether the response is an SSE stream, sniffing +// the body when the Content-Type header is absent: the ChatGPT-login Codex +// backend (chatgpt.com/backend-api/codex/responses) streams SSE with no +// Content-Type at all. Sniffing may wrap resp.Body to preserve the peeked +// bytes, so callers must keep using resp.Body afterwards (not a saved copy). +func responseLooksLikeSSE(resp *http.Response) bool { + if isEventStream(resp) { + return true + } + if resp.Header.Get("Content-Type") != "" { + return false // upstream declared something else; believe it + } + + br := bufio.NewReader(resp.Body) + body := resp.Body + resp.Body = &sniffedBody{Reader: br, Closer: body} + + // Peek the first field name of the body. An SSE stream starts with an + // "event:"/"data:"/"id:"/"retry:" field or a ":" comment. Peek blocks only + // until the first bytes arrive, which for a stream is the first event. + peek, _ := br.Peek(len("event:")) + s := string(peek) + for _, prefix := range []string{"event:", "data:", "id:", "retry:", ":"} { + if strings.HasPrefix(s, prefix) { + return true + } + } + return false +} + +// sniffedBody re-joins a buffered reader (holding peeked bytes) with the +// original body's Closer. +type sniffedBody struct { + io.Reader + io.Closer +} + // streamCodec restores masked PII inside one provider's SSE token stream. The // engine below (streamSSEResponse) owns the transport-level framing and event // splitting; a codec only knows how to rewrite a single, complete SSE event for diff --git a/src/backend/proxy/streaming_test.go b/src/backend/proxy/streaming_test.go index 2360ca39..6750a776 100644 --- a/src/backend/proxy/streaming_test.go +++ b/src/backend/proxy/streaming_test.go @@ -354,3 +354,81 @@ func TestStreamSSEResponse_EndToEnd(t *testing.T) { t.Errorf("terminal event missing:\n%s", text) } } + +// The Codex backend streams SSE with no Content-Type header; sniffing the body +// start must classify it as SSE, while JSON bodies without a header stay +// non-SSE, and a declared non-SSE Content-Type is believed without sniffing. +func TestResponseLooksLikeSSE(t *testing.T) { + mk := func(contentType, body string) *http.Response { + h := http.Header{} + if contentType != "" { + h.Set("Content-Type", contentType) + } + return &http.Response{Header: h, Body: io.NopCloser(strings.NewReader(body))} + } + + tests := []struct { + name string + resp *http.Response + want bool + wantBodyRaw string // body readable after the sniff, byte-for-byte + }{ + {"declared SSE", mk("text/event-stream; charset=utf-8", "event: x\n\n"), true, "event: x\n\n"}, + {"declared JSON not sniffed", mk("application/json", "data: looks like SSE"), false, "data: looks like SSE"}, + {"no content-type, SSE event body", mk("", "event: response.created\ndata: {}\n\n"), true, "event: response.created\ndata: {}\n\n"}, + {"no content-type, SSE data body", mk("", "data: {}\n\n"), true, "data: {}\n\n"}, + {"no content-type, JSON body", mk("", `{"ok":true}`), false, `{"ok":true}`}, + {"no content-type, empty body", mk("", ""), false, ""}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := responseLooksLikeSSE(tt.resp); got != tt.want { + t.Errorf("responseLooksLikeSSE() = %v, want %v", got, tt.want) + } + body, err := io.ReadAll(tt.resp.Body) + if err != nil { + t.Fatalf("body unreadable after sniff: %v", err) + } + if string(body) != tt.wantBodyRaw { + t.Errorf("body after sniff = %q, want %q (peeked bytes lost?)", body, tt.wantBodyRaw) + } + }) + } +} + +// Codex (and other Responses-API clients) render their final message from the +// nested payload copies — content_part.done's part.text, output_item.done's +// item.content[].text, and response.completed's response.output[].content[]. +// text — so those must be restored too, not only the flat .delta/.done fields. +func TestOpenAICodec_RestoresNestedDonePayloads(t *testing.T) { + events := []string{ + "event: response.content_part.done\n" + + `data: {"type":"response.content_part.done","output_index":0,"content_index":0,"part":{"type":"output_text","text":"Hi Miguel,"}}` + "\n\n", + "event: response.output_item.done\n" + + `data: {"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg_MiguelX","content":[{"type":"output_text","text":"Hi Miguel,"}]}}` + "\n\n", + "event: response.completed\n" + + `data: {"type":"response.completed","response":{"id":"resp_1","output":[{"type":"message","content":[{"type":"output_text","text":"Hi Miguel,"}]}]}}` + "\n\n", + } + + for _, raw := range events { + codec := newOpenAICodec(map[string]string{"Miguel": "David"}) + out := string(codec.transformEvent(sseLines(raw))) + if !strings.Contains(out, "Hi David,") { + t.Errorf("nested text not restored in %q: %s", raw[:40], out) + } + if strings.Contains(out, "Hi Miguel,") { + t.Errorf("dummy leaked in %q: %s", raw[:40], out) + } + } + + // Opaque strings (ids) that contain a dummy as a substring must NOT be + // rewritten. + codec := newOpenAICodec(map[string]string{"Miguel": "David"}) + out := string(codec.transformEvent(sseLines( + "event: response.output_item.done\n" + + `data: {"type":"response.output_item.done","output_index":0,"item":{"id":"msg_MiguelX","content":[]}}` + "\n\n"))) + if !strings.Contains(out, "msg_MiguelX") { + t.Errorf("opaque id was rewritten: %s", out) + } +} diff --git a/src/backend/proxy/transparent.go b/src/backend/proxy/transparent.go index 87895e9e..3b5bd31e 100644 --- a/src/backend/proxy/transparent.go +++ b/src/backend/proxy/transparent.go @@ -119,8 +119,9 @@ func (tp *TransparentProxy) handleHTTPRequest(w http.ResponseWriter, r *http.Req targetHost = r.URL.Host } - // Check if we should intercept - if !tp.router.ShouldIntercept(targetHost) { + // Check if we should intercept (host must be an intercept domain and, for + // hosts with a path allowlist, the path must match) + if !tp.router.ShouldInterceptRequest(targetHost, r.URL.Path) { // Passthrough - forward directly tp.passthroughHTTP(w, r) return @@ -377,7 +378,11 @@ func (tp *TransparentProxy) interceptCONNECT(w http.ResponseWriter, _ *http.Requ Conn: tlsConn, } - // Handle HTTP requests over the TLS connection + // Handle HTTP requests over the TLS connection. The buffered reader is + // created once for the connection: a per-iteration reader would silently + // drop any bytes of the next request it had already buffered while reading + // the previous one, desyncing the keep-alive connection. + connReader := bufio.NewReader(conn) for { // Set read deadline if err := conn.SetReadDeadline(time.Now().Add(30 * time.Second)); err != nil { @@ -386,7 +391,7 @@ func (tp *TransparentProxy) interceptCONNECT(w http.ResponseWriter, _ *http.Requ } // Read HTTP request - req, err := http.ReadRequest(bufio.NewReader(conn)) + req, err := http.ReadRequest(connReader) if err != nil { if err == io.EOF { break @@ -399,8 +404,16 @@ func (tp *TransparentProxy) interceptCONNECT(w http.ResponseWriter, _ *http.Requ req.URL.Scheme = "https" req.URL.Host = host - // Process the request - tp.interceptHTTPOverTLS(conn, req, host, provider) + // Process the request. Hosts with a path allowlist (chatgpt.com) only + // mask matching paths; everything else on the (already MITM'd) + // connection is forwarded verbatim so streaming side channels like the + // Codex MCP transport survive. + if tp.router.ShouldInterceptRequest(host, req.URL.Path) { + tp.interceptHTTPOverTLS(conn, req, host, provider) + } else if !tp.passthroughHTTPOverTLS(conn, req, host) { + // Response framing can't keep the connection in sync; stop reusing it. + break + } } } @@ -416,6 +429,17 @@ func (tp *TransparentProxy) interceptHTTPOverTLS(conn net.Conn, r *http.Request, return } + // Reject protocol upgrades (e.g. WebSocket) on intercepted paths: PII inside + // a switched protocol cannot be masked, so fail closed — and fast, so the + // client's fallback/retry logic isn't stalled by a 30s hang. (Codex can be + // steered to SSE via a provider config with supports_websockets = false.) + if upgrade := r.Header.Get("Upgrade"); upgrade != "" { + log.Printf("[TransparentProxy] ❌ Rejecting %q upgrade on intercepted path %s (cannot mask non-HTTP protocols)", upgrade, r.URL.Path) + drainAndClose(r.Body) + tp.writeErrorResponse(conn, http.StatusNotImplemented, "Kiji proxy: protocol upgrades are not supported on masked endpoints") + return + } + // Read request body body, err := io.ReadAll(r.Body) if err != nil { @@ -474,11 +498,15 @@ func (tp *TransparentProxy) interceptHTTPOverTLS(conn net.Conn, r *http.Request, } defer resp.Body.Close() + log.Printf("[TransparentProxy] Upstream response for %s: status=%d content-type=%q content-encoding=%q", + r.URL.Path, resp.StatusCode, resp.Header.Get("Content-Type"), resp.Header.Get("Content-Encoding")) + // Stream Server-Sent Events straight through to the client, restoring PII // incrementally. Buffering an SSE stream (as the non-streaming path below // does) breaks streaming clients like Claude Code and hangs until the - // upstream timeout fires. - if wantStream && isEventStream(resp) { + // upstream timeout fires. responseLooksLikeSSE also body-sniffs, because + // the Codex backend streams SSE without a Content-Type header. + if wantStream && responseLooksLikeSSE(resp) { log.Printf("[TransparentProxy] Streaming SSE response for %s", r.URL.Path) codec := codecForProvider(provider, processed.MaskedToOriginal) if streamErr := streamSSEResponse(conn, resp, codec); streamErr != nil { @@ -509,13 +537,17 @@ func (tp *TransparentProxy) interceptHTTPOverTLS(conn net.Conn, r *http.Request, modifiedBody := tp.handler.ProcessResponseBody(ctx, respBody, resp.Header.Get("Content-Type"), processed.MaskedToOriginal, processed.TransactionID, provider) responseRestoreTime := time.Since(restoreStart) - // Create new response with modified body + // Create new response with modified body. The proto is pinned to HTTP/1.1 + // regardless of what the upstream leg negotiated: the client side of the + // tunnel speaks HTTP/1.1, and relaying an upstream "HTTP/2.0" proto into + // the status line makes strict clients (e.g. Codex's hyper) drop the + // connection as malformed. newResp := &http.Response{ StatusCode: resp.StatusCode, Status: resp.Status, - Proto: resp.Proto, - ProtoMajor: resp.ProtoMajor, - ProtoMinor: resp.ProtoMinor, + Proto: protoHTTP11, + ProtoMajor: 1, + ProtoMinor: 1, Header: resp.Header, Body: io.NopCloser(bytes.NewReader(modifiedBody)), ContentLength: int64(len(modifiedBody)), @@ -550,7 +582,7 @@ func (tp *TransparentProxy) writeErrorResponse(conn net.Conn, statusCode int, me resp := &http.Response{ StatusCode: statusCode, Status: http.StatusText(statusCode), - Proto: "HTTP/1.1", + Proto: protoHTTP11, ProtoMajor: 1, ProtoMinor: 1, Header: make(http.Header), diff --git a/src/backend/proxy/transparent_factory.go b/src/backend/proxy/transparent_factory.go index cb95342c..0524e958 100644 --- a/src/backend/proxy/transparent_factory.go +++ b/src/backend/proxy/transparent_factory.go @@ -20,7 +20,7 @@ func NewTransparentProxyFromConfig(cfg *config.Config, handler *Handler) (*Trans } // Create router - router := NewRouter(cfg.Providers.GetInterceptDomains()) + router := NewRouter(cfg.Providers.GetInterceptDomains(), cfg.Providers.GetInterceptPathPrefixes()) // Create transparent proxy, reusing the existing Handler transparentProxy := NewTransparentProxy( From be948f86dcb5c216f65c32ff8b39bd30670585d4 Mon Sep 17 00:00:00 2001 From: Davidnet Date: Fri, 3 Jul 2026 15:18:08 -0400 Subject: [PATCH 09/13] fix: restore PII in a single pass to prevent chained substitution The response-restoration logic replaced each masked (dummy) value with its original via a sequence of strings.ReplaceAll calls. When the fake-value generator produced a dummy that coincides with a real original from another mapping (e.g. "Priya"->"Nicole" alongside "Claude"->"Priya"), the passes chained: restoring "Nicole"->"Priya" and then re-replacing that "Priya"-> "Claude", so the client received the wrong PII ("Hi Claude" where the model wrote "Hi Nicole" for Priya). The bug was also map-order dependent. Replace both restore paths (streaming codecs and the buffered ResponseProcessor) with a shared single-pass strings.Replacer built by processor.BuildRestorer, keys ordered longest-first. strings.Replacer scans the input once and never re-substitutes its own output, so restored text is never re-replaced. The streaming restoreCore builds the replacer once per stream instead of iterating the map per delta. Adds regression tests for the exact collision and longest-match precedence. --- src/backend/processor/response.go | 35 ++++++++++++++++--- src/backend/processor/response_test.go | 48 ++++++++++++++++++++++++++ src/backend/proxy/streaming.go | 23 ++++++------ src/backend/proxy/streaming_test.go | 17 +++++++++ 4 files changed, 106 insertions(+), 17 deletions(-) create mode 100644 src/backend/processor/response_test.go diff --git a/src/backend/processor/response.go b/src/backend/processor/response.go index 2e5ae445..4e0356d5 100644 --- a/src/backend/processor/response.go +++ b/src/backend/processor/response.go @@ -3,6 +3,7 @@ package processor import ( "encoding/json" "log" + "sort" "strings" "time" @@ -72,11 +73,35 @@ func (rp *ResponseProcessor) ProcessResponse(body []byte, contentType string, ma return modifiedBody } -// RestorePII restores masked PII text back to original text using the provided mapping +// RestorePII restores masked PII text back to original text using the provided mapping. func (rp *ResponseProcessor) RestorePII(text string, maskedToOriginal map[string]string) string { - // Replace all occurrences of masked text with original text - for maskedText, originalText := range maskedToOriginal { - text = strings.ReplaceAll(text, maskedText, originalText) + return BuildRestorer(maskedToOriginal).Replace(text) +} + +// BuildRestorer builds a single-pass replacer that maps each masked (dummy) +// value back to its original. +// +// It must be a single simultaneous pass, not a sequence of ReplaceAll calls: +// a generated dummy can coincide with a real original from another mapping +// (e.g. "Priya"→"Nicole" alongside "Claude"→"Priya"). Sequential replacement +// chains — restoring "Nicole"→"Priya" and then re-replacing that "Priya"→ +// "Claude" — corrupting the output. strings.Replacer scans the input once and +// never re-examines what it has emitted, so restored text is never +// re-substituted. Keys are ordered longest-first so that when one dummy is a +// prefix of another the longer match wins. +func BuildRestorer(maskedToOriginal map[string]string) *strings.Replacer { + keys := make([]string, 0, len(maskedToOriginal)) + for masked := range maskedToOriginal { + if masked == "" { + continue // an empty key would match everywhere; skip defensively + } + keys = append(keys, masked) + } + sort.Slice(keys, func(i, j int) bool { return len(keys[i]) > len(keys[j]) }) + + pairs := make([]string, 0, len(keys)*2) + for _, masked := range keys { + pairs = append(pairs, masked, maskedToOriginal[masked]) } - return text + return strings.NewReplacer(pairs...) } diff --git a/src/backend/processor/response_test.go b/src/backend/processor/response_test.go new file mode 100644 index 00000000..4d749994 --- /dev/null +++ b/src/backend/processor/response_test.go @@ -0,0 +1,48 @@ +package processor + +import "testing" + +// A generated dummy can coincide with a real original from another mapping +// ("Priya"→"Nicole" alongside "Claude"→"Priya"). Restoration must be a single +// simultaneous pass: restoring the model's "Nicole" to "Priya" must not then +// chain through the "Priya"→"Claude" mapping. Regression for the sequential +// ReplaceAll bug that corrupted restored PII. +func TestRestorePII_NoChainedSubstitution(t *testing.T) { + rp := &ResponseProcessor{} + got := rp.RestorePII("Hi Nicole, regards Priya.", map[string]string{ + "Nicole": "Priya", + "Priya": "Claude", + }) + want := "Hi Priya, regards Claude." + if got != want { + t.Errorf("RestorePII = %q, want %q", got, want) + } +} + +// When one dummy is a prefix of another, the longest match must win so a +// shorter dummy doesn't partially consume a longer one. +func TestRestorePII_LongestMatchWins(t *testing.T) { + rp := &ResponseProcessor{} + got := rp.RestorePII("value abc here", map[string]string{ + "ab": "SHORT", + "abc": "LONG", + }) + want := "value LONG here" + if got != want { + t.Errorf("RestorePII = %q, want %q", got, want) + } +} + +func TestRestorePII_EmptyAndPlainCases(t *testing.T) { + rp := &ResponseProcessor{} + if got := rp.RestorePII("nothing to do", nil); got != "nothing to do" { + t.Errorf("nil mapping = %q, want unchanged", got) + } + got := rp.RestorePII("email dummy@x.test twice: dummy@x.test", map[string]string{ + "dummy@x.test": "real@example.com", + }) + want := "email real@example.com twice: real@example.com" + if got != want { + t.Errorf("RestorePII = %q, want %q", got, want) + } +} diff --git a/src/backend/proxy/streaming.go b/src/backend/proxy/streaming.go index afc10fd7..7b867efd 100644 --- a/src/backend/proxy/streaming.go +++ b/src/backend/proxy/streaming.go @@ -14,6 +14,7 @@ import ( "strings" "time" + "github.com/dataiku/kiji-proxy/src/backend/processor" "github.com/dataiku/kiji-proxy/src/backend/providers" ) @@ -124,13 +125,13 @@ type streamCodec interface { } // restoreCore holds the provider-agnostic restore state shared by every codec: -// the dummy→original mapping, the hold-back length, and the pre-restore model +// the dummy→original replacer, the hold-back length, and the pre-restore model // output accumulated for audit logging. Codecs embed it so they only implement // their own event grammar. type restoreCore struct { - mapping map[string]string - keep int // bytes to hold back = longest dummy length - 1 - masked strings.Builder // raw model output (pre-restore), accumulated for logging + replacer *strings.Replacer + keep int // bytes to hold back = longest dummy length - 1 + masked strings.Builder // raw model output (pre-restore), accumulated for logging } func newRestoreCore(mapping map[string]string) restoreCore { @@ -143,17 +144,15 @@ func newRestoreCore(mapping map[string]string) restoreCore { if keep > 0 { keep-- } - return restoreCore{mapping: mapping, keep: keep} + return restoreCore{replacer: processor.BuildRestorer(mapping), keep: keep} } -// restore replaces every masked (dummy) value with its original. Replacing -// already-restored text is idempotent provided an original value does not -// contain a dummy placeholder as a substring. +// restore replaces every masked (dummy) value with its original in a single +// pass. See processor.BuildRestorer for why this must not be a sequence of +// ReplaceAll calls (chained substitution corrupts restoration when a dummy +// coincides with another mapping's original). func (c *restoreCore) restore(text string) string { - for masked, original := range c.mapping { - text = strings.ReplaceAll(text, masked, original) - } - return text + return c.replacer.Replace(text) } func (c *restoreCore) maskedOutput() string { diff --git a/src/backend/proxy/streaming_test.go b/src/backend/proxy/streaming_test.go index 6750a776..0d8d39ab 100644 --- a/src/backend/proxy/streaming_test.go +++ b/src/backend/proxy/streaming_test.go @@ -60,6 +60,23 @@ func TestRestoreCore(t *testing.T) { } } +// A generated dummy can coincide with a real original from another mapping +// ("Priya"→"Nicole" alongside "Claude"→"Priya"). Restoration must be a single +// pass: the model's "Nicole" restores to "Priya" and must NOT then chain +// through the "Priya"→"Claude" mapping. Regression for the sequential +// ReplaceAll bug that corrupted restored PII. +func TestRestoreCore_NoChainedSubstitution(t *testing.T) { + core := newRestoreCore(map[string]string{ + "Nicole": "Priya", // Priya was masked to the dummy Nicole + "Priya": "Claude", // Claude was masked to the dummy Priya + }) + got := core.restore("Hi Nicole, regards Priya.") + want := "Hi Priya, regards Claude." + if got != want { + t.Errorf("restore = %q, want %q", got, want) + } +} + // --- helpers --- // event turns raw SSE text into the line slices transformEvent receives. From 8a934cdb7f0019cb3986904957f4886517d04f39 Mon Sep 17 00:00:00 2001 From: Davidnet Date: Mon, 6 Jul 2026 13:59:43 -0400 Subject: [PATCH 10/13] docs: explain why the proxy env vars use http:// not https:// --- docs/09-coding-agents.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/09-coding-agents.md b/docs/09-coding-agents.md index 5da7c592..6049d3cf 100644 --- a/docs/09-coding-agents.md +++ b/docs/09-coding-agents.md @@ -33,6 +33,8 @@ For **any** agent, two things must be true: Throughout this chapter the proxy endpoint is `http://127.0.0.1:8081` (the transparent proxy port). Adjust if you changed `proxy_port`. +> **Why `http://` and not `https://`?** The scheme in `HTTP_PROXY` / `HTTPS_PROXY` describes how the agent reaches the *proxy*, not the traffic being proxied. `HTTPS_PROXY=http://127.0.0.1:8081` means "for HTTPS destinations, use the proxy reachable over plain HTTP at `127.0.0.1:8081`". The agent still reaches the proxy by opening an `HTTP CONNECT` tunnel (an unencrypted control message), and the proxy then TLS-terminates the tunnel with its own CA. This is not a security downgrade: the `http://` hop is **loopback only** — it never leaves your machine, so there is no wire to intercept, and traffic from the proxy onward to the provider is full TLS. Using `https://127.0.0.1:8081` would fail, because the listener speaks plain HTTP on that port (this is the standard convention for local proxies such as mitmproxy, Charles, and Burp). Note that TLS interception is deliberate here — the proxy decrypts requests in order to mask PII, so the security boundary that matters is your local machine and user account, not the proxy scheme. + ## Claude Code Claude Code is a Node.js application, so it uses the standard Node proxy and CA variables. From 0e9b02dc5e3adb6095a6a54d65c79416b19ea8d8 Mon Sep 17 00:00:00 2001 From: Davidnet Date: Mon, 6 Jul 2026 15:17:07 -0400 Subject: [PATCH 11/13] test: add failing regression tests for streaming PII-restore bugs Pin six defects surfaced in review of the streaming SSE path. Each test asserts the correct behavior and currently fails: - carry buffer re-restores already-restored text (chained substitution) in both the OpenAI and Anthropic codecs - OpenAI codec fails open for /v1/chat/completions (chat.completion.chunk) - splitSafe cuts multi-byte UTF-8 runes -> json.Marshal inserts U+FFFD - restored PII spliced unescaped into partial_json tool-arg fragments - codec carry tail dropped when a stream ends without .done/content_block_stop --- src/backend/proxy/streaming_test.go | 165 ++++++++++++++++++++++++++++ 1 file changed, 165 insertions(+) diff --git a/src/backend/proxy/streaming_test.go b/src/backend/proxy/streaming_test.go index 0d8d39ab..5e9913d1 100644 --- a/src/backend/proxy/streaming_test.go +++ b/src/backend/proxy/streaming_test.go @@ -449,3 +449,168 @@ func TestOpenAICodec_RestoresNestedDonePayloads(t *testing.T) { t.Errorf("opaque id was rewritten: %s", out) } } + +// --- regression tests for known streaming-restore bugs (currently FAILING) --- +// +// Each test below asserts the behavior the codec SHOULD have. They fail against +// the current implementation and pin the bugs surfaced in review. + +// BUG 1 (OpenAI): the per-channel carry buffer stores POST-restore text and +// re-restores it on the next delta, so a restored original that is itself a +// dummy key gets substituted a second time — the exact chained-substitution +// corruption processor.BuildRestorer was written to prevent. The buffered path +// guards this in TestRestoreCore_NoChainedSubstitution; the streaming path does +// not. +func TestOpenAICodec_NoChainedSubstitutionAcrossDeltas(t *testing.T) { + // "Priya" was masked to the dummy "Nicole"; "Claude" was masked to the + // dummy "Priya". Restoring the model's "Nicole" must yield "Priya" and stop. + codec := newOpenAICodec(map[string]string{ + "Nicole": "Priya", + "Priya": "Claude", + }) + ev1 := codec.transformEvent(sseLines( + "event: response.output_text.delta\n" + + `data: {"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"Hi Nicole"}` + "\n\n")) + ev2 := codec.transformEvent(sseLines( + "event: response.output_text.delta\n" + + `data: {"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"!"}` + "\n\n")) + done := codec.transformEvent(sseLines( + "event: response.output_text.done\n" + + `data: {"type":"response.output_text.done","output_index":0,"content_index":0,"text":"Hi Nicole!"}` + "\n\n")) + + assembled := deltaPayloads(t, string(ev1)+string(ev2)+string(done)) + if strings.Contains(assembled, "Claude") { + t.Errorf("restored original was re-substituted (chained): assembled=%q contains \"Claude\"", assembled) + } + if assembled != "Hi Priya!" { + t.Errorf("client-assembled deltas = %q, want %q", assembled, "Hi Priya!") + } +} + +// BUG 1 (Anthropic): same defect in the per-content-block carry buffer. +func TestAnthropicCodec_NoChainedSubstitutionAcrossDeltas(t *testing.T) { + codec := newAnthropicCodec(map[string]string{ + "Nicole": "Priya", + "Priya": "Claude", + }) + ev1 := codec.transformEvent(sseLines( + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hi Nicole"}}` + "\n\n")) + ev2 := codec.transformEvent(sseLines( + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"!"}}` + "\n\n")) + stop := codec.transformEvent(sseLines( + "event: content_block_stop\n" + + `data: {"type":"content_block_stop","index":0}` + "\n\n")) + + assembled := deltaPayloads(t, string(ev1)+string(ev2)+string(stop)) + if strings.Contains(assembled, "Claude") { + t.Errorf("restored original was re-substituted (chained): assembled=%q contains \"Claude\"", assembled) + } + if assembled != "Hi Priya!" { + t.Errorf("client-assembled deltas = %q, want %q", assembled, "Hi Priya!") + } +} + +// BUG 2: the OpenAI codec only recognizes the Responses-API grammar +// (type == *.delta / *.done). A /v1/chat/completions stream uses +// chat.completion.chunk with no top-level "type", so every chunk falls through +// to passthrough and the dummy value reaches the client unrestored (fail-open). +func TestOpenAICodec_ChatCompletionsChunkRestored(t *testing.T) { + codec := newOpenAICodec(map[string]string{"Nicole": "Priya"}) + out := string(codec.transformEvent(sseLines( + `data: {"id":"c1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"Hi Nicole"}}]}` + "\n\n"))) + + if strings.Contains(out, "Nicole") { + t.Errorf("chat.completion.chunk not restored — dummy leaked to client: %s", out) + } + if !strings.Contains(out, "Priya") { + t.Errorf("chat.completion.chunk should restore to \"Priya\": %s", out) + } +} + +// BUG 3: splitSafe splits the restored string on a raw byte index, which can cut +// a multi-byte UTF-8 rune. The emit half is then json.Marshal'd, which replaces +// the truncated bytes with U+FFFD, corrupting non-ASCII restored PII. +func TestAnthropicCodec_MultibyteRuneNotCorrupted(t *testing.T) { + // keep = len("XX") - 1 = 1, so the emit/hold boundary lands inside the + // 2-byte 'é' of the restored "José". + codec := newAnthropicCodec(map[string]string{"XX": "José"}) + ev := codec.transformEvent(sseLines( + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"XX"}}` + "\n\n")) + stop := codec.transformEvent(sseLines( + "event: content_block_stop\n" + + `data: {"type":"content_block_stop","index":0}` + "\n\n")) + + assembled := deltaPayloads(t, string(ev)+string(stop)) + if strings.Contains(assembled, "�") { + t.Errorf("multi-byte rune corrupted to U+FFFD: assembled=%q", assembled) + } + if assembled != "José" { + t.Errorf("client-assembled deltas = %q, want %q", assembled, "José") + } +} + +// BUG 4: tool-call arguments arrive as fragments of a JSON string +// (input_json_delta.partial_json). Restoring a value containing a JSON +// metacharacter (here backslashes in a Windows path) splices it in unescaped, so +// the arguments JSON the client reconstructs is invalid. +func TestAnthropicCodec_ToolArgsRemainValidJSON(t *testing.T) { + codec := newAnthropicCodec(map[string]string{"MASKED_PATH": `C:\Temp\x`}) + ev := codec.transformEvent(sseLines( + "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"path\":\"MASKED_PATH\"}"}}` + "\n\n")) + stop := codec.transformEvent(sseLines( + "event: content_block_stop\n" + + `data: {"type":"content_block_stop","index":0}` + "\n\n")) + + // The client concatenates partial_json fragments and parses them as the + // tool-call arguments. + assembled := deltaPayloads(t, string(ev)+string(stop)) + var args map[string]interface{} + if err := json.Unmarshal([]byte(assembled), &args); err != nil { + t.Errorf("reconstructed tool-call arguments are not valid JSON: %v\nassembled=%q", err, assembled) + } +} + +// BUG 6: streamSSEResponse flushes the last buffered event on EOF but never +// flushes each codec's held-back carry tail. If the upstream ends without a +// content_block_stop / *.done for an open channel, the trailing keep bytes are +// silently dropped, truncating the client's output. +func TestStreamSSEResponse_FlushesCarryTailOnEOF(t *testing.T) { + // keep = len("SECRET") - 1 = 5, so " Jane" is held back after emitting + // "Hello". With no content_block_stop, the tail must still reach the client. + upstream := "event: content_block_delta\n" + + `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello Jane"}}` + "\n\n" + + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(upstream)), + } + + client, server := net.Pipe() + codec := newAnthropicCodec(map[string]string{"SECRET": "unused"}) + errCh := make(chan error, 1) + go func() { + errCh <- streamSSEResponse(server, resp, codec) + server.Close() + }() + + parsed, err := http.ReadResponse(bufio.NewReader(client), nil) + if err != nil { + t.Fatalf("client failed to parse response: %v", err) + } + body, err := io.ReadAll(parsed.Body) + if err != nil { + t.Fatalf("client failed to read body: %v", err) + } + if err := <-errCh; err != nil { + t.Fatalf("streamSSEResponse returned error: %v", err) + } + + if assembled := deltaPayloads(t, string(body)); assembled != "Hello Jane" { + t.Errorf("client-assembled deltas = %q, want %q (carry tail dropped on EOF)", assembled, "Hello Jane") + } +} From dca148f6de95d23f644a90601a2ff5dc847635ee Mon Sep 17 00:00:00 2001 From: Davidnet Date: Mon, 6 Jul 2026 15:44:44 -0400 Subject: [PATCH 12/13] fix: rework streaming PII restore to a raw-carry single-pass matcher Fixes the six failing regression tests from 0e9b02d: - Carry buffers now hold RAW (pre-restore) text and restoration runs as a single left-to-right pass (streamRestore) that emits originals exactly once, so a restored original that coincides with another mapping's dummy can no longer be re-substituted across delta boundaries. - Text is only split on rune boundaries; an incomplete trailing UTF-8 rune is held back instead of being emitted truncated (no more U+FFFD). - Only a suffix that is a proper prefix of a dummy is held back, instead of a fixed longest-dummy-1 tail, which also reduces streaming latency. - Tool-call argument fragments (input_json_delta, function_call_arguments, chat tool_calls) restore with JSON-escaped originals so the arguments the client reassembles stay valid JSON; .done arguments fields likewise. - The OpenAI codec now handles /v1/chat/completions chunks (chat.completion.chunk: choices[].delta.content and tool_calls), which previously passed through unrestored. - Codecs expose flushTail(); streamSSEResponse flushes held carry tails at EOF so streams ending without .done/content_block_stop are not truncated. - Tail flushes after *.completed reuse the channel's recorded delta type instead of deriving a bogus response.completed.delta. --- src/backend/proxy/codec_anthropic.go | 30 ++++- src/backend/proxy/codec_openai.go | 167 ++++++++++++++++++++++++--- src/backend/proxy/streaming.go | 136 ++++++++++++++++++++-- 3 files changed, 308 insertions(+), 25 deletions(-) diff --git a/src/backend/proxy/codec_anthropic.go b/src/backend/proxy/codec_anthropic.go index 2181c827..a1feb3c0 100644 --- a/src/backend/proxy/codec_anthropic.go +++ b/src/backend/proxy/codec_anthropic.go @@ -19,7 +19,7 @@ const ( // block stops. type anthropicCodec struct { restoreCore - carry map[int]string // un-emitted tail per content-block index + carry map[int]string // un-emitted RAW (pre-restore) tail per content-block index kind map[int]string // delta type per index ("text_delta" / "input_json_delta") } @@ -73,8 +73,12 @@ func (s *anthropicCodec) transformEvent(lines [][]byte) []byte { } s.masked.WriteString(raw) // raw model output (text or tool args), for audit s.kind[evt.Index] = evt.Delta.Type - buf := s.carry[evt.Index] + raw - emit, hold := splitSafe(s.restore(buf), s.keep) + // input_json_delta fragments are pieces of a JSON string, so restored + // originals must be JSON-escaped to keep the reassembled arguments valid. + jsonCtx := evt.Delta.Type == deltaTypeInputJSON + // The carry holds RAW (pre-restore) text: restored output is emitted + // once and never rescanned, so restoration cannot chain across deltas. + emit, hold := s.streamRestore(s.carry[evt.Index]+raw, jsonCtx) s.carry[evt.Index] = hold // Rewrite only the data: line; keep the original event:/blank lines. out := make([][]byte, len(lines)) @@ -85,7 +89,9 @@ func (s *anthropicCodec) transformEvent(lines [][]byte) []byte { case evt.Type == "content_block_stop": var out []byte if tail := s.carry[evt.Index]; tail != "" { - out = append(out, deltaEvent(evt.Index, s.kind[evt.Index], tail)...) + kind := s.kind[evt.Index] + restored := s.flushCarry(tail, kind == deltaTypeInputJSON) + out = append(out, deltaEvent(evt.Index, kind, restored)...) delete(s.carry, evt.Index) delete(s.kind, evt.Index) } @@ -96,6 +102,22 @@ func (s *anthropicCodec) transformEvent(lines [][]byte) []byte { } } +// flushTail emits every non-empty carry as a synthetic delta so a stream that +// ends (EOF) without content_block_stop events doesn't truncate the output. +func (s *anthropicCodec) flushTail() []byte { + var out []byte + for idx, tail := range s.carry { + if tail == "" { + continue + } + kind := s.kind[idx] + restored := s.flushCarry(tail, kind == deltaTypeInputJSON) + out = append(out, deltaEvent(idx, kind, restored)...) + } + s.carry = map[int]string{} + return out +} + // deltaField returns the JSON field that carries the payload for a delta type. func deltaField(deltaType string) string { if deltaType == deltaTypeInputJSON { diff --git a/src/backend/proxy/codec_openai.go b/src/backend/proxy/codec_openai.go index cfe8ba1a..7fb96f8e 100644 --- a/src/backend/proxy/codec_openai.go +++ b/src/backend/proxy/codec_openai.go @@ -25,16 +25,24 @@ import ( // does not recognise pass through byte-for-byte. type openaiCodec struct { restoreCore - carry map[string]string // un-emitted (already-restored) tail per channel + carry map[string]string // un-emitted RAW (pre-restore) tail per channel + kind map[string]string // delta event type per channel, for tail-flush framing } func newOpenAICodec(mapping map[string]string) *openaiCodec { return &openaiCodec{ restoreCore: newRestoreCore(mapping), carry: map[string]string{}, + kind: map[string]string{}, } } +// argumentsChannel reports whether an event type carries a fragment of a JSON +// string (tool-call arguments), whose restored originals must be JSON-escaped. +func argumentsChannel(typ string) bool { + return strings.Contains(typ, "function_call_arguments") || strings.Contains(typ, "custom_tool_call_input") +} + // doneFields are the fields a Responses `*.done` event uses to carry the full // value, by event kind: output text, tool-call arguments, and refusals. var doneFields = []string{jsonKeyText, "arguments", "refusal"} @@ -59,6 +67,19 @@ func (c *openaiCodec) transformEvent(lines [][]byte) []byte { return concatLines(lines) // non-JSON payload (e.g. "[DONE]"): pass through } typ, _ := obj[jsonKeyType].(string) + + // The /v1/chat/completions stream grammar has no top-level "type": chunks + // are identified by object == "chat.completion.chunk" and carry text in + // choices[].delta.content (and tool args in choices[].delta.tool_calls[]. + // function.arguments). Restore those too — otherwise chat-completions + // streams fail open and dummy values reach the client. + if object, _ := obj["object"].(string); object == "chat.completion.chunk" { + if c.transformChatChunk(obj) { + return rewriteDataLine(lines, dataIdx, obj) + } + return concatLines(lines) + } + key := channelKey(obj) switch { @@ -68,7 +89,10 @@ func (c *openaiCodec) transformEvent(lines [][]byte) []byte { return concatLines(lines) // non-text delta (e.g. audio): pass through } c.masked.WriteString(delta) // raw model output, for audit - emit, hold := splitSafe(c.restore(c.carry[key]+delta), c.keep) + c.kind[key] = typ + // The carry holds RAW (pre-restore) text: restored output is emitted + // once and never rescanned, so restoration cannot chain across deltas. + emit, hold := c.streamRestore(c.carry[key]+delta, argumentsChannel(typ)) c.carry[key] = hold obj["delta"] = emit return rewriteDataLine(lines, dataIdx, obj) @@ -78,8 +102,10 @@ func (c *openaiCodec) transformEvent(lines [][]byte) []byte { // Flush any held-back delta tail as a synthetic delta so incremental // renderers see the complete text before the terminating .done event. if tail := c.carry[key]; tail != "" { + restored := c.flushCarry(tail, argumentsChannel(c.kind[key])) + out = append(out, c.tailDeltaEvent(typ, key, obj, restored)...) delete(c.carry, key) - out = append(out, c.tailDeltaEvent(typ, obj, tail)...) + delete(c.kind, key) } // Restore the full values the event repeats, including nested copies: // content_part.done nests part.text, output_item.done nests @@ -87,7 +113,7 @@ func (c *openaiCodec) transformEvent(lines [][]byte) []byte { // with output[].content[].text — clients (e.g. Codex) render their final // message from these, not from the deltas. (Not added to the audit // accumulator: the deltas above already captured this text.) - restoreTextFields(obj, c.restore) + restoreTextFields(obj, c.restore, c.restoreJSON) return append(out, rewriteDataLine(lines, dataIdx, obj)...) default: @@ -95,37 +121,152 @@ func (c *openaiCodec) transformEvent(lines [][]byte) []byte { } } +// transformChatChunk restores PII in a chat.completion.chunk in place, +// returning whether anything changed. Assistant text streams per choice; +// tool-call arguments stream per (choice, tool-call) and are JSON-string +// fragments, so their originals are JSON-escaped. +func (c *openaiCodec) transformChatChunk(obj map[string]interface{}) bool { + choices, _ := obj["choices"].([]interface{}) + changed := false + for _, ch := range choices { + chm, ok := ch.(map[string]interface{}) + if !ok { + continue + } + idx, _ := chm["index"].(float64) + key := fmt.Sprintf("chat:%d", int(idx)) + delta, _ := chm["delta"].(map[string]interface{}) + if delta != nil { + if content, ok := delta["content"].(string); ok && content != "" { + c.masked.WriteString(content) + emit, hold := c.streamRestore(c.carry[key]+content, false) + c.carry[key] = hold + delta["content"] = emit + changed = true + } + if tcs, ok := delta["tool_calls"].([]interface{}); ok { + for _, tc := range tcs { + tcm, ok := tc.(map[string]interface{}) + if !ok { + continue + } + ti, _ := tcm["index"].(float64) + fn, _ := tcm["function"].(map[string]interface{}) + if fn == nil { + continue + } + if args, ok := fn["arguments"].(string); ok && args != "" { + c.masked.WriteString(args) + tcKey := fmt.Sprintf("chat:%d:tc:%d", int(idx), int(ti)) + emit, hold := c.streamRestore(c.carry[tcKey]+args, true) + c.carry[tcKey] = hold + fn["arguments"] = emit + changed = true + } + } + } + } + // The final chunk for a choice carries finish_reason: flush the held + // tail into it so the client's assembled text is complete. + if fr, ok := chm["finish_reason"].(string); ok && fr != "" { + if tail := c.carry[key]; tail != "" { + delete(c.carry, key) + if delta == nil { + delta = map[string]interface{}{} + chm["delta"] = delta + } + s, _ := delta["content"].(string) + delta["content"] = s + c.flushCarry(tail, false) + changed = true + } + } + } + return changed +} + +// flushTail emits every non-empty carry as a synthetic delta so a stream that +// ends (EOF) without .done events doesn't truncate the output. Channels whose +// delta event type is unknown (chat.completion.chunk) are flushed as a chunk. +func (c *openaiCodec) flushTail() []byte { + var out []byte + for key, tail := range c.carry { + if tail == "" { + continue + } + kind := c.kind[key] + restored := c.flushCarry(tail, argumentsChannel(kind)) + if kind != "" { + d := map[string]interface{}{jsonKeyType: kind, "delta": restored} + b, _ := json.Marshal(d) + out = append(out, []byte("event: "+kind+"\n")...) + out = append(out, append(append([]byte("data: "), b...), '\n')...) + out = append(out, '\n') + continue + } + // chat.completion.chunk channel: key is "chat:" (tool-call + // argument tails stay held — a truncated arguments fragment is not + // actionable by the client anyway). + var idx int + if _, err := fmt.Sscanf(key, "chat:%d", &idx); err != nil || strings.Contains(key, ":tc:") { + continue + } + chunk := map[string]interface{}{ + "object": "chat.completion.chunk", + "choices": []interface{}{map[string]interface{}{ + "index": idx, + "delta": map[string]interface{}{"content": restored}, + }}, + } + b, _ := json.Marshal(chunk) + out = append(out, append(append([]byte("data: "), b...), '\n')...) + out = append(out, '\n') + } + c.carry = map[string]string{} + return out +} + // restoreTextFields recursively restores PII in every string value held by a // known text-carrying field (doneFields) anywhere in the event payload. Only // those fields are rewritten — ids, signatures, and other opaque strings that -// could contain a dummy value as a substring pass through untouched. -func restoreTextFields(v interface{}, restore func(string) string) { +// could contain a dummy value as a substring pass through untouched. The +// "arguments" field is the full tool-call JSON document as a string, so its +// originals must be JSON-escaped (restoreJSON) to keep it parseable. +func restoreTextFields(v interface{}, restore, restoreJSON func(string) string) { switch t := v.(type) { case map[string]interface{}: for k, val := range t { if s, ok := val.(string); ok { for _, field := range doneFields { if k == field { - t[k] = restore(s) + if k == "arguments" { + t[k] = restoreJSON(s) + } else { + t[k] = restore(s) + } break } } continue } - restoreTextFields(val, restore) + restoreTextFields(val, restore, restoreJSON) } case []interface{}: for _, item := range t { - restoreTextFields(item, restore) + restoreTextFields(item, restore, restoreJSON) } } } // tailDeltaEvent builds a synthetic `*.delta` event mirroring a `*.done` event's -// channel, carrying the already-restored held-back tail. The payload is emitted -// as-is (the carry buffer stores post-restore text). -func (c *openaiCodec) tailDeltaEvent(doneType string, obj map[string]interface{}, tail string) []byte { - deltaType := strings.TrimSuffix(doneType, ".done") + ".delta" +// channel, carrying the already-restored held-back tail. The delta event type +// is taken from the channel's recorded delta kind when available (a +// "*.completed" event's type cannot be derived by suffix-trimming — it would +// yield a bogus "response.completed.delta"). +func (c *openaiCodec) tailDeltaEvent(doneType, key string, obj map[string]interface{}, tail string) []byte { + deltaType := c.kind[key] + if deltaType == "" { + deltaType = strings.TrimSuffix(doneType, ".done") + ".delta" + } d := map[string]interface{}{jsonKeyType: deltaType, "delta": tail} for _, f := range []string{"item_id", "output_index", "content_index"} { if v, ok := obj[f]; ok { diff --git a/src/backend/proxy/streaming.go b/src/backend/proxy/streaming.go index 7b867efd..f11d4910 100644 --- a/src/backend/proxy/streaming.go +++ b/src/backend/proxy/streaming.go @@ -11,8 +11,10 @@ import ( "net/http" "os" "path/filepath" + "sort" "strings" "time" + "unicode/utf8" "github.com/dataiku/kiji-proxy/src/backend/processor" "github.com/dataiku/kiji-proxy/src/backend/providers" @@ -116,6 +118,11 @@ type sniffedBody struct { // (up to and including the blank terminator) and returns the bytes to write. type streamCodec interface { transformEvent(lines [][]byte) []byte + // flushTail emits any text still held in carry buffers as synthetic delta + // events. Called when the upstream stream ends (EOF) so a stream that never + // sent its terminating .done/content_block_stop doesn't silently truncate + // the client's output. + flushTail() []byte // restore replaces every masked (dummy) value with its original. Exposed so // the caller can restore the accumulated audit text after the stream ends. restore(text string) string @@ -125,26 +132,61 @@ type streamCodec interface { } // restoreCore holds the provider-agnostic restore state shared by every codec: -// the dummy→original replacer, the hold-back length, and the pre-restore model -// output accumulated for audit logging. Codecs embed it so they only implement -// their own event grammar. +// the dummy→original mapping (plus a JSON-escaped variant for tool-argument +// fragments), the hold-back length, and the pre-restore model output +// accumulated for audit logging. Codecs embed it so they only implement their +// own event grammar. type restoreCore struct { - replacer *strings.Replacer - keep int // bytes to hold back = longest dummy length - 1 - masked strings.Builder // raw model output (pre-restore), accumulated for logging + keys []string // dummy values, longest-first + vals map[string]string // dummy → original + jsonVals map[string]string // dummy → original, JSON-string-escaped + replacer *strings.Replacer // single-pass plain restorer (audit / final flush) + jsonReplacer *strings.Replacer // single-pass restorer with JSON-escaped originals + keep int // longest dummy length - 1 (max possible hold-back) + masked strings.Builder // raw model output (pre-restore), accumulated for logging +} + +// jsonEscapeString returns s escaped as JSON string content (without the +// surrounding quotes), for splicing into a fragment of a JSON document. +func jsonEscapeString(s string) string { + b, err := json.Marshal(s) + if err != nil { + return s + } + return string(b[1 : len(b)-1]) } func newRestoreCore(mapping map[string]string) restoreCore { keep := 0 - for masked := range mapping { + keys := make([]string, 0, len(mapping)) + vals := make(map[string]string, len(mapping)) + jsonVals := make(map[string]string, len(mapping)) + jsonMapping := make(map[string]string, len(mapping)) + for masked, original := range mapping { + if masked == "" { + continue + } if len(masked) > keep { keep = len(masked) } + keys = append(keys, masked) + vals[masked] = original + esc := jsonEscapeString(original) + jsonVals[masked] = esc + jsonMapping[masked] = esc } if keep > 0 { keep-- } - return restoreCore{replacer: processor.BuildRestorer(mapping), keep: keep} + sort.Slice(keys, func(i, j int) bool { return len(keys[i]) > len(keys[j]) }) + return restoreCore{ + keys: keys, + vals: vals, + jsonVals: jsonVals, + replacer: processor.BuildRestorer(mapping), + jsonReplacer: processor.BuildRestorer(jsonMapping), + keep: keep, + } } // restore replaces every masked (dummy) value with its original in a single @@ -155,6 +197,76 @@ func (c *restoreCore) restore(text string) string { return c.replacer.Replace(text) } +// restoreJSON is restore for text that is a fragment (or whole) of a JSON +// string value, e.g. tool-call arguments: originals are JSON-escaped so a +// restored quote or backslash cannot break the JSON the client reassembles. +func (c *restoreCore) restoreJSON(text string) string { + return c.jsonReplacer.Replace(text) +} + +// streamRestore restores dummies in buf — the raw (pre-restore) text +// accumulated for one channel — in a single left-to-right pass, returning the +// restored text that is safe to emit now and the RAW tail that must be held +// back and re-fed on the next delta. +// +// Holding back raw text (not restored text) is what preserves the single-pass +// guarantee across deltas: restored originals are emitted immediately and never +// rescanned, so an original that coincides with another mapping's dummy cannot +// be substituted a second time. The hold is also minimal — only a suffix that +// is a proper prefix of some dummy (a placeholder possibly still arriving) or +// an incomplete trailing UTF-8 rune is withheld, and text is only ever split on +// rune boundaries so multi-byte characters are never corrupted. +func (c *restoreCore) streamRestore(buf string, jsonCtx bool) (emit, hold string) { + vals := c.vals + if jsonCtx { + vals = c.jsonVals + } + var out strings.Builder + i := 0 + for i < len(buf) { + rest := buf[i:] + // A dummy may be arriving split across deltas: if everything that + // remains is a proper prefix of some dummy, hold it raw. + for _, k := range c.keys { + if len(rest) < len(k) && strings.HasPrefix(k, rest) { + return out.String(), rest + } + } + // Complete dummy at this position: emit its original (longest wins). + matched := false + for _, k := range c.keys { + if strings.HasPrefix(rest, k) { + out.WriteString(vals[k]) + i += len(k) + matched = true + break + } + } + if matched { + continue + } + // No match: emit one rune. An incomplete trailing rune (a multi-byte + // character split across deltas) is held back, never emitted truncated. + r, size := utf8.DecodeRuneInString(rest) + if r == utf8.RuneError && size == 1 && !utf8.FullRuneInString(rest) { + return out.String(), rest + } + out.WriteString(rest[:size]) + i += size + } + return out.String(), "" +} + +// flushCarry restores a raw held-back tail in full — used when its channel +// terminates (content_block_stop / *.done / stream EOF) and no more input can +// complete a partial placeholder. +func (c *restoreCore) flushCarry(raw string, jsonCtx bool) string { + if jsonCtx { + return c.jsonReplacer.Replace(raw) + } + return c.replacer.Replace(raw) +} + func (c *restoreCore) maskedOutput() string { return c.masked.String() } @@ -261,6 +373,14 @@ func streamSSEResponse(conn net.Conn, resp *http.Response, codec streamCodec) er if ferr := flush(); ferr != nil { // trailing event w/o blank line return ferr } + // The upstream ended without terminating every channel (no + // .done / content_block_stop): flush any held-back carry tails + // so the client's output is not silently truncated. + if tail := codec.flushTail(); len(tail) > 0 { + if werr := writeChunk(tail); werr != nil { + return werr + } + } break } return err From 479c7a7dc183b947581bbede6be4e1d272c33d76 Mon Sep 17 00:00:00 2001 From: Davidnet Date: Mon, 6 Jul 2026 16:50:06 -0400 Subject: [PATCH 13/13] fix: extract repeated delta/content JSON keys into constants golangci-lint (goconst) flagged the literal "delta" and "content" map keys repeated across the SSE codecs. --- src/backend/proxy/codec_anthropic.go | 6 +++--- src/backend/proxy/codec_openai.go | 24 ++++++++++++------------ src/backend/proxy/handler.go | 12 +++++++----- 3 files changed, 22 insertions(+), 20 deletions(-) diff --git a/src/backend/proxy/codec_anthropic.go b/src/backend/proxy/codec_anthropic.go index a1feb3c0..a5dbfda1 100644 --- a/src/backend/proxy/codec_anthropic.go +++ b/src/backend/proxy/codec_anthropic.go @@ -130,9 +130,9 @@ func deltaField(deltaType string) string { // input_json_delta carrying the given (restored) payload. func deltaDataLine(index int, deltaType, payload string) []byte { b, _ := json.Marshal(map[string]interface{}{ - jsonKeyType: "content_block_delta", - "index": index, - "delta": map[string]interface{}{jsonKeyType: deltaType, deltaField(deltaType): payload}, + jsonKeyType: "content_block_delta", + "index": index, + jsonKeyDelta: map[string]interface{}{jsonKeyType: deltaType, deltaField(deltaType): payload}, }) out := append([]byte("data: "), b...) return append(out, '\n') diff --git a/src/backend/proxy/codec_openai.go b/src/backend/proxy/codec_openai.go index 7fb96f8e..d9b592c3 100644 --- a/src/backend/proxy/codec_openai.go +++ b/src/backend/proxy/codec_openai.go @@ -84,7 +84,7 @@ func (c *openaiCodec) transformEvent(lines [][]byte) []byte { switch { case strings.HasSuffix(typ, ".delta"): - delta, ok := obj["delta"].(string) + delta, ok := obj[jsonKeyDelta].(string) if !ok { return concatLines(lines) // non-text delta (e.g. audio): pass through } @@ -94,7 +94,7 @@ func (c *openaiCodec) transformEvent(lines [][]byte) []byte { // once and never rescanned, so restoration cannot chain across deltas. emit, hold := c.streamRestore(c.carry[key]+delta, argumentsChannel(typ)) c.carry[key] = hold - obj["delta"] = emit + obj[jsonKeyDelta] = emit return rewriteDataLine(lines, dataIdx, obj) case strings.HasSuffix(typ, ".done"), strings.HasSuffix(typ, ".completed"): @@ -135,13 +135,13 @@ func (c *openaiCodec) transformChatChunk(obj map[string]interface{}) bool { } idx, _ := chm["index"].(float64) key := fmt.Sprintf("chat:%d", int(idx)) - delta, _ := chm["delta"].(map[string]interface{}) + delta, _ := chm[jsonKeyDelta].(map[string]interface{}) if delta != nil { - if content, ok := delta["content"].(string); ok && content != "" { + if content, ok := delta[jsonKeyContent].(string); ok && content != "" { c.masked.WriteString(content) emit, hold := c.streamRestore(c.carry[key]+content, false) c.carry[key] = hold - delta["content"] = emit + delta[jsonKeyContent] = emit changed = true } if tcs, ok := delta["tool_calls"].([]interface{}); ok { @@ -173,10 +173,10 @@ func (c *openaiCodec) transformChatChunk(obj map[string]interface{}) bool { delete(c.carry, key) if delta == nil { delta = map[string]interface{}{} - chm["delta"] = delta + chm[jsonKeyDelta] = delta } - s, _ := delta["content"].(string) - delta["content"] = s + c.flushCarry(tail, false) + s, _ := delta[jsonKeyContent].(string) + delta[jsonKeyContent] = s + c.flushCarry(tail, false) changed = true } } @@ -196,7 +196,7 @@ func (c *openaiCodec) flushTail() []byte { kind := c.kind[key] restored := c.flushCarry(tail, argumentsChannel(kind)) if kind != "" { - d := map[string]interface{}{jsonKeyType: kind, "delta": restored} + d := map[string]interface{}{jsonKeyType: kind, jsonKeyDelta: restored} b, _ := json.Marshal(d) out = append(out, []byte("event: "+kind+"\n")...) out = append(out, append(append([]byte("data: "), b...), '\n')...) @@ -213,8 +213,8 @@ func (c *openaiCodec) flushTail() []byte { chunk := map[string]interface{}{ "object": "chat.completion.chunk", "choices": []interface{}{map[string]interface{}{ - "index": idx, - "delta": map[string]interface{}{"content": restored}, + "index": idx, + jsonKeyDelta: map[string]interface{}{jsonKeyContent: restored}, }}, } b, _ := json.Marshal(chunk) @@ -267,7 +267,7 @@ func (c *openaiCodec) tailDeltaEvent(doneType, key string, obj map[string]interf if deltaType == "" { deltaType = strings.TrimSuffix(doneType, ".done") + ".delta" } - d := map[string]interface{}{jsonKeyType: deltaType, "delta": tail} + d := map[string]interface{}{jsonKeyType: deltaType, jsonKeyDelta: tail} for _, f := range []string{"item_id", "output_index", "content_index"} { if v, ok := obj[f]; ok { d[f] = v diff --git a/src/backend/proxy/handler.go b/src/backend/proxy/handler.go index dcb03110..5771c7cd 100644 --- a/src/backend/proxy/handler.go +++ b/src/backend/proxy/handler.go @@ -29,8 +29,10 @@ const paramLimit = "limit" // JSON field names reused across the response wrappers and the SSE codecs. const ( - jsonKeyType = "type" - jsonKeyText = "text" + jsonKeyType = "type" + jsonKeyText = "text" + jsonKeyDelta = "delta" + jsonKeyContent = "content" ) // Handler handles HTTP requests and proxies them to LLM provider @@ -593,7 +595,7 @@ func (h *Handler) MaskPIIInTextWithLogging(ctx context.Context, text, site strin func wrapPIICheckMessage(text, site, transactionID string) string { envelope := map[string]any{ "messages": []map[string]string{ - {"role": "user", "content": text}, + {"role": "user", jsonKeyContent: text}, }, "_transaction_id": transactionID, } @@ -696,8 +698,8 @@ func (h *Handler) LogStreamedResponse(ctx context.Context, transactionID, masked // transaction ID can be attached and the log UI can parse it consistently. wrap := func(text string) string { b, err := json.Marshal(map[string]interface{}{ - "streamed": true, - "content": []map[string]interface{}{{jsonKeyType: jsonKeyText, jsonKeyText: text}}, + "streamed": true, + jsonKeyContent: []map[string]interface{}{{jsonKeyType: jsonKeyText, jsonKeyText: text}}, }) if err != nil { return text