Merge pull request #1509 from trheyi/main

Enhance Claude integration, and fix sandbox lifecycle bugs
This commit is contained in:
Max 2026-03-30 11:53:16 +08:00 committed by GitHub
commit 255e6ec6c3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
56 changed files with 2330 additions and 532 deletions

View file

@ -171,11 +171,12 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
var v2Computer infraV2.Computer
var v2LoadingMsgID string
var v2Cfg *sandboxTypes.SandboxConfig
if ast.HasSandboxV2() {
ctx.Logger.Phase("Sandbox V2")
var err error
var v2Cleanup func()
v2Runner, v2Computer, v2Cleanup, v2LoadingMsgID, err = ast.initSandboxV2(ctx, opts)
v2Runner, v2Computer, v2Cfg, v2Cleanup, v2LoadingMsgID, err = ast.initSandboxV2(ctx, opts)
if err != nil {
ast.traceAgentFail(agentNode, err)
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
@ -189,7 +190,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
if ci.BoxID != "" {
ctx.Logger.Trace("Computer: %s", ci.BoxID)
}
ctx.Logger.Trace("Workspace: %s", ast.SandboxV2.WorkspaceID)
ctx.Logger.Trace("Workspace: %s", v2Cfg.WorkspaceID)
if conn, _, err := ast.GetConnector(ctx, opts); err == nil && conn != nil {
ctx.Logger.Trace("Connector: %s", conn.ID())
}
@ -337,6 +338,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
Handler: streamHandler,
Runner: v2Runner,
Computer: v2Computer,
Config: v2Cfg,
LoadingMsgID: v2LoadingMsgID,
Options: opts,
})

View file

