Merge pull request #1481 from trheyi/main

Refactor delivery/OAuth handling, assistant data & MCP tags
This commit is contained in:
Max 2026-03-02 15:07:12 +08:00 committed by GitHub
commit 59435f3d6a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 1221 additions and 476 deletions

View file

@ -146,6 +146,7 @@ func (ast *Assistant) Map() map[string]interface{} {
"locales": ast.Locales,
"uses": ast.Uses,
"search": ast.Search,
"dependencies": ast.Dependencies,
"created_at": store.ToMySQLTime(ast.CreatedAt),
"updated_at": store.ToMySQLTime(ast.UpdatedAt),
}
@ -455,6 +456,14 @@ func (ast *Assistant) Clone() *Assistant {
}
}
// Deep copy dependencies
if ast.Dependencies != nil {
clone.Dependencies = make(map[string]string, len(ast.Dependencies))
for k, v := range ast.Dependencies {
clone.Dependencies[k] = v
}
}
return clone
}
@ -639,6 +648,26 @@ func (ast *Assistant) Update(data map[string]interface{}) error {
ast.Search = search
}
// Dependencies
if v, has := data["dependencies"]; has {
if v == nil {
ast.Dependencies = nil
} else {
switch d := v.(type) {
case map[string]string:
ast.Dependencies = d
case map[string]interface{}:
deps := make(map[string]string, len(d))
for k, val := range d {
if s, ok := val.(string); ok {
deps[k] = s
}
}
ast.Dependencies = deps
}
}
}
return ast.Validate()
}

View file

