yao/openapi/agent/assistant.go
Max 0f17db74ca Implement GetAssistant endpoint with permission verification
- Added a new GetAssistant function to retrieve assistant details by ID, incorporating permission checks to ensure authorized access.
- Updated the assistant route to use the new GetAssistant handler, enhancing security by verifying user permissions before returning assistant data.
- Introduced a FilterBuiltInAssistant function to clear sensitive fields for built-in assistants, ensuring compliance with data protection standards.
- Enhanced test coverage for the GetAssistant functionality, including tests for successful retrieval, permission filtering, and error handling scenarios.
2025-11-08 13:46:29 +08:00

368 lines
11 KiB
Go

package agent
import (
"strconv"
"strings"
"github.com/gin-gonic/gin"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/agent"
agenttypes "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/openapi/oauth/authorized"
"github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/openapi/response"
)
// ListAssistants lists assistants with pagination and filtering
func ListAssistants(c *gin.Context) {
// Get authorized information
authInfo := authorized.GetInfo(c)
// Get Agent instance from global variable
agentInstance := agent.GetAgent()
if agentInstance == nil || agentInstance.Store == nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Agent store not initialized",
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Parse pagination parameters
page := 1
if pageStr := c.Query("page"); pageStr != "" {
if p, err := strconv.Atoi(pageStr); err == nil && p > 0 {
page = p
}
}
pagesize := 20
if pagesizeStr := c.Query("pagesize"); pagesizeStr != "" {
if ps, err := strconv.Atoi(pagesizeStr); err == nil && ps > 0 && ps <= 100 {
pagesize = ps
}
}
// Validate pagination
if err := ValidatePagination(page, pagesize); err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrInvalidRequest.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusBadRequest, errorResp)
return
}
// Parse select parameter
var selectFields []string
if selectParam := strings.TrimSpace(c.Query("select")); selectParam != "" {
requestedFields := strings.Split(selectParam, ",")
for _, field := range requestedFields {
field = strings.TrimSpace(field)
if field != "" && availableAssistantFields[field] {
selectFields = append(selectFields, field)
}
}
// If no valid fields found, use default
if len(selectFields) == 0 {
selectFields = defaultAssistantFields
}
} else {
selectFields = defaultAssistantFields
}
// Parse filter parameters
keywords := strings.TrimSpace(c.Query("keywords"))
typeParam := strings.TrimSpace(c.Query("type"))
if typeParam == "" {
typeParam = "assistant" // Default type
}
connector := strings.TrimSpace(c.Query("connector"))
assistantID := strings.TrimSpace(c.Query("assistant_id"))
// Parse assistant IDs (multiple)
var assistantIDs []string
if assistantIDsParam := c.Query("assistant_ids"); assistantIDsParam != "" {
assistantIDs = strings.Split(assistantIDsParam, ",")
// Trim spaces
for i, id := range assistantIDs {
assistantIDs[i] = strings.TrimSpace(id)
}
}
// Parse tags
var tags []string
if tagsParam := c.Query("tags"); tagsParam != "" {
tags = strings.Split(tagsParam, ",")
// Trim spaces
for i, tag := range tags {
tags[i] = strings.TrimSpace(tag)
}
}
// Parse boolean filters
var builtIn, mentionable, automated *bool
if builtInParam := c.Query("built_in"); builtInParam != "" {
builtIn = parseBoolValue(builtInParam)
}
if mentionableParam := c.Query("mentionable"); mentionableParam != "" {
mentionable = parseBoolValue(mentionableParam)
}
if automatedParam := c.Query("automated"); automatedParam != "" {
automated = parseBoolValue(automatedParam)
}
// Note: public and share filters are not yet supported in AssistantFilter
// They would need to be added to the store layer for proper filtering
// Parse locale
locale := "en-us" // Default locale
if loc := c.Query("locale"); loc != "" {
locale = strings.ToLower(strings.TrimSpace(loc))
}
// Build filter using the existing AssistantFilter structure
filter := BuildAssistantFilter(AssistantFilterParams{
Page: page,
PageSize: pagesize,
Keywords: keywords,
Type: typeParam,
Connector: connector,
AssistantID: assistantID,
AssistantIDs: assistantIDs,
Tags: tags,
SelectFields: selectFields,
BuiltIn: builtIn,
Mentionable: mentionable,
Automated: automated,
})
// Apply permission-based filtering (Scope filtering)
filter.QueryFilter = AuthQueryFilter(c, authInfo)
// Use the existing GetAssistants method from agent.Store
result, err := agentInstance.Store.GetAssistants(filter, locale)
if err != nil {
log.Error("Failed to list assistants: %v", err)
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Failed to list assistants: " + err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Filter sensitive fields for built-in assistants
// For built-in assistants, clear code-level fields (prompts, workflow, tools, kb, mcp, options)
FilterBuiltInFields(result.Data)
// Return the result with standard response format
response.RespondWithSuccess(c, response.StatusOK, result)
}
// GetAssistant retrieves a single assistant by ID with permission verification
func GetAssistant(c *gin.Context) {
// Get authorized information
authInfo := authorized.GetInfo(c)
// Get Agent instance from global variable
agentInstance := agent.GetAgent()
if agentInstance == nil || agentInstance.Store == nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Agent store not initialized",
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Get assistant ID from URL parameter
assistantID := c.Param("id")
if assistantID == "" {
errorResp := &response.ErrorResponse{
Code: response.ErrInvalidRequest.Code,
ErrorDescription: "assistant_id is required",
}
response.RespondWithError(c, response.StatusBadRequest, errorResp)
return
}
// Parse locale (optional - if not provided, returns raw data without i18n translation)
// This is useful for form editing scenarios where you need the original values
var assistant *agenttypes.AssistantModel
var err error
if loc := c.Query("locale"); loc != "" {
// If locale is specified, get assistant with translation
locale := strings.ToLower(strings.TrimSpace(loc))
assistant, err = agentInstance.Store.GetAssistant(assistantID, locale)
} else {
// If no locale specified, get raw data without translation
assistant, err = agentInstance.Store.GetAssistant(assistantID)
}
if err != nil {
log.Error("Failed to get assistant %s: %v", assistantID, err)
errorResp := &response.ErrorResponse{
Code: response.ErrInvalidRequest.Code,
ErrorDescription: "Assistant not found: " + err.Error(),
}
response.RespondWithError(c, response.StatusNotFound, errorResp)
return
}
// Check permission
hasPermission, err := checkAssistantPermission(authInfo, assistant)
if err != nil {
log.Error("Failed to check permission for assistant %s: %v", assistantID, err)
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Failed to check permission: " + err.Error(),
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
if !hasPermission {
errorResp := &response.ErrorResponse{
Code: response.ErrAccessDenied.Code,
ErrorDescription: "Forbidden: No permission to access this assistant",
}
response.RespondWithError(c, response.StatusForbidden, errorResp)
return
}
// Filter sensitive fields for built-in assistants
FilterBuiltInAssistant(assistant)
// Return the result with standard response format
response.RespondWithSuccess(c, response.StatusOK, map[string]interface{}{
"data": assistant,
})
}
// ListAssistantTags lists assistant tags with permission-based filtering
func ListAssistantTags(c *gin.Context) {
// Get authorized information
authInfo := authorized.GetInfo(c)
// Get Agent instance from global variable
agentInstance := agent.GetAgent()
if agentInstance == nil || agentInstance.Store == nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Agent store not initialized",
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Parse locale
locale := "en-us" // Default locale
if loc := c.Query("locale"); loc != "" {
locale = strings.ToLower(strings.TrimSpace(loc))
}
// Parse filter parameters
typeParam := strings.TrimSpace(c.Query("type"))
if typeParam == "" {
typeParam = "assistant" // Default type
}
connector := strings.TrimSpace(c.Query("connector"))
keywords := strings.TrimSpace(c.Query("keywords"))
// Parse boolean filters
var builtIn, mentionable, automated *bool
if builtInParam := c.Query("built_in"); builtInParam != "" {
builtIn = parseBoolValue(builtInParam)
}
if mentionableParam := c.Query("mentionable"); mentionableParam != "" {
mentionable = parseBoolValue(mentionableParam)
}
if automatedParam := c.Query("automated"); automatedParam != "" {
automated = parseBoolValue(automatedParam)
}
// Build filter
filter := BuildAssistantFilter(AssistantFilterParams{
Type: typeParam,
Connector: connector,
Keywords: keywords,
BuiltIn: builtIn,
Mentionable: mentionable,
Automated: automated,
})
// Apply permission-based filtering (Scope filtering)
filter.QueryFilter = AuthQueryFilter(c, authInfo)
// Get tags with filter
tags, err := agentInstance.Store.GetAssistantTags(filter, locale)
if err != nil {
log.Error("Failed to get assistant tags: %v", err)
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Failed to get assistant tags: " + err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Return the result with standard response format
response.RespondWithSuccess(c, response.StatusOK, map[string]interface{}{
"data": tags,
})
}
// checkAssistantPermission checks if the user has permission to access the assistant
// Similar logic to checkFilePermission in openapi/file/file.go
func checkAssistantPermission(authInfo *types.AuthorizedInfo, assistant *agenttypes.AssistantModel) (bool, error) {
// No auth info, allow access
if authInfo == nil {
return true, nil
}
// No constraints, allow access
if !authInfo.Constraints.TeamOnly && !authInfo.Constraints.OwnerOnly {
return true, nil
}
// If assistant is public, allow access
if assistant.Public {
return true, nil
}
// Combined Team and Owner permission validation
if authInfo.Constraints.TeamOnly && authInfo.Constraints.OwnerOnly {
if assistant.YaoCreatedBy == authInfo.UserID && assistant.YaoTeamID == authInfo.TeamID {
return true, nil
}
return false, nil
}
// Team only permission validation
if authInfo.Constraints.TeamOnly {
// Check if assistant belongs to the same team
if assistant.YaoTeamID == authInfo.TeamID {
// Allow if user created it or if it's shared with team
if assistant.YaoCreatedBy == authInfo.UserID || assistant.Share == "team" {
return true, nil
}
}
return false, nil
}
// Owner only permission validation
if authInfo.Constraints.OwnerOnly {
// Check if user created the assistant and team_id is empty (not a team resource)
if assistant.YaoCreatedBy == authInfo.UserID && assistant.YaoTeamID == "" {
return true, nil
}
return false, nil
}
return false, nil
}