From 632ba621bfc07721baf9207e49ebb9324b5122b4 Mon Sep 17 00:00:00 2001 From: Christian Rocha Date: Wed, 14 Jan 2026 15:28:01 -0500 Subject: [PATCH] chore: tidy agent package (#1857) --- internal/agent/agent.go | 23 +++++++++++++++-------- internal/agent/coordinator.go | 2 +- internal/agent/tools/diagnostics.go | 2 +- internal/agent/tools/edit.go | 19 ++++++++++++------- internal/agent/tools/fetch.go | 22 +++++++++++++--------- internal/agent/tools/ls.go | 17 ++++++++++++----- internal/agent/tools/view.go | 7 ------- internal/agent/tools/write.go | 9 +-------- 8 files changed, 55 insertions(+), 46 deletions(-) diff --git a/internal/agent/agent.go b/internal/agent/agent.go index c0b9080bb640085c6fd0fdbde8db0fbfe7f476dd..4377ab56cc9e5666100378d6166321051774b26a 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -40,7 +40,14 @@ import ( "github.com/charmbracelet/crush/internal/stringext" ) -const defaultSessionName = "Untitled Session" +const ( + defaultSessionName = "Untitled Session" + + // Constants for auto-summarization thresholds + largeContextWindowThreshold = 200_000 + largeContextWindowBuffer = 20_000 + smallContextWindowRatio = 0.2 +) //go:embed templates/title.md var titlePrompt []byte @@ -383,10 +390,10 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy tokens := currentSession.CompletionTokens + currentSession.PromptTokens remaining := cw - tokens var threshold int64 - if cw > 200_000 { - threshold = 20_000 + if cw > largeContextWindowThreshold { + threshold = largeContextWindowBuffer } else { - threshold = int64(float64(cw) * 0.2) + threshold = int64(float64(cw) * smallContextWindowRatio) } if (remaining <= threshold) && !a.disableAutoSummarize { shouldSummarize = true @@ -720,15 +727,15 @@ func (a *sessionAgent) getSessionMessages(ctx context.Context, session session.S } if session.SummaryMessageID != "" { - summaryMsgInex := -1 + summaryMsgIndex := -1 for i, msg := range msgs { if msg.ID == session.SummaryMessageID { - summaryMsgInex = i + summaryMsgIndex = i break } } - if summaryMsgInex != -1 { - msgs = msgs[summaryMsgInex:] + if summaryMsgIndex != -1 { + msgs = msgs[summaryMsgIndex:] msgs[0].Role = message.User } } diff --git a/internal/agent/coordinator.go b/internal/agent/coordinator.go index d48c074ff5d9c02db1427aa44ee13aed4c9af7e2..943c3efc41b33ea9f261b4ffc7256b6f544beff9 100644 --- a/internal/agent/coordinator.go +++ b/internal/agent/coordinator.go @@ -489,7 +489,7 @@ func (c *coordinator) buildAgentModels(ctx context.Context, isSubAgent bool) (Mo } if smallCatwalkModel == nil { - return Model{}, Model{}, errors.New("snall model not found in provider config") + return Model{}, Model{}, errors.New("small model not found in provider config") } largeModelID := largeModelCfg.Model diff --git a/internal/agent/tools/diagnostics.go b/internal/agent/tools/diagnostics.go index de1fc9a13c95296bb8e637f56b4d0abc4f25c34b..85e8b8d0f7d997f8db83f0d6176ce30c644b86f0 100644 --- a/internal/agent/tools/diagnostics.go +++ b/internal/agent/tools/diagnostics.go @@ -16,7 +16,7 @@ import ( ) type DiagnosticsParams struct { - FilePath string `json:"file_path,omitempty" description:"The path to the file to get diagnostics for (leave w empty for project diagnostics)"` + FilePath string `json:"file_path,omitempty" description:"The path to the file to get diagnostics for (leave empty for project diagnostics)"` } const DiagnosticsToolName = "lsp_diagnostics" diff --git a/internal/agent/tools/edit.go b/internal/agent/tools/edit.go index a3680d009c6d76f8bcb3e39f1c1ddd2041aa1e52..8d8bb87be59dbfb0e038ba40e2b09a3c14d7624b 100644 --- a/internal/agent/tools/edit.go +++ b/internal/agent/tools/edit.go @@ -44,6 +44,11 @@ type EditResponseMetadata struct { const EditToolName = "edit" +var ( + oldStringNotFoundErr = fantasy.NewTextErrorResponse("old_string not found in file. Make sure it matches exactly, including whitespace and line breaks.") + oldStringMultipleMatchesErr = fantasy.NewTextErrorResponse("old_string appears multiple times in the file. Please provide more context to ensure a unique match, or set replace_all to true") +) + //go:embed edit.md var editDescription []byte @@ -217,12 +222,12 @@ func deleteContent(edit editContext, filePath, oldString string, replaceAll bool newContent = strings.ReplaceAll(oldContent, oldString, "") deletionCount = strings.Count(oldContent, oldString) if deletionCount == 0 { - return fantasy.NewTextErrorResponse("old_string not found in file. Make sure it matches exactly, including whitespace and line breaks"), nil + return oldStringNotFoundErr, nil } } else { index := strings.Index(oldContent, oldString) if index == -1 { - return fantasy.NewTextErrorResponse("old_string not found in file. Make sure it matches exactly, including whitespace and line breaks"), nil + return oldStringNotFoundErr, nil } lastIndex := strings.LastIndex(oldContent, oldString) @@ -287,7 +292,7 @@ func deleteContent(edit editContext, filePath, oldString string, replaceAll bool } } if file.Content != oldContent { - // User Manually changed the content store an intermediate version + // User manually changed the content; store an intermediate version _, err = edit.files.CreateVersion(edit.ctx, sessionID, filePath, oldContent) if err != nil { slog.Error("Error creating file history version", "error", err) @@ -353,17 +358,17 @@ func replaceContent(edit editContext, filePath, oldString, newString string, rep newContent = strings.ReplaceAll(oldContent, oldString, newString) replacementCount = strings.Count(oldContent, oldString) if replacementCount == 0 { - return fantasy.NewTextErrorResponse("old_string not found in file. Make sure it matches exactly, including whitespace and line breaks"), nil + return oldStringNotFoundErr, nil } } else { index := strings.Index(oldContent, oldString) if index == -1 { - return fantasy.NewTextErrorResponse("old_string not found in file. Make sure it matches exactly, including whitespace and line breaks"), nil + return oldStringNotFoundErr, nil } lastIndex := strings.LastIndex(oldContent, oldString) if index != lastIndex { - return fantasy.NewTextErrorResponse("old_string appears multiple times in the file. Please provide more context to ensure a unique match, or set replace_all to true"), nil + return oldStringMultipleMatchesErr, nil } newContent = oldContent[:index] + newString + oldContent[index+len(oldString):] @@ -425,7 +430,7 @@ func replaceContent(edit editContext, filePath, oldString, newString string, rep } } if file.Content != oldContent { - // User Manually changed the content store an intermediate version + // User manually changed the content; store an intermediate version _, err = edit.files.CreateVersion(edit.ctx, sessionID, filePath, oldContent) if err != nil { slog.Debug("Error creating file history version", "error", err) diff --git a/internal/agent/tools/fetch.go b/internal/agent/tools/fetch.go index 29fa6f15b5a90fe2dd8d34ef383990b892b742c3..fdb63f057958e5e5a67affe0783a452c27febf41 100644 --- a/internal/agent/tools/fetch.go +++ b/internal/agent/tools/fetch.go @@ -73,12 +73,14 @@ func NewFetchTool(permissions permission.Service, workingDir string, client *htt return fantasy.ToolResponse{}, permission.ErrorPermissionDenied } + // maxFetchTimeoutSeconds is the maximum allowed timeout for fetch requests (2 minutes) + const maxFetchTimeoutSeconds = 120 + // Handle timeout with context requestCtx := ctx if params.Timeout > 0 { - maxTimeout := 120 // 2 minutes - if params.Timeout > maxTimeout { - params.Timeout = maxTimeout + if params.Timeout > maxFetchTimeoutSeconds { + params.Timeout = maxFetchTimeoutSeconds } var cancel context.CancelFunc requestCtx, cancel = context.WithTimeout(ctx, time.Duration(params.Timeout)*time.Second) @@ -102,7 +104,10 @@ func NewFetchTool(permissions permission.Service, workingDir string, client *htt return fantasy.NewTextErrorResponse(fmt.Sprintf("Request failed with status code: %d", resp.StatusCode)), nil } - maxSize := int64(5 * 1024 * 1024) // 5MB + // maxFetchResponseSizeBytes is the maximum size of response body to read (5MB) + const maxFetchResponseSizeBytes = int64(5 * 1024 * 1024) + + maxSize := maxFetchResponseSizeBytes body, err := io.ReadAll(io.LimitReader(resp.Body, maxSize)) if err != nil { return fantasy.NewTextErrorResponse("Failed to read response body: " + err.Error()), nil @@ -110,8 +115,8 @@ func NewFetchTool(permissions permission.Service, workingDir string, client *htt content := string(body) - isValidUt8 := utf8.ValidString(content) - if !isValidUt8 { + validUTF8 := utf8.ValidString(content) + if !validUTF8 { return fantasy.NewTextErrorResponse("Response content is not valid UTF-8"), nil } contentType := resp.Header.Get("Content-Type") @@ -154,9 +159,8 @@ func NewFetchTool(permissions permission.Service, workingDir string, client *htt content = "\n\n" + body + "\n\n" } } - // calculate byte size of content - contentSize := int64(len(content)) - if contentSize > MaxReadSize { + // truncate content if it exceeds max read size + if int64(len(content)) > MaxReadSize { content = content[:MaxReadSize] content += fmt.Sprintf("\n\n[Content truncated to %d bytes]", MaxReadSize) } diff --git a/internal/agent/tools/ls.go b/internal/agent/tools/ls.go index eff7bac0757b5956f669a752c378cab548affb85..20bb1bad4c2d92bb02e564045d4553206ffd12a1 100644 --- a/internal/agent/tools/ls.go +++ b/internal/agent/tools/ls.go @@ -28,10 +28,17 @@ type LSPermissionsParams struct { Depth int `json:"depth"` } +type NodeType string + +const ( + NodeTypeFile NodeType = "file" + NodeTypeDirectory NodeType = "directory" +) + type TreeNode struct { Name string `json:"name"` Path string `json:"path"` - Type string `json:"type"` // "file" or "directory" + Type NodeType `json:"type"` Children []*TreeNode `json:"children,omitempty"` } @@ -179,9 +186,9 @@ func createFileTree(sortedPaths []string, rootPath string) []*TreeNode { isLastPart := i == len(parts)-1 isDir := !isLastPart || strings.HasSuffix(relativePath, string(filepath.Separator)) - nodeType := "file" + nodeType := NodeTypeFile if isDir { - nodeType = "directory" + nodeType = NodeTypeDirectory } newNode := &TreeNode{ Name: part, @@ -228,13 +235,13 @@ func printNode(builder *strings.Builder, node *TreeNode, level int) { indent := strings.Repeat(" ", level) nodeName := node.Name - if node.Type == "directory" { + if node.Type == NodeTypeDirectory { nodeName = nodeName + "/" } fmt.Fprintf(builder, "%s- %s\n", indent, nodeName) - if node.Type == "directory" && len(node.Children) > 0 { + if node.Type == NodeTypeDirectory && len(node.Children) > 0 { for _, child := range node.Children { printNode(builder, child, level+1) } diff --git a/internal/agent/tools/view.go b/internal/agent/tools/view.go index 96150669c292d457f7e8f1bb514f2a209bb73b6e..35865cf43f7c587d60764b3ed177374940bbe2dc 100644 --- a/internal/agent/tools/view.go +++ b/internal/agent/tools/view.go @@ -35,13 +35,6 @@ type ViewPermissionsParams struct { Limit int `json:"limit"` } -type viewTool struct { - lspClients *csync.Map[string, *lsp.Client] - workingDir string - permissions permission.Service - skillsPaths []string -} - type ViewResponseMetadata struct { FilePath string `json:"file_path"` Content string `json:"content"` diff --git a/internal/agent/tools/write.go b/internal/agent/tools/write.go index bbd5e50cf863d4d13503f6cee926b57df80f69bc..8becaea3c08157897dcece7b3d5d4de5cb2ee929 100644 --- a/internal/agent/tools/write.go +++ b/internal/agent/tools/write.go @@ -36,13 +36,6 @@ type WritePermissionsParams struct { NewContent string `json:"new_content,omitempty"` } -type writeTool struct { - lspClients *csync.Map[string, *lsp.Client] - permissions permission.Service - files history.Service - workingDir string -} - type WriteResponseMetadata struct { Diff string `json:"diff"` Additions int `json:"additions"` @@ -148,7 +141,7 @@ func NewWriteTool(lspClients *csync.Map[string, *lsp.Client], permissions permis } } if file.Content != oldContent { - // User Manually changed the content store an intermediate version + // User manually changed the content; store an intermediate version _, err = files.CreateVersion(ctx, sessionID, filePath, oldContent) if err != nil { slog.Error("Error creating file history version", "error", err)