perf(mcp): reuse decoded response envelope views

internal/mcp/jsonrpc/jsonrpc.go

@@ -408,14 +408,22 @@ func ExtractVisibleStringsFromJSONResultWithMediaBudget(raw json.RawMessage, bud
 }
 
 func extractVisibleStrings(raw json.RawMessage, budget *MediaTextBudget) ExtractStringsResult {
-	if budget == nil {
-		budget = &MediaTextBudget{}
-	}
 	var parsed interface{}
 	if err := extract.DecodeJSON(raw, &parsed); err != nil {
 		return ExtractStringsResult{}
 	}
+	return ExtractVisibleStringsFromValueResult(parsed, budget)
+}
 
+// ExtractVisibleStringsFromValueResult walks a value already decoded by
+// extract.DecodeJSON, optionally drawing on a shared response media budget.
+// The value must contain only the decoder's JSON types, including json.Number,
+// with any removed fields replaced by nil. It must not be mutated during the
+// walk. This skips parsing, not the extraction depth or media limits.
+func ExtractVisibleStringsFromValueResult(parsed any, budget *MediaTextBudget) ExtractStringsResult {
+	if budget == nil {
+		budget = &MediaTextBudget{}
+	}
 	var result []string
 	var atoms []extract.DLPAtom
 	var positions []int
@@ -933,6 +941,20 @@ func ExtractStringsFromJSONResult(raw json.RawMessage) ExtractStringsResult {
 // ignores keys because they are normally structural. Callers that surface a
 // JSON object as agent-visible tool metadata can opt into scanning its keys.
 func ExtractKeysFromJSONResult(raw json.RawMessage) ExtractKeysResult {
+	var parsed interface{}
+	dec := json.NewDecoder(bytes.NewReader(raw))
+	dec.UseNumber()
+	if err := dec.Decode(&parsed); err != nil || !json.Valid(raw) {
+		return ExtractKeysResult{}
+	}
+	return ExtractKeysFromValueResult(parsed)
+}
+
+// ExtractKeysFromValueResult walks a value already decoded by extract.DecodeJSON
+// (or the equivalent complete-document, UseNumber decoder above). The value
+// must contain only JSON decoder types, with any removed fields replaced by nil,
+// and must not be mutated during the walk. Depth and key limits still apply.
+func ExtractKeysFromValueResult(parsed any) ExtractKeysResult {
 	var result []string
 	truncated := false
 	var extract func(v interface{}, depth int)
@@ -975,11 +997,6 @@ func ExtractKeysFromJSONResult(raw json.RawMessage) ExtractKeysResult {
 			}
 		}
 	}
-	var parsed interface{}
-	dec := json.NewDecoder(bytes.NewReader(raw))
-	dec.UseNumber()
-	if err := dec.Decode(&parsed); err == nil && json.Valid(raw) {
-		extract(parsed, 0)
-	}
+	extract(parsed, 0)
 	return ExtractKeysResult{Keys: result, Truncated: truncated}
 }

internal/mcp/jsonrpc/value_views_test.go

