ragflow/internal/handler/agent_test.go

//
//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//

package handler

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"gorm.io/gorm"

	"ragflow/internal/common"
	"ragflow/internal/dao"
	"ragflow/internal/entity"
	"ragflow/internal/service"
)

// setupHandlerAgentsTestDB sets up SQLite in-memory DB with tables needed for agent handler tests.
func setupHandlerAgentsTestDB(t *testing.T) *gorm.DB {
	t.Helper()

	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{
		TranslateError: true,
	})
	if err != nil {
		t.Fatalf("failed to open sqlite: %v", err)
	}

	if err := db.AutoMigrate(
		&entity.User{},
		&entity.UserCanvas{},
		&entity.UserCanvasVersion{},
	); err != nil {
		t.Fatalf("failed to migrate: %v", err)
	}

	return db
}

// setupGinContextWithUserAndDB creates a gin context with pre-authenticated user
// and swaps dao.DB to the test database. Returns cleanup function.
func setupGinContextWithUserAndDB(t *testing.T, method, path string) (*gin.Context, *httptest.ResponseRecorder, *gorm.DB) {
	t.Helper()
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Request = httptest.NewRequest(method, path, nil)
	c.Set("user", &entity.User{ID: "user-1"})
	c.Set("user_id", "user-1")

	db := setupHandlerAgentsTestDB(t)
	orig := dao.DB
	dao.DB = db
	t.Cleanup(func() { dao.DB = orig })

	return c, w, db
}

// draw a box with a slot for the old TestListAgents test.
// TestListAgents_Success verifies the ListAgents handler returns a valid response.

// TestListAgentVersionsHandler_Success verifies the happy path with real DB.
func TestListAgentVersionsHandler_Success(t *testing.T) {
	c, w, db := setupGinContextWithUserAndDB(t, "GET", "/api/v1/agents/canvas-1/versions")
	c.Params = gin.Params{{Key: "agent_id", Value: "canvas-1"}}

	// Insert canvas owned by user-1
	db.Create(&entity.UserCanvas{
		ID:     "canvas-1",
		UserID: "user-1",
		Title:  sptr("Test Agent"),
	})

	// Insert 2 versions with staggered timestamps
	now := time.Now()
	db.Create(&entity.UserCanvasVersion{
		ID:           "v2",
		UserCanvasID: "canvas-1",
		Title:        sptr("v2"),
		BaseModel: entity.BaseModel{
			UpdateTime: ptr(now.UnixMilli()),
		},
	})
	db.Create(&entity.UserCanvasVersion{
		ID:           "v1",
		UserCanvasID: "canvas-1",
		Title:        sptr("v1"),
		BaseModel: entity.BaseModel{
			UpdateTime: ptr(now.Add(-time.Hour).UnixMilli()),
		},
	})

	h := NewAgentHandler(service.NewAgentService(), nil)
	h.ListAgentVersions(c)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
	}

	var resp map[string]interface{}
	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
		t.Fatalf("failed to parse response: %v", err)
	}

	code, _ := resp["code"].(float64)
	if code != float64(common.CodeSuccess) {
		t.Fatalf("expected code 0, got %v: %v", code, resp["message"])
	}

	data, ok := resp["data"].([]interface{})
	if !ok {
		t.Fatalf("expected data array, got %T", resp["data"])
	}
	if len(data) != 2 {
		t.Fatalf("expected 2 versions, got %d", len(data))
	}

	v2 := data[0].(map[string]interface{})
	if v2["title"] != "v2" {
		t.Errorf("expected v2 first, got %s", v2["title"])
	}
}

// TestListAgentVersionsHandler_NoPermission verifies cross-user access is denied.
func TestListAgentVersionsHandler_NoPermission(t *testing.T) {
	gin.SetMode(gin.TestMode)

	db := setupHandlerAgentsTestDB(t)
	orig := dao.DB
	dao.DB = db
	t.Cleanup(func() { dao.DB = orig })

	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Request = httptest.NewRequest("GET", "/api/v1/agents/canvas-b/versions", nil)
	c.Set("user", &entity.User{ID: "user-a"})
	c.Set("user_id", "user-a")
	c.Params = gin.Params{{Key: "agent_id", Value: "canvas-b"}}

	// Canvas owned by user-b
	db.Create(&entity.UserCanvas{ID: "canvas-b", UserID: "user-b", Title: sptr("Not Yours")})

	h := NewAgentHandler(service.NewAgentService(), nil)
	h.ListAgentVersions(c)

	var resp map[string]interface{}
	json.Unmarshal(w.Body.Bytes(), &resp)
	code, _ := resp["code"].(float64)
	if code != float64(common.CodeOperatingError) {
		t.Errorf("expected operating error code %d, got %v", common.CodeOperatingError, code)
	}
}

