From 0835df8da99764c1212ab8f77b0a564b216a07aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Wed, 29 Jul 2026 19:22:36 +0200 Subject: [PATCH 1/3] feat(#3996): capture safe request-shape diagnostics for image requests Diagnose the observed google/gemini-2.5-flash-image HTTP 400 without dumping raw requests/prompts/schemas. RequestShape summarizes a built GenerateContentConfig as booleans/counts/fixed enums only (modalities, tool/toolconfig presence, structured-output presence, thinking config, API surface), safe to log at Debug and to hand to bug reports. Track apiSurface (gemini_api/vertex_ai/gateway) on Client so the shape can distinguish backends. Replace the old per-tool debug log in CreateChatCompletionStream, which logged full Tool structs including function names/descriptions/parameter schemas, with the bounded shape. Trace the Split Diff local-only classification from the actual request-building boundary rather than just this package: the regression scans this package (where the request is serialized) plus pkg/tools (where the []tools.Tool list handed to every provider is built) and pkg/runtime (where messages/tools are assembled before any provider call) for any reference to the TUI-local Split Diff toggle, and finds none. MCP and other dynamically registered tools are outside the reach of a static scan; that boundary is documented on the test. --- pkg/model/provider/gemini/client.go | 29 +- pkg/model/provider/gemini/diagnostics.go | 155 ++++++++ pkg/model/provider/gemini/diagnostics_test.go | 333 ++++++++++++++++++ 3 files changed, 508 insertions(+), 9 deletions(-) create mode 100644 pkg/model/provider/gemini/diagnostics.go create mode 100644 pkg/model/provider/gemini/diagnostics_test.go diff --git a/pkg/model/provider/gemini/client.go b/pkg/model/provider/gemini/client.go index d51fb43d9..1a50f7800 100644 --- a/pkg/model/provider/gemini/client.go +++ b/pkg/model/provider/gemini/client.go @@ -34,6 +34,12 @@ type Client struct { base.Config clientFn func(context.Context) (*genai.Client, error) + + // apiSurface classifies which backend/transport this client talks to + // (see the apiSurface* constants in diagnostics.go), for safe + // request-shape diagnostics. Never exposed to the model or logged + // alongside anything provider-supplied. + apiSurface string } // NewClient creates a new Gemini client from the provided configuration @@ -49,6 +55,7 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro globalOptions := options.Apply(opts...) var clientFn func(context.Context) (*genai.Client, error) + var apiSurface string if gateway := globalOptions.Gateway(); gateway == "" { var ( httpClient *http.Client @@ -111,6 +118,12 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro httpClient = httpclient.NewHTTPClient(ctx) } + if backend == genai.BackendVertexAI { + apiSurface = apiSurfaceVertexAI + } else { + apiSurface = apiSurfaceGeminiAPI + } + globalOptions.WrapTransport(ctx, httpClient) client, err := genai.NewClient(ctx, &genai.ClientConfig{ @@ -131,6 +144,8 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro return client, nil } } else { + apiSurface = apiSurfaceGateway + // When using a Gateway targeting a Docker domain, tokens are short-lived. // Only require and inject the Docker JWT if the gateway is a .docker.com URL. if err := base.VerifyDockerGatewayAuth(ctx, env, gateway); err != nil { @@ -184,7 +199,8 @@ func NewClient(ctx context.Context, cfg *latest.ModelConfig, env environment.Pro ModelOptions: globalOptions, Env: env, }, - clientFn: clientFn, + clientFn: clientFn, + apiSurface: apiSurface, }, nil } @@ -749,16 +765,11 @@ func (c *Client) CreateChatCompletionStream( if len(config.Tools) > len(allTools) { config.ToolConfig.IncludeServerSideToolInvocations = new(true) } - - // Debug: Log the tools we're sending - slog.DebugContext(ctx, "Gemini tools config", "tools", config.Tools) - for _, tool := range config.Tools { - for _, fn := range tool.FunctionDeclarations { - slog.DebugContext(ctx, "Function", "name", fn.Name, "desc", fn.Description, "params", fn.Parameters) - } - } } + shape := newRequestShape(c, config, len(requestTools)) + slog.DebugContext(ctx, "Gemini request shape", shape.LogAttrs()...) + contents := convertMessagesToGemini(ctx, messages, c.ID(), c.ModelOptions.ModelsDevStore(), c.CapsOverride()) // Debug: Log the messages we're sending diff --git a/pkg/model/provider/gemini/diagnostics.go b/pkg/model/provider/gemini/diagnostics.go new file mode 100644 index 000000000..01fda9283 --- /dev/null +++ b/pkg/model/provider/gemini/diagnostics.go @@ -0,0 +1,155 @@ +package gemini + +import ( + "strings" + + "google.golang.org/genai" +) + +// apiSurface* classifies which backend/transport a Client talks to. Used +// only for diagnostics; never derived from or containing request content. +const ( + apiSurfaceGeminiAPI = "gemini_api" + apiSurfaceVertexAI = "vertex_ai" + apiSurfaceGateway = "gateway" +) + +// RequestShape is a safe, allowlisted summary of a Gemini +// GenerateContentConfig request. Every field is a boolean, count, or a value +// drawn from a small fixed enum (e.g. a modality name or tool kind) — never +// a prompt, tool description/schema, media payload, token, credential, or +// any other provider- or user-supplied free text. It is safe to log at +// Debug level and to include in diagnostics/bug reports. +type RequestShape struct { + // ResponseModalitiesSet reports whether the request specified any + // response modalities at all. ResponseModalities holds the normalized + // (trimmed, upper-cased, de-duplicated) values, e.g. ["TEXT", "IMAGE"]. + ResponseModalitiesSet bool + ResponseModalities []string + + // BuiltInToolKinds lists which fixed-kind built-in tools (google_search, + // google_maps, code_execution) were enabled. BuiltInToolCount is + // len(BuiltInToolKinds). + BuiltInToolKinds []string + BuiltInToolCount int + + // FunctionToolCount is the number of caller-supplied (MCP/custom) + // function tools converted for this request. Never their names, + // descriptions, or parameter schemas. + FunctionToolCount int + + // HasToolConfig, FunctionCallingMode, and ServerSideToolInvocation + // describe genai.ToolConfig without exposing AllowedFunctionNames. + HasToolConfig bool + FunctionCallingMode string + ServerSideToolInvocation bool + + // StructuredOutputPresent reports whether a structured-output response + // (MIME type and/or schema) was requested, never the schema itself. + StructuredOutputPresent bool + + // ThinkingConfigSet and NoThinkingRequested describe the thinking + // configuration shape without exposing budgets or levels (already safe + // enums/numbers, but kept out to keep this struct minimal). + ThinkingConfigSet bool + NoThinkingRequested bool + + // APISurface classifies the backend/transport: "gemini_api", "vertex_ai", + // or "gateway". + APISurface string + + // OutputCapabilityKnown and OutputCapabilityEnabled report whether an + // authoritative source for the model's image/media *output* capability + // was consulted for this request. No such source exists yet — it is the + // subject of a later step — so both fields are always false today, + // deliberately reporting "unknown" rather than guessing from the model + // ID string. + OutputCapabilityKnown bool + OutputCapabilityEnabled bool +} + +// newRequestShape captures a [RequestShape] from a fully-built +// genai.GenerateContentConfig (i.e. after tools/ToolConfig have been +// attached) and the client that built it. +func newRequestShape(c *Client, config *genai.GenerateContentConfig, functionToolCount int) RequestShape { + modalities := normalizeResponseModalities(config.ResponseModalities) + kinds := builtInToolKinds(config.Tools) + + shape := RequestShape{ + ResponseModalitiesSet: len(modalities) > 0, + ResponseModalities: modalities, + BuiltInToolKinds: kinds, + BuiltInToolCount: len(kinds), + FunctionToolCount: functionToolCount, + HasToolConfig: config.ToolConfig != nil, + StructuredOutputPresent: config.ResponseMIMEType != "" || config.ResponseSchema != nil || config.ResponseJsonSchema != nil, + ThinkingConfigSet: config.ThinkingConfig != nil, + NoThinkingRequested: c.ModelOptions.NoThinking(), + APISurface: c.apiSurface, + } + + if config.ToolConfig != nil { + if fc := config.ToolConfig.FunctionCallingConfig; fc != nil { + shape.FunctionCallingMode = string(fc.Mode) + } + shape.ServerSideToolInvocation = config.ToolConfig.IncludeServerSideToolInvocations != nil && + *config.ToolConfig.IncludeServerSideToolInvocations + } + + return shape +} + +// normalizeResponseModalities trims, upper-cases, and de-duplicates raw +// modality strings, dropping empty entries. +func normalizeResponseModalities(raw []string) []string { + out := make([]string, 0, len(raw)) + seen := make(map[string]bool, len(raw)) + for _, m := range raw { + norm := strings.ToUpper(strings.TrimSpace(m)) + if norm == "" || seen[norm] { + continue + } + seen[norm] = true + out = append(out, norm) + } + return out +} + +// builtInToolKinds returns the fixed-kind names of built-in tools present in +// toolsList (e.g. "google_search"). The caller-supplied function-declarations +// tool (added separately by convertToolsToGemini) has none of these fields +// set, so it is naturally excluded. +func builtInToolKinds(toolsList []*genai.Tool) []string { + var kinds []string + for _, t := range toolsList { + switch { + case t.GoogleSearch != nil: + kinds = append(kinds, "google_search") + case t.GoogleMaps != nil: + kinds = append(kinds, "google_maps") + case t.CodeExecution != nil: + kinds = append(kinds, "code_execution") + } + } + return kinds +} + +// LogAttrs renders the shape as a flat slog key/value list. +func (s RequestShape) LogAttrs() []any { + return []any{ + "response_modalities_set", s.ResponseModalitiesSet, + "response_modalities", s.ResponseModalities, + "built_in_tool_kinds", s.BuiltInToolKinds, + "built_in_tool_count", s.BuiltInToolCount, + "function_tool_count", s.FunctionToolCount, + "has_tool_config", s.HasToolConfig, + "function_calling_mode", s.FunctionCallingMode, + "server_side_tool_invocation", s.ServerSideToolInvocation, + "structured_output_present", s.StructuredOutputPresent, + "thinking_config_set", s.ThinkingConfigSet, + "no_thinking_requested", s.NoThinkingRequested, + "api_surface", s.APISurface, + "output_capability_known", s.OutputCapabilityKnown, + "output_capability_enabled", s.OutputCapabilityEnabled, + } +} diff --git a/pkg/model/provider/gemini/diagnostics_test.go b/pkg/model/provider/gemini/diagnostics_test.go new file mode 100644 index 000000000..b25bd56f9 --- /dev/null +++ b/pkg/model/provider/gemini/diagnostics_test.go @@ -0,0 +1,333 @@ +package gemini + +import ( + "bytes" + "io/fs" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/genai" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config/latest" + "github.com/docker/docker-agent/pkg/environment" + "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/docker/docker-agent/pkg/model/provider/options" + "github.com/docker/docker-agent/pkg/tools" +) + +// TestNewRequestShape_Minimal verifies that a request with no tools, no +// thinking override, and no structured output produces an all-zero-value +// shape (aside from the API surface), matching the actual gemini-2.5-flash-image +// image-generation request captured in production diagnostics. +func TestNewRequestShape_Minimal(t *testing.T) { + t.Parallel() + + client := &Client{ + Config: base.Config{ + ModelConfig: latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash-image"}, + }, + apiSurface: apiSurfaceGeminiAPI, + } + + config := client.buildConfig() + config.Tools = client.builtInTools() + + shape := newRequestShape(client, config, 0) + + assert.False(t, shape.ResponseModalitiesSet) + assert.Empty(t, shape.ResponseModalities) + assert.Equal(t, 0, shape.BuiltInToolCount) + assert.Empty(t, shape.BuiltInToolKinds) + assert.Equal(t, 0, shape.FunctionToolCount) + assert.False(t, shape.HasToolConfig) + assert.False(t, shape.StructuredOutputPresent) + assert.False(t, shape.ThinkingConfigSet) + assert.False(t, shape.NoThinkingRequested) + assert.Equal(t, apiSurfaceGeminiAPI, shape.APISurface) + // No authoritative output-capability source exists yet: diagnostics must + // report "unknown", never guess from the model ID. + assert.False(t, shape.OutputCapabilityKnown) + assert.False(t, shape.OutputCapabilityEnabled) +} + +// TestNewRequestShape_ToolsAndBuiltIns mirrors the tool-attachment logic in +// CreateChatCompletionStream to verify the shape captures built-in tool +// kinds/count, custom function count, ToolConfig mode, and the server-side +// tool invocation flag Gemini requires when mixing built-ins with functions. +func TestNewRequestShape_ToolsAndBuiltIns(t *testing.T) { + t.Parallel() + + client := &Client{ + Config: base.Config{ + ModelConfig: latest.ModelConfig{ + Provider: "google", + Model: "gemini-2.5-flash", + ProviderOpts: map[string]any{"google_search": true, "google_maps": true}, + }, + }, + apiSurface: apiSurfaceVertexAI, + } + + config := client.buildConfig() + config.Tools = client.builtInTools() + + requestTools := []tools.Tool{ + {Name: "read_file", Description: "reads a file", Parameters: map[string]any{"type": "object"}}, + {Name: "write_file", Description: "writes a file", Parameters: map[string]any{"type": "object"}}, + } + allTools, err := convertToolsToGemini(requestTools) + require.NoError(t, err) + config.Tools = append(config.Tools, allTools...) + config.ToolConfig = &genai.ToolConfig{ + FunctionCallingConfig: &genai.FunctionCallingConfig{Mode: genai.FunctionCallingConfigModeAuto}, + } + if len(config.Tools) > len(allTools) { + config.ToolConfig.IncludeServerSideToolInvocations = new(true) + } + + shape := newRequestShape(client, config, len(requestTools)) + + assert.Equal(t, 2, shape.BuiltInToolCount) + assert.ElementsMatch(t, []string{"google_search", "google_maps"}, shape.BuiltInToolKinds) + assert.Equal(t, 2, shape.FunctionToolCount) + assert.True(t, shape.HasToolConfig) + assert.Equal(t, string(genai.FunctionCallingConfigModeAuto), shape.FunctionCallingMode) + assert.True(t, shape.ServerSideToolInvocation, "mixing built-in and function tools must set the server-side invocation flag") + assert.Equal(t, apiSurfaceVertexAI, shape.APISurface) +} + +// TestNewRequestShape_ResponseModalitiesNormalized verifies modality values +// are trimmed, upper-cased, and de-duplicated, and that "set" reflects +// presence independent of the normalized values. +func TestNewRequestShape_ResponseModalitiesNormalized(t *testing.T) { + t.Parallel() + + client := &Client{Config: base.Config{ModelConfig: latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash-image"}}} + config := client.buildConfig() + config.ResponseModalities = []string{" text ", "IMAGE", "text", ""} + + shape := newRequestShape(client, config, 0) + + assert.True(t, shape.ResponseModalitiesSet) + assert.Equal(t, []string{"TEXT", "IMAGE"}, shape.ResponseModalities) +} + +// TestNewRequestShape_NoThinkingRequested verifies the title-generation path +// (options.WithNoThinking) is captured: a non-Gemini-3 model gets an +// explicit ThinkingConfig even though the model config itself set no +// thinking budget. +func TestNewRequestShape_NoThinkingRequested(t *testing.T) { + t.Parallel() + + client := &Client{ + Config: base.Config{ + ModelConfig: latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash-image"}, + ModelOptions: options.Apply(options.WithNoThinking()), + }, + } + config := client.buildConfig() + + shape := newRequestShape(client, config, 0) + + assert.True(t, shape.ThinkingConfigSet) + assert.True(t, shape.NoThinkingRequested) +} + +// TestNewRequestShape_StructuredOutputPresent verifies structured-output +// presence is captured as a boolean only — never the schema contents. +func TestNewRequestShape_StructuredOutputPresent(t *testing.T) { + t.Parallel() + + client := &Client{ + Config: base.Config{ + ModelConfig: latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash"}, + ModelOptions: options.Apply(options.WithStructuredOutput(&latest.StructuredOutput{ + Schema: map[string]any{"type": "object", "properties": map[string]any{"secret_field_name": map[string]any{"type": "string"}}}, + })), + }, + } + config := client.buildConfig() + + shape := newRequestShape(client, config, 0) + + require.True(t, shape.StructuredOutputPresent) + + // The shape must never carry the schema itself. + for _, attr := range shape.LogAttrs() { + if s, ok := attr.(string); ok { + assert.NotContains(t, s, "secret_field_name") + } + } +} + +// TestRequestShape_LogAttrsNeverLeaksToolSchemas is the core safety +// regression for this diagnostic: it builds a request with a function tool +// carrying a marker description and parameter schema, then verifies that +// nothing serialized by LogAttrs (or logged through CreateChatCompletionStream) +// contains that marker. This is exactly the kind of leak the previous +// per-tool debug logging in CreateChatCompletionStream produced. +func TestRequestShape_LogAttrsNeverLeaksToolSchemas(t *testing.T) { + t.Parallel() + + const marker = "SECRET_TOOL_DESCRIPTION_MARKER" + + client := &Client{Config: base.Config{ModelConfig: latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash"}}} + config := client.buildConfig() + + requestTools := []tools.Tool{ + {Name: "danger_tool", Description: marker, Parameters: map[string]any{"type": "object", "properties": map[string]any{marker: map[string]any{"type": "string"}}}}, + } + allTools, err := convertToolsToGemini(requestTools) + require.NoError(t, err) + config.Tools = allTools + + shape := newRequestShape(client, config, len(requestTools)) + + for _, attr := range flattenAttrs(shape.LogAttrs()) { + assert.NotContains(t, attr, marker, "RequestShape must never carry tool descriptions or schemas") + } +} + +// flattenAttrs renders slog-style key/value pairs to strings for substring +// assertions in tests. +func flattenAttrs(attrs []any) []string { + out := make([]string, 0, len(attrs)) + for _, a := range attrs { + switch v := a.(type) { + case string: + out = append(out, v) + case []string: + out = append(out, v...) + } + } + return out +} + +// TestCreateChatCompletionStream_DebugLogNeverLeaksToolSchemas is an +// end-to-end regression at the CreateChatCompletionStream boundary: it +// swaps the default slog logger for a buffer-backed one, issues a real +// request (against a local httptest server) with a marker-carrying tool, +// and asserts the marker never reaches the debug log — the leak the old +// per-tool "Function" debug log would have produced. Not parallel: it +// swaps process-global slog state. +func TestCreateChatCompletionStream_DebugLogNeverLeaksToolSchemas(t *testing.T) { + const marker = "SECRET_TOOL_DESCRIPTION_MARKER_E2E" + + var buf bytes.Buffer + prev := slog.Default() + slog.SetDefault(slog.New(slog.NewTextHandler(&buf, &slog.HandlerOptions{Level: slog.LevelDebug}))) + t.Cleanup(func() { slog.SetDefault(prev) }) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + writeGeminiSSEResponse(w) + })) + defer server.Close() + + cfg := &latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash", BaseURL: server.URL} + env := environment.NewMapEnvProvider(map[string]string{"GOOGLE_API_KEY": "test-key"}) + + client, err := NewClient(t.Context(), cfg, env) + require.NoError(t, err) + + stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{ + {Role: chat.MessageRoleUser, Content: "hello"}, + }, []tools.Tool{ + {Name: marker, Description: marker, Parameters: map[string]any{"type": "object", "properties": map[string]any{marker: map[string]any{"type": "string"}}}}, + }) + require.NoError(t, err) + defer stream.Close() + + for { + if _, err := stream.Recv(); err != nil { + break + } + } + + assert.NotContains(t, buf.String(), marker, "debug logs must never contain a tool name, description, or schema value") +} + +// splitDiffNeedles are case-insensitive substrings that would appear in Go +// source if Split Diff View state ever leaked into a name, description, or +// argument sent to a model provider. +var splitDiffNeedles = []string{"splitdiff", "split-diff", "split_diff"} + +// splitDiffTraceRoots are the source directories a Split Diff reference +// would have to pass through to ever reach a Gemini request: this package +// (where the request is finally serialized) and the two layers upstream of +// it that assemble what every provider call receives — pkg/tools (where the +// []tools.Tool list handed to CreateChatCompletionStream is built) and +// pkg/runtime (where messages and tools are gathered before any provider +// call). Dynamic MCP/external tool registrations are out of reach of a +// static scan; the local built-in tool definitions and the request-building +// path itself are not. +var splitDiffTraceRoots = []string{".", "../../../tools", "../../../runtime"} + +// TestSplitDiffView_NeverReferencedInToolOrRequestBuildingSource is a static +// source-scan classification for the owner's screenshot, which showed +// "Split Diff View" active in the sidebar alongside an observed Gemini 400. +// It is a source scan, not a runtime trace: it proves that no non-test .go +// file under [splitDiffTraceRoots] mentions Split Diff View in any form, +// tracing the actual boundary a reference would have to cross — tool +// registration (pkg/tools) and request assembly (pkg/runtime) — rather than +// just this leaf package. It cannot see into MCP/other dynamically +// registered tools, which no static scan can enumerate. +// +// That, combined with reading pkg/tui/service/sessionstate.go (SplitDiffView +// is a plain bool getter/setter backed by session/userconfig persistence, +// wired only into TUI rendering — pkg/tui/components/sidebar/sidebar.go, +// pkg/tui/components/tool/editfile) and confirming pkg/tools and +// pkg/runtime contain no reference to it either, is the evidence for +// classifying Split Diff View as local-only TUI/session state with no path +// into any provider request, Gemini included. +// +// If a future change ever wires it into a tool/config sent to a provider, +// this test fails wherever that reference lands, and the classification +// must move from "local-only" to "serialized-tool". +func TestSplitDiffView_NeverReferencedInToolOrRequestBuildingSource(t *testing.T) { + t.Parallel() + + for _, root := range splitDiffTraceRoots { + for _, path := range goSourceFiles(t, root) { + assertNoSplitDiffReference(t, path) + } + } +} + +// goSourceFiles returns every non-test .go file under root, recursively. +func goSourceFiles(t *testing.T, root string) []string { + t.Helper() + + var files []string + err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + if d.IsDir() || !strings.HasSuffix(d.Name(), ".go") || strings.HasSuffix(d.Name(), "_test.go") { + return nil + } + files = append(files, path) + return nil + }) + require.NoError(t, err) + return files +} + +func assertNoSplitDiffReference(t *testing.T, path string) { + t.Helper() + + content, err := os.ReadFile(path) + require.NoError(t, err) + lower := strings.ToLower(string(content)) + for _, needle := range splitDiffNeedles { + assert.NotContains(t, lower, needle, + "%s must never reference split-diff: it is local-only TUI state, never a tool/request field", path) + } +} From 995765271489d7e405ea2622f264bd6dd8b1b109 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Sun, 16 Aug 2026 10:59:54 +0200 Subject: [PATCH 2/3] feat(#3996): add shared UTF-8-safe display-name sanitization helpers Introduce chat.SanitizeDisplayName and chat.TruncateUTF8Bytes with the canonical MaxSanitizedFieldBytes (128) field bound. Provider-supplied display names are untrusted model output: the sanitizer neutralizes control characters, path separators, traversal-like sequences, and angle brackets (so a name can never forge an XML/tag boundary in prompts), and the truncation helper enforces byte bounds without splitting multi-byte runes. These helpers are shared infrastructure: the Gemini 400 classifier uses them to bound diagnostic fields, and later generated-media commits use them for every persisted or displayed metadata field. --- pkg/chat/display_name.go | 72 ++++++++++++++++++++++++++++++++ pkg/chat/display_name_test.go | 77 +++++++++++++++++++++++++++++++++++ 2 files changed, 149 insertions(+) create mode 100644 pkg/chat/display_name.go create mode 100644 pkg/chat/display_name_test.go diff --git a/pkg/chat/display_name.go b/pkg/chat/display_name.go new file mode 100644 index 000000000..8ab25398b --- /dev/null +++ b/pkg/chat/display_name.go @@ -0,0 +1,72 @@ +package chat + +import ( + "strings" + "unicode/utf8" +) + +// MaxSanitizedFieldBytes is the canonical UTF-8-safe upper bound applied to +// every sanitized display-name and MIME-type field (see +// [SanitizeDisplayName] and the MIME sanitizer in pkg/runtime), independent +// of the separate, larger bound applied to a fully formatted placeholder or +// warning line. It exists so a single overlong provider-supplied field +// cannot itself balloon persisted metadata, session history, or prompts. +const MaxSanitizedFieldBytes = 128 + +// TruncateUTF8Bytes returns the longest prefix of s that is at most max +// bytes and still valid UTF-8 — it never splits a multi-byte rune, so the +// result is always safe to display or re-encode. Used by every canonical +// sanitizer (display name, MIME type, and formatted placeholder/warning +// text) to enforce a byte bound without corrupting non-ASCII content. +func TruncateUTF8Bytes(s string, maxBytes int) string { + if len(s) <= maxBytes { + return s + } + for maxBytes > 0 && !utf8.RuneStart(s[maxBytes]) { + maxBytes-- + } + return s[:maxBytes] +} + +// SanitizeDisplayName neutralizes a provider-supplied display name before +// it is stored in a [Document] or any other session-visible metadata. The +// name is untrusted input (e.g. Gemini's InlineData.DisplayName): it is +// never used verbatim to build a filesystem path (path-bearing consumers +// validate the requested name themselves), but it does end up in +// UI, warnings, placeholder text, and interpolated harness/prompt text +// (see pkg/runtime/harness.go's "..."-delimited blocks), so +// it must not carry control characters, path separators, traversal-like +// sequences, or angle brackets that could confuse a terminal, log line, +// forge a fake XML/tag boundary, or trick a human copy-pasting it into a +// shell. +// +// Control characters, path separators, and '<'/'>' are rewritten to '_'; +// any residual ".." sequence (which could still read as a traversal hint +// even without separators, e.g. "..name") is rewritten too. The result is +// capped at [MaxSanitizedFieldBytes] (UTF-8-safe) and trimmed of +// surrounding whitespace. An empty or all-whitespace input returns "" — +// callers are responsible for substituting their own fallback name. +func SanitizeDisplayName(name string) string { + name = strings.TrimSpace(name) + + var b strings.Builder + b.Grow(len(name)) + for _, r := range name { + switch { + case r == '/' || r == '\\': + b.WriteRune('_') + case r == '<' || r == '>': + b.WriteRune('_') + case r < 0x20 || r == 0x7f: + b.WriteRune('_') + default: + b.WriteRune(r) + } + } + sanitized := b.String() + for strings.Contains(sanitized, "..") { + sanitized = strings.ReplaceAll(sanitized, "..", "_") + } + sanitized = TruncateUTF8Bytes(strings.TrimSpace(sanitized), MaxSanitizedFieldBytes) + return strings.TrimSpace(sanitized) +} diff --git a/pkg/chat/display_name_test.go b/pkg/chat/display_name_test.go new file mode 100644 index 000000000..c39cf828e --- /dev/null +++ b/pkg/chat/display_name_test.go @@ -0,0 +1,77 @@ +package chat + +import ( + "strings" + "testing" + "unicode/utf8" + + "github.com/stretchr/testify/assert" +) + +func TestSanitizeDisplayName(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input string + want string + }{ + {"unchanged for a plain name", "cat.png", "cat.png"}, + {"path separators rewritten", "a/b\\c.png", "a_b_c.png"}, + {"control chars rewritten", "cat\x00\x01name.png", "cat__name.png"}, + {"DEL rewritten", "cat\x7fname.png", "cat_name.png"}, + {"traversal sequence neutralized", "../../etc/passwd", "____etc_passwd"}, + {"traversal without separators neutralized", "..name", "_name"}, + {"leading/trailing whitespace trimmed", " cat.png ", "cat.png"}, + {"empty input yields empty output", "", ""}, + {"whitespace-only input yields empty output", " \t\n ", ""}, + {"control chars only collapse to empty after trim", " \x00\x01 ", "__"}, + {"unicode name preserved", "猫.png", "猫.png"}, + {"angle brackets rewritten", "