feat(sandbox): enhance sandbox initialization and image management

- Updated Docker run commands in CI workflows to include the `-direct` flag for improved server operation.
- Implemented image existence checks and automatic pulling for sandbox environments, enhancing reliability during initialization.
- Added loading status updates for sandbox operations, providing better feedback during the setup process.
- Refactored lifecycle management to ensure accurate tracking of sandbox states and improved error handling.

Made-with: Cursor
This commit is contained in:
Max 2026-03-15 13:17:58 +08:00
parent 3797091d37
commit a5c8109db8
11 changed files with 355 additions and 51 deletions

View file

@ -1088,7 +1088,7 @@ jobs:
docker run -d --name tai-docker \
-v /var/run/docker.sock:/var/run/docker.sock \
-p 19100:19100 -p 8099:8099 -p 12375:12375 -p 16080:16080 \
yaoapp/tai:latest server \
yaoapp/tai:latest server -direct \
-grpc 0.0.0.0:19100 -http 0.0.0.0:8099 -vnc 0.0.0.0:16080 -docker 0.0.0.0:12375
for i in $(seq 1 30); do

View file

@ -799,7 +799,7 @@ jobs:
docker run -d --name tai-docker \
-v /var/run/docker.sock:/var/run/docker.sock \
-p 19100:19100 -p 8099:8099 -p 12375:12375 -p 16080:16080 \
yaoapp/tai:latest server \
yaoapp/tai:latest server -direct \
-grpc 0.0.0.0:19100 -http 0.0.0.0:8099 -vnc 0.0.0.0:16080 -docker 0.0.0.0:12375
for i in $(seq 1 30); do

View file

