8fdd1a3a6c
- Introduced a new `workspace` field across various robot-related structures, including `CreateRobotRequest`, `UpdateRobotRequest`, and `RobotResponse`, allowing for better organization and management of robots within specific workspaces. - Updated database queries and response mappings to accommodate the new workspace field, ensuring seamless integration with existing functionalities. - Enhanced agent execution context to include workspace information, improving the contextual awareness of agents during operations. - Added tests to validate the creation and updating of robots with workspace data, ensuring robust functionality and backward compatibility.
394 lines
13 KiB
Go
394 lines
13 KiB
Go
package sandboxv2
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"fmt"
|
|
mathrand "math/rand"
|
|
"strings"
|
|
|
|
"github.com/yaoapp/kun/log"
|
|
agentContext "github.com/yaoapp/yao/agent/context"
|
|
"github.com/yaoapp/yao/agent/sandbox/v2/types"
|
|
infra "github.com/yaoapp/yao/sandbox/v2"
|
|
"github.com/yaoapp/yao/tai"
|
|
"github.com/yaoapp/yao/tai/registry"
|
|
"github.com/yaoapp/yao/workspace"
|
|
)
|
|
|
|
// BuildIdentifier determines the Computer identifier based on lifecycle policy
|
|
// and optional metadata override. Returns "" for oneshot (always new).
|
|
func BuildIdentifier(cfg *types.SandboxConfig, ownerID, chatID, assistantID, workspaceID string, metadata map[string]any) string {
|
|
if cfg.Lifecycle == "oneshot" {
|
|
return ""
|
|
}
|
|
|
|
switch cfg.Lifecycle {
|
|
case "session":
|
|
return fmt.Sprintf("%s-%s-%s", ownerID, assistantID, chatID)
|
|
case "longrunning", "persistent":
|
|
return fmt.Sprintf("%s-%s.%s", ownerID, assistantID, workspaceID)
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
// ResolveNodeID determines the target nodeID and computer kind based on
|
|
// metadata and DSL configuration, without creating or acquiring a container.
|
|
// Returns (nodeID, kind, error). kind is "box" or "host".
|
|
func ResolveNodeID(ctx *agentContext.Context, cfg *types.SandboxConfig, manager *infra.Manager) (string, string, error) {
|
|
computerID := ""
|
|
if ctx.Metadata != nil {
|
|
if cid, ok := ctx.Metadata["computer_id"].(string); ok && cid != "" {
|
|
computerID = cid
|
|
}
|
|
}
|
|
|
|
workspaceID := ""
|
|
if ctx.Metadata != nil {
|
|
if ws, ok := ctx.Metadata["workspace_id"].(string); ok && ws != "" {
|
|
workspaceID = ws
|
|
}
|
|
}
|
|
ownerID := resolveOwnerID(ctx)
|
|
|
|
log.Trace("[sandbox/v2] ResolveNodeID: computerID=%q workspaceID=%q ownerID=%q image=%q", computerID, workspaceID, ownerID, cfg.Computer.Image)
|
|
|
|
if workspaceID != "" {
|
|
wsNode, err := workspace.M().NodeForWorkspace(context.Background(), workspaceID)
|
|
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)
|
|
}
|
|
}
|
|
|
|
if computerID == "" {
|
|
pickedID, err := pickNodeByFilter(cfg.Filter, cfg.Computer.Image)
|
|
if err != nil {
|
|
return "", "", fmt.Errorf("auto-select node for ResolveNodeID: %w", err)
|
|
}
|
|
log.Trace("[sandbox/v2] ResolveNodeID: pickNodeByFilter -> %s", pickedID)
|
|
computerID = pickedID
|
|
cfg.NodeID = pickedID
|
|
}
|
|
|
|
if node, ok := tai.GetNodeMeta(computerID); ok {
|
|
hasContainerRuntime := node.Capabilities.Docker || node.Capabilities.K8s
|
|
log.Trace("[sandbox/v2] ResolveNodeID: node=%q HostExec=%v Docker=%v K8s=%v hasContainer=%v", computerID, node.Capabilities.HostExec, node.Capabilities.Docker, node.Capabilities.K8s, hasContainerRuntime)
|
|
if node.Capabilities.HostExec && !hasContainerRuntime {
|
|
log.Trace("[sandbox/v2] ResolveNodeID: -> host (host-only node)")
|
|
return computerID, "host", nil
|
|
}
|
|
if node.Capabilities.HostExec && hasContainerRuntime && cfg.Computer.Image == "" {
|
|
log.Trace("[sandbox/v2] ResolveNodeID: -> host (dual-capable, no image)")
|
|
return computerID, "host", nil
|
|
}
|
|
if !hasContainerRuntime {
|
|
return "", "", fmt.Errorf("node %q has no container runtime and no host_exec capability", computerID)
|
|
}
|
|
log.Trace("[sandbox/v2] ResolveNodeID: -> box")
|
|
return computerID, "box", nil
|
|
}
|
|
log.Trace("[sandbox/v2] ResolveNodeID: node %q not found in registry, assuming box", computerID)
|
|
return computerID, "box", nil
|
|
}
|
|
|
|
// GetComputer obtains or creates a Computer for the current request.
|
|
// Connector config is injected per-execution inside ClaudeRunner.Stream
|
|
// via "tai a2o config put".
|
|
// Returns the Computer, the resolved identifier, and any error.
|
|
func GetComputer(ctx *agentContext.Context, cfg *types.SandboxConfig, manager *infra.Manager) (infra.Computer, string, error) {
|
|
ownerID := resolveOwnerID(ctx)
|
|
|
|
workspaceID := ""
|
|
if ctx.Metadata != nil {
|
|
if ws, ok := ctx.Metadata["workspace_id"].(string); ok && ws != "" {
|
|
workspaceID = ws
|
|
}
|
|
}
|
|
|
|
identifier := BuildIdentifier(cfg, ownerID, ctx.ChatID, ctx.AssistantID, workspaceID, ctx.Metadata)
|
|
|
|
// Fill runtime fields.
|
|
cfg.Owner = ownerID
|
|
cfg.ID = identifier
|
|
cfg.WorkspaceID = workspaceID
|
|
|
|
// Resolve computer_id from metadata to determine kind and nodeID.
|
|
computerID := ""
|
|
if ctx.Metadata != nil {
|
|
if cid, ok := ctx.Metadata["computer_id"].(string); ok && cid != "" {
|
|
computerID = cid
|
|
}
|
|
}
|
|
|
|
// Workspace-wins rule: when workspace_id is present,
|
|
// the workspace's bound node takes precedence over computer_id.
|
|
if workspaceID != "" {
|
|
wsNode, err := workspace.M().NodeForWorkspace(context.Background(), workspaceID)
|
|
if err == nil && wsNode != "" {
|
|
if computerID != "" && computerID != wsNode {
|
|
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)
|
|
}
|
|
}
|
|
|
|
log.Trace("[sandbox/v2] GetComputer: computerID=%q workspaceID=%q ownerID=%q cfgNodeID=%q image=%q", computerID, workspaceID, ownerID, cfg.NodeID, cfg.Computer.Image)
|
|
|
|
if computerID != "" {
|
|
log.Trace("[sandbox/v2] GetComputer: -> resolveComputerByID(%s)", computerID)
|
|
return resolveComputerByID(cfg, manager, computerID, ownerID, identifier, workspaceID)
|
|
}
|
|
|
|
log.Trace("[sandbox/v2] GetComputer: -> resolveComputerByDSL (no computerID)")
|
|
return resolveComputerByDSL(cfg, manager, ownerID, identifier, workspaceID)
|
|
}
|
|
|
|
// resolveComputerByID dispatches based on the runtime computer_id from metadata.
|
|
// It queries the registry and sandbox manager to determine the computer kind.
|
|
func resolveComputerByID(
|
|
cfg *types.SandboxConfig, manager *infra.Manager,
|
|
computerID, ownerID, identifier, workspaceID string,
|
|
) (infra.Computer, string, error) {
|
|
|
|
// 1) Check if computer_id is a known Tai node (host or node kind).
|
|
if node, ok := tai.GetNodeMeta(computerID); ok {
|
|
cfg.NodeID = computerID
|
|
hasContainerRuntime := node.Capabilities.Docker || node.Capabilities.K8s
|
|
log.Trace("[sandbox/v2] resolveComputerByID: node=%q found=true HostExec=%v Docker=%v K8s=%v hasContainer=%v image=%q", computerID, node.Capabilities.HostExec, node.Capabilities.Docker, node.Capabilities.K8s, hasContainerRuntime, cfg.Computer.Image)
|
|
|
|
if node.Capabilities.HostExec && !hasContainerRuntime {
|
|
log.Trace("[sandbox/v2] resolveComputerByID: -> host (host-only node)")
|
|
cfg.Kind = "host"
|
|
host, err := manager.Host(context.Background(), computerID)
|
|
if err != nil {
|
|
return nil, identifier, fmt.Errorf("get host computer: %w", err)
|
|
}
|
|
host.BindWorkplace(workspaceID)
|
|
return host, identifier, nil
|
|
}
|
|
|
|
if node.Capabilities.HostExec && hasContainerRuntime && cfg.Computer.Image == "" {
|
|
// Dual-capable node with no image in DSL: prefer host mode.
|
|
cfg.Kind = "host"
|
|
host, err := manager.Host(context.Background(), computerID)
|
|
if err != nil {
|
|
return nil, identifier, fmt.Errorf("get host computer: %w", err)
|
|
}
|
|
host.BindWorkplace(workspaceID)
|
|
return host, identifier, nil
|
|
}
|
|
|
|
if !hasContainerRuntime {
|
|
return nil, identifier, fmt.Errorf("node %q has no container runtime and no host_exec capability", computerID)
|
|
}
|
|
|
|
// Node with container runtime and DSL has image: create/reuse a box.
|
|
cfg.Kind = "box"
|
|
return resolveBox(cfg, manager, ownerID, identifier, workspaceID)
|
|
}
|
|
|
|
// 2) Check if computer_id is an existing box ID.
|
|
if manager != nil {
|
|
box, err := manager.Get(context.Background(), computerID)
|
|
if err == nil && box != nil {
|
|
cfg.Kind = "box"
|
|
box.BindWorkplace(workspaceID)
|
|
return box, computerID, nil
|
|
}
|
|
}
|
|
|
|
return nil, identifier, fmt.Errorf("computer %q not found in registry or sandbox manager", computerID)
|
|
}
|
|
|
|
// resolveComputerByDSL dispatches based on DSL static configuration (cfg.Computer.Image).
|
|
func resolveComputerByDSL(
|
|
cfg *types.SandboxConfig, manager *infra.Manager,
|
|
ownerID, identifier, workspaceID string,
|
|
) (infra.Computer, string, error) {
|
|
|
|
log.Trace("[sandbox/v2] resolveComputerByDSL: cfgNodeID=%q image=%q", cfg.NodeID, cfg.Computer.Image)
|
|
|
|
if cfg.NodeID == "" {
|
|
pickedID, err := pickNodeByFilter(cfg.Filter, cfg.Computer.Image)
|
|
if err != nil {
|
|
return nil, identifier, fmt.Errorf("auto-select node: %w", err)
|
|
}
|
|
log.Trace("[sandbox/v2] resolveComputerByDSL: pickNodeByFilter -> %s", pickedID)
|
|
cfg.NodeID = pickedID
|
|
}
|
|
|
|
log.Trace("[sandbox/v2] resolveComputerByDSL: -> resolveComputerByID(%s)", cfg.NodeID)
|
|
return resolveComputerByID(cfg, manager, cfg.NodeID, ownerID, identifier, workspaceID)
|
|
}
|
|
|
|
// resolveBox reuses or creates a box container.
|
|
func resolveBox(
|
|
cfg *types.SandboxConfig, manager *infra.Manager,
|
|
ownerID, identifier, workspaceID string,
|
|
) (infra.Computer, string, error) {
|
|
|
|
if workspaceID == "" && cfg.NodeID != "" {
|
|
workspaceID = workspace.DefaultWorkspaceID(ownerID, cfg.NodeID)
|
|
cfg.WorkspaceID = workspaceID
|
|
if dot := strings.LastIndex(identifier, "."); dot >= 0 {
|
|
identifier = identifier[:dot+1] + workspaceID
|
|
cfg.ID = identifier
|
|
}
|
|
}
|
|
|
|
// Reuse: non-empty identifier → try Get first.
|
|
if identifier != "" {
|
|
box, err := manager.Get(context.Background(), identifier)
|
|
if err == nil && box != nil {
|
|
if box.IsStopped() {
|
|
if startErr := manager.StartBox(context.Background(), identifier); startErr != nil {
|
|
log.Trace("[sandbox/v2] auto-start stopped box %s failed: %v, creating new", identifier, startErr)
|
|
} else {
|
|
box.BindWorkplace(workspaceID)
|
|
return box, identifier, nil
|
|
}
|
|
} else {
|
|
box.BindWorkplace(workspaceID)
|
|
return box, identifier, nil
|
|
}
|
|
}
|
|
}
|
|
|
|
// Create new box.
|
|
createOpts, err := BuildCreateOptions(cfg, identifier, ownerID, workspaceID)
|
|
if err != nil {
|
|
return nil, identifier, fmt.Errorf("build create options: %w", err)
|
|
}
|
|
log.Trace("[sandbox/v2] resolveBox: createOpts NodeID=%q Image=%q WorkspaceID=%q ID=%q Owner=%q", createOpts.NodeID, createOpts.Image, createOpts.WorkspaceID, createOpts.ID, createOpts.Owner)
|
|
|
|
// Oneshot with empty identifier: generate a random one.
|
|
if createOpts.ID == "" {
|
|
createOpts.ID = randomID()
|
|
identifier = createOpts.ID
|
|
cfg.ID = identifier
|
|
}
|
|
|
|
box, err := manager.Create(context.Background(), createOpts)
|
|
if err != nil {
|
|
return nil, identifier, fmt.Errorf("create computer: %w", err)
|
|
}
|
|
return box, identifier, nil
|
|
}
|
|
|
|
// LifecycleAction performs the post-request lifecycle operation based on policy.
|
|
// Called in defer after executeSandboxStream completes.
|
|
func LifecycleAction(ctx context.Context, cfg *types.SandboxConfig, computer infra.Computer, manager *infra.Manager) {
|
|
if computer == nil || cfg == nil {
|
|
return
|
|
}
|
|
|
|
info := computer.ComputerInfo()
|
|
|
|
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)
|
|
}
|
|
}
|
|
|
|
case "session", "longrunning":
|
|
if info.Kind == "box" && manager != nil {
|
|
manager.Heartbeat(cfg.ID, false, 0) // active=false: request finished, start idle timer
|
|
}
|
|
|
|
case "persistent":
|
|
// No action — persistent boxes survive indefinitely.
|
|
}
|
|
}
|
|
|
|
// resolveOwnerID returns teamID if available, otherwise userID.
|
|
func resolveOwnerID(ctx *agentContext.Context) string {
|
|
if ctx.Authorized != nil {
|
|
if ctx.Authorized.TeamID != "" {
|
|
return ctx.Authorized.TeamID
|
|
}
|
|
if ctx.Authorized.UserID != "" {
|
|
return ctx.Authorized.UserID
|
|
}
|
|
}
|
|
return "anonymous"
|
|
}
|
|
|
|
// pickNodeByFilter selects a random online node that satisfies the given filter
|
|
// and image requirement. If image is non-empty, candidate nodes must have a
|
|
// container runtime (Docker or K8s).
|
|
func pickNodeByFilter(filter *types.ComputerFilter, image string) (string, error) {
|
|
reg := registry.Global()
|
|
if reg == nil {
|
|
return "", fmt.Errorf("tai registry not initialized")
|
|
}
|
|
|
|
nodes := reg.List()
|
|
var candidates []string
|
|
for _, n := range nodes {
|
|
if n.Status != "online" && n.Status != "" {
|
|
continue
|
|
}
|
|
|
|
if filter != nil {
|
|
if filter.OS != "" && !strings.EqualFold(n.System.OS, filter.OS) {
|
|
continue
|
|
}
|
|
if filter.Arch != "" && !strings.EqualFold(n.System.Arch, filter.Arch) {
|
|
continue
|
|
}
|
|
if len(filter.Kind) > 0 {
|
|
matched := false
|
|
for _, k := range filter.Kind {
|
|
switch strings.ToLower(k) {
|
|
case "host":
|
|
if n.Capabilities.HostExec {
|
|
matched = true
|
|
}
|
|
case "box":
|
|
if n.Capabilities.Docker || n.Capabilities.K8s {
|
|
matched = true
|
|
}
|
|
}
|
|
}
|
|
if !matched {
|
|
continue
|
|
}
|
|
}
|
|
}
|
|
|
|
if image != "" && !(n.Capabilities.Docker || n.Capabilities.K8s) {
|
|
continue
|
|
}
|
|
|
|
candidates = append(candidates, n.TaiID)
|
|
}
|
|
|
|
if len(candidates) == 0 {
|
|
kind := ""
|
|
os := ""
|
|
arch := ""
|
|
if filter != nil {
|
|
kind = fmt.Sprintf("%v", []string(filter.Kind))
|
|
os = filter.OS
|
|
arch = filter.Arch
|
|
}
|
|
return "", fmt.Errorf("no online node matches filter (kind=%s os=%s arch=%s image=%s)", kind, os, arch, image)
|
|
}
|
|
|
|
return candidates[mathrand.Intn(len(candidates))], nil
|
|
}
|
|
|
|
func randomID() string {
|
|
b := make([]byte, 8)
|
|
_, _ = rand.Read(b)
|
|
return hex.EncodeToString(b)
|
|
}
|