Merge pull request #1481 from trheyi/main
Refactor delivery/OAuth handling, assistant data & MCP tags
This commit is contained in:
commit
59435f3d6a
20 changed files with 1221 additions and 476 deletions
|
|
@ -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()
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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"}}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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])
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
190
agent/store/xun/utils_test.go
Normal file
190
agent/store/xun/utils_test.go
Normal 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)
|
||||
})
|
||||
}
|
||||
748
data/bindata.go
748
data/bindata.go
File diff suppressed because one or more lines are too long
|
|
@ -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,
|
||||
})
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue