From 9124dfb09de20b0f9eee384f71c4b1d894d1f19f Mon Sep 17 00:00:00 2001 From: Hung Nguyen Date: Sun, 31 May 2026 21:10:51 +0700 Subject: [PATCH 1/2] feat(builtin): add WithLLMPreviewRows to cap rows the sub-agent LLM reads Decouples LLM-visible rows from the full set emitted to the OnSQL hook and tools.Result.Structured. Opt-in (n<=0 = no-op, no behavior change); when set, only the first n rows are marshaled into the model context with the true RowCount and Truncated set so it knows it saw a sample. --- pkg/tools/builtin/sql_agent.go | 38 +++++++++- pkg/tools/builtin/sql_agent_preview_test.go | 84 +++++++++++++++++++++ 2 files changed, 121 insertions(+), 1 deletion(-) create mode 100644 pkg/tools/builtin/sql_agent_preview_test.go diff --git a/pkg/tools/builtin/sql_agent.go b/pkg/tools/builtin/sql_agent.go index ebddf5e..f9c6821 100644 --- a/pkg/tools/builtin/sql_agent.go +++ b/pkg/tools/builtin/sql_agent.go @@ -75,6 +75,7 @@ type CallSQLAgentTool struct { examples []SQLExample businessRules []string maxRows int + llmPreviewRows int queryTimeout time.Duration selfConsistency int sessionManager agent.SessionManager @@ -152,6 +153,24 @@ func (t *CallSQLAgentTool) WithMaxRows(n int) *CallSQLAgentTool { return t } +// WithLLMPreviewRows caps the number of result rows formatted into the text +// the sub-agent LLM reads, independent of WithMaxRows. The query still runs +// and returns up to WithMaxRows rows: the full set is emitted on the OnSQL / +// SQLQueryEvent hook (so a host grid stays rich) and attached on +// tools.Result.Structured, while only the first n rows are serialized into +// the model's context — with the true RowCount preserved and Truncated set so +// the model knows it saw a sample and should aggregate or ask the host to +// export rather than assume it read everything. +// +// Use it to keep a large grid (e.g. WithMaxRows(1000)) without paying LLM +// tokens for 1000 wide rows on every "run it and tell me about it" turn. +// n <= 0 disables the cap (default): the model sees the full WithMaxRows set, +// preserving pre-feature behaviour. +func (t *CallSQLAgentTool) WithLLMPreviewRows(n int) *CallSQLAgentTool { + t.llmPreviewRows = n + return t +} + // WithQueryTimeout caps the wall-clock time of each underlying QueryContext // call. d <= 0 disables the timeout (default). Separate from the agent's // overall request context, which may be much longer. @@ -393,6 +412,7 @@ func (t *CallSQLAgentTool) runOnce(ctx context.Context, query string, idx int) s onSQL: t.onSQL, sessionKey: subSessionKey, maxRows: t.maxRows, + llmPreviewRows: t.llmPreviewRows, queryTimeout: t.queryTimeout, allowMutations: t.allowMutations, allowDDL: t.allowDDL, @@ -611,6 +631,7 @@ type executeSQLTool struct { sessionKey string onSQL func(context.Context, SQLQueryEvent) maxRows int + llmPreviewRows int queryTimeout time.Duration allowMutations bool allowDDL bool @@ -714,7 +735,7 @@ func (t *executeSQLTool) makeEmitFunc(ctx context.Context) func(SQLResult) (tool Truncated: res.Truncated, }) } - b, err := json.Marshal(res) + b, err := json.Marshal(previewForLLM(res, t.llmPreviewRows)) if err != nil { return tools.Result{}, fmt.Errorf("tools: marshal result: %w", err) } @@ -722,6 +743,21 @@ func (t *executeSQLTool) makeEmitFunc(ctx context.Context) func(SQLResult) (tool } } +// previewForLLM returns the SQLResult to serialize into the model's context. +// When n > 0 and the result holds more than n rows it returns a shallow copy +// keeping only the first n rows, with the true RowCount preserved and +// Truncated set so the model knows it saw a sample. The original res is left +// intact for the OnSQL hook and tools.Result.Structured. n <= 0 is a no-op. +func previewForLLM(res SQLResult, n int) SQLResult { + if n <= 0 || len(res.Rows) <= n { + return res + } + preview := res + preview.Rows = res.Rows[:n] + preview.Truncated = true + return preview +} + // executeRead runs the read-only path: optional LIMIT injection, then // QueryContext + row iteration into a structured SQLResult. func (t *executeSQLTool) executeRead(ctx context.Context, sqlStr string, emit func(SQLResult) (tools.Result, error)) (tools.Result, error) { diff --git a/pkg/tools/builtin/sql_agent_preview_test.go b/pkg/tools/builtin/sql_agent_preview_test.go new file mode 100644 index 0000000..9b6a567 --- /dev/null +++ b/pkg/tools/builtin/sql_agent_preview_test.go @@ -0,0 +1,84 @@ +package builtin + +import ( + "context" + "encoding/json" + "testing" +) + +func fiveRows() []map[string]any { + rows := make([]map[string]any, 5) + for i := range rows { + rows[i] = map[string]any{"id": i} + } + return rows +} + +// TestEmit_LLMPreviewRows_CapsTextNotHookOrStructured asserts WithLLMPreviewRows +// truncates ONLY the LLM-visible result text; the OnSQL hook (host grid) and +// the tools.Result.Structured payload still carry the full row set. +func TestEmit_LLMPreviewRows_CapsTextNotHookOrStructured(t *testing.T) { + var hookEv SQLQueryEvent + exec := &executeSQLTool{ + sessionKey: "s", + onSQL: func(_ context.Context, ev SQLQueryEvent) { hookEv = ev }, + llmPreviewRows: 2, + } + res := SQLResult{SQL: "SELECT id FROM t", Columns: []string{"id"}, Rows: fiveRows(), RowCount: 5} + + out, err := exec.makeEmitFunc(context.Background())(res) + if err != nil { + t.Fatalf("emit: %v", err) + } + + // Host grid (hook) still sees all 5 rows. + if len(hookEv.Rows) != 5 { + t.Fatalf("hook rows: got %d, want 5 (full set)", len(hookEv.Rows)) + } + // Structured (host via OnToolResult) carries the full set. + if sr, ok := out.Structured.(SQLResult); !ok || len(sr.Rows) != 5 { + t.Fatalf("Structured should carry full 5 rows, got %#v", out.Structured) + } + // LLM-visible Text carries only the preview, with the true RowCount and a + // truncation signal so the model knows it saw a sample. + var seen SQLResult + if err := json.Unmarshal([]byte(out.Text), &seen); err != nil { + t.Fatalf("unmarshal text: %v", err) + } + if len(seen.Rows) != 2 { + t.Fatalf("LLM text rows: got %d, want 2 (preview cap)", len(seen.Rows)) + } + if seen.RowCount != 5 { + t.Fatalf("LLM text RowCount should stay 5 (true count), got %d", seen.RowCount) + } + if !seen.Truncated { + t.Fatal("LLM text should be flagged Truncated so the model knows it's a sample") + } +} + +// TestEmit_LLMPreviewRows_DisabledByDefault asserts the default (0) is a no-op: +// the LLM sees the full set, byte-identical to pre-feature behavior. +func TestEmit_LLMPreviewRows_DisabledByDefault(t *testing.T) { + exec := &executeSQLTool{sessionKey: "s"} // llmPreviewRows defaults to 0 + res := SQLResult{Columns: []string{"id"}, Rows: fiveRows(), RowCount: 5} + + out, err := exec.makeEmitFunc(context.Background())(res) + if err != nil { + t.Fatalf("emit: %v", err) + } + var seen SQLResult + if err := json.Unmarshal([]byte(out.Text), &seen); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if len(seen.Rows) != 5 || seen.Truncated { + t.Fatalf("default must be a no-op: rows=%d truncated=%v", len(seen.Rows), seen.Truncated) + } +} + +// TestWithLLMPreviewRows_SetsField asserts the builder wires the cap through. +func TestWithLLMPreviewRows_SetsField(t *testing.T) { + tool := NewCallSQLAgentTool(nil, "", nil, nil).WithLLMPreviewRows(50) + if tool.llmPreviewRows != 50 { + t.Fatalf("WithLLMPreviewRows: got %d, want 50", tool.llmPreviewRows) + } +} From 7af6f332cbbeb05366ab29cbfa58d6dc9f7479d3 Mon Sep 17 00:00:00 2001 From: Hung Nguyen Date: Sun, 31 May 2026 21:10:51 +0700 Subject: [PATCH 2/2] refactor(builtin): fmt.Fprintf over WriteString(fmt.Sprintf) (QF1012) Pre-existing staticcheck finding surfaced when the package cache was invalidated; no behavior change. --- pkg/tools/builtin/async_tools.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/tools/builtin/async_tools.go b/pkg/tools/builtin/async_tools.go index fadfc38..42a7d18 100644 --- a/pkg/tools/builtin/async_tools.go +++ b/pkg/tools/builtin/async_tools.go @@ -275,7 +275,7 @@ func (t *ListAsyncTasksTool) Execute(ctx context.Context, inputJSON string) (too var sb strings.Builder for _, task := range tasks { - sb.WriteString(fmt.Sprintf("- %s (Agent: %s, Status: %s)\n", task.TaskID, task.AgentName, task.Status)) + fmt.Fprintf(&sb, "- %s (Agent: %s, Status: %s)\n", task.TaskID, task.AgentName, task.Status) } return tools.Text(sb.String()), nil }