@ -4,6 +4,7 @@ import (
"fmt"
"reflect"
jsoniter "github.com/json-iterator/go"
agentcontext "github.com/yaoapp/yao/agent/context"
storetypes "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/trace/types"
@ -19,6 +20,18 @@ type HistoryResult struct {
FullMessages []agentcontext.Message // Full messages (history + clean input)
}
// getHistorySize returns the history size with priority: opts.HistorySize > storeSetting.MaxSize > default (20)
func getHistorySize(opts *agentcontext.Options) int {
const defaultHistorySize = 20
if opts != nil && opts.HistorySize > 0 {
return opts.HistorySize
}
if setting := GetStoreSetting(); setting != nil && setting.MaxSize > 0 {
return setting.MaxSize
}
return defaultHistorySize
}
// WithHistory merges the input messages with chat history and traces it
// Returns HistoryResult containing:
// - InputMessages: cleaned input (overlap removed)
@ -41,14 +54,11 @@ func (ast *Assistant) WithHistory(ctx *agentcontext.Context, input []agentcontex
return result, nil
}
// Get MaxSize from store setting
maxSize := 20 // default
if storeSetting := GetStoreSetting(); storeSetting != nil && storeSetting.MaxSize > 0 {
maxSize = storeSetting.MaxSize
}
// Resolve history size: opts.HistorySize > storeSetting.MaxSize > default (20)
maxSize := getHistorySize(opts)
// Load history from store
historyMessages, err := ast.loadHistory(ctx)
historyMessages, err := ast.loadHistory(ctx, maxSize)
if err != nil {
// Log warning but continue without history
ctx.Logger.Warn("Failed to load history for chat=%s: %v", ctx.ChatID, err)
@ -102,8 +112,8 @@ func (ast *Assistant) WithHistory(ctx *agentcontext.Context, input []agentcontex
}
// loadHistory loads chat history from the store
// Returns the most recent MaxSize messages, ordered by time (oldest first)
func (ast *Assistant) loadHistory(ctx *agentcontext.Context) ([]agentcontext.Message, error) {
// Returns the most recent maxSize messages, ordered by time (oldest first)
func (ast *Assistant) loadHistory(ctx *agentcontext.Context, maxSize int) ([]agentcontext.Message, error) {
// Check if chat ID is available
if ctx.ChatID == "" {
return nil, nil
@ -115,13 +125,6 @@ func (ast *Assistant) loadHistory(ctx *agentcontext.Context) ([]agentcontext.Mes
return nil, nil
}
// Get store setting for MaxSize
setting := GetStoreSetting()
maxSize := 20 // default
if setting != nil && setting.MaxSize > 0 {
maxSize = setting.MaxSize
}
// Load messages from store with limit
filter := storetypes.MessageFilter{
Limit: maxSize,
@ -161,11 +164,16 @@ func (ast *Assistant) convertStoreMessageToContext(msg *storetypes.Message) *age
return nil
}
// Skip internal message types that should not be included in LLM context
// These types are for UI/internal use only and can confuse the LLM
// Note: "error" is kept so LLM can help troubleshoot issues
// Handle special message types:
// - tool_call/action: convert to historical summary text for LLM context
// - loading/event: skip (pure UI/lifecycle signals, no semantic value)
// - error: kept as-is so LLM can help troubleshoot issues
switch msg.Type {
case "tool_call", "loading", "action", "event":
case "tool_call":
return ast.convertToolCallToContext(msg)
case "action":
return ast.convertActionToContext(msg)
case "loading", "event":
return nil
}
@ -224,6 +232,137 @@ func (ast *Assistant) extractContentFromProps(props map[string]interface{}, msgT
return nil
}
// convertToolCallToContext converts a tool_call store message to a historical summary text message.
// This allows the LLM to understand what tools were previously called without re-invoking them.
//
// Supports two Props formats:
// - Standard ToolCallProps: {"name": "tool_name", "arguments": "{...}"}
// - Raw stream chunks: {"content": "[{\"index\":0,\"id\":\"call_...\",\"function\":{\"name\":\"tool\"}}][...]"}
func (ast *Assistant) convertToolCallToContext(msg *storetypes.Message) *agentcontext.Message {
if msg.Props == nil {
return nil
}
// Try standard ToolCallProps format first
if name, ok := msg.Props["name"].(string); ok && name != "" {
args, _ := msg.Props["arguments"].(string)
const maxArgsLen = 500
if len(args) > maxArgsLen {
args = args[:maxArgsLen] + "..."
}
return &agentcontext.Message{
Role: agentcontext.RoleAssistant,
Content: fmt.Sprintf("[Historical Tool Call Summary] Called tool \"%s\" with arguments: %s", name, args),
}
}
// Try raw stream chunk format: {"content": "[...][...]..."}
// Each chunk is a JSON array like [{"index":0,"id":"call_...","function":{"name":"echo__ping"}}]
// Subsequent chunks append arguments: [{"index":0,"function":{"arguments":"..."}}]
if raw, ok := msg.Props["content"].(string); ok && raw != "" {
name, args := parseToolCallRawChunks(raw)
if name == "" {
return nil
}
const maxArgsLen = 500
if len(args) > maxArgsLen {
args = args[:maxArgsLen] + "..."
}
return &agentcontext.Message{
Role: agentcontext.RoleAssistant,
Content: fmt.Sprintf("[Historical Tool Call Summary] Called tool \"%s\" with arguments: %s", name, args),
}
}
return nil
}
// parseToolCallRawChunks parses concatenated raw stream chunks to extract tool name and arguments.
// Input format: "[{...}][{...}][{...}]" — multiple JSON arrays concatenated without separator.
func parseToolCallRawChunks(raw string) (name, args string) {
// Split concatenated JSON arrays: "][" is the boundary
// e.g. "[{...}][{...}]" → ["[{...}]", "[{...}]"]
chunks := splitJSONArrays(raw)
var argParts []string
for _, chunk := range chunks {
var items []map[string]interface{}
if err := jsoniter.UnmarshalFromString(chunk, &items); err != nil || len(items) == 0 {
continue
}
item := items[0]
if fn, ok := item["function"].(map[string]interface{}); ok {
if n, ok := fn["name"].(string); ok && n != "" && name == "" {
name = n
}
if a, ok := fn["arguments"].(string); ok && a != "" {
argParts = append(argParts, a)
}
}
}
args = ""
for _, part := range argParts {
args += part
}
return name, args
}
// splitJSONArrays splits a string of concatenated JSON arrays "[...][...][...]" into individual arrays.
func splitJSONArrays(s string) []string {
var result []string
depth := 0
start := -1
for i, ch := range s {
switch ch {
case '[':
if depth == 0 {
start = i
}
depth++
case ']':
depth--
if depth == 0 && start >= 0 {
result = append(result, s[start:i+1])
start = -1
}
}
}
return result
}
// convertActionToContext converts an action store message to a historical summary text message.
// This allows the LLM to understand what system actions were previously executed.
func (ast *Assistant) convertActionToContext(msg *storetypes.Message) *agentcontext.Message {
if msg.Props == nil {
return nil
}
name, _ := msg.Props["name"].(string)
if name == "" {
return nil
}
payload := ""
if msg.Props["payload"] != nil {
if payloadStr, err := jsoniter.MarshalToString(msg.Props["payload"]); err == nil {
const maxPayloadLen = 500
if len(payloadStr) > maxPayloadLen {
payloadStr = payloadStr[:maxPayloadLen] + "..."
}
payload = payloadStr
}
}
if payload != "" {
return &agentcontext.Message{
Role: agentcontext.RoleAssistant,
Content: fmt.Sprintf("[Historical Action Summary] Executed action \"%s\" with payload: %s", name, payload),
}
}
return &agentcontext.Message{
Role: agentcontext.RoleAssistant,
Content: fmt.Sprintf("[Historical Action Summary] Executed action \"%s\"", name),
}
}
// findOverlapIndex finds the index in input where history messages end
// Returns the number of input messages that overlap with history
func (ast *Assistant) findOverlapIndex(history, input []agentcontext.Message) int {

View file

@ -415,7 +415,10 @@ func TestHistoryLoading(t *testing.T) {
}()
// Add various message types (only user/assistant roles allowed by DB constraint)
// loadHistory filters by role (user/assistant only) and converts based on type
// loadHistory filters by role (user/assistant only) and converts based on type:
// - loading/event: skipped (no semantic value)
// - tool_call/action: converted to historical summary text
// - text/user_input/error: kept as-is
err = chatStore.SaveMessages(chatID, []*storetypes.Message{
{
MessageID: fmt.Sprintf("filter_1_%s", reqID),
@ -461,11 +464,9 @@ func TestHistoryLoading(t *testing.T) {
require.NoError(t, err)
require.NotNil(t, result)
// User and assistant roles are included
// loading type messages with role=assistant are included (role-based filtering)
// History contains: 1 user + 2 assistant (loading + text) = 3 messages
// Plus 1 new input = 4 total
assert.GreaterOrEqual(t, len(result.FullMessages), 3) // At least 1 user + 1 assistant from history + 1 new
// loading type is skipped (no semantic value)
// History: user_input + text = 2 messages; plus 1 new input = 3 total
assert.Len(t, result.FullMessages, 3)
// Verify only user and assistant roles
for _, msg := range result.FullMessages {
@ -473,7 +474,217 @@ func TestHistoryLoading(t *testing.T) {
"Expected user or assistant role, got: %s", msg.Role)
}
t.Log("✓ Only user and assistant roles included in history")
t.Log("✓ Loading type filtered, user/assistant roles kept")
})
t.Run("ToolCallConvertedToSummary", func(t *testing.T) {
chatID := fmt.Sprintf("test_toolcall_%s", uuid.New().String()[:8])
ctx := newHistoryTestContext(chatID)
reqID := uuid.New().String()[:8]
err := chatStore.CreateChat(&storetypes.Chat{
ChatID: chatID,
AssistantID: ast.ID,
Status: "active",
Share: "private",
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
})
require.NoError(t, err)
defer func() {
chatStore.DeleteMessages(chatID, nil)
chatStore.DeleteChat(chatID)
}()
// Add tool_call messages in both formats
err = chatStore.SaveMessages(chatID, []*storetypes.Message{
{
MessageID: fmt.Sprintf("tc_1_%s", reqID),
ChatID: chatID,
RequestID: fmt.Sprintf("req_tc_%s", reqID),
Role: "user",
Type: "user_input",
Props: map[string]interface{}{"content": "echo 3 ping 4"},
Sequence: 1,
AssistantID: ast.ID,
CreatedAt: time.Now().Add(-3 * time.Minute),
},
// Raw stream chunk format (actual DB format)
{
MessageID: fmt.Sprintf("tc_2_%s", reqID),
ChatID: chatID,
RequestID: fmt.Sprintf("req_tc_%s", reqID),
Role: "assistant",
Type: "tool_call",
Props: map[string]interface{}{"content": `[{"index":0,"id":"call_abc","type":"function","function":{"name":"echo__ping"}}][{"index":0,"function":{"arguments":"{\"count\":3}"}}]`},
Sequence: 2,
AssistantID: ast.ID,
CreatedAt: time.Now().Add(-2 * time.Minute),
},
// Standard ToolCallProps format
{
MessageID: fmt.Sprintf("tc_3_%s", reqID),
ChatID: chatID,
RequestID: fmt.Sprintf("req_tc_%s", reqID),
Role: "assistant",
Type: "tool_call",
Props: map[string]interface{}{"name": "echo__echo", "arguments": `{"message":"hello"}`},
Sequence: 3,
AssistantID: ast.ID,
CreatedAt: time.Now().Add(-1 * time.Minute),
},
})
require.NoError(t, err)
input := []agentcontext.Message{
{Role: agentcontext.RoleUser, Content: "echo 5 ping 6"},
}
result, err := ast.WithHistory(ctx, input, nil)
require.NoError(t, err)
require.NotNil(t, result)
// 1 user_input + 2 tool_call summaries + 1 new input = 4
assert.Len(t, result.FullMessages, 4)
// Verify tool_call messages are converted to summary text
tcMsg1 := result.FullMessages[1]
assert.Equal(t, agentcontext.RoleAssistant, tcMsg1.Role)
assert.Contains(t, tcMsg1.Content, "[Historical Tool Call Summary]")
assert.Contains(t, tcMsg1.Content, "echo__ping")
assert.Contains(t, tcMsg1.Content, `{"count":3}`)
tcMsg2 := result.FullMessages[2]
assert.Equal(t, agentcontext.RoleAssistant, tcMsg2.Role)
assert.Contains(t, tcMsg2.Content, "[Historical Tool Call Summary]")
assert.Contains(t, tcMsg2.Content, "echo__echo")
assert.Contains(t, tcMsg2.Content, `{"message":"hello"}`)
t.Log("✓ Tool call messages converted to historical summaries")
})
t.Run("ActionConvertedToSummary", func(t *testing.T) {
chatID := fmt.Sprintf("test_action_%s", uuid.New().String()[:8])
ctx := newHistoryTestContext(chatID)
reqID := uuid.New().String()[:8]
err := chatStore.CreateChat(&storetypes.Chat{
ChatID: chatID,
AssistantID: ast.ID,
Status: "active",
Share: "private",
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
})
require.NoError(t, err)
defer func() {
chatStore.DeleteMessages(chatID, nil)
chatStore.DeleteChat(chatID)
}()
err = chatStore.SaveMessages(chatID, []*storetypes.Message{
{
MessageID: fmt.Sprintf("act_1_%s", reqID),
ChatID: chatID,
RequestID: fmt.Sprintf("req_act_%s", reqID),
Role: "user",
Type: "user_input",
Props: map[string]interface{}{"content": "Do something"},
Sequence: 1,
AssistantID: ast.ID,
CreatedAt: time.Now().Add(-2 * time.Minute),
},
// Action with payload
{
MessageID: fmt.Sprintf("act_2_%s", reqID),
ChatID: chatID,
RequestID: fmt.Sprintf("req_act_%s", reqID),
Role: "assistant",
Type: "action",
Props: map[string]interface{}{
"name": "robot.execute",
"payload": map[string]interface{}{"goals": "test goal", "robot_id": "12345"},
},
Sequence: 2,
AssistantID: ast.ID,
CreatedAt: time.Now().Add(-1 * time.Minute),
},
})
require.NoError(t, err)
input := []agentcontext.Message{
{Role: agentcontext.RoleUser, Content: "What happened?"},
}
result, err := ast.WithHistory(ctx, input, nil)
require.NoError(t, err)
require.NotNil(t, result)
// 1 user_input + 1 action summary + 1 new input = 3
assert.Len(t, result.FullMessages, 3)
actMsg := result.FullMessages[1]
assert.Equal(t, agentcontext.RoleAssistant, actMsg.Role)
assert.Contains(t, actMsg.Content, "[Historical Action Summary]")
assert.Contains(t, actMsg.Content, "robot.execute")
assert.Contains(t, actMsg.Content, "test goal")
assert.Contains(t, actMsg.Content, "12345")
t.Log("✓ Action messages converted to historical summaries with payload")
})
t.Run("ActionWithoutPayload", func(t *testing.T) {
chatID := fmt.Sprintf("test_action_nopay_%s", uuid.New().String()[:8])
ctx := newHistoryTestContext(chatID)
reqID := uuid.New().String()[:8]
err := chatStore.CreateChat(&storetypes.Chat{
ChatID: chatID,
AssistantID: ast.ID,
Status: "active",
Share: "private",
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
})
require.NoError(t, err)
defer func() {
chatStore.DeleteMessages(chatID, nil)
chatStore.DeleteChat(chatID)
}()
err = chatStore.SaveMessages(chatID, []*storetypes.Message{
{
MessageID: fmt.Sprintf("actnp_1_%s", reqID),
ChatID: chatID,
RequestID: fmt.Sprintf("req_actnp_%s", reqID),
Role: "assistant",
Type: "action",
Props: map[string]interface{}{"name": "navigate"},
Sequence: 1,
AssistantID: ast.ID,
CreatedAt: time.Now(),
},
})
require.NoError(t, err)
input := []agentcontext.Message{
{Role: agentcontext.RoleUser, Content: "What happened?"},
}
result, err := ast.WithHistory(ctx, input, nil)
require.NoError(t, err)
require.NotNil(t, result)
// 1 action summary + 1 new input = 2
assert.Len(t, result.FullMessages, 2)
actMsg := result.FullMessages[0]
assert.Equal(t, agentcontext.RoleAssistant, actMsg.Role)
assert.Contains(t, actMsg.Content, "[Historical Action Summary]")
assert.Contains(t, actMsg.Content, "navigate")
assert.NotContains(t, actMsg.Content, "payload")
t.Log("✓ Action without payload handled correctly")
})
t.Run("ContentExtraction", func(t *testing.T) {

View file

@ -729,6 +729,30 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
assistant.Sandbox = sb
}
// dependencies (name -> version constraint, like npm dependencies)
if deps, has := data["dependencies"]; has {
switch v := deps.(type) {
case map[string]string:
assistant.Dependencies = v
case map[string]interface{}:
d := make(map[string]string, len(v))
for k, val := range v {
d[k] = cast.ToString(val)
}
assistant.Dependencies = d
default:
raw, err := jsoniter.Marshal(v)
if err != nil {
return nil, err
}
var d map[string]string
if err := jsoniter.Unmarshal(raw, &d); err != nil {
return nil, err
}
assistant.Dependencies = d
}
}
// uses (wrapper configurations for vision, audio, etc.)
// Merge hierarchy: global uses < assistant uses
if uses, has := data["uses"]; has {

View file

@ -494,6 +494,10 @@ func TestLoadStoreWithAllFields(t *testing.T) {
Description: "This is a test placeholder",
Prompts: []string{"Test prompt 1", "Test prompt 2"},
},
Dependencies: map[string]string{
"echo": "^1.0.0",
"customer": ">=2.0.0",
},
Source: `
// @ts-nocheck
function Create(ctx: any, messages: any[]): any {
@ -571,6 +575,12 @@ function Create(ctx: any, messages: any[]): any {
assert.NotNil(t, loaded.HookScript)
assert.NotEmpty(t, loaded.Source)
// Dependencies
require.NotNil(t, loaded.Dependencies)
assert.Len(t, loaded.Dependencies, 2)
assert.Equal(t, "^1.0.0", loaded.Dependencies["echo"])
assert.Equal(t, ">=2.0.0", loaded.Dependencies["customer"])
// Execute the Create hook to verify it works
ctx := newStoreTestContext("test-chat-all-fields", assistantID)
messages := []context.Message{{Role: "user", Content: "Test message"}}

View file

@ -193,6 +193,18 @@ func TestLoadPath(t *testing.T) {
assert.NotNil(t, zhLocale)
})
t.Run("LoadDependencies", func(t *testing.T) {
assistant, err := assistant.LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
require.NotNil(t, assistant)
// Dependencies
assert.NotNil(t, assistant.Dependencies)
assert.Len(t, assistant.Dependencies, 2)
assert.Equal(t, "^1.0.0", assistant.Dependencies["echo"])
assert.Equal(t, ">=2.0.0", assistant.Dependencies["customer"])
})
t.Run("LoadNonExistentAssistant", func(t *testing.T) {
_, err := assistant.LoadPath("/assistants/non-existent")
assert.Error(t, err)
@ -333,6 +345,13 @@ func TestClone(t *testing.T) {
assert.False(t, exists, "Clone should not have modified key")
delete(original.Options, "test_key") // cleanup
}
if original.Dependencies != nil {
original.Dependencies["test_dep"] = "^9.9.9"
_, exists := clone.Dependencies["test_dep"]
assert.False(t, exists, "Clone dependencies should not have modified key")
delete(original.Dependencies, "test_dep") // cleanup
}
})
t.Run("CloneNil", func(t *testing.T) {
@ -457,6 +476,7 @@ func TestMap(t *testing.T) {
assert.Equal(t, assistant.ConnectorOptions, m["connector_options"])
assert.Equal(t, assistant.PromptPresets, m["prompt_presets"])
assert.Equal(t, assistant.Source, m["source"])
assert.Equal(t, assistant.Dependencies, m["dependencies"])
}
// TestLoadSystemAgents tests loading system agents from bindata

View file

@ -319,6 +319,11 @@ type Options struct {
// Metadata for passing custom data to hooks (e.g., scenario selection)
Metadata map[string]any `json:"metadata,omitempty"` // Custom metadata passed to Create/Next hooks
// HistorySize controls the max number of history messages loaded for LLM context.
// Priority: HistorySize > StoreSetting.MaxSize > default (20)
// 0 means use StoreSetting or default.
HistorySize int `json:"history_size,omitempty"`
// OnMessage is called for each message sent via ctx.Send()
// Used by ctx.agent.Call with onChunk callback to receive SSE messages
// Returns: 0 = continue, non-zero = stop

View file

@ -91,25 +91,27 @@ func (h *robotHandler) handleDelivery(ctx context.Context, ev *eventtypes.Event,
}
}
// Push delivery to integration channels (Telegram, etc.) via ReplyFunc
if reply := getReplyFunc(); reply != nil {
msg := buildDeliveryMessage(content)
if msg != nil {
channel, chatID := splitChannelChatID(payload.ChatID)
extra := map[string]any{
"member_id": payload.MemberID,
"execution_id": payload.ExecutionID,
}
for k, v := range payload.Extra {
extra[k] = v
}
metadata := &MessageMetadata{
Channel: channel,
ChatID: chatID,
Extra: extra,
}
if err := reply(ctx, msg, metadata); err != nil {
log.Error("delivery handler: integration reply failed execution=%s: %v", payload.ExecutionID, err)
// Push delivery to integration channels only when the task originated from one
if reply := getReplyFunc(); reply != nil && payload.ChatID != "" {
channel, chatID := splitChannelChatID(payload.ChatID)
if channel != "" && chatID != "" {
msg := buildDeliveryMessage(content)
if msg != nil {
extra := map[string]any{
"member_id": payload.MemberID,
"execution_id": payload.ExecutionID,
}
for k, v := range payload.Extra {
extra[k] = v
}
metadata := &MessageMetadata{
Channel: channel,
ChatID: chatID,
Extra: extra,
}
if err := reply(ctx, msg, metadata); err != nil {
log.Error("delivery handler: integration reply failed channel=%s execution=%s: %v", channel, payload.ExecutionID, err)
}
}
}
}

View file

@ -454,6 +454,17 @@ func ToAssistantModel(v interface{}) (*AssistantModel, error) {
}
}
// Dependencies
if deps, ok := data["dependencies"]; ok && deps != nil {
raw, err := jsoniter.Marshal(deps)
if err == nil {
var d map[string]string
if err := jsoniter.Unmarshal(raw, &d); err == nil {
model.Dependencies = d
}
}
}
// Permission fields
if createdBy, ok := data["__yao_created_by"].(string); ok {
model.YaoCreatedBy = createdBy

View file

@ -583,6 +583,10 @@ func TestToAssistantModel(t *testing.T) {
"name": "English Name",
},
},
"dependencies": map[string]interface{}{
"echo": "^1.0.0",
"customer": ">=2.0.0",
},
}
result, err := ToAssistantModel(data)
@ -711,6 +715,19 @@ func TestToAssistantModel(t *testing.T) {
if result.Locales == nil {
t.Error("Expected Locales to be set")
}
if result.Dependencies == nil {
t.Error("Expected Dependencies to be set")
} else {
if len(result.Dependencies) != 2 {
t.Errorf("Expected 2 dependencies, got %d", len(result.Dependencies))
}
if result.Dependencies["echo"] != "^1.0.0" {
t.Errorf("Expected echo dependency '^1.0.0', got '%s'", result.Dependencies["echo"])
}
if result.Dependencies["customer"] != ">=2.0.0" {
t.Errorf("Expected customer dependency '>=2.0.0', got '%s'", result.Dependencies["customer"])
}
}
})
t.Run("MapWithFloatNumbers", func(t *testing.T) {
@ -750,6 +767,7 @@ func TestToAssistantModel(t *testing.T) {
"workflow": nil,
"placeholder": nil,
"locales": nil,
"dependencies": nil,
}
result, err := ToAssistantModel(data)
@ -764,6 +782,9 @@ func TestToAssistantModel(t *testing.T) {
if result.Tags != nil {
t.Error("Expected Tags to be nil")
}
if result.Dependencies != nil {
t.Error("Expected Dependencies to be nil")
}
if result.Modes != nil {
t.Error("Expected Modes to be nil")
}

View file

@ -36,6 +36,7 @@ var AssistantAllowedFields = map[string]bool{
"locales": true,
"uses": true,
"search": true,
"dependencies": true,
"automated": true,
"mentionable": true,
"created_at": true,
@ -66,10 +67,11 @@ var AssistantDefaultFields = []string{
"share",
"automated",
"mentionable",
"sandbox", // Sandbox configuration presence (lightweight)
"kb", // Knowledge base configuration (lightweight)
"db", // Database configuration (lightweight)
"mcp", // MCP servers configuration (lightweight)
"sandbox", // Sandbox configuration presence (lightweight)
"kb", // Knowledge base configuration (lightweight)
"db", // Database configuration (lightweight)
"mcp", // MCP servers configuration (lightweight)
"dependencies", // Dependencies on other MCP Clients (lightweight)
"created_at",
"updated_at",
"__yao_created_by", // Permission: creator user ID
@ -112,6 +114,7 @@ var AssistantFullFields = []string{
"locales",
"uses",
"search",
"dependencies",
"automated",
"mentionable",
"created_at",

View file

@ -454,6 +454,7 @@ type AssistantModel struct {
Locales i18n.Map `json:"locales,omitempty"` // Assistant Locales
Uses *context.Uses `json:"uses,omitempty"` // Assistant-specific wrapper configurations for vision, audio, etc. If not set, use global settings
Search *searchTypes.Config `json:"search,omitempty"` // Search configuration (web, kb, db, citation, weights, etc.)
Dependencies map[string]string `json:"dependencies,omitempty"` // Dependencies on other MCP Clients (name -> version constraint)
CreatedAt int64 `json:"created_at"` // Creation timestamp
UpdatedAt int64 `json:"updated_at"` // Last update timestamp

View file

@ -145,31 +145,6 @@ func (store *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
data["__yao_tenant_id"] = nil
}
// Handle simple types
if assistant.Options != nil {
jsonStr, err := jsoniter.MarshalToString(assistant.Options)
if err != nil {
return "", fmt.Errorf("failed to marshal options: %w", err)
}
data["options"] = jsonStr
}
if assistant.Tags != nil {
jsonStr, err := jsoniter.MarshalToString(assistant.Tags)
if err != nil {
return "", fmt.Errorf("failed to marshal tags: %w", err)
}
data["tags"] = jsonStr
}
if assistant.Modes != nil {
jsonStr, err := jsoniter.MarshalToString(assistant.Modes)
if err != nil {
return "", fmt.Errorf("failed to marshal modes: %w", err)
}
data["modes"] = jsonStr
}
// DefaultMode is a simple string field
if assistant.DefaultMode != "" {
data["default_mode"] = assistant.DefaultMode
@ -177,8 +152,12 @@ func (store *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
data["default_mode"] = nil
}
// Handle interface{} fields - they should already be in the correct format
// Handle all JSON fields uniformly via marshalJSONFields.
// Uses isNil() to correctly skip typed nils stored in interface{}.
jsonFields := map[string]interface{}{
"options": assistant.Options,
"tags": assistant.Tags,
"modes": assistant.Modes,
"prompts": assistant.Prompts,
"prompt_presets": assistant.PromptPresets,
"connector_options": assistant.ConnectorOptions,
@ -191,16 +170,11 @@ func (store *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
"locales": assistant.Locales,
"uses": assistant.Uses,
"search": assistant.Search,
"dependencies": assistant.Dependencies,
}
for field, value := range jsonFields {
if value != nil {
jsonStr, err := jsoniter.MarshalToString(value)
if err != nil {
return "", fmt.Errorf("failed to marshal %s: %w", field, err)
}
data[field] = jsonStr
}
if err := marshalJSONFields(data, jsonFields); err != nil {
return "", err
}
// Update or insert
@ -249,7 +223,7 @@ func (store *Xun) UpdateAssistant(assistantID string, updates map[string]interfa
data := make(map[string]interface{})
// List of fields that need JSON marshaling
jsonFields := []string{"options", "tags", "modes", "prompts", "prompt_presets", "connector_options", "kb", "db", "mcp", "workflow", "sandbox", "placeholder", "locales", "uses", "search"}
jsonFields := []string{"options", "tags", "modes", "prompts", "prompt_presets", "connector_options", "kb", "db", "mcp", "workflow", "sandbox", "placeholder", "locales", "uses", "search", "dependencies"}
jsonFieldSet := make(map[string]bool)
for _, field := range jsonFields {
jsonFieldSet[field] = true
@ -271,14 +245,14 @@ func (store *Xun) UpdateAssistant(assistantID string, updates map[string]interfa
// Handle JSON fields
if jsonFieldSet[key] {
if value != nil {
if isNil(value) {
data[key] = nil
} else {
jsonStr, err := jsoniter.MarshalToString(value)
if err != nil {
return fmt.Errorf("failed to marshal %s: %w", key, err)
}
data[key] = jsonStr
} else {
data[key] = nil
}
} else {
// Handle regular fields
@ -470,7 +444,7 @@ func (store *Xun) GetAssistants(filter types.AssistantFilter, locale ...string)
// Convert rows to types.AssistantModel slice
assistants := make([]*types.AssistantModel, 0, len(rows))
jsonFields := []string{"tags", "options", "prompts", "prompt_presets", "connector_options", "workflow", "sandbox", "kb", "mcp", "placeholder", "locales", "uses", "search"}
jsonFields := []string{"tags", "options", "prompts", "prompt_presets", "connector_options", "workflow", "sandbox", "kb", "mcp", "placeholder", "locales", "uses", "search", "dependencies"}
for _, row := range rows {
data := row.ToMap()
@ -543,7 +517,7 @@ func (store *Xun) GetAssistant(assistantID string, fields []string, locale ...st
}
// Parse JSON fields
jsonFields := []string{"tags", "modes", "options", "prompts", "prompt_presets", "connector_options", "workflow", "sandbox", "kb", "db", "mcp", "placeholder", "locales", "uses", "search"}
jsonFields := []string{"tags", "modes", "options", "prompts", "prompt_presets", "connector_options", "workflow", "sandbox", "kb", "db", "mcp", "placeholder", "locales", "uses", "search", "dependencies"}
store.parseJSONFields(data, jsonFields)
// Convert map to types.AssistantModel
@ -706,6 +680,16 @@ func (store *Xun) GetAssistant(assistantID string, fields []string, locale ...st
}
}
if deps, has := data["dependencies"]; has && deps != nil {
raw, err := jsoniter.Marshal(deps)
if err == nil {
var d map[string]string
if err := jsoniter.Unmarshal(raw, &d); err == nil {
model.Dependencies = d
}
}
}
// Apply i18n translation if locale is provided
if len(locale) > 0 && locale[0] != "" {
store.translate(model, assistantID, locale[0])

View file

@ -2,9 +2,44 @@ package xun
import (
"fmt"
"reflect"
"time"
jsoniter "github.com/json-iterator/go"
)
// isNil checks whether a value is truly nil, handling the Go typed-nil-in-interface pitfall.
// A nil map, slice, or pointer stored in an interface{} is not == nil in Go;
// this helper uses reflect to detect that case.
func isNil(v interface{}) bool {
if v == nil {
return true
}
rv := reflect.ValueOf(v)
switch rv.Kind() {
case reflect.Ptr, reflect.Map, reflect.Slice, reflect.Interface, reflect.Chan, reflect.Func:
return rv.IsNil()
}
return false
}
// marshalJSONFields serialises each value in fields to a JSON string and writes
// it into data. Truly-nil values (including typed nils) are skipped so the
// database column keeps its SQL NULL / default.
func marshalJSONFields(data map[string]interface{}, fields map[string]interface{}) error {
for field, value := range fields {
if isNil(value) {
continue
}
jsonStr, err := jsoniter.MarshalToString(value)
if err != nil {
return fmt.Errorf("failed to marshal %s: %w", field, err)
}
data[field] = jsonStr
}
return nil
}
// Helper functions for type conversion
func getString(data map[string]interface{}, key string) string {
if v, ok := data[key].(string); ok {

View file

@ -0,0 +1,190 @@
package xun
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
type testStruct struct{ Name string }
func TestIsNil(t *testing.T) {
// Untyped nil
t.Run("UntypedNil", func(t *testing.T) {
assert.True(t, isNil(nil))
})
// Typed nil pointer
t.Run("TypedNilPointer", func(t *testing.T) {
var p *testStruct
assert.True(t, isNil(p))
})
// Typed nil map
t.Run("TypedNilMap", func(t *testing.T) {
var m map[string]string
assert.True(t, isNil(m))
})
// Typed nil slice
t.Run("TypedNilSlice", func(t *testing.T) {
var s []string
assert.True(t, isNil(s))
})
// Non-nil pointer
t.Run("NonNilPointer", func(t *testing.T) {
p := &testStruct{Name: "test"}
assert.False(t, isNil(p))
})
// Non-nil map (empty)
t.Run("NonNilEmptyMap", func(t *testing.T) {
m := map[string]string{}
assert.False(t, isNil(m))
})
// Non-nil map with values
t.Run("NonNilMap", func(t *testing.T) {
m := map[string]string{"a": "1"}
assert.False(t, isNil(m))
})
// Non-nil slice (empty)
t.Run("NonNilEmptySlice", func(t *testing.T) {
s := []string{}
assert.False(t, isNil(s))
})
// Non-nil slice with values
t.Run("NonNilSlice", func(t *testing.T) {
s := []string{"a"}
assert.False(t, isNil(s))
})
// Scalar types (never nil)
t.Run("String", func(t *testing.T) {
assert.False(t, isNil("hello"))
})
t.Run("EmptyString", func(t *testing.T) {
assert.False(t, isNil(""))
})
t.Run("Int", func(t *testing.T) {
assert.False(t, isNil(42))
})
t.Run("Bool", func(t *testing.T) {
assert.False(t, isNil(false))
})
}
func TestMarshalJSONFields(t *testing.T) {
t.Run("SkipUntypedNil", func(t *testing.T) {
data := make(map[string]interface{})
err := marshalJSONFields(data, map[string]interface{}{
"field1": nil,
})
require.NoError(t, err)
_, exists := data["field1"]
assert.False(t, exists, "untyped nil should be skipped")
})
t.Run("SkipTypedNilMap", func(t *testing.T) {
data := make(map[string]interface{})
var m map[string]string
err := marshalJSONFields(data, map[string]interface{}{
"deps": m,
})
require.NoError(t, err)
_, exists := data["deps"]
assert.False(t, exists, "typed nil map should be skipped")
})
t.Run("SkipTypedNilSlice", func(t *testing.T) {
data := make(map[string]interface{})
var s []string
err := marshalJSONFields(data, map[string]interface{}{
"tags": s,
})
require.NoError(t, err)
_, exists := data["tags"]
assert.False(t, exists, "typed nil slice should be skipped")
})
t.Run("SkipTypedNilPointer", func(t *testing.T) {
data := make(map[string]interface{})
var p *testStruct
err := marshalJSONFields(data, map[string]interface{}{
"kb": p,
})
require.NoError(t, err)
_, exists := data["kb"]
assert.False(t, exists, "typed nil pointer should be skipped")
})
t.Run("MarshalNonNilMap", func(t *testing.T) {
data := make(map[string]interface{})
err := marshalJSONFields(data, map[string]interface{}{
"deps": map[string]string{"echo": "^1.0.0"},
})
require.NoError(t, err)
assert.Equal(t, `{"echo":"^1.0.0"}`, data["deps"])
})
t.Run("MarshalEmptyMap", func(t *testing.T) {
data := make(map[string]interface{})
err := marshalJSONFields(data, map[string]interface{}{
"deps": map[string]string{},
})
require.NoError(t, err)
assert.Equal(t, `{}`, data["deps"])
})
t.Run("MarshalSlice", func(t *testing.T) {
data := make(map[string]interface{})
err := marshalJSONFields(data, map[string]interface{}{
"tags": []string{"ai", "bot"},
})
require.NoError(t, err)
assert.Equal(t, `["ai","bot"]`, data["tags"])
})
t.Run("MarshalPointer", func(t *testing.T) {
data := make(map[string]interface{})
err := marshalJSONFields(data, map[string]interface{}{
"kb": &testStruct{Name: "test"},
})
require.NoError(t, err)
assert.Equal(t, `{"Name":"test"}`, data["kb"])
})
t.Run("MixedNilAndNonNil", func(t *testing.T) {
data := make(map[string]interface{})
var nilMap map[string]string
var nilSlice []string
var nilPtr *testStruct
err := marshalJSONFields(data, map[string]interface{}{
"nil_map": nilMap,
"nil_slice": nilSlice,
"nil_ptr": nilPtr,
"nil_raw": nil,
"good_map": map[string]string{"k": "v"},
"good_list": []string{"a"},
})
require.NoError(t, err)
assert.Len(t, data, 2, "only non-nil fields should be written")
assert.Equal(t, `{"k":"v"}`, data["good_map"])
assert.Equal(t, `["a"]`, data["good_list"])
_, exists := data["nil_map"]
assert.False(t, exists)
_, exists = data["nil_slice"]
assert.False(t, exists)
_, exists = data["nil_ptr"]
assert.False(t, exists)
_, exists = data["nil_raw"]
assert.False(t, exists)
})
}

File diff suppressed because one or more lines are too long

View file

@ -9,12 +9,13 @@ import (
// Server represents an MCP server option (from user perspective)
type Server struct {
Label string `json:"label"`
Value string `json:"value"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
Transport string `json:"transport,omitempty"` // "stdio", "sse", "http"
Builtin bool `json:"builtin"` // true for system built-in, false for user-defined
Label string `json:"label"`
Value string `json:"value"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
Tags []string `json:"tags,omitempty"`
Transport string `json:"transport,omitempty"` // "stdio", "sse", "http"
Builtin bool `json:"builtin"` // true for system built-in, false for user-defined
}
// Attach attaches the MCP server management handlers to the router with OAuth protection
@ -58,6 +59,7 @@ func listServers(c *gin.Context) {
Value: id,
Name: name,
Description: meta.Description,
Tags: meta.Tags,
Transport: transport,
Builtin: meta.Builtin,
})

View file

@ -1,9 +1,11 @@
package oauth
import (
"errors"
"fmt"
"net/http"
"strings"
"sync"
"time"
"github.com/gin-gonic/gin"
@ -14,6 +16,16 @@ import (
"github.com/yaoapp/yao/openapi/response"
)
var (
errRefreshInProgress = errors.New("refresh in progress")
errRefreshAlreadyDone = errors.New("refresh already done")
refreshGates sync.Map // refreshToken → *refreshGate
)
type refreshGate struct {
done chan struct{} // closed when rotation completes
}
// Guard is the OAuth guard middleware
func (s *Service) Guard(c *gin.Context) {
// Authenticate first (validates token and sets authorized info)
@ -69,12 +81,17 @@ func (s *Service) Authenticate(c *gin.Context) bool {
if !expiredClaims.ExpiresAt.IsZero() && expiredClaims.ExpiresAt.Before(time.Now()) {
newClaims, refreshErr := s.TryRefreshToken(c, expiredClaims)
if refreshErr != nil {
log.Error("[OAuth] Token refresh failed: %v", refreshErr)
response.RespondWithError(c, http.StatusUnauthorized, types.ErrInvalidRefreshToken)
c.Abort()
return false
if errors.Is(refreshErr, errRefreshInProgress) || errors.Is(refreshErr, errRefreshAlreadyDone) {
claims = expiredClaims
} else {
log.Error("[OAuth] Token refresh failed: %v", refreshErr)
response.RespondWithError(c, http.StatusUnauthorized, types.ErrInvalidRefreshToken)
c.Abort()
return false
}
} else {
claims = newClaims
}
claims = newClaims
} else {
response.RespondWithError(c, http.StatusUnauthorized, types.ErrInvalidToken)
c.Abort()
@ -104,6 +121,29 @@ func (s *Service) TryRefreshToken(c *gin.Context, expiredClaims *types.TokenClai
return nil, fmt.Errorf("refresh token missing")
}
gate := &refreshGate{done: make(chan struct{})}
if actual, loaded := refreshGates.LoadOrStore(refreshToken, gate); loaded {
// Another goroutine owns the rotation for this refresh token.
// It may still be running or already finished.
existing := actual.(*refreshGate)
select {
case <-existing.done:
return nil, errRefreshAlreadyDone
default:
return nil, errRefreshInProgress
}
}
// We own the gate — clean up when finished.
defer func() {
close(gate.done)
// Keep the gate in the map for 30 s so late arrivals see "done"
// instead of starting a new rotation with the now-revoked token.
time.AfterFunc(30*time.Second, func() {
refreshGates.CompareAndDelete(refreshToken, gate)
})
}()
refreshClaims, err := s.VerifyRefreshToken(refreshToken)
if err != nil {
return nil, fmt.Errorf("invalid or expired refresh token: %w", err)
@ -244,6 +284,12 @@ func (s *Service) GetRefreshToken(c *gin.Context) string {
return s.getRefreshToken(c)
}
// IsRefreshInProgress checks whether an error signals that another goroutine
// is already rotating (or has just rotated) the same refresh token.
func IsRefreshInProgress(err error) bool {
return errors.Is(err, errRefreshInProgress) || errors.Is(err, errRefreshAlreadyDone)
}
// GetSessionID gets the session ID from the request (public method)
func (s *Service) GetSessionID(c *gin.Context) string {
return s.getSessionID(c)

View file

@ -124,9 +124,14 @@ func guardOAuth(r *Request) error {
!expiredClaims.ExpiresAt.IsZero() && expiredClaims.ExpiresAt.Before(time.Now()) {
refreshed, refreshErr := oauth.OAuth.TryRefreshToken(c, expiredClaims)
if refreshErr != nil {
return fmt.Errorf("Exception|401:Token expired and refresh failed")
if oauth.IsRefreshInProgress(refreshErr) {
claims = expiredClaims
} else {
return fmt.Errorf("Exception|401:Token expired and refresh failed")
}
} else {
claims = refreshed
}
claims = refreshed
} else {
return fmt.Errorf("Exception|401:Invalid token")
}

View file

@ -258,6 +258,13 @@
"comment": "Search configuration (web, kb, db, citation, weights, etc.)",
"nullable": true
},
{
"name": "dependencies",
"type": "json",
"label": "Dependencies",
"comment": "Dependencies on other MCP Clients (name -> version constraint)",
"nullable": true
},
{
"name": "automated",
"type": "boolean",