diff --git a/internal/agent/agent.go b/internal/agent/agent.go index 25186e84..f33ad93f 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "io" + "slices" "strings" "sync" "time" @@ -1049,12 +1050,12 @@ func splitPromptAndHistory(messages []fantasy.Message) (string, []fantasy.FilePa } // Walk backwards to find the last user message - for i := len(messages) - 1; i >= 0; i-- { - if messages[i].Role == fantasy.MessageRoleUser { + for i, message := range slices.Backward(messages) { + if message.Role == fantasy.MessageRoleUser { // Extract text and file parts from the user message var prompt string var files []fantasy.FilePart - for _, part := range messages[i].Content { + for _, part := range message.Content { switch p := part.(type) { case fantasy.TextPart: if prompt == "" { @@ -1215,8 +1216,12 @@ func (a *Agent) GetMCPToolNames() []string { return names } -// GetExtensionToolCount returns the number of tools registered by extensions. -func (a *Agent) GetExtensionToolCount() int { +// GetExtraToolCount returns the number of extra tools on the agent — that is, +// every tool that is neither a core tool nor an MCP tool. The bucket mixes +// extension-registered tools, the built-in activate_skill tool and tools +// supplied by SDK callers, so it must not be reported as an extension count. +// Ask the extension runner for that (see Kit.GetExtensionToolCount). +func (a *Agent) GetExtraToolCount() int { a.toolsMu.RLock() defer a.toolsMu.RUnlock() return len(a.extraTools) diff --git a/internal/app/app.go b/internal/app/app.go index be72bf20..9453bf27 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -6,6 +6,7 @@ import ( "fmt" "log" "os" + "slices" "sync" "sync/atomic" "time" @@ -504,8 +505,8 @@ func (a *App) PopLastUserMessage() (string, []kit.LLMFilePart, error) { // Walk the current branch backwards to find the most recent user message. branch := ts.GetBranch("") var target *session.MessageEntry - for i := len(branch) - 1; i >= 0; i-- { - me, ok := branch[i].(*session.MessageEntry) + for _, b := range slices.Backward(branch) { + me, ok := b.(*session.MessageEntry) if !ok { continue } diff --git a/internal/compaction/compaction.go b/internal/compaction/compaction.go index e8a3f38b..8691cc57 100644 --- a/internal/compaction/compaction.go +++ b/internal/compaction/compaction.go @@ -22,6 +22,7 @@ import ( "context" "encoding/json" "fmt" + "slices" "strings" "charm.land/fantasy" @@ -308,8 +309,8 @@ func FindCutPoint(messages []fantasy.Message, keepRecentTokens int) int { accumulated := 0 - for i := len(messages) - 1; i >= 0; i-- { - accumulated += estimateSingleMessageTokens(messages[i]) + for i, message := range slices.Backward(messages) { + accumulated += estimateSingleMessageTokens(message) if accumulated > keepRecentTokens { cut := i + 1 diff --git a/internal/core/bash.go b/internal/core/bash.go index 4026e146..64425db8 100644 --- a/internal/core/bash.go +++ b/internal/core/bash.go @@ -8,6 +8,7 @@ import ( "os" "os/exec" "regexp" + "slices" "strings" "sync" "time" @@ -141,8 +142,7 @@ func rewriteSudoForStdin(command string) string { // Build result from end to start to preserve indices result := command - for i := len(matches) - 1; i >= 0; i-- { - match := matches[i] + for _, match := range slices.Backward(matches) { start, end := match[0], match[1] matchedText := result[start:end] @@ -389,6 +389,28 @@ func executeBashStreaming(cmdCtx context.Context, call fantasy.ToolCall, cmd *ex } mu.Unlock() } + + // A scan error ends the loop early and leaves data in the pipe. The + // common case is bufio.ErrTooLong: a single line longer than the 1 MB + // limit set above, which any minified JSON or packed asset produces. + // + // Draining is not optional. If the remainder is left unread the child + // blocks writing to a full pipe and cmd.Wait() below never returns, so + // the whole call hangs until the command timeout fires and every byte + // of output is lost. Discard the rest so the process can exit, and + // report the truncation instead of failing silently. + if err := scanner.Err(); err != nil { + discarded, _ := io.Copy(io.Discard, reader) + notice := fmt.Sprintf("[output truncated: %v; discarded %d further bytes]", err, discarded) + outputCallback(call.ID, "bash", notice, isStderr) + mu.Lock() + if isStderr { + stderrChunks = append(stderrChunks, notice) + } else { + stdoutChunks = append(stdoutChunks, notice) + } + mu.Unlock() + } } wg.Add(2) diff --git a/internal/core/bash_streaming_test.go b/internal/core/bash_streaming_test.go new file mode 100644 index 00000000..61926c3c --- /dev/null +++ b/internal/core/bash_streaming_test.go @@ -0,0 +1,122 @@ +package core + +import ( + "context" + "os/exec" + "strings" + "testing" + "time" + + "charm.land/fantasy" +) + +// runStreaming drives executeBashStreaming directly for a shell command and +// returns the tool response plus every chunk pushed to the output callback. +// +// It fails the test rather than blocking forever if the call does not return, +// because the defect it guards against is a deadlock. +func runStreaming(t *testing.T, command string) (fantasy.ToolResponse, []string) { + t.Helper() + + type result struct { + resp fantasy.ToolResponse + chunks []string + err error + } + done := make(chan result, 1) + + go func() { + ctx := context.Background() + cmd := exec.CommandContext(ctx, "bash", "-c", command) + var chunks []string + cb := func(_, _, chunk string, _ bool) { + chunks = append(chunks, chunk) + } + resp, err := executeBashStreaming(ctx, bashCall(command, 0), cmd, cb, "") + done <- result{resp, chunks, err} + }() + + select { + case r := <-done: + if r.err != nil { + t.Fatalf("executeBashStreaming: %v", r.err) + } + return r.resp, r.chunks + case <-time.After(30 * time.Second): + t.Fatal("executeBashStreaming did not return within 30s (deadlocked on an undrained pipe?)") + return fantasy.ToolResponse{}, nil + } +} + +// TestBashStreaming_ReportsOversizedLine is a regression test for a deadlock. +// The streaming scanner caps a single line at 1 MB; a longer one ends the scan +// early. The loop used to ignore scanner.Err() and leave the rest of the pipe +// unread, so the child process blocked writing to a full pipe and cmd.Wait() +// never returned. The call hung until the command timeout fired, and every +// byte of output was lost. +// +// A single line over the limit is not exotic — minified JSON, a packed bundle +// or `cat` of a binary all produce one. +func TestBashStreaming_ReportsOversizedLine(t *testing.T) { + // One line of 2 MB, comfortably over the 1 MB scanner limit. + resp, chunks := runStreaming(t, `head -c 2000000 /dev/zero | tr '\0' 'a'`) + + if !strings.Contains(resp.Content, "output truncated") { + t.Errorf("oversized line should be reported in the result, got %d bytes: %.120q", + len(resp.Content), resp.Content) + } + + var sawNotice bool + for _, c := range chunks { + if strings.Contains(c, "output truncated") { + sawNotice = true + break + } + } + if !sawNotice { + t.Error("oversized line should be reported to the streaming callback too") + } +} + +// TestBashStreaming_NormalOutputUnaffected is the control: ordinary output +// must stream through unchanged, with no truncation notice. +func TestBashStreaming_NormalOutputUnaffected(t *testing.T) { + resp, chunks := runStreaming(t, "printf 'alpha\\nbeta\\ngamma\\n'") + + if strings.Contains(resp.Content, "output truncated") { + t.Errorf("normal output must not be flagged as truncated: %q", resp.Content) + } + for _, want := range []string{"alpha", "beta", "gamma"} { + if !strings.Contains(resp.Content, want) { + t.Errorf("result missing %q, got %q", want, resp.Content) + } + } + if len(chunks) != 3 { + t.Errorf("want 3 streamed chunks, got %d: %v", len(chunks), chunks) + } +} + +// TestBashStreaming_LongButUnderLimit guards the boundary: a line near but +// below the 1 MB scanner cap must stream without a truncation notice, so the +// fix does not over-report. +// +// The response content is still shorter than the 500 KB produced, because +// buildBashResponse applies its own deliberate display truncation +// (defaultMaxLineLen caps a line at 2000 characters). That is a separate, +// intended mechanism; what matters here is that the scanner did not error. +func TestBashStreaming_LongButUnderLimit(t *testing.T) { + resp, chunks := runStreaming(t, `head -c 500000 /dev/zero | tr '\0' 'b'`) + + if strings.Contains(resp.Content, "output truncated") { + t.Error("a 500 KB line is under the 1 MB limit and must not be flagged") + } + + // The scanner should have delivered the line whole, before any + // display-level truncation is applied downstream. + if len(chunks) != 1 { + t.Fatalf("want 1 streamed chunk, got %d", len(chunks)) + } + if got := len(chunks[0]); got != 500000 { + t.Errorf("scanner should deliver the full 500000-byte line, got %d", got) + } +} diff --git a/internal/core/edit.go b/internal/core/edit.go index dd624543..e1935f35 100644 --- a/internal/core/edit.go +++ b/internal/core/edit.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "os" + "slices" "sort" "strings" "unicode" @@ -207,8 +208,7 @@ func applyEdits(content string, edits []replacement) (string, []matchedReplaceme // Apply edits in reverse order (end to start) to maintain stable offsets result := normalizedContent - for i := len(matched) - 1; i >= 0; i-- { - m := matched[i] + for _, m := range slices.Backward(matched) { result = result[:m.start] + m.newText + result[m.end:] } diff --git a/internal/session/tree_manager.go b/internal/session/tree_manager.go index a38df64e..799856e7 100644 --- a/internal/session/tree_manager.go +++ b/internal/session/tree_manager.go @@ -7,6 +7,7 @@ import ( "io" "os" "path/filepath" + "slices" "strings" "sync" "time" @@ -841,8 +842,8 @@ func (tm *TreeManager) BuildContext() (messages []fantasy.Message, provider stri // which older messages are replaced by the summary. var lastCompaction *CompactionEntry var compactionIndex = -1 - for i := len(branch) - 1; i >= 0; i-- { - if c, ok := branch[i].(*CompactionEntry); ok { + for i, b := range slices.Backward(branch) { + if c, ok := b.(*CompactionEntry); ok { lastCompaction = c compactionIndex = i break @@ -1118,8 +1119,8 @@ func (tm *TreeManager) GetContextEntryIDs() []string { // Find the last compaction entry for skip logic. var lastCompaction *CompactionEntry var compactionIndex = -1 - for i := len(branch) - 1; i >= 0; i-- { - if c, ok := branch[i].(*CompactionEntry); ok { + for i, b := range slices.Backward(branch) { + if c, ok := b.(*CompactionEntry); ok { lastCompaction = c compactionIndex = i break @@ -1234,8 +1235,8 @@ func (tm *TreeManager) GetLastCompaction() *CompactionEntry { } branch := tm.getBranchLocked(tm.leafID) - for i := len(branch) - 1; i >= 0; i-- { - if c, ok := branch[i].(*CompactionEntry); ok { + for _, b := range slices.Backward(branch) { + if c, ok := b.(*CompactionEntry); ok { return c } } diff --git a/internal/ui/activity_test.go b/internal/ui/activity_test.go index 55665d7f..7892fa8e 100644 --- a/internal/ui/activity_test.go +++ b/internal/ui/activity_test.go @@ -177,7 +177,7 @@ func TestComposerFillsWidth(t *testing.T) { ic := NewInputComponent(width, nil) rendered := ic.View().Content - firstLine := strings.SplitN(rendered, "\n", 2)[0] + firstLine, _, _ := strings.Cut(rendered, "\n") if got := lipgloss.Width(firstLine); got != width { t.Errorf("expected composer to span %d columns, got %d: %q", width, got, firstLine) diff --git a/internal/ui/block_contract_test.go b/internal/ui/block_contract_test.go index 6742917a..b85d4dae 100644 --- a/internal/ui/block_contract_test.go +++ b/internal/ui/block_contract_test.go @@ -2,6 +2,7 @@ package ui import ( "regexp" + "slices" "strings" "testing" "time" @@ -52,8 +53,8 @@ func widestColumn(rendered string) int { func trailingBlankLines(rendered string) int { lines := visibleLines(rendered) n := 0 - for i := len(lines) - 1; i >= 0; i-- { - if strings.TrimSpace(lines[i]) != "" { + for _, line := range slices.Backward(lines) { + if strings.TrimSpace(line) != "" { break } n++ diff --git a/internal/ui/imagepreview/imagepreview_test.go b/internal/ui/imagepreview/imagepreview_test.go index f70ef137..89134a48 100644 --- a/internal/ui/imagepreview/imagepreview_test.go +++ b/internal/ui/imagepreview/imagepreview_test.go @@ -146,7 +146,7 @@ func TestColumnCountWithinBounds(t *testing.T) { if err != nil { t.Fatalf("unexpected error: %v", err) } - firstRow := strings.SplitN(out, "\n", 2)[0] + firstRow, _, _ := strings.Cut(out, "\n") cols := strings.Count(firstRow, upperHalfBlock) if cols > 8 { t.Errorf("expected at most 8 columns, got %d", cols) diff --git a/internal/ui/model.go b/internal/ui/model.go index 5da5bd53..2466634e 100644 --- a/internal/ui/model.go +++ b/internal/ui/model.go @@ -10,6 +10,7 @@ import ( "os" "os/exec" "path/filepath" + "slices" "sort" "strings" "time" @@ -5906,8 +5907,8 @@ func (m *AppModel) handleCopyCommand() tea.Cmd { text string role string ) - for i := len(m.messages) - 1; i >= 0; i-- { - switch msg := m.messages[i].(type) { + for _, v := range slices.Backward(m.messages) { + switch msg := v.(type) { case *TextMessageItem: if msg.role == "user" || msg.role == "assistant" { text = msg.content diff --git a/internal/ui/scrolllist.go b/internal/ui/scrolllist.go index 0053f1cb..b916735a 100644 --- a/internal/ui/scrolllist.go +++ b/internal/ui/scrolllist.go @@ -1,6 +1,7 @@ package ui import ( + "slices" "strings" "time" @@ -548,7 +549,7 @@ func (s *ScrollList) bottomOffset() (offsetIdx, offsetLine int) { } budget := s.height - for idx := len(s.items) - 1; idx >= 0; idx-- { + for idx := range slices.Backward(s.items) { ih := s.itemHeight(idx) // Account for gap *above* this item (gap between idx-1 and idx). diff --git a/internal/ui/tree_selector.go b/internal/ui/tree_selector.go index 7fb79f66..806e79bc 100644 --- a/internal/ui/tree_selector.go +++ b/internal/ui/tree_selector.go @@ -2,6 +2,7 @@ package ui import ( "fmt" + "slices" "strings" "charm.land/bubbles/v2/key" @@ -107,8 +108,8 @@ func NewTreeSelectorForFork(roots []app.TreeNodeView, leafID string, width, heig ts.initPopup() ts.rebuild() // Position cursor at the last user message before the leaf. - for i := len(ts.flatNodes) - 1; i >= 0; i-- { - if ts.flatNodes[i].Node.IsUserMessage() { + for i, v := range slices.Backward(ts.flatNodes) { + if v.Node.IsUserMessage() { ts.popup.SetCursor(i) break } diff --git a/pkg/extensions/test/harness.go b/pkg/extensions/test/harness.go index cfd06287..e10acadc 100644 --- a/pkg/extensions/test/harness.go +++ b/pkg/extensions/test/harness.go @@ -37,6 +37,7 @@ package test import ( "os" + "reflect" "testing" "github.com/mark3labs/kit/internal/extensions" @@ -120,7 +121,7 @@ func (h *Harness) loadSource(src string, path string) *extensions.LoadedExtensio h.t.Fatalf("extension has no Init function: %v", err) } - initFn, ok := initVal.Interface().(func(extensions.API)) + initFn, ok := reflect.TypeAssert[func(extensions.API)](initVal) if !ok { h.t.Fatalf("Init has wrong signature (want func(ext.API), got %T)", initVal.Interface()) } diff --git a/pkg/kit/kit.go b/pkg/kit/kit.go index 3d6140fc..d4abb929 100644 --- a/pkg/kit/kit.go +++ b/pkg/kit/kit.go @@ -8,6 +8,7 @@ import ( "log" "os" "path/filepath" + "slices" "sort" "strings" "sync" @@ -450,8 +451,16 @@ func (m *Kit) ListMCPServers() []MCPServerStatus { } // GetExtensionToolCount returns the number of tools registered by extensions. +// +// The count comes from the extension runner rather than the agent's extra-tool +// list. That list is a shared bucket also holding the built-in activate_skill +// tool, SDK-supplied Options.ExtraTools and tools added at runtime via +// AddTools, so counting it would attribute all of those to extensions. func (m *Kit) GetExtensionToolCount() int { - return m.agent.GetExtensionToolCount() + if m.extRunner == nil { + return 0 + } + return len(m.extRunner.RegisteredTools()) } // -------------------------------------------------------------------------- @@ -3281,8 +3290,8 @@ func (m *Kit) applyPromptOptions(ctx context.Context, opts PromptOptions) (func( m.promptOptsMu.Lock() var restores []func() restore := func() { - for i := len(restores) - 1; i >= 0; i-- { - restores[i]() + for _, restore := range slices.Backward(restores) { + restore() } m.promptOptsMu.Unlock() } diff --git a/pkg/kit/tool_count_test.go b/pkg/kit/tool_count_test.go new file mode 100644 index 00000000..bb6d7528 --- /dev/null +++ b/pkg/kit/tool_count_test.go @@ -0,0 +1,104 @@ +package kit_test + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/mark3labs/kit/pkg/kit" +) + +// writeSkill creates a minimal skill directory and returns its path. +func writeSkill(t *testing.T, name string) string { + t.Helper() + dir := t.TempDir() + body := "---\nname: " + name + "\ndescription: A skill used to test tool counting.\n---\n\nBody.\n" + path := filepath.Join(dir, name+".md") + if err := os.WriteFile(path, []byte(body), 0o644); err != nil { + t.Fatal(err) + } + return path +} + +// newEchoTool returns a trivial tool used to populate the extra-tool bucket. +func newEchoTool(name string) kit.Tool { + type input struct { + Text string `json:"text"` + } + return kit.NewTool(name, "Echo the input back", + func(ctx context.Context, in input) (kit.ToolOutput, error) { + return kit.TextResult(in.Text), nil + }, + ) +} + +// TestExtensionToolCount_ExcludesSkillTool is a regression test for the +// startup banner reporting "extensions 1 tools" when no extension is loaded. +// +// Loading any skill registers the built-in activate_skill tool into the +// agent's extra-tool bucket. GetExtensionToolCount used to return the size of +// that whole bucket, so a skill was reported as an extension. +func TestExtensionToolCount_ExcludesSkillTool(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "sk-test") + + host, err := kit.New(context.Background(), &kit.Options{ + Model: "openai/gpt-4o-mini", + Quiet: true, + NoSession: true, + NoExtensions: true, + DisableCoreTools: true, + SkipConfig: true, + NoContextFiles: true, + Skills: []string{writeSkill(t, "counted-skill")}, + }) + if err != nil { + t.Fatalf("kit.New: %v", err) + } + defer func() { _ = host.Close() }() + + // Guard the premise: the skill really did load, so a zero count below + // cannot pass for the wrong reason. + if got := len(host.GetSkills()); got != 1 { + t.Fatalf("premise failed: want 1 skill loaded, got %d", got) + } + if got := host.GetExtensionToolCount(); got != 0 { + t.Errorf("no extensions loaded: want extension tool count 0, got %d", got) + } +} + +// TestExtensionToolCount_ExcludesSDKExtraTools covers the same bucket-conflation +// bug from the SDK side: tools supplied via Options.ExtraTools or AddTools are +// caller-provided, not extension-provided, and must not inflate the count. +func TestExtensionToolCount_ExcludesSDKExtraTools(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "sk-test") + + host, err := kit.New(context.Background(), &kit.Options{ + Model: "openai/gpt-4o-mini", + Quiet: true, + NoSession: true, + NoExtensions: true, + DisableCoreTools: true, + SkipConfig: true, + NoSkills: true, + NoContextFiles: true, + }) + if err != nil { + t.Fatalf("kit.New: %v", err) + } + defer func() { _ = host.Close() }() + + if got := host.GetExtensionToolCount(); got != 0 { + t.Fatalf("baseline: want 0, got %d", got) + } + + host.AddTools(newEchoTool("sdk_tool_a"), newEchoTool("sdk_tool_b")) + + // Premise guard: the tools really were added to the agent. + if got := len(host.GetExtraTools()); got != 2 { + t.Fatalf("premise failed: want 2 extra tools, got %d", got) + } + if got := host.GetExtensionToolCount(); got != 0 { + t.Errorf("SDK-supplied tools counted as extensions: want 0, got %d", got) + } +}