@@ -0,0 +1,196 @@
+// Copyright 2026 Josh Waldrep
+// SPDX-License-Identifier: Apache-2.0
+
+package jsonrpc
+
+import (
+	"encoding/base64"
+	"encoding/binary"
+	"encoding/json"
+	"fmt"
+	"reflect"
+	"strings"
+	"testing"
+
+	"github.com/luckyPipewrench/pipelock/internal/extract"
+)
+
+// A decoded tree and its JSON round trip must expose identical values, card
+// barriers, key order and media budget use. This is the round trip the response
+// remainder used before consuming the decoded tree directly.
+func TestExtractValueViewsRoundTrip(t *testing.T) {
+	for _, raw := range []string{
+		`null`, `""`, `[]`, `{}`, `[null,true,false,"",-0,1e10000,123456789012345678901]`,
+		`{"z":["last",{"b":null,"a":"first"}],"a":"<>&\u2028\u2029","\u0062":"escaped"}`,
+		`{"data":{"note":"visible"},"blob":["plain text",{"data":"more text"}],"raw":null}`,
+		`{"unicode":"\ud800 \udc00 \ud83d\ude00","nul":"a\u0000b"}`,
+	} {
+		t.Run(raw, func(t *testing.T) {
+			var value any
+			if err := extract.DecodeJSON([]byte(raw), &value); err != nil {
+				t.Fatal(err)
+			}
+			encoded, err := json.Marshal(value)
+			if err != nil {
+				t.Fatal(err)
+			}
+			var directBudget, encodedBudget MediaTextBudget
+			got := ExtractVisibleStringsFromValueResult(value, &directBudget)
+			want := ExtractVisibleStringsFromJSONResultWithMediaBudget(encoded, &encodedBudget)
+			if !reflect.DeepEqual(got, want) || directBudget != encodedBudget {
+				t.Fatalf("value = %+v, budget = %+v; round trip = %+v, budget = %+v", got, directBudget, want, encodedBudget)
+			}
+			if got, want := ExtractKeysFromValueResult(value), ExtractKeysFromJSONResult(encoded); !reflect.DeepEqual(got, want) {
+				t.Fatalf("keys = %+v, round trip = %+v", got, want)
+			}
+		})
+	}
+}
+
+func TestExtractValueViewsScalars(t *testing.T) {
+	var value any
+	if err := extract.DecodeJSON([]byte(`{"z":"last","a":["",null,true,false,-0,1e10000,123456789012345678901]}`), &value); err != nil {
+		t.Fatal(err)
+	}
+	got := ExtractVisibleStringsFromValueResult(value, nil)
+	if !reflect.DeepEqual(got.Strings, []string{"", "last"}) || !reflect.DeepEqual(got.StringPositions, []int{0, 7}) {
+		t.Fatalf("strings or positions changed: %+v", got)
+	}
+	wantText := []string{"", Null, "true", "false", "-0", "1e10000", "123456789012345678901", "last"}
+	if len(got.CardAtoms) != len(wantText) {
+		t.Fatalf("atoms = %+v", got.CardAtoms)
+	}
+	for i, atom := range got.CardAtoms {
+		if atom.Text != wantText[i] || atom.Number != (i >= 4 && i <= 6) || atom.Key {
+			t.Errorf("atom %d = %+v, want text %q", i, atom, wantText[i])
+		}
+	}
+	if keys := ExtractKeysFromValueResult(value); keys.Truncated || !reflect.DeepEqual(keys.Keys, []string{"a", "z"}) {
+		t.Fatalf("sorted keys = %+v", keys)
+	}
+	for _, raw := range []string{`[]`, `{}`} {
+		if err := extract.DecodeJSON([]byte(raw), &value); err != nil {
+			t.Fatal(err)
+		}
+		if got := ExtractVisibleStringsFromValueResult(value, nil); !reflect.DeepEqual(got, ExtractStringsResult{}) {
+			t.Fatalf("empty value = %+v", got)
+		}
+		if got := ExtractKeysFromValueResult(value); !reflect.DeepEqual(got, ExtractKeysResult{}) {
+			t.Fatalf("empty keys = %+v", got)
+		}
+	}
+}
+
+func TestExtractValueViewsBounds(t *testing.T) {
+	for _, depth := range []int{maxExtractDepth, maxExtractDepth + 1} {
+		var value any
+		raw := strings.Repeat("[", depth) + `"leaf"` + strings.Repeat("]", depth)
+		if err := extract.DecodeJSON([]byte(raw), &value); err != nil {
+			t.Fatal(err)
+		}
+		stringsResult := ExtractVisibleStringsFromValueResult(value, nil)
+		keys := ExtractKeysFromValueResult(value)
+		wantTruncated := depth > maxExtractDepth
+		if stringsResult.Truncated != wantTruncated || keys.Truncated != wantTruncated {
+			t.Fatalf("depth %d: strings = %+v, keys = %+v", depth, stringsResult, keys)
+		}
+		if wantTruncated && len(stringsResult.Strings) != 0 {
+			t.Fatal("depth-limited extraction exposed an uninspected leaf")
+		}
+		if !wantTruncated && !reflect.DeepEqual(stringsResult.Strings, []string{"leaf"}) {
+			t.Fatalf("boundary leaf missing: %+v", stringsResult)
+		}
+	}
+	for _, width := range []int{maxExtractKeys - 1, maxExtractKeys, maxExtractKeys + 1} {
+		object := make(map[string]any, width)
+		for i := range width {
+			object[fmt.Sprintf("k%05d", i)] = nil
+		}
+		// Construct the same types as DecodeJSON without serializing a wide map.
+		got := ExtractKeysFromValueResult(object)
+		if got.Truncated != (width >= maxExtractKeys) {
+			t.Fatalf("width %d: truncated = %v", width, got.Truncated)
+		}
+		wantLength := width
+		if width > maxExtractKeys {
+			wantLength = 0 // Refuse before allocating/sorting an oversized map.
+		}
+		if len(got.Keys) != wantLength {
+			t.Fatalf("width %d: keys = %d, want %d", width, len(got.Keys), wantLength)
+		}
+	}
+	// A nested object can use the remaining allowance before a parent visits
+	// its next array element or object member. Both loops must stop there.
+	nested := make(map[string]any, maxExtractKeys-1)
+	for i := range maxExtractKeys - 1 {
+		nested[fmt.Sprintf("k%05d", i)] = nil
+	}
+	for _, value := range []any{
+		[]any{map[string]any{"a": nested}, "unvisited"},
+		map[string]any{"a": nested, "z": "unvisited"},
+	} {
+		got := ExtractKeysFromValueResult(value)
+		if !got.Truncated || len(got.Keys) != maxExtractKeys || got.Keys[len(got.Keys)-1] == "z" {
+			t.Fatalf("nested object exceeded the shared key allowance: %+v", got)
+		}
+	}
+}
+
+func TestExtractValueViewsMedia(t *testing.T) {
+	for name, payload := range map[string][]byte{
+		"plain":    []byte(mediaTestPhrase),
+		"utf16":    utf16Encode(mediaTestPhrase, binary.LittleEndian),
+		"controls": interleaved("decoded "+mediaTestPhrase, 0),
+	} {
+		t.Run(name, func(t *testing.T) {
+			encoded := base64.StdEncoding.EncodeToString(append([]byte("GIF89a"), payload...))
+			var value any
+			raw, err := json.Marshal(map[string]any{"data": []any{encoded}, "raw": map[string]any{"note": "visible"}})
+			if err != nil {
+				t.Fatal(err)
+			}
+			if err := extract.DecodeJSON(raw, &value); err != nil {
+				t.Fatal(err)
+			}
+			var directBudget, encodedBudget MediaTextBudget
+			got := ExtractVisibleStringsFromValueResult(value, &directBudget)
+			want := ExtractVisibleStringsFromJSONResultWithMediaBudget(raw, &encodedBudget)
+			if !reflect.DeepEqual(got, want) || directBudget != encodedBudget || directBudget.used == 0 {
+				t.Fatalf("media parity: got %+v/%+v, want %+v/%+v", got, directBudget, want, encodedBudget)
+			}
+			if !strings.Contains(strings.Join(got.Strings, " "), mediaTestPhrase) || got.IncompleteReason != "" {
+				t.Fatalf("decoded text missing: %+v", got)
+			}
+			if !strings.Contains(strings.Join(got.Strings, " "), "visible") {
+				t.Fatal("a media key hid a nested object")
+			}
+			budget := MediaTextBudget{used: MaxMediaTextBytes - 1}
+			if got := ExtractVisibleStringsFromValueResult(value, &budget); got.IncompleteReason != mediaTextBudgetExceeded || got.Strings != nil || got.CardAtoms != nil || got.StringPositions != nil {
+				t.Fatalf("budget exhaustion returned partial output: %+v", got)
+			}
+		})
+	}
+	encoded := base64.StdEncoding.EncodeToString(append([]byte("GIF89a"), make([]byte, 64)...))
+	got := ExtractVisibleStringsFromValueResult(map[string]any{"blob": encoded}, nil)
+	if got.Strings != nil || got.StringPositions != nil || len(got.CardAtoms) != 1 || got.CardAtoms[0].Card != Null {
+		t.Fatalf("nontext media became text or lost its barrier: %+v", got)
+	}
+	for _, field := range []string{"plain instruction text", "data:image/gif;base64,broken!", "%%%not-base64%%%"} {
+		got := ExtractVisibleStringsFromValueResult(map[string]any{"data": field}, nil)
+		if !reflect.DeepEqual(got.Strings, []string{field}) || got.IncompleteReason != "" {
+			t.Fatalf("unrecognized media was hidden: %+v", got)
+		}
+	}
+}
+
+func TestExtractValueViewsInvalidRaw(t *testing.T) {
+	for _, raw := range []string{"", "{", "{} {}", "[1,]"} {
+		budget := MediaTextBudget{used: MaxMediaTextBytes, reason: mediaTextBudgetExceeded}
+		if got := ExtractVisibleStringsFromJSONResultWithMediaBudget([]byte(raw), &budget); !reflect.DeepEqual(got, ExtractStringsResult{}) {
+			t.Fatalf("invalid raw input %q returned %+v", raw, got)
+		}
+		if got := ExtractKeysFromJSONResult([]byte(raw)); !reflect.DeepEqual(got, ExtractKeysResult{}) {
+			t.Fatalf("invalid raw input %q returned keys %+v", raw, got)
+		}
+	}
+}

