opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #1474 from trheyi/main

Browse source 
Enhance Assistant model with capabilities and sandbox configuration
This commit is contained in:
Max 2026-02-24 12:24:15 +08:00 • committed by GitHub
commit c022c3cc3e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
42 changed files with 2465 additions and 712 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
agent/assistant/agent.go
View file

@ -6,7 +6,7 @@ import (
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
goullm "github.com/yaoapp/gou/llm"
"github.com/yaoapp/yao/agent/assistant/handlers"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
@ -570,7 +570,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// GetConnector get the connector object, capabilities, and error with priority: opts.Connector > ast.Connector
// Note: opts.Connector may be set by Create hook's applyOptionsAdjustments
// Returns: (connector, capabilities, error)
func (ast *Assistant) GetConnector(ctx *context.Context, opts ...*context.Options) (connector.Connector, *openai.Capabilities, error) {
func (ast *Assistant) GetConnector(ctx *context.Context, opts ...*context.Options) (connector.Connector, *goullm.Capabilities, error) {
// Determine connector ID with priority: opts.Connector > ast.Connector
connectorID := ast.Connector
if len(opts) > 0 && opts[0] != nil && opts[0].Connector != "" {
@ -588,9 +588,7 @@ func (ast *Assistant) GetConnector(ctx *context.Context, opts ...*context.Option
return nil, nil, err
}
// Get connector capabilities from settings
// Uses unified capability getter: 1. User-defined models.yml, 2. connector's Setting()["capabilities"], 3. default
capabilities := llm.GetCapabilitiesFromConn(conn, modelCapabilities)
capabilities := llm.GetCapabilitiesFromConn(conn)
return conn, capabilities, nil
}

41
agent/assistant/load.go
View file

@ -9,7 +9,6 @@ import (
jsoniter "github.com/json-iterator/go"
"github.com/spf13/cast"
"github.com/yaoapp/gou/application"
gouOpenAI "github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/fs"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
@ -21,8 +20,7 @@ import (
// loaded the loaded assistant
var loaded = NewCache(200) // 200 is the default capacity
var storage store.Store = nil
var storeSetting *store.Setting = nil // store setting from agent.yml
var modelCapabilities map[string]gouOpenAI.Capabilities = map[string]gouOpenAI.Capabilities{}
var storeSetting *store.Setting = nil // store setting from agent.yml
var defaultConnector string = "" // default connector
var globalUses *context.Uses = nil // global uses configuration from agent.yml
var globalPrompts []store.Prompt = nil // global prompts from agent/prompts.yml
@ -140,11 +138,6 @@ func GetStorage() store.Store {
return storage
}
// SetModelCapabilities set the model capabilities configuration
func SetModelCapabilities(capabilities map[string]gouOpenAI.Capabilities) {
modelCapabilities = capabilities
}
// SetConnector set the connector
func SetConnector(c string) {
defaultConnector = c
@ -564,6 +557,11 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
assistant.Description = v
}
// capabilities
if v, ok := data["capabilities"].(string); ok {
assistant.Capabilities = v
}
// locales
if locales, ok := data["locales"].(i18n.Map); ok {
assistant.Locales = locales
@ -575,30 +573,39 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
if i18nObj.Messages == nil {
i18nObj.Messages = make(map[string]any)
}
// Add name and description if not already present
// Add name, description, and capabilities if not already present
if _, exists := i18nObj.Messages["name"]; !exists && assistant.Name != "" {
i18nObj.Messages["name"] = assistant.Name
}
if _, exists := i18nObj.Messages["description"]; !exists && assistant.Description != "" {
i18nObj.Messages["description"] = assistant.Description
}
if _, exists := i18nObj.Messages["capabilities"]; !exists && assistant.Capabilities != "" {
i18nObj.Messages["capabilities"] = assistant.Capabilities
}
flattened[locale] = i18nObj
}
i18n.Locales[id] = flattened
} else {
// No locales defined, create default with name and description for all common locales
if assistant.Name != "" || assistant.Description != "" {
// No locales defined, create default with name, description, and capabilities for all common locales
if assistant.Name != "" || assistant.Description != "" || assistant.Capabilities != "" {
defaultLocales := make(map[string]i18n.I18n)
// Create entries for all common locales so {{name}} can be resolved
commonLocales := []string{"en", "en-us", "zh", "zh-cn", "zh-tw"}
for _, locale := range commonLocales {
messages := map[string]any{}
if assistant.Name != "" {
messages["name"] = assistant.Name
}
if assistant.Description != "" {
messages["description"] = assistant.Description
}
if assistant.Capabilities != "" {
messages["capabilities"] = assistant.Capabilities
}
defaultLocales[locale] = i18n.I18n{
Locale: locale,
Messages: map[string]any{
"name": assistant.Name,
"description": assistant.Description,
},
Locale: locale,
Messages: messages,
}
}
i18n.Locales[id] = defaultLocales

10
agent/assistant/load_system.go
View file

@ -256,20 +256,11 @@ func resolveSystemConnector(agentID string) string {
// findCapableConnector finds the first connector that supports tool calling
func findCapableConnector() string {
// Get all registered connectors
for id, conn := range connector.Connectors {
if !conn.Is(connector.OPENAI) {
continue
}
// Check from modelCapabilities (user-defined in models.yml)
if caps, exists := modelCapabilities[id]; exists {
if caps.ToolCalls {
return id
}
}
// Check capabilities from connector's Options
if connOpenAI, ok := conn.(*gouOpenAI.Connector); ok {
if connOpenAI.Options.Capabilities != nil && connOpenAI.Options.Capabilities.ToolCalls {
return id
@ -277,7 +268,6 @@ func findCapableConnector() string {
}
}
// No capable connector found, return empty
return ""
}

4
agent/assistant/sandbox_integration_test.go
View file

@ -181,8 +181,8 @@ func TestHasSandboxMethod(t *testing.T) {
require.NoError(t, err)
assert.True(t, astWithSandbox.HasSandbox(), "Assistant with sandbox config should return true")
// Test assistant without sandbox (fullfields doesn't have sandbox)
astWithoutSandbox, err := assistant.LoadPath("/assistants/tests/fullfields")
// Test assistant without sandbox
astWithoutSandbox, err := assistant.LoadPath("/assistants/tests/simple-greeting")
require.NoError(t, err)
assert.False(t, astWithoutSandbox.HasSandbox(), "Assistant without sandbox config should return false")
}

5
agent/context/output.go
View file

@ -3,7 +3,7 @@ package context
import (
"time"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/llm"
"github.com/yaoapp/yao/agent/output"
"github.com/yaoapp/yao/agent/output/message"
)
@ -571,9 +571,8 @@ func (ctx *Context) getOutput() (*output.Output, error) {
Accept: string(ctx.Accept),
}
// Set ModelCapabilities (now using openai.Capabilities directly)
if ctx.Capabilities != nil {
caps := openai.Capabilities(*ctx.Capabilities)
caps := llm.Capabilities(*ctx.Capabilities)
options.Capabilities = &caps
}

4
agent/context/types.go
View file

@ -5,7 +5,7 @@ import (
"sync"
"time"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/llm"
"github.com/yaoapp/gou/store"
"github.com/yaoapp/yao/agent/memory"
"github.com/yaoapp/yao/agent/output"
@ -253,7 +253,7 @@ type Context struct {
sandboxExecutor SandboxExecutor `json:"-"` // Sandbox executor for hooks (set by assistant when sandbox is configured)
// Model capabilities (set by assistant, used by output adapters)
Capabilities *openai.Capabilities `json:"-"` // Model capabilities for the current connector
Capabilities *llm.Capabilities `json:"-"` // Model capabilities for the current connector
// Interrupt control (all interrupt-related logic is encapsulated in InterruptController)
Interrupt *InterruptController `json:"-"` // Interrupt controller for handling user interrupts during streaming

6
agent/context/types_llm.go
View file

@ -1,7 +1,7 @@
package context
import (
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/llm"
"github.com/yaoapp/yao/agent/output/message"
)
@ -38,7 +38,7 @@ const (
)
// GetVisionSupport returns whether vision is supported and the format
func GetVisionSupport(cap *openai.Capabilities) (bool, VisionFormat) {
func GetVisionSupport(cap *llm.Capabilities) (bool, VisionFormat) {
if cap == nil || cap.Vision == nil {
return false, VisionFormatNone
}
@ -70,7 +70,7 @@ func GetVisionSupport(cap *openai.Capabilities) (bool, VisionFormat) {
type CompletionOptions struct {
// Model capabilities (used by LLM to select appropriate provider)
// nil means capabilities are not specified/checked
Capabilities *openai.Capabilities `json:"capabilities,omitempty"`
Capabilities *llm.Capabilities `json:"capabilities,omitempty"`
// User-specified tools for vision, audio, search, and fetch processing
Uses *Uses `json:"uses,omitempty"`

94
agent/llm/capabilities.go
View file

@ -2,92 +2,48 @@ package llm
import (
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/anthropic"
"github.com/yaoapp/gou/connector/openai"
goullm "github.com/yaoapp/gou/llm"
)
// GetCapabilities get the capabilities of a connector by connector ID
// This is a unified function to get connector capabilities with proper priority:
// 1. User-defined model capabilities from agent/models.yml (passed via modelCapabilities map)
// 2. Connector's Setting()["capabilities"] (default capabilities from connector)
// 3. Fallback to minimal default capabilities
//
// Usage in Agent with user-defined models:
//
// capabilities := llm.GetCapabilities(connectorID, modelCapabilities)
//
// Usage in API (without user-defined models):
//
// capabilities := llm.GetCapabilities(connectorID, nil)
func GetCapabilities(connectorID string, modelCapabilities map[string]openai.Capabilities) *openai.Capabilities {
// Reads capabilities from connector's Setting()["capabilities"], with fallback to defaults.
func GetCapabilities(connectorID string) *goullm.Capabilities {
if connectorID == "" {
return getDefaultCapabilities()
}
// Priority 1: Check user-defined model capabilities from agent/models.yml
if modelCapabilities != nil {
if modelCaps, exists := modelCapabilities[connectorID]; exists {
return &modelCaps
}
}
// Priority 2: Get connector and extract capabilities from Setting()
conn, err := connector.Select(connectorID)
if err != nil {
// If connector not found, return default
return getDefaultCapabilities()
}
return GetCapabilitiesFromConn(conn, modelCapabilities)
return GetCapabilitiesFromConn(conn)
}
// GetCapabilitiesFromConn get the capabilities from a connector instance
// This is useful when you already have the connector object
func GetCapabilitiesFromConn(conn connector.Connector, modelCapabilities map[string]openai.Capabilities) *openai.Capabilities {
func GetCapabilitiesFromConn(conn connector.Connector) *goullm.Capabilities {
if conn == nil {
return getDefaultCapabilities()
}
connectorID := conn.ID()
// Priority 1: Check user-defined model capabilities from agent/models.yml
if modelCapabilities != nil {
if modelCaps, exists := modelCapabilities[connectorID]; exists {
return &modelCaps
}
}
// Priority 2: Get capabilities from connector's Setting() method
settings := conn.Setting()
if settings != nil {
if caps, ok := settings["capabilities"]; ok {
// Try to convert to *openai.Capabilities
if capabilities, ok := caps.(*openai.Capabilities); ok {
if capabilities, ok := caps.(*goullm.Capabilities); ok {
return capabilities
}
// Try to convert to openai.Capabilities (value type)
if capabilities, ok := caps.(openai.Capabilities); ok {
if capabilities, ok := caps.(goullm.Capabilities); ok {
return &capabilities
}
// Try to convert from *anthropic.Capabilities
if capabilities, ok := caps.(*anthropic.Capabilities); ok {
return convertAnthropicCaps(capabilities)
}
// Try to convert from anthropic.Capabilities (value type)
if capabilities, ok := caps.(anthropic.Capabilities); ok {
return convertAnthropicCaps(&capabilities)
}
}
}
// Priority 3: Fallback to minimal default capabilities
return getDefaultCapabilities()
}
// getDefaultCapabilities returns minimal default capabilities
// This should rarely be used as modern connectors provide capabilities via Setting()
func getDefaultCapabilities() *openai.Capabilities {
return &openai.Capabilities{
func getDefaultCapabilities() *goullm.Capabilities {
return &goullm.Capabilities{
Vision: false,
ToolCalls: false,
Audio: false,
@ -95,14 +51,13 @@ func getDefaultCapabilities() *openai.Capabilities {
Streaming: false,
JSON: false,
Multimodal: false,
TemperatureAdjustable: true, // Default to true for non-reasoning models
TemperatureAdjustable: true,
}
}
// GetCapabilitiesMap get capabilities as map[string]interface{} for API responses
// This is useful for OpenAPI responses that need JSON-serializable format
func GetCapabilitiesMap(connectorID string, modelCapabilities map[string]openai.Capabilities) map[string]interface{} {
caps := GetCapabilities(connectorID, modelCapabilities)
func GetCapabilitiesMap(connectorID string) map[string]interface{} {
caps := GetCapabilities(connectorID)
if caps == nil {
return nil
}
@ -110,16 +65,14 @@ func GetCapabilitiesMap(connectorID string, modelCapabilities map[string]openai.
return ToMap(caps)
}
// ToMap converts openai.Capabilities to map[string]interface{}
// This is useful for JSON serialization in API responses
func ToMap(caps *openai.Capabilities) map[string]interface{} {
// ToMap converts Capabilities to map[string]interface{}
func ToMap(caps *goullm.Capabilities) map[string]interface{} {
if caps == nil {
return nil
}
result := make(map[string]interface{})
// Handle Vision field specially as it can be bool or string
if caps.Vision != nil {
result["vision"] = caps.Vision
}
@ -135,22 +88,3 @@ func ToMap(caps *openai.Capabilities) map[string]interface{} {
return result
}
// convertAnthropicCaps converts anthropic.Capabilities to openai.Capabilities
// This provides a unified capabilities interface across connector types
func convertAnthropicCaps(caps *anthropic.Capabilities) *openai.Capabilities {
if caps == nil {
return getDefaultCapabilities()
}
return &openai.Capabilities{
Vision: caps.Vision,
Audio: caps.Audio,
STT: caps.STT,
ToolCalls: caps.ToolCalls,
Reasoning: caps.Reasoning,
Streaming: caps.Streaming,
JSON: caps.JSON,
Multimodal: caps.Multimodal,
TemperatureAdjustable: caps.TemperatureAdjustable,
}
}

8
agent/llm/jsapi.go
View file

@ -203,7 +203,7 @@ func parseToolCalls(toolCalls []interface{}) []agentContext.ToolCall {
// buildCompletionOptions creates CompletionOptions from JS opts map
func buildCompletionOptions(conn connector.Connector, opts map[string]interface{}) *agentContext.CompletionOptions {
// Get capabilities from connector
capabilities := GetCapabilitiesFromConn(conn, nil)
capabilities := GetCapabilitiesFromConn(conn)
completionOptions := &agentContext.CompletionOptions{
Capabilities: capabilities,
@ -737,12 +737,6 @@ func (api *JSAPI) executeRace(requests []*Request) []interface{} {
return []interface{}{result}
}
// executeSingleRequest executes a single LLM request using the original context
// This is used for single calls (not batch)
func (api *JSAPI) executeSingleRequest(request *Request) interface{} {
return api.StreamWithHandler(request.Connector, request.Messages, request.Options, request.Handler)
}
// executeSingleRequestWithForkedContext executes a single LLM request with a forked context
// This is used by batch operations (All/Any/Race) to avoid race conditions
// when multiple goroutines access shared context state

26
agent/llm/providers/anthropic/anthropic.go
View file

@ -9,9 +9,8 @@ import (
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/connector"
gouAnthropicConn "github.com/yaoapp/gou/connector/anthropic"
gouOpenAI "github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/http"
goullm "github.com/yaoapp/gou/llm"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
@ -27,34 +26,15 @@ type Provider struct {
}
// New create a new Anthropic provider
func New(conn connector.Connector, capabilities *gouOpenAI.Capabilities) *Provider {
func New(conn connector.Connector, capabilities *goullm.Capabilities) *Provider {
return &Provider{
Provider: base.NewProvider(conn, capabilities),
adapters: buildAdapters(capabilities),
}
}
// NewFromAnthropicCaps create a new Anthropic provider from Anthropic capabilities
func NewFromAnthropicCaps(conn connector.Connector, caps *gouAnthropicConn.Capabilities) *Provider {
// Convert anthropic capabilities to openai capabilities for base provider compatibility
openaiCaps := &gouOpenAI.Capabilities{
Vision: caps.Vision,
Audio: caps.Audio,
ToolCalls: caps.ToolCalls,
Reasoning: caps.Reasoning,
Streaming: caps.Streaming,
JSON: caps.JSON,
Multimodal: caps.Multimodal,
TemperatureAdjustable: caps.TemperatureAdjustable,
}
return &Provider{
Provider: base.NewProvider(conn, openaiCaps),
adapters: buildAdapters(openaiCaps),
}
}
// buildAdapters builds capability adapters based on model capabilities
func buildAdapters(cap *gouOpenAI.Capabilities) []adapters.CapabilityAdapter {
func buildAdapters(cap *goullm.Capabilities) []adapters.CapabilityAdapter {
if cap == nil {
return []adapters.CapabilityAdapter{}
}

9
agent/llm/providers/anthropic/anthropic_test.go
View file

@ -7,7 +7,7 @@ import (
"testing"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
goullm "github.com/yaoapp/gou/llm"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/llm"
"github.com/yaoapp/yao/agent/output/message"
@ -34,9 +34,8 @@ func TestAnthropicStreamBasic(t *testing.T) {
t.Fatal("Connector is not ANTHROPIC type")
}
// Use openai.Capabilities — SelectProvider auto-detects Anthropic format from connector type
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Capabilities: &goullm.Capabilities{
Streaming: true,
ToolCalls: true,
},
@ -111,7 +110,7 @@ func TestAnthropicStreamWithToolCalls(t *testing.T) {
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Capabilities: &goullm.Capabilities{
Streaming: true,
ToolCalls: true,
},
@ -226,7 +225,7 @@ func TestAnthropicStreamRetry(t *testing.T) {
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Capabilities: &goullm.Capabilities{
Streaming: true,
ToolCalls: true,
},

6
agent/llm/providers/base/base.go
View file

@ -4,7 +4,7 @@ import (
"fmt"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/llm"
"github.com/yaoapp/yao/agent/context"
)
@ -12,11 +12,11 @@ import (
// Provides common functionality for all LLM providers
type Provider struct {
Connector connector.Connector
Capabilities *openai.Capabilities
Capabilities *llm.Capabilities
}
// NewProvider create a new base provider
func NewProvider(conn connector.Connector, capabilities *openai.Capabilities) *Provider {
func NewProvider(conn connector.Connector, capabilities *llm.Capabilities) *Provider {
return &Provider{
Connector: conn,
Capabilities: capabilities,

8
agent/llm/providers/factory.go
View file

@ -4,7 +4,6 @@ import (
"fmt"
"github.com/yaoapp/gou/connector"
gouAnthropicConn "github.com/yaoapp/gou/connector/anthropic"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/llm/providers/anthropic"
"github.com/yaoapp/yao/agent/llm/providers/openai"
@ -43,13 +42,6 @@ func SelectProvider(conn connector.Connector, options *context.CompletionOptions
return openai.New(conn, options.Capabilities), nil
case "anthropic":
// Anthropic Messages API (Claude, Kimi Code, etc.)
// Check if connector has native Anthropic capabilities
settings := conn.Setting()
if caps, ok := settings["capabilities"].(*gouAnthropicConn.Capabilities); ok {
return anthropic.NewFromAnthropicCaps(conn, caps), nil
}
// Fallback: use OpenAI capabilities (converted from connector settings)
return anthropic.New(conn, options.Capabilities), nil
default:

8
agent/llm/providers/openai/openai.go
View file

@ -8,8 +8,8 @@ import (
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/connector"
gouOpenAI "github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/http"
goullm "github.com/yaoapp/gou/llm"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
@ -137,7 +137,7 @@ func buildAPIURL(host, endpoint string) string {
}
// New create a new OpenAI provider with capability adapters
func New(conn connector.Connector, capabilities *gouOpenAI.Capabilities) *Provider {
func New(conn connector.Connector, capabilities *goullm.Capabilities) *Provider {
return &Provider{
Provider: base.NewProvider(conn, capabilities),
adapters: buildAdapters(capabilities),
@ -145,7 +145,7 @@ func New(conn connector.Connector, capabilities *gouOpenAI.Capabilities) *Provid
}
// buildAdapters builds capability adapters based on model capabilities
func buildAdapters(cap *gouOpenAI.Capabilities) []adapters.CapabilityAdapter {
func buildAdapters(cap *goullm.Capabilities) []adapters.CapabilityAdapter {
if cap == nil {
return []adapters.CapabilityAdapter{}
}
@ -182,7 +182,7 @@ func buildAdapters(cap *gouOpenAI.Capabilities) []adapters.CapabilityAdapter {
}
// detectReasoningFormat detects the reasoning format based on capabilities
func detectReasoningFormat(cap *gouOpenAI.Capabilities) adapters.ReasoningFormat {
func detectReasoningFormat(cap *goullm.Capabilities) adapters.ReasoningFormat {
// TODO: Implement better detection logic
// For now, default to OpenAI o1 format if reasoning is supported
if cap.Reasoning {

71
agent/load.go
View file

@ -6,7 +6,7 @@ import (
"github.com/yaoapp/gou/application"
"github.com/yaoapp/gou/connector"
gouOpenAI "github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/helper"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
@ -47,6 +47,9 @@ func Load(cfg config.Config) error {
setting.StoreSetting.MaxSize = 20 // default is 20
}
// Resolve $ENV.XXX references in system and uses fields
resolveEnvStrings(&setting)
// Default Assistant, Agent is the developer name, Mohe is the brand name of the assistant
if setting.Uses == nil {
setting.Uses = &types.Uses{Default: "mohe"} // Agent is the developer name, Mohe is the brand name of the assistant
@ -75,12 +78,6 @@ func Load(cfg config.Config) error {
return err
}
// Initialize model capabilities
err = initModelCapabilities()
if err != nil {
return err
}
// Initialize Global I18n
err = initGlobalI18n()
if err != nil {
@ -148,29 +145,6 @@ func GetGlobalPrompts(ctx map[string]string) []store.Prompt {
return store.Prompts(agentDSL.GlobalPrompts).Parse(ctx)
}
// initModelCapabilities initialize the model capabilities configuration
func initModelCapabilities() error {
path := filepath.Join("agent", "models.yml")
if exists, _ := application.App.Exists(path); !exists {
return nil
}
// Read the model capabilities configuration
bytes, err := application.App.Read(path)
if err != nil {
return err
}
var models map[string]gouOpenAI.Capabilities = map[string]gouOpenAI.Capabilities{}
err = application.Parse("models.yml", bytes, &models)
if err != nil {
return err
}
agentDSL.Models = models
return nil
}
// initStore initialize the store
func initStore() error {
@ -231,10 +205,6 @@ func initAssistant() error {
assistant.SetGlobalPrompts(agentDSL.GlobalPrompts)
}
if agentDSL.Models != nil {
assistant.SetModelCapabilities(agentDSL.Models)
}
if agentDSL.KB != nil {
assistant.SetGlobalKBSetting(agentDSL.KB)
}
@ -478,3 +448,36 @@ func defaultAssistant() (*assistant.Assistant, error) {
}
return assistant.Get(agentDSL.Uses.Default)
}
// resolveEnvStrings resolves $ENV.XXX references in agent.yml string fields.
// agent.yml is parsed via yaml.Unmarshal which does not handle $ENV substitution,
// unlike connector files which call helper.EnvString explicitly during Register.
func resolveEnvStrings(setting *types.DSL) {
if setting.System != nil {
setting.System.Default = helper.EnvString(setting.System.Default)
setting.System.Keyword = helper.EnvString(setting.System.Keyword)
setting.System.QueryDSL = helper.EnvString(setting.System.QueryDSL)
setting.System.Title = helper.EnvString(setting.System.Title)
setting.System.Prompt = helper.EnvString(setting.System.Prompt)
setting.System.RobotPrompt = helper.EnvString(setting.System.RobotPrompt)
setting.System.NeedSearch = helper.EnvString(setting.System.NeedSearch)
setting.System.Entity = helper.EnvString(setting.System.Entity)
}
if setting.Uses != nil {
setting.Uses.Default = helper.EnvString(setting.Uses.Default)
setting.Uses.Title = helper.EnvString(setting.Uses.Title)
setting.Uses.Prompt = helper.EnvString(setting.Uses.Prompt)
setting.Uses.RobotPrompt = helper.EnvString(setting.Uses.RobotPrompt)
setting.Uses.Vision = helper.EnvString(setting.Uses.Vision)
setting.Uses.Audio = helper.EnvString(setting.Uses.Audio)
setting.Uses.Search = helper.EnvString(setting.Uses.Search)
setting.Uses.Fetch = helper.EnvString(setting.Uses.Fetch)
setting.Uses.Web = helper.EnvString(setting.Uses.Web)
setting.Uses.Keyword = helper.EnvString(setting.Uses.Keyword)
setting.Uses.QueryDSL = helper.EnvString(setting.Uses.QueryDSL)
setting.Uses.Rerank = helper.EnvString(setting.Uses.Rerank)
}
setting.Cache = helper.EnvString(setting.Cache)
}

113
agent/load_test.go
View file

@ -8,6 +8,7 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/types"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/test"
)
@ -54,12 +55,6 @@ func TestLoad(t *testing.T) {
assert.Contains(t, agent.GlobalPrompts[0].Content, "$SYS.")
})
t.Run("LoadModelCapabilities", func(t *testing.T) {
// Model capabilities should be loaded from agent/models.yml
assert.NotNil(t, agent.Models)
assert.Greater(t, len(agent.Models), 0)
})
t.Run("LoadKBConfig", func(t *testing.T) {
// KB configuration should be loaded from agent/kb.yml
assert.NotNil(t, agent.KB)
@ -217,6 +212,112 @@ func TestGetGlobalPromptsWithDisableFlag(t *testing.T) {
})
}
func TestResolveEnvStrings(t *testing.T) {
t.Setenv("TEST_CONNECTOR", "openai.gpt-5")
t.Setenv("TEST_ASSISTANT", "my-assistant")
t.Setenv("TEST_CACHE", "my-cache")
t.Run("SystemFields", func(t *testing.T) {
setting := &types.DSL{
System: &types.System{
Default: "$ENV.TEST_CONNECTOR",
Keyword: "$ENV.TEST_CONNECTOR",
QueryDSL: "$ENV.TEST_CONNECTOR",
Title: "$ENV.TEST_CONNECTOR",
Prompt: "$ENV.TEST_CONNECTOR",
RobotPrompt: "$ENV.TEST_CONNECTOR",
NeedSearch: "$ENV.TEST_CONNECTOR",
Entity: "$ENV.TEST_CONNECTOR",
},
}
resolveEnvStrings(setting)
assert.Equal(t, "openai.gpt-5", setting.System.Default)
assert.Equal(t, "openai.gpt-5", setting.System.Keyword)
assert.Equal(t, "openai.gpt-5", setting.System.QueryDSL)
assert.Equal(t, "openai.gpt-5", setting.System.Title)
assert.Equal(t, "openai.gpt-5", setting.System.Prompt)
assert.Equal(t, "openai.gpt-5", setting.System.RobotPrompt)
assert.Equal(t, "openai.gpt-5", setting.System.NeedSearch)
assert.Equal(t, "openai.gpt-5", setting.System.Entity)
})
t.Run("UsesFields", func(t *testing.T) {
setting := &types.DSL{
Uses: &types.Uses{
Default: "$ENV.TEST_ASSISTANT",
Title: "$ENV.TEST_ASSISTANT",
Prompt: "$ENV.TEST_ASSISTANT",
RobotPrompt: "$ENV.TEST_ASSISTANT",
Vision: "$ENV.TEST_ASSISTANT",
Audio: "$ENV.TEST_ASSISTANT",
Search: "$ENV.TEST_ASSISTANT",
Fetch: "$ENV.TEST_ASSISTANT",
Web: "$ENV.TEST_ASSISTANT",
Keyword: "$ENV.TEST_ASSISTANT",
QueryDSL: "$ENV.TEST_ASSISTANT",
Rerank: "$ENV.TEST_ASSISTANT",
},
}
resolveEnvStrings(setting)
assert.Equal(t, "my-assistant", setting.Uses.Default)
assert.Equal(t, "my-assistant", setting.Uses.Title)
assert.Equal(t, "my-assistant", setting.Uses.Prompt)
assert.Equal(t, "my-assistant", setting.Uses.RobotPrompt)
assert.Equal(t, "my-assistant", setting.Uses.Vision)
assert.Equal(t, "my-assistant", setting.Uses.Audio)
assert.Equal(t, "my-assistant", setting.Uses.Search)
assert.Equal(t, "my-assistant", setting.Uses.Fetch)
assert.Equal(t, "my-assistant", setting.Uses.Web)
assert.Equal(t, "my-assistant", setting.Uses.Keyword)
assert.Equal(t, "my-assistant", setting.Uses.QueryDSL)
assert.Equal(t, "my-assistant", setting.Uses.Rerank)
})
t.Run("CacheField", func(t *testing.T) {
setting := &types.DSL{Cache: "$ENV.TEST_CACHE"}
resolveEnvStrings(setting)
assert.Equal(t, "my-cache", setting.Cache)
})
t.Run("PlainStringsUnchanged", func(t *testing.T) {
setting := &types.DSL{
Cache: "plain-cache",
System: &types.System{
Default: "openai.gpt-5",
},
Uses: &types.Uses{
Default: "mohe",
Title: "__yao.title",
},
}
resolveEnvStrings(setting)
assert.Equal(t, "plain-cache", setting.Cache)
assert.Equal(t, "openai.gpt-5", setting.System.Default)
assert.Equal(t, "mohe", setting.Uses.Default)
assert.Equal(t, "__yao.title", setting.Uses.Title)
})
t.Run("NilSystemAndUses", func(t *testing.T) {
setting := &types.DSL{Cache: "test"}
assert.NotPanics(t, func() {
resolveEnvStrings(setting)
})
})
t.Run("UndefinedEnvReturnsEmpty", func(t *testing.T) {
setting := &types.DSL{
System: &types.System{
Default: "$ENV.UNDEFINED_VAR_12345",
},
}
resolveEnvStrings(setting)
assert.Equal(t, "", setting.System.Default)
})
}
func TestGlobalPromptsContent(t *testing.T) {
prepare(t)
defer test.Clean()

4
agent/output/message/types.go
View file

@ -3,7 +3,7 @@ package message
import (
"net/http"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/llm"
traceTypes "github.com/yaoapp/yao/trace/types"
)
@ -13,7 +13,7 @@ type Options struct {
Accept string
Writer http.ResponseWriter
Trace traceTypes.Manager
Capabilities *openai.Capabilities
Capabilities *llm.Capabilities
Locale string
}

11
agent/store/types/convert.go
View file

@ -291,6 +291,9 @@ func ToAssistantModel(v interface{}) (*AssistantModel, error) {
if description, ok := data["description"].(string); ok {
model.Description = description
}
if capabilities, ok := data["capabilities"].(string); ok {
model.Capabilities = capabilities
}
if share, ok := data["share"].(string); ok {
model.Share = share
}
@ -421,6 +424,14 @@ func ToAssistantModel(v interface{}) (*AssistantModel, error) {
}
}
// Sandbox
if sandbox, ok := data["sandbox"]; ok && sandbox != nil {
sb, err := ToSandbox(sandbox)
if err == nil {
model.Sandbox = sb
}
}
// Placeholder
if placeholder, ok := data["placeholder"]; ok && placeholder != nil {
raw, err := jsoniter.Marshal(placeholder)

12
agent/store/types/fields.go
View file

@ -20,10 +20,12 @@ var AssistantAllowedFields = map[string]bool{
"prompts": true,
"prompt_presets": true,
"disable_global_prompts": true,
"capabilities": true,
"workflow": true,
"kb": true,
"db": true,
"mcp": true,
"sandbox": true,
"source": true,
"tags": true,
"modes": true,
@ -53,6 +55,7 @@ var AssistantDefaultFields = []string{
"avatar",
"connector",
"description",
"capabilities", // Capabilities description for Robot orchestration (lightweight)
"tags", // Tags for categorization (lightweight)
"modes", // Supported modes (lightweight)
"default_mode", // Default mode (lightweight)
@ -63,9 +66,10 @@ var AssistantDefaultFields = []string{
"share",
"automated",
"mentionable",
"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)
"created_at",
"updated_at",
"__yao_created_by", // Permission: creator user ID
@ -84,6 +88,7 @@ var AssistantFullFields = []string{
"connector",
"connector_options",
"description",
"capabilities",
"path",
"sort",
"built_in",
@ -96,6 +101,7 @@ var AssistantFullFields = []string{
"kb",
"db",
"mcp",
"sandbox",
"source",
"tags",
"modes",

2
agent/store/types/types.go
View file

@ -183,6 +183,7 @@ type AssistantFilter struct {
Mentionable *bool `json:"mentionable,omitempty"` // Filter by mentionable status
Automated *bool `json:"automated,omitempty"` // Filter by automation status
BuiltIn *bool `json:"built_in,omitempty"` // Filter by built-in status
Sandbox *bool `json:"sandbox,omitempty"` // Filter by sandbox configuration (true=has sandbox, false=no sandbox)
Page int `json:"page,omitempty"` // Page number, starting from 1
PageSize int `json:"pagesize,omitempty"` // Items per page
Select []string `json:"select,omitempty"` // Fields to return, returns all fields if empty
@ -429,6 +430,7 @@ type AssistantModel struct {
BuiltIn bool `json:"built_in,omitempty"` // Whether this is a built-in assistant
Sort int `json:"sort,omitempty"` // Assistant Sort
Description string `json:"description,omitempty"` // Assistant Description
Capabilities string `json:"capabilities,omitempty"` // Assistant capabilities description (useful for Robot orchestration)
Tags []string `json:"tags,omitempty"` // Assistant Tags
Modes []string `json:"modes,omitempty"` // Supported modes (e.g., ["task", "chat"]), null means all modes are supported
DefaultMode string `json:"default_mode,omitempty"` // Default mode, can be empty

45
agent/store/xun/assistant.go
View file

@ -99,6 +99,11 @@ func (store *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
} else {
data["description"] = nil
}
if assistant.Capabilities != "" {
data["capabilities"] = assistant.Capabilities
} else {
data["capabilities"] = nil
}
if assistant.Path != "" {
data["path"] = assistant.Path
} else {
@ -181,6 +186,7 @@ func (store *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
"db": assistant.DB,
"mcp": assistant.MCP,
"workflow": assistant.Workflow,
"sandbox": assistant.Sandbox,
"placeholder": assistant.Placeholder,
"locales": assistant.Locales,
"uses": assistant.Uses,
@ -243,14 +249,14 @@ 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", "placeholder", "locales", "uses", "search"}
jsonFields := []string{"options", "tags", "modes", "prompts", "prompt_presets", "connector_options", "kb", "db", "mcp", "workflow", "sandbox", "placeholder", "locales", "uses", "search"}
jsonFieldSet := make(map[string]bool)
for _, field := range jsonFields {
jsonFieldSet[field] = true
}
// List of nullable string fields
nullableStringFields := []string{"name", "avatar", "description", "path", "source", "default_mode", "__yao_created_by", "__yao_updated_by", "__yao_team_id", "__yao_tenant_id"}
nullableStringFields := []string{"name", "avatar", "description", "capabilities", "path", "source", "default_mode", "__yao_created_by", "__yao_updated_by", "__yao_team_id", "__yao_tenant_id"}
nullableFieldSet := make(map[string]bool)
for _, field := range nullableStringFields {
nullableFieldSet[field] = true
@ -349,6 +355,7 @@ func (store *Xun) GetAssistants(filter types.AssistantFilter, locale ...string)
qb.Where(func(qb query.Query) {
qb.Where("name", "like", fmt.Sprintf("%%%s%%", filter.Keywords)).
OrWhere("description", "like", fmt.Sprintf("%%%s%%", filter.Keywords)).
OrWhere("capabilities", "like", fmt.Sprintf("%%%s%%", filter.Keywords)).
OrWhere("locales", "like", fmt.Sprintf("%%%s%%", filter.Keywords))
})
}
@ -393,6 +400,21 @@ func (store *Xun) GetAssistants(filter types.AssistantFilter, locale ...string)
qb.Where("built_in", *filter.BuiltIn)
}
// Apply sandbox filter (true = has sandbox config, false = no sandbox config)
// MySQL JSON columns distinguish between SQL NULL and JSON literal null.
// CAST(sandbox AS CHAR) returns 'null' for JSON null and NULL for SQL NULL.
if filter.Sandbox != nil {
if *filter.Sandbox {
qb.WhereNotNull("sandbox").
WhereRaw("CAST(`sandbox` AS CHAR) <> 'null'")
} else {
qb.Where(func(qb query.Query) {
qb.WhereNull("sandbox").
OrWhereRaw("CAST(`sandbox` AS CHAR) = 'null'")
})
}
}
// Apply custom query filter function (for permission filtering)
if filter.QueryFilter != nil {
qb.Where(filter.QueryFilter)
@ -448,7 +470,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", "kb", "mcp", "placeholder", "locales", "uses", "search"}
jsonFields := []string{"tags", "options", "prompts", "prompt_presets", "connector_options", "workflow", "sandbox", "kb", "mcp", "placeholder", "locales", "uses", "search"}
for _, row := range rows {
data := row.ToMap()
@ -521,7 +543,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", "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"}
store.parseJSONFields(data, jsonFields)
// Convert map to types.AssistantModel
@ -536,6 +558,7 @@ func (store *Xun) GetAssistant(assistantID string, fields []string, locale ...st
BuiltIn: getBool(data, "built_in"),
Sort: getInt(data, "sort"),
Description: getString(data, "description"),
Capabilities: getString(data, "capabilities"),
DefaultMode: getString(data, "default_mode"),
Readonly: getBool(data, "readonly"),
Public: getBool(data, "public"),
@ -636,6 +659,13 @@ func (store *Xun) GetAssistant(assistantID string, fields []string, locale ...st
}
}
if sandbox, has := data["sandbox"]; has && sandbox != nil {
sb, err := types.ToSandbox(sandbox)
if err == nil {
model.Sandbox = sb
}
}
if placeholder, has := data["placeholder"]; has && placeholder != nil {
raw, err := jsoniter.Marshal(placeholder)
if err == nil {
@ -839,6 +869,13 @@ func (store *Xun) translate(model *types.AssistantModel, assistantID string, loc
}
}
// Translate capabilities
if translated := i18n.Translate(assistantID, locale, model.Capabilities); translated != nil {
if s, ok := translated.(string); ok {
model.Capabilities = s
}
}
// Translate prompts
if model.Prompts != nil {
for i := range model.Prompts {

461
agent/store/xun/assistant_test.go
View file

@ -788,6 +788,467 @@ func TestSaveAssistant(t *testing.T) {
t.Logf("Successfully saved and retrieved source code for assistant %s", id)
})
t.Run("SandboxConfiguration", func(t *testing.T) {
assistant := &types.AssistantModel{
Name: "Sandbox Test Assistant",
Type: "assistant",
Connector: "openai",
Share: "private",
Sandbox: &types.Sandbox{
Command: "claude",
Timeout: "5m",
Arguments: map[string]interface{}{
"max_turns": 10,
"permission_mode": "bypassPermissions",
},
},
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with sandbox: %v", err)
}
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.Sandbox == nil {
t.Fatal("Expected sandbox to be set")
}
if retrieved.Sandbox.Command != "claude" {
t.Errorf("Expected command 'claude', got '%s'", retrieved.Sandbox.Command)
}
if retrieved.Sandbox.Timeout != "5m" {
t.Errorf("Expected timeout '5m', got '%s'", retrieved.Sandbox.Timeout)
}
if retrieved.Sandbox.Arguments == nil {
t.Fatal("Expected sandbox arguments to be set")
}
if maxTurns, ok := retrieved.Sandbox.Arguments["max_turns"].(float64); !ok || maxTurns != 10 {
t.Errorf("Expected max_turns 10, got %v", retrieved.Sandbox.Arguments["max_turns"])
}
t.Logf("Successfully saved and retrieved sandbox configuration for assistant %s", id)
})
t.Run("SandboxWithImage", func(t *testing.T) {
assistant := &types.AssistantModel{
Name: "Sandbox Image Test",
Type: "assistant",
Connector: "openai",
Share: "private",
Sandbox: &types.Sandbox{
Command: "claude",
Image: "yaoapp/sandbox-claude-desktop:latest",
Timeout: "20m",
MaxMemory: "4g",
MaxCPU: 2.0,
Arguments: map[string]interface{}{
"max_turns": 500,
"permission_mode": "bypassPermissions",
"disallowed_tools": "WebSearch",
},
Secrets: map[string]string{
"GITHUB_TOKEN": "$ENV.GITHUB_TOKEN",
},
},
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with sandbox image: %v", err)
}
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.Sandbox == nil {
t.Fatal("Expected sandbox to be set")
}
if retrieved.Sandbox.Image != "yaoapp/sandbox-claude-desktop:latest" {
t.Errorf("Expected image 'yaoapp/sandbox-claude-desktop:latest', got '%s'", retrieved.Sandbox.Image)
}
if retrieved.Sandbox.MaxMemory != "4g" {
t.Errorf("Expected max_memory '4g', got '%s'", retrieved.Sandbox.MaxMemory)
}
if retrieved.Sandbox.MaxCPU != 2.0 {
t.Errorf("Expected max_cpu 2.0, got %f", retrieved.Sandbox.MaxCPU)
}
if retrieved.Sandbox.Secrets == nil || retrieved.Sandbox.Secrets["GITHUB_TOKEN"] != "$ENV.GITHUB_TOKEN" {
t.Errorf("Expected secrets to contain GITHUB_TOKEN, got %v", retrieved.Sandbox.Secrets)
}
t.Logf("Successfully saved and retrieved sandbox with image for assistant %s", id)
})
t.Run("NilSandbox", func(t *testing.T) {
assistant := &types.AssistantModel{
Name: "No Sandbox Assistant",
Type: "assistant",
Connector: "openai",
Share: "private",
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant without sandbox: %v", err)
}
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.Sandbox != nil {
t.Errorf("Expected sandbox to be nil, got %+v", retrieved.Sandbox)
}
})
t.Run("CapabilitiesField", func(t *testing.T) {
assistant := &types.AssistantModel{
Name: "Capabilities Test",
Type: "assistant",
Connector: "openai",
Share: "private",
Description: "A test assistant",
Capabilities: "Can search the web, analyze data, write code, and summarize documents.",
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with capabilities: %v", err)
}
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.Capabilities != "Can search the web, analyze data, write code, and summarize documents." {
t.Errorf("Expected capabilities to match, got '%s'", retrieved.Capabilities)
}
if retrieved.Description != "A test assistant" {
t.Errorf("Expected description 'A test assistant', got '%s'", retrieved.Description)
}
t.Logf("Successfully saved and retrieved capabilities for assistant %s", id)
})
t.Run("EmptyCapabilities", func(t *testing.T) {
assistant := &types.AssistantModel{
Name: "No Capabilities Assistant",
Type: "assistant",
Connector: "openai",
Share: "private",
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant without capabilities: %v", err)
}
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.Capabilities != "" {
t.Errorf("Expected empty capabilities, got '%s'", retrieved.Capabilities)
}
})
t.Run("CapabilitiesWithI18n", func(t *testing.T) {
assistant := &types.AssistantModel{
Name: "{{name}}",
Type: "assistant",
Connector: "openai",
Share: "private",
Description: "{{description}}",
Capabilities: "{{capabilities}}",
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with i18n capabilities: %v", err)
}
// Setup i18n
i18n.Locales[id] = map[string]i18n.I18n{
"en": {
Locale: "en",
Messages: map[string]any{
"name": "i18n Test",
"description": "Description in English",
"capabilities": "Can do X, Y, and Z",
},
},
"zh-cn": {
Locale: "zh-cn",
Messages: map[string]any{
"name": "国际化测试",
"description": "中文描述",
"capabilities": "可以做X、Y和Z",
},
},
}
retrievedEN, err := store.GetAssistant(id, types.AssistantFullFields, "en")
if err != nil {
t.Fatalf("Failed to get assistant with EN locale: %v", err)
}
if retrievedEN.Capabilities != "Can do X, Y, and Z" {
t.Errorf("Expected capabilities 'Can do X, Y, and Z', got '%s'", retrievedEN.Capabilities)
}
retrievedZH, err := store.GetAssistant(id, types.AssistantFullFields, "zh-cn")
if err != nil {
t.Fatalf("Failed to get assistant with ZH locale: %v", err)
}
if retrievedZH.Capabilities != "可以做X、Y和Z" {
t.Errorf("Expected capabilities '可以做X、Y和Z', got '%s'", retrievedZH.Capabilities)
}
// Cleanup
delete(i18n.Locales, id)
t.Logf("Successfully tested capabilities i18n for assistant %s", id)
})
t.Run("CapabilitiesInKeywordSearch", func(t *testing.T) {
uniqueCapability := fmt.Sprintf("unique-cap-%d", time.Now().UnixNano())
assistant := &types.AssistantModel{
Name: "Capabilities Search Test",
Type: "assistant",
Connector: "openai",
Share: "private",
Capabilities: uniqueCapability,
}
_, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant: %v", err)
}
response, err := store.GetAssistants(types.AssistantFilter{
Keywords: uniqueCapability,
Page: 1,
PageSize: 20,
})
if err != nil {
t.Fatalf("Failed to search by capabilities keyword: %v", err)
}
if len(response.Data) < 1 {
t.Error("Expected to find assistant by capabilities keyword search")
}
found := false
for _, a := range response.Data {
if a.Capabilities == uniqueCapability {
found = true
break
}
}
if !found {
t.Error("Expected to find assistant with matching capabilities")
}
})
t.Run("UpdateSandbox", func(t *testing.T) {
assistant := &types.AssistantModel{
Name: "Update Sandbox Test",
Type: "assistant",
Connector: "openai",
Share: "private",
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to create assistant: %v", err)
}
// Update with sandbox
updates := map[string]interface{}{
"sandbox": &types.Sandbox{
Command: "claude",
Timeout: "10m",
},
}
err = store.UpdateAssistant(id, updates)
if err != nil {
t.Fatalf("Failed to update sandbox: %v", err)
}
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.Sandbox == nil {
t.Fatal("Expected sandbox to be set")
}
if retrieved.Sandbox.Command != "claude" {
t.Errorf("Expected command 'claude', got '%s'", retrieved.Sandbox.Command)
}
// Update to remove sandbox
updates2 := map[string]interface{}{
"sandbox": nil,
}
err = store.UpdateAssistant(id, updates2)
if err != nil {
t.Fatalf("Failed to remove sandbox: %v", err)
}
retrieved2, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved2.Sandbox != nil {
t.Errorf("Expected sandbox to be nil, got %+v", retrieved2.Sandbox)
}
})
t.Run("UpdateCapabilities", func(t *testing.T) {
assistant := &types.AssistantModel{
Name: "Update Capabilities Test",
Type: "assistant",
Connector: "openai",
Share: "private",
Capabilities: "Original capabilities",
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to create assistant: %v", err)
}
// Update capabilities
updates := map[string]interface{}{
"capabilities": "Updated capabilities: can search, analyze, and write code",
}
err = store.UpdateAssistant(id, updates)
if err != nil {
t.Fatalf("Failed to update capabilities: %v", err)
}
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.Capabilities != "Updated capabilities: can search, analyze, and write code" {
t.Errorf("Expected updated capabilities, got '%s'", retrieved.Capabilities)
}
// Update to clear capabilities
updates2 := map[string]interface{}{
"capabilities": "",
}
err = store.UpdateAssistant(id, updates2)
if err != nil {
t.Fatalf("Failed to clear capabilities: %v", err)
}
retrieved2, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved2.Capabilities != "" {
t.Errorf("Expected empty capabilities, got '%s'", retrieved2.Capabilities)
}
})
t.Run("FilterBySandbox", func(t *testing.T) {
// Create one assistant with sandbox
withSandbox := &types.AssistantModel{
Name: "Filter Sandbox Yes",
Type: "assistant",
Connector: "openai",
Share: "private",
Sandbox: &types.Sandbox{
Command: "claude",
Timeout: "5m",
},
}
idWith, err := store.SaveAssistant(withSandbox)
if err != nil {
t.Fatalf("Failed to save assistant with sandbox: %v", err)
}
// Create one assistant without sandbox
withoutSandbox := &types.AssistantModel{
Name: "Filter Sandbox No",
Type: "assistant",
Connector: "openai",
Share: "private",
}
idWithout, err := store.SaveAssistant(withoutSandbox)
if err != nil {
t.Fatalf("Failed to save assistant without sandbox: %v", err)
}
testIDs := []string{idWith, idWithout}
// Filter: sandbox=true, scoped to test IDs
trueVal := true
result, err := store.GetAssistants(types.AssistantFilter{
Page: 1,
PageSize: 100,
Sandbox: &trueVal,
AssistantIDs: testIDs,
})
if err != nil {
t.Fatalf("Failed to filter with sandbox=true: %v", err)
}
if len(result.Data) != 1 {
t.Errorf("Expected 1 result for sandbox=true, got %d", len(result.Data))
} else if result.Data[0].ID != idWith {
t.Errorf("Expected assistant %s, got %s", idWith, result.Data[0].ID)
}
// Filter: sandbox=false, scoped to test IDs
falseVal := false
result2, err := store.GetAssistants(types.AssistantFilter{
Page: 1,
PageSize: 100,
Sandbox: &falseVal,
AssistantIDs: testIDs,
})
if err != nil {
t.Fatalf("Failed to filter with sandbox=false: %v", err)
}
if len(result2.Data) != 1 {
t.Errorf("Expected 1 result for sandbox=false, got %d", len(result2.Data))
} else if result2.Data[0].ID != idWithout {
t.Errorf("Expected assistant %s, got %s", idWithout, result2.Data[0].ID)
}
t.Logf("Sandbox filter test passed: sandbox=true returned %d, sandbox=false returned %d", len(result.Data), len(result2.Data))
})
t.Run("AllNewFieldsTogether", func(t *testing.T) {
// Test assistant with all new fields together
optionalFalse := false

13
agent/store/xun/xun.go
View file

@ -168,9 +168,18 @@ func (store *Xun) getAssistantTable() string {
func (store *Xun) parseJSONFields(data map[string]interface{}, fields []string) {
for _, field := range fields {
if val := data[field]; val != nil {
if strVal, ok := val.(string); ok && strVal != "" {
var jsonStr string
switch v := val.(type) {
case string:
jsonStr = v
case []byte:
jsonStr = string(v)
default:
continue
}
if jsonStr != "" {
var parsed interface{}
if err := jsoniter.UnmarshalFromString(strVal, &parsed); err == nil {
if err := jsoniter.UnmarshalFromString(jsonStr, &parsed); err == nil {
data[field] = parsed
}
}

8
agent/types/types.go
View file

@ -1,7 +1,6 @@
package types
import (
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/yao/agent/assistant"
searchTypes "github.com/yaoapp/yao/agent/search/types"
store "github.com/yaoapp/yao/agent/store/types"
@ -23,11 +22,10 @@ type DSL struct {
// If not set, fallback to the first connector that supports the required capabilities
System *System `json:"system,omitempty" yaml:"system,omitempty"`
// Global External Settings - model capabilities, tools, etc.
// Global External Settings
// ===============================
Models map[string]openai.Capabilities `json:"models,omitempty" yaml:"models,omitempty"` // The model capabilities configuration
KB *store.KBSetting `json:"kb,omitempty" yaml:"kb,omitempty"` // The knowledge base configuration loaded from agent/kb.yml
Search *searchTypes.Config `json:"search,omitempty" yaml:"search,omitempty"` // The search configuration loaded from agent/search.yao
KB *store.KBSetting `json:"kb,omitempty" yaml:"kb,omitempty"` // The knowledge base configuration loaded from agent/kb.yml
Search *searchTypes.Config `json:"search,omitempty" yaml:"search,omitempty"` // The search configuration loaded from agent/search.yao
// Internal
// ===============================

1717
data/bindata.go
View file

File diff suppressed because one or more lines are too long

41
event/sub.go
View file

@ -54,14 +54,25 @@ func (sm *subManager) subscribe(pattern string, ch chan<- *types.Event, opts ...
return id
}
// unsubscribe removes a subscriber by ID.
// unsubscribe removes a subscriber by ID and closes its channel
// so that any goroutine blocked on `range ch` will unblock and exit.
func (sm *subManager) unsubscribe(id string) {
sm.mu.Lock()
defer sm.mu.Unlock()
entry, ok := sm.entries[id]
delete(sm.entries, id)
sm.mu.Unlock()
if ok && entry.ch != nil {
func() {
defer func() { recover() }()
close(entry.ch)
}()
}
}
// notify sends an event to all matching subscribers (non-blocking).
// Recovers from send-on-closed-channel panics that may occur if
// unsubscribe closes a channel concurrently.
func (sm *subManager) notify(ev *types.Event) {
sm.mu.RLock()
defer sm.mu.RUnlock()
@ -73,19 +84,31 @@ func (sm *subManager) notify(ev *types.Event) {
if entry.filter != nil && !entry.filter(ev) {
continue
}
select {
case entry.ch <- ev:
default:
// Subscriber chan full, skip (non-blocking)
}
func() {
defer func() { recover() }()
select {
case entry.ch <- ev:
default:
}
}()
}
}
// clear removes all subscribers. Used during Stop.
// clear removes all subscribers and closes their channels. Used during Stop.
func (sm *subManager) clear() {
sm.mu.Lock()
defer sm.mu.Unlock()
old := sm.entries
sm.entries = make(map[string]*subEntry)
sm.mu.Unlock()
for _, entry := range old {
if entry.ch != nil {
func() {
defer func() { recover() }()
close(entry.ch)
}()
}
}
}
// Subscribe dynamically subscribes to events matching the given pattern.

6
event/sub_test.go
View file

@ -164,6 +164,7 @@ func TestSubscribe_StopClearsSubscribers(t *testing.T) {
}
// drainChan reads up to n events from ch within timeout.
// Stops early if the channel is closed.
func drainChan(ch chan *types.Event, n int, timeout time.Duration) []*types.Event {
var result []*types.Event
timer := time.NewTimer(timeout)
@ -171,7 +172,10 @@ func drainChan(ch chan *types.Event, n int, timeout time.Duration) []*types.Even
for range n {
select {
case ev := <-ch:
case ev, ok := <-ch:
if !ok {
return result
}
result = append(result, ev)
case <-timer.C:
return result

27
openapi/agent/assistant.go
View file

@ -117,7 +117,7 @@ func ListAssistants(c *gin.Context) {
}
// Parse boolean filters
var builtIn, mentionable, automated *bool
var builtIn, mentionable, automated, sandbox *bool
if builtInParam := c.Query("built_in"); builtInParam != "" {
builtIn = parseBoolValue(builtInParam)
}
@ -127,6 +127,9 @@ func ListAssistants(c *gin.Context) {
if automatedParam := c.Query("automated"); automatedParam != "" {
automated = parseBoolValue(automatedParam)
}
if sandboxParam := c.Query("sandbox"); sandboxParam != "" {
sandbox = parseBoolValue(sandboxParam)
}
// Note: public and share filters are not yet supported in AssistantFilter
// They would need to be added to the store layer for proper filtering
@ -152,6 +155,7 @@ func ListAssistants(c *gin.Context) {
BuiltIn: builtIn,
Mentionable: mentionable,
Automated: automated,
Sandbox: sandbox,
})
// Apply permission-based filtering (Scope filtering)
@ -169,12 +173,19 @@ func ListAssistants(c *gin.Context) {
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)
// Convert sandbox to boolean and filter built-in sensitive fields
resp := map[string]interface{}{
"data": AssistantsToResponse(result.Data),
"total": result.Total,
"page": result.Page,
"pagesize": result.PageSize,
"pagecount": result.PageCount,
"next": result.Next,
"prev": result.Prev,
}
// Return the result with standard response format
response.RespondWithSuccess(c, response.StatusOK, result)
response.RespondWithSuccess(c, response.StatusOK, resp)
}
// GetAssistant retrieves a single assistant by ID with permission verification
@ -260,11 +271,13 @@ func GetAssistant(c *gin.Context) {
return
}
// Filter sensitive fields for built-in assistants
// Convert sandbox to boolean and filter built-in sensitive fields
hasSandbox := assistant.Sandbox != nil
FilterBuiltInAssistant(assistant)
resp := AssistantToResponse(assistant, hasSandbox)
// Return the result with standard response format
response.RespondWithSuccess(c, response.StatusOK, assistant)
response.RespondWithSuccess(c, response.StatusOK, resp)
}
// ListAssistantTags lists assistant tags with permission-based filtering

42
openapi/agent/filter.go
View file

@ -1,6 +1,8 @@
package agent
import (
"encoding/json"
"github.com/gin-gonic/gin"
"github.com/yaoapp/gou/model"
"github.com/yaoapp/xun/dbal/query"
@ -150,9 +152,49 @@ func FilterBuiltInAssistant(assistant *agenttypes.AssistantModel) {
assistant.Prompts = nil
assistant.PromptPresets = nil
assistant.Workflow = nil
assistant.Sandbox = nil
assistant.KB = nil
assistant.MCP = nil
assistant.Options = nil
assistant.Source = ""
}
}
// AssistantToResponse converts an AssistantModel to a response map,
// replacing the sandbox JSON object with a boolean indicating whether sandbox is configured.
// hasSandbox must be captured before FilterBuiltInAssistant clears the Sandbox field.
func AssistantToResponse(assistant *agenttypes.AssistantModel, hasSandbox bool) map[string]interface{} {
if assistant == nil {
return nil
}
raw, err := json.Marshal(assistant)
if err != nil {
return nil
}
var result map[string]interface{}
if err := json.Unmarshal(raw, &result); err != nil {
return nil
}
result["sandbox"] = hasSandbox
return result
}
// AssistantsToResponse converts a slice of AssistantModel to response maps,
// replacing sandbox with a boolean for each assistant.
// Captures sandbox state before filtering, then applies FilterBuiltInAssistant.
func AssistantsToResponse(assistants []*agenttypes.AssistantModel) []map[string]interface{} {
if assistants == nil {
return nil
}
result := make([]map[string]interface{}, 0, len(assistants))
for _, a := range assistants {
hasSandbox := a.Sandbox != nil
FilterBuiltInAssistant(a)
result = append(result, AssistantToResponse(a, hasSandbox))
}
return result
}

10
openapi/agent/types.go
View file

@ -12,9 +12,9 @@ var (
// availableAssistantFields defines all available fields for security filtering
availableAssistantFields = map[string]bool{
"id": true, "assistant_id": true, "type": true, "name": true, "avatar": true,
"connector": true, "description": true, "path": true, "sort": true,
"connector": true, "description": true, "capabilities": true, "path": true, "sort": true,
"built_in": true, "placeholder": true, "options": true, "prompts": true,
"workflow": true, "kb": true, "mcp": true, "tools": true, "tags": true,
"workflow": true, "sandbox": true, "kb": true, "mcp": true, "tools": true, "tags": true,
"readonly": true, "public": true, "share": true, "locales": true,
"automated": true, "mentionable": true,
"created_at": true, "updated_at": true, "deleted_at": true,
@ -23,9 +23,9 @@ var (
// defaultAssistantFields defines the default compact field list
defaultAssistantFields = []string{
"assistant_id", "type", "name", "avatar", "connector", "description",
"assistant_id", "type", "name", "avatar", "connector", "description", "capabilities",
"sort", "built_in", "tags", "readonly", "public", "share",
"automated", "mentionable", "created_at", "updated_at",
"automated", "mentionable", "sandbox", "created_at", "updated_at",
}
)
@ -59,6 +59,7 @@ type AssistantFilterParams struct {
BuiltIn *bool
Mentionable *bool
Automated *bool
Sandbox *bool
Public *bool
Share string
}
@ -79,6 +80,7 @@ func BuildAssistantFilter(params AssistantFilterParams) agenttypes.AssistantFilt
BuiltIn: params.BuiltIn,
Mentionable: params.Mentionable,
Automated: params.Automated,
Sandbox: params.Sandbox,
}
// Set default type if not specified (only when Types is also empty)

31
openapi/llm/llm.go
View file

@ -5,8 +5,6 @@ import (
"github.com/gin-gonic/gin"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/yao/agent"
agentllm "github.com/yaoapp/yao/agent/llm"
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/openapi/response"
@ -21,16 +19,6 @@ type Provider struct {
Capabilities map[string]interface{} `json:"capabilities"` // Model capabilities from connector settings
}
// getModelCapabilities returns user-defined model capabilities from agent DSL
// Returns nil if agent not initialized or no models configured
func getModelCapabilities() map[string]openai.Capabilities {
agentDSL := agent.GetAgent()
if agentDSL != nil && agentDSL.Models != nil {
return agentDSL.Models
}
return nil
}
// Attach attaches the LLM management handlers to the router with OAuth protection
func Attach(group *gin.RouterGroup, oauth oauthTypes.OAuth) {
@ -56,23 +44,15 @@ func listProviders(c *gin.Context) {
}
}
// Get user-defined model capabilities once at the start of request
modelCapabilities := getModelCapabilities()
// Get all LLM connectors from AIConnectors
// Note: All AI type connectors (openai, anthropic, fastembed) are automatically added to AIConnectors during loading
// See gou/connector/connector.go LoadSource() for details
for _, opt := range connector.AIConnectors {
connType := getConnectorType(opt.Value)
// Include OpenAI-compatible and Anthropic LLM connectors
if connType == "openai" || connType == "anthropic" {
conn, ok := connector.Connectors[opt.Value]
if !ok {
continue
}
// Get capabilities from connector settings
capabilities := getCapabilitiesWithModels(conn, modelCapabilities)
capabilities := getCapabilitiesFromConn(conn)
// Apply capability filters
if len(filters) > 0 && !matchesFilters(capabilities, filters) {
@ -110,16 +90,13 @@ func getConnectorType(id string) string {
return "unknown"
}
// getCapabilitiesWithModels extracts capabilities from connector settings
// Uses the unified capability getter from agent/llm package
// Takes modelCapabilities as parameter to avoid repeated calls to getModelCapabilities()
func getCapabilitiesWithModels(conn connector.Connector, modelCapabilities map[string]openai.Capabilities) map[string]interface{} {
// getCapabilitiesFromConn extracts capabilities from connector settings
func getCapabilitiesFromConn(conn connector.Connector) map[string]interface{} {
if conn == nil {
return nil
}
// Use unified capability getter with user-defined model capabilities
caps := agentllm.GetCapabilitiesFromConn(conn, modelCapabilities)
caps := agentllm.GetCapabilitiesFromConn(conn)
return agentllm.ToMap(caps)
}

65
openapi/tests/agent/assistant_test.go
View file

@ -264,6 +264,67 @@ func TestListAssistants(t *testing.T) {
}
})
t.Run("ListAssistantsWithSandboxFilter", func(t *testing.T) {
// Test with sandbox=true filter
req, err := http.NewRequest("GET", serverURL+baseURL+"/agent/assistants?sandbox=true&types=assistant", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
data, hasData := response["data"].([]interface{})
if hasData {
for _, item := range data {
a, ok := item.(map[string]interface{})
if ok {
sandboxVal, exists := a["sandbox"]
assert.True(t, exists, "sandbox field should be present in list response")
assert.Equal(t, true, sandboxVal, "sandbox should be true when filtering sandbox=true")
}
}
t.Logf("Successfully retrieved %d assistants with sandbox filter", len(data))
}
})
t.Run("ListAssistantsSandboxReturnsBool", func(t *testing.T) {
// Verify sandbox field is returned as boolean (not JSON object) in list response
req, err := http.NewRequest("GET", serverURL+baseURL+"/agent/assistants?pagesize=5&types=assistant", nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var response map[string]interface{}
err = json.NewDecoder(resp.Body).Decode(&response)
assert.NoError(t, err)
data, hasData := response["data"].([]interface{})
if hasData && len(data) > 0 {
a, ok := data[0].(map[string]interface{})
if ok {
sandboxVal, exists := a["sandbox"]
assert.True(t, exists, "sandbox field should be present in default list fields")
_, isBool := sandboxVal.(bool)
assert.True(t, isBool, "sandbox should be a boolean value, got %T", sandboxVal)
t.Logf("sandbox field correctly returned as bool: %v", sandboxVal)
}
}
})
t.Run("ListAssistantsWithSelectFields", func(t *testing.T) {
// Test with select parameter to limit returned fields
req, err := http.NewRequest("GET", serverURL+baseURL+"/agent/assistants?select=assistant_id,name,avatar,type", nil)
@ -1413,6 +1474,10 @@ func TestGetAssistantResponseStructure(t *testing.T) {
assert.Contains(t, responseAssistant, "name", "Assistant should have name")
assert.Contains(t, responseAssistant, "type", "Assistant should have type")
// Verify capabilities field is present in response (may be empty/null)
// capabilities is a default field that should always be returned
t.Logf("capabilities field value: %v", responseAssistant["capabilities"])
responseAssistantID := responseAssistant["assistant_id"].(string)
t.Logf("Response structure is correct for assistant: %s", responseAssistantID)
})

7
openapi/trace/events.go
View file

@ -58,13 +58,14 @@ func handleStreamMode(c *gin.Context, manager types.Manager, info *types.TraceIn
c.Header("X-Accel-Buffering", "no")
// Subscribe to trace updates
updates, err := manager.Subscribe()
updates, cancel, err := manager.Subscribe()
if err != nil {
// Send error as SSE event
fmt.Fprintf(c.Writer, "event: error\ndata: {\"error\":\"Failed to subscribe: %s\"}\n\n", err.Error())
c.Writer.Flush()
return
}
defer cancel()
// Stream events
ctx := c.Request.Context()
@ -73,22 +74,18 @@ func handleStreamMode(c *gin.Context, manager types.Manager, info *types.TraceIn
for {
select {
case <-clientGone:
// Client disconnected
return
case update, ok := <-updates:
if !ok {
// Channel closed
return
}
// Format and send SSE event
err := sendSSEEvent(c.Writer, *update)
if err != nil {
return
}
// Check if trace is complete
if update.Type == types.UpdateTypeComplete {
return
}

32
trace/subscription.go
View file

@ -2,6 +2,7 @@ package trace
import (
"fmt"
"sync"
"github.com/yaoapp/yao/event"
eventTypes "github.com/yaoapp/yao/event/types"
@ -12,13 +13,16 @@ func dedupKey(u *types.TraceUpdate) string {
return fmt.Sprintf("%s:%s:%d", u.Type, u.NodeID, u.Timestamp)
}
// Subscribe creates a new subscription for trace updates (replays all historical events from the beginning)
func (m *manager) Subscribe() (<-chan *types.TraceUpdate, error) {
// Subscribe creates a new subscription for trace updates (replays all historical events from the beginning).
// Returns the update channel and a cancel function. The caller MUST call
// cancel when done (e.g., client disconnect) to release the goroutine.
func (m *manager) Subscribe() (<-chan *types.TraceUpdate, func(), error) {
return m.subscribe(0)
}
// SubscribeFrom creates a subscription starting from a specific timestamp
func (m *manager) SubscribeFrom(since int64) (<-chan *types.TraceUpdate, error) {
// SubscribeFrom creates a subscription starting from a specific timestamp.
// Returns the update channel and a cancel function.
func (m *manager) SubscribeFrom(since int64) (<-chan *types.TraceUpdate, func(), error) {
return m.subscribe(since)
}
@ -26,12 +30,14 @@ func (m *manager) SubscribeFrom(since int64) (<-chan *types.TraceUpdate, error)
// updates, then streams live events via the event service's Subscriber.
// The subscriber is registered BEFORE reading historical state to prevent
// missing events that occur between the state snapshot and subscriber setup.
func (m *manager) subscribe(since int64) (<-chan *types.TraceUpdate, error) {
//
// The returned cancel function triggers event.Unsubscribe which closes
// liveCh, causing the goroutine to exit via `for range liveCh`.
func (m *manager) subscribe(since int64) (<-chan *types.TraceUpdate, func(), error) {
bufferSize := 1000
out := make(chan *types.TraceUpdate, bufferSize)
// Register live subscriber FIRST to avoid missing events between snapshot and subscribe.
liveCh := make(chan *eventTypes.Event, bufferSize)
traceID := m.traceID
subID := event.Subscribe("trace.*", liveCh, event.Filter(func(ev *eventTypes.Event) bool {
@ -42,19 +48,23 @@ func (m *manager) subscribe(since int64) (<-chan *types.TraceUpdate, error) {
return update.TraceID == traceID
}))
// THEN snapshot historical updates (may overlap with live events).
historical := m.stateGetUpdates(since)
// Build a set of historical event identifiers for dedup.
// Key: "type:nodeID:timestamp" is unique enough for trace events.
histSeen := make(map[string]struct{}, len(historical))
for _, u := range historical {
histSeen[dedupKey(u)] = struct{}{}
}
var cancelOnce sync.Once
cancel := func() {
cancelOnce.Do(func() {
event.Unsubscribe(subID)
})
}
go func() {
defer close(out)
defer event.Unsubscribe(subID)
defer cancel()
for _, update := range historical {
out <- update
@ -77,5 +87,5 @@ func (m *manager) subscribe(since int64) (<-chan *types.TraceUpdate, error) {
}
}()
return out, nil
return out, cancel, nil
}

3
trace/trace_autocomplete_test.go
View file

@ -331,8 +331,9 @@ func TestAutoCompleteParentEvents(t *testing.T) {
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
// Subscribe to updates
updates, err := manager.Subscribe()
updates, cancel, err := manager.Subscribe()
assert.NoError(t, err)
defer cancel()
// Collect updates in background
var receivedUpdates []*types.TraceUpdate

3
trace/trace_bench_test.go
View file

@ -266,7 +266,7 @@ func BenchmarkSubscription(b *testing.B) {
}
// Subscribe
updates, err := manager.Subscribe()
updates, cancel, err := manager.Subscribe()
if err != nil {
b.Fatalf("Failed to subscribe: %s", err.Error())
}
@ -301,6 +301,7 @@ func BenchmarkSubscription(b *testing.B) {
}
}
cancel()
trace.Release(traceID)
trace.Remove(ctx, trace.Local, traceID)
}

11
trace/trace_concurrent_test.go
View file

@ -179,6 +179,7 @@ func TestConcurrentSubscribers(t *testing.T) {
var wg sync.WaitGroup
numSubscribers := 5
subscribers := make([]<-chan *types.TraceUpdate, numSubscribers)
cancels := make([]func(), numSubscribers)
var mu sync.Mutex
for i := 0; i < numSubscribers; i++ {
@ -186,14 +187,22 @@ func TestConcurrentSubscribers(t *testing.T) {
go func(idx int) {
defer wg.Done()
sub, err := manager.Subscribe()
sub, cancelSub, err := manager.Subscribe()
assert.NoError(t, err)
mu.Lock()
subscribers[idx] = sub
cancels[idx] = cancelSub
mu.Unlock()
}(i)
}
defer func() {
for _, c := range cancels {
if c != nil {
c()
}
}
}()
wg.Wait()
// Verify all subscriptions were created

6
trace/trace_mem_test.go
View file

@ -259,10 +259,11 @@ func TestMemoryLeakComplexScenarios(t *testing.T) {
{
name: "WithSubscription",
execute: func(m types.Manager) error {
updates, err := m.Subscribe()
updates, cancel, err := m.Subscribe()
if err != nil {
return err
}
defer cancel()
// Drain updates in background with timeout
done := make(chan bool)
@ -563,7 +564,7 @@ func TestGoroutineLeak(t *testing.T) {
}
// Subscribe (creates goroutines)
updates, err := manager.Subscribe()
updates, cancel, err := manager.Subscribe()
if err != nil {
t.Errorf("Subscribe failed at iteration %d: %s", i, err.Error())
}
@ -598,6 +599,7 @@ func TestGoroutineLeak(t *testing.T) {
}
}
cancel()
trace.Release(traceID)
trace.Remove(ctx, trace.Local, traceID)
}

157
trace/trace_subscription_leak_test.go Normal file
View file

@ -0,0 +1,157 @@
package trace_test
import (
"context"
"runtime"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/event"
"github.com/yaoapp/yao/trace"
"github.com/yaoapp/yao/trace/types"
)
// stableGoroutines waits for runtime to settle and returns goroutine count.
func stableGoroutines() int {
for i := 0; i < 5; i++ {
runtime.GC()
runtime.Gosched()
time.Sleep(10 * time.Millisecond)
}
return runtime.NumGoroutine()
}
// TestLeak_SubscriptionClientDisconnect reproduces the goroutine leak that
// occurs when an SSE client subscribes to a trace and then disconnects
// without the trace ever completing (no UpdateTypeComplete sent).
//
// The subscription goroutine in subscription.go blocks on
// `for ev := range liveCh` and never exits because:
// 1. liveCh is never closed (event.Unsubscribe only deletes the map entry)
// 2. The goroutine only returns on UpdateTypeComplete
// 3. No context/cancellation mechanism exists
//
// This simulates the real-world scenario: SSE handler returns on client
// disconnect, but the subscription goroutine keeps running forever.
func TestLeak_SubscriptionClientDisconnect(t *testing.T) {
drivers := trace.GetTestDrivers()
for _, d := range drivers {
t.Run(d.Name, func(t *testing.T) {
ctx := context.Background()
before := stableGoroutines()
const numClients = 10
for i := 0; i < numClients; i++ {
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
assert.NoError(t, err)
// Client subscribes (like SSE handler calling manager.Subscribe())
updates, cancel, err := manager.Subscribe()
assert.NoError(t, err)
assert.NotNil(t, updates)
// Simulate some trace activity
_, err = manager.Add("step", types.TraceNodeOption{Label: "Processing"})
assert.NoError(t, err)
// Read a couple of events (like the SSE handler would)
timeout := time.After(500 * time.Millisecond)
drain:
for {
select {
case _, ok := <-updates:
if !ok {
break drain
}
case <-timeout:
break drain
}
}
// Client disconnects: SSE handler calls cancel (deferred).
// This triggers event.Unsubscribe which closes liveCh,
// allowing the subscription goroutine to exit.
cancel()
trace.Release(traceID)
}
// Wait for goroutines to settle
time.Sleep(1 * time.Second)
after := stableGoroutines()
leaked := after - before
t.Logf("goroutines: before=%d after=%d leaked=%d (over %d simulated client disconnects)", before, after, leaked, numClients)
// Each Subscribe() spawns a goroutine that should eventually exit.
// If it doesn't, we'll see roughly numClients leaked goroutines.
if leaked >= numClients {
t.Errorf("goroutine leak detected: %d goroutines leaked after %d client disconnects. "+
"Subscription goroutines are not cleaned up when clients disconnect without trace completion.",
leaked, numClients)
}
})
}
}
// TestLeak_SubscriptionEventServiceStop reproduces the goroutine leak when
// event.Stop() is called (e.g., during shutdown) while subscriptions are active.
//
// event.Stop() calls smgr.clear() which deletes all subscriber entries but
// does NOT close their channels, leaving goroutines blocked on `range liveCh`.
func TestLeak_SubscriptionEventServiceStop(t *testing.T) {
drivers := trace.GetTestDrivers()
for _, d := range drivers {
t.Run(d.Name, func(t *testing.T) {
ctx := context.Background()
before := stableGoroutines()
const numSubs = 5
traceID, manager, err := trace.New(ctx, d.DriverType, nil, d.DriverOptions...)
assert.NoError(t, err)
// Create multiple subscriptions (simulating multiple SSE clients)
for i := 0; i < numSubs; i++ {
_, _, err := manager.Subscribe()
assert.NoError(t, err)
}
// Simulate some activity
_, err = manager.Add("work", types.TraceNodeOption{Label: "Working"})
assert.NoError(t, err)
time.Sleep(200 * time.Millisecond)
// Stop event service (like during server shutdown)
err = event.Stop(ctx)
assert.NoError(t, err)
// Restart event service for other tests
err = event.Start()
if err != nil && err != event.ErrAlreadyStart {
t.Fatalf("Failed to restart event service: %v", err)
}
trace.Release(traceID)
time.Sleep(1 * time.Second)
after := stableGoroutines()
leaked := after - before
t.Logf("goroutines: before=%d after=%d leaked=%d (over %d subscriptions + event.Stop)", before, after, leaked, numSubs)
if leaked >= numSubs {
t.Errorf("goroutine leak detected: %d goroutines leaked after event.Stop() with %d active subscriptions. "+
"smgr.clear() does not close subscriber channels.",
leaked, numSubs)
}
})
}
}

15
trace/trace_subscription_test.go
View file

@ -24,9 +24,10 @@ func TestSubscription(t *testing.T) {
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
// Subscribe to updates
updates, err := manager.Subscribe()
updates, cancel, err := manager.Subscribe()
assert.NoError(t, err)
assert.NotNil(t, updates)
defer cancel()
// Collect updates in background
var receivedUpdates []*types.TraceUpdate
@ -146,9 +147,10 @@ func TestSubscribeFrom(t *testing.T) {
// Real scenario: User refreshes page and resumes from last known timestamp
// This should replay events from resumeTimestamp onwards
updates, err := manager.SubscribeFrom(resumeTimestamp)
updates, cancelSub, err := manager.SubscribeFrom(resumeTimestamp)
assert.NoError(t, err)
assert.NotNil(t, updates)
defer cancelSub()
// Collect updates
var receivedUpdates []*types.TraceUpdate
@ -233,14 +235,17 @@ func TestMultipleSubscribers(t *testing.T) {
defer trace.Remove(ctx, d.DriverType, traceID, d.DriverOptions...)
// Create multiple subscribers
sub1, err := manager.Subscribe()
sub1, cancel1, err := manager.Subscribe()
assert.NoError(t, err)
defer cancel1()
sub2, err := manager.Subscribe()
sub2, cancel2, err := manager.Subscribe()
assert.NoError(t, err)
defer cancel2()
sub3, err := manager.Subscribe()
sub3, cancel3, err := manager.Subscribe()
assert.NoError(t, err)
defer cancel3()
// Collect updates from all subscribers
var wg sync.WaitGroup

11
trace/types/interfaces.go
View file

@ -48,10 +48,13 @@ type Manager interface {
MarkComplete() error
// Subscription Operations
// Subscribe subscribes to trace updates (replay history + real-time)
Subscribe() (<-chan *TraceUpdate, error)
// SubscribeFrom subscribes from a specific timestamp (for resume)
SubscribeFrom(since int64) (<-chan *TraceUpdate, error)
// Subscribe subscribes to trace updates (replay history + real-time).
// Returns the update channel and a cancel function. The caller MUST call
// cancel when done (e.g., client disconnect) to release the goroutine.
Subscribe() (<-chan *TraceUpdate, func(), error)
// SubscribeFrom subscribes from a specific timestamp (for resume).
// Returns the update channel and a cancel function.
SubscribeFrom(since int64) (<-chan *TraceUpdate, func(), error)
// IsComplete checks if the trace is completed
IsComplete() bool

16
yao/models/agent/assistant.mod.yao
View file

@ -73,6 +73,15 @@
"nullable": true,
"index": true
},
{
"name": "capabilities",
"type": "string",
"label": "Capabilities",
"comment": "Assistant capabilities description, useful for Robot orchestration",
"length": 600,
"nullable": true,
"index": true
},
{
"name": "path",
"type": "string",
@ -161,6 +170,13 @@
"comment": "MCP servers available for the assistant to use",
"nullable": true
},
{
"name": "sandbox",
"type": "json",
"label": "Sandbox",
"comment": "Sandbox configuration for coding agents (command, image, timeout, etc.)",
"nullable": true
},
{
"name": "source",
"type": "text",