// TestListAgentVersionsHandler_CanvasNotFound verifies behavior for non-existent canvas.
func TestListAgentVersionsHandler_CanvasNotFound(t *testing.T) {
	gin.SetMode(gin.TestMode)

	db := setupHandlerAgentsTestDB(t)
	orig := dao.DB
	dao.DB = db
	t.Cleanup(func() { dao.DB = orig })

	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	c.Request = httptest.NewRequest("GET", "/api/v1/agents/non-existent/versions", nil)
	c.Set("user", &entity.User{ID: "user-1"})
	c.Set("user_id", "user-1")
	c.Params = gin.Params{{Key: "agent_id", Value: "non-existent"}}

	h := NewAgentHandler(service.NewAgentService(), nil)
	h.ListAgentVersions(c)

	var resp map[string]interface{}
	json.Unmarshal(w.Body.Bytes(), &resp)
	code, _ := resp["code"].(float64)
	if code != float64(common.CodeOperatingError) {
		t.Errorf("expected operating error code %d, got %v", common.CodeOperatingError, code)
	}
}

// TestGetAgentVersionHandler_Success verifies getting a specific version.
func TestGetAgentVersionHandler_Success(t *testing.T) {
	c, w, db := setupGinContextWithUserAndDB(t, "GET", "/api/v1/agents/canvas-1/versions/v1")
	c.Params = gin.Params{{Key: "agent_id", Value: "canvas-1"}, {Key: "version_id", Value: "v1"}}

	db.Create(&entity.UserCanvas{
		ID:     "canvas-1",
		UserID: "user-1",
		Title:  sptr("Test Agent"),
	})
	db.Create(&entity.UserCanvasVersion{
		ID:           "v1",
		UserCanvasID: "canvas-1",
		Title:        sptr("version-1"),
		DSL:          entity.JSONMap{"key": "value"},
	})

	h := NewAgentHandler(service.NewAgentService(), nil)
	h.GetAgentVersion(c)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
	}

	var resp map[string]interface{}
	json.Unmarshal(w.Body.Bytes(), &resp)
	code, _ := resp["code"].(float64)
	if code != float64(common.CodeSuccess) {
		t.Fatalf("expected code 0, got %v: %v", code, resp["message"])
	}

	data, ok := resp["data"].(map[string]interface{})
	if !ok {
		t.Fatalf("expected data object, got %T", resp["data"])
	}
	if data["title"] != "version-1" {
		t.Errorf("expected title 'version-1', got %v", data["title"])
	}
	if _, ok := data["dsl"]; !ok {
		t.Errorf("expected dsl field in version detail response")
	}
}

// TestGetAgentVersionHandler_VersionNotFound verifies 404 for missing version.
func TestGetAgentVersionHandler_VersionNotFound(t *testing.T) {
	c, w, db := setupGinContextWithUserAndDB(t, "GET", "/api/v1/agents/canvas-1/versions/non-existent")
	c.Params = gin.Params{{Key: "agent_id", Value: "canvas-1"}, {Key: "version_id", Value: "non-existent"}}

	db.Create(&entity.UserCanvas{
		ID:     "canvas-1",
		UserID: "user-1",
		Title:  sptr("Test Agent"),
	})

	h := NewAgentHandler(service.NewAgentService(), nil)
	h.GetAgentVersion(c)

	var resp map[string]interface{}
	json.Unmarshal(w.Body.Bytes(), &resp)
	code, _ := resp["code"].(float64)
	if code != float64(common.CodeNotFound) {
		t.Errorf("expected not found code %d, got %v", common.CodeNotFound, code)
	}
}

