git.iko.soy/crush
on 94de6bc help sign in
crush/internal/agent/tools/download.go
157 lines 5.0 kB history raw
1package tools
3import (
4 "cmp"
5 "context"
6 _ "embed"
7 "fmt"
8 "io"
9 "net/http"
10 "os"
11 "path/filepath"
12 "strings"
13 "time"
15 "charm.land/fantasy"
16 "github.com/charmbracelet/crush/internal/filepathext"
17 "github.com/charmbracelet/crush/internal/permission"
20type DownloadParams struct {
21 URL string `json:"url" description:"The URL to download from"`
22 FilePath string `json:"file_path" description:"The local file path where the downloaded content should be saved"`
23 Timeout int `json:"timeout,omitempty" description:"Optional timeout in seconds (max 600)"`
26type DownloadPermissionsParams struct {
27 URL string `json:"url"`
28 FilePath string `json:"file_path"`
29 Timeout int `json:"timeout,omitempty"`
32const DownloadToolName = "download"
34//go:embed download.md
35var downloadDescription []byte
37func NewDownloadTool(permissions permission.Service, workingDir string, client *http.Client) fantasy.AgentTool {
38 if client == nil {
39 client = &http.Client{
40 Timeout: 5 * time.Minute, // Default 5 minute timeout for downloads
41 Transport: &http.Transport{
42 MaxIdleConns: 100,
43 MaxIdleConnsPerHost: 10,
44 IdleConnTimeout: 90 * time.Second,
45 },
46 }
47 }
48 return fantasy.NewParallelAgentTool(
49 DownloadToolName,
50 string(downloadDescription),
51 func(ctx context.Context, params DownloadParams, call fantasy.ToolCall) (fantasy.ToolResponse, error) {
52 if params.URL == "" {
53 return fantasy.NewTextErrorResponse("URL parameter is required"), nil
54 }
56 if params.FilePath == "" {
57 return fantasy.NewTextErrorResponse("file_path parameter is required"), nil
58 }
60 if !strings.HasPrefix(params.URL, "http://") && !strings.HasPrefix(params.URL, "https://") {
61 return fantasy.NewTextErrorResponse("URL must start with http:// or https://"), nil
62 }
64 filePath := filepathext.SmartJoin(workingDir, params.FilePath)
65 relPath, _ := filepath.Rel(workingDir, filePath)
66 relPath = filepath.ToSlash(cmp.Or(relPath, filePath))
68 sessionID := GetSessionFromContext(ctx)
69 if sessionID == "" {
70 return fantasy.ToolResponse{}, fmt.Errorf("session ID is required for downloading files")
71 }
73 p := permissions.Request(
74 permission.CreatePermissionRequest{
75 SessionID: sessionID,
76 Path: filePath,
77 ToolName: DownloadToolName,
78 Action: "download",
79 Description: fmt.Sprintf("Download file from URL: %s to %s", params.URL, filePath),
80 Params: DownloadPermissionsParams(params),
81 },
82 )
84 if !p {
85 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied
86 }
88 // Handle timeout with context
89 requestCtx := ctx
90 if params.Timeout > 0 {
91 maxTimeout := 600 // 10 minutes
92 if params.Timeout > maxTimeout {
93 params.Timeout = maxTimeout
94 }
95 var cancel context.CancelFunc
96 requestCtx, cancel = context.WithTimeout(ctx, time.Duration(params.Timeout)*time.Second)
97 defer cancel()
98 }
100 req, err := http.NewRequestWithContext(requestCtx, "GET", params.URL, nil)
101 if err != nil {
102 return fantasy.ToolResponse{}, fmt.Errorf("failed to create request: %w", err)
103 }
105 req.Header.Set("User-Agent", "crush/1.0")
107 resp, err := client.Do(req)
108 if err != nil {
109 return fantasy.ToolResponse{}, fmt.Errorf("failed to download from URL: %w", err)
110 }
111 defer resp.Body.Close()
113 if resp.StatusCode != http.StatusOK {
114 return fantasy.NewTextErrorResponse(fmt.Sprintf("Request failed with status code: %d", resp.StatusCode)), nil
115 }
117 // Check content length if available
118 maxSize := int64(100 * 1024 * 1024) // 100MB
119 if resp.ContentLength > maxSize {
120 return fantasy.NewTextErrorResponse(fmt.Sprintf("File too large: %d bytes (max %d bytes)", resp.ContentLength, maxSize)), nil
121 }
123 // Create parent directories if they don't exist
124 if err := os.MkdirAll(filepath.Dir(filePath), 0o755); err != nil {
125 return fantasy.ToolResponse{}, fmt.Errorf("failed to create parent directories: %w", err)
126 }
128 // Create the output file
129 outFile, err := os.Create(filePath)
130 if err != nil {
131 return fantasy.ToolResponse{}, fmt.Errorf("failed to create output file: %w", err)
132 }
133 defer outFile.Close()
135 // Copy data with size limit
136 limitedReader := io.LimitReader(resp.Body, maxSize)
137 bytesWritten, err := io.Copy(outFile, limitedReader)
138 if err != nil {
139 return fantasy.ToolResponse{}, fmt.Errorf("failed to write file: %w", err)
140 }
142 // Check if we hit the size limit
143 if bytesWritten == maxSize {
144 // Clean up the file since it might be incomplete
145 os.Remove(filePath)
146 return fantasy.NewTextErrorResponse(fmt.Sprintf("File too large: exceeded %d bytes limit", maxSize)), nil
147 }
149 contentType := resp.Header.Get("Content-Type")
150 responseMsg := fmt.Sprintf("Successfully downloaded %d bytes to %s", bytesWritten, relPath)
151 if contentType != "" {
152 responseMsg += fmt.Sprintf(" (Content-Type: %s)", contentType)
153 }
155 return fantasy.NewTextResponse(responseMsg), nil
156 })