Feat/sql agent llm preview rows

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
 }

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,14 +735,29 @@ 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)
 		}
 		return tools.Result{Text: string(b), Structured: res}, nil
 	}
 }
 
+// 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) {

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)
+	}
+}