632ba62 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

chore: tidy agent package (#1857)

632ba62Christian Rocha9 months ago

internal/agent/agent.go+15−8

compare with main

unifiedsplit

@@ -40,7 +40,14 @@ import (

40 "github.com/charmbracelet/crush/internal/stringext"

41)

43+const (

43−const defaultSessionName = "Untitled Session"

44+ defaultSessionName = "Untitled Session"

45+

46+ // Constants for auto-summarization thresholds

47+ largeContextWindowThreshold = 200_000

48+ largeContextWindowBuffer = 20_000

49+ smallContextWindowRatio = 0.2

50+)

45//go:embed templates/title.md

52//go:embed templates/title.md

46var titlePrompt []byte

53var titlePrompt []byte

@@ -383,10 +390,10 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy

383 tokens := currentSession.CompletionTokens + currentSession.PromptTokens

390 tokens := currentSession.CompletionTokens + currentSession.PromptTokens

384 remaining := cw - tokens

391 remaining := cw - tokens

385 var threshold int64

392 var threshold int64

386− if cw > 200_000 {

393+ if cw > largeContextWindowThreshold {

387− threshold = 20_000

394+ threshold = largeContextWindowBuffer

388 } else {

395 } else {

389− threshold = int64(float64(cw) * 0.2)

396+ threshold = int64(float64(cw) * smallContextWindowRatio)

390 }

397 }

391 if (remaining <= threshold) && !a.disableAutoSummarize {

398 if (remaining <= threshold) && !a.disableAutoSummarize {

392 shouldSummarize = true

399 shouldSummarize = true

@@ -720,15 +727,15 @@ func (a *sessionAgent) getSessionMessages(ctx context.Context, session session.S

720 }

727 }

721

728

722 if session.SummaryMessageID != "" {

729 if session.SummaryMessageID != "" {

723− summaryMsgInex := -1

730+ summaryMsgIndex := -1

724 for i, msg := range msgs {

731 for i, msg := range msgs {

725 if msg.ID == session.SummaryMessageID {

732 if msg.ID == session.SummaryMessageID {

726− summaryMsgInex = i

733+ summaryMsgIndex = i

727 break

734 break

728 }

735 }

729 }

736 }

730− if summaryMsgInex != -1 {

737+ if summaryMsgIndex != -1 {

731− msgs = msgs[summaryMsgInex:]

738+ msgs = msgs[summaryMsgIndex:]

732 msgs[0].Role = message.User

739 msgs[0].Role = message.User

733 }

740 }

734 }

741 }

internal/agent/coordinator.go+1−1

@@ -489,7 +489,7 @@ func (c *coordinator) buildAgentModels(ctx context.Context, isSubAgent bool) (Mo

489 }

490

491 if smallCatwalkModel == nil {

492− return Model{}, Model{}, errors.New("snall model not found in provider config")

492+ return Model{}, Model{}, errors.New("small model not found in provider config")

493 }

494

495 largeModelID := largeModelCfg.Model

internal/agent/tools/diagnostics.go+1−1

@@ -16,7 +16,7 @@ import (

16)

18type DiagnosticsParams struct {

19− FilePath string `json:"file_path,omitempty" description:"The path to the file to get diagnostics for (leave w empty for project diagnostics)"`

19+ FilePath string `json:"file_path,omitempty" description:"The path to the file to get diagnostics for (leave empty for project diagnostics)"`

20}

22const DiagnosticsToolName = "lsp_diagnostics"

internal/agent/tools/edit.go+12−7

@@ -44,6 +44,11 @@ type EditResponseMetadata struct {

45const EditToolName = "edit"

47+var (

48+ oldStringNotFoundErr = fantasy.NewTextErrorResponse("old_string not found in file. Make sure it matches exactly, including whitespace and line breaks.")

49+ 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")

50+)

51+

47//go:embed edit.md

52//go:embed edit.md

48var editDescription []byte

53var editDescription []byte

@@ -217,12 +222,12 @@ func deleteContent(edit editContext, filePath, oldString string, replaceAll bool

217 newContent = strings.ReplaceAll(oldContent, oldString, "")

222 newContent = strings.ReplaceAll(oldContent, oldString, "")

218 deletionCount = strings.Count(oldContent, oldString)

223 deletionCount = strings.Count(oldContent, oldString)

219 if deletionCount == 0 {

224 if deletionCount == 0 {

220− return fantasy.NewTextErrorResponse("old_string not found in file. Make sure it matches exactly, including whitespace and line breaks"), nil

225+ return oldStringNotFoundErr, nil

221 }

226 }

222 } else {

227 } else {

223 index := strings.Index(oldContent, oldString)

228 index := strings.Index(oldContent, oldString)

224 if index == -1 {

229 if index == -1 {

225− return fantasy.NewTextErrorResponse("old_string not found in file. Make sure it matches exactly, including whitespace and line breaks"), nil

230+ return oldStringNotFoundErr, nil

226 }

231 }

227

232

228 lastIndex := strings.LastIndex(oldContent, oldString)

233 lastIndex := strings.LastIndex(oldContent, oldString)

@@ -287,7 +292,7 @@ func deleteContent(edit editContext, filePath, oldString string, replaceAll bool

287 }

292 }

288 }

293 }

289 if file.Content != oldContent {

294 if file.Content != oldContent {

290− // User Manually changed the content store an intermediate version

295+ // User manually changed the content; store an intermediate version

291 _, err = edit.files.CreateVersion(edit.ctx, sessionID, filePath, oldContent)

296 _, err = edit.files.CreateVersion(edit.ctx, sessionID, filePath, oldContent)

292 if err != nil {

297 if err != nil {

293 slog.Error("Error creating file history version", "error", err)

298 slog.Error("Error creating file history version", "error", err)

@@ -353,17 +358,17 @@ func replaceContent(edit editContext, filePath, oldString, newString string, rep

353 newContent = strings.ReplaceAll(oldContent, oldString, newString)

358 newContent = strings.ReplaceAll(oldContent, oldString, newString)

354 replacementCount = strings.Count(oldContent, oldString)

359 replacementCount = strings.Count(oldContent, oldString)

355 if replacementCount == 0 {

360 if replacementCount == 0 {

356− return fantasy.NewTextErrorResponse("old_string not found in file. Make sure it matches exactly, including whitespace and line breaks"), nil

361+ return oldStringNotFoundErr, nil

357 }

362 }

358 } else {

363 } else {

359 index := strings.Index(oldContent, oldString)

364 index := strings.Index(oldContent, oldString)

360 if index == -1 {

365 if index == -1 {

361− return fantasy.NewTextErrorResponse("old_string not found in file. Make sure it matches exactly, including whitespace and line breaks"), nil

366+ return oldStringNotFoundErr, nil

362 }

367 }

363

368

364 lastIndex := strings.LastIndex(oldContent, oldString)

369 lastIndex := strings.LastIndex(oldContent, oldString)

365 if index != lastIndex {

370 if index != lastIndex {

366− 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

371+ return oldStringMultipleMatchesErr, nil

367 }

372 }

368

373

369 newContent = oldContent[:index] + newString + oldContent[index+len(oldString):]

374 newContent = oldContent[:index] + newString + oldContent[index+len(oldString):]

@@ -425,7 +430,7 @@ func replaceContent(edit editContext, filePath, oldString, newString string, rep

425 }

430 }

426 }

431 }

427 if file.Content != oldContent {

432 if file.Content != oldContent {

428− // User Manually changed the content store an intermediate version

433+ // User manually changed the content; store an intermediate version

429 _, err = edit.files.CreateVersion(edit.ctx, sessionID, filePath, oldContent)

434 _, err = edit.files.CreateVersion(edit.ctx, sessionID, filePath, oldContent)

430 if err != nil {

435 if err != nil {

431 slog.Debug("Error creating file history version", "error", err)

436 slog.Debug("Error creating file history version", "error", err)

internal/agent/tools/fetch.go+13−9

@@ -73,12 +73,14 @@ func NewFetchTool(permissions permission.Service, workingDir string, client *htt

73 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

74 }

76+ // maxFetchTimeoutSeconds is the maximum allowed timeout for fetch requests (2 minutes)

77+ const maxFetchTimeoutSeconds = 120

78+

76 // Handle timeout with context

79 // Handle timeout with context

77 requestCtx := ctx

80 requestCtx := ctx

78 if params.Timeout > 0 {

81 if params.Timeout > 0 {

79− maxTimeout := 120 // 2 minutes

80− if params.Timeout > maxTimeout {

82+ if params.Timeout > maxFetchTimeoutSeconds {

81− params.Timeout = maxTimeout

83+ params.Timeout = maxFetchTimeoutSeconds

82 }

84 }

83 var cancel context.CancelFunc

85 var cancel context.CancelFunc

84 requestCtx, cancel = context.WithTimeout(ctx, time.Duration(params.Timeout)*time.Second)

86 requestCtx, cancel = context.WithTimeout(ctx, time.Duration(params.Timeout)*time.Second)

@@ -102,7 +104,10 @@ func NewFetchTool(permissions permission.Service, workingDir string, client *htt

102 return fantasy.NewTextErrorResponse(fmt.Sprintf("Request failed with status code: %d", resp.StatusCode)), nil

104 return fantasy.NewTextErrorResponse(fmt.Sprintf("Request failed with status code: %d", resp.StatusCode)), nil

103 }

105 }

104

106

107+ // maxFetchResponseSizeBytes is the maximum size of response body to read (5MB)

105− maxSize := int64(5 * 1024 * 1024) // 5MB

108+ const maxFetchResponseSizeBytes = int64(5 * 1024 * 1024)

109+

110+ maxSize := maxFetchResponseSizeBytes

106 body, err := io.ReadAll(io.LimitReader(resp.Body, maxSize))

111 body, err := io.ReadAll(io.LimitReader(resp.Body, maxSize))

107 if err != nil {

112 if err != nil {

108 return fantasy.NewTextErrorResponse("Failed to read response body: " + err.Error()), nil

113 return fantasy.NewTextErrorResponse("Failed to read response body: " + err.Error()), nil

@@ -110,8 +115,8 @@ func NewFetchTool(permissions permission.Service, workingDir string, client *htt

110

115

111 content := string(body)

116 content := string(body)

112

117

113− isValidUt8 := utf8.ValidString(content)

118+ validUTF8 := utf8.ValidString(content)

114− if !isValidUt8 {

119+ if !validUTF8 {

115 return fantasy.NewTextErrorResponse("Response content is not valid UTF-8"), nil

120 return fantasy.NewTextErrorResponse("Response content is not valid UTF-8"), nil

116 }

121 }

117 contentType := resp.Header.Get("Content-Type")

122 contentType := resp.Header.Get("Content-Type")

@@ -154,9 +159,8 @@ func NewFetchTool(permissions permission.Service, workingDir string, client *htt

154 content = "<html>\n<body>\n" + body + "\n</body>\n</html>"

159 content = "<html>\n<body>\n" + body + "\n</body>\n</html>"

155 }

160 }

156 }

161 }

157− // calculate byte size of content

162+ // truncate content if it exceeds max read size

158− contentSize := int64(len(content))

163+ if int64(len(content)) > MaxReadSize {

159− if contentSize > MaxReadSize {

160 content = content[:MaxReadSize]

164 content = content[:MaxReadSize]

161 content += fmt.Sprintf("\n\n[Content truncated to %d bytes]", MaxReadSize)

165 content += fmt.Sprintf("\n\n[Content truncated to %d bytes]", MaxReadSize)

162 }

166 }

internal/agent/tools/ls.go+12−5

@@ -28,10 +28,17 @@ type LSPermissionsParams struct {

28 Depth int `json:"depth"`

29}

31+type NodeType string

32+

33+const (

34+ NodeTypeFile NodeType = "file"

35+ NodeTypeDirectory NodeType = "directory"

36+)

37+

31type TreeNode struct {

38type TreeNode struct {

32 Name string `json:"name"`

39 Name string `json:"name"`

33 Path string `json:"path"`

40 Path string `json:"path"`

34− Type string `json:"type"` // "file" or "directory"

41+ Type NodeType `json:"type"`

35 Children []*TreeNode `json:"children,omitempty"`

42 Children []*TreeNode `json:"children,omitempty"`

36}

43}

@@ -179,9 +186,9 @@ func createFileTree(sortedPaths []string, rootPath string) []*TreeNode {

179

186

180 isLastPart := i == len(parts)-1

187 isLastPart := i == len(parts)-1

181 isDir := !isLastPart || strings.HasSuffix(relativePath, string(filepath.Separator))

188 isDir := !isLastPart || strings.HasSuffix(relativePath, string(filepath.Separator))

182− nodeType := "file"

189+ nodeType := NodeTypeFile

183 if isDir {

190 if isDir {

184− nodeType = "directory"

191+ nodeType = NodeTypeDirectory

185 }

192 }

186 newNode := &TreeNode{

193 newNode := &TreeNode{

187 Name: part,

194 Name: part,

@@ -228,13 +235,13 @@ func printNode(builder *strings.Builder, node *TreeNode, level int) {

228 indent := strings.Repeat(" ", level)

235 indent := strings.Repeat(" ", level)

229

236

230 nodeName := node.Name

237 nodeName := node.Name

231− if node.Type == "directory" {

238+ if node.Type == NodeTypeDirectory {

232 nodeName = nodeName + "/"

239 nodeName = nodeName + "/"

233 }

240 }

234

241

235 fmt.Fprintf(builder, "%s- %s\n", indent, nodeName)

242 fmt.Fprintf(builder, "%s- %s\n", indent, nodeName)

236

243

237− if node.Type == "directory" && len(node.Children) > 0 {

244+ if node.Type == NodeTypeDirectory && len(node.Children) > 0 {

238 for _, child := range node.Children {

245 for _, child := range node.Children {

239 printNode(builder, child, level+1)

246 printNode(builder, child, level+1)

240 }

247 }

internal/agent/tools/view.go−7

@@ -35,13 +35,6 @@ type ViewPermissionsParams struct {

35 Limit int `json:"limit"`

36}

38−type viewTool struct {

39− lspClients *csync.Map[string, *lsp.Client]

40− workingDir string

41− permissions permission.Service

42− skillsPaths []string

43−}

44−

45type ViewResponseMetadata struct {

38type ViewResponseMetadata struct {

46 FilePath string `json:"file_path"`

39 FilePath string `json:"file_path"`

47 Content string `json:"content"`

40 Content string `json:"content"`

internal/agent/tools/write.go+1−8

@@ -36,13 +36,6 @@ type WritePermissionsParams struct {

36 NewContent string `json:"new_content,omitempty"`

37}

39−type writeTool struct {

40− lspClients *csync.Map[string, *lsp.Client]

41− permissions permission.Service

42− files history.Service

43− workingDir string

44−}

45−

46type WriteResponseMetadata struct {

39type WriteResponseMetadata struct {

47 Diff string `json:"diff"`

40 Diff string `json:"diff"`

48 Additions int `json:"additions"`

41 Additions int `json:"additions"`

@@ -148,7 +141,7 @@ func NewWriteTool(lspClients *csync.Map[string, *lsp.Client], permissions permis

148 }

141 }

149 }

142 }

150 if file.Content != oldContent {

143 if file.Content != oldContent {

151− // User Manually changed the content store an intermediate version

144+ // User manually changed the content; store an intermediate version

152 _, err = files.CreateVersion(ctx, sessionID, filePath, oldContent)

145 _, err = files.CreateVersion(ctx, sessionID, filePath, oldContent)

153 if err != nil {

146 if err != nil {

154 slog.Error("Error creating file history version", "error", err)

147 slog.Error("Error creating file history version", "error", err)