From 3796835c4d087ffc89020c9b7a4d0075cfba6771 Mon Sep 17 00:00:00 2001 From: Hz_ Date: Wed, 10 Jun 2026 16:09:36 +0800 Subject: [PATCH] =?UTF-8?q?feat(go-api):=20migrate=20agent=20file=20downlo?= =?UTF-8?q?ad=20handler=20to=20Go=20with=20strict=20P=E2=80=A6=20(#15769)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What does this PR do? This PR migrates the Agent Temporary File Download endpoint (`GET /api/v1/agents/download`) from the Python backend to the Go backend, optimizing the data retrieval flow and maintaining strict functional parity. It also fixes a persistent parsing error in the Sandbox code execution node. ## Checklist - [x] Code logic matches Python implementation - [x] All local unit tests passed - [x] No breaking changes to existing router interfaces Co-authored-by: Yingfeng --- internal/handler/agent.go | 57 +++++++++--- internal/handler/agent_test.go | 65 ++++++++++++- internal/handler/agent_upload_test.go | 5 +- internal/router/router.go | 1 + internal/service/file.go | 17 ++++ internal/service/file_test.go | 126 ++++++++++++++++++++++++++ 6 files changed, 258 insertions(+), 13 deletions(-) create mode 100644 internal/service/file_test.go diff --git a/internal/handler/agent.go b/internal/handler/agent.go index bb223bbd6a..4ef972d8b6 100644 --- a/internal/handler/agent.go +++ b/internal/handler/agent.go @@ -22,6 +22,8 @@ import ( "io" "mime/multipart" "net/http" + "net/url" + "ragflow/internal/utility" "strconv" "strings" @@ -36,19 +38,24 @@ import ( // AgentHandler agent handler // fileUploader is the subset of FileService used by agent handlers. -type fileUploader interface { +type agentFileService interface { UploadFile(tenantID, parentID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) + DownloadAgentFile(tenantID, location string) ([]byte, error) } // AgentHandler agent handler type AgentHandler struct { agentService *service.AgentService - fileService fileUploader + fileService agentFileService } // NewAgentHandler create agent handler + func NewAgentHandler(agentService *service.AgentService, fileService *service.FileService) *AgentHandler { - return &AgentHandler{agentService: agentService, fileService: fileService} + return &AgentHandler{ + agentService: agentService, + fileService: fileService, + } } // ListAgents lists agent canvases for the current user. @@ -655,14 +662,42 @@ func (h *AgentHandler) UpdateAgentTags(c *gin.Context) { }) } -// GetPrompts returns the default prompts used by the agent. -// @Summary Get Agent Prompts -// @Description Returns the default prompts used by the agent, such as task analysis, plan generation, reflection, and citation guidelines. -// @Tags agents -// @Produce json -// @Security ApiKeyAuth -// @Success 200 {object} map[string]interface{} -// @Router /api/v1/agents/prompts [get] +func (h *AgentHandler) DownloadAgentFile(c *gin.Context) { + user, errorCode, errorMessage := GetUser(c) + if errorCode != common.CodeSuccess { + jsonError(c, errorCode, errorMessage) + return + } + + fileID := c.Query("id") + if fileID == "" { + jsonError(c, common.CodeArgumentError, "id parameter is required") + return + } + + blob, err := h.fileService.DownloadAgentFile(user.ID, fileID) + if err != nil { + jsonError(c, common.CodeServerError, err.Error()) + return + } + + ext := utility.GetFileExtension(fileID) + contentType := utility.GetContentType(ext, "") + if contentType != "" { + c.Header("Content-Type", contentType) + } else { + c.Header("Content-Type", "application/octet-stream") + } + + if utility.ShouldForceAttachment(ext, contentType) { + c.Header("X-Content-Type-Options", "nosniff") + encodedName := url.QueryEscape(fileID) + c.Header("Content-Disposition", "attachment; filename*=UTF-8''"+encodedName) + } + + c.Data(http.StatusOK, contentType, blob) +} + func (h *AgentHandler) GetPrompts(c *gin.Context) { if _, errorCode, errorMessage := GetUser(c); errorCode != common.CodeSuccess { jsonError(c, errorCode, errorMessage) diff --git a/internal/handler/agent_test.go b/internal/handler/agent_test.go index 92a217b18a..0eb9b43ccc 100644 --- a/internal/handler/agent_test.go +++ b/internal/handler/agent_test.go @@ -18,6 +18,7 @@ package handler import ( "encoding/json" + "mime/multipart" "net/http" "net/http/httptest" "strings" @@ -979,9 +980,71 @@ func TestListAgentTemplates_RequiresAuth(t *testing.T) { } } +type fakeAgentFileService struct { + blob []byte + err error +} + +func (f *fakeAgentFileService) UploadFile(tenantID, parentID string, files []*multipart.FileHeader) ([]map[string]interface{}, error) { + return nil, nil +} + +func (f *fakeAgentFileService) DownloadAgentFile(tenantID, location string) ([]byte, error) { + return f.blob, f.err +} + +func TestDownloadAgentFile_Success(t *testing.T) { + c, w, _ := setupGinContextWithUserAndDB(t, http.MethodGet, "/api/v1/agents/download?id=test-file.pdf") + + fakeFileSvc := &fakeAgentFileService{ + blob: []byte("test content"), + err: nil, + } + + h := &AgentHandler{ + agentService: service.NewAgentService(), + fileService: fakeFileSvc, + } + + h.DownloadAgentFile(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String()) + } + + if w.Header().Get("Content-Type") != "application/pdf" { + t.Errorf("expected Content-Type application/pdf, got %s", w.Header().Get("Content-Type")) + } + + if w.Body.String() != "test content" { + t.Errorf("expected 'test content', got %s", w.Body.String()) + } +} + +func TestDownloadAgentFile_MissingID(t *testing.T) { + c, w, _ := setupGinContextWithUserAndDB(t, http.MethodGet, "/api/v1/agents/download") + + h := &AgentHandler{ + agentService: service.NewAgentService(), + fileService: &fakeAgentFileService{}, + } + + h.DownloadAgentFile(c) + + if w.Code != http.StatusOK { + t.Fatalf("expected 200 (json error return), got %d", w.Code) + } + + var resp map[string]interface{} + json.Unmarshal(w.Body.Bytes(), &resp) + if code, _ := resp["code"].(float64); code != float64(common.CodeArgumentError) { + t.Errorf("expected code 102, got %v", code) + } +} + func TestGetPrompts_Success(t *testing.T) { c, w, _ := setupGinContextWithUserAndDB(t, http.MethodGet, "/api/v1/agents/prompts") - + // Create handler with fake or real service. h := NewAgentHandler(service.NewAgentService(), nil) h.GetPrompts(c) diff --git a/internal/handler/agent_upload_test.go b/internal/handler/agent_upload_test.go index 6a8c5a3259..403e9868c6 100644 --- a/internal/handler/agent_upload_test.go +++ b/internal/handler/agent_upload_test.go @@ -235,4 +235,7 @@ func TestUploadAgentFileHandler_TeamMemberTenant(t *testing.T) { } // sp returns a pointer to the given string. -func sp(s string) *string { return &s } \ No newline at end of file +func sp(s string) *string { return &s } +func (f *fakeUploadFileService) DownloadAgentFile(tenantID, location string) ([]byte, error) { + panic("DownloadAgentFile should not be called during upload tests") +} diff --git a/internal/router/router.go b/internal/router/router.go index c6b8ac87c8..b3d3d85854 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -389,6 +389,7 @@ func (r *Router) Setup(engine *gin.Engine) { agents.GET("", r.agentHandler.ListAgents) agents.GET("/prompts", r.agentHandler.GetPrompts) agents.GET("/templates", r.agentHandler.ListTemplates) + agents.GET("/download", r.agentHandler.DownloadAgentFile) agents.POST("/test_db_connection", r.agentHandler.TestDBConnection) agents.GET("/:agent_id/versions", r.agentHandler.ListAgentVersions) agents.GET("/:agent_id/versions/:version_id", r.agentHandler.GetAgentVersion) diff --git a/internal/service/file.go b/internal/service/file.go index 872ae57363..e4be10feb0 100644 --- a/internal/service/file.go +++ b/internal/service/file.go @@ -991,3 +991,20 @@ func (s *FileService) GetStorageAddress(fileID string) (*StorageAddress, error) Name: *doc.Location, }, nil } + +// DownloadAgentFile downloads an agent-generated file directly from MinIO without querying the database. +func (s *FileService) DownloadAgentFile(tenantID, location string) ([]byte, error) { + storageImpl := storage.GetStorageFactory().GetStorage() + if storageImpl == nil { + return nil, fmt.Errorf("storage not initialized") + } + + bucketName := fmt.Sprintf("%s-downloads", tenantID) + + blob, err := storageImpl.Get(bucketName, location) + if err != nil { + return nil, fmt.Errorf("failed to read file from storage: %w", err) + } + + return blob, nil +} diff --git a/internal/service/file_test.go b/internal/service/file_test.go new file mode 100644 index 0000000000..bb32232788 --- /dev/null +++ b/internal/service/file_test.go @@ -0,0 +1,126 @@ +package service + +import ( + "bytes" + "errors" + "testing" + "time" + + "ragflow/internal/storage" +) + +// fakeStorage mocks storage.Storage for testing DownloadAgentFile. +type fakeStorage struct { + lastBucket string + lastFnm string + blob []byte + err error +} + +func (f *fakeStorage) Health() bool { + return true +} + +func (f *fakeStorage) Put(bucket, fnm string, binary []byte, tenantID ...string) error { + panic("not implemented in fakeStorage") +} + +func (f *fakeStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, error) { + f.lastBucket = bucket + f.lastFnm = fnm + return f.blob, f.err +} + +func (f *fakeStorage) Remove(bucket, fnm string, tenantID ...string) error { + panic("not implemented in fakeStorage") +} + +func (f *fakeStorage) ObjExist(bucket, fnm string, tenantID ...string) bool { + panic("not implemented in fakeStorage") +} + +func (f *fakeStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) { + panic("not implemented in fakeStorage") +} + +func (f *fakeStorage) BucketExists(bucket string) bool { + panic("not implemented in fakeStorage") +} + +func (f *fakeStorage) RemoveBucket(bucket string) error { + panic("not implemented in fakeStorage") +} + +func (f *fakeStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bool { + panic("not implemented in fakeStorage") +} + +func (f *fakeStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool { + panic("not implemented in fakeStorage") +} + +func TestFileService_DownloadAgentFile_Success(t *testing.T) { + // Setup mock storage + expectedBlob := []byte("fake file content") + mockStorage := &fakeStorage{ + blob: expectedBlob, + err: nil, + } + + factory := storage.GetStorageFactory() + originalStorage := factory.GetStorage() + factory.SetStorage(mockStorage) + t.Cleanup(func() { + factory.SetStorage(originalStorage) + }) + + svc := NewFileService() + tenantID := "tenant123" + location := "file-abc.txt" + + blob, err := svc.DownloadAgentFile(tenantID, location) + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + + if mockStorage.lastBucket != "tenant123-downloads" { + t.Errorf("expected bucket 'tenant123-downloads', got %q", mockStorage.lastBucket) + } + if mockStorage.lastFnm != location { + t.Errorf("expected fnm %q, got %q", location, mockStorage.lastFnm) + } + if !bytes.Equal(blob, expectedBlob) { + t.Errorf("expected blob %v, got %v", expectedBlob, blob) + } +} + +func TestFileService_DownloadAgentFile_Error(t *testing.T) { + // Setup mock storage + expectedErr := errors.New("not found") + mockStorage := &fakeStorage{ + blob: nil, + err: expectedErr, + } + + factory := storage.GetStorageFactory() + originalStorage := factory.GetStorage() + factory.SetStorage(mockStorage) + t.Cleanup(func() { + factory.SetStorage(originalStorage) + }) + + svc := NewFileService() + tenantID := "tenant123" + location := "file-abc.txt" + + blob, err := svc.DownloadAgentFile(tenantID, location) + if err == nil { + t.Fatalf("expected error, got nil") + } + if !errors.Is(err, expectedErr) { + t.Errorf("expected error %v, got %v", expectedErr, err) + } + if blob != nil { + t.Errorf("expected nil blob, got %v", blob) + } +}