1  // Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
   2  //
   3  // Licensed under the Apache License, Version 2.0 (the "License");
   4  // you may not use this file except in compliance with the License.
   5  // You may obtain a copy of the License at
   6  //
   7  //     http://www.apache.org/licenses/LICENSE-2.0
   8  //
   9  // Unless required by applicable law or agreed to in writing, software
  10  // distributed under the License is distributed on an "AS IS" BASIS,
  11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12  // See the License for the specific language governing permissions and
  13  // limitations under the License.
  14  //
  15  // Requirement: Any integration or derivative work must explicitly attribute
  16  // Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
  17  // documentation or user interface, as detailed in the NOTICE file.
  18  
  19  package websocket
  20  
  21  import (
  22  	"encoding/json"
  23  	"errors"
  24  	"fmt"
  25  	"net/http"
  26  	"os"
  27  	"path/filepath"
  28  	"sort"
  29  	"strings"
  30  
  31  	"github.com/Tencent/AI-Infra-Guard/common/utils"
  32  	"github.com/Tencent/AI-Infra-Guard/internal/gologger"
  33  	"github.com/Tencent/AI-Infra-Guard/internal/mcp"
  34  	"github.com/gin-gonic/gin"
  35  	"gopkg.in/yaml.v3"
  36  )
  37  
  38  func HandleList(root string, loadFile func(filePath string) (interface{}, error)) gin.HandlerFunc {
  39  	return func(c *gin.Context) {
  40  		var allItems []interface{}
  41  		err := filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error {
  42  			if err != nil {
  43  				return nil // skip entry on error
  44  			}
  45  			if !d.IsDir() {
  46  				item, err := loadFile(path)
  47  				if err != nil {
  48  					return err
  49  				}
  50  				allItems = append(allItems, item)
  51  			}
  52  			return nil
  53  		})
  54  		if err != nil {
  55  			c.JSON(http.StatusInternalServerError, gin.H{
  56  				"status":  1,
  57  				"message": err.Error(),
  58  			})
  59  			return
  60  		}
  61  		c.JSON(http.StatusOK, gin.H{
  62  			"status":  0,
  63  			"message": "success",
  64  			"data": gin.H{
  65  				"total": len(allItems),
  66  				"items": allItems,
  67  			},
  68  		})
  69  	}
  70  }
  71  func HandleCreate(readAndSave func(content string) error) gin.HandlerFunc {
  72  	return func(c *gin.Context) {
  73  		var request struct {
  74  			Content string `json:"content" binding:"required"`
  75  		}
  76  		if err := c.ShouldBindJSON(&request); err != nil {
  77  			c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "content parameter is required"})
  78  			return
  79  		}
  80  		if err := readAndSave(request.Content); err != nil {
  81  			c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "failed to save config: " + err.Error()})
  82  			return
  83  		}
  84  		c.JSON(http.StatusOK, gin.H{"status": 0, "message": "created"})
  85  	}
  86  }
  87  
  88  // HandleEdit returns a HandlerFunc for edit requests
  89  func HandleEdit(updateFunc func(id string, content string) error) gin.HandlerFunc {
  90  	return func(c *gin.Context) {
  91  		name := c.Param("id")
  92  		if name == "" {
  93  			c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "name must not be empty"})
  94  			return
  95  		}
  96  
  97  		var request struct {
  98  			Content string `json:"content" binding:"required"`
  99  		}
 100  		if err := c.ShouldBindJSON(&request); err != nil {
 101  			c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "content parameter is required"})
 102  			return
 103  		}
 104  
 105  		if err := updateFunc(c.Param("id"), request.Content); err != nil {
 106  			c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "update failed: " + err.Error()})
 107  			return
 108  		}
 109  
 110  		c.JSON(http.StatusOK, gin.H{"status": 0, "message": "updated"})
 111  	}
 112  }
 113  
 114  // HandleDelete returns a HandlerFunc for delete requests
 115  func HandleDelete(deleteFunc func(id string) error) gin.HandlerFunc {
 116  	return func(c *gin.Context) {
 117  		name := c.Param("id")
 118  		if name == "" {
 119  			c.JSON(http.StatusBadRequest, gin.H{"status": 1, "message": "name must not be empty"})
 120  			return
 121  		}
 122  
 123  		if err := deleteFunc(name); err != nil {
 124  			c.JSON(http.StatusInternalServerError, gin.H{"status": 1, "message": "delete failed: " + err.Error()})
 125  			return
 126  		}
 127  
 128  		c.JSON(http.StatusOK, gin.H{"status": 0, "message": "deleted"})
 129  	}
 130  }
 131  
 132  // MCP prompt management
 133  const MCPROOT = "data/mcp"
 134  
 135  func McpLoadFile(filePath string) (interface{}, error) {
 136  	if filePath == "" {
 137  		return nil, nil
 138  	}
 139  	if !strings.HasSuffix(filePath, ".yaml") {
 140  		return nil, nil
 141  	}
 142  	var ret struct {
 143  		mcp.PluginConfig `yaml:",inline"`
 144  		RawData          string `yaml:"raw_data"`
 145  	}
 146  	data, err := os.ReadFile(filePath)
 147  	if err != nil {
 148  		return nil, err
 149  	}
 150  
 151  	var config mcp.PluginConfig
 152  	err = yaml.Unmarshal(data, &config)
 153  	if err != nil {
 154  		return nil, err
 155  	}
 156  	ret.RawData = string(data)
 157  	ret.PluginConfig = config
 158  	return ret, nil
 159  }
 160  
 161  func mcpReadAndSave(content string) error {
 162  	// Ensure directory exists
 163  	if err := os.MkdirAll(MCPROOT, 0755); err != nil {
 164  		return fmt.Errorf("failed to create directory: %w", err)
 165  	}
 166  
 167  	// Parse YAML to validate format
 168  	var config mcp.PluginConfig
 169  	err := yaml.Unmarshal([]byte(content), &config)
 170  	if err != nil {
 171  		return fmt.Errorf("failed to parse YAML: %w", err)
 172  	}
 173  
 174  	// Extract ID
 175  	id := config.Info.ID
 176  	if id == "" {
 177  		return errors.New("missing info.id field")
 178  	}
 179  
 180  	// Safety check
 181  	if strings.Contains(id, "..") || strings.ContainsAny(id, "/\\<>:\"|?*") {
 182  		return errors.New("invalid filename")
 183  	}
 184  
 185  	filename, err := safeJoinPath(MCPROOT, id+".yaml")
 186  	if err != nil {
 187  		return err
 188  	}
 189  	return os.WriteFile(filename, []byte(content), 0644)
 190  }
 191  
 192  func mcpUpdateFunc(id string, content string) error {
 193  	// Parse YAML to validate content format
 194  	var config mcp.PluginConfig
 195  	if err := yaml.Unmarshal([]byte(content), &config); err != nil {
 196  		return fmt.Errorf("failed to parse YAML: %w", err)
 197  	}
 198  
 199  	// Validate filename safety
 200  	if strings.Contains(id, "..") || strings.ContainsAny(id, "/\\<>:\"|?*") {
 201  		return errors.New("invalid filename")
 202  	}
 203  
 204  	// Use provided name as filename; allow file update without forcing a rename
 205  	filePath, err := safeJoinPath(MCPROOT, id+".yaml")
 206  	if err != nil {
 207  		return err
 208  	}
 209  	return os.WriteFile(filePath, []byte(content), 0644)
 210  }
 211  
 212  func mcpDeleteFunc(id string) error {
 213  	// Validate filename safety
 214  	if strings.Contains(id, "..") || strings.ContainsAny(id, "/\\<>:\"|?*") {
 215  		return errors.New("invalid filename")
 216  	}
 217  
 218  	filePath, pathErr := safeJoinPath(MCPROOT, id+".yaml")
 219  	if pathErr != nil {
 220  		return pathErr
 221  	}
 222  	// Check if file exists
 223  	if _, err := os.Stat(filePath); os.IsNotExist(err) {
 224  		return errors.New("file not found")
 225  	}
 226  	return os.Remove(filePath)
 227  }
 228  
 229  // AI application inspector management
 230  const PromptCollectionsRoot = "data/prompt_collections"
 231  
 232  type PromptCollection struct {
 233  	CodeExec     bool   `json:"code_exec"`
 234  	UploadFile   bool   `json:"upload_file"`
 235  	Product      string `json:"product"`
 236  	MultiModal   bool   `json:"multi_modal"`
 237  	ModelVersion string `json:"model_version"`
 238  	Prompt       string `json:"prompt"`
 239  	UpdateDate   string `json:"update_date"`
 240  	WebSearch    bool   `json:"web_search"`
 241  	SecPolicies  bool   `json:"sec_policies"`
 242  	Affiliation  string `json:"affiliation"`
 243  	Id           string `json:"id"`
 244  }
 245  
 246  func promptCollectionLoadFile(filePath string) (interface{}, error) {
 247  	if filePath == "" {
 248  		return nil, nil
 249  	}
 250  	if !strings.HasSuffix(filePath, ".json") {
 251  		return nil, nil
 252  	}
 253  	data, err := os.ReadFile(filePath)
 254  	if err != nil {
 255  		return nil, err
 256  	}
 257  	var config PromptCollection
 258  	err = json.Unmarshal(data, &config)
 259  	if err != nil {
 260  		return nil, err
 261  	}
 262  	base := filepath.Base(filePath)
 263  	config.Id = strings.Split(base, ".")[0]
 264  	return config, nil
 265  }
 266  
 267  func promptCollectionReadAndSave(content string) error {
 268  	// Validate JSON format
 269  	var collection map[string]interface{}
 270  	err := json.Unmarshal([]byte(content), &collection)
 271  	if err != nil {
 272  		return fmt.Errorf("failed to parse JSON: %w", err)
 273  	}
 274  
 275  	// Use ID as filename
 276  	id, ok := collection["id"].(string)
 277  	if !ok || id == "" {
 278  		return errors.New("missing id field")
 279  	}
 280  
 281  	// Safety check
 282  	if strings.Contains(id, "..") || strings.ContainsAny(id, "/\\<>:\"|?*") {
 283  		return errors.New("invalid filename")
 284  	}
 285  
 286  	filename, err := safeJoinPath(PromptCollectionsRoot, id+".json")
 287  	if err != nil {
 288  		return err
 289  	}
 290  	return os.WriteFile(filename, []byte(content), 0644)
 291  }
 292  
 293  func promptCollectionUpdateFunc(id string, content string) error {
 294  	// Validate JSON format
 295  	var collection map[string]interface{}
 296  	err := json.Unmarshal([]byte(content), &collection)
 297  	if err != nil {
 298  		return fmt.Errorf("invalid JSON format: %w", err)
 299  	}
 300  
 301  	// Validate filename safety
 302  	if strings.Contains(id, "..") || strings.ContainsAny(id, "/\\<>:\"|?*") {
 303  		return errors.New("invalid filename")
 304  	}
 305  
 306  	filename, err := safeJoinPath(PromptCollectionsRoot, id+".json")
 307  	if err != nil {
 308  		return err
 309  	}
 310  	return os.WriteFile(filename, []byte(content), 0644)
 311  }
 312  
 313  func promptCollectionDeleteFunc(id string) error {
 314  	// Validate filename safety
 315  	if strings.Contains(id, "..") || strings.ContainsAny(id, "/\\<>:\"|?*") {
 316  		return errors.New("invalid filename")
 317  	}
 318  
 319  	filePath, err := safeJoinPath(PromptCollectionsRoot, id+".json")
 320  	if err != nil {
 321  		return err
 322  	}
 323  
 324  	// Check if file exists
 325  	if _, err := os.Stat(filePath); os.IsNotExist(err) {
 326  		return errors.New("file not found")
 327  	}
 328  
 329  	return os.Remove(filePath)
 330  }
 331  func GetJailBreak(c *gin.Context) {
 332  	promptSecurityDir, err := utils.ResolvePromptSecurityDir()
 333  	if err != nil {
 334  		c.JSON(http.StatusOK, gin.H{
 335  			"status":  1,
 336  			"message": "Failed to resolve prompt security directory: " + err.Error(),
 337  		})
 338  		return
 339  	}
 340  	dataPath := filepath.Join(promptSecurityDir, "utils", "strategy_map.json")
 341  	data, err := os.ReadFile(dataPath)
 342  	if err != nil {
 343  		c.JSON(http.StatusOK, gin.H{
 344  			"status":  1,
 345  			"message": "Failed to read strategy map: " + err.Error(),
 346  		})
 347  		return
 348  	}
 349  	var data1 interface{}
 350  	err = json.Unmarshal(data, &data1)
 351  	if err != nil {
 352  		c.JSON(http.StatusOK, gin.H{
 353  			"status":  1,
 354  			"message": "Failed to parse strategy map: " + err.Error(),
 355  		})
 356  		return
 357  	}
 358  	c.JSON(http.StatusOK, gin.H{
 359  		"status":  0,
 360  		"message": "success",
 361  		"data":    data1,
 362  	})
 363  }
 364  
 365  // ============== Agent Scan Config Management ==============
 366  const AgentConfigRoot = "data/agents"
 367  const PublicUser = "public_user"
 368  
 369  // getAgentUserDir returns the agent config directory for a user.
 370  // username must have passed validateUsername; safeJoinPath provides an extra defence layer.
 371  func getAgentUserDir(username string) string {
 372  	p, err := safeJoinPath(AgentConfigRoot, username)
 373  	if err != nil {
 374  		// fallback to public user dir; validateUsername should have caught this
 375  		return filepath.Join(AgentConfigRoot, PublicUser)
 376  	}
 377  	return p
 378  }
 379  
 380  // validateUsername checks that a username is safe to use in file paths (prevents path traversal)
 381  func validateUsername(username string) bool {
 382  	if username == "" {
 383  		return false
 384  	}
 385  	if strings.Contains(username, "..") || strings.ContainsAny(username, "/\\<>:\"|?*") {
 386  		return false
 387  	}
 388  	return true
 389  }
 390  
 391  func HandleListAgentNames(c *gin.Context) {
 392  	username := c.GetString("username")
 393  	if !validateUsername(username) {
 394  		username = PublicUser
 395  	}
 396  
 397  	names, err := listAgentConfigNames(username)
 398  	if err != nil {
 399  		c.JSON(http.StatusInternalServerError, gin.H{
 400  			"status":  1,
 401  			"message": "failed to retrieve: " + err.Error(),
 402  		})
 403  		return
 404  	}
 405  	c.JSON(http.StatusOK, gin.H{
 406  		"status":  0,
 407  		"message": "success",
 408  		"data":    names,
 409  	})
 410  }
 411  
 412  func HandleGetAgentConfig(c *gin.Context) {
 413  	username := c.GetString("username")
 414  	if !validateUsername(username) {
 415  		username = PublicUser
 416  	}
 417  
 418  	name := strings.TrimSpace(c.Param("name"))
 419  	if name == "" || !isValidName(name) {
 420  		c.JSON(http.StatusBadRequest, gin.H{
 421  			"status":  1,
 422  			"message": "invalid config name",
 423  		})
 424  		return
 425  	}
 426  
 427  	data, err := readAgentConfigContent(username, name)
 428  	if err != nil {
 429  		if errors.Is(err, os.ErrNotExist) {
 430  			c.JSON(http.StatusNotFound, gin.H{
 431  				"status":  1,
 432  				"message": "config not found",
 433  			})
 434  			return
 435  		}
 436  		c.JSON(http.StatusInternalServerError, gin.H{
 437  			"status":  1,
 438  			"message": "failed to read config: " + err.Error(),
 439  		})
 440  		return
 441  	}
 442  
 443  	c.JSON(http.StatusOK, gin.H{
 444  		"status":  0,
 445  		"message": "success",
 446  		"data":    string(data),
 447  	})
 448  }
 449  
 450  // testAgentConnectivity verifies Agent configuration connectivity.
 451  // Returns (success, message, error).
 452  func testAgentConnectivity(content string) (bool, string, error) {
 453  	agentScanDir, err := utils.ResolveAgentScanDir()
 454  	if err != nil {
 455  		return false, "", fmt.Errorf("failed to resolve agent-scan directory: %v", err)
 456  	}
 457  	uvBin, err := utils.ResolveUvBin()
 458  	if err != nil {
 459  		return false, "", fmt.Errorf("failed to resolve uv binary: %v", err)
 460  	}
 461  	// Create temporary file for the YAML content
 462  	tmpFile, err := os.CreateTemp("", "agent_connect_*.yaml")
 463  	if err != nil {
 464  		return false, "", fmt.Errorf("failed to create temporary config file: %v", err)
 465  	}
 466  	defer os.Remove(tmpFile.Name())
 467  
 468  	// Write YAML content to temp file
 469  	if _, err := tmpFile.WriteString(content); err != nil {
 470  		tmpFile.Close()
 471  		return false, "", fmt.Errorf("failed to write config file: %v", err)
 472  	}
 473  	tmpFile.Close()
 474  
 475  	// Run Python connectivity test script using uv
 476  	var lastLine string
 477  	err = utils.RunCmd(
 478  		agentScanDir,
 479  		uvBin,
 480  		[]string{"run", "test_client_connect.py", "--client_file", tmpFile.Name()},
 481  		func(line string) {
 482  			lastLine += line
 483  		},
 484  	)
 485  
 486  	if err != nil {
 487  		return false, "", fmt.Errorf("connectivity test execution failed: %v", err)
 488  	}
 489  	if lastLine != "" {
 490  		gologger.Infoln("test_agent_connect", lastLine)
 491  	}
 492  
 493  	// Parse the JSON output from Python script
 494  	var result ConnectResultUpdate
 495  	if err := json.Unmarshal([]byte(lastLine), &result); err != nil {
 496  		return false, "", fmt.Errorf("failed to parse connectivity test result: %v", err)
 497  	}
 498  
 499  	return result.Content.Success, result.Content.Message, nil
 500  }
 501  
 502  func HandleSaveAgentConfig(c *gin.Context) {
 503  	username := c.GetString("username")
 504  	if !validateUsername(username) {
 505  		username = PublicUser
 506  	}
 507  
 508  	name := strings.TrimSpace(c.Param("name"))
 509  	if name == "" || !isValidName(name) {
 510  		c.JSON(http.StatusBadRequest, gin.H{
 511  			"status":  1,
 512  			"message": "invalid config name",
 513  		})
 514  		return
 515  	}
 516  
 517  	var req struct {
 518  		Content string `json:"content" binding:"required"`
 519  	}
 520  	if err := c.ShouldBindJSON(&req); err != nil {
 521  		c.JSON(http.StatusBadRequest, gin.H{
 522  			"status":  1,
 523  			"message": "content parameter is required",
 524  		})
 525  		return
 526  	}
 527  	content := strings.TrimSpace(req.Content)
 528  	if content == "" {
 529  		c.JSON(http.StatusBadRequest, gin.H{
 530  			"status":  1,
 531  			"message": "content must not be empty",
 532  		})
 533  		return
 534  	}
 535  
 536  	// Validate Agent connectivity before saving the config (skip when ?verify=false).
 537  	skipVerify := strings.ToLower(strings.TrimSpace(c.Query("verify"))) == "false"
 538  	if !skipVerify {
 539  		success, message, err := testAgentConnectivity(content)
 540  		if err != nil {
 541  			c.JSON(http.StatusInternalServerError, gin.H{
 542  				"status":  1,
 543  				"message": "connectivity check failed: " + err.Error(),
 544  			})
 545  			return
 546  		}
 547  		if !success {
 548  			c.JSON(http.StatusOK, gin.H{
 549  				"status":  1,
 550  				"message": "connectivity check failed: " + message,
 551  			})
 552  			return
 553  		}
 554  	}
 555  
 556  	// Create user-specific directory
 557  	userDir := getAgentUserDir(username)
 558  	if err := os.MkdirAll(userDir, 0755); err != nil {
 559  		c.JSON(http.StatusInternalServerError, gin.H{
 560  			"status":  1,
 561  			"message": "failed to create directory: " + err.Error(),
 562  		})
 563  		return
 564  	}
 565  
 566  	targetPath, err := resolveAgentConfigPathForWrite(username, name)
 567  	if err != nil {
 568  		c.JSON(http.StatusInternalServerError, gin.H{
 569  			"status":  1,
 570  			"message": "failed to save config: " + err.Error(),
 571  		})
 572  		return
 573  	}
 574  
 575  	if err := os.WriteFile(targetPath, []byte(content), 0644); err != nil {
 576  		c.JSON(http.StatusInternalServerError, gin.H{
 577  			"status":  1,
 578  			"message": "failed to save config: " + err.Error(),
 579  		})
 580  		return
 581  	}
 582  
 583  	c.JSON(http.StatusOK, gin.H{
 584  		"status":  0,
 585  		"message": "Saved successfully and connectivity check passed",
 586  	})
 587  }
 588  
 589  func HandleDeleteAgentConfig(c *gin.Context) {
 590  	username := c.GetString("username")
 591  	if !validateUsername(username) {
 592  		username = PublicUser
 593  	}
 594  
 595  	name := strings.TrimSpace(c.Param("name"))
 596  	if name == "" || !isValidName(name) {
 597  		c.JSON(http.StatusBadRequest, gin.H{
 598  			"status":  1,
 599  			"message": "invalid config name",
 600  		})
 601  		return
 602  	}
 603  
 604  	deleted, err := deleteAgentConfig(username, name)
 605  	if err != nil {
 606  		c.JSON(http.StatusInternalServerError, gin.H{
 607  			"status":  1,
 608  			"message": "delete failed: " + err.Error(),
 609  		})
 610  		return
 611  	}
 612  	if !deleted {
 613  		c.JSON(http.StatusNotFound, gin.H{
 614  			"status":  1,
 615  			"message": "config not found",
 616  		})
 617  		return
 618  	}
 619  
 620  	c.JSON(http.StatusOK, gin.H{
 621  		"status":  0,
 622  		"message": "deleted",
 623  	})
 624  }
 625  
 626  // listAgentConfigNamesFromDir reads config names from a given directory
 627  func listAgentConfigNamesFromDir(dir string) ([]string, error) {
 628  	entries, err := os.ReadDir(dir)
 629  	if err != nil {
 630  		if errors.Is(err, os.ErrNotExist) {
 631  			return []string{}, nil
 632  		}
 633  		return nil, err
 634  	}
 635  
 636  	var names []string
 637  	for _, entry := range entries {
 638  		if entry.IsDir() {
 639  			continue
 640  		}
 641  		switch {
 642  		case strings.HasSuffix(entry.Name(), ".yaml"):
 643  			names = append(names, strings.TrimSuffix(entry.Name(), ".yaml"))
 644  		case strings.HasSuffix(entry.Name(), ".yml"):
 645  			names = append(names, strings.TrimSuffix(entry.Name(), ".yml"))
 646  		}
 647  	}
 648  	return names, nil
 649  }
 650  
 651  // listAgentConfigNames lists config names for a user, merging user-dir and public-dir entries with deduplication
 652  func listAgentConfigNames(username string) ([]string, error) {
 653  	// Read configs from user directory
 654  	userDir := getAgentUserDir(username)
 655  	userNames, err := listAgentConfigNamesFromDir(userDir)
 656  	if err != nil {
 657  		return nil, err
 658  	}
 659  
 660  	// If not the public user, also merge configs from the public directory
 661  	if username != PublicUser {
 662  		publicDir := getAgentUserDir(PublicUser)
 663  		publicNames, err := listAgentConfigNamesFromDir(publicDir)
 664  		if err != nil {
 665  			return nil, err
 666  		}
 667  
 668  		// Merge and deduplicate
 669  		nameSet := make(map[string]struct{})
 670  		for _, name := range userNames {
 671  			nameSet[name] = struct{}{}
 672  		}
 673  		for _, name := range publicNames {
 674  			nameSet[name] = struct{}{}
 675  		}
 676  
 677  		userNames = make([]string, 0, len(nameSet))
 678  		for name := range nameSet {
 679  			userNames = append(userNames, name)
 680  		}
 681  	}
 682  
 683  	sort.Strings(userNames)
 684  	return userNames, nil
 685  }
 686  
 687  // readAgentConfigContentFromDir reads config file content from a directory
 688  func readAgentConfigContentFromDir(dir, name string) ([]byte, error) {
 689  	name = strings.TrimSpace(name)
 690  	if name == "" || !isValidName(name) {
 691  		return nil, os.ErrNotExist
 692  	}
 693  
 694  	cleanDir := filepath.Clean(dir)
 695  	rootDir := filepath.Clean(AgentConfigRoot)
 696  	relDir, err := filepath.Rel(rootDir, cleanDir)
 697  	if err != nil || relDir == ".." || strings.HasPrefix(relDir, ".."+string(filepath.Separator)) {
 698  		return nil, os.ErrNotExist
 699  	}
 700  
 701  	for _, ext := range []string{".yaml", ".yml"} {
 702  		path := filepath.Join(cleanDir, name+ext)
 703  		data, err := os.ReadFile(path)
 704  		if err == nil {
 705  			return data, nil
 706  		}
 707  		if !errors.Is(err, os.ErrNotExist) {
 708  			return nil, err
 709  		}
 710  	}
 711  	return nil, os.ErrNotExist
 712  }
 713  
 714  // readAgentConfigContent reads config content, preferring the user directory with fallback to the public directory
 715  func readAgentConfigContent(username, name string) ([]byte, error) {
 716  	// Ensure the username is safe for path construction; fall back to public user directory otherwise
 717  	safeUsername := username
 718  	if !validateUsername(safeUsername) {
 719  		safeUsername = PublicUser
 720  	}
 721  	name = strings.TrimSpace(name)
 722  	if name == "" || !isValidName(name) {
 723  		return nil, os.ErrNotExist
 724  	}
 725  
 726  	// Prefer user directory
 727  	userDir := getAgentUserDir(safeUsername)
 728  	data, err := readAgentConfigContentFromDir(userDir, name)
 729  	if err == nil {
 730  		return data, nil
 731  	}
 732  	if !errors.Is(err, os.ErrNotExist) {
 733  		return nil, err
 734  	}
 735  
 736  	// Fall back to public directory when user directory has no match
 737  	if safeUsername != PublicUser {
 738  		publicDir := getAgentUserDir(PublicUser)
 739  		return readAgentConfigContentFromDir(publicDir, name)
 740  	}
 741  
 742  	return nil, os.ErrNotExist
 743  }
 744  
 745  // resolveAgentConfigPathForWrite resolves the write path (always writes to user directory)
 746  func resolveAgentConfigPathForWrite(username, name string) (string, error) {
 747  	userDir := getAgentUserDir(username)
 748  	p1, err1 := safeJoinPath(userDir, name+".yaml")
 749  	p2, err2 := safeJoinPath(userDir, name+".yml")
 750  	if err1 != nil && err2 != nil {
 751  		return "", err1
 752  	}
 753  	candidates := []string{}
 754  	if err1 == nil {
 755  		candidates = append(candidates, p1)
 756  	}
 757  	if err2 == nil {
 758  		candidates = append(candidates, p2)
 759  	}
 760  	for _, path := range candidates {
 761  		_, statErr := os.Stat(path)
 762  		if statErr == nil {
 763  			return path, nil
 764  		}
 765  		if statErr != nil && !errors.Is(statErr, os.ErrNotExist) {
 766  			return "", statErr
 767  		}
 768  	}
 769  	return p1, nil
 770  }
 771  
 772  // deleteAgentConfig deletes a config entry (only from the user directory)
 773  func deleteAgentConfig(username, name string) (bool, error) {
 774  	userDir := getAgentUserDir(username)
 775  	for _, ext := range []string{".yaml", ".yml"} {
 776  		path, pathErr := safeJoinPath(userDir, name+ext)
 777  		if pathErr != nil {
 778  			continue
 779  		}
 780  		err := os.Remove(path)
 781  		if err == nil {
 782  			return true, nil
 783  		}
 784  		if errors.Is(err, os.ErrNotExist) {
 785  			continue
 786  		}
 787  		return false, err
 788  	}
 789  	return false, nil
 790  }
 791  
 792  // AgentConnectRequest represents the request body for agent connect test
 793  type AgentConnectRequest struct {
 794  	Content string `json:"content"`
 795  }
 796  
 797  // AgentPromptTestRequest represents the request body for agent prompt test
 798  type AgentPromptTestRequest struct {
 799  	Content string `json:"content"`
 800  	Prompt  string `json:"prompt"`
 801  }
 802  
 803  // ProviderResponse represents the provider_response field in result
 804  type ProviderResponse struct {
 805  	Raw    interface{} `json:"raw"`
 806  	Output *string     `json:"output"`
 807  	Error  *string     `json:"error"`
 808  }
 809  
 810  // ConnectResultContent represents the content of resultUpdate response
 811  type ConnectResultContent struct {
 812  	Success          bool              `json:"success"`
 813  	Message          string            `json:"message"`
 814  	ProviderResponse *ProviderResponse `json:"provider_response"`
 815  }
 816  
 817  // ConnectResultUpdate represents the resultUpdate response from Python script
 818  type ConnectResultUpdate struct {
 819  	Type    string               `json:"type"`
 820  	Content ConnectResultContent `json:"content"`
 821  }
 822  
 823  func HandleAgentConnect(c *gin.Context) {
 824  	var req AgentConnectRequest
 825  	if err := c.ShouldBindJSON(&req); err != nil {
 826  		c.JSON(http.StatusBadRequest, gin.H{
 827  			"status":  1,
 828  			"message": "Invalid request body: " + err.Error(),
 829  		})
 830  		return
 831  	}
 832  
 833  	if req.Content == "" {
 834  		c.JSON(http.StatusBadRequest, gin.H{
 835  			"status":  1,
 836  			"message": "Content cannot be empty",
 837  		})
 838  		return
 839  	}
 840  
 841  	// Delegate to shared connectivity test helper
 842  	success, message, err := testAgentConnectivity(req.Content)
 843  	if err != nil {
 844  		c.JSON(http.StatusInternalServerError, gin.H{
 845  			"status":  1,
 846  			"message": "Failed to run connectivity test: " + err.Error(),
 847  		})
 848  		return
 849  	}
 850  
 851  	if success {
 852  		c.JSON(http.StatusOK, gin.H{
 853  			"status":  0,
 854  			"message": message,
 855  		})
 856  	} else {
 857  		c.JSON(http.StatusOK, gin.H{
 858  			"status":  1,
 859  			"message": message,
 860  		})
 861  	}
 862  }
 863  
 864  func HandleAgentPromptTest(c *gin.Context) {
 865  	var req AgentPromptTestRequest
 866  	if err := c.ShouldBindJSON(&req); err != nil {
 867  		c.JSON(http.StatusBadRequest, gin.H{
 868  			"status":  1,
 869  			"message": "Invalid request body: " + err.Error(),
 870  		})
 871  		return
 872  	}
 873  
 874  	if req.Content == "" {
 875  		c.JSON(http.StatusBadRequest, gin.H{
 876  			"status":  1,
 877  			"message": "Content cannot be empty",
 878  		})
 879  		return
 880  	}
 881  
 882  	if req.Prompt == "" {
 883  		c.JSON(http.StatusBadRequest, gin.H{
 884  			"status":  1,
 885  			"message": "Prompt cannot be empty",
 886  		})
 887  		return
 888  	}
 889  
 890  	// Create temporary file for the YAML content
 891  	tmpFile, err := os.CreateTemp("", "agent_prompt_test_*.yaml")
 892  	if err != nil {
 893  		c.JSON(http.StatusInternalServerError, gin.H{
 894  			"status":  1,
 895  			"message": "Failed to create temporary file: " + err.Error(),
 896  		})
 897  		return
 898  	}
 899  	defer os.Remove(tmpFile.Name())
 900  
 901  	// Write YAML content to temp file
 902  	if _, err := tmpFile.WriteString(req.Content); err != nil {
 903  		tmpFile.Close()
 904  		c.JSON(http.StatusInternalServerError, gin.H{
 905  			"status":  1,
 906  			"message": "Failed to write config file: " + err.Error(),
 907  		})
 908  		return
 909  	}
 910  	tmpFile.Close()
 911  
 912  	// Run Python prompt test script using uv
 913  	agentScanDir, err := utils.ResolveAgentScanDir()
 914  	if err != nil {
 915  		c.JSON(http.StatusInternalServerError, gin.H{
 916  			"status":  1,
 917  			"message": "Failed to resolve agent-scan directory: " + err.Error(),
 918  		})
 919  		return
 920  	}
 921  	uvBin, err := utils.ResolveUvBin()
 922  	if err != nil {
 923  		c.JSON(http.StatusInternalServerError, gin.H{
 924  			"status":  1,
 925  			"message": "Failed to resolve uv binary: " + err.Error(),
 926  		})
 927  		return
 928  	}
 929  	var lastLine string
 930  	err = utils.RunCmd(
 931  		agentScanDir,
 932  		uvBin,
 933  		[]string{"run", "test_client_connect.py", "--client_file", tmpFile.Name(), "--prompt", req.Prompt},
 934  		func(line string) {
 935  			lastLine += line
 936  		},
 937  	)
 938  
 939  	if err != nil {
 940  		c.JSON(http.StatusInternalServerError, gin.H{
 941  			"status":  1,
 942  			"message": "Failed to run prompt test: " + err.Error(),
 943  		})
 944  		return
 945  	}
 946  	gologger.Infof("prompt test result: %s", lastLine)
 947  
 948  	// Parse the JSON output from Python script
 949  	var result ConnectResultUpdate
 950  	if err := json.Unmarshal([]byte(lastLine), &result); err != nil {
 951  		c.JSON(http.StatusInternalServerError, gin.H{
 952  			"status":  1,
 953  			"message": "Failed to parse result: " + err.Error(),
 954  		})
 955  		return
 956  	}
 957  
 958  	// Return result based on prompt test outcome
 959  	if result.Content.Success {
 960  		// Extract output from provider_response
 961  		var output string
 962  		if result.Content.ProviderResponse != nil {
 963  			if result.Content.ProviderResponse.Output != nil && *result.Content.ProviderResponse.Output != "" {
 964  				output = *result.Content.ProviderResponse.Output
 965  			} else if result.Content.ProviderResponse.Raw != nil {
 966  				// Fallback to raw response
 967  				rawBytes, _ := json.Marshal(result.Content.ProviderResponse.Raw)
 968  				output = string(rawBytes)
 969  			}
 970  		}
 971  		if output == "" {
 972  			output = result.Content.Message
 973  		}
 974  		c.JSON(http.StatusOK, gin.H{
 975  			"status":  0,
 976  			"message": output,
 977  		})
 978  	} else {
 979  		c.JSON(http.StatusOK, gin.H{
 980  			"status":  1,
 981  			"message": result.Content.Message,
 982  		})
 983  	}
 984  }
 985  
 986  func HandleAgentTemplate(c *gin.Context) {
 987  	enConfig := "agent-scan/config/provider_config_en.json"
 988  	zhConfig := "agent-scan/config/provider_config_zh.json"
 989  	language := c.DefaultQuery("language", "zh")
 990  	var data []byte
 991  	var err error
 992  	if language == "zh" {
 993  		data, err = os.ReadFile(zhConfig)
 994  		if err != nil {
 995  			gologger.WithError(err).Errorln("read zh config")
 996  		}
 997  	} else {
 998  		data, err = os.ReadFile(enConfig)
 999  		if err != nil {
1000  			gologger.WithError(err).Errorln("read en config")
1001  		}
1002  	}
1003  	c.Data(http.StatusOK, "application/json", data)
1004  }