From c62079c053b98e61a5700b33095475edf4213ebd Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Wed, 5 Aug 2026 11:59:37 +0100 Subject: [PATCH] refactor(coderd): optimize chatdebug (#27129) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a bund of optimizations to chatdebug: In `coderd/x/chatd/chatdebug`: - Adds a benchmark (excluding LLM and database stuff) - Replaces string concatenation with strings.Builder when accumulating stream parts (~105,000ns -> ~50,00ns) - Removes double JSON encode in RecordingTransport (114,000ns -> 64,000ns) In `coderd/util/strings`: - Adds a benchmark for Truncate - Removes unnecessary allocations in Truncate (~110,000ns -> 1,550ns in truncation case, 1 alloc -> 0 allocs in no truncation case) > 🤖 Claude helped with this. --- coderd/util/strings/strings.go | 16 +- coderd/util/strings/strings_test.go | 24 ++ .../x/chatd/chatdebug/bench_internal_test.go | 350 ++++++++++++++++++ coderd/x/chatd/chatdebug/model.go | 112 ++++-- .../model_normalization_internal_test.go | 55 ++- coderd/x/chatd/chatdebug/redaction.go | 75 ++-- coderd/x/chatd/chatdebug/transport.go | 90 +++-- .../chatdebug/transport_internal_test.go | 43 +++ 8 files changed, 664 insertions(+), 101 deletions(-) create mode 100644 coderd/x/chatd/chatdebug/bench_internal_test.go diff --git a/coderd/util/strings/strings.go b/coderd/util/strings/strings.go index d2594b80a0..11a05cb82a 100644 --- a/coderd/util/strings/strings.go +++ b/coderd/util/strings/strings.go @@ -64,8 +64,19 @@ func Truncate(s string, n int, opts ...TruncateOption) string { if n < 1 { return "" } - runes := []rune(s) - if len(runes) <= n { + + // Find the byte offset of the (n+1)th rune, if any; early exit + // avoids decoding s past what's needed. + runeCount := 0 + cutoff := -1 + for i := range s { + runeCount++ + if runeCount > n { + cutoff = i + break + } + } + if cutoff < 0 { return s } @@ -73,6 +84,7 @@ func Truncate(s string, n int, opts ...TruncateOption) string { if options&TruncateWithEllipsis != 0 { maxLen-- } + runes := []rune(s[:cutoff]) var sb strings.Builder if options&TruncateWithFullWords != 0 { // Convert the rune-safe prefix to a string, then find diff --git a/coderd/util/strings/strings_test.go b/coderd/util/strings/strings_test.go index 494246c6cf..b6f8a3204b 100644 --- a/coderd/util/strings/strings_test.go +++ b/coderd/util/strings/strings_test.go @@ -1,6 +1,7 @@ package strings_test import ( + "bytes" "fmt" "testing" @@ -32,6 +33,7 @@ func TestTruncate(t *testing.T) { {"foo", 1, "f", nil}, {"foo", 0, "", nil}, {"foo", -1, "", nil}, + {"", 5, "", nil}, {"foo bar", 7, "foo bar", []strings.TruncateOption{strings.TruncateWithEllipsis}}, {"foo bar", 6, "foo b…", []strings.TruncateOption{strings.TruncateWithEllipsis}}, {"foo bar", 5, "foo …", []strings.TruncateOption{strings.TruncateWithEllipsis}}, @@ -81,6 +83,28 @@ func TestTruncate(t *testing.T) { } } +func BenchmarkTruncate(b *testing.B) { + b.Run("NoTruncationNeeded", func(b *testing.B) { + s := "a short string well under the limit" + b.ReportAllocs() + for b.Loop() { + strings.Truncate(s, 1000) + } + }) + + b.Run("ActualTruncation", func(b *testing.B) { + var buf bytes.Buffer + for range 2000 { + buf.WriteString("日本語テスト word ") + } + s := buf.String() + b.ReportAllocs() + for b.Loop() { + strings.Truncate(s, 100, strings.TruncateWithEllipsis, strings.TruncateWithFullWords) + } + }) +} + func TestUISanitize(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chatdebug/bench_internal_test.go b/coderd/x/chatd/chatdebug/bench_internal_test.go new file mode 100644 index 0000000000..d5b6a4ca31 --- /dev/null +++ b/coderd/x/chatd/chatdebug/bench_internal_test.go @@ -0,0 +1,350 @@ +package chatdebug + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "slices" + "strings" + "testing" + + "charm.land/fantasy" + "github.com/google/uuid" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/pubsub" + "github.com/coder/coder/v2/coderd/x/chatd/chattest" +) + +// noopStore implements database.Store by embedding it (nil) and +// overriding only the methods chatdebug's hot paths call. Every +// override returns immediately with a canned value and does no I/O, +// so any time spent inside Service methods that use this store is +// chatdebug's own CPU/allocation cost, not Postgres's. +type noopStore struct { + database.Store +} + +func (noopStore) GetChatDebugLoggingAllowUsers(context.Context) (bool, error) { + return true, nil +} + +func (noopStore) GetUserChatDebugLoggingEnabled(context.Context, uuid.UUID) (bool, error) { + return true, nil +} + +func (noopStore) InsertChatDebugStep( + _ context.Context, + arg database.InsertChatDebugStepParams, +) (database.ChatDebugStep, error) { + return database.ChatDebugStep{ + ID: uuid.New(), + RunID: arg.RunID, + ChatID: arg.ChatID, + StepNumber: arg.StepNumber, + Operation: arg.Operation, + Status: arg.Status, + }, nil +} + +func (noopStore) UpdateChatDebugStep( + _ context.Context, + arg database.UpdateChatDebugStepParams, +) (database.ChatDebugStep, error) { + return database.ChatDebugStep{ + ID: arg.ID, + ChatID: arg.ChatID, + }, nil +} + +// TouchChatDebugStepAndRun and GetChatDebugStepsByRunID aren't on the +// hot path these benchmarks exercise (heartbeat interval and step +// retry-collision handling respectively), but overriding them avoids +// a nil-dereference panic if a future benchmark configuration change +// reaches them through the embedded nil database.Store. +func (noopStore) TouchChatDebugStepAndRun(context.Context, database.TouchChatDebugStepAndRunParams) error { + return nil +} + +func (noopStore) GetChatDebugStepsByRunID(context.Context, uuid.UUID) ([]database.ChatDebugStep, error) { + return nil, nil +} + +var _ database.Store = noopStore{} + +// benchService builds a Service backed by noopStore and an in-memory +// pubsub so publishEvent still marshals and dispatches DebugEvent +// payloads, mirroring a real deployment's wiring, without any network +// or disk I/O. +func benchService(b *testing.B) *Service { + b.Helper() + ps := pubsub.NewInMemory() + b.Cleanup(func() { _ = ps.Close() }) + return NewService(noopStore{}, slog.Make(), ps) +} + +// benchCall builds a realistic multi-turn prompt with tool +// definitions so normalizeCall walks a representative message/tool +// shape. +func benchCall(nMessages, nTools int) fantasy.Call { + prompt := make(fantasy.Prompt, 0, nMessages) + for i := range nMessages { + if i%2 == 0 { + prompt = append(prompt, fantasy.NewUserMessage( + fmt.Sprintf("Please investigate issue #%d and summarize the relevant log lines around the failure.", i))) + continue + } + prompt = append(prompt, fantasy.Message{ + Role: fantasy.MessageRoleAssistant, + Content: []fantasy.MessagePart{ + fantasy.TextPart{Text: strings.Repeat("analysis details ", 40)}, + }, + }) + } + + tools := make([]fantasy.Tool, 0, nTools) + for i := range nTools { + tools = append(tools, fantasy.FunctionTool{ + Name: fmt.Sprintf("tool_%d", i), + Description: "Runs a workspace command and returns its output", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "command": map[string]any{"type": "string"}, + "timeout": map[string]any{"type": "integer"}, + }, + "required": []string{"command"}, + }, + }) + } + + return fantasy.Call{ + Prompt: prompt, + Tools: tools, + } +} + +func benchStreamParts(nDeltas int) []fantasy.StreamPart { + parts := make([]fantasy.StreamPart, 0, nDeltas+2) + for i := range nDeltas { + parts = append(parts, fantasy.StreamPart{ + Type: fantasy.StreamPartTypeTextDelta, + ID: "text-1", + Delta: fmt.Sprintf("token-%d ", i), + }) + } + parts = append(parts, + fantasy.StreamPart{Type: fantasy.StreamPartTypeToolCall, ID: "tool-1", ToolCallName: "tool_0"}, + fantasy.StreamPart{ + Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop, + Usage: fantasy.Usage{InputTokens: 512, OutputTokens: int64(nDeltas), TotalTokens: 512 + int64(nDeltas)}, + }, + ) + return parts +} + +// runContext returns a fresh RunContext/context pair for one +// benchmark iteration, mirroring a new chat turn. +func runContext(chatID uuid.UUID) context.Context { + return ContextWithRun(context.Background(), &RunContext{RunID: uuid.New(), ChatID: chatID}) +} + +func BenchmarkWrapModel_Generate(b *testing.B) { + svc := benchService(b) + chatID, ownerID := uuid.New(), uuid.New() + call := benchCall(6, 4) + resp := &fantasy.Response{ + Content: fantasy.ResponseContent{ + fantasy.TextContent{Text: strings.Repeat("response text ", 100)}, + fantasy.ToolCallContent{ToolCallID: "tool-1", ToolName: "tool_0", Input: `{"command":"ls"}`}, + }, + FinishReason: fantasy.FinishReasonStop, + Usage: fantasy.Usage{InputTokens: 512, OutputTokens: 128, TotalTokens: 640}, + } + inner := &chattest.FakeModel{GenerateFn: func(context.Context, fantasy.Call) (*fantasy.Response, error) { + return resp, nil + }} + model := WrapModel(inner, svc, RecorderOptions{ChatID: chatID, OwnerID: ownerID}) + + b.ReportAllocs() + for b.Loop() { + ctx := runContext(chatID) + if _, err := model.Generate(ctx, call); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkWrapModel_Stream(b *testing.B) { + svc := benchService(b) + chatID, ownerID := uuid.New(), uuid.New() + call := benchCall(6, 4) + parts := benchStreamParts(200) + inner := &chattest.FakeModel{StreamFn: func(context.Context, fantasy.Call) (fantasy.StreamResponse, error) { + return slices.Values(parts), nil + }} + model := WrapModel(inner, svc, RecorderOptions{ChatID: chatID, OwnerID: ownerID}) + + b.ReportAllocs() + for b.Loop() { + ctx := runContext(chatID) + seq, err := model.Stream(ctx, call) + if err != nil { + b.Fatal(err) + } + for range seq { //nolint:revive // draining the stream is the point of the benchmark. + } + } +} + +// cannedRoundTripper replays a fixed response body for every request, +// isolating RecordingTransport's own redaction/buffering cost from any +// real network or provider latency. +type cannedRoundTripper struct { + status int + header http.Header + body []byte + chunkSize int // 0 means return the whole body in one Read. +} + +func (rt *cannedRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + var body io.ReadCloser + if rt.chunkSize > 0 { + body = io.NopCloser(&chunkedReader{data: rt.body, chunkSize: rt.chunkSize}) + } else { + body = io.NopCloser(bytes.NewReader(rt.body)) + } + header := rt.header.Clone() + return &http.Response{ + StatusCode: rt.status, + Header: header, + Body: body, + ContentLength: -1, + Request: req, + }, nil +} + +// chunkedReader splits data into fixed-size Read() calls so the +// benchmark exercises recordingBody's incremental accumulation path +// the way a real SSE stream would, instead of handing back the whole +// body in a single Read. +type chunkedReader struct { + data []byte + chunkSize int +} + +func (c *chunkedReader) Read(p []byte) (int, error) { + if len(c.data) == 0 { + return 0, io.EOF + } + n := min(c.chunkSize, len(c.data), len(p)) + copy(p, c.data[:n]) + c.data = c.data[n:] + return n, nil +} + +func benchJSONResponseBody(n int) []byte { + type choice struct { + Index int `json:"index"` + Text string `json:"text"` + } + payload := struct { + ID string `json:"id"` + Object string `json:"object"` + Model string `json:"model"` + Choices []choice `json:"choices"` + Usage struct { + PromptTokens int `json:"prompt_tokens"` + CompletionTokens int `json:"completion_tokens"` + } `json:"usage"` + }{ + ID: "resp-bench", + Object: "chat.completion", + Model: "bench-model", + } + for i := range n { + payload.Choices = append(payload.Choices, choice{ + Index: i, + Text: strings.Repeat("word ", 50), + }) + } + data, err := json.Marshal(payload) + if err != nil { + panic(err) + } + return data +} + +func benchSSEResponseBody(nEvents int) []byte { + var buf bytes.Buffer + for i := range nEvents { + // bytes.Buffer.Write* never returns an error. + _, _ = fmt.Fprintf(&buf, "data: {\"delta\":\"token-%d \",\"index\":%d}\n\n", i, i) + } + _, _ = buf.WriteString("data: [DONE]\n\n") + return buf.Bytes() +} + +func newBenchRequest(b *testing.B, sink *attemptSink) *http.Request { + b.Helper() + req, err := http.NewRequestWithContext( + withAttemptSink(context.Background(), sink), + http.MethodPost, + "https://api.example.com/v1/chat/completions", + bytes.NewReader([]byte(`{"model":"bench-model","messages":[{"role":"user","content":"hi"}]}`)), + ) + if err != nil { + b.Fatal(err) + } + req.Header.Set("Authorization", "Bearer sk-should-be-redacted") + req.Header.Set("Content-Type", "application/json") + return req +} + +func BenchmarkRecordingTransport_RoundTrip_JSON(b *testing.B) { + header := http.Header{"Content-Type": []string{"application/json"}} + body := benchJSONResponseBody(20) + transport := &RecordingTransport{Base: &cannedRoundTripper{status: 200, header: header, body: body}} + + b.ReportAllocs() + for b.Loop() { + sink := &attemptSink{} + req := newBenchRequest(b, sink) + resp, err := transport.RoundTrip(req) + if err != nil { + b.Fatal(err) + } + if _, err := io.Copy(io.Discard, resp.Body); err != nil { + b.Fatal(err) + } + if err := resp.Body.Close(); err != nil { + b.Fatal(err) + } + } +} + +func BenchmarkRecordingTransport_RoundTrip_SSE(b *testing.B) { + header := http.Header{"Content-Type": []string{"text/event-stream"}} + body := benchSSEResponseBody(200) + transport := &RecordingTransport{Base: &cannedRoundTripper{status: 200, header: header, body: body, chunkSize: 256}} + + b.ReportAllocs() + for b.Loop() { + sink := &attemptSink{} + req := newBenchRequest(b, sink) + resp, err := transport.RoundTrip(req) + if err != nil { + b.Fatal(err) + } + if _, err := io.Copy(io.Discard, resp.Body); err != nil { + b.Fatal(err) + } + if err := resp.Body.Close(); err != nil { + b.Fatal(err) + } + } +} diff --git a/coderd/x/chatd/chatdebug/model.go b/coderd/x/chatd/chatdebug/model.go index e30a8a21e5..080a80ec8c 100644 --- a/coderd/x/chatd/chatdebug/model.go +++ b/coderd/x/chatd/chatdebug/model.go @@ -7,6 +7,7 @@ import ( "fmt" "iter" "reflect" + "strings" "sync" "sync/atomic" "unicode/utf8" @@ -460,6 +461,9 @@ func wrapStreamSeq( finishSeen bool finishReason fantasy.FinishReason content []normalizedContentPart + currentText *strings.Builder + currentTextIdx int + argBuilders = make(map[string]*strings.Builder) warnings []normalizedWarning streamDebugBytes int streamError any @@ -480,6 +484,7 @@ func wrapStreamSeq( close(heartbeatDone) summary.FinishReason = string(finishReason) + materializeStreamContent(content, currentText, currentTextIdx, argBuilders) resp := normalizedResponsePayload{ Content: content, @@ -528,7 +533,8 @@ func wrapStreamSeq( streamComplete.Store(true) } - content = appendNormalizedStreamContent(content, part, &streamDebugBytes) + content, currentText, currentTextIdx = appendNormalizedStreamContent( + content, currentText, currentTextIdx, argBuilders, part, &streamDebugBytes) if part.Type == fantasy.StreamPartTypeError || part.Error != nil { summary.ErrorCount++ @@ -758,14 +764,18 @@ func safeMarshalJSON(label string, value any) json.RawMessage { return append(json.RawMessage(nil), data...) } +// appendStreamContentText accumulates a delta into the current +// text/reasoning run. func appendStreamContentText( content []normalizedContentPart, + currentText *strings.Builder, + currentTextIdx int, partType string, delta string, streamDebugBytes *int, -) []normalizedContentPart { +) ([]normalizedContentPart, *strings.Builder, int) { if delta == "" { - return content + return content, currentText, currentTextIdx } remaining := maxStreamDebugTextBytes @@ -773,7 +783,7 @@ func appendStreamContentText( remaining -= *streamDebugBytes } if remaining <= 0 { - return content + return content, currentText, currentTextIdx } if len(delta) > remaining { cut := 0 @@ -790,18 +800,26 @@ func appendStreamContentText( delta = delta[:cut] } if delta == "" { - return content + return content, currentText, currentTextIdx } if len(content) == 0 || content[len(content)-1].Type != partType { + // A different-typed part is interrupting (or this is the + // first part): flush the outgoing run into content now, since + // currentText is about to be replaced and materializeStreamContent + // only ever sees whichever run is still open at the end. + if currentText != nil && currentTextIdx < len(content) { + content[currentTextIdx].Text = currentText.String() + } content = append(content, normalizedContentPart{Type: partType}) + currentText = &strings.Builder{} + currentTextIdx = len(content) - 1 } - last := &content[len(content)-1] - last.Text += delta + _, _ = currentText.WriteString(delta) if streamDebugBytes != nil { *streamDebugBytes += len(delta) } - return content + return content, currentText, currentTextIdx } // appendStreamToolInput accumulates incremental tool-input deltas @@ -809,6 +827,7 @@ func appendStreamContentText( // remain distinguishable in interrupted stream debug payloads. func appendStreamToolInput( content []normalizedContentPart, + argBuilders map[string]*strings.Builder, part fantasy.StreamPart, streamDebugBytes *int, ) []normalizedContentPart { @@ -842,26 +861,19 @@ func appendStreamToolInput( return content } - // Find the existing tool_input part for this specific tool call ID. - // Scan backwards through all content; tool_input deltas for the - // same call may be separated by text, reasoning, or source parts - // when streams interleave multiple tool invocations. - for i := len(content) - 1; i >= 0; i-- { - if content[i].Type == "tool_input" && content[i].ToolCallID == part.ID { - content[i].Arguments += delta - if streamDebugBytes != nil { - *streamDebugBytes += len(delta) - } - return content - } + // argBuilders indexes by ToolCallID directly, so resuming an + // interleaved call needs no scan through content. + b := argBuilders[part.ID] + if b == nil { + b = &strings.Builder{} + argBuilders[part.ID] = b + content = append(content, normalizedContentPart{ + Type: "tool_input", + ToolCallID: part.ID, + ToolName: part.ToolCallName, + }) } - - content = append(content, normalizedContentPart{ - Type: "tool_input", - ToolCallID: part.ID, - ToolName: part.ToolCallName, - Arguments: delta, - }) + _, _ = b.WriteString(delta) if streamDebugBytes != nil { *streamDebugBytes += len(delta) } @@ -879,16 +891,42 @@ func canonicalContentType(partType string) string { } } +// materializeStreamContent copies currentText (if any) and each +// argBuilders entry into the corresponding content part's +// Text/Arguments. It should be called once, after streaming completes +// and before content is read. It is safe to call more than once. +func materializeStreamContent( + content []normalizedContentPart, + currentText *strings.Builder, + currentTextIdx int, + argBuilders map[string]*strings.Builder, +) { + if currentText != nil && currentTextIdx < len(content) { + content[currentTextIdx].Text = currentText.String() + } + for i := range content { + if content[i].Type != "tool_input" { + continue + } + if b, ok := argBuilders[content[i].ToolCallID]; ok { + content[i].Arguments = b.String() + } + } +} + func appendNormalizedStreamContent( content []normalizedContentPart, + currentText *strings.Builder, + currentTextIdx int, + argBuilders map[string]*strings.Builder, part fantasy.StreamPart, streamDebugBytes *int, -) []normalizedContentPart { +) ([]normalizedContentPart, *strings.Builder, int) { switch part.Type { case fantasy.StreamPartTypeTextDelta: - return appendStreamContentText(content, "text", part.Delta, streamDebugBytes) + return appendStreamContentText(content, currentText, currentTextIdx, "text", part.Delta, streamDebugBytes) case fantasy.StreamPartTypeReasoningStart, fantasy.StreamPartTypeReasoningDelta: - return appendStreamContentText(content, "reasoning", part.Delta, streamDebugBytes) + return appendStreamContentText(content, currentText, currentTextIdx, "reasoning", part.Delta, streamDebugBytes) case fantasy.StreamPartTypeToolInputStart, fantasy.StreamPartTypeToolInputDelta, fantasy.StreamPartTypeToolInputEnd: @@ -896,31 +934,35 @@ func appendNormalizedStreamContent( // tool_call summary. Attribute each chunk to its tool call // so interrupted streams can reconstruct which partial input // belonged to which invocation. - return appendStreamToolInput(content, part, streamDebugBytes) + content = appendStreamToolInput(content, argBuilders, part, streamDebugBytes) + return content, currentText, currentTextIdx case fantasy.StreamPartTypeToolCall: - return append(content, normalizedContentPart{ + content = append(content, normalizedContentPart{ Type: canonicalContentType(string(part.Type)), ToolCallID: part.ID, ToolName: part.ToolCallName, Arguments: boundText(part.ToolCallInput), InputLength: utf8.RuneCountInString(part.ToolCallInput), }) + return content, currentText, currentTextIdx case fantasy.StreamPartTypeToolResult: - return append(content, normalizedContentPart{ + content = append(content, normalizedContentPart{ Type: canonicalContentType(string(part.Type)), ToolCallID: part.ID, ToolName: part.ToolCallName, Result: boundText(part.ToolCallInput), }) + return content, currentText, currentTextIdx case fantasy.StreamPartTypeSource: - return append(content, normalizedContentPart{ + content = append(content, normalizedContentPart{ Type: string(part.Type), SourceType: string(part.SourceType), Title: part.Title, URL: part.URL, }) + return content, currentText, currentTextIdx default: - return content + return content, currentText, currentTextIdx } } diff --git a/coderd/x/chatd/chatdebug/model_normalization_internal_test.go b/coderd/x/chatd/chatdebug/model_normalization_internal_test.go index 0f80806d36..29bc1b2295 100644 --- a/coderd/x/chatd/chatdebug/model_normalization_internal_test.go +++ b/coderd/x/chatd/chatdebug/model_normalization_internal_test.go @@ -132,16 +132,30 @@ func TestNormalizers_SkipTypedNilInterfaceValues(t *testing.T) { func TestAppendNormalizedStreamContent_PreservesOrderAndCanonicalTypes(t *testing.T) { t.Parallel() - var content []normalizedContentPart + var ( + content []normalizedContentPart + currentText *strings.Builder + currentTextIdx int + ) + argBuilders := make(map[string]*strings.Builder) streamDebugBytes := 0 for _, part := range []fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeTextDelta, Delta: "before "}, + // "before " and "after" each arrive as two consecutive + // same-type deltas (rather than one) so this also exercises + // accumulating into an already-existing text part across + // multiple appendNormalizedStreamContent calls, not just + // starting a fresh one. + {Type: fantasy.StreamPartTypeTextDelta, Delta: "bef"}, + {Type: fantasy.StreamPartTypeTextDelta, Delta: "ore "}, {Type: fantasy.StreamPartTypeToolCall, ID: "call-1", ToolCallName: "search_docs", ToolCallInput: `{"query":"debug"}`}, {Type: fantasy.StreamPartTypeToolResult, ID: "call-1", ToolCallName: "search_docs", ToolCallInput: `{"matches":1}`}, - {Type: fantasy.StreamPartTypeTextDelta, Delta: "after"}, + {Type: fantasy.StreamPartTypeTextDelta, Delta: "aft"}, + {Type: fantasy.StreamPartTypeTextDelta, Delta: "er"}, } { - content = appendNormalizedStreamContent(content, part, &streamDebugBytes) + content, currentText, currentTextIdx = appendNormalizedStreamContent( + content, currentText, currentTextIdx, argBuilders, part, &streamDebugBytes) } + materializeStreamContent(content, currentText, currentTextIdx, argBuilders) require.Equal(t, []normalizedContentPart{ {Type: "text", Text: "before "}, @@ -154,7 +168,12 @@ func TestAppendNormalizedStreamContent_PreservesOrderAndCanonicalTypes(t *testin func TestAppendNormalizedStreamContent_ToolInputAttributionPerCall(t *testing.T) { t.Parallel() - var content []normalizedContentPart + var ( + content []normalizedContentPart + currentText *strings.Builder + currentTextIdx int + ) + argBuilders := make(map[string]*strings.Builder) streamDebugBytes := 0 for _, part := range []fantasy.StreamPart{ {Type: fantasy.StreamPartTypeToolInputStart, ID: "call-a", ToolCallName: "search", Delta: `{"q`}, @@ -164,8 +183,10 @@ func TestAppendNormalizedStreamContent_ToolInputAttributionPerCall(t *testing.T) {Type: fantasy.StreamPartTypeToolInputDelta, ID: "call-a", ToolCallName: "search", Delta: `":"x"}`}, {Type: fantasy.StreamPartTypeToolInputEnd, ID: "call-b", ToolCallName: "calc", Delta: `":"add"}`}, } { - content = appendNormalizedStreamContent(content, part, &streamDebugBytes) + content, currentText, currentTextIdx = appendNormalizedStreamContent( + content, currentText, currentTextIdx, argBuilders, part, &streamDebugBytes) } + materializeStreamContent(content, currentText, currentTextIdx, argBuilders) require.Equal(t, []normalizedContentPart{ {Type: "tool_input", ToolCallID: "call-a", ToolName: "search", Arguments: `{"query":"x"}`}, @@ -176,7 +197,12 @@ func TestAppendNormalizedStreamContent_ToolInputAttributionPerCall(t *testing.T) func TestAppendNormalizedStreamContent_ToolInputAcrossInterleavedText(t *testing.T) { t.Parallel() - var content []normalizedContentPart + var ( + content []normalizedContentPart + currentText *strings.Builder + currentTextIdx int + ) + argBuilders := make(map[string]*strings.Builder) streamDebugBytes := 0 for _, part := range []fantasy.StreamPart{ {Type: fantasy.StreamPartTypeToolInputStart, ID: "call-a", ToolCallName: "search", Delta: `{"q`}, @@ -184,8 +210,10 @@ func TestAppendNormalizedStreamContent_ToolInputAcrossInterleavedText(t *testing {Type: fantasy.StreamPartTypeTextDelta, Delta: "thinking..."}, {Type: fantasy.StreamPartTypeToolInputDelta, ID: "call-a", ToolCallName: "search", Delta: `uery":"x"}`}, } { - content = appendNormalizedStreamContent(content, part, &streamDebugBytes) + content, currentText, currentTextIdx = appendNormalizedStreamContent( + content, currentText, currentTextIdx, argBuilders, part, &streamDebugBytes) } + materializeStreamContent(content, currentText, currentTextIdx, argBuilders) require.Equal(t, []normalizedContentPart{ {Type: "tool_input", ToolCallID: "call-a", ToolName: "search", Arguments: `{"query":"x"}`}, @@ -198,14 +226,21 @@ func TestAppendNormalizedStreamContent_GlobalTextCap(t *testing.T) { streamDebugBytes := 0 long := strings.Repeat("a", maxStreamDebugTextBytes) - var content []normalizedContentPart + var ( + content []normalizedContentPart + currentText *strings.Builder + currentTextIdx int + ) + argBuilders := make(map[string]*strings.Builder) for _, part := range []fantasy.StreamPart{ {Type: fantasy.StreamPartTypeTextDelta, Delta: long}, {Type: fantasy.StreamPartTypeToolCall, ID: "call-1", ToolCallName: "search_docs", ToolCallInput: `{}`}, {Type: fantasy.StreamPartTypeTextDelta, Delta: "tail"}, } { - content = appendNormalizedStreamContent(content, part, &streamDebugBytes) + content, currentText, currentTextIdx = appendNormalizedStreamContent( + content, currentText, currentTextIdx, argBuilders, part, &streamDebugBytes) } + materializeStreamContent(content, currentText, currentTextIdx, argBuilders) require.Len(t, content, 2) require.Equal(t, strings.Repeat("a", maxStreamDebugTextBytes), content[0].Text) diff --git a/coderd/x/chatd/chatdebug/redaction.go b/coderd/x/chatd/chatdebug/redaction.go index fc4677c710..b204c997cd 100644 --- a/coderd/x/chatd/chatdebug/redaction.go +++ b/coderd/x/chatd/chatdebug/redaction.go @@ -88,6 +88,44 @@ func RedactHeaders(h http.Header) map[string]string { return redacted } +var ( + errNotValidJSON = xerrors.New("chatdebug: body is not valid JSON") + errExtraJSONData = xerrors.New("chatdebug: body contains extra JSON values") +) + +func decodeCompleteJSON(data []byte) (any, error) { + decoder := json.NewDecoder(bytes.NewReader(data)) + // Preserve precision: callers may re-marshal the result. + decoder.UseNumber() + + var value any + if err := decoder.Decode(&value); err != nil { + return nil, errNotValidJSON + } + var extra any + if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) { + return nil, errExtraJSONData + } + return value, nil +} + +// redactDecodedJSON redacts an already-decoded JSON value. ok is +// false when nothing needed redacting or re-marshaling the redacted +// value failed; either way the caller should fall back to the +// original bytes. +func redactDecodedJSON(value any) (encoded []byte, ok bool) { + redacted, changed := redactJSONValue(value) + if !changed { + return nil, false + } + + data, err := json.Marshal(redacted) + if err != nil { + return nil, false + } + return data, true +} + // RedactJSONSecrets redacts sensitive JSON values by key name. When // the input is not valid JSON (truncated body, HTML error page, etc.) // the raw bytes are replaced entirely with a diagnostic placeholder @@ -97,29 +135,20 @@ func RedactJSONSecrets(data []byte) []byte { return data } - decoder := json.NewDecoder(bytes.NewReader(data)) - decoder.UseNumber() - - var value any - if err := decoder.Decode(&value); err != nil { + value, err := decodeCompleteJSON(data) + if err != nil { + if errors.Is(err, errExtraJSONData) { + return []byte(`{"error":"chatdebug: body contains extra JSON values, redacted for safety"}`) + } // Cannot parse: replace entirely to prevent credential leaks // from non-JSON error responses (HTML pages, partial bodies). return []byte(`{"error":"chatdebug: body is not valid JSON, redacted for safety"}`) } - if err := consumeJSONEOF(decoder); err != nil { - return []byte(`{"error":"chatdebug: body contains extra JSON values, redacted for safety"}`) - } - redacted, changed := redactJSONValue(value) - if !changed { - return data + if encoded, changed := redactDecodedJSON(value); changed { + return encoded } - - encoded, err := json.Marshal(redacted) - if err != nil { - return data - } - return encoded + return data } // RedactNDJSONSecrets redacts sensitive values in newline-delimited @@ -150,18 +179,6 @@ func RedactNDJSONSecrets(data []byte) []byte { return bytes.Join(lines, []byte("\n")) } -func consumeJSONEOF(decoder *json.Decoder) error { - var extra any - err := decoder.Decode(&extra) - if errors.Is(err, io.EOF) { - return nil - } - if err == nil { - return xerrors.New("chatdebug: extra JSON values") - } - return err -} - // safeRateLimitHeaderNames lists rate-limit headers that contain // "token" in the name but carry numeric usage counters, not // credentials. They are checked in isSensitiveName before the diff --git a/coderd/x/chatd/chatdebug/transport.go b/coderd/x/chatd/chatdebug/transport.go index 07cdb92568..d892d0afdd 100644 --- a/coderd/x/chatd/chatdebug/transport.go +++ b/coderd/x/chatd/chatdebug/transport.go @@ -2,7 +2,6 @@ package chatdebug import ( "bytes" - "encoding/json" "errors" "io" "mime" @@ -177,6 +176,7 @@ func captureRequestBody(req *http.Request) ([]byte, error) { if len(limited) > maxRecordedRequestBodyBytes { return []byte("[TRUNCATED]"), nil } + // Request bodies have no completeness pre-check to reuse; always decode fresh. return RedactJSONSecrets(limited), nil } } @@ -316,12 +316,12 @@ func (r *recordingBody) Close() error { // The SSE EOF path already appended a completed attempt. // inner.Close() surfaced a transport error, so upgrade // that entry to failed instead of losing the close error. - upgraded := r.buildAttemptLocked(closeErr) + upgraded := r.buildAttemptLocked(closeErr, nil) r.sink.replaceByNumber(upgraded.Number, upgraded) r.recordedProvisional = false } else { r.recordOnce.Do(func() { - r.sink.record(r.buildAttemptLocked(closeErr)) + r.sink.record(r.buildAttemptLocked(closeErr, nil)) }) } r.mu.Unlock() @@ -344,20 +344,35 @@ func (r *recordingBody) Close() error { responseBody = append([]byte(nil), r.buf.Bytes()...) r.mu.Unlock() + // decodedJSON is nil unless the completeness check below decoded + // these bytes; non-nil even for a stored nil value (JSON "null"). + // A truncated buffer is an incomplete prefix, so skip the check. + var decodedJSON *any + if contentLength < 0 && !truncated { + if value, ok := decodeUnknownLengthJSONBody(contentType, responseBody); ok { + decodedJSON = &value + } + } + + recordReusingDecodedJSON := func(err error) { + if decodedJSON != nil { + r.recordJSON(err, *decodedJSON) + return + } + r.record(err) + } + switch { - // Only check JSON completeness when the recording buffer is - // not truncated. A truncated buffer is an incomplete prefix - // of the body, so the completeness check would false-positive. - case sawEOF && !truncated && contentLength < 0 && isJSONLikeContentType(contentType) && !isCompleteUnknownLengthJSONBody(contentType, responseBody): + case sawEOF && !truncated && contentLength < 0 && isJSONLikeContentType(contentType) && decodedJSON == nil: r.record(io.ErrUnexpectedEOF) case sawEOF: - r.record(io.EOF) + recordReusingDecodedJSON(io.EOF) case responseHasNoBody(r.base.Method, r.base.ResponseStatus): r.record(nil) case contentLength >= 0 && bytesRead >= contentLength: r.record(nil) - case contentLength < 0 && !truncated && isCompleteUnknownLengthJSONBody(contentType, responseBody): - r.record(nil) + case contentLength < 0 && !truncated && decodedJSON != nil: + recordReusingDecodedJSON(nil) // Truncated unknown-length bodies: the caller consumed the // response successfully but the recording buffer exceeded // maxRecordedResponseBodyBytes. This is not a transport @@ -437,30 +452,39 @@ func (r *recordingBody) drainToEOF() error { } func isCompleteUnknownLengthJSONBody(contentType string, body []byte) bool { + _, ok := decodeUnknownLengthJSONBody(contentType, body) + return ok +} + +// decodeUnknownLengthJSONBody reports whether body is a single, +// complete JSON document for the given content type. It returns the +// decoded value too, to avoid re-decoding during redaction. +func decodeUnknownLengthJSONBody(contentType string, body []byte) (any, bool) { if !isJSONLikeContentType(contentType) { - return false + return nil, false } trimmed := bytes.TrimSpace(body) if len(trimmed) == 0 { - return false + return nil, false } - decoder := json.NewDecoder(bytes.NewReader(trimmed)) - var value any - if err := decoder.Decode(&value); err != nil { - return false + value, err := decodeCompleteJSON(trimmed) + if err != nil { + return nil, false } - var extra any - return errors.Is(decoder.Decode(&extra), io.EOF) + return value, true } // buildAttemptLocked materializes the final Attempt from the current -// buffered response data plus err. Callers use this from both the -// record-once append path and the provisional-upgrade replace path so -// both sites apply the same redaction and status rules. The caller -// must hold r.mu for the duration of the call. -func (r *recordingBody) buildAttemptLocked(err error) Attempt { +// buffered response data plus err. The caller must hold r.mu for the +// duration of the call. +// +// decodedJSON is nil when nothing was decoded ahead of time. A +// non-nil decodedJSON - even one pointing at a stored nil interface, +// e.g. a JSON "null" body - is redacted directly instead of decoding +// the raw bytes again. +func (r *recordingBody) buildAttemptLocked(err error, decodedJSON *any) Attempt { finishedAt := time.Now() truncated := r.truncated @@ -472,6 +496,12 @@ func (r *recordingBody) buildAttemptLocked(err error) Attempt { switch { case truncated: base.ResponseBody = []byte("[TRUNCATED]") + case decodedJSON != nil: + if encoded, changed := redactDecodedJSON(*decodedJSON); changed { + base.ResponseBody = encoded + } else { + base.ResponseBody = responseBody + } case isNDJSONContentType(contentType): base.ResponseBody = RedactNDJSONSecrets(responseBody) case contentType == "" || isJSONLikeContentType(contentType): @@ -508,7 +538,17 @@ func (r *recordingBody) record(err error) { r.mu.Lock() defer r.mu.Unlock() r.recordOnce.Do(func() { - r.sink.record(r.buildAttemptLocked(err)) + r.sink.record(r.buildAttemptLocked(err, nil)) + }) +} + +// recordJSON behaves like record but reuses an already-decoded JSON +// value, avoiding a second full decode of the same bytes. +func (r *recordingBody) recordJSON(err error, value any) { + r.mu.Lock() + defer r.mu.Unlock() + r.recordOnce.Do(func() { + r.sink.record(r.buildAttemptLocked(err, &value)) }) } @@ -523,7 +563,7 @@ func (r *recordingBody) recordProvisional(err error) { r.mu.Lock() defer r.mu.Unlock() r.recordOnce.Do(func() { - r.sink.record(r.buildAttemptLocked(err)) + r.sink.record(r.buildAttemptLocked(err, nil)) r.recordedProvisional = true }) } diff --git a/coderd/x/chatd/chatdebug/transport_internal_test.go b/coderd/x/chatd/chatdebug/transport_internal_test.go index abe2ff616c..86217d7f7c 100644 --- a/coderd/x/chatd/chatdebug/transport_internal_test.go +++ b/coderd/x/chatd/chatdebug/transport_internal_test.go @@ -615,6 +615,46 @@ func TestRecordingTransport_CloseAfterDecoderConsumesUnknownLengthJSONSucceeds(t require.Len(t, attempts, 1) require.Equal(t, attemptStatusCompleted, attempts[0].Status) require.Empty(t, attempts[0].Error) + require.JSONEq(t, `{"token":"[REDACTED]","safe":"ok"}`, string(attempts[0].ResponseBody)) +} + +// TestRecordingTransport_CloseAfterDecoderConsumesUnknownLengthJSONWithNoSensitiveKeysSucceeds +// exercises buildAttemptLocked's decodedJSON!=nil-but-changed==false path: +// the completeness check precomputes a decoded value, but nothing in it +// needs redacting, so the original bytes must be recorded unchanged. +func TestRecordingTransport_CloseAfterDecoderConsumesUnknownLengthJSONWithNoSensitiveKeysSucceeds(t *testing.T) { + t.Parallel() + + ctx, sink := newTestSinkContext(t) + client := &http.Client{ + Transport: &RecordingTransport{ + Base: roundTripFunc(func(req *http.Request) (*http.Response, error) { + return &http.Response{ //nolint:exhaustruct // Test response exercises unknown-length close semantics. + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: &scriptedReadCloser{chunks: [][]byte{[]byte(`{"safe":"ok","count":1}`)}}, + ContentLength: -1, + }, nil + }), + }, + } + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example.invalid", nil) + require.NoError(t, err) + + resp, err := client.Do(req) + require.NoError(t, err) + + var decoded map[string]any + require.NoError(t, json.NewDecoder(resp.Body).Decode(&decoded)) + require.Equal(t, "ok", decoded["safe"]) + require.NoError(t, resp.Body.Close()) + + attempts := sink.snapshot() + require.Len(t, attempts, 1) + require.Equal(t, attemptStatusCompleted, attempts[0].Status) + require.Empty(t, attempts[0].Error) + require.JSONEq(t, `{"safe":"ok","count":1}`, string(attempts[0].ResponseBody)) } func TestRecordingTransport_CloseAfterDecoderConsumesUnknownLengthJSONWithTrailingDocumentMarksFailed(t *testing.T) { @@ -649,6 +689,9 @@ func TestRecordingTransport_CloseAfterDecoderConsumesUnknownLengthJSONWithTraili require.Len(t, attempts, 1) require.Equal(t, attemptStatusFailed, attempts[0].Status) require.Equal(t, io.ErrUnexpectedEOF.Error(), attempts[0].Error) + require.JSONEq(t, + `{"error":"chatdebug: body contains extra JSON values, redacted for safety"}`, + string(attempts[0].ResponseBody)) } func TestRecordingTransport_CloseAfterDecoderConsumesUnknownLengthNDJSONMarksFailed(t *testing.T) {