@ -224,35 +224,41 @@ func (s *streamState) handleToolCall(data []byte) int {
var deltaPath string
if len(toolCallArray) == 1 {
// Single tool call - flatten to props root level
tc := toolCallArray[0]
props = map[string]interface{}{}
// Static fields (only in first chunk): use merge
hasStaticFields := false
if id, ok := tc["id"].(string); ok {
props["id"] = id
hasStaticFields = true
}
if typ, ok := tc["type"].(string); ok {
props["type"] = typ
hasStaticFields = true
}
if index, ok := tc["index"].(float64); ok {
props["index"] = int(index)
hasStaticFields = true
}
if fn, ok := tc["function"].(map[string]interface{}); ok {
if name, ok := fn["name"].(string); ok {
props["name"] = name
hasStaticFields = true
}
// Arguments field: use append
if args, ok := fn["arguments"].(string); ok {
props["arguments"] = args
// If this chunk has arguments, use append action for arguments field
deltaAction = "append"
deltaPath = "arguments"
}
}
// If no arguments in this chunk, use merge for other fields
if deltaAction == "" {
if hasStaticFields {
// First chunk with id/name/type: merge so all fields are applied.
// arguments="" is included but that's fine — subsequent appends build on it.
deltaAction = "merge"
} else if _, ok := props["arguments"]; ok {
// Subsequent chunk with only arguments fragment: append to arguments.
deltaAction = "append"
deltaPath = "arguments"
} else {
deltaAction = "merge"
}
} else {

View file

@ -51,7 +51,34 @@ func (ast *Assistant) initSandboxV2(ctx *context.Context, opts *context.Options)
// 2. Build human-readable DisplayName from real Agent name + Workspace name.
cfg.DisplayName = buildBoxDisplayName(ctx, ast.ID, ast.Name)
// 2.5. Image existence check + pull (for box mode).
if cfg.Computer.Image != "" && manager != nil {
nodeID, kind, _ := sandboxv2.ResolveNodeID(ctx, cfg, manager)
if kind == "box" && nodeID != "" {
updateLoadingV2(ctx, loadingMsgID, "sandbox.starting")
exists, existsErr := manager.ImageExists(stdCtx, nodeID, cfg.Computer.Image)
if existsErr != nil {
log.Printf("[sandbox/v2] image exists check failed on node %s: %v", nodeID, existsErr)
}
if existsErr == nil && !exists {
updateLoadingV2(ctx, loadingMsgID, "sandbox.pulling_image")
ch, pullErr := manager.PullImage(stdCtx, nodeID, cfg.Computer.Image, infraV2.ImagePullOptions{})
if pullErr != nil {
log.Printf("[sandbox/v2] image pull failed on node %s: %v (will retry in Create)", nodeID, pullErr)
} else if ch != nil {
for p := range ch {
if p.Error != "" {
log.Printf("[sandbox/v2] image pull progress error: %s", p.Error)
break
}
}
}
}
}
}
// 3. Obtain Computer (passes connector for OPENAI_PROXY_* env injection).
updateLoadingV2(ctx, loadingMsgID, "sandbox.starting")
computer, identifier, err := sandboxv2.GetComputer(ctx, cfg, manager, conn)
if err != nil {
closeLoadingV2(ctx, loadingMsgID, "sandbox.failed")
@ -89,6 +116,7 @@ func (ast *Assistant) initSandboxV2(ctx *context.Context, opts *context.Options)
}
// 7. Runner.Prepare (standard context).
updateLoadingV2(ctx, loadingMsgID, "sandbox.configuring")
err = runner.Prepare(stdCtx, &sandboxTypes.PrepareRequest{
Computer: computer,
Config: cfg,
@ -133,11 +161,6 @@ func (ast *Assistant) executeSandboxV2Stream(
cfg := ast.SandboxV2
manager := infraV2.M()
// Close the "preparing" loading on first output.
if loadingMsgID != "" {
closeLoadingV2(ctx, loadingMsgID, "")
}
// Build system prompt.
var systemPrompt string
if len(ast.Prompts) > 0 {
@ -172,11 +195,12 @@ func (ast *Assistant) executeSandboxV2Stream(
}
execReq := &sandboxv2.ExecuteRequest{
Computer: computer,
Runner: runner,
Config: cfg,
StreamReq: streamReq,
Manager: manager,
Computer: computer,
Runner: runner,
Config: cfg,
StreamReq: streamReq,
Manager: manager,
LoadingMsgID: loadingMsgID,
}
return sandboxv2.ExecuteSandboxStream(ctx, execReq, streamHandler)
@ -230,6 +254,22 @@ func buildBoxDisplayName(ctx *context.Context, assistantID, rawName string) stri
return ""
}
func updateLoadingV2(ctx *context.Context, loadingMsgID, msgKey string) {
if loadingMsgID == "" || ctx == nil || msgKey == "" {
return
}
msg := &message.Message{
MessageID: loadingMsgID,
Delta: true,
DeltaAction: message.DeltaReplace,
Type: message.TypeLoading,
Props: map[string]any{
"message": i18n.T(ctx.Locale, msgKey),
},
}
ctx.Send(msg)
}
func closeLoadingV2(ctx *context.Context, loadingMsgID, msgKey string) {
if loadingMsgID == "" || ctx == nil {
return

View file

@ -100,11 +100,14 @@ func init() {
"kb.chat.description": "Auto-created knowledge base collection for chat sessions",
// Sandbox: assistant/sandbox.go - Sandbox status messages
"sandbox.preparing": "Preparing sandbox environment",
"sandbox.ready": "Sandbox ready",
"sandbox.working": "Working on your request",
"sandbox.completed": "Completed",
"sandbox.failed": "Execution failed",
"sandbox.preparing": "Preparing sandbox environment",
"sandbox.ready": "Sandbox ready",
"sandbox.working": "Working on your request",
"sandbox.completed": "Completed",
"sandbox.failed": "Execution failed",
"sandbox.starting": "Starting sandbox environment",
"sandbox.configuring": "Configuring runtime environment",
"sandbox.pulling_image": "Pulling container image",
// Sandbox: claude/executor.go - Tool execution messages
"sandbox.tool.read": "Reading file",
@ -224,11 +227,14 @@ func init() {
"kb.chat.description": "自动为聊天会话创建的知识库集合",
// Sandbox: assistant/sandbox.go - Sandbox status messages
"sandbox.preparing": "正在准备沙箱环境",
"sandbox.ready": "沙箱环境就绪",
"sandbox.working": "正在处理您的请求",
"sandbox.completed": "处理完成",
"sandbox.failed": "执行失败",
"sandbox.preparing": "正在准备沙箱环境",
"sandbox.ready": "沙箱环境就绪",
"sandbox.working": "正在处理您的请求",
"sandbox.completed": "处理完成",
"sandbox.failed": "执行失败",
"sandbox.starting": "正在启动沙箱环境",
"sandbox.configuring": "正在配置运行环境",
"sandbox.pulling_image": "正在拉取容器镜像",
// Sandbox: claude/executor.go - Tool execution messages
"sandbox.tool.read": "正在读取文件",
@ -376,11 +382,14 @@ func init() {
"kb.chat.description": "自动为聊天会话创建的知识库集合",
// Sandbox: assistant/sandbox.go - Sandbox status messages
"sandbox.preparing": "正在准备沙箱环境",
"sandbox.ready": "沙箱环境就绪",
"sandbox.working": "正在处理您的请求",
"sandbox.completed": "处理完成",
"sandbox.failed": "执行失败",
"sandbox.preparing": "正在准备沙箱环境",
"sandbox.ready": "沙箱环境就绪",
"sandbox.working": "正在处理您的请求",
"sandbox.completed": "处理完成",
"sandbox.failed": "执行失败",
"sandbox.starting": "正在启动沙箱环境",
"sandbox.configuring": "正在配置运行环境",
"sandbox.pulling_image": "正在拉取容器镜像",
// Sandbox: claude/executor.go - Tool execution messages
"sandbox.tool.read": "正在读取文件",

View file

@ -23,9 +23,13 @@ func parseStreamJSON(_ context.Context, stdout io.ReadCloser, handler message.St
scanner.Buffer(buf, 1024*1024)
messageStarted := false
toolBlockActive := false
toolIndex := 0
type toolState struct {
id string
name string
index int
inputJSON strings.Builder
}
var currentTool *toolState
@ -66,10 +70,37 @@ func parseStreamJSON(_ context.Context, stdout io.ReadCloser, handler message.St
blockType, _ := cb["type"].(string)
if blockType == "tool_use" {
toolName, _ := cb["name"].(string)
currentTool = &toolState{name: toolName}
toolID, _ := cb["id"].(string)
if toolID == "" {
toolID = fmt.Sprintf("tool_%d_%d", toolIndex, time.Now().UnixNano())
}
currentTool = &toolState{id: toolID, name: toolName, index: toolIndex}
toolIndex++
if handler != nil {
data, _ := json.Marshal(map[string]any{"tool": toolName})
if handler(message.ChunkToolCall, data) != 0 {
if !toolBlockActive {
startData := message.EventMessageStartData{
MessageID: fmt.Sprintf("sandbox-tool-%d", time.Now().UnixNano()),
Type: "tool_call",
Timestamp: time.Now().UnixMilli(),
}
sd, _ := json.Marshal(startData)
if handler(message.ChunkMessageStart, sd) != 0 {
stopped = true
break
}
toolBlockActive = true
}
tcData, _ := json.Marshal([]map[string]any{{
"index": currentTool.index,
"id": currentTool.id,
"type": "function",
"function": map[string]any{
"name": toolName,
"arguments": "",
},
}})
if handler(message.ChunkToolCall, tcData) != 0 {
stopped = true
}
}
@ -82,7 +113,13 @@ func parseStreamJSON(_ context.Context, stdout io.ReadCloser, handler message.St
switch deltaType {
case "text_delta":
if text, ok := delta["text"].(string); ok && text != "" {
text = strings.ReplaceAll(text, "\r\n", "\n")
text = strings.ReplaceAll(text, "\r", "\n")
if handler != nil {
if toolBlockActive {
handler(message.ChunkMessageEnd, nil)
toolBlockActive = false
}
if !messageStarted {
startData := message.EventMessageStartData{
MessageID: fmt.Sprintf("sandbox-%d", time.Now().UnixNano()),
@ -105,6 +142,17 @@ func parseStreamJSON(_ context.Context, stdout io.ReadCloser, handler message.St
if currentTool != nil {
if partial, ok := delta["partial_json"].(string); ok {
currentTool.inputJSON.WriteString(partial)
if handler != nil {
tcData, _ := json.Marshal([]map[string]any{{
"index": currentTool.index,
"function": map[string]any{
"arguments": partial,
},
}})
if handler(message.ChunkToolCall, tcData) != 0 {
stopped = true
}
}
}
}
}
@ -125,8 +173,53 @@ func parseStreamJSON(_ context.Context, stdout io.ReadCloser, handler message.St
continue
}
itemType, _ := ci["type"].(string)
if itemType == "tool_use" && handler != nil {
toolName, _ := ci["name"].(string)
toolID, _ := ci["id"].(string)
if toolID == "" {
toolID = fmt.Sprintf("tool_%d_%d", toolIndex, time.Now().UnixNano())
}
inputRaw, _ := json.Marshal(ci["input"])
idx := toolIndex
toolIndex++
if !toolBlockActive {
startData := message.EventMessageStartData{
MessageID: fmt.Sprintf("sandbox-tool-%d", time.Now().UnixNano()),
Type: "tool_call",
Timestamp: time.Now().UnixMilli(),
}
sd, _ := json.Marshal(startData)
if handler(message.ChunkMessageStart, sd) != 0 {
stopped = true
break
}
toolBlockActive = true
}
tcData, _ := json.Marshal([]map[string]any{{
"index": idx,
"id": toolID,
"type": "function",
"function": map[string]any{
"name": toolName,
"arguments": string(inputRaw),
},
}})
if handler(message.ChunkToolCall, tcData) != 0 {
stopped = true
break
}
}
if itemType == "text" {
if text, ok := ci["text"].(string); ok && text != "" && handler != nil && !messageStarted {
text = strings.ReplaceAll(text, "\r\n", "\n")
text = strings.ReplaceAll(text, "\r", "\n")
if toolBlockActive {
handler(message.ChunkMessageEnd, nil)
toolBlockActive = false
}
startData := message.EventMessageStartData{
MessageID: fmt.Sprintf("sandbox-%d", time.Now().UnixNano()),
Type: "text",
@ -159,8 +252,14 @@ func parseStreamJSON(_ context.Context, stdout io.ReadCloser, handler message.St
return fmt.Errorf("Claude CLI error: %s", result)
}
}
if handler != nil && messageStarted {
handler(message.ChunkMessageEnd, nil)
if handler != nil {
if toolBlockActive {
handler(message.ChunkMessageEnd, nil)
toolBlockActive = false
}
if messageStarted {
handler(message.ChunkMessageEnd, nil)
}
}
case "error":

View file

@ -227,6 +227,18 @@ func (r *ClaudeRunner) buildCLICommand(req *types.StreamRequest, oe *osEnv, isCo
env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = model
env["CLAUDE_CODE_SUBAGENT_MODEL"] = model
}
if thinking, ok := setting["thinking"].(map[string]interface{}); ok {
thinkType, _ := thinking["type"].(string)
switch thinkType {
case "disabled":
env["MAX_THINKING_TOKENS"] = "0"
case "enabled":
if budget, ok := thinking["budget_tokens"].(float64); ok && budget > 0 {
env["MAX_THINKING_TOKENS"] = fmt.Sprintf("%d", int(budget))
}
}
}
}
if req.Config != nil && len(req.Config.Secrets) > 0 {

View file

@ -32,6 +32,61 @@ func BuildIdentifier(cfg *types.SandboxConfig, ownerID, chatID, assistantID, wor
}
}
// 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)
if workspaceID == "" {
workspaceID = ownerID
}
if workspaceID != "" && workspaceID != ownerID {
wsNode, err := workspace.M().NodeForWorkspace(context.Background(), workspaceID)
if err == nil && wsNode != "" {
computerID = wsNode
}
}
if computerID != "" {
if node, ok := tai.GetNodeMeta(computerID); ok {
hasContainerRuntime := node.Capabilities.Docker || node.Capabilities.K8s
if node.Capabilities.HostExec && !hasContainerRuntime {
return computerID, "host", nil
}
if node.Capabilities.HostExec && hasContainerRuntime && cfg.Computer.Image == "" {
return computerID, "host", nil
}
if !hasContainerRuntime {
return "", "", fmt.Errorf("node %q has no container runtime and no host_exec capability", computerID)
}
return computerID, "box", nil
}
return computerID, "box", nil
}
if cfg.Computer.Image == "" {
nodeID := cfg.NodeID
return nodeID, "host", nil
}
nodeID := cfg.NodeID
return nodeID, "box", nil
}
// GetComputer obtains or creates a Computer for the current request.
// An optional connector may be passed to inject OPENAI_PROXY_* env vars.
// Returns the Computer, the resolved identifier, and any error.

View file

@ -29,8 +29,8 @@ func TestBuildIdentifier_Oneshot(t *testing.T) {
func TestBuildIdentifier_Session(t *testing.T) {
cfg := &types.SandboxConfig{Lifecycle: "session"}
id := sandboxv2.BuildIdentifier(cfg, "owner1", "chat42", "ast1", "ws1", nil)
if id != "owner1-chat42" {
t.Errorf("session: got %q, want %q", id, "owner1-chat42")
if id != "owner1-ast1-chat42" {
t.Errorf("session: got %q, want %q", id, "owner1-ast1-chat42")
}
}
@ -54,8 +54,9 @@ func TestBuildIdentifier_MetadataOverride(t *testing.T) {
cfg := &types.SandboxConfig{Lifecycle: "session"}
meta := map[string]any{"computer_id": "custom-box"}
id := sandboxv2.BuildIdentifier(cfg, "owner1", "chat1", "ast1", "ws1", meta)
if id != "owner1-custom-box.ws1" {
t.Errorf("metadata override: got %q, want %q", id, "owner1-custom-box.ws1")
// computer_id is used for routing only, not for identifier generation.
if id != "owner1-ast1-chat1" {
t.Errorf("metadata override: got %q, want %q", id, "owner1-ast1-chat1")
}
}
@ -63,7 +64,7 @@ func TestBuildIdentifier_MetadataEmptyIgnored(t *testing.T) {
cfg := &types.SandboxConfig{Lifecycle: "session"}
meta := map[string]any{"computer_id": ""}
id := sandboxv2.BuildIdentifier(cfg, "owner1", "chat42", "ast1", "ws1", meta)
if id != "owner1-chat42" {
if id != "owner1-ast1-chat42" {
t.Errorf("empty metadata should fall through to session, got %q", id)
}
}

View file

@ -15,11 +15,12 @@ import (
// ExecuteRequest consolidates all parameters for ExecuteSandboxStream.
type ExecuteRequest struct {
Computer infra.Computer
Runner types.Runner
Config *types.SandboxConfig
StreamReq *types.StreamRequest
Manager *infra.Manager
Computer infra.Computer
Runner types.Runner
Config *types.SandboxConfig
StreamReq *types.StreamRequest
Manager *infra.Manager
LoadingMsgID string
}
// ExecuteSandboxStream is the V2 replacement for executeSandboxStream.
@ -106,7 +107,14 @@ func ExecuteSandboxStream(
}()
var textContent []byte
loadingClosed := false
wrappedHandler := func(chunkType message.StreamChunkType, data []byte) int {
if !loadingClosed && req.LoadingMsgID != "" {
if chunkType == message.ChunkText || chunkType == message.ChunkToolCall || chunkType == message.ChunkMessageStart {
closeLoading(ctx, req.LoadingMsgID)
loadingClosed = true
}
}
if chunkType == message.ChunkText {
textContent = append(textContent, data...)
}
@ -118,6 +126,10 @@ func ExecuteSandboxStream(
err := req.Runner.Stream(runnerCtx, req.StreamReq, wrappedHandler)
if !loadingClosed && req.LoadingMsgID != "" {
closeLoading(ctx, req.LoadingMsgID)
}
panicked = false // Normal exit reached.
if err != nil {
@ -136,3 +148,20 @@ func ExecuteSandboxStream(
}
return resp, nil
}
func closeLoading(ctx *agentContext.Context, loadingMsgID string) {
if loadingMsgID == "" || ctx == nil {
return
}
msg := &message.Message{
MessageID: loadingMsgID,
Delta: true,
DeltaAction: message.DeltaReplace,
Type: message.TypeLoading,
Props: map[string]any{
"done": true,
"message": "",
},
}
ctx.Send(msg)
}

View file

@ -31,7 +31,7 @@ func (d *dockerCore) create(ctx context.Context, opts CreateOptions, addVNCPorts
}
hostCfg := &container.HostConfig{
Binds: opts.Binds,
Binds: normalizeBinds(opts.Binds),
ExtraHosts: []string{"host.tai.internal:host-gateway"},
}
@ -282,3 +282,56 @@ func (d *dockerCore) list(ctx context.Context, opts ListOptions) ([]ContainerInf
}
return result, nil
}
// normalizeBinds converts Windows-style host paths in Docker bind-mount
// specifications to WSL2 mount paths that Docker (running in WSL2) accepts.
// e.g. "D:\volumes\ws-abc:/workspace:rw" -> "/mnt/d/volumes/ws-abc:/workspace:rw"
//
// Detection is based on the path content (drive-letter prefix), not runtime.GOOS,
// because the path may originate from a remote Tai node (Windows) while Yao
// runs on macOS/Linux.
func normalizeBinds(binds []string) []string {
if len(binds) == 0 {
return binds
}
out := make([]string, len(binds))
changed := false
for i, b := range binds {
out[i] = normalizeWindowsBind(b)
if out[i] != b {
changed = true
}
}
if !changed {
return binds
}
return out
}
// normalizeWindowsBind handles a single bind spec "hostPath:containerPath[:mode]".
// When Yao runs on Windows and Docker runs in WSL2, Windows paths like
// "D:\volumes\ws-abc" must be converted to "/mnt/d/volumes/ws-abc" because
// WSL2 mounts Windows drives under /mnt/<lowercase-letter>/.
func normalizeWindowsBind(bind string) string {
if len(bind) < 3 {
return bind
}
// Detect drive-letter prefix: "X:\" or "X:/"
if bind[1] != ':' || (bind[2] != '\\' && bind[2] != '/') {
return bind
}
// Find the next colon after the drive letter colon (the bind separator)
idx := strings.Index(bind[2:], ":")
if idx < 0 {
return bind
}
hostPath := bind[:2+idx]
rest := bind[2+idx:] // starts with ":"
// Convert "D:\foo\bar" -> "/mnt/d/foo/bar"
drive := strings.ToLower(string(hostPath[0]))
tail := strings.ReplaceAll(hostPath[2:], `\`, `/`)
return "/mnt/" + drive + tail + rest
}