Merge pull request #15 from BackendStack21/fix/tool-result-envelope-display

fix(tui): decode JSON tool envelopes and fold untrusted wrappers for display

Author: jkyberneees

internal/tui/events.go

@@ -3,6 +3,9 @@ package tui
 import (
 	"encoding/json"
 	"fmt"
+	"regexp"
+	"sort"
+	"strconv"
 	"strings"
 	"time"
 	"unicode"
@@ -225,6 +228,93 @@ func eventTail(ev client.Event) string {
 	return s
 }
 
+// normalizeToolResult prepares raw tool output for DISPLAY ONLY — the stored
+// and forwarded data is never touched. It handles two render-noise shapes:
+//  1. JSON envelopes (read_file, search_files, …): Go's encoding/json
+//     HTML-escapes <, >, & in string values, so the odek prompt-injection
+//     wrapper shows up as <untrusted_content_… — the envelope is
+//     decoded and the real content rendered, with remaining scalar metadata
+//     as a one-line footer.
+//  2. untrusted_content wrappers, folded into a "⚠ untrusted: <source>"
+//     badge line. Both the literal tag form (live stream events) and the
+//     < escaped form (undecoded JSON envelopes) are recognized.
+//
+// Fail-safe: anything that does not match a known shape exactly is returned
+// unchanged — normalization is never lossy.
+func normalizeToolResult(data string) string {
+	return foldUntrustedWrappers(decodeToolEnvelope(data))
+}
+
+// decodeToolEnvelope unwraps a JSON object with a string "content" field,
+// returning the decoded content plus a footer line of the remaining scalar
+// metadata (e.g. "total_lines: 430"). Nested objects/arrays, missing
+// content, or invalid JSON mean an unknown shape — returned unchanged.
+func decodeToolEnvelope(data string) string {
+	t := strings.TrimSpace(data)
+	if len(t) < 2 || t[0] != '{' {
+		return data
+	}
+	var env map[string]any
+	if err := json.Unmarshal([]byte(t), &env); err != nil {
+		return data
+	}
+	content, ok := env["content"].(string)
+	if !ok {
+		return data
+	}
+	meta := make(map[string]string, len(env)-1)
+	for k, v := range env {
+		if k == "content" || v == nil {
+			continue
+		}
+		switch x := v.(type) {
+		case string:
+			meta[k] = x
+		case float64:
+			meta[k] = strconv.FormatFloat(x, 'f', -1, 64)
+		case bool:
+			meta[k] = strconv.FormatBool(x)
+		default:
+			return data // nested value: not a known envelope shape
+		}
+	}
+	if len(meta) == 0 {
+		return content
+	}
+	keys := make([]string, 0, len(meta))
+	for k := range meta {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+	parts := make([]string, 0, len(keys))
+	for _, k := range keys {
+		parts = append(parts, k+": "+meta[k])
+	}
+	return content + "\n(" + strings.Join(parts, " · ") + ")"
+}
+
+// untrustedWrapperRe matches an odek prompt-injection wrapper in either its
+// literal form or the < escaped form produced by encoding/json.
+// Capture groups: 1 = source attribute, 2 = wrapped body.
+var untrustedWrapperRe = regexp.MustCompile(
+	`(?s)(?:<|\\u003c)untrusted_content_[0-9a-f]+ source=\\?"([^"]*)\\?"(?:>|\\u003e)\n?(.*?)\n?(?:<|\\u003c)/untrusted_content_[0-9a-f]+(?:>|\\u003e)`)
+
+// foldUntrustedWrappers replaces each untrusted_content wrapper with a badge
+// line naming the source, followed by the wrapped body.
+func foldUntrustedWrappers(s string) string {
+	if !strings.Contains(s, "untrusted_content_") {
+		return s
+	}
+	return untrustedWrapperRe.ReplaceAllStringFunc(s, func(m string) string {
+		sub := untrustedWrapperRe.FindStringSubmatch(m)
+		badge := "⚠ untrusted: " + sub[1]
+		if body := sub[2]; body != "" {
+			return badge + "\n" + body
+		}
+		return badge
+	})
+}
+
 // cur returns the index of the active streaming assistant message, or -1.
 func (m *Model) cur() int {
 	if m.curIdx >= 0 && m.curIdx < len(m.msgs) {
@@ -331,7 +421,7 @@ func argPreview(data string) string {
 // lines, so the transcript can show a useful excerpt (rendered by renderSteps)
 // without retaining the unbounded output of a chatty tool.
 func resultPreview(data string) string {
-	s := sanitize(data)
+	s := sanitize(normalizeToolResult(data))
 	lines := strings.Split(s, "\n")
 	const cap = 200
 	if len(lines) > cap {

internal/tui/normalize_test.go

@@ -0,0 +1,66 @@
+package tui
+
+import (
+	"strings"
+	"testing"
+)
+
+// read_file-style JSON envelope: Go's encoding/json HTML-escapes <, >, & in
+// the wrapped content, so the raw result exposes < untrusted_content_…
+// noise. resultPreview must decode the envelope and fold the wrapper.
+func TestResultPreviewJSONEnvelope(t *testing.T) {
+	raw := `{"content":"<untrusted_content_dd96e237d3be7447 source=\"/path/to/file.go\">\n217|\n218|\tta := textarea.New()\n</untrusted_content_dd96e237d3be7447>","total_lines":300}`
+	got := resultPreview(raw)
+	if strings.Contains(got, `<`) || strings.Contains(got, "untrusted_content_") {
+		t.Errorf("escaped wrapper still visible:\n%s", got)
+	}
+	if !strings.Contains(got, "⚠ untrusted: /path/to/file.go") {
+		t.Errorf("missing source badge:\n%s", got)
+	}
+	if !strings.Contains(got, "218|\tta := textarea.New()") {
+		t.Errorf("body line numbers not intact:\n%s", got)
+	}
+	if !strings.Contains(got, "total_lines: 300") {
+		t.Errorf("missing scalar metadata footer:\n%s", got)
+	}
+}
+
+// Live (non-JSON) tool results carry the wrapper literally; it folds into a
+// badge line plus the body.
+func TestResultPreviewLiteralWrapper(t *testing.T) {
+	raw := "<untrusted_content_ab12cd34 source=\"/etc/hosts\">\n127.0.0.1 localhost\n</untrusted_content_ab12cd34>"
+	got := resultPreview(raw)
+	if strings.Contains(got, "untrusted_content_") {
+		t.Errorf("wrapper tags still visible:\n%s", got)
+	}
+	if !strings.Contains(got, "⚠ untrusted: /etc/hosts") {
+		t.Errorf("missing source badge:\n%s", got)
+	}
+	if !strings.Contains(got, "127.0.0.1 localhost") {
+		t.Errorf("body lost:\n%s", got)
+	}
+}
+
+// Plain, non-JSON tool output renders byte-identical.
+func TestResultPreviewPlainUnchanged(t *testing.T) {
+	for _, s := range []string{
+		"",
+		"plain output",
+		"exit status 1\nFAIL TestX",
+		"line with <angle> & \"quotes\" kept as-is",
+		`{"matches":[{"path":"a.go"}]}`,    // JSON without a string content field
+		`{"content":"x","nested":{"a":1}}`, // non-scalar metadata: not a known envelope
+	} {
+		if got := resultPreview(s); got != s {
+			t.Errorf("resultPreview(%q) = %q, want unchanged", s, got)
+		}
+	}
+}
+
+// A malformed envelope (truncated JSON) renders raw, never lossy.
+func TestResultPreviewMalformedEnvelope(t *testing.T) {
+	raw := `{"content":"<untrusted_content_ab`
+	if got := resultPreview(raw); got != raw {
+		t.Errorf("malformed envelope rewritten: got %q", got)
+	}
+}