download.go

  1package tools
  2
  3import (
  4	"context"
  5	_ "embed"
  6	"fmt"
  7	"io"
  8	"net/http"
  9	"os"
 10	"path/filepath"
 11	"strings"
 12	"time"
 13
 14	"github.com/charmbracelet/crush/internal/permission"
 15	"github.com/charmbracelet/fantasy/ai"
 16)
 17
 18type DownloadParams struct {
 19	URL      string `json:"url" description:"The URL to download from"`
 20	FilePath string `json:"file_path" description:"The local file path where the downloaded content should be saved"`
 21	Timeout  int    `json:"timeout,omitempty" description:"Optional timeout in seconds (max 600)"`
 22}
 23
 24type DownloadPermissionsParams struct {
 25	URL      string `json:"url"`
 26	FilePath string `json:"file_path"`
 27	Timeout  int    `json:"timeout,omitempty"`
 28}
 29
 30const DownloadToolName = "download"
 31
 32//go:embed download.md
 33var downloadDescription []byte
 34
 35func NewDownloadTool(permissions permission.Service, workingDir string) ai.AgentTool {
 36	client := &http.Client{
 37		Timeout: 5 * time.Minute, // Default 5 minute timeout for downloads
 38		Transport: &http.Transport{
 39			MaxIdleConns:        100,
 40			MaxIdleConnsPerHost: 10,
 41			IdleConnTimeout:     90 * time.Second,
 42		},
 43	}
 44	return ai.NewAgentTool(
 45		DownloadToolName,
 46		string(downloadDescription),
 47		func(ctx context.Context, params DownloadParams, call ai.ToolCall) (ai.ToolResponse, error) {
 48			if params.URL == "" {
 49				return ai.NewTextErrorResponse("URL parameter is required"), nil
 50			}
 51
 52			if params.FilePath == "" {
 53				return ai.NewTextErrorResponse("file_path parameter is required"), nil
 54			}
 55
 56			if !strings.HasPrefix(params.URL, "http://") && !strings.HasPrefix(params.URL, "https://") {
 57				return ai.NewTextErrorResponse("URL must start with http:// or https://"), nil
 58			}
 59
 60			// Convert relative path to absolute path
 61			var filePath string
 62			if filepath.IsAbs(params.FilePath) {
 63				filePath = params.FilePath
 64			} else {
 65				filePath = filepath.Join(workingDir, params.FilePath)
 66			}
 67
 68			sessionID := GetSessionFromContext(ctx)
 69			if sessionID == "" {
 70				return ai.ToolResponse{}, fmt.Errorf("session ID is required for downloading files")
 71			}
 72
 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			)
 83
 84			if !p {
 85				return ai.ToolResponse{}, permission.ErrorPermissionDenied
 86			}
 87
 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			}
 99
100			req, err := http.NewRequestWithContext(requestCtx, "GET", params.URL, nil)
101			if err != nil {
102				return ai.ToolResponse{}, fmt.Errorf("failed to create request: %w", err)
103			}
104
105			req.Header.Set("User-Agent", "crush/1.0")
106
107			resp, err := client.Do(req)
108			if err != nil {
109				return ai.ToolResponse{}, fmt.Errorf("failed to download from URL: %w", err)
110			}
111			defer resp.Body.Close()
112
113			if resp.StatusCode != http.StatusOK {
114				return ai.NewTextErrorResponse(fmt.Sprintf("Request failed with status code: %d", resp.StatusCode)), nil
115			}
116
117			// Check content length if available
118			maxSize := int64(100 * 1024 * 1024) // 100MB
119			if resp.ContentLength > maxSize {
120				return ai.NewTextErrorResponse(fmt.Sprintf("File too large: %d bytes (max %d bytes)", resp.ContentLength, maxSize)), nil
121			}
122
123			// Create parent directories if they don't exist
124			if err := os.MkdirAll(filepath.Dir(filePath), 0o755); err != nil {
125				return ai.ToolResponse{}, fmt.Errorf("failed to create parent directories: %w", err)
126			}
127
128			// Create the output file
129			outFile, err := os.Create(filePath)
130			if err != nil {
131				return ai.ToolResponse{}, fmt.Errorf("failed to create output file: %w", err)
132			}
133			defer outFile.Close()
134
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 ai.ToolResponse{}, fmt.Errorf("failed to write file: %w", err)
140			}
141
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 ai.NewTextErrorResponse(fmt.Sprintf("File too large: exceeded %d bytes limit", maxSize)), nil
147			}
148
149			contentType := resp.Header.Get("Content-Type")
150			responseMsg := fmt.Sprintf("Successfully downloaded %d bytes to %s", bytesWritten, filePath)
151			if contentType != "" {
152				responseMsg += fmt.Sprintf(" (Content-Type: %s)", contentType)
153			}
154
155			return ai.NewTextResponse(responseMsg), nil
156		})
157}