func TestUpdateAgentTagsHandlerSuccess(t *testing.T) {
	c, w, db := setupGinContextWithUserAndDB(t, http.MethodPut, "/api/v1/agents/canvas-1/tags")
	c.Request = httptest.NewRequest(http.MethodPut, "/api/v1/agents/canvas-1/tags", strings.NewReader(`{"tags":["alpha","beta","alpha"]}`))
	c.Request.Header.Set("Content-Type", "application/json")
	c.Params = gin.Params{{Key: "agent_id", Value: "canvas-1"}}

	db.Create(&entity.UserCanvas{
		ID:     "canvas-1",
		UserID: "user-1",
		Title:  sptr("Test Agent"),
	})

	h := NewAgentHandler(service.NewAgentService(), nil)
	h.UpdateAgentTags(c)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
	}

	var resp map[string]interface{}
	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
		t.Fatalf("failed to parse response: %v", err)
	}
	code, _ := resp["code"].(float64)
	if code != float64(common.CodeSuccess) {
		t.Fatalf("expected code %d, got %v: %v", common.CodeSuccess, code, resp["message"])
	}
	if resp["data"] != true {
		t.Fatalf("expected data true, got %v", resp["data"])
	}

	var canvas entity.UserCanvas
	if err := db.Where("id = ?", "canvas-1").First(&canvas).Error; err != nil {
		t.Fatalf("failed to reload canvas: %v", err)
	}
	if canvas.Tags != "alpha,beta" {
		t.Fatalf("expected normalized tags alpha,beta, got %q", canvas.Tags)
	}
}

func TestUpdateAgentTagsHandlerNoPermission(t *testing.T) {
	c, w, db := setupGinContextWithUserAndDB(t, http.MethodPut, "/api/v1/agents/canvas-b/tags")
	c.Request = httptest.NewRequest(http.MethodPut, "/api/v1/agents/canvas-b/tags", strings.NewReader(`{"tags":["alpha"]}`))
	c.Request.Header.Set("Content-Type", "application/json")
	c.Params = gin.Params{{Key: "agent_id", Value: "canvas-b"}}

	db.Create(&entity.UserCanvas{
		ID:         "canvas-b",
		UserID:     "user-b",
		Title:      sptr("Private Agent"),
		Permission: "me",
	})

	h := NewAgentHandler(service.NewAgentService(), nil)
	h.UpdateAgentTags(c)

	var resp map[string]interface{}
	if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
		t.Fatalf("failed to parse response: %v", err)
	}
	code, _ := resp["code"].(float64)
	if code != float64(common.CodeOperatingError) {
		t.Fatalf("expected code %d, got %v: %v", common.CodeOperatingError, code, resp["message"])
	}
	if resp["data"] != false {
		t.Fatalf("expected data false, got %v", resp["data"])
	}
	if resp["message"] != "Agent not found or no permission." {
		t.Fatalf("unexpected message: %v", resp["message"])
	}
}

// sptr returns a pointer to the given string.
// ptr returns a pointer to the given int64.
func ptr(v int64) *int64 { return &v }

// fakeAgentService satisfies the subset of AgentService used by the handler.
// It is injected via a wrapper to avoid importing the real DAO (which requires a DB).
type fakeAgentService struct {
	result       *service.ListAgentsResponse
	code         common.ErrorCode
	err          error
	templates    []*entity.CanvasTemplate
	templatesErr error
}

// agentServiceIface is the minimum interface the handler depends on.
type agentServiceIface interface {
	ListAgents(userID, keywords string, page, pageSize int, orderby string, desc bool, ownerIDs []string, canvasCategory string) (*service.ListAgentsResponse, common.ErrorCode, error)
	ListTemplates() ([]*entity.CanvasTemplate, error)
}

// agentHandlerTestable is a version of AgentHandler that accepts the interface.
type agentHandlerTestable struct {
	svc agentServiceIface
}