internal/mcp/scan_remainder.go

@@ -27,15 +27,15 @@ type responseEnvelopeViews struct {
 // allowance here. Callers must check budget.Reason() before using the text: an
 // exhausted budget leaves it partial.
 func responseEnvelopeRemainder(raw []byte, rpc jsonrpc.RPCResponse, budget *jsonrpc.MediaTextBudget) (responseEnvelopeViews, bool) {
-	// Check the complete message before removing any consumed subtrees. In
-	// particular, pruning structuredContent must not relax the depth/key bound.
-	if jsonrpc.ExtractKeysFromJSONResult(raw).Truncated {
-		return responseEnvelopeViews{}, false
-	}
 	var envelope map[string]any
 	if extract.DecodeJSON(raw, &envelope) != nil {
 		return responseEnvelopeViews{}, false
 	}
+	// Check the complete decoded message before removing any consumed subtrees.
+	// In particular, pruning structuredContent must not relax the depth/key bound.
+	if jsonrpc.ExtractKeysFromValueResult(envelope).Truncated {
+		return responseEnvelopeViews{}, false
+	}
 	var uris []string
 	for _, key := range jsonrpc.SortedKeys(envelope) {
 		value := envelope[key]
@@ -76,12 +76,8 @@ func responseEnvelopeRemainder(raw []byte, rpc jsonrpc.RPCResponse, budget *json
 			}
 		}
 	}
-	encoded, err := json.Marshal(envelope)
-	if err != nil {
-		return responseEnvelopeViews{}, false
-	}
-	values := jsonrpc.ExtractVisibleStringsFromJSONResultWithMediaBudget(encoded, budget)
-	keys := jsonrpc.ExtractKeysFromJSONResult(encoded)
+	values := jsonrpc.ExtractVisibleStringsFromValueResult(envelope, budget)
+	keys := jsonrpc.ExtractKeysFromValueResult(envelope)
 	if values.Truncated || keys.Truncated {
 		return responseEnvelopeViews{}, false
 	}