@ -13,6 +13,7 @@ import (
"github.com/yaoapp/yao/agent/output/message"
sandboxv2 "github.com/yaoapp/yao/agent/sandbox/v2"
sandboxTypes "github.com/yaoapp/yao/agent/sandbox/v2/types"
store "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/config"
infraV2 "github.com/yaoapp/yao/sandbox/v2"
traceTypes "github.com/yaoapp/yao/trace/types"
@ -25,12 +26,16 @@ func (ast *Assistant) HasSandboxV2() bool {
}
// initSandboxV2 initializes the V2 sandbox: obtains a Computer, gets a Runner,
// runs Prepare, and returns the runner, computer, cleanup closure, loading
// message ID, and any error.
// runs Prepare, and returns the runner, computer, a per-request copy of the
// SandboxConfig, cleanup closure, loading message ID, and any error.
//
// A shallow copy of ast.SandboxV2 is made so that concurrent requests to the
// same assistant each get their own mutable config (Owner, ID, NodeID, etc.).
func (ast *Assistant) initSandboxV2(ctx *context.Context, opts *context.Options) (
sandboxTypes.Runner, infraV2.Computer, func(), string, error,
sandboxTypes.Runner, infraV2.Computer, *sandboxTypes.SandboxConfig, func(), string, error,
) {
cfg := ast.SandboxV2
cfgCopy := *ast.SandboxV2
cfg := &cfgCopy
manager := infraV2.M()
loadingMsg := &message.Message{
@ -47,7 +52,7 @@ func (ast *Assistant) initSandboxV2(ctx *context.Context, opts *context.Options)
conn, _, err := ast.GetConnector(ctx, opts)
if err != nil && cfg.Runner.Name != "yao" {
closeLoadingV2(ctx, loadingMsgID, "sandbox.failed")
return nil, nil, nil, "", fmt.Errorf("get connector: %w", err)
return nil, nil, nil, nil, "", fmt.Errorf("get connector: %w", err)
}
// 2. Build human-readable DisplayName from real Agent name + Workspace name.
@ -84,7 +89,7 @@ func (ast *Assistant) initSandboxV2(ctx *context.Context, opts *context.Options)
computer, identifier, err := sandboxv2.GetComputer(ctx, cfg, manager)
if err != nil {
closeLoadingV2(ctx, loadingMsgID, "sandbox.failed")
return nil, nil, nil, "", fmt.Errorf("getComputer failed: %w", err)
return nil, nil, nil, nil, "", fmt.Errorf("getComputer failed: %w", err)
}
_ = identifier
@ -93,7 +98,7 @@ func (ast *Assistant) initSandboxV2(ctx *context.Context, opts *context.Options)
if err != nil {
sandboxv2.LifecycleAction(stdCtx, cfg, computer, manager)
closeLoadingV2(ctx, loadingMsgID, "sandbox.failed")
return nil, nil, nil, "", fmt.Errorf("get runner %q: %w", cfg.Runner.Name, err)
return nil, nil, nil, nil, "", fmt.Errorf("get runner %q: %w", cfg.Runner.Name, err)
}
// 5. Resolve assistant directory and skills subdirectory.
@ -124,6 +129,7 @@ func (ast *Assistant) initSandboxV2(ctx *context.Context, opts *context.Options)
Computer: computer,
Config: cfg,
Connector: conn,
AssistantID: ast.ID,
SkillsDir: skillsDir,
AssistantDir: assistantDir,
MCPServers: mcpServers,
@ -134,7 +140,7 @@ func (ast *Assistant) initSandboxV2(ctx *context.Context, opts *context.Options)
runner.Cleanup(stdCtx, computer)
sandboxv2.LifecycleAction(stdCtx, cfg, computer, manager)
closeLoadingV2(ctx, loadingMsgID, "sandbox.failed")
return nil, nil, nil, "", fmt.Errorf("runner.Prepare: %w", err)
return nil, nil, nil, nil, "", fmt.Errorf("runner.Prepare: %w", err)
}
// Inject computer + workspace into context so Create/Next hooks
@ -148,7 +154,7 @@ func (ast *Assistant) initSandboxV2(ctx *context.Context, opts *context.Options)
sandboxv2.LifecycleAction(cleanCtx, cfg, computer, manager)
}
return runner, computer, cleanup, loadingMsgID, nil
return runner, computer, cfg, cleanup, loadingMsgID, nil
}
// sandboxV2StreamParams groups arguments for executeSandboxV2Stream.
@ -158,6 +164,7 @@ type sandboxV2StreamParams struct {
Handler message.StreamFunc
Runner sandboxTypes.Runner
Computer infraV2.Computer
Config *sandboxTypes.SandboxConfig
LoadingMsgID string
Options *context.Options
}
@ -169,13 +176,15 @@ func (ast *Assistant) executeSandboxV2Stream(
) (*context.CompletionResponse, error) {
_ = p.AgentNode
cfg := ast.SandboxV2
cfg := p.Config
manager := infraV2.M()
// Build system prompt.
// Build system prompt (parse $CTX variables the same way as buildSystemPrompts).
var systemPrompt string
if len(ast.Prompts) > 0 {
for _, pr := range ast.Prompts {
ctxVars := ast.buildContextVariables(ctx)
parsed := store.Prompts(ast.Prompts).Parse(ctxVars)
for _, pr := range parsed {
if pr.Role == "system" && pr.Content != "" {
systemPrompt = pr.Content
break
@ -199,6 +208,7 @@ func (ast *Assistant) executeSandboxV2Stream(
Computer: p.Computer,
Config: cfg,
Connector: conn,
AssistantID: ast.ID,
Messages: p.Messages,
SystemPrompt: systemPrompt,
ChatID: ctx.ChatID,

View file

@ -199,7 +199,7 @@ func loadRobotFromDB(memberID string) (*types.Robot, error) {
"id", "member_id", "team_id", "display_name", "bio",
"system_prompt", "robot_status", "autonomous_mode",
"robot_config", "robot_email", "agents", "mcp_servers",
"manager_id", "language_model",
"manager_id", "language_model", "workspace",
},
Wheres: []model.QueryWhere{
{Column: "member_id", Value: memberID},
@ -268,7 +268,7 @@ func ListRobotsFromDB(query *ListQuery) (*ListResult, error) {
"id", "member_id", "team_id", "display_name", "bio",
"system_prompt", "robot_status", "autonomous_mode",
"robot_config", "robot_email", "agents", "mcp_servers",
"language_model",
"language_model", "workspace",
},
Wheres: wheres,
Orders: orders,
@ -444,6 +444,7 @@ func CreateRobot(ctx *types.Context, req *CreateRobotRequest) (*RobotResponse, e
Agents: req.Agents,
MCPServers: req.MCPServers,
LanguageModel: req.LanguageModel,
Workspace: req.Workspace,
// Limits
CostLimit: req.CostLimit,
@ -559,6 +560,9 @@ func UpdateRobot(ctx *types.Context, memberID string, req *UpdateRobotRequest) (
if req.LanguageModel != nil {
existing.LanguageModel = *req.LanguageModel
}
if req.Workspace != nil {
existing.Workspace = *req.Workspace
}
// Limits
if req.CostLimit != nil {
@ -684,6 +688,7 @@ func recordToResponse(record *store.RobotRecord) *RobotResponse {
Agents: record.Agents,
MCPServers: record.MCPServers,
LanguageModel: record.LanguageModel,
Workspace: record.Workspace,
CostLimit: record.CostLimit,
InvitedBy: record.InvitedBy,

View file

@ -144,7 +144,7 @@ type CreateRobotRequest struct {
AutonomousMode *bool `json:"autonomous_mode,omitempty"` // Whether autonomous mode is enabled
// Communication
RobotEmail string `json:"robot_email,omitempty"` // Robot email address
RobotEmail string `json:"robot_email,omitempty"` // Deprecated: Robot email address
AuthorizedSenders interface{} `json:"authorized_senders,omitempty"` // Email whitelist (JSON array)
EmailFilterRules interface{} `json:"email_filter_rules,omitempty"` // Email filter rules (JSON array)
@ -153,6 +153,7 @@ type CreateRobotRequest struct {
Agents interface{} `json:"agents,omitempty"` // Accessible agents (JSON array)
MCPServers interface{} `json:"mcp_servers,omitempty"` // MCP servers (JSON array)
LanguageModel string `json:"language_model,omitempty"` // Language model name
Workspace string `json:"workspace,omitempty"` // Workspace ID bound to this robot
// Limits
CostLimit float64 `json:"cost_limit,omitempty"` // Monthly cost limit USD
@ -179,7 +180,7 @@ type UpdateRobotRequest struct {
AutonomousMode *bool `json:"autonomous_mode,omitempty"` // Autonomous mode
// Communication
RobotEmail *string `json:"robot_email,omitempty"` // Robot email address
RobotEmail *string `json:"robot_email,omitempty"` // Deprecated: Robot email address
AuthorizedSenders interface{} `json:"authorized_senders,omitempty"` // Email whitelist
EmailFilterRules interface{} `json:"email_filter_rules,omitempty"` // Email filter rules
@ -188,6 +189,7 @@ type UpdateRobotRequest struct {
Agents interface{} `json:"agents,omitempty"` // Accessible agents
MCPServers interface{} `json:"mcp_servers,omitempty"` // MCP servers
LanguageModel *string `json:"language_model,omitempty"` // Language model name
Workspace *string `json:"workspace,omitempty"` // Workspace ID (nil=no change, ""=unbind)
// Limits
CostLimit *float64 `json:"cost_limit,omitempty"` // Monthly cost limit USD
@ -226,6 +228,7 @@ type RobotResponse struct {
Agents interface{} `json:"agents,omitempty"`
MCPServers interface{} `json:"mcp_servers,omitempty"`
LanguageModel string `json:"language_model,omitempty"`
Workspace string `json:"workspace,omitempty"`
// Limits
CostLimit float64 `json:"cost_limit,omitempty"`

View file

@ -28,6 +28,7 @@ var memberFields = []interface{}{
"mcp_servers",
"manager_id",
"language_model",
"workspace",
}
// SetMemberModel sets the member model name

View file

@ -49,6 +49,10 @@ type AgentCaller struct {
// When non-empty, passed as opts.Connector to ast.Stream so the agent uses the Robot's model.
Connector string
// Workspace is the workspace ID bound to the Robot.
// When non-empty, injected into agentCtx.Metadata["workspace_id"] for sandbox node resolution.
Workspace string
// log is an optional structured logger; when set, Call emits agent-call logs.
log *execLogger
}
@ -444,6 +448,13 @@ func (c *AgentCaller) buildAgentContext(ctx *robottypes.Context, assistantID str
}
agentCtx.Logger = agentcontext.Noop()
if c.Workspace != "" {
if agentCtx.Metadata == nil {
agentCtx.Metadata = map[string]interface{}{}
}
agentCtx.Metadata["workspace_id"] = c.Workspace
}
kunlog.Trace("[robot-agent] context built: assistantID=%s chatID=%s contextID=%s", assistantID, c.ChatID, agentCtx.ID)
return agentCtx
}

View file

@ -43,6 +43,7 @@ func (e *Executor) RunDelivery(ctx *robottypes.Context, exec *robottypes.Executi
caller := NewAgentCaller()
caller.Connector = robot.LanguageModel
caller.Workspace = robot.Workspace
result, err := caller.CallWithMessages(ctx, agentID, userContent)
if err != nil {
return fmt.Errorf("delivery agent (%s) call failed: %w", agentID, err)

View file

@ -85,6 +85,7 @@ func (e *Executor) RunGoals(ctx *robottypes.Context, exec *robottypes.Execution,
// Call agent
caller := NewAgentCaller()
caller.Connector = robot.LanguageModel
caller.Workspace = robot.Workspace
result, err := caller.CallWithMessages(ctx, agentID, userContent)
if err != nil {
return fmt.Errorf("goals agent (%s) call failed: %w", agentID, err)

View file

@ -32,6 +32,8 @@ func (e *Executor) CallHostAgent(ctx *robottypes.Context, robot *robottypes.Robo
kunlog.Info("calling Host Agent %s for scenario=%s chatID=%s", agentID, input.Scenario, chatID)
caller := NewConversationCaller(chatID)
caller.Connector = robot.LanguageModel
caller.Workspace = robot.Workspace
result, err := caller.CallWithMessages(ctx, agentID, string(inputJSON))
if err != nil {
return nil, fmt.Errorf("host agent (%s) call failed: %w", agentID, err)

View file

@ -54,6 +54,7 @@ func (e *Executor) RunInspiration(ctx *robottypes.Context, exec *robottypes.Exec
// Call agent
caller := NewAgentCaller()
caller.Connector = robot.LanguageModel
caller.Workspace = robot.Workspace
result, err := caller.CallWithMessages(ctx, agentID, userContent)
if err != nil {
return fmt.Errorf("inspiration agent (%s) call failed: %w", agentID, err)

View file

@ -37,6 +37,13 @@ func (l *execLogger) connector() string {
return ""
}
func (l *execLogger) workspace() string {
if l.robot != nil {
return l.robot.Workspace
}
return ""
}
// ---------------------------------------------------------------------------
// P2: Task Overview
// ---------------------------------------------------------------------------
@ -51,6 +58,7 @@ func (l *execLogger) logTaskOverview(tasks []robottypes.Task) {
"phase": "tasks",
"task_count": len(tasks),
"language_model": l.connector(),
"workspace": l.workspace(),
}).Info("P2 task overview: %d tasks generated", len(tasks))
}
@ -69,6 +77,9 @@ func (l *execLogger) devTaskOverview(tasks []robottypes.Task) {
if l.connector() != "" {
sb.WriteString(fmt.Sprintf("%s Model: %s%s%s\n", w, v, l.connector(), r))
}
if l.workspace() != "" {
sb.WriteString(fmt.Sprintf("%s Workspace: %s%s%s\n", w, v, l.workspace(), r))
}
sb.WriteString(fmt.Sprintf("%s%s%s\n", w, strings.Repeat("─", 60), r))
for i, t := range tasks {
desc := t.Description

View file

@ -134,6 +134,7 @@ func (r *Runner) executeAssistantTask(task *robottypes.Task, taskCtx *RunnerCont
caller := NewAgentCaller()
caller.log = r.log
caller.Connector = r.robot.LanguageModel
caller.Workspace = r.robot.Workspace
caller.ChatID = r.chatID
messages := r.BuildAssistantMessages(task, taskCtx)

View file

@ -55,6 +55,7 @@ func (e *Executor) RunTasks(ctx *robottypes.Context, exec *robottypes.Execution,
caller := NewAgentCaller()
caller.log = newExecLogger(robot, exec.ID)
caller.Connector = robot.LanguageModel
caller.Workspace = robot.Workspace
result, err := caller.CallWithMessages(ctx, agentID, userContent)
if err != nil {
return fmt.Errorf("tasks agent (%s) call failed: %w", agentID, err)

View file

@ -434,6 +434,7 @@ func (v *Validator) validateSemantic(task *robottypes.Task, output interface{})
// Call validation agent
caller := NewAgentCaller()
caller.Connector = v.robot.LanguageModel
caller.Workspace = v.robot.Workspace
result, err := caller.CallWithMessages(v.ctx, validationAgentID, validationPrompt)
if err != nil {
return &robottypes.ValidationResult{
@ -641,6 +642,7 @@ func (av *robotAgentValidator) Validate(agentID string, output, input, criteria
// Call agent
caller := NewAgentCaller()
caller.Connector = av.v.robot.LanguageModel
caller.Workspace = av.v.robot.Workspace
callResult, err := caller.CallWithMessages(av.v.ctx, agentID, string(inputJSON))
if err != nil {
result.Passed = false

View file

@ -326,17 +326,19 @@ func (m *Manager) callHostAgentForScenario(ctx *types.Context, robot *types.Robo
Scenario: scenario,
Messages: []agentcontext.Message{{Role: "user", Content: message}},
Context: hostCtx,
}, chatID)
}, chatID, robot)
}
// callHostAgent calls the Host Agent assistant and parses output.
func (m *Manager) callHostAgent(ctx *types.Context, agentID string, input *types.HostInput, chatID string) (*types.HostOutput, error) {
func (m *Manager) callHostAgent(ctx *types.Context, agentID string, input *types.HostInput, chatID string, robot *types.Robot) (*types.HostOutput, error) {
inputJSON, err := json.Marshal(input)
if err != nil {
return nil, fmt.Errorf("failed to marshal host input: %w", err)
}
caller := standard.NewConversationCaller(chatID)
caller.Connector = robot.LanguageModel
caller.Workspace = robot.Workspace
result, err := caller.CallWithMessages(ctx, agentID, string(inputJSON))
if err != nil {
return nil, fmt.Errorf("host agent (%s) call failed: %w", agentID, err)
@ -728,16 +730,18 @@ func (m *Manager) callHostAgentForScenarioStream(ctx *types.Context, robot *type
Scenario: scenario,
Messages: []agentcontext.Message{{Role: "user", Content: msg}},
Context: hostCtx,
}, chatID, streamFn)
}, chatID, robot, streamFn)
}
func (m *Manager) callHostAgentStream(ctx *types.Context, agentID string, input *types.HostInput, chatID string, streamFn standard.StreamCallback) (*types.HostOutput, error) {
func (m *Manager) callHostAgentStream(ctx *types.Context, agentID string, input *types.HostInput, chatID string, robot *types.Robot, streamFn standard.StreamCallback) (*types.HostOutput, error) {
inputJSON, err := json.Marshal(input)
if err != nil {
return nil, fmt.Errorf("failed to marshal host input: %w", err)
}
caller := standard.NewConversationCaller(chatID)
caller.Connector = robot.LanguageModel
caller.Workspace = robot.Workspace
result, err := caller.CallWithMessagesStream(ctx, agentID, string(inputJSON), streamFn)
if err != nil {
return nil, fmt.Errorf("host agent (%s) call failed: %w", agentID, err)
@ -895,7 +899,7 @@ func (m *Manager) callHostAgentForScenarioStreamRaw(ctx *types.Context, robot *t
Scenario: scenario,
Messages: []agentcontext.Message{{Role: "user", Content: msg}},
Context: hostCtx,
}, chatID, onMessage)
}, chatID, robot, onMessage)
}
// callHostAgentStreamRaw calls the Host Agent with CUI raw message streaming.
@ -903,7 +907,7 @@ func (m *Manager) callHostAgentForScenarioStreamRaw(ctx *types.Context, robot *t
// so the frontend never sees raw decision JSON. If the final result is a decision,
// the buffered chunks are discarded and a clean reply is sent instead. If the
// result is a normal conversation turn, buffered chunks are flushed through.
func (m *Manager) callHostAgentStreamRaw(ctx *types.Context, agentID string, input *types.HostInput, chatID string, onMessage agentcontext.OnMessageFunc) (*types.HostOutput, error) {
func (m *Manager) callHostAgentStreamRaw(ctx *types.Context, agentID string, input *types.HostInput, chatID string, robot *types.Robot, onMessage agentcontext.OnMessageFunc) (*types.HostOutput, error) {
inputJSON, err := json.Marshal(input)
if err != nil {
return nil, fmt.Errorf("failed to marshal host input: %w", err)
@ -956,6 +960,8 @@ func (m *Manager) callHostAgentStreamRaw(ctx *types.Context, agentID string, inp
}
caller := standard.NewConversationCaller(chatID)
caller.Connector = robot.LanguageModel
caller.Workspace = robot.Workspace
result, err := caller.CallWithMessagesStreamRaw(ctx, agentID, string(inputJSON), wrappedOnMessage)
if err != nil {
return nil, fmt.Errorf("host agent (%s) call failed: %w", agentID, err)

View file

@ -33,7 +33,7 @@ type RobotRecord struct {
ManagerID string `json:"manager_id"` // Direct manager user_id (who manages this robot)
// Communication
RobotEmail string `json:"robot_email"` // Robot email address
RobotEmail string `json:"robot_email"` // Deprecated: Robot email address
AuthorizedSenders interface{} `json:"authorized_senders,omitempty"` // Email whitelist (JSON array)
EmailFilterRules interface{} `json:"email_filter_rules,omitempty"` // Email filter rules (JSON array)
@ -42,6 +42,7 @@ type RobotRecord struct {
Agents interface{} `json:"agents,omitempty"` // Accessible agents (JSON array)
MCPServers interface{} `json:"mcp_servers,omitempty"` // MCP servers (JSON array)
LanguageModel string `json:"language_model,omitempty"` // Language model name
Workspace string `json:"workspace,omitempty"` // Workspace ID bound to this robot
// Limits
CostLimit float64 `json:"cost_limit,omitempty"` // Monthly cost limit USD
@ -117,6 +118,7 @@ var robotFields = []interface{}{
"agents",
"mcp_servers",
"language_model",
"workspace",
// Limits
"cost_limit",
@ -435,6 +437,9 @@ func (s *RobotStore) recordToMap(record *RobotRecord) map[string]interface{} {
if record.LanguageModel != "" {
data["language_model"] = record.LanguageModel
}
if record.Workspace != "" {
data["workspace"] = record.Workspace
}
// Limits
if record.CostLimit > 0 {
@ -547,6 +552,9 @@ func (s *RobotStore) mapToRecord(row map[string]interface{}) (*RobotRecord, erro
if v, ok := row["language_model"].(string); ok {
record.LanguageModel = v
}
if v, ok := row["workspace"].(string); ok {
record.Workspace = v
}
// Limits
if v := row["cost_limit"]; v != nil {
@ -596,6 +604,8 @@ func (r *RobotRecord) ToRobot() (*types.Robot, error) {
SystemPrompt: r.SystemPrompt,
AutonomousMode: r.AutonomousMode,
RobotEmail: r.RobotEmail,
LanguageModel: r.LanguageModel,
Workspace: r.Workspace,
}
// Parse robot_status
@ -678,6 +688,8 @@ func FromRobot(robot *types.Robot) *RobotRecord {
RobotStatus: string(robot.Status),
AutonomousMode: robot.AutonomousMode,
RobotEmail: robot.RobotEmail,
LanguageModel: robot.LanguageModel,
Workspace: robot.Workspace,
MemberType: "robot",
Status: "active",
}

View file

@ -21,8 +21,9 @@ type Robot struct {
SystemPrompt string `json:"system_prompt"`
Status RobotStatus `json:"robot_status"`
AutonomousMode bool `json:"autonomous_mode"`
RobotEmail string `json:"robot_email"` // Robot's email address for sending emails
RobotEmail string `json:"robot_email"` // Deprecated: Robot's email address for sending emails
LanguageModel string `json:"language_model"` // LLM connector override (from __yao.member.language_model)
Workspace string `json:"workspace"` // Workspace ID bound to this robot (nullable in DB)
// Manager info (from __yao.member)
ManagerID string `json:"manager_id"` // Direct manager user_id (who manages this robot)
@ -509,6 +510,7 @@ func NewRobotFromMap(m map[string]interface{}) (*Robot, error) {
ManagerID: getString(m, "manager_id"),
ManagerEmail: getString(m, "manager_email"),
LanguageModel: getString(m, "language_model"),
Workspace: getString(m, "workspace"),
}
// Parse robot_status

View file

@ -11,6 +11,7 @@ import (
"github.com/google/uuid"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/store"
"github.com/yaoapp/kun/str"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/sandbox/v2/types"
infra "github.com/yaoapp/yao/sandbox/v2"
@ -53,13 +54,10 @@ type command struct {
workDir string
}
func (r *ClaudeRunner) buildCommand(ctx context.Context, req *types.StreamRequest, p platform) command {
func (r *Runner) buildCommand(ctx context.Context, req *types.StreamRequest, p platform) command {
computer := req.Computer
workDir := computer.GetWorkDir()
assistantID := ""
if req.Config != nil {
assistantID = req.Config.ID
}
assistantID := req.AssistantID
chatID := req.ChatID
var isContinuation bool
@ -110,14 +108,17 @@ func buildEnv(req *types.StreamRequest, p platform) map[string]string {
for k, v := range p.HomeEnv(workDir) {
env[k] = v
}
env["WORKDIR"] = workDir
assistantID := ""
if req.Config != nil {
assistantID = req.Config.ID
}
assistantID := req.AssistantID
if assistantID != "" {
configDir := p.PathJoin(workDir, ".yao", "assistants", assistantID)
env["CLAUDE_CONFIG_DIR"] = configDir
env["CTX_ASSISTANT_ID"] = assistantID
// CTX_SKILLS_DIR: absolute path to the skills directory inside the sandbox.
// Use this in skill scripts instead of constructing the path manually,
// so it works correctly on all platforms (Linux, macOS, Windows).
env["CTX_SKILLS_DIR"] = p.PathJoin(workDir, ".yao", "assistants", assistantID, "skills")
}
if req.Connector != nil {
@ -165,7 +166,7 @@ func buildEnv(req *types.StreamRequest, p platform) map[string]string {
if req.Config != nil && len(req.Config.Secrets) > 0 {
for k, v := range req.Config.Secrets {
env[k] = v
env[k] = str.EnvVar(v)
}
}
@ -181,7 +182,7 @@ func buildEnv(req *types.StreamRequest, p platform) map[string]string {
return env
}
func buildArgs(req *types.StreamRequest, r *ClaudeRunner, p platform, isContinuation bool, assistantID, chatID string) []string {
func buildArgs(req *types.StreamRequest, r *Runner, p platform, isContinuation bool, assistantID, chatID string) []string {
var args []string
permMode := ""
@ -248,12 +249,18 @@ func buildSandboxEnvPrompt(p platform, workDir string) string {
shellNote := p.EnvPromptNote()
envVarSyntax := "$VAR_NAME"
if osName == "windows" {
envVarSyntax = "$env:VAR_NAME"
}
return fmt.Sprintf(`## Sandbox Environment
- **Operating System**: %[2]s
- **Shell**: %[3]s
- **Working Directory**: %[1]s
- **File Access**: You have full read/write access to %[1]s%[4]s
- **File Access**: You have full read/write access to %[1]s
- **Environment variable syntax**: `+"`%[5]s`"+` (e.g. `+"`$CTX_SKILLS_DIR`"+` on POSIX, `+"`$env:CTX_SKILLS_DIR`"+` on Windows)%[4]s
## User Attachments
@ -261,7 +268,7 @@ User-uploaded files (images, documents, code files, etc.) are placed in %[1]s/.a
Each chat session has its own subdirectory to avoid conflicts.
When the user references an attached file, read it from this directory using the Read or Bash tool.
For image files, you can view them directly as Claude supports vision on local files.
`, workDir, osName, shell, shellNote)
`, workDir, osName, shell, shellNote, envVarSyntax)
}
func hasExistingSession(ctx context.Context, computer infra.Computer, p platform, assistantID string) bool {

View file

@ -33,7 +33,8 @@ func TestBuildEnv_HomeEnv(t *testing.T) {
func TestBuildEnv_ConfigDirIsolation(t *testing.T) {
req := &types.StreamRequest{
Config: &types.SandboxConfig{ID: "my-assistant"},
Config: &types.SandboxConfig{ID: "my-assistant"},
AssistantID: "my-assistant",
}
req.Computer = newFakeComputer("/workspace")
p := testPlatform()
@ -89,7 +90,7 @@ func TestBuildEnv_Secrets(t *testing.T) {
func TestBuildArgs_Default(t *testing.T) {
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
req.Computer = newFakeComputer("/workspace")
r := &ClaudeRunner{}
r := &Runner{}
p := testPlatform()
args := buildArgs(req, r, p, false, "", "")
@ -103,7 +104,7 @@ func TestBuildArgs_Default(t *testing.T) {
func TestBuildArgs_Continuation(t *testing.T) {
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
req.Computer = newFakeComputer("/workspace")
r := &ClaudeRunner{}
r := &Runner{}
p := testPlatform()
args := buildArgs(req, r, p, true, "", "")
@ -121,7 +122,7 @@ func TestBuildArgs_PermissionMode(t *testing.T) {
},
}
req.Computer = newFakeComputer("/workspace")
r := &ClaudeRunner{}
r := &Runner{}
p := testPlatform()
args := buildArgs(req, r, p, false, "", "")
@ -132,7 +133,7 @@ func TestBuildArgs_PermissionMode(t *testing.T) {
func TestBuildArgs_MCP(t *testing.T) {
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
req.Computer = newFakeComputer("/workspace")
r := &ClaudeRunner{hasMCP: true, mcpToolPattern: "mcp__yao__*"}
r := &Runner{hasMCP: true, mcpToolPattern: "mcp__yao__*"}
p := testPlatform()
args := buildArgs(req, r, p, false, "test-assistant", "")
@ -163,7 +164,7 @@ func TestBuildArgs_WhitelistOptions(t *testing.T) {
},
}
req.Computer = newFakeComputer("/workspace")
r := &ClaudeRunner{}
r := &Runner{}
p := testPlatform()
args := buildArgs(req, r, p, false, "", "")
@ -378,7 +379,7 @@ func TestSanitizeSessionName_Empty(t *testing.T) {
func TestBuildArgs_SessionID_NewSession(t *testing.T) {
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
req.Computer = newFakeComputer("/workspace")
r := &ClaudeRunner{}
r := &Runner{}
p := testPlatform()
args := buildArgs(req, r, p, false, "asst-1", "robot_m1_e1")
@ -402,7 +403,7 @@ func TestBuildArgs_SessionID_NewSession(t *testing.T) {
func TestBuildArgs_SessionID_Continuation(t *testing.T) {
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
req.Computer = newFakeComputer("/workspace")
r := &ClaudeRunner{}
r := &Runner{}
p := testPlatform()
args := buildArgs(req, r, p, true, "asst-1", "robot_m1_e1")
@ -425,7 +426,7 @@ func TestBuildArgs_SessionID_Continuation(t *testing.T) {
func TestBuildArgs_EmptyChatID_Continuation(t *testing.T) {
req := &types.StreamRequest{Config: &types.SandboxConfig{}}
req.Computer = newFakeComputer("/workspace")
r := &ClaudeRunner{}
r := &Runner{}
p := testPlatform()
args := buildArgs(req, r, p, true, "", "")

View file

@ -172,6 +172,30 @@ func (p *streamParser) closeTextMessage() {
}
}
// closeCurrentTool closes the in-flight streaming tool message (if any),
// flushing its accumulated input and emitting message_end. This must be
// called before opening a new message group so that the downstream handler
// never sees interleaved message_start/message_end pairs.
func (p *streamParser) closeCurrentTool() {
if p.curTool == nil {
return
}
toolID := p.curTool.id
inputStr := p.curTool.inputJSON.String()
if inputStr != "" {
p.toolInputs[toolID] = inputStr
summary := extractSummary(p.curTool.name, inputStr)
if summary != "" {
p.toolSummaries[toolID] = summary
p.emitExecute(map[string]any{
"summary": summary,
})
}
}
p.endMessage()
p.curTool = nil
}
func (p *streamParser) ensureTextMessage() (stopped bool) {
if !p.textActive {
_, stopped = p.beginMessage("text")
@ -318,22 +342,7 @@ func (p *streamParser) onContentBlockStart(event map[string]any) (stopped bool)
}
func (p *streamParser) onContentBlockStop() (stopped bool) {
if p.curTool != nil {
toolID := p.curTool.id
inputStr := p.curTool.inputJSON.String()
if inputStr != "" {
p.toolInputs[toolID] = inputStr
summary := extractSummary(p.curTool.name, inputStr)
if summary != "" {
p.toolSummaries[toolID] = summary
p.emitExecute(map[string]any{
"summary": summary,
})
}
}
p.endMessage()
p.curTool = nil
}
p.closeCurrentTool()
return false
}
@ -417,6 +426,7 @@ func (p *streamParser) handleAssistant(msg map[string]any) (stopped bool) {
}
p.closeTextMessage()
p.closeCurrentTool()
toolName, _ := ci["name"].(string)
if toolID == "" {
@ -488,12 +498,17 @@ func (p *streamParser) handleUser(msg map[string]any) (stopped bool) {
continue
}
// Close any open text message before opening an execute message,
// otherwise textActive stays true while currentGroupID gets
// overwritten by the execute message lifecycle, causing subsequent
// text chunks to be emitted without a message_id.
// Close any open text message before opening an execute message.
p.closeTextMessage()
// When Claude CLI executes tools in parallel, tool_result messages
// can arrive while a new tool_use is still streaming. The downstream
// handler (stream.go) tracks only a single currentGroupID, so we
// must close the in-flight streaming tool message before opening
// the result message — otherwise the message_start/message_end
// pairs become interleaved and chunks lose their message_id.
p.closeCurrentTool()
toolUseID, _ := ci["tool_use_id"].(string)
content := ci["content"]
isError, _ := ci["is_error"].(bool)

View file

@ -15,8 +15,8 @@ import (
infra "github.com/yaoapp/yao/sandbox/v2"
)
// ClaudeRunner implements the Runner interface for Claude CLI (mode=cli).
type ClaudeRunner struct {
// Runner implements the sandbox Runner interface for Claude CLI (mode=cli).
type Runner struct {
mode string
hasMCP bool
mcpToolPattern string
@ -25,21 +25,22 @@ type ClaudeRunner struct {
logger *agentContext.RequestLogger
}
// New creates a new ClaudeRunner.
func New() *ClaudeRunner {
return &ClaudeRunner{mode: "cli"}
// New creates a new Runner.
func New() *Runner {
return &Runner{mode: "cli"}
}
func (r *ClaudeRunner) Name() string { return "claude" }
// Name returns the runner identifier.
func (r *Runner) Name() string { return "claude" }
// Prepare executes user-defined and runner-specific prepare steps.
func (r *ClaudeRunner) Prepare(ctx context.Context, req *types.PrepareRequest) error {
func (r *Runner) Prepare(ctx context.Context, req *types.PrepareRequest) error {
r.mode = req.Config.Runner.Mode
if r.mode == "" {
r.mode = "cli"
}
assistantID := req.Config.ID
assistantID := req.AssistantID
prefix := ".yao/assistants/" + assistantID
if assistantID == "" {
prefix = ".claude"
@ -70,7 +71,7 @@ func (r *ClaudeRunner) Prepare(ctx context.Context, req *types.PrepareRequest) e
}
if req.RunSteps != nil && len(steps) > 0 {
if err := req.RunSteps(ctx, steps, req.Computer, req.Config.ID, req.ConfigHash, req.AssistantDir); err != nil {
if err := req.RunSteps(ctx, steps, req.Computer, req.AssistantID, req.ConfigHash, req.AssistantDir); err != nil {
return fmt.Errorf("claude prepare steps: %w", err)
}
}
@ -79,7 +80,7 @@ func (r *ClaudeRunner) Prepare(ctx context.Context, req *types.PrepareRequest) e
}
// Stream executes the Claude CLI and streams output to handler.
func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, handler message.StreamFunc) error {
func (r *Runner) Stream(ctx context.Context, req *types.StreamRequest, handler message.StreamFunc) error {
computer := req.Computer
if computer == nil {
return fmt.Errorf("computer is nil")
@ -111,12 +112,17 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han
chatID := req.ChatID
r.lastChatID = chatID
assistantID := ""
if req.Config != nil {
assistantID = req.Config.ID
}
assistantID := req.AssistantID
log.Trace("[claude-runner] Stream started: assistantID=%s chatID=%s promptLen=%d", assistantID, chatID, len(cmd.shell))
r.logger.Debug("env vars passed to session (%d total):", len(cmd.env))
for k, v := range cmd.env {
if strings.HasPrefix(k, "CTX_") || k == "CLAUDE_CONFIG_DIR" || k == "HOME" || k == "WORKDIR" {
r.logger.Debug(" %s=%s", k, v)
} else {
r.logger.Debug(" %s=(set, len=%d)", k, len(v))
}
}
sess, err := startSession(ctx, computer, p, cmd, chatID, r.logger)
if err != nil {
@ -132,10 +138,6 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han
if completed {
sess.shutdown()
if chatID != "" {
assistantID := ""
if req.Config != nil {
assistantID = req.Config.ID
}
storeKey := "claude-session:" + assistantID + ":" + chatID
sessionUUID := chatIDToSessionUUID(assistantID, chatID)
markChatSession(storeKey, sessionUUID, 90*24*time.Hour)
@ -146,7 +148,7 @@ func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, han
// Cleanup kills any remaining claude processes. If the stream completed
// normally (received "result"), child processes are preserved.
func (r *ClaudeRunner) Cleanup(ctx context.Context, computer infra.Computer) error {
func (r *Runner) Cleanup(ctx context.Context, computer infra.Computer) error {
if computer == nil {
return nil
}

View file

@ -60,6 +60,8 @@ func ResolveNodeID(ctx *agentContext.Context, cfg *types.SandboxConfig, manager
if err == nil && wsNode != "" {
log.Trace("[sandbox/v2] ResolveNodeID: workspace %s -> node %s", workspaceID, wsNode)
computerID = wsNode
} else if err != nil {
log.Warn("[sandbox/v2] ResolveNodeID: workspace %s not found or deleted, falling back to auto-select: %v", workspaceID, err)
}
}
@ -132,6 +134,8 @@ func GetComputer(ctx *agentContext.Context, cfg *types.SandboxConfig, manager *i
log.Trace("[sandbox/v2] workspace %s bound to node %s overrides computer_id %s", workspaceID, wsNode, computerID)
}
computerID = wsNode
} else if err != nil {
log.Warn("[sandbox/v2] GetComputer: workspace %s not found or deleted, falling back: %v", workspaceID, err)
}
}
@ -280,24 +284,30 @@ func resolveBox(
// LifecycleAction performs the post-request lifecycle operation based on policy.
// Called in defer after executeSandboxStream completes.
//
// NOTE: cfg.ID must NOT be used here — it is a shared mutable field on the
// Assistant struct and is overwritten by concurrent requests. The authoritative
// box ID is computer.ComputerInfo().BoxID, which is set once when the Box is
// created and never changes.
func LifecycleAction(ctx context.Context, cfg *types.SandboxConfig, computer infra.Computer, manager *infra.Manager) {
if computer == nil || cfg == nil {
return
}
info := computer.ComputerInfo()
boxID := info.BoxID // use the box's own immutable ID, not cfg.ID
switch cfg.Lifecycle {
case "oneshot":
if info.Kind == "box" && manager != nil {
if err := manager.Remove(ctx, cfg.ID); err != nil {
log.Trace("[sandbox/v2] oneshot remove %s: %v", cfg.ID, err)
if err := manager.Remove(ctx, boxID); err != nil {
log.Trace("[sandbox/v2] oneshot remove %s: %v", boxID, err)
}
}
case "session", "longrunning":
if info.Kind == "box" && manager != nil {
manager.Heartbeat(cfg.ID, false, 0) // active=false: request finished, start idle timer
manager.Heartbeat(boxID, false, 0) // active=false: request finished, start idle timer
}
case "persistent":

View file

@ -2,22 +2,14 @@ package sandboxv2
import (
"fmt"
"os"
"strings"
"time"
"github.com/yaoapp/kun/str"
"github.com/yaoapp/yao/agent/sandbox/v2/types"
infra "github.com/yaoapp/yao/sandbox/v2"
)
// resolveEnvRef resolves $ENV.XXX references to os.Getenv("XXX").
func resolveEnvRef(value string) string {
if strings.HasPrefix(value, "$ENV.") {
return os.Getenv(value[5:])
}
return value
}
// BuildCreateOptions converts a SandboxConfig into the V2 infrastructure
// CreateOptions. Connector config injection is handled separately via the
// a2o HTTP API (POST /config) after the container starts.
@ -120,10 +112,10 @@ func BuildCreateOptions(cfg *types.SandboxConfig, identifier, ownerID, workspace
if envSize > 0 {
opts.Env = make(map[string]string, envSize)
for k, v := range cfg.Environment {
opts.Env[k] = resolveEnvRef(v)
opts.Env[k] = str.EnvVar(v)
}
for k, v := range cfg.Secrets {
opts.Env[k] = resolveEnvRef(v)
opts.Env[k] = str.EnvVar(v)
}
}
@ -135,7 +127,7 @@ func BuildCreateOptions(cfg *types.SandboxConfig, identifier, ownerID, workspace
if cfg.Computer.VNC.Enabled {
opts.Env["VNC_ENABLED"] = "true"
if cfg.Computer.VNC.Password != "" {
opts.Env["VNC_PASSWORD"] = resolveEnvRef(cfg.Computer.VNC.Password)
opts.Env["VNC_PASSWORD"] = str.EnvVar(cfg.Computer.VNC.Password)
}
if cfg.Computer.VNC.Resolution != "" {
opts.Env["VNC_RESOLUTION"] = cfg.Computer.VNC.Resolution

View file

@ -93,6 +93,7 @@ func runFileStep(ws workspace.FS, step types.PrepareStep) error {
return fmt.Errorf("file step requires workspace")
}
step.Path = expandTilde(step.Path)
dir := path.Dir(step.Path)
if dir != "." && dir != "/" {
if err := ws.MkdirAll(dir, 0755); err != nil {
@ -126,8 +127,9 @@ func runCopyStep(ws workspace.FS, step types.PrepareStep, assistantDir string) e
src = "local:///" + pathpkg.Join(assistantDir, src)
}
if _, err := ws.Copy(src, step.Dst); err != nil {
return fmt.Errorf("copy %s -> %s: %w", src, step.Dst, err)
dst := expandTilde(step.Dst)
if _, err := ws.Copy(src, dst); err != nil {
return fmt.Errorf("copy %s -> %s: %w", src, dst, err)
}
return nil
}
@ -136,6 +138,19 @@ func isHostURI(s string) bool {
return strings.HasPrefix(s, "local:///") || strings.HasPrefix(s, "tmp:///")
}
// expandTilde replaces a leading "~/" with the empty string so the path
// becomes relative to the workspace root (which is HOME inside the sandbox).
// Paths without "~/" are returned unchanged.
func expandTilde(p string) string {
if strings.HasPrefix(p, "~/") {
return p[2:]
}
if p == "~" {
return "."
}
return p
}
func runExecStep(ctx context.Context, computer infra.Computer, step types.PrepareStep) error {
if step.Cmd == "" {
return fmt.Errorf("exec step requires cmd")

View file

@ -4,7 +4,6 @@ import (
"context"
"errors"
"fmt"
"log"
"time"
agentContext "github.com/yaoapp/yao/agent/context"
@ -24,9 +23,11 @@ type ExecuteRequest struct {
LoadingMsgID string
}
// ExecuteSandboxStream is the V2 replacement for executeSandboxStream.
// It calls runner.Stream, handles interrupts, and performs cleanup/lifecycle
// in defer.
// ExecuteSandboxStream runs runner.Stream and bridges agentContext interrupts.
//
// Cleanup (runner.Cleanup + LifecycleAction) is NOT performed here; the caller
// (agent.go sandboxCleanup closure) is responsible for all lifecycle management
// so that cleanup happens exactly once regardless of code path.
func ExecuteSandboxStream(
ctx *agentContext.Context,
req *ExecuteRequest,
@ -38,42 +39,6 @@ func ExecuteSandboxStream(
}
stdCtx := ctx.Context
panicked := true // Assume panic; set false on normal exit.
// Resolve stop timeout from config (default 2s).
stopTimeout := 2 * time.Second
if req.Config != nil && req.Config.StopTimeout != "" {
if d, err := time.ParseDuration(req.Config.StopTimeout); err == nil {
stopTimeout = d
}
}
// Panic recovery (registered first, executes last in LIFO order).
defer func() {
if r := recover(); r != nil {
log.Printf("[sandbox/v2] panic in stream: %v", r)
cleanCtx, cancel := context.WithTimeout(context.Background(), stopTimeout)
defer cancel()
req.Runner.Cleanup(cleanCtx, req.Computer)
LifecycleAction(cleanCtx, req.Config, req.Computer, req.Manager)
}
}()
// Lifecycle action (registered second, executes second-to-last).
defer func() {
if !panicked {
LifecycleAction(stdCtx, req.Config, req.Computer, req.Manager)
}
}()
// Runner cleanup (registered last, executes first).
defer func() {
if !panicked {
cleanCtx, cancel := context.WithTimeout(context.Background(), stopTimeout)
defer cancel()
req.Runner.Cleanup(cleanCtx, req.Computer)
}
}()
// Build a cancellable runnerCtx that bridges agentContext interrupts.
runnerCtx, cancelRunner := context.WithCancel(stdCtx)
@ -144,8 +109,6 @@ func ExecuteSandboxStream(
closeLoading(ctx, req.LoadingMsgID)
}
panicked = false // Normal exit reached.
if err != nil {
if errors.Is(err, context.Canceled) {
return nil, err

View file

@ -38,6 +38,7 @@ type PrepareRequest struct {
Computer infra.Computer
Config *SandboxConfig
Connector connector.Connector
AssistantID string // the assistant's own ID (e.g. "yao/postman")
SkillsDir string
AssistantDir string // absolute host path to the assistant source directory
MCPServers []MCPServer
@ -50,6 +51,7 @@ type StreamRequest struct {
Computer infra.Computer
Config *SandboxConfig
Connector connector.Connector
AssistantID string // the assistant's own ID (e.g. "yao/postman")
Messages []agentContext.Message
SystemPrompt string
ChatID string

View file

@ -4,6 +4,7 @@ import (
"github.com/spf13/cobra"
"github.com/yaoapp/kun/maps"
"github.com/yaoapp/kun/utils"
"github.com/yaoapp/yao/commercial"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/engine"
"github.com/yaoapp/yao/share"
@ -15,11 +16,13 @@ var inspectCmd = &cobra.Command{
Long: L("Show app configure"),
Run: func(cmd *cobra.Command, args []string) {
Boot()
commercial.Load(config.Conf.Root, "yao")
engine.InspectExtTools()
res := maps.Map{
"version": share.VERSION,
"config": config.Conf,
}
res["license"] = commercial.License
if share.Tools != nil {
res["tools"] = share.Tools
}

View file

@ -36,8 +36,8 @@ var migrateCmd = &cobra.Command{
exception.New(L("Migrate is not allowed on production mode."), 403).Throw()
}
// 加载数据模型
loadWarnings, err := engine.Load(config.Conf, engine.LoadOption{Action: "migrate"})
// 仅加载 Application、DB 连接和 Model含自动 migrate不启动完整 Engine
loadWarnings, err := engine.LoadForMigrate(config.Conf)
if err != nil {
fmt.Println(color.RedString(L("Fatal: %s"), err.Error()))
os.Exit(1)

View file

@ -0,0 +1,121 @@
// Command generate-test-cert creates a test root CA and a license certificate
// for local development and testing.
//
// Usage:
//
// go run ./commercial/cmd/generate-test-cert \
// -out-cert /path/to/yao-dev-app/license.pem \
// -out-root-ca ./commercial/roots/root-ca-1.pem
package main
import (
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/x509"
"crypto/x509/pkix"
"encoding/pem"
"flag"
"fmt"
"math/big"
"os"
"time"
"github.com/yaoapp/yao/commercial"
)
func main() {
outCert := flag.String("out-cert", "license.pem", "path to write the license certificate PEM")
outRootCA := flag.String("out-root-ca", "", "path to write the root CA PEM (optional; for injecting into commercial/roots/)")
flag.Parse()
rootKey, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
if err != nil {
fmt.Fprintf(os.Stderr, "generate root key: %v\n", err)
os.Exit(1)
}
rootSerial, _ := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
rootTmpl := &x509.Certificate{
SerialNumber: rootSerial,
Subject: pkix.Name{
CommonName: "Yao Test Root CA",
Organization: []string{"Infinite Wisdom Software"},
Country: []string{"CN"},
},
NotBefore: time.Now().Add(-1 * time.Hour),
NotAfter: time.Now().Add(10 * 365 * 24 * time.Hour),
KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
BasicConstraintsValid: true,
IsCA: true,
MaxPathLen: 1,
}
rootDER, err := x509.CreateCertificate(rand.Reader, rootTmpl, rootTmpl, &rootKey.PublicKey, rootKey)
if err != nil {
fmt.Fprintf(os.Stderr, "create root cert: %v\n", err)
os.Exit(1)
}
rootCert, _ := x509.ParseCertificate(rootDER)
rootPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: rootDER})
leafKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if err != nil {
fmt.Fprintf(os.Stderr, "generate leaf key: %v\n", err)
os.Exit(1)
}
leafTmpl := &x509.Certificate{
SerialNumber: big.NewInt(20001),
Subject: pkix.Name{
CommonName: "Yao Dev App",
Organization: []string{"Dev Testing"},
Country: []string{"CN"},
},
EmailAddresses: []string{"dev@yaoapps.com"},
NotBefore: time.Now().Add(-1 * time.Hour),
NotAfter: time.Now().Add(365 * 24 * time.Hour),
KeyUsage: x509.KeyUsageDigitalSignature,
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageAny},
ExtraExtensions: []pkix.Extension{
{Id: commercial.OIDProduct, Value: []byte("yao,tai")},
{Id: commercial.OIDEdition, Value: []byte("enterprise")},
{Id: commercial.OIDMaxUsers, Value: []byte("0")},
{Id: commercial.OIDMaxTaiNodes, Value: []byte("0")},
{Id: commercial.OIDMaxAgents, Value: []byte("0")},
{Id: commercial.OIDMaxSandboxes, Value: []byte("0")},
{Id: commercial.OIDMaxAPIRPM, Value: []byte("0")},
{Id: commercial.OIDMaxStorageGB, Value: []byte("0")},
{Id: commercial.OIDAllowBrandingRemoval, Value: []byte("true")},
{Id: commercial.OIDAllowWhiteLabel, Value: []byte("true")},
{Id: commercial.OIDAllowMultiTenant, Value: []byte("true")},
{Id: commercial.OIDAllowCustomDomain, Value: []byte("true")},
{Id: commercial.OIDAllowHostExec, Value: []byte("true")},
{Id: commercial.OIDAllowSSO, Value: []byte("true")},
{Id: commercial.OIDSupportLevel, Value: []byte("dedicated")},
},
}
leafDER, err := x509.CreateCertificate(rand.Reader, leafTmpl, rootCert, &leafKey.PublicKey, rootKey)
if err != nil {
fmt.Fprintf(os.Stderr, "create leaf cert: %v\n", err)
os.Exit(1)
}
leafPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: leafDER})
if err := os.WriteFile(*outCert, leafPEM, 0644); err != nil {
fmt.Fprintf(os.Stderr, "write cert: %v\n", err)
os.Exit(1)
}
fmt.Printf("Wrote license certificate to %s\n", *outCert)
if *outRootCA != "" {
if err := os.WriteFile(*outRootCA, rootPEM, 0644); err != nil {
fmt.Fprintf(os.Stderr, "write root CA: %v\n", err)
os.Exit(1)
}
fmt.Printf("Wrote root CA to %s\n", *outRootCA)
}
}

296
commercial/commercial.go Normal file
View file

@ -0,0 +1,296 @@
package commercial
import (
"crypto/sha256"
"crypto/x509"
"encoding/asn1"
"fmt"
"log"
"os"
"path/filepath"
"strconv"
"strings"
"time"
)
const envLicenseCert = "YAO_LICENSE_CERT"
// Load discovers and verifies the commercial license certificate,
// writing the result into the global License variable.
// It never returns an error — failures degrade to community defaults.
func Load(appRoot, product string) {
License = DefaultLicense()
License.Product = []string{product}
pemData, source := findCert(appRoot)
if pemData == nil {
log.Printf("[License] No license certificate found, running with community defaults")
return
}
info, err := verify(pemData, product)
if err != nil {
License.Source = source
License.Error = err.Error()
log.Printf("[License] %v — running with community defaults", err)
return
}
info.Source = source
info.LoadedAt = time.Now().Unix()
License = *info
if License.Valid {
remaining := time.Until(time.Unix(License.NotAfter, 0))
log.Printf("[License] Loaded: %s (%s) — valid until %s",
License.LicenseeName, License.Edition,
time.Unix(License.NotAfter, 0).UTC().Format("2006-01-02"))
if remaining < 90*24*time.Hour {
log.Printf("[License] WARNING: Certificate expires in %d days", int(remaining.Hours()/24))
}
}
}
// findCert locates the license PEM data.
// Search order:
// 1. YAO_LICENSE_CERT env (PEM content or file path)
// 2. <appRoot>/license.pem
// 3. <appRoot>/certs/license.pem
func findCert(appRoot string) (pemData []byte, source string) {
if v := os.Getenv(envLicenseCert); v != "" {
if strings.HasPrefix(v, "-----BEGIN") {
return []byte(v), "env"
}
data, err := os.ReadFile(v)
if err == nil {
return data, "env"
}
log.Printf("[License] env %s points to unreadable file: %v", envLicenseCert, err)
}
candidates := []string{
filepath.Join(appRoot, "license.pem"),
filepath.Join(appRoot, "certs", "license.pem"),
}
for _, path := range candidates {
data, err := os.ReadFile(path)
if err == nil {
return data, "file"
}
}
return nil, "none"
}
// verify parses PEM data, validates the certificate chain against built-in
// roots, checks revocation, time validity, product scope, and extracts
// custom extension fields.
func verify(pemData []byte, product string) (*LicenseInfo, error) {
certs, err := ParsePEMChain(pemData)
if err != nil {
return nil, fmt.Errorf("parse PEM: %w", err)
}
if len(certs) == 0 {
return nil, fmt.Errorf("no certificates found in PEM data")
}
leaf := certs[0]
pool := RootPool()
if pool == nil {
return nil, fmt.Errorf("no root certificates available (development build)")
}
// Verify the trust chain with a synthetic time within the leaf's validity
// window. This lets us extract structured info from expired/future
// certificates instead of returning an opaque x509 error.
// We use NotAfter-1s (just before expiry) to maximize overlap with CA validity.
opts := x509.VerifyOptions{
Roots: pool,
CurrentTime: leaf.NotAfter.Add(-time.Second),
KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsageAny},
}
if len(certs) > 1 {
intermediates := x509.NewCertPool()
for _, c := range certs[1:] {
intermediates.AddCert(c)
}
opts.Intermediates = intermediates
}
if _, err := leaf.Verify(opts); err != nil {
return nil, fmt.Errorf("certificate verification failed: %w", err)
}
if IsRevoked(leaf.SerialNumber) {
return nil, fmt.Errorf("certificate serial %s has been revoked", leaf.SerialNumber.Text(16))
}
info := extractIdentity(leaf)
parseExtensions(leaf, info)
now := time.Now()
if now.Before(leaf.NotBefore) {
info.Valid = false
info.Error = fmt.Sprintf("certificate not yet valid (starts %s)",
leaf.NotBefore.UTC().Format("2006-01-02"))
return info, nil
}
if now.After(leaf.NotAfter) {
info.Valid = false
info.IsExpired = true
info.Error = fmt.Sprintf("certificate expired on %s",
leaf.NotAfter.UTC().Format("2006-01-02"))
return info, nil
}
if !info.HasProduct(product) {
info.Valid = false
info.Error = fmt.Sprintf("certificate not licensed for product %q (licensed: %v)",
product, info.Product)
return info, nil
}
// Machine ID binding: if the certificate specifies a machine ID,
// it must match the current runtime machine ID.
// Empty machine_id means no binding — runs on any machine.
if info.MachineID != "" {
if got := currentMachineID(); got != info.MachineID {
info.Valid = false
info.Error = "certificate machine_id does not match this host"
return info, nil
}
}
info.Valid = true
return info, nil
}
func extractIdentity(cert *x509.Certificate) *LicenseInfo {
info := &LicenseInfo{
LicenseeName: cert.Subject.CommonName,
SerialNumber: cert.SerialNumber.Text(16),
NotBefore: cert.NotBefore.Unix(),
NotAfter: cert.NotAfter.Unix(),
Issuer: cert.Issuer.CommonName,
Edition: "community",
Product: []string{},
Permissions: Permissions{SupportLevel: "none"},
}
if len(cert.Subject.Organization) > 0 {
info.LicenseeOrg = cert.Subject.Organization[0]
}
if len(cert.Subject.Country) > 0 {
info.LicenseeCountry = cert.Subject.Country[0]
}
if len(cert.EmailAddresses) > 0 {
info.LicenseeEmail = cert.EmailAddresses[0]
}
return info
}
func parseExtensions(cert *x509.Certificate, info *LicenseInfo) {
for _, ext := range cert.Extensions {
val := string(ext.Value)
switch {
// Scope
case ext.Id.Equal(OIDProduct):
info.Product = splitCSV(val)
case ext.Id.Equal(OIDEdition):
info.Edition = val
case ext.Id.Equal(OIDEnv):
if val != "" {
info.Env = splitCSV(val)
}
case ext.Id.Equal(OIDDomain):
info.Domain = val
case ext.Id.Equal(OIDAppID):
info.AppID = val
case ext.Id.Equal(OIDMachineID):
info.MachineID = val
// Quota
case ext.Id.Equal(OIDMaxUsers):
info.MaxUsers = atoi(val)
case ext.Id.Equal(OIDMaxTaiNodes):
info.MaxTaiNodes = atoi(val)
case ext.Id.Equal(OIDMaxAgents):
info.MaxAgents = atoi(val)
case ext.Id.Equal(OIDMaxSandboxes):
info.MaxSandboxes = atoi(val)
case ext.Id.Equal(OIDMaxAPIRPM):
info.MaxAPIRPM = atoi(val)
case ext.Id.Equal(OIDMaxStorageGB):
info.MaxStorageGB = atoi(val)
// Permissions
case ext.Id.Equal(OIDAllowBrandingRemoval):
info.Permissions.AllowBrandingRemoval = toBool(val)
case ext.Id.Equal(OIDAllowWhiteLabel):
info.Permissions.AllowWhiteLabel = toBool(val)
case ext.Id.Equal(OIDAllowMultiTenant):
info.Permissions.AllowMultiTenant = toBool(val)
case ext.Id.Equal(OIDAllowCustomDomain):
info.Permissions.AllowCustomDomain = toBool(val)
case ext.Id.Equal(OIDAllowHostExec):
info.Permissions.AllowHostExec = toBool(val)
case ext.Id.Equal(OIDAllowSSO):
info.Permissions.AllowSSO = toBool(val)
case ext.Id.Equal(OIDSupportLevel):
info.Permissions.SupportLevel = val
}
}
}
// MakeExtension creates a pkix.Extension for embedding in a certificate.
func MakeExtension(oid asn1.ObjectIdentifier, value string) ExtensionValue {
return ExtensionValue{OID: oid, Value: value}
}
// ExtensionValue pairs an OID with its string value for certificate generation.
type ExtensionValue struct {
OID asn1.ObjectIdentifier
Value string
}
func splitCSV(s string) []string {
parts := strings.Split(s, ",")
var result []string
for _, p := range parts {
if t := strings.TrimSpace(p); t != "" {
result = append(result, t)
}
}
return result
}
func atoi(s string) int {
n, _ := strconv.Atoi(strings.TrimSpace(s))
return n
}
func toBool(s string) bool {
s = strings.TrimSpace(strings.ToLower(s))
return s == "true" || s == "1" || s == "yes"
}
// currentMachineID returns a deterministic identifier for the current host,
// using the same algorithm as tai/machine.ID():
// - macOS: IOPlatformUUID via ioreg
// - Linux: /etc/machine-id
// - Windows: HKLM MachineGuid registry value
// - fallback: sha256("tai-fallback:" + hostname)[:16]
//
// Implemented via platform-specific machine_{os}.go files in this package.
func currentMachineID() string {
if id := platformMachineID(); id != "" {
return id
}
hostname, err := os.Hostname()
if err != nil || hostname == "" {
hostname = "unknown"
}
h := sha256.Sum256([]byte("tai-fallback:" + hostname))
return fmt.Sprintf("%x", h[:16])
}

View file

@ -0,0 +1,587 @@
package commercial
import (
"crypto/x509"
_ "embed"
"math/big"
"os"
"path/filepath"
"sync"
"testing"
"time"
)
//go:embed testdata/test-intermediate-ca.pem
var testIntermediateCAPEM []byte
//go:embed testdata/test-license.pem
var testLicensePEM []byte
// withTestRootPool temporarily replaces the root pool and revocation list
// for testing, restoring originals on cleanup.
func withTestRootPool(t *testing.T, ca *testCA, revoked []*big.Int) {
t.Helper()
origPool := rootPool
origOnce := rootPoolOnce
origSerials := revokedSerials
origRevokedOnce := revokedOnce
pool := x509.NewCertPool()
pool.AddCert(ca.Cert)
rootPool = pool
doneOnce := &sync.Once{}
doneOnce.Do(func() {}) // pre-mark as done so RootPool() returns our pool
rootPoolOnce = doneOnce
revokedSerials = revoked
doneOnce2 := &sync.Once{}
doneOnce2.Do(func() {})
revokedOnce = doneOnce2
t.Cleanup(func() {
rootPool = origPool
rootPoolOnce = origOnce
revokedSerials = origSerials
revokedOnce = origRevokedOnce
})
}
func TestNoCertificate(t *testing.T) {
dir := t.TempDir()
License = DefaultLicense()
Load(dir, "yao")
if License.Valid {
t.Fatal("expected Valid=false with no certificate")
}
if License.Source != "none" {
t.Fatalf("expected Source=none, got %s", License.Source)
}
if License.Edition != "community" {
t.Fatalf("expected Edition=community, got %s", License.Edition)
}
if License.MaxUsers != 100 {
t.Fatalf("expected MaxUsers=100, got %d", License.MaxUsers)
}
}
func TestValidCertificate(t *testing.T) {
root, err := generateRootCA("Test Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
withTestRootPool(t, root, nil)
opts := defaultLicenseOpts()
leaf, err := generateLicenseCert(root, opts)
if err != nil {
t.Fatal(err)
}
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "license.pem"), leaf.CertPEM, 0644); err != nil {
t.Fatal(err)
}
Load(dir, "yao")
if !License.Valid {
t.Fatalf("expected Valid=true, got error: %s", License.Error)
}
if License.Source != "file" {
t.Fatalf("expected Source=file, got %s", License.Source)
}
if License.Edition != "pro" {
t.Fatalf("expected Edition=pro, got %s", License.Edition)
}
if License.LicenseeName != "Test Corp" {
t.Fatalf("expected LicenseeName=Test Corp, got %s", License.LicenseeName)
}
if License.MaxUsers != 500 {
t.Fatalf("expected MaxUsers=500, got %d", License.MaxUsers)
}
}
func TestExpiredCertificate(t *testing.T) {
root, err := generateRootCA("Test Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
withTestRootPool(t, root, nil)
opts := defaultLicenseOpts()
opts.NotBefore = time.Now().Add(-30 * time.Minute)
opts.NotAfter = time.Now().Add(-1 * time.Minute)
leaf, err := generateLicenseCert(root, opts)
if err != nil {
t.Fatal(err)
}
dir := t.TempDir()
os.WriteFile(filepath.Join(dir, "license.pem"), leaf.CertPEM, 0644)
Load(dir, "yao")
if License.Valid {
t.Fatal("expected Valid=false for expired certificate")
}
if !License.IsExpired {
t.Fatal("expected IsExpired=true")
}
}
func TestNotYetValidCertificate(t *testing.T) {
root, err := generateRootCA("Test Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
withTestRootPool(t, root, nil)
opts := defaultLicenseOpts()
opts.NotBefore = time.Now().Add(24 * time.Hour)
opts.NotAfter = time.Now().Add(365 * 24 * time.Hour)
leaf, err := generateLicenseCert(root, opts)
if err != nil {
t.Fatal(err)
}
dir := t.TempDir()
os.WriteFile(filepath.Join(dir, "license.pem"), leaf.CertPEM, 0644)
Load(dir, "yao")
if License.Valid {
t.Fatal("expected Valid=false for not-yet-valid certificate")
}
if License.IsExpired {
t.Fatal("expected IsExpired=false for future certificate")
}
}
func TestTamperedCertificate(t *testing.T) {
root, err := generateRootCA("Test Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
withTestRootPool(t, root, nil)
// Generate with a different root (not in our pool) to simulate tampering
fakeRoot, err := generateRootCA("Fake Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
opts := defaultLicenseOpts()
leaf, err := generateLicenseCert(fakeRoot, opts)
if err != nil {
t.Fatal(err)
}
dir := t.TempDir()
os.WriteFile(filepath.Join(dir, "license.pem"), leaf.CertPEM, 0644)
Load(dir, "yao")
if License.Valid {
t.Fatal("expected Valid=false for tampered certificate")
}
if License.Error == "" {
t.Fatal("expected Error to be set")
}
}
func TestWrongProduct(t *testing.T) {
root, err := generateRootCA("Test Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
withTestRootPool(t, root, nil)
opts := defaultLicenseOpts()
// Only licensed for "tai", not "yao"
for i, ext := range opts.Extensions {
if ext.OID.Equal(OIDProduct) {
opts.Extensions[i].Value = "tai"
}
}
leaf, err := generateLicenseCert(root, opts)
if err != nil {
t.Fatal(err)
}
dir := t.TempDir()
os.WriteFile(filepath.Join(dir, "license.pem"), leaf.CertPEM, 0644)
Load(dir, "yao")
if License.Valid {
t.Fatal("expected Valid=false for wrong product")
}
}
func TestCertificateChainWithIntermediate(t *testing.T) {
root, err := generateRootCA("Test Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
withTestRootPool(t, root, nil)
intermediate, err := generateIntermediateCA("Test Intermediate CA", 3*365*24*time.Hour, root)
if err != nil {
t.Fatal(err)
}
opts := defaultLicenseOpts()
leaf, err := generateLicenseCert(intermediate, opts)
if err != nil {
t.Fatal(err)
}
// PEM chain: leaf + intermediate
chainPEM := append(leaf.CertPEM, intermediate.CertPEM...)
dir := t.TempDir()
os.WriteFile(filepath.Join(dir, "license.pem"), chainPEM, 0644)
Load(dir, "yao")
if !License.Valid {
t.Fatalf("expected Valid=true with intermediate chain, got error: %s", License.Error)
}
}
func TestRevokedCertificate(t *testing.T) {
root, err := generateRootCA("Test Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
opts := defaultLicenseOpts()
opts.Serial = big.NewInt(99999)
leaf, err := generateLicenseCert(root, opts)
if err != nil {
t.Fatal(err)
}
withTestRootPool(t, root, []*big.Int{big.NewInt(99999)})
dir := t.TempDir()
os.WriteFile(filepath.Join(dir, "license.pem"), leaf.CertPEM, 0644)
Load(dir, "yao")
if License.Valid {
t.Fatal("expected Valid=false for revoked certificate")
}
}
func TestEnvVarLoading(t *testing.T) {
root, err := generateRootCA("Test Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
withTestRootPool(t, root, nil)
opts := defaultLicenseOpts()
leaf, err := generateLicenseCert(root, opts)
if err != nil {
t.Fatal(err)
}
// Write cert to a temp file and point env var to it
certFile := filepath.Join(t.TempDir(), "test-license.pem")
os.WriteFile(certFile, leaf.CertPEM, 0644)
t.Setenv(envLicenseCert, certFile)
Load(t.TempDir(), "yao")
if !License.Valid {
t.Fatalf("expected Valid=true via env, got error: %s", License.Error)
}
if License.Source != "env" {
t.Fatalf("expected Source=env, got %s", License.Source)
}
}
func TestAllExtensions(t *testing.T) {
root, err := generateRootCA("Test Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
withTestRootPool(t, root, nil)
opts := defaultLicenseOpts()
opts.Extensions = []ExtensionValue{
{OID: OIDProduct, Value: "yao,tai"},
{OID: OIDEdition, Value: "enterprise"},
{OID: OIDEnv, Value: "production,staging"},
{OID: OIDDomain, Value: "*.acme.com"},
{OID: OIDAppID, Value: "acme-crm"},
{OID: OIDMaxUsers, Value: "0"},
{OID: OIDMaxTaiNodes, Value: "0"},
{OID: OIDMaxAgents, Value: "0"},
{OID: OIDMaxSandboxes, Value: "0"},
{OID: OIDMaxAPIRPM, Value: "0"},
{OID: OIDMaxStorageGB, Value: "0"},
{OID: OIDAllowBrandingRemoval, Value: "true"},
{OID: OIDAllowWhiteLabel, Value: "true"},
{OID: OIDAllowMultiTenant, Value: "true"},
{OID: OIDAllowCustomDomain, Value: "true"},
{OID: OIDAllowHostExec, Value: "true"},
{OID: OIDAllowSSO, Value: "true"},
{OID: OIDSupportLevel, Value: "dedicated"},
}
leaf, err := generateLicenseCert(root, opts)
if err != nil {
t.Fatal(err)
}
dir := t.TempDir()
os.WriteFile(filepath.Join(dir, "license.pem"), leaf.CertPEM, 0644)
Load(dir, "yao")
if !License.Valid {
t.Fatalf("expected Valid=true, got error: %s", License.Error)
}
if License.Edition != "enterprise" {
t.Fatalf("expected enterprise, got %s", License.Edition)
}
if !License.HasProduct("yao") || !License.HasProduct("tai") {
t.Fatalf("expected product yao,tai, got %v", License.Product)
}
if len(License.Env) != 2 {
t.Fatalf("expected 2 envs, got %v", License.Env)
}
if License.Domain != "*.acme.com" {
t.Fatalf("expected domain *.acme.com, got %s", License.Domain)
}
if License.AppID != "acme-crm" {
t.Fatalf("expected app_id acme-crm, got %s", License.AppID)
}
if License.MaxUsers != 0 {
t.Fatalf("expected MaxUsers=0 (unlimited), got %d", License.MaxUsers)
}
if !License.Permissions.AllowBrandingRemoval {
t.Fatal("expected AllowBrandingRemoval=true")
}
if !License.Permissions.AllowWhiteLabel {
t.Fatal("expected AllowWhiteLabel=true")
}
if !License.Permissions.AllowMultiTenant {
t.Fatal("expected AllowMultiTenant=true")
}
if !License.Permissions.AllowSSO {
t.Fatal("expected AllowSSO=true")
}
if License.Permissions.SupportLevel != "dedicated" {
t.Fatalf("expected SupportLevel=dedicated, got %s", License.Permissions.SupportLevel)
}
}
func TestDefaultLicenseAndHelpers(t *testing.T) {
def := DefaultLicense()
if def.Valid {
t.Fatal("default should not be Valid")
}
if def.Edition != "community" {
t.Fatalf("expected community, got %s", def.Edition)
}
if !def.IsLevel("community") {
t.Fatal("community should satisfy IsLevel(community)")
}
if def.IsLevel("starter") {
t.Fatal("community should not satisfy IsLevel(starter)")
}
if def.IsLevel("pro") {
t.Fatal("community should not satisfy IsLevel(pro)")
}
pro := LicenseInfo{Edition: "pro"}
if !pro.IsLevel("community") {
t.Fatal("pro should satisfy IsLevel(community)")
}
if !pro.IsLevel("starter") {
t.Fatal("pro should satisfy IsLevel(starter)")
}
if !pro.IsLevel("pro") {
t.Fatal("pro should satisfy IsLevel(pro)")
}
if pro.IsLevel("enterprise") {
t.Fatal("pro should not satisfy IsLevel(enterprise)")
}
multi := LicenseInfo{Product: []string{"yao", "tai"}}
if !multi.HasProduct("yao") {
t.Fatal("expected HasProduct(yao)=true")
}
if !multi.HasProduct("tai") {
t.Fatal("expected HasProduct(tai)=true")
}
if multi.HasProduct("other") {
t.Fatal("expected HasProduct(other)=false")
}
}
// TestRealChainIntermediate verifies a certificate signed by a real test intermediate CA,
// which in turn is signed by the real Root CA 1 embedded in the binary.
// Uses testdata/test-intermediate-ca.pem and testdata/test-license.pem generated by
// /Volumes/DATA/Work/Yaobots/keys/gen-test-certs.go.
func TestRealChainIntermediate(t *testing.T) {
// Use the real embedded root pool (not overridden)
dir := t.TempDir()
if err := os.WriteFile(filepath.Join(dir, "license.pem"), testLicensePEM, 0644); err != nil {
t.Fatal(err)
}
License = DefaultLicense()
Load(dir, "yao")
if !License.Valid {
t.Fatalf("expected Valid=true with real chain, got error: %s", License.Error)
}
if License.Edition != "enterprise" {
t.Fatalf("expected Edition=enterprise, got %s", License.Edition)
}
if !License.HasProduct("yao") {
t.Fatalf("expected HasProduct(yao)=true, got products: %v", License.Product)
}
if !License.Permissions.AllowBrandingRemoval {
t.Fatal("expected AllowBrandingRemoval=true")
}
if !License.Permissions.AllowWhiteLabel {
t.Fatal("expected AllowWhiteLabel=true")
}
if !License.IsLevel("enterprise") {
t.Fatalf("expected IsLevel(enterprise)=true, edition=%s", License.Edition)
}
}
// TestRealIntermediateCACert parses the embedded test intermediate CA cert
// and verifies it is signed by the real Root CA 1.
func TestRealIntermediateCACert(t *testing.T) {
intCert, err := ParsePEMChain(testIntermediateCAPEM)
if err != nil || len(intCert) == 0 {
t.Fatalf("failed to parse test intermediate CA: %v", err)
}
opts := x509.VerifyOptions{
Roots: RootPool(),
CurrentTime: time.Now(),
}
// Intermediate CA certs are not end-entity; relax KeyUsages check
opts.KeyUsages = []x509.ExtKeyUsage{x509.ExtKeyUsageAny}
if _, err := intCert[0].Verify(opts); err != nil {
t.Fatalf("test intermediate CA not verified by real root pool: %v", err)
}
}
func TestCertsSubdirectoryFallback(t *testing.T) {
root, err := generateRootCA("Test Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
withTestRootPool(t, root, nil)
opts := defaultLicenseOpts()
leaf, err := generateLicenseCert(root, opts)
if err != nil {
t.Fatal(err)
}
dir := t.TempDir()
certsDir := filepath.Join(dir, "certs")
os.MkdirAll(certsDir, 0755)
os.WriteFile(filepath.Join(certsDir, "license.pem"), leaf.CertPEM, 0644)
Load(dir, "yao")
if !License.Valid {
t.Fatalf("expected Valid=true from certs/ fallback, got error: %s", License.Error)
}
if License.Source != "file" {
t.Fatalf("expected Source=file, got %s", License.Source)
}
}
func TestMachineIDEmpty(t *testing.T) {
// No machine_id in cert → valid on any machine
root, err := generateRootCA("Test Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
withTestRootPool(t, root, nil)
opts := defaultLicenseOpts()
// defaultLicenseOpts has no OIDMachineID → empty
leaf, err := generateLicenseCert(root, opts)
if err != nil {
t.Fatal(err)
}
dir := t.TempDir()
os.WriteFile(filepath.Join(dir, "license.pem"), leaf.CertPEM, 0644)
Load(dir, "yao")
if !License.Valid {
t.Fatalf("expected Valid=true when machine_id is empty, got: %s", License.Error)
}
}
func TestMachineIDMatch(t *testing.T) {
// machine_id in cert matches current machine → valid
root, err := generateRootCA("Test Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
withTestRootPool(t, root, nil)
thisID := currentMachineID()
opts := defaultLicenseOpts()
opts.Extensions = append(opts.Extensions, ExtensionValue{OID: OIDMachineID, Value: thisID})
leaf, err := generateLicenseCert(root, opts)
if err != nil {
t.Fatal(err)
}
dir := t.TempDir()
os.WriteFile(filepath.Join(dir, "license.pem"), leaf.CertPEM, 0644)
Load(dir, "yao")
if !License.Valid {
t.Fatalf("expected Valid=true when machine_id matches, got: %s", License.Error)
}
if License.MachineID != thisID {
t.Fatalf("expected MachineID=%s, got %s", thisID, License.MachineID)
}
}
func TestMachineIDMismatch(t *testing.T) {
// machine_id in cert does not match current machine → invalid
root, err := generateRootCA("Test Root CA", 10*365*24*time.Hour)
if err != nil {
t.Fatal(err)
}
withTestRootPool(t, root, nil)
opts := defaultLicenseOpts()
opts.Extensions = append(opts.Extensions, ExtensionValue{OID: OIDMachineID, Value: "000000000000000000000000deadbeef"})
leaf, err := generateLicenseCert(root, opts)
if err != nil {
t.Fatal(err)
}
dir := t.TempDir()
os.WriteFile(filepath.Join(dir, "license.pem"), leaf.CertPEM, 0644)
Load(dir, "yao")
if License.Valid {
t.Fatal("expected Valid=false when machine_id does not match")
}
if License.Error == "" {
t.Fatal("expected Error to be set")
}
}

View file

@ -0,0 +1,27 @@
//go:build darwin
package commercial
import (
"os/exec"
"strings"
)
func platformMachineID() string {
out, err := exec.Command("ioreg", "-rd1", "-c", "IOPlatformExpertDevice").Output()
if err != nil {
return ""
}
for _, line := range strings.Split(string(out), "\n") {
if strings.Contains(line, "IOPlatformUUID") {
parts := strings.SplitN(line, "=", 2)
if len(parts) == 2 {
uuid := strings.Trim(strings.TrimSpace(parts[1]), "\"")
if uuid != "" {
return strings.TrimSpace(strings.ToLower(uuid))
}
}
}
}
return ""
}

View file

@ -0,0 +1,23 @@
//go:build linux
package commercial
import (
"os"
"strings"
)
func platformMachineID() string {
data, err := os.ReadFile("/etc/machine-id")
if err != nil {
data, err = os.ReadFile("/var/lib/dbus/machine-id")
if err != nil {
return ""
}
}
id := strings.TrimSpace(strings.ToLower(string(data)))
if id == "" {
return ""
}
return id
}

View file

@ -0,0 +1,28 @@
//go:build windows
package commercial
import (
"os/exec"
"strings"
)
func platformMachineID() string {
out, err := exec.Command("reg", "query",
`HKLM\SOFTWARE\Microsoft\Cryptography`,
"/v", "MachineGuid",
).Output()
if err != nil {
return ""
}
for _, line := range strings.Split(string(out), "\n") {
line = strings.TrimSpace(line)
if strings.Contains(line, "MachineGuid") {
fields := strings.Fields(line)
if len(fields) >= 3 {
return strings.TrimSpace(strings.ToLower(fields[len(fields)-1]))
}
}
}
return ""
}

41
commercial/oid.go Normal file
View file

@ -0,0 +1,41 @@
package commercial
import "encoding/asn1"
// OID prefix: 1.3.6.1.4.1.15099.1
// 15099 is Yao's internal port number, used as a recognizable enterprise number
// for this closed-loop certificate system. Not registered with IANA.
var (
// Scope
OIDProduct = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 1, 1}
OIDEdition = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 1, 2}
OIDEnv = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 1, 3}
OIDDomain = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 1, 4}
OIDAppID = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 1, 5}
// Quota
OIDMaxUsers = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 2, 1}
OIDMaxTaiNodes = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 2, 2}
OIDMaxAgents = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 2, 3}
OIDMaxSandboxes = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 2, 4}
OIDMaxAPIRPM = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 2, 5}
OIDMaxStorageGB = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 2, 6}
// Permissions
OIDAllowBrandingRemoval = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 3, 1}
OIDAllowWhiteLabel = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 3, 2}
OIDAllowMultiTenant = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 3, 3}
OIDAllowCustomDomain = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 3, 4}
OIDAllowHostExec = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 3, 5}
OIDAllowSSO = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 3, 6}
OIDSupportLevel = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 3, 7}
// Binding (optional — if present, must match at runtime)
OIDMachineID = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 5, 1}
// Issuance (internal tracking)
OIDIssuerID = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 4, 1}
OIDOrderID = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 4, 2}
OIDNote = asn1.ObjectIdentifier{1, 3, 6, 1, 4, 1, 15099, 1, 4, 3}
)

98
commercial/roots.go Normal file
View file

@ -0,0 +1,98 @@
package commercial
import (
"crypto/x509"
_ "embed"
"encoding/json"
"encoding/pem"
"math/big"
"sync"
)
//go:embed roots/root-ca-1.pem
var rootCA1PEM []byte
//go:embed roots/root-ca-2.pem
var rootCA2PEM []byte
//go:embed roots/revoked.json
var revokedJSON []byte
var (
rootPoolOnce = &sync.Once{}
rootPool *x509.CertPool
revokedOnce = &sync.Once{}
revokedSerials []*big.Int
)
// RootPool returns the built-in root certificate pool (primary + backup).
// Returns nil if neither root certificate could be parsed (e.g. dev placeholder).
func RootPool() *x509.CertPool {
rootPoolOnce.Do(func() {
pool := x509.NewCertPool()
added := false
for _, pemData := range [][]byte{rootCA1PEM, rootCA2PEM} {
if pool.AppendCertsFromPEM(pemData) {
added = true
}
}
if added {
rootPool = pool
}
})
return rootPool
}
// RevokedSerials returns the list of revoked certificate serial numbers
// embedded in the binary.
func RevokedSerials() []*big.Int {
revokedOnce.Do(func() {
var data struct {
Serials []string `json:"serials"`
}
if err := json.Unmarshal(revokedJSON, &data); err != nil {
return
}
for _, s := range data.Serials {
n := new(big.Int)
if _, ok := n.SetString(s, 0); ok {
revokedSerials = append(revokedSerials, n)
}
}
})
return revokedSerials
}
// IsRevoked checks whether the given serial number is in the revocation list.
func IsRevoked(serial *big.Int) bool {
for _, s := range RevokedSerials() {
if s.Cmp(serial) == 0 {
return true
}
}
return false
}
// ParsePEMChain parses a PEM bundle into a list of x509 certificates.
// The first certificate is treated as the leaf; remaining are intermediates.
func ParsePEMChain(pemData []byte) ([]*x509.Certificate, error) {
var certs []*x509.Certificate
rest := pemData
for {
var block *pem.Block
block, rest = pem.Decode(rest)
if block == nil {
break
}
if block.Type != "CERTIFICATE" {
continue
}
cert, err := x509.ParseCertificate(block.Bytes)
if err != nil {
return nil, err
}
certs = append(certs, cert)
}
return certs, nil
}

View file

@ -0,0 +1 @@
{"serials": []}

View file

@ -0,0 +1,13 @@
-----BEGIN CERTIFICATE-----
MIICATCCAYegAwIBAgIBATAKBggqhkjOPQQDAzBIMQswCQYDVQQGEwJDTjEhMB8G
A1UEChMYSW5maW5pdGUgV2lzZG9tIFNvZnR3YXJlMRYwFAYDVQQDEw1ZYW8gUm9v
dCBDQSAxMB4XDTI2MDMyNzA2MzYxMloXDTM2MDMyNDA3MzYxMlowSDELMAkGA1UE
BhMCQ04xITAfBgNVBAoTGEluZmluaXRlIFdpc2RvbSBTb2Z0d2FyZTEWMBQGA1UE
AxMNWWFvIFJvb3QgQ0EgMTB2MBAGByqGSM49AgEGBSuBBAAiA2IABDg0J9BxEjOn
Os5D7i3fzFxFOlLR13nanGIjde3bqaohZoD7fCGtfVd5+2vq7y/UOQRAvbu5/RzK
e1S9LFlWjFppdpgbeegd0dMCDEe/9tJknAO35MpWCJxmU7f3Jo4QE6NFMEMwDgYD
VR0PAQH/BAQDAgEGMBIGA1UdEwEB/wQIMAYBAf8CAQIwHQYDVR0OBBYEFB6XvPEQ
D+g1yTpwyPj3Quwu9wuXMAoGCCqGSM49BAMDA2gAMGUCMFc2K+Xggn5GRo181pAU
aU43Sud8kyXlOuMMpQ/+mfqVNKJhPN75VditQoGiau8VjQIxAO4lJGpuG8f/j4mE
6WaF1x4x45GKc6zaHtgpVUTaX2FazAqkZgSYy57Ipnx0UWrD+A==
-----END CERTIFICATE-----

View file

@ -0,0 +1,13 @@
-----BEGIN CERTIFICATE-----
MIICATCCAYegAwIBAgIBAjAKBggqhkjOPQQDAzBIMQswCQYDVQQGEwJDTjEhMB8G
A1UEChMYSW5maW5pdGUgV2lzZG9tIFNvZnR3YXJlMRYwFAYDVQQDEw1ZYW8gUm9v
dCBDQSAyMB4XDTI2MDMyNzA2MzYxMloXDTM2MDMyNDA3MzYxMlowSDELMAkGA1UE
BhMCQ04xITAfBgNVBAoTGEluZmluaXRlIFdpc2RvbSBTb2Z0d2FyZTEWMBQGA1UE
AxMNWWFvIFJvb3QgQ0EgMjB2MBAGByqGSM49AgEGBSuBBAAiA2IABB01XlJON6el
vpif5oCD4Mmdvj5muXeq02ZGHSWSrqrxkzUbCbI5LpwURiGObOVnGMw2MT76hW0P
Eq28XDTiXxnVqrjkE95ppl5npgnqMv3PTODizkak2HSXqwuIESLip6NFMEMwDgYD
VR0PAQH/BAQDAgEGMBIGA1UdEwEB/wQIMAYBAf8CAQIwHQYDVR0OBBYEFNHmPR9D
aXyJZQng5BcwhaB9H9CsMAoGCCqGSM49BAMDA2gAMGUCMQDkQlbGE1yZq7cLele8
4S+RiYhj3sTsZejU5215f0doOfDka18iv5BeclELP2aQOeQCMFqfeeSKS9SFrHly
SRfZV9o/qxQM3i9EgV1l/D7T9qzPZWR2YcTUSS8TTcwxyIiLAA==
-----END CERTIFICATE-----

View file

@ -0,0 +1,14 @@
-----BEGIN CERTIFICATE-----
MIICLzCCAbSgAwIBAgICIykwCgYIKoZIzj0EAwMwSDELMAkGA1UEBhMCQ04xITAf
BgNVBAoTGEluZmluaXRlIFdpc2RvbSBTb2Z0d2FyZTEWMBQGA1UEAxMNWWFvIFJv
b3QgQ0EgMTAeFw0yNjAzMjcwNjQyNTlaFw0yODAzMjYwNzQyNTlaMFMxCzAJBgNV
BAYTAkNOMSEwHwYDVQQKExhJbmZpbml0ZSBXaXNkb20gU29mdHdhcmUxITAfBgNV
BAMTGFlhbyBUZXN0IEludGVybWVkaWF0ZSBDQTB2MBAGByqGSM49AgEGBSuBBAAi
A2IABHG2U+9ScPK/NBfGW3jgg8fNVH/NfGcPOkypKHonsIad1fG4H+kUtEyyQTvP
4of6FO8FmnvyezMdtaEPlUxRLmYPW9dqsnCZswZasgQLOd5b+qxe4bFY86yla2Iq
HxWIyKNmMGQwDgYDVR0PAQH/BAQDAgEGMBIGA1UdEwEB/wQIMAYBAf8CAQAwHQYD
VR0OBBYEFHoGH8BliDziOpSBraPs5+ehyWBMMB8GA1UdIwQYMBaAFB6XvPEQD+g1
yTpwyPj3Quwu9wuXMAoGCCqGSM49BAMDA2kAMGYCMQCkN/W+unUuHyP/2phruVv4
itASDUZfm3ZDGS450ixZO9OMQZ9t/dNv1HyfN0M9wCMCMQDoYKwveqllF7XvSPQp
fBfRbxsR+t8LpJqzNXMGgzUpKuZ/KLBAjnbZgwCCGCR1zN0=
-----END CERTIFICATE-----

34
commercial/testdata/test-license.pem vendored Normal file
View file

@ -0,0 +1,34 @@
-----BEGIN CERTIFICATE-----
MIIDSzCCAtGgAwIBAgIDAV+RMAoGCCqGSM49BAMDMFMxCzAJBgNVBAYTAkNOMSEw
HwYDVQQKExhJbmZpbml0ZSBXaXNkb20gU29mdHdhcmUxITAfBgNVBAMTGFlhbyBU
ZXN0IEludGVybWVkaWF0ZSBDQTAeFw0yNjAzMjcwNjQyNTlaFw0yNzAzMjcwNzQy
NTlaMD4xCzAJBgNVBAYTAkNOMRkwFwYDVQQKExBZYW8gQXBwIFRlc3QgT3JnMRQw
EgYDVQQDEwtZYW8gRGV2IEFwcDBZMBMGByqGSM49AgEGCCqGSM49AwEHA0IABLBN
Ikb7MXULLpHymC2qV2YjB2pbUjuB7hP6tQa/5aN82o5W7lxMUF2jP4kuptMeaiXj
QyKrtlhDWZzAJ/HshhmjggGnMIIBozAOBgNVHQ8BAf8EBAMCB4AwDwYDVR0lBAgw
BgYEVR0lADAfBgNVHSMEGDAWgBR6Bh/AZYg84jqUga2j7OfnoclgTDAaBgNVHREE
EzARgQ9kZXZAeWFvYXBwcy5jb20wFQYKKwYBBAH1ewEBAQQHeWFvLHRhaTAYBgor
BgEEAfV7AQECBAplbnRlcnByaXNlMA8GCisGAQQB9XsBAgEEATAwDwYKKwYBBAH1
ewECAgQBMDAPBgorBgEEAfV7AQIDBAEwMA8GCisGAQQB9XsBAgQEATAwDwYKKwYB
BAH1ewECBQQBMDAPBgorBgEEAfV7AQIGBAEwMBIGCisGAQQB9XsBAwEEBHRydWUw
EgYKKwYBBAH1ewEDAgQEdHJ1ZTASBgorBgEEAfV7AQMDBAR0cnVlMBIGCisGAQQB
9XsBAwQEBHRydWUwEgYKKwYBBAH1ewEDBQQEdHJ1ZTASBgorBgEEAfV7AQMGBAR0
cnVlMBcGCisGAQQB9XsBAwcECWRlZGljYXRlZDAbBgorBgEEAfV7AQQCBA1URVNU
LTIwMjYtMDAxMAoGCCqGSM49BAMDA2gAMGUCMDKJ9oL04HJa0Djh5PbB/bEDEpZA
nS/UWyfaXCbgqQeI+rIhjIXFsEC9CXeFXfa5iwIxAKmTybVyE0QkIQqEifB50LLy
8FXcL1owzKWB5JHn2rTVGr2cSf2SrBQJMefFWiLEUQ==
-----END CERTIFICATE-----
-----BEGIN CERTIFICATE-----
MIICLzCCAbSgAwIBAgICIykwCgYIKoZIzj0EAwMwSDELMAkGA1UEBhMCQ04xITAf
BgNVBAoTGEluZmluaXRlIFdpc2RvbSBTb2Z0d2FyZTEWMBQGA1UEAxMNWWFvIFJv
b3QgQ0EgMTAeFw0yNjAzMjcwNjQyNTlaFw0yODAzMjYwNzQyNTlaMFMxCzAJBgNV
BAYTAkNOMSEwHwYDVQQKExhJbmZpbml0ZSBXaXNkb20gU29mdHdhcmUxITAfBgNV
BAMTGFlhbyBUZXN0IEludGVybWVkaWF0ZSBDQTB2MBAGByqGSM49AgEGBSuBBAAi
A2IABHG2U+9ScPK/NBfGW3jgg8fNVH/NfGcPOkypKHonsIad1fG4H+kUtEyyQTvP
4of6FO8FmnvyezMdtaEPlUxRLmYPW9dqsnCZswZasgQLOd5b+qxe4bFY86yla2Iq
HxWIyKNmMGQwDgYDVR0PAQH/BAQDAgEGMBIGA1UdEwEB/wQIMAYBAf8CAQAwHQYD
VR0OBBYEFHoGH8BliDziOpSBraPs5+ehyWBMMB8GA1UdIwQYMBaAFB6XvPEQD+g1
yTpwyPj3Quwu9wuXMAoGCCqGSM49BAMDA2kAMGYCMQCkN/W+unUuHyP/2phruVv4
itASDUZfm3ZDGS450ixZO9OMQZ9t/dNv1HyfN0M9wCMCMQDoYKwveqllF7XvSPQp
fBfRbxsR+t8LpJqzNXMGgzUpKuZ/KLBAjnbZgwCCGCR1zN0=
-----END CERTIFICATE-----

View file

@ -0,0 +1,177 @@
package commercial
import (
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/x509"
"crypto/x509/pkix"
"encoding/pem"
"math/big"
"time"
)
// testCA holds a generated CA certificate and its private key.
type testCA struct {
Cert *x509.Certificate
Key *ecdsa.PrivateKey
CertPEM []byte
}
// testLicenseCert holds a generated leaf (license) certificate.
type testLicenseCert struct {
Cert *x509.Certificate
CertPEM []byte
}
// generateRootCA creates a self-signed ECDSA P-384 root CA for testing.
func generateRootCA(cn string, validity time.Duration) (*testCA, error) {
key, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
if err != nil {
return nil, err
}
serial, _ := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
tmpl := &x509.Certificate{
SerialNumber: serial,
Subject: pkix.Name{
CommonName: cn,
Organization: []string{"Infinite Wisdom Software"},
Country: []string{"CN"},
},
NotBefore: time.Now().Add(-1 * time.Hour),
NotAfter: time.Now().Add(validity),
KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
BasicConstraintsValid: true,
IsCA: true,
MaxPathLen: 1,
}
certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
if err != nil {
return nil, err
}
cert, err := x509.ParseCertificate(certDER)
if err != nil {
return nil, err
}
certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
return &testCA{Cert: cert, Key: key, CertPEM: certPEM}, nil
}
// generateIntermediateCA creates an intermediate CA signed by the given parent.
func generateIntermediateCA(cn string, validity time.Duration, parent *testCA) (*testCA, error) {
key, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader)
if err != nil {
return nil, err
}
serial, _ := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
tmpl := &x509.Certificate{
SerialNumber: serial,
Subject: pkix.Name{
CommonName: cn,
Organization: []string{"Test Partner Inc."},
Country: []string{"US"},
},
NotBefore: time.Now().Add(-1 * time.Hour),
NotAfter: time.Now().Add(validity),
KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
BasicConstraintsValid: true,
IsCA: true,
MaxPathLen: 0,
MaxPathLenZero: true,
}
certDER, err := x509.CreateCertificate(rand.Reader, tmpl, parent.Cert, &key.PublicKey, parent.Key)
if err != nil {
return nil, err
}
cert, err := x509.ParseCertificate(certDER)
if err != nil {
return nil, err
}
certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
return &testCA{Cert: cert, Key: key, CertPEM: certPEM}, nil
}
// licenseOpts configures a test license certificate.
type licenseOpts struct {
CN string
Org string
Country string
Email string
NotBefore time.Time
NotAfter time.Time
Serial *big.Int
Extensions []ExtensionValue
}
func defaultLicenseOpts() licenseOpts {
return licenseOpts{
CN: "Test Corp",
Org: "Test Inc.",
Country: "CN",
Email: "test@example.com",
NotBefore: time.Now().Add(-1 * time.Hour),
NotAfter: time.Now().Add(365 * 24 * time.Hour),
Serial: big.NewInt(10042),
Extensions: []ExtensionValue{
{OID: OIDProduct, Value: "yao"},
{OID: OIDEdition, Value: "pro"},
{OID: OIDMaxUsers, Value: "500"},
{OID: OIDMaxTaiNodes, Value: "10"},
{OID: OIDMaxAgents, Value: "50"},
{OID: OIDMaxSandboxes, Value: "20"},
{OID: OIDMaxAPIRPM, Value: "10000"},
{OID: OIDMaxStorageGB, Value: "100"},
{OID: OIDSupportLevel, Value: "priority"},
},
}
}
// generateLicenseCert creates a leaf license certificate signed by the given CA.
func generateLicenseCert(signer *testCA, opts licenseOpts) (*testLicenseCert, error) {
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
if err != nil {
return nil, err
}
tmpl := &x509.Certificate{
SerialNumber: opts.Serial,
Subject: pkix.Name{
CommonName: opts.CN,
Organization: []string{opts.Org},
Country: []string{opts.Country},
},
EmailAddresses: []string{opts.Email},
NotBefore: opts.NotBefore,
NotAfter: opts.NotAfter,
KeyUsage: x509.KeyUsageDigitalSignature,
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageAny},
}
for _, ext := range opts.Extensions {
tmpl.ExtraExtensions = append(tmpl.ExtraExtensions, pkix.Extension{
Id: ext.OID,
Value: []byte(ext.Value),
})
}
certDER, err := x509.CreateCertificate(rand.Reader, tmpl, signer.Cert, &key.PublicKey, signer.Key)
if err != nil {
return nil, err
}
cert, err := x509.ParseCertificate(certDER)
if err != nil {
return nil, err
}
certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
return &testLicenseCert{Cert: cert, CertPEM: certPEM}, nil
}

115
commercial/types.go Normal file
View file

@ -0,0 +1,115 @@
package commercial
import "time"
// License is the global commercial license state, populated by Load().
// Read-only after initialization; all modules may read without synchronization.
var License LicenseInfo
// LicenseInfo holds the parsed result of a commercial license certificate.
type LicenseInfo struct {
Valid bool `json:"valid"`
Source string `json:"source"` // "none" | "file" | "env"
LoadedAt int64 `json:"loaded_at"`
Error string `json:"error,omitempty"`
// Identity (from X.509 Subject)
LicenseeName string `json:"licensee_name"`
LicenseeOrg string `json:"licensee_org"`
LicenseeCountry string `json:"licensee_country,omitempty"`
LicenseeEmail string `json:"licensee_email,omitempty"`
SerialNumber string `json:"serial_number"`
NotBefore int64 `json:"not_before"`
NotAfter int64 `json:"not_after"`
IsExpired bool `json:"is_expired"`
Issuer string `json:"issuer"`
// Scope
Product []string `json:"product"`
Edition string `json:"edition"` // "community" | "starter" | "pro" | "enterprise"
Env []string `json:"env,omitempty"`
Domain string `json:"domain,omitempty"`
AppID string `json:"app_id,omitempty"`
MachineID string `json:"machine_id,omitempty"` // if set, must match runtime machine ID
// Quota (0 = unlimited)
MaxUsers int `json:"max_users"`
MaxTaiNodes int `json:"max_tai_nodes"`
MaxAgents int `json:"max_agents"`
MaxSandboxes int `json:"max_sandboxes"`
MaxAPIRPM int `json:"max_api_rpm"`
MaxStorageGB int `json:"max_storage_gb"`
// Permissions
Permissions Permissions `json:"permissions"`
}
// Permissions controls feature switches.
type Permissions struct {
AllowBrandingRemoval bool `json:"allow_branding_removal"`
AllowWhiteLabel bool `json:"allow_white_label"`
AllowMultiTenant bool `json:"allow_multi_tenant"`
AllowCustomDomain bool `json:"allow_custom_domain"`
AllowHostExec bool `json:"allow_host_exec"`
AllowSSO bool `json:"allow_sso"`
SupportLevel string `json:"support_level"` // "none" | "email" | "priority" | "dedicated"
}
// PublicInfo is the subset safe for well-known / public API exposure.
type PublicInfo struct {
Valid bool `json:"valid"`
Edition string `json:"edition"`
NotAfter int64 `json:"not_after,omitempty"`
Product []string `json:"product"`
}
// DefaultLicense returns community-level defaults when no certificate is present.
func DefaultLicense() LicenseInfo {
return LicenseInfo{
Source: "none",
LoadedAt: time.Now().Unix(),
Edition: "community",
Product: []string{"yao"},
MaxUsers: 100,
MaxTaiNodes: 1,
MaxAgents: 3,
MaxSandboxes: 1,
MaxAPIRPM: 1000,
MaxStorageGB: 10,
Permissions: Permissions{
SupportLevel: "none",
},
}
}
// GetPublicInfo returns the public-safe subset of the current license state.
func GetPublicInfo() *PublicInfo {
return &PublicInfo{
Valid: License.Valid,
Edition: License.Edition,
NotAfter: License.NotAfter,
Product: License.Product,
}
}
var editionRank = map[string]int{
"community": 0,
"starter": 1,
"pro": 2,
"enterprise": 3,
}
// IsLevel reports whether the license meets or exceeds the given minimum edition.
func (l LicenseInfo) IsLevel(minEdition string) bool {
return editionRank[l.Edition] >= editionRank[minEdition]
}
// HasProduct reports whether the license covers the given product name.
func (l LicenseInfo) HasProduct(product string) bool {
for _, p := range l.Product {
if p == product {
return true
}
}
return false
}

File diff suppressed because one or more lines are too long

View file

@ -23,6 +23,7 @@ import (
"github.com/yaoapp/yao/api"
"github.com/yaoapp/yao/attachment"
"github.com/yaoapp/yao/cert"
"github.com/yaoapp/yao/commercial"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/connector"
"github.com/yaoapp/yao/data"
@ -94,6 +95,27 @@ func loadStep(name string, loadFunc func() error, callback func(string, string))
return err
}
// LoadForMigrate loads only the minimal modules needed for schema migration:
// application config, database connection, and models (with auto-migrate).
func LoadForMigrate(cfg config.Config) (warnings []Warning, err error) {
defer func() { err = exception.Catch(recover()) }()
exception.Mode = cfg.Mode
if err = loadApp(cfg.AppSource); err != nil {
return append(warnings, Warning{Widget: "Load Application", Error: err}), err
}
if err = share.DBConnect(cfg.DB); err != nil {
return append(warnings, Warning{Widget: "DB", Error: err}), err
}
if err = model.Load(cfg); err != nil {
warnings = append(warnings, Warning{Widget: "Model", Error: err})
}
return warnings, err
}
// Load application engine
func Load(cfg config.Config, options LoadOption, progressCallback ...func(string, string)) (warnings []Warning, err error) {
@ -149,6 +171,12 @@ func Load(cfg config.Config, options LoadOption, progressCallback ...func(string
warnings = append(warnings, Warning{Widget: "Registry", Error: err})
}
// Load Commercial License
loadStep("License", func() error {
commercial.Load(cfg.Root, "yao")
return nil
}, callback)
// Load Certs
err = loadStep("Cert", func() error {
return cert.Load(cfg)

View file

@ -41,6 +41,7 @@ type CreateRobotRequest struct {
Agents interface{} `json:"agents,omitempty"` // Accessible agents (JSON array)
MCPServers interface{} `json:"mcp_servers,omitempty"` // MCP servers (JSON array)
LanguageModel string `json:"language_model,omitempty"` // Language model name
Workspace string `json:"workspace,omitempty"` // Workspace ID bound to this robot
// Limits
CostLimit float64 `json:"cost_limit,omitempty"` // Monthly cost limit USD
@ -73,6 +74,7 @@ type UpdateRobotRequest struct {
Agents interface{} `json:"agents,omitempty"` // Accessible agents
MCPServers interface{} `json:"mcp_servers,omitempty"` // MCP servers
LanguageModel *string `json:"language_model,omitempty"` // Language model name
Workspace *string `json:"workspace,omitempty"` // Workspace ID (nil=no change, ""=unbind)
// Limits
CostLimit *float64 `json:"cost_limit,omitempty"` // Monthly cost limit USD
@ -115,6 +117,7 @@ type Response struct {
Agents interface{} `json:"agents,omitempty"`
MCPServers interface{} `json:"mcp_servers,omitempty"`
LanguageModel string `json:"language_model,omitempty"`
Workspace string `json:"workspace,omitempty"`
// Limits
CostLimit float64 `json:"cost_limit,omitempty"`
@ -186,6 +189,7 @@ func NewResponse(r *robotapi.RobotResponse) *Response {
Agents: r.Agents,
MCPServers: r.MCPServers,
LanguageModel: r.LanguageModel,
Workspace: r.Workspace,
CostLimit: r.CostLimit,
InvitedBy: r.InvitedBy,
JoinedAt: r.JoinedAt,
@ -215,6 +219,7 @@ func (r *CreateRobotRequest) ToAPICreateRequest() *robotapi.CreateRobotRequest {
Agents: r.Agents,
MCPServers: r.MCPServers,
LanguageModel: r.LanguageModel,
Workspace: r.Workspace,
CostLimit: r.CostLimit,
}
}
@ -238,6 +243,7 @@ func (r *UpdateRobotRequest) ToAPIUpdateRequest() *robotapi.UpdateRobotRequest {
Agents: r.Agents,
MCPServers: r.MCPServers,
LanguageModel: r.LanguageModel,
Workspace: r.Workspace,
CostLimit: r.CostLimit,
}
}

View file

@ -167,7 +167,7 @@ var (
"member_id", "team_id", "user_id", "member_type", "display_name", "bio", "avatar", "email", "role_id", "is_owner", "status",
"system_prompt", "manager_id", "robot_email", "authorized_senders", "email_filter_rules",
"robot_config", "agents", "mcp_servers",
"language_model", "cost_limit", "autonomous_mode", "last_robot_activity", "robot_status",
"language_model", "workspace", "cost_limit", "autonomous_mode", "last_robot_activity", "robot_status",
"invitation_id", "invited_by", "invited_at", "joined_at", "invitation_token",
"invitation_expires_at", "last_active_at",
"login_count", "notes", "metadata", "created_at", "updated_at",

View file

@ -362,7 +362,7 @@ func (u *DefaultUser) CreateRobotMember(ctx context.Context, teamID string, robo
robotFields := []string{
"role_id", "system_prompt", "manager_id", "robot_email", "authorized_senders", "email_filter_rules",
"robot_config", "agents", "mcp_servers",
"language_model", "cost_limit", "autonomous_mode", "robot_status",
"language_model", "workspace", "cost_limit", "autonomous_mode", "robot_status",
"notes", "metadata",
"__yao_created_by", "__yao_updated_by", "__yao_team_id", "__yao_tenant_id",
}
@ -444,7 +444,7 @@ func (u *DefaultUser) UpdateRobotMember(ctx context.Context, memberID string, ro
robotFields := []string{
"role_id", "system_prompt", "manager_id", "robot_email", "authorized_senders", "email_filter_rules",
"robot_config", "agents", "mcp_servers",
"language_model", "cost_limit", "autonomous_mode", "robot_status",
"language_model", "workspace", "cost_limit", "autonomous_mode", "robot_status",
"notes", "metadata", "status",
"__yao_updated_by", "__yao_team_id", "__yao_tenant_id",
}

View file

@ -1152,6 +1152,7 @@ func TestMemberCreateRobot(t *testing.T) {
"mcp_tools": []string{"filesystem", "database"},
"autonomous_mode": "enabled",
"cost_limit": 100.50,
"workspace": "ws-test-create",
},
map[string]string{
"Authorization": "Bearer " + tokenInfo.AccessToken,
@ -1374,6 +1375,9 @@ func TestMemberCreateRobot(t *testing.T) {
if tc.body["prompt"] != nil {
assert.Equal(t, tc.body["prompt"], member["system_prompt"], "Should have correct system_prompt")
}
if tc.body["workspace"] != nil {
assert.Equal(t, "ws-test-create", member["workspace"], "Should have correct workspace")
}
}
}
}
@ -1595,6 +1599,7 @@ func TestMemberUpdateRobot(t *testing.T) {
"llm": "gpt-3.5-turbo",
"autonomous_mode": "disabled",
"cost_limit": 50.0,
"workspace": "ws-initial",
}
robotBodyBytes, _ := json.Marshal(robotBody)
robotReq, _ := http.NewRequest("POST", serverURL+baseURL+"/user/teams/"+teamID+"/members/robots", bytes.NewBuffer(robotBodyBytes))
@ -1655,6 +1660,7 @@ func TestMemberUpdateRobot(t *testing.T) {
"cost_limit": 100.0,
"status": "active",
"robot_status": "working",
"workspace": "ws-updated",
},
map[string]string{
"Authorization": "Bearer " + tokenInfo.AccessToken,
@ -1679,6 +1685,7 @@ func TestMemberUpdateRobot(t *testing.T) {
assert.Equal(t, fmt.Sprintf("https://example.com/avatars/full-%s.png", testUUID), member["avatar"])
assert.Equal(t, "Updated system prompt", member["system_prompt"])
assert.Equal(t, "gpt-4", member["language_model"])
assert.Equal(t, "ws-updated", member["workspace"], "Should have correct workspace")
}
}
},
@ -1973,6 +1980,35 @@ func TestMemberUpdateRobot(t *testing.T) {
}
},
},
{
"update workspace to unbind",
func() (string, string) { return createTestRobot("15") },
map[string]interface{}{
"workspace": "",
},
map[string]string{
"Authorization": "Bearer " + tokenInfo.AccessToken,
},
200,
"should unbind workspace by setting to empty string",
func(t *testing.T, memberID string) {
getMemberURL := serverURL + baseURL + "/user/teams/" + teamID + "/members/" + memberID
getReq, _ := http.NewRequest("GET", getMemberURL, nil)
getReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
client := &http.Client{}
getResp, err := client.Do(getReq)
assert.NoError(t, err)
if getResp != nil {
defer getResp.Body.Close()
if getResp.StatusCode == 200 {
var member map[string]interface{}
body, _ := io.ReadAll(getResp.Body)
json.Unmarshal(body, &member)
assert.Empty(t, member["workspace"], "Workspace should be empty after unbinding")
}
}
},
},
}
for _, tc := range testCases {

View file

@ -315,6 +315,9 @@ func GinMemberCreateRobot(c *gin.Context) {
if req.LanguageModel != "" {
baseData["language_model"] = req.LanguageModel
}
if req.Workspace != "" {
baseData["workspace"] = req.Workspace
}
if len(req.Agents) > 0 {
baseData["agents"] = req.Agents
}
@ -431,6 +434,9 @@ func GinMemberUpdateRobot(c *gin.Context) {
if req.LanguageModel != "" {
updateData["language_model"] = req.LanguageModel
}
if req.Workspace != nil {
updateData["workspace"] = *req.Workspace
}
if req.Status != "" {
updateData["status"] = req.Status
}
@ -1591,6 +1597,7 @@ func mapToMemberDetailResponse(data maps.MapStr) MemberDetailResponse {
SystemPrompt: utils.ToString(data["system_prompt"]),
ManagerID: utils.ToString(data["manager_id"]),
LanguageModel: utils.ToString(data["language_model"]),
Workspace: utils.ToString(data["workspace"]),
CostLimit: utils.ToFloat64(data["cost_limit"]),
AutonomousMode: data["autonomous_mode"], // Keep original type (bool or string)
LastRobotActivity: utils.ToTimeString(data["last_robot_activity"]),

View file

@ -457,6 +457,7 @@ type MemberDetailResponse struct {
Agents []string `json:"agents,omitempty"`
MCPServers []string `json:"mcp_servers,omitempty"`
LanguageModel string `json:"language_model,omitempty"`
Workspace string `json:"workspace,omitempty"`
CostLimit float64 `json:"cost_limit,omitempty"`
AutonomousMode interface{} `json:"autonomous_mode,omitempty"` // Can be bool or string
LastRobotActivity string `json:"last_robot_activity,omitempty"`
@ -480,6 +481,7 @@ type CreateRobotMemberRequest struct {
ManagerID string `json:"report_to,omitempty"` // Direct manager user ID
SystemPrompt string `json:"prompt" binding:"required"` // Identity & role prompt
LanguageModel string `json:"llm,omitempty"` // Language model (e.g., "gpt-4")
Workspace string `json:"workspace,omitempty"` // Workspace ID bound to this robot
Agents []string `json:"agents,omitempty"` // Accessible agents
MCPServers []string `json:"mcp_tools,omitempty"` // MCP servers/tools
AutonomousMode string `json:"autonomous_mode,omitempty"` // "enabled" or "disabled"
@ -499,6 +501,7 @@ type UpdateRobotMemberRequest struct {
ManagerID string `json:"report_to,omitempty"` // Direct manager user ID
SystemPrompt string `json:"prompt,omitempty"` // Identity & role prompt
LanguageModel string `json:"llm,omitempty"` // Language model (e.g., "gpt-4")
Workspace *string `json:"workspace"` // Workspace ID (nil=no change, ""=unbind)
Agents []string `json:"agents,omitempty"` // Accessible agents
MCPServers []string `json:"mcp_tools,omitempty"` // MCP servers/tools
AutonomousMode string `json:"autonomous_mode,omitempty"` // "enabled" or "disabled"

View file

@ -8,6 +8,7 @@ import (
"strings"
"github.com/gin-gonic/gin"
"github.com/yaoapp/yao/commercial"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/share"
)
@ -52,6 +53,9 @@ type YaoMetadata struct {
GRPC string `json:"grpc,omitempty"` // gRPC server address (e.g., "127.0.0.1:9099")
Optional map[string]interface{} `json:"optional,omitempty"` // Optional settings
// Commercial license
License *commercial.PublicInfo `json:"license,omitempty"`
// Developer information
Developer *share.Developer `json:"developer,omitempty"`
}
@ -76,6 +80,9 @@ func (openapi *OpenAPI) yaoMetadata(c *gin.Context) {
Optional: share.App.Optional,
}
// Include license info
metadata.License = commercial.GetPublicInfo()
// Include developer info if available
if share.App.Developer.ID != "" || share.App.Developer.Name != "" {
metadata.Developer = &share.App.Developer

View file

@ -75,6 +75,7 @@ const (
DefaultStopTimeout = 2 * time.Second
DefaultSessionIdleTimeout = 30 * time.Minute
DefaultLongRunningIdleTimeout = 2 * time.Hour
DefaultOneShotMaxAge = 8 * time.Hour
)
// ---------------------------------------------------------------------------

View file

@ -71,6 +71,31 @@ func (w *sandboxWatcher) Check(ctx context.Context) []monitor.Alert {
}
}
// OneShot safety net: these containers should have been removed by
// LifecycleAction right after execution. If they still exist after
// DefaultOneShotMaxAge it means cleanup failed (e.g. process crash,
// cfg.ID race before the fix). Remove them based on createdAt so we
// never kill a container that is still actively executing.
if b.policy == OneShot {
age := time.Since(b.createdAt)
if age > DefaultOneShotMaxAge {
alerts = append(alerts, monitor.Alert{
Level: monitor.Warn,
Target: "box:" + b.id,
Message: fmt.Sprintf("oneshot exceeded max age (%s > %s), removing", age.Round(time.Second), DefaultOneShotMaxAge),
Action: func(ctx context.Context) { mgr.Remove(ctx, b.id) },
})
} else {
alerts = append(alerts, monitor.Alert{
Level: monitor.Trace,
Target: "box:" + b.id,
Message: fmt.Sprintf("oneshot age %s (max=%s)",
age.Round(time.Second), DefaultOneShotMaxAge),
})
}
return true
}
idle := time.Since(b.idleSince())
timeout := b.idleTimeout()

View file

@ -185,12 +185,18 @@ func TestRegister_Ping(t *testing.T) {
t.Fatal(err)
}
pong, err := stream.Recv()
if err != nil {
t.Fatal(err)
}
if pong.Type != "pong" {
t.Errorf("expected pong, got %q", pong.Type)
// Loop until we receive "pong"; skip "open" frames that may arrive from
// the asynchronous connectTunnelNode goroutine if a real gRPC endpoint
// happens to be reachable in the test environment.
for {
pong, err := stream.Recv()
if err != nil {
t.Fatal(err)
}
if pong.Type == "pong" {
break
}
// skip unexpected frames (e.g. "open" from connectTunnelNode)
}
stream.CloseSend()

View file

@ -214,6 +214,15 @@
"length": 100,
"nullable": true
},
{
"name": "workspace",
"type": "string",
"label": "Workspace",
"comment": "Workspace ID bound to this robot member (nullable = not bound)",
"length": 255,
"nullable": true,
"index": true
},
{
"name": "cost_limit",
"type": "decimal",