func (h *agentHandlerTestable) listAgents(c *gin.Context) {
	user, errorCode, errorMessage := GetUser(c)
	if errorCode != common.CodeSuccess {
		jsonError(c, errorCode, errorMessage)
		return
	}
	result, code, err := h.svc.ListAgents(user.ID, "", 0, 0, "create_time", true, nil, "")
	if err != nil {
		c.JSON(http.StatusOK, gin.H{"code": code, "data": false, "message": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"code": common.CodeSuccess, "data": result, "message": "success"})
}

func (h *agentHandlerTestable) listTemplates(c *gin.Context) {
	if _, errorCode, errorMessage := GetUser(c); errorCode != common.CodeSuccess {
		jsonError(c, errorCode, errorMessage)
		return
	}
	templates, err := h.svc.ListTemplates()
	if err != nil {
		jsonError(c, common.CodeServerError, err.Error())
		return
	}
	if templates == nil {
		templates = []*entity.CanvasTemplate{}
	}
	c.JSON(http.StatusOK, gin.H{"code": common.CodeSuccess, "data": templates, "message": "success"})
}

func (f *fakeAgentService) ListAgents(userID, keywords string, page, pageSize int, orderby string, desc bool, ownerIDs []string, canvasCategory string) (*service.ListAgentsResponse, common.ErrorCode, error) {
	return f.result, f.code, f.err
}

func (f *fakeAgentService) ListTemplates() ([]*entity.CanvasTemplate, error) {
	return f.templates, f.templatesErr
}

func setupAgentRouter(svc agentServiceIface) *gin.Engine {
	gin.SetMode(gin.TestMode)
	r := gin.New()
	h := &agentHandlerTestable{svc: svc}
	r.GET("/api/v1/agents", func(c *gin.Context) {
		c.Set("user", &entity.User{ID: "user-abc"})
		h.listAgents(c)
	})
	r.GET("/api/v1/agents/templates", func(c *gin.Context) {
		c.Set("user", &entity.User{ID: "user-abc"})
		h.listTemplates(c)
	})
	r.GET("/api/v1/agents/templates_anon", func(c *gin.Context) {
		// no user set → unauthenticated probe
		h.listTemplates(c)
	})
	return r
}

func TestListAgents_Success(t *testing.T) {
	title := "My Agent"
	svc := &fakeAgentService{
		result: &service.ListAgentsResponse{
			Canvas: []*service.AgentItem{{ID: "canvas-1", Title: &title, Permission: "me", CanvasCategory: "agent_canvas"}},
			Total:  1,
		},
		code: common.CodeSuccess,
	}

	w := httptest.NewRecorder()
	req, _ := http.NewRequest(http.MethodGet, "/api/v1/agents", nil)
	setupAgentRouter(svc).ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d", w.Code)
	}
	var body map[string]interface{}
	if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
		t.Fatal(err)
	}
	if body["code"] != float64(common.CodeSuccess) {
		t.Errorf("expected code %d, got %v", common.CodeSuccess, body["code"])
	}
	data, ok := body["data"].(map[string]interface{})
	if !ok {
		t.Fatalf("data is not a map: %v", body["data"])
	}
	if data["total"] != float64(1) {
		t.Errorf("expected total=1, got %v", data["total"])
	}
}

func TestListAgentTemplates_Success(t *testing.T) {
	cnvType := "agent"
	svc := &fakeAgentService{
		templates: []*entity.CanvasTemplate{
			{
				ID:             "template-1",
				CanvasType:     &cnvType,
				CanvasCategory: "agent_canvas",
				Title:          entity.JSONMap{"en": "Sample"},
				Description:    entity.JSONMap{"en": "Sample desc"},
			},
		},
	}

	w := httptest.NewRecorder()
	req, _ := http.NewRequest(http.MethodGet, "/api/v1/agents/templates", nil)
	setupAgentRouter(svc).ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("expected 200, got %d body=%s", w.Code, w.Body.String())
	}
	var body map[string]interface{}
	if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
		t.Fatalf("unmarshal: %v body=%s", err, w.Body.String())
	}
	if body["code"] != float64(common.CodeSuccess) {
		t.Errorf("code=%v want %d", body["code"], common.CodeSuccess)
	}
	data, ok := body["data"].([]interface{})
	if !ok {
		t.Fatalf("data is not an array: %v", body["data"])
	}
	if len(data) != 1 {
		t.Fatalf("expected 1 template, got %d", len(data))
	}
	first := data[0].(map[string]interface{})
	if first["id"] != "template-1" {
		t.Errorf("id=%v want template-1", first["id"])
	}
	if first["canvas_category"] != "agent_canvas" {
		t.Errorf("canvas_category=%v want agent_canvas", first["canvas_category"])
	}
}

func TestListAgentTemplates_EmptyIsArrayNotNull(t *testing.T) {
	svc := &fakeAgentService{templates: nil}

	w := httptest.NewRecorder()
	req, _ := http.NewRequest(http.MethodGet, "/api/v1/agents/templates", nil)
	setupAgentRouter(svc).ServeHTTP(w, req)

	var body map[string]interface{}
	_ = json.Unmarshal(w.Body.Bytes(), &body)
	// JSON shape contract: never null - frontends do .map() on it.
	if _, ok := body["data"].([]interface{}); !ok {
		t.Fatalf("data is not an array when templates empty: %v (raw=%s)", body["data"], w.Body.String())
	}
}

func TestListAgentTemplates_RequiresAuth(t *testing.T) {
	svc := &fakeAgentService{templates: []*entity.CanvasTemplate{}}

	w := httptest.NewRecorder()
	req, _ := http.NewRequest(http.MethodGet, "/api/v1/agents/templates_anon", nil)
	setupAgentRouter(svc).ServeHTTP(w, req)

	var body map[string]interface{}
	if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
		t.Fatalf("unmarshal: %v", err)
	}
	if code, _ := body["code"].(float64); int(code) == int(common.CodeSuccess) {
		t.Errorf("expected non-success without auth, got body=%v", body)
	}
}