diff --git a/.gitignore b/.gitignore index 24fe2d1c..6f3a811a 100644 --- a/.gitignore +++ b/.gitignore @@ -74,4 +74,5 @@ tg-login tg-send registry/data/ registry/manager/DESIGN*.md -tai/testdata/ \ No newline at end of file +tai/testdata/ +agent/sandbox/docs/*.md \ No newline at end of file diff --git a/Makefile b/Makefile index 7844d513..88a19449 100644 --- a/Makefile +++ b/Makefile @@ -10,11 +10,11 @@ NOW := $(shell date +"%FT%T%z") OS := $(shell uname) # ROOT_DIR := $(shell dirname $(realpath $(firstword $(MAKEFILE_LIST)))) -TESTFOLDER := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|registry' | awk '!/\/tests\// || /openapi\/tests/') +TESTFOLDER := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|registry|agent/sandbox/v2' | awk '!/\/tests\// || /openapi\/tests/') # Core tests (exclude AI-related: agent, aigc, openai, KB, sandbox, registry, grpc, and integrations which require external services) TESTFOLDER_CORE := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|agent|kb|sandbox|integrations|registry|tai|grpc' | awk '!/\/tests\// || /openapi\/tests/') -# Agent tests (agent, aigc) - exclude agent/search/handlers/web (requires external API keys) and robot packages (tested in robot job) -TESTFOLDER_AGENT := $(shell $(GO) list ./agent/... ./aigc/... | grep -vE 'agent/search/handlers/web|agent/robot/') +# Agent tests (agent, aigc) - exclude agent/search/handlers/web (requires external API keys), robot packages (tested in robot job), and agent/sandbox/v2 (WIP, has its own job) +TESTFOLDER_AGENT := $(shell $(GO) list ./agent/... ./aigc/... | grep -vE 'agent/search/handlers/web|agent/robot/|agent/sandbox/v2') # KB tests (kb) TESTFOLDER_KB := $(shell $(GO) list ./kb/...) # Robot tests (agent/robot/... packages, excluding events/integrations which require Telegram etc.) diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index fca12ab8..9e49226b 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -14,6 +14,8 @@ import ( "github.com/yaoapp/yao/agent/llm" "github.com/yaoapp/yao/agent/output/message" agentsandbox "github.com/yaoapp/yao/agent/sandbox" + sandboxTypes "github.com/yaoapp/yao/agent/sandbox/v2/types" + infraV2 "github.com/yaoapp/yao/sandbox/v2" ) // Stream stream the agent @@ -163,7 +165,25 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa var sandboxExecutor agentsandbox.Executor var sandboxCleanup func() var sandboxLoadingMsgID string - if ast.HasSandbox() { + + // V2 sandbox state + var v2Runner sandboxTypes.Runner + var v2Computer infraV2.Computer + var v2LoadingMsgID string + + if ast.HasSandboxV2() { + ctx.Logger.Phase("Sandbox V2") + var err error + var v2Cleanup func() + v2Runner, v2Computer, v2Cleanup, v2LoadingMsgID, err = ast.initSandboxV2(ctx, opts) + if err != nil { + ast.traceAgentFail(agentNode, err) + ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err) + return nil, err + } + sandboxCleanup = v2Cleanup + ctx.Logger.PhaseComplete("Sandbox V2") + } else if ast.HasSandbox() { ctx.Logger.Phase("Sandbox") var err error sandboxExecutor, sandboxCleanup, sandboxLoadingMsgID, err = ast.initSandbox(ctx, opts) @@ -289,8 +309,17 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa // Execute the LLM streaming call // Choose between sandbox execution or direct LLM execution - if ast.HasSandbox() { - // Sandbox execution path (Claude CLI, Cursor CLI, etc.) + if ast.HasSandboxV2() && v2Runner != nil && v2Computer != nil && v2Runner.Name() != "yao" { + // V2 Sandbox execution path (non-yao runners replace LLM.Stream) + completionResponse, err = ast.executeSandboxV2Stream(ctx, completionMessages, agentNode, streamHandler, v2Runner, v2Computer, v2LoadingMsgID) + } else if ast.HasSandboxV2() && v2Runner != nil && v2Runner.Name() == "yao" { + // V2 yao runner: Prepare is done, close loading, fall through to LLM + if v2LoadingMsgID != "" { + closeLoadingV2(ctx, v2LoadingMsgID, "") + } + completionResponse, err = ast.executeLLMStream(ctx, completionMessages, completionOptions, agentNode, streamHandler, opts) + } else if ast.HasSandbox() { + // V1 Sandbox execution path (Claude CLI, Cursor CLI, etc.) completionResponse, err = ast.executeSandboxStream(ctx, completionMessages, agentNode, streamHandler, sandboxExecutor, sandboxLoadingMsgID) } else { // Direct LLM execution path diff --git a/agent/assistant/load.go b/agent/assistant/load.go index 0333ac58..75092b47 100644 --- a/agent/assistant/load.go +++ b/agent/assistant/load.go @@ -12,8 +12,10 @@ import ( "github.com/yaoapp/gou/fs" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/i18n" + sandboxTypes "github.com/yaoapp/yao/agent/sandbox/v2/types" searchTypes "github.com/yaoapp/yao/agent/search/types" store "github.com/yaoapp/yao/agent/store/types" + "github.com/yaoapp/yao/config" "gopkg.in/yaml.v3" ) @@ -378,7 +380,45 @@ func LoadPath(path string) (*Assistant, error) { return nil, err } data["locales"] = locales - return loadMap(data) + + // V2 sandbox: load standalone sandbox.yao if present (Path A). + sandboxFile := filepath.Join(path, "sandbox.yao") + if has, _ := app.Exists(sandboxFile); has { + absFile := filepath.Join(config.Conf.AppSource, sandboxFile) + sbCfg, sbErr := store.LoadSandboxConfig(absFile) + if sbErr != nil { + return nil, fmt.Errorf("load sandbox.yao: %w", sbErr) + } + data["__sandbox_v2"] = sbCfg + } + + ast, err := loadMap(data) + if err != nil { + return nil, err + } + + // If V2 sandbox was loaded via Path A, assign it now. + if sbCfg, ok := data["__sandbox_v2"].(*sandboxTypes.SandboxConfig); ok && sbCfg != nil { + ast.SandboxV2 = sbCfg + } + + // Compute config hash for V2 sandbox. + if ast.SandboxV2 != nil { + var mcpServers []store.MCPServerConfig + if ast.MCP != nil { + mcpServers = ast.MCP.Servers + } + skillsDir := "" + if ast.Path != "" { + dir := filepath.Join(config.Conf.AppSource, ast.Path, "skills") + if info, e := os.Stat(dir); e == nil && info.IsDir() { + skillsDir = dir + } + } + ast.ConfigHash = store.ComputeConfigHash(ast.SandboxV2, mcpServers, skillsDir) + } + + return ast, nil } func loadMap(data map[string]interface{}) (*Assistant, error) { @@ -721,12 +761,25 @@ func loadMap(data map[string]interface{}) (*Assistant, error) { } // sandbox (for coding agents like Claude CLI, Cursor CLI) - if sandbox, has := data["sandbox"]; has { - sb, err := store.ToSandbox(sandbox) - if err != nil { - return nil, err + // V2 sandbox via independent sandbox.yao is loaded in LoadPath (below). + // This block handles the package.yao embedded "sandbox" field with version dispatch. + if assistant.SandboxV2 == nil { + if sandbox, has := data["sandbox"]; has { + version := extractSandboxVersion(sandbox) + if version == sandboxTypes.SandboxVersionV2 { + sb, err := store.ToSandboxV2(sandbox) + if err != nil { + return nil, err + } + assistant.SandboxV2 = sb + } else { + sb, err := store.ToSandbox(sandbox) + if err != nil { + return nil, err + } + assistant.Sandbox = sb + } } - assistant.Sandbox = sb } // dependencies (name -> version constraint, like npm dependencies) @@ -1036,3 +1089,13 @@ func mergeSearchConfig(base, override *searchTypes.Config) *searchTypes.Config { return &result } + +// extractSandboxVersion tries to read the "version" field from a sandbox config value. +func extractSandboxVersion(v any) string { + if m, ok := v.(map[string]any); ok { + if ver, ok := m["version"].(string); ok { + return ver + } + } + return "" +} diff --git a/agent/assistant/load_test.go b/agent/assistant/load_test.go index a15406dd..818ee334 100644 --- a/agent/assistant/load_test.go +++ b/agent/assistant/load_test.go @@ -594,6 +594,162 @@ func TestLoadSystemAgents(t *testing.T) { }) } +// TestLoadPathSandboxV2 tests loading assistants with V2 sandbox configuration (standalone sandbox.yao) +func TestLoadPathSandboxV2(t *testing.T) { + prepare(t) + defer test.Clean() + + t.Run("OneshotCLI", func(t *testing.T) { + ast, err := assistant.LoadPath("/assistants/tests/sandbox-v2/oneshot-cli") + require.NoError(t, err) + require.NotNil(t, ast) + + assert.Equal(t, "Sandbox V2 Oneshot CLI", ast.Name) + assert.Contains(t, ast.Tags, "SandboxV2") + + // V2 sandbox should be loaded from sandbox.yao + require.NotNil(t, ast.SandboxV2, "SandboxV2 should be loaded") + assert.Equal(t, "2.0", ast.SandboxV2.Version) + assert.Equal(t, "yaoapp/tai-sandbox-claude:latest", ast.SandboxV2.Computer.Image) + assert.Equal(t, "2GB", ast.SandboxV2.Computer.Memory) + assert.Equal(t, float64(2), ast.SandboxV2.Computer.CPUs) + assert.Equal(t, "/workspace", ast.SandboxV2.Computer.WorkDir) + assert.Equal(t, "claude", ast.SandboxV2.Runner.Name) + assert.Equal(t, "cli", ast.SandboxV2.Runner.Mode) + assert.Equal(t, "oneshot", ast.SandboxV2.Lifecycle) + + // Runner options + assert.NotNil(t, ast.SandboxV2.Runner.Options) + assert.Equal(t, float64(5), ast.SandboxV2.Runner.Options["max_turns"]) + + // V1 Sandbox should be nil + assert.Nil(t, ast.Sandbox, "V1 Sandbox should be nil when V2 is present") + + // ConfigHash should be computed + assert.NotEmpty(t, ast.ConfigHash, "ConfigHash should be computed for V2 sandbox") + + // HasSandboxV2 helper + assert.True(t, ast.HasSandboxV2()) + }) + + t.Run("SessionCLI", func(t *testing.T) { + ast, err := assistant.LoadPath("/assistants/tests/sandbox-v2/session-cli") + require.NoError(t, err) + require.NotNil(t, ast) + + require.NotNil(t, ast.SandboxV2) + assert.Equal(t, "session", ast.SandboxV2.Lifecycle) + assert.Equal(t, "10m", ast.SandboxV2.IdleTimeout) + + // Prepare steps + require.Len(t, ast.SandboxV2.Prepare, 1) + assert.Equal(t, "exec", ast.SandboxV2.Prepare[0].Action) + assert.True(t, ast.SandboxV2.Prepare[0].Once) + }) + + t.Run("LongrunningCLI", func(t *testing.T) { + ast, err := assistant.LoadPath("/assistants/tests/sandbox-v2/longrunning-cli") + require.NoError(t, err) + require.NotNil(t, ast) + + require.NotNil(t, ast.SandboxV2) + assert.Equal(t, "longrunning", ast.SandboxV2.Lifecycle) + assert.Equal(t, "15m", ast.SandboxV2.IdleTimeout) + assert.Equal(t, "2h", ast.SandboxV2.MaxLifetime) + assert.Equal(t, "5s", ast.SandboxV2.StopTimeout) + assert.Equal(t, "4GB", ast.SandboxV2.Computer.Memory) + assert.Equal(t, "rw", ast.SandboxV2.Computer.MountMode) + + // Environment + assert.Equal(t, "test", ast.SandboxV2.Environment["NODE_ENV"]) + assert.Equal(t, "longrunning", ast.SandboxV2.Environment["V2_TEST_MODE"]) + + // Secrets + assert.Equal(t, "sandbox-v2-longrunning-secret", ast.SandboxV2.Secrets["TEST_SECRET"]) + + // Prepare steps + require.Len(t, ast.SandboxV2.Prepare, 3) + assert.True(t, ast.SandboxV2.Prepare[2].IgnoreError) + + // MCP (from package.yao) + require.NotNil(t, ast.MCP) + require.Len(t, ast.MCP.Servers, 1) + assert.Equal(t, "echo", ast.MCP.Servers[0].ServerID) + + // ConfigHash should include MCP servers + hashWithMCP := ast.ConfigHash + assert.NotEmpty(t, hashWithMCP) + }) + + t.Run("HooksOnly_YaoRunner", func(t *testing.T) { + ast, err := assistant.LoadPath("/assistants/tests/sandbox-v2/hooks-only") + require.NoError(t, err) + require.NotNil(t, ast) + + require.NotNil(t, ast.SandboxV2) + assert.Equal(t, "yao", ast.SandboxV2.Runner.Name) + assert.Equal(t, "oneshot", ast.SandboxV2.Lifecycle) + assert.Equal(t, float64(1), ast.SandboxV2.Computer.CPUs) + + // Runner mode should be empty (yao runner ignores mode) + assert.Empty(t, ast.SandboxV2.Runner.Mode) + }) + + t.Run("FullPrepare", func(t *testing.T) { + ast, err := assistant.LoadPath("/assistants/tests/sandbox-v2/full-prepare") + require.NoError(t, err) + require.NotNil(t, ast) + + require.NotNil(t, ast.SandboxV2) + assert.Equal(t, "session", ast.SandboxV2.Lifecycle) + assert.Equal(t, "15m", ast.SandboxV2.IdleTimeout) + + // Prepare: 5 steps with mixed actions + require.Len(t, ast.SandboxV2.Prepare, 5) + assert.Equal(t, "copy", ast.SandboxV2.Prepare[0].Action) + assert.Equal(t, "skills", ast.SandboxV2.Prepare[0].Src) + assert.Equal(t, "~/.claude/skills", ast.SandboxV2.Prepare[0].Dst) + assert.Equal(t, "exec", ast.SandboxV2.Prepare[1].Action) + assert.True(t, ast.SandboxV2.Prepare[1].Once) + assert.True(t, ast.SandboxV2.Prepare[3].IgnoreError) + + // Environment + Secrets + assert.Equal(t, "full", ast.SandboxV2.Environment["V2_PREPARE_TEST"]) + assert.Equal(t, "v2-full-prepare-key", ast.SandboxV2.Secrets["TEST_API_KEY"]) + + // Runner options + assert.Equal(t, "acceptEdits", ast.SandboxV2.Runner.Options["permission_mode"]) + }) + + t.Run("HostMode", func(t *testing.T) { + ast, err := assistant.LoadPath("/assistants/tests/sandbox-v2/host-mode") + require.NoError(t, err) + require.NotNil(t, ast) + + require.NotNil(t, ast.SandboxV2) + // Host mode: no image + assert.Empty(t, ast.SandboxV2.Computer.Image) + assert.Equal(t, "/tmp/yao-sandbox-v2-host-test", ast.SandboxV2.Computer.WorkDir) + assert.Equal(t, "session", ast.SandboxV2.Lifecycle) + }) + + t.Run("ConfigHashDeterministic", func(t *testing.T) { + ast1, err := assistant.LoadPath("/assistants/tests/sandbox-v2/oneshot-cli") + require.NoError(t, err) + ast2, err := assistant.LoadPath("/assistants/tests/sandbox-v2/oneshot-cli") + require.NoError(t, err) + assert.Equal(t, ast1.ConfigHash, ast2.ConfigHash, "same config should produce same hash") + }) + + t.Run("ConfigHashDiffers", func(t *testing.T) { + ast1, err := assistant.LoadPath("/assistants/tests/sandbox-v2/oneshot-cli") + require.NoError(t, err) + ast2, err := assistant.LoadPath("/assistants/tests/sandbox-v2/longrunning-cli") + require.NoError(t, err) + assert.NotEqual(t, ast1.ConfigHash, ast2.ConfigHash, "different configs should produce different hashes") + }) +} + // TestValidate tests the assistant Validate method func TestValidate(t *testing.T) { tests := []struct { diff --git a/agent/assistant/sandbox_v2.go b/agent/assistant/sandbox_v2.go new file mode 100644 index 00000000..a6e2b97e --- /dev/null +++ b/agent/assistant/sandbox_v2.go @@ -0,0 +1,189 @@ +package assistant + +import ( + "fmt" + "os" + "path/filepath" + + "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/i18n" + "github.com/yaoapp/yao/agent/output/message" + sandboxv2 "github.com/yaoapp/yao/agent/sandbox/v2" + sandboxTypes "github.com/yaoapp/yao/agent/sandbox/v2/types" + "github.com/yaoapp/yao/config" + infraV2 "github.com/yaoapp/yao/sandbox/v2" + traceTypes "github.com/yaoapp/yao/trace/types" +) + +// HasSandboxV2 returns true if the assistant has a V2 sandbox configuration. +func (ast *Assistant) HasSandboxV2() bool { + return ast.SandboxV2 != nil +} + +// 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. +func (ast *Assistant) initSandboxV2(ctx *context.Context, opts *context.Options) ( + sandboxTypes.Runner, infraV2.Computer, func(), string, error, +) { + cfg := ast.SandboxV2 + manager := infraV2.M() + + loadingMsg := &message.Message{ + Type: message.TypeLoading, + Props: map[string]any{ + "message": i18n.T(ctx.Locale, "sandbox.preparing"), + }, + } + loadingMsgID, _ := ctx.SendStream(loadingMsg) + + stdCtx := ctx.Context + + // 1. Obtain Computer. + 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) + } + _ = identifier + + // 2. Get Runner. + runner, err := sandboxv2.Get(cfg.Runner.Name) + 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) + } + + // 3. Resolve connector. + conn, _, err := ast.GetConnector(ctx, opts) + if err != nil && cfg.Runner.Name != "yao" { + sandboxv2.LifecycleAction(stdCtx, cfg, computer, manager) + closeLoadingV2(ctx, loadingMsgID, "sandbox.failed") + return nil, nil, nil, "", fmt.Errorf("get connector: %w", err) + } + + // 4. Resolve skills directory. + skillsDir := "" + if ast.Path != "" { + dir := filepath.Join(config.Conf.AppSource, ast.Path, "skills") + if info, e := os.Stat(dir); e == nil && info.IsDir() { + skillsDir = dir + } + } + + // 5. Convert MCP servers. + var mcpServers []sandboxTypes.MCPServer + if ast.MCP != nil { + for _, s := range ast.MCP.Servers { + mcpServers = append(mcpServers, sandboxTypes.MCPServer{ + ServerID: s.ServerID, + Resources: s.Resources, + Tools: s.Tools, + }) + } + } + + // 6. Runner.Prepare (standard context). + err = runner.Prepare(stdCtx, &sandboxTypes.PrepareRequest{ + Computer: computer, + Config: cfg, + Connector: conn, + SkillsDir: skillsDir, + MCPServers: mcpServers, + ConfigHash: ast.ConfigHash, + RunSteps: sandboxv2.RunPrepareSteps, + }) + if err != nil { + 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) + } + + // Inject computer + workspace into context so Create/Next hooks + // can access ctx.computer and ctx.workspace. + ctx.SetComputer(computer) + + cleanup := func() { + // Defensive fallback — executeSandboxV2Stream defer handles the + // normal case; this covers paths that never reach execution. + } + + return runner, computer, cleanup, loadingMsgID, nil +} + +// executeSandboxV2Stream calls the V2 Runner.Stream and wraps it in the +// standard completion response. +func (ast *Assistant) executeSandboxV2Stream( + ctx *context.Context, + completionMessages []context.Message, + agentNode traceTypes.Node, + streamHandler message.StreamFunc, + runner sandboxTypes.Runner, + computer infraV2.Computer, + loadingMsgID string, +) (*context.CompletionResponse, error) { + _ = agentNode + + 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 { + for _, p := range ast.Prompts { + if p.Role == "system" && p.Content != "" { + systemPrompt = p.Content + break + } + } + } + + // Resolve connector for Stream. + conn, _, _ := ast.GetConnector(ctx) + + streamReq := &sandboxTypes.StreamRequest{ + Computer: computer, + Config: cfg, + Connector: conn, + Messages: completionMessages, + SystemPrompt: systemPrompt, + ChatID: ctx.ChatID, + } + + execReq := &sandboxv2.ExecuteRequest{ + Computer: computer, + Runner: runner, + Config: cfg, + StreamReq: streamReq, + Manager: manager, + } + + return sandboxv2.ExecuteSandboxStream(ctx, execReq, streamHandler) +} + +func closeLoadingV2(ctx *context.Context, loadingMsgID, msgKey string) { + if loadingMsgID == "" || ctx == nil { + return + } + props := map[string]any{"done": true} + if msgKey != "" { + props["message"] = i18n.T(ctx.Locale, msgKey) + } else { + props["message"] = "" + } + doneMsg := &message.Message{ + MessageID: loadingMsgID, + Delta: true, + DeltaAction: message.DeltaReplace, + Type: message.TypeLoading, + Props: props, + } + ctx.Send(doneMsg) +} diff --git a/agent/context/jsapi.go b/agent/context/jsapi.go index 1081ba7d..9aa767fc 100644 --- a/agent/context/jsapi.go +++ b/agent/context/jsapi.go @@ -152,7 +152,7 @@ func (ctx *Context) NewObject(v8ctx *v8go.Context) (*v8go.Value, error) { memoryObj.Release() } - // Sandbox object - only set if sandbox executor is available + // Sandbox object - only set if sandbox executor is available (V1) if ctx.sandboxExecutor != nil { sandboxObj := ctx.createSandboxInstance(v8ctx) if sandboxObj != nil { @@ -161,6 +161,24 @@ func (ctx *Context) NewObject(v8ctx *v8go.Context) (*v8go.Value, error) { } } + // Computer object - only set if V2 computer is available + if ctx.computer != nil { + computerObj := ctx.createComputerInstance(v8ctx) + if computerObj != nil { + obj.Set("computer", computerObj) + computerObj.Release() + } + } + + // Workspace object - only set if V2 workspace is available + if ctx.workspace != nil { + wsObj := ctx.createWorkspaceInstance(v8ctx) + if wsObj != nil { + obj.Set("workspace", wsObj) + wsObj.Release() + } + } + return instance.Value, nil } diff --git a/agent/context/jsapi_computer.go b/agent/context/jsapi_computer.go new file mode 100644 index 00000000..ae3c57c1 --- /dev/null +++ b/agent/context/jsapi_computer.go @@ -0,0 +1,228 @@ +package context + +import ( + "context" + "fmt" + "strings" + + "github.com/yaoapp/gou/runtime/v8/bridge" + infraV2 "github.com/yaoapp/yao/sandbox/v2" + "github.com/yaoapp/yao/tai/workspace" + "rogchap.com/v8go" +) + +// SetComputer sets the V2 computer and its workspace for this context. +// Should be called after Runner.Prepare succeeds in initSandboxV2. +func (ctx *Context) SetComputer(computer infraV2.Computer) { + ctx.computer = computer + if computer != nil { + ctx.workspace = computer.Workplace() + } +} + +// GetComputer returns the V2 computer if available. +func (ctx *Context) GetComputer() infraV2.Computer { + return ctx.computer +} + +// GetWorkspace returns the V2 workspace FS if available. +func (ctx *Context) GetWorkspace() workspace.FS { + return ctx.workspace +} + +// HasComputer returns true if V2 computer is available. +func (ctx *Context) HasComputer() bool { + return ctx.computer != nil +} + +// createComputerInstance creates the ctx.computer JavaScript object. +func (ctx *Context) createComputerInstance(v8ctx *v8go.Context) *v8go.Value { + if ctx.computer == nil { + return nil + } + + iso := v8ctx.Isolate() + objTpl := v8go.NewObjectTemplate(iso) + + info := ctx.computer.ComputerInfo() + id := info.BoxID + if id == "" { + id = info.NodeID + } + objTpl.Set("id", id) + + objTpl.Set("Exec", ctx.computerExecMethod(iso)) + objTpl.Set("VNC", ctx.computerVNCMethod(iso)) + objTpl.Set("Proxy", ctx.computerProxyMethod(iso)) + objTpl.Set("Info", ctx.computerInfoMethod(iso)) + + instance, err := objTpl.NewInstance(v8ctx) + if err != nil { + return nil + } + return instance.Value +} + +// computerExecMethod implements ctx.computer.Exec(cmd) +// cmd can be a string or an array of strings. +// Returns: { stdout, stderr, exit_code } +func (ctx *Context) computerExecMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + args := info.Args() + + if ctx.computer == nil { + return bridge.JsException(v8ctx, "computer not available") + } + if len(args) < 1 { + return bridge.JsException(v8ctx, "Exec requires a command argument") + } + + cmd, err := parseCommandArg(v8ctx, args[0]) + if err != nil { + return bridge.JsException(v8ctx, err.Error()) + } + + result, err := ctx.computer.Exec(context.Background(), cmd) + if err != nil { + return bridge.JsException(v8ctx, "Exec failed: "+err.Error()) + } + + res := map[string]interface{}{ + "stdout": result.Stdout, + "stderr": result.Stderr, + "exit_code": int32(result.ExitCode), + } + jsVal, err := bridge.JsValue(v8ctx, res) + if err != nil { + return bridge.JsException(v8ctx, err.Error()) + } + return jsVal + }) +} + +// computerVNCMethod implements ctx.computer.VNC() +// Returns the VNC URL string. +func (ctx *Context) computerVNCMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + + if ctx.computer == nil { + return bridge.JsException(v8ctx, "computer not available") + } + + url, err := ctx.computer.VNC(context.Background()) + if err != nil { + return bridge.JsException(v8ctx, "VNC failed: "+err.Error()) + } + + jsVal, err := v8go.NewValue(iso, url) + if err != nil { + return bridge.JsException(v8ctx, err.Error()) + } + return jsVal + }) +} + +// computerProxyMethod implements ctx.computer.Proxy(port, path?) +// Returns the proxy URL string. +func (ctx *Context) computerProxyMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + args := info.Args() + + if ctx.computer == nil { + return bridge.JsException(v8ctx, "computer not available") + } + if len(args) < 1 || !args[0].IsNumber() { + return bridge.JsException(v8ctx, "Proxy requires a port number") + } + + port := int(args[0].Integer()) + path := "" + if len(args) >= 2 && args[1].IsString() { + path = args[1].String() + } + + url, err := ctx.computer.Proxy(context.Background(), port, path) + if err != nil { + return bridge.JsException(v8ctx, "Proxy failed: "+err.Error()) + } + + jsVal, err := v8go.NewValue(iso, url) + if err != nil { + return bridge.JsException(v8ctx, err.Error()) + } + return jsVal + }) +} + +// computerInfoMethod implements ctx.computer.Info() +// Returns a JS object with computer identity and system information. +func (ctx *Context) computerInfoMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + + if ctx.computer == nil { + return bridge.JsException(v8ctx, "computer not available") + } + + ci := ctx.computer.ComputerInfo() + result := map[string]interface{}{ + "kind": ci.Kind, + "node_id": ci.NodeID, + "tai_id": ci.TaiID, + "status": ci.Status, + "system": map[string]interface{}{ + "os": ci.System.OS, + "arch": ci.System.Arch, + "hostname": ci.System.Hostname, + "num_cpu": int32(ci.System.NumCPU), + "shell": ci.System.Shell, + }, + } + if ci.BoxID != "" { + result["box_id"] = ci.BoxID + result["container_id"] = ci.ContainerID + result["image"] = ci.Image + result["policy"] = string(ci.Policy) + } + + jsVal, err := bridge.JsValue(v8ctx, result) + if err != nil { + return bridge.JsException(v8ctx, err.Error()) + } + return jsVal + }) +} + +// parseCommandArg converts a JS value (string or string array) to []string. +func parseCommandArg(v8ctx *v8go.Context, val *v8go.Value) ([]string, error) { + if val.IsString() { + raw := val.String() + return strings.Fields(raw), nil + } + + if val.IsArray() { + obj, err := val.AsObject() + if err != nil { + return nil, err + } + lengthVal, err := obj.Get("length") + if err != nil { + return nil, err + } + length := int(lengthVal.Integer()) + cmd := make([]string, length) + for i := 0; i < length; i++ { + item, err := obj.GetIdx(uint32(i)) + if err != nil { + return nil, err + } + cmd[i] = item.String() + } + return cmd, nil + } + + return nil, fmt.Errorf("command must be a string or array of strings") +} diff --git a/agent/context/jsapi_workspace.go b/agent/context/jsapi_workspace.go new file mode 100644 index 00000000..d02c3168 --- /dev/null +++ b/agent/context/jsapi_workspace.go @@ -0,0 +1,291 @@ +package context + +import ( + "io/fs" + "os" + + "github.com/yaoapp/gou/runtime/v8/bridge" + "rogchap.com/v8go" +) + +// createWorkspaceInstance creates the ctx.workspace JavaScript object. +func (ctx *Context) createWorkspaceInstance(v8ctx *v8go.Context) *v8go.Value { + if ctx.workspace == nil { + return nil + } + + iso := v8ctx.Isolate() + objTpl := v8go.NewObjectTemplate(iso) + + objTpl.Set("ReadFile", ctx.wsReadFileMethod(iso)) + objTpl.Set("WriteFile", ctx.wsWriteFileMethod(iso)) + objTpl.Set("ReadDir", ctx.wsReadDirMethod(iso)) + objTpl.Set("MkdirAll", ctx.wsMkdirAllMethod(iso)) + objTpl.Set("Remove", ctx.wsRemoveMethod(iso)) + objTpl.Set("RemoveAll", ctx.wsRemoveAllMethod(iso)) + objTpl.Set("Rename", ctx.wsRenameMethod(iso)) + objTpl.Set("Copy", ctx.wsCopyMethod(iso)) + objTpl.Set("Stat", ctx.wsStatMethod(iso)) + objTpl.Set("Exists", ctx.wsExistsMethod(iso)) + + instance, err := objTpl.NewInstance(v8ctx) + if err != nil { + return nil + } + return instance.Value +} + +// wsReadFileMethod implements ctx.workspace.ReadFile(path) +func (ctx *Context) wsReadFileMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + args := info.Args() + + if ctx.workspace == nil { + return bridge.JsException(v8ctx, "workspace not available") + } + if len(args) < 1 { + return bridge.JsException(v8ctx, "ReadFile requires a path argument") + } + + data, err := ctx.workspace.ReadFile(args[0].String()) + if err != nil { + return bridge.JsException(v8ctx, "ReadFile failed: "+err.Error()) + } + + jsVal, err := v8go.NewValue(iso, string(data)) + if err != nil { + return bridge.JsException(v8ctx, err.Error()) + } + return jsVal + }) +} + +// wsWriteFileMethod implements ctx.workspace.WriteFile(path, content) +func (ctx *Context) wsWriteFileMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + args := info.Args() + + if ctx.workspace == nil { + return bridge.JsException(v8ctx, "workspace not available") + } + if len(args) < 2 { + return bridge.JsException(v8ctx, "WriteFile requires path and content arguments") + } + + path := args[0].String() + content := args[1].String() + + if err := ctx.workspace.WriteFile(path, []byte(content), 0o644); err != nil { + return bridge.JsException(v8ctx, "WriteFile failed: "+err.Error()) + } + return v8go.Undefined(iso) + }) +} + +// wsReadDirMethod implements ctx.workspace.ReadDir(path) +func (ctx *Context) wsReadDirMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + args := info.Args() + + if ctx.workspace == nil { + return bridge.JsException(v8ctx, "workspace not available") + } + + path := "." + if len(args) >= 1 && args[0].IsString() { + path = args[0].String() + } + + entries, err := ctx.workspace.ReadDir(path) + if err != nil { + return bridge.JsException(v8ctx, "ReadDir failed: "+err.Error()) + } + + result := make([]map[string]interface{}, 0, len(entries)) + for _, e := range entries { + fi, _ := e.Info() + item := map[string]interface{}{ + "name": e.Name(), + "is_dir": e.IsDir(), + } + if fi != nil { + item["size"] = int32(fi.Size()) + } + result = append(result, item) + } + + jsVal, err := bridge.JsValue(v8ctx, result) + if err != nil { + return bridge.JsException(v8ctx, err.Error()) + } + return jsVal + }) +} + +// wsMkdirAllMethod implements ctx.workspace.MkdirAll(path) +func (ctx *Context) wsMkdirAllMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + args := info.Args() + + if ctx.workspace == nil { + return bridge.JsException(v8ctx, "workspace not available") + } + if len(args) < 1 { + return bridge.JsException(v8ctx, "MkdirAll requires a path argument") + } + + if err := ctx.workspace.MkdirAll(args[0].String(), 0o755); err != nil { + return bridge.JsException(v8ctx, "MkdirAll failed: "+err.Error()) + } + return v8go.Undefined(iso) + }) +} + +// wsRemoveMethod implements ctx.workspace.Remove(path) +func (ctx *Context) wsRemoveMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + args := info.Args() + + if ctx.workspace == nil { + return bridge.JsException(v8ctx, "workspace not available") + } + if len(args) < 1 { + return bridge.JsException(v8ctx, "Remove requires a path argument") + } + + if err := ctx.workspace.Remove(args[0].String()); err != nil { + return bridge.JsException(v8ctx, "Remove failed: "+err.Error()) + } + return v8go.Undefined(iso) + }) +} + +// wsRemoveAllMethod implements ctx.workspace.RemoveAll(path) +func (ctx *Context) wsRemoveAllMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + args := info.Args() + + if ctx.workspace == nil { + return bridge.JsException(v8ctx, "workspace not available") + } + if len(args) < 1 { + return bridge.JsException(v8ctx, "RemoveAll requires a path argument") + } + + if err := ctx.workspace.RemoveAll(args[0].String()); err != nil { + return bridge.JsException(v8ctx, "RemoveAll failed: "+err.Error()) + } + return v8go.Undefined(iso) + }) +} + +// wsRenameMethod implements ctx.workspace.Rename(oldName, newName) +func (ctx *Context) wsRenameMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + args := info.Args() + + if ctx.workspace == nil { + return bridge.JsException(v8ctx, "workspace not available") + } + if len(args) < 2 { + return bridge.JsException(v8ctx, "Rename requires oldName and newName arguments") + } + + if err := ctx.workspace.Rename(args[0].String(), args[1].String()); err != nil { + return bridge.JsException(v8ctx, "Rename failed: "+err.Error()) + } + return v8go.Undefined(iso) + }) +} + +// wsCopyMethod implements ctx.workspace.Copy(src, dst) +func (ctx *Context) wsCopyMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + args := info.Args() + + if ctx.workspace == nil { + return bridge.JsException(v8ctx, "workspace not available") + } + if len(args) < 2 { + return bridge.JsException(v8ctx, "Copy requires src and dst arguments") + } + + if _, err := ctx.workspace.Copy(args[0].String(), args[1].String()); err != nil { + return bridge.JsException(v8ctx, "Copy failed: "+err.Error()) + } + return v8go.Undefined(iso) + }) +} + +// wsStatMethod implements ctx.workspace.Stat(path) +func (ctx *Context) wsStatMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + args := info.Args() + + if ctx.workspace == nil { + return bridge.JsException(v8ctx, "workspace not available") + } + if len(args) < 1 { + return bridge.JsException(v8ctx, "Stat requires a path argument") + } + + fi, err := ctx.workspace.Stat(args[0].String()) + if err != nil { + return bridge.JsException(v8ctx, "Stat failed: "+err.Error()) + } + + result := map[string]interface{}{ + "name": fi.Name(), + "size": int32(fi.Size()), + "is_dir": fi.IsDir(), + "mode": int32(fi.Mode()), + "mtime": fi.ModTime().UnixMilli(), + } + jsVal, err := bridge.JsValue(v8ctx, result) + if err != nil { + return bridge.JsException(v8ctx, err.Error()) + } + return jsVal + }) +} + +// wsExistsMethod implements ctx.workspace.Exists(path) +func (ctx *Context) wsExistsMethod(iso *v8go.Isolate) *v8go.FunctionTemplate { + return v8go.NewFunctionTemplate(iso, func(info *v8go.FunctionCallbackInfo) *v8go.Value { + v8ctx := info.Context() + args := info.Args() + + if ctx.workspace == nil { + return bridge.JsException(v8ctx, "workspace not available") + } + if len(args) < 1 { + return bridge.JsException(v8ctx, "Exists requires a path argument") + } + + _, err := ctx.workspace.Stat(args[0].String()) + exists := err == nil || !isNotExist(err) + + jsVal, _ := v8go.NewValue(iso, exists) + return jsVal + }) +} + +func isNotExist(err error) bool { + if os.IsNotExist(err) { + return true + } + pathErr, ok := err.(*fs.PathError) + if ok && os.IsNotExist(pathErr.Err) { + return true + } + return false +} diff --git a/agent/context/types.go b/agent/context/types.go index 5f5f39ac..10f3aa40 100644 --- a/agent/context/types.go +++ b/agent/context/types.go @@ -11,6 +11,8 @@ import ( "github.com/yaoapp/yao/agent/output" "github.com/yaoapp/yao/agent/output/message" "github.com/yaoapp/yao/openapi/oauth/types" + infraV2 "github.com/yaoapp/yao/sandbox/v2" + "github.com/yaoapp/yao/tai/workspace" traceTypes "github.com/yaoapp/yao/trace/types" ) @@ -251,6 +253,8 @@ type Context struct { trace traceTypes.Manager `json:"-"` // Trace manager, lazy initialized on first access messageMetadata *messageMetadataStore `json:"-"` // Thread-safe message metadata store for delta operations sandboxExecutor SandboxExecutor `json:"-"` // Sandbox executor for hooks (set by assistant when sandbox is configured) + computer infraV2.Computer `json:"-"` // V2 sandbox computer (set by assistant when V2 sandbox is configured) + workspace workspace.FS `json:"-"` // V2 workspace FS (derived from computer.Workplace()) // Model capabilities (set by assistant, used by output adapters) Capabilities *llm.Capabilities `json:"-"` // Model capabilities for the current connector diff --git a/agent/sandbox/v2/claude/attachments.go b/agent/sandbox/v2/claude/attachments.go new file mode 100644 index 00000000..5e38438e --- /dev/null +++ b/agent/sandbox/v2/claude/attachments.go @@ -0,0 +1,225 @@ +package claude + +import ( + "context" + "fmt" + "path/filepath" + "strings" + + agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/attachment" + workspace "github.com/yaoapp/yao/tai/workspace" +) + +// prepareAttachments resolves __yao.attachment:// URLs in messages, +// copies actual files into the workspace .attachments/{chatID}/ directory via ws.Copy, +// and replaces multimodal content parts with text references. +func prepareAttachments(ctx context.Context, messages []agentContext.Message, chatID string, ws workspace.FS) ([]agentContext.Message, error) { + usedNames := make(map[string]int) + attachDir := ".attachments/" + chatID + + result := make([]agentContext.Message, len(messages)) + copy(result, messages) + + for i, msg := range result { + if msg.Role != "user" { + continue + } + + parts, ok := msg.Content.([]interface{}) + if !ok { + if typedParts, ok := msg.Content.([]agentContext.ContentPart); ok { + iparts := make([]interface{}, len(typedParts)) + for j, p := range typedParts { + m := map[string]interface{}{"type": string(p.Type)} + if p.Text != "" { + m["text"] = p.Text + } + if p.ImageURL != nil { + m["image_url"] = map[string]interface{}{ + "url": p.ImageURL.URL, + "detail": string(p.ImageURL.Detail), + } + } + if p.File != nil { + m["file"] = map[string]interface{}{ + "url": p.File.URL, + "filename": p.File.Filename, + } + } + iparts[j] = m + } + parts = iparts + } else { + continue + } + } + + if len(parts) == 0 { + continue + } + + var textParts []string + + for _, item := range parts { + m, ok := item.(map[string]interface{}) + if !ok { + continue + } + partType, _ := m["type"].(string) + + switch partType { + case "text": + if text, ok := m["text"].(string); ok && text != "" { + textParts = append(textParts, text) + } + + case "image_url": + imgData, _ := m["image_url"].(map[string]interface{}) + if imgData == nil { + continue + } + url, _ := imgData["url"].(string) + if url == "" { + continue + } + uploaderName, fileID, isWrapper := attachment.Parse(url) + if !isWrapper { + textParts = append(textParts, fmt.Sprintf("[Image: %s]", url)) + continue + } + ref, err := resolveAttachment(ctx, uploaderName, fileID, "", attachDir, usedNames, ws) + if err != nil { + textParts = append(textParts, "[Attached image: failed to load]") + continue + } + textParts = append(textParts, ref) + + case "file": + fileData, _ := m["file"].(map[string]interface{}) + if fileData == nil { + continue + } + url, _ := fileData["url"].(string) + hintName, _ := fileData["filename"].(string) + if url == "" { + continue + } + uploaderName, fileID, isWrapper := attachment.Parse(url) + if !isWrapper { + textParts = append(textParts, fmt.Sprintf("[File: %s]", url)) + continue + } + ref, err := resolveAttachment(ctx, uploaderName, fileID, hintName, attachDir, usedNames, ws) + if err != nil { + textParts = append(textParts, "[Attached file: failed to load]") + continue + } + textParts = append(textParts, ref) + } + } + + if len(textParts) > 0 { + newMsg := result[i] + newMsg.Content = strings.Join(textParts, "\n\n") + result[i] = newMsg + } + } + + return result, nil +} + +// resolveAttachment gets the local path of an attachment and copies it into +// the workspace via ws.Copy("local:///abs/path", ".attachments/{chatID}/filename"). +func resolveAttachment( + ctx context.Context, + uploaderName, fileID, hintName, attachDir string, + usedNames map[string]int, + ws workspace.FS, +) (string, error) { + manager, exists := attachment.Managers[uploaderName] + if !exists { + return "", fmt.Errorf("attachment manager not found: %s", uploaderName) + } + + fileInfo, err := manager.Info(ctx, fileID) + if err != nil { + return "", fmt.Errorf("failed to get file info: %w", err) + } + + absPath, _, err := manager.LocalPath(ctx, fileID) + if err != nil { + return "", fmt.Errorf("failed to get local path: %w", err) + } + + filename := fileInfo.Filename + if filename == "" && hintName != "" { + filename = hintName + } + if filename == "" { + ext := extensionFromContentType(fileInfo.ContentType) + filename = fileID + ext + } + + baseName := filename + if count, exists := usedNames[baseName]; exists { + ext := filepath.Ext(filename) + name := strings.TrimSuffix(filename, ext) + filename = fmt.Sprintf("%s_%d%s", name, count+1, ext) + usedNames[baseName] = count + 1 + } else { + usedNames[baseName] = 0 + } + + dstPath := attachDir + "/" + filename + src := "local:///" + absPath + + if _, err := ws.Copy(src, dstPath); err != nil { + return "", fmt.Errorf("failed to copy attachment to workspace: %w", err) + } + + sizeStr := formatFileSize(fileInfo.Bytes) + return fmt.Sprintf("[Attached file: %s (%s, %s)]", dstPath, fileInfo.ContentType, sizeStr), nil +} + +func extensionFromContentType(contentType string) string { + switch contentType { + case "image/png": + return ".png" + case "image/jpeg": + return ".jpg" + case "image/gif": + return ".gif" + case "image/webp": + return ".webp" + case "image/svg+xml": + return ".svg" + case "application/pdf": + return ".pdf" + case "text/plain": + return ".txt" + case "text/html": + return ".html" + case "text/css": + return ".css" + case "text/javascript", "application/javascript": + return ".js" + case "application/json": + return ".json" + case "application/zip": + return ".zip" + default: + return "" + } +} + +func formatFileSize(bytes int) string { + switch { + case bytes >= 1024*1024: + return fmt.Sprintf("%.1fMB", float64(bytes)/(1024*1024)) + case bytes >= 1024: + return fmt.Sprintf("%.1fKB", float64(bytes)/1024) + default: + return fmt.Sprintf("%dB", bytes) + } +} diff --git a/agent/sandbox/v2/claude/parse.go b/agent/sandbox/v2/claude/parse.go new file mode 100644 index 00000000..a960a224 --- /dev/null +++ b/agent/sandbox/v2/claude/parse.go @@ -0,0 +1,239 @@ +package claude + +import ( + "bufio" + "context" + "encoding/json" + "fmt" + "io" + "log" + "strings" + "time" + + goujson "github.com/yaoapp/gou/json" + agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/output/message" +) + +// parseStreamJSON reads stream-json lines from Claude CLI stdout and +// pushes them through handler as standard StreamChunkType events. +func parseStreamJSON(_ context.Context, stdout io.ReadCloser, handler message.StreamFunc) error { + scanner := bufio.NewScanner(stdout) + buf := make([]byte, 0, 64*1024) + scanner.Buffer(buf, 1024*1024) + + messageStarted := false + + type toolState struct { + name string + inputJSON strings.Builder + } + var currentTool *toolState + + for scanner.Scan() { + line := scanner.Text() + if line == "" { + continue + } + + var msg map[string]any + if err := json.Unmarshal([]byte(line), &msg); err != nil { + continue + } + + msgType, _ := msg["type"].(string) + stopped := false + + switch msgType { + case "system": + if handler != nil { + data, _ := json.Marshal(msg) + if handler(message.ChunkMetadata, data) != 0 { + stopped = true + } + } + + case "stream_event": + event, _ := msg["event"].(map[string]any) + if event == nil { + continue + } + eventType, _ := event["type"].(string) + + switch eventType { + case "content_block_start": + if cb, ok := event["content_block"].(map[string]any); ok { + blockType, _ := cb["type"].(string) + if blockType == "tool_use" { + toolName, _ := cb["name"].(string) + currentTool = &toolState{name: toolName} + if handler != nil { + data, _ := json.Marshal(map[string]any{"tool": toolName}) + if handler(message.ChunkToolCall, data) != 0 { + stopped = true + } + } + } + } + + case "content_block_delta": + if delta, ok := event["delta"].(map[string]any); ok { + deltaType, _ := delta["type"].(string) + switch deltaType { + case "text_delta": + if text, ok := delta["text"].(string); ok && text != "" { + if handler != nil { + if !messageStarted { + startData := message.EventMessageStartData{ + MessageID: fmt.Sprintf("sandbox-%d", time.Now().UnixNano()), + Type: "text", + Timestamp: time.Now().UnixMilli(), + } + sd, _ := json.Marshal(startData) + if handler(message.ChunkMessageStart, sd) != 0 { + stopped = true + break + } + messageStarted = true + } + if handler(message.ChunkText, []byte(text)) != 0 { + stopped = true + } + } + } + case "input_json_delta": + if currentTool != nil { + if partial, ok := delta["partial_json"].(string); ok { + currentTool.inputJSON.WriteString(partial) + } + } + } + } + + case "content_block_stop": + currentTool = nil + } + + case "assistant": + if msgData, ok := msg["message"].(map[string]any); ok { + stopReason, _ := msgData["stop_reason"].(string) + if stopReason != "" { + if contentArr, ok := msgData["content"].([]any); ok { + for _, item := range contentArr { + ci, ok := item.(map[string]any) + if !ok { + continue + } + itemType, _ := ci["type"].(string) + if itemType == "text" { + if text, ok := ci["text"].(string); ok && text != "" && handler != nil && !messageStarted { + startData := message.EventMessageStartData{ + MessageID: fmt.Sprintf("sandbox-%d", time.Now().UnixNano()), + Type: "text", + Timestamp: time.Now().UnixMilli(), + } + sd, _ := json.Marshal(startData) + if handler(message.ChunkMessageStart, sd) != 0 { + stopped = true + break + } + if handler(message.ChunkText, []byte(text)) != 0 { + stopped = true + break + } + messageStarted = true + } + } + } + } + } + } + + case "result": + isError, _ := msg["is_error"].(bool) + if isError { + if result, ok := msg["result"].(string); ok { + if handler != nil { + handler(message.ChunkError, []byte(result)) + } + return fmt.Errorf("Claude CLI error: %s", result) + } + } + if handler != nil && messageStarted { + handler(message.ChunkMessageEnd, nil) + } + + case "error": + var errMsg string + switch e := msg["error"].(type) { + case string: + errMsg = e + case map[string]any: + errMsg, _ = e["message"].(string) + } + if errMsg != "" { + if handler != nil { + handler(message.ChunkError, []byte(errMsg)) + } + return fmt.Errorf("Claude CLI error: %s", errMsg) + } + } + + if stopped { + break + } + } + + return scanner.Err() +} + +// buildFirstRequestJSONL builds JSONL with all messages for the first request. +func buildFirstRequestJSONL(messages []agentContext.Message) string { + var lines []string + for _, msg := range messages { + if msg.Role == "system" { + continue + } + content := msg.Content + if content == nil { + content = "" + } + streamMsg := map[string]any{ + "type": string(msg.Role), + "message": map[string]any{ + "role": string(msg.Role), + "content": content, + }, + } + data, _ := json.Marshal(streamMsg) + lines = append(lines, string(data)) + } + return strings.Join(lines, "\n") +} + +// buildLastUserMessageJSONL builds JSONL with only the last user message. +func buildLastUserMessageJSONL(messages []agentContext.Message) string { + for i := len(messages) - 1; i >= 0; i-- { + if messages[i].Role == "user" { + content := messages[i].Content + if content == nil { + content = "" + } + msg := map[string]any{ + "type": "user", + "message": map[string]any{ + "role": "user", + "content": content, + }, + } + data, _ := json.Marshal(msg) + return string(data) + } + } + return "" +} + +// Suppress unused import warnings — goujson.Parse is used for tool description +// parsing in V1 and will be used for detailed tool descriptions in future. +var _ = goujson.Parse +var _ = log.Printf diff --git a/agent/sandbox/v2/claude/runner.go b/agent/sandbox/v2/claude/runner.go new file mode 100644 index 00000000..ea9516ab --- /dev/null +++ b/agent/sandbox/v2/claude/runner.go @@ -0,0 +1,458 @@ +package claude + +import ( + "context" + "encoding/json" + "fmt" + "path" + "strings" + "time" + + "github.com/yaoapp/gou/connector" + "github.com/yaoapp/yao/agent/output/message" + "github.com/yaoapp/yao/agent/sandbox/v2/types" + infra "github.com/yaoapp/yao/sandbox/v2" +) + +const ( + defaultWorkDir = "/workspace" + defaultUser = "sandbox" + defaultUserHome = "/home/sandbox" + defaultProxyPort = 3456 +) + +// ClaudeRunner implements the Runner interface for Claude CLI (mode=cli). +type ClaudeRunner struct { + mode string + proxyReady bool + hasMCP bool + mcpToolPattern string // e.g. "mcp__yao__*,mcp__github__*" + servicePort int + servicePath string + serviceProtocol string +} + +// New creates a new ClaudeRunner. +func New() *ClaudeRunner { + return &ClaudeRunner{mode: "cli"} +} + +func (r *ClaudeRunner) Name() string { return "claude" } + +// Prepare executes user-defined and runner-specific prepare steps. +func (r *ClaudeRunner) Prepare(ctx context.Context, req *types.PrepareRequest) error { + r.mode = req.Config.Runner.Mode + if r.mode == "" { + r.mode = "cli" + } + + workDir := resolveWorkDir(req.Config) + + // Merge user-defined steps with runner-specific steps. + steps := append([]types.PrepareStep{}, req.Config.Prepare...) + + // Runner-specific: ensure .claude directory in workDir. + if req.SkillsDir != "" { + steps = append(steps, types.PrepareStep{ + Action: "exec", + Cmd: fmt.Sprintf("mkdir -p %s/.claude", workDir), + Once: true, + }) + } + + // Runner-specific: write proxy config and start proxy (for non-anthropic connectors). + if req.Connector != nil && !req.Connector.Is(connector.ANTHROPIC) { + setting := req.Connector.Setting() + host, _ := setting["host"].(string) + key, _ := setting["key"].(string) + model, _ := setting["model"].(string) + if host != "" && key != "" { + proxyJSON := buildProxyConfig(host, key, model, setting) + steps = append(steps, types.PrepareStep{ + Action: "file", + Path: ".yao/proxy.json", + Content: proxyJSON, + Once: true, + }) + steps = append(steps, types.PrepareStep{ + Action: "exec", + Cmd: "which start-claude-proxy && start-claude-proxy || true", + Once: true, + IgnoreError: true, + }) + r.proxyReady = true + } + } + + // Runner-specific: write MCP config. + if len(req.MCPServers) > 0 { + r.hasMCP = true + r.mcpToolPattern = buildMCPAllowedTools(req.MCPServers) + mcpJSON := buildMCPConfig(req.MCPServers) + steps = append(steps, types.PrepareStep{ + Action: "file", + Path: path.Join(workDir, ".mcp.json"), + Content: mcpJSON, + }) + } + + // Execute all steps via the injected callback. + if req.RunSteps != nil && len(steps) > 0 { + if err := req.RunSteps(ctx, steps, req.Computer, req.Config.ID, req.ConfigHash); err != nil { + return fmt.Errorf("claude prepare steps: %w", err) + } + } + + return nil +} + +// Stream executes the Claude CLI and streams output to handler. +func (r *ClaudeRunner) Stream(ctx context.Context, req *types.StreamRequest, handler message.StreamFunc) error { + computer := req.Computer + if computer == nil { + return fmt.Errorf("computer is nil") + } + + workDir := resolveWorkDir(req.Config) + + // Prepare attachments: resolve __yao.attachment:// URLs, copy files to workspace. + if req.ChatID != "" { + ws := computer.Workplace() + if ws != nil { + processed, err := prepareAttachments(ctx, req.Messages, req.ChatID, ws) + if err != nil { + return fmt.Errorf("prepareAttachments: %w", err) + } + req.Messages = processed + } + } + + // Detect continuation (existing .claude/projects/ directory). + isContinuation := hasExistingSession(ctx, computer, workDir) + + // Build CLI command and env. + cmd, env := r.buildCLICommand(req, isContinuation) + + // Create stream. + execStream, err := computer.Stream(ctx, cmd, infra.WithWorkDir(workDir), infra.WithEnv(env)) + if err != nil { + return fmt.Errorf("computer.Stream: %w", err) + } + + // Monitor for context cancellation — kill the process. + done := make(chan struct{}) + defer func() { + close(done) + }() + + go func() { + select { + case <-ctx.Done(): + killCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + computer.Exec(killCtx, []string{"pkill", "-f", "claude"}) + execStream.Cancel() + case <-done: + } + }() + + // Parse streaming output. + parseErr := parseStreamJSON(ctx, execStream.Stdout, handler) + + // Wait for process exit. + exitCode, waitErr := execStream.Wait() + if parseErr != nil { + return parseErr + } + if waitErr != nil { + return waitErr + } + if exitCode != 0 { + return fmt.Errorf("claude CLI exited with code %d", exitCode) + } + return nil +} + +// Cleanup kills any remaining claude processes. +// mode=service: don't kill the service daemon (lifecycle manages it), only clean proxy. +// mode=cli: kill all claude CLI processes. +func (r *ClaudeRunner) Cleanup(ctx context.Context, computer infra.Computer) error { + if computer == nil { + return nil + } + + if r.mode != "service" { + computer.Exec(ctx, []string{"sh", "-c", "pkill -f 'claude' || true"}) + } + + if r.proxyReady { + computer.Exec(ctx, []string{"sh", "-c", "pkill -f 'claude-proxy' || true"}) + } + + return nil +} + +// hasExistingSession checks if a Claude CLI session exists in the workspace. +func hasExistingSession(ctx context.Context, computer infra.Computer, workDir string) bool { + sessionDir := path.Join(workDir, ".claude/projects") + result, err := computer.Exec(ctx, []string{"ls", sessionDir}) + if err != nil || result.ExitCode != 0 { + return false + } + return strings.TrimSpace(result.Stdout) != "" +} + +// buildCLICommand constructs the Claude CLI command and environment variables. +func (r *ClaudeRunner) buildCLICommand(req *types.StreamRequest, isContinuation bool) ([]string, map[string]string) { + workDir := resolveWorkDir(req.Config) + userHome := resolveUserHome(req.Config) + + env := make(map[string]string) + env["HOME"] = workDir + + // User-specific paths (only set when running as non-root user inside container). + if userHome != "" { + env["XAUTHORITY"] = path.Join(userHome, ".Xauthority") + } + + // Connector environment. + if req.Connector != nil { + setting := req.Connector.Setting() + host, _ := setting["host"].(string) + key, _ := setting["key"].(string) + model, _ := setting["model"].(string) + + if req.Connector.Is(connector.ANTHROPIC) { + env["ANTHROPIC_BASE_URL"] = host + env["ANTHROPIC_API_KEY"] = key + } else { + env["ANTHROPIC_BASE_URL"] = fmt.Sprintf("http://127.0.0.1:%d", defaultProxyPort) + env["ANTHROPIC_API_KEY"] = "dummy" + } + + if model != "" { + env["ANTHROPIC_MODEL"] = model + env["ANTHROPIC_DEFAULT_OPUS_MODEL"] = model + env["ANTHROPIC_DEFAULT_SONNET_MODEL"] = model + env["ANTHROPIC_DEFAULT_HAIKU_MODEL"] = model + env["CLAUDE_CODE_SUBAGENT_MODEL"] = model + } + } + + // Secrets from config. + if req.Config != nil && len(req.Config.Secrets) > 0 { + for k, v := range req.Config.Secrets { + env[k] = v + } + } + + // Build system prompt. + var systemPrompt string + envPrompt := buildSandboxEnvPrompt(workDir) + if !isContinuation && req.SystemPrompt != "" { + systemPrompt = req.SystemPrompt + "\n\n" + envPrompt + } else if !isContinuation { + systemPrompt = envPrompt + } + + // Build input JSONL. + var inputJSONL string + if isContinuation { + inputJSONL = buildLastUserMessageJSONL(req.Messages) + } else { + inputJSONL = buildFirstRequestJSONL(req.Messages) + } + + // CLI args. + var args []string + args = append(args, "--dangerously-skip-permissions") + args = append(args, "--permission-mode", "bypassPermissions") + args = append(args, "--input-format", "stream-json") + args = append(args, "--output-format", "stream-json") + args = append(args, "--include-partial-messages") + args = append(args, "--verbose") + + if isContinuation { + args = append(args, "--continue") + } + + // Runner options pass-through. + if req.Config != nil && req.Config.Runner.Options != nil { + for key, val := range req.Config.Runner.Options { + if flag, ok := claudeArgWhitelist[key]; ok { + args = append(args, flag, fmt.Sprintf("%v", val)) + } + } + } + + // MCP config (set by Prepare if MCPServers were present). + if r.hasMCP { + args = append(args, "--mcp-config", path.Join(workDir, ".mcp.json")) + if r.mcpToolPattern != "" { + args = append(args, "--allowedTools", r.mcpToolPattern) + } + } + + // Build bash command with heredoc. + var bash strings.Builder + if userHome != "" { + bash.WriteString(fmt.Sprintf("touch %s/.Xauthority 2>/dev/null; ", userHome)) + } + bash.WriteString("touch \"$HOME/.Xauthority\" 2>/dev/null\n") + + if systemPrompt != "" { + promptFile := path.Join(workDir, ".yao/.system-prompt.txt") + bash.WriteString(fmt.Sprintf("mkdir -p %s/.yao\n", workDir)) + bash.WriteString(fmt.Sprintf("cat << 'PROMPTEOF' > %s\n", promptFile)) + bash.WriteString(systemPrompt) + bash.WriteString("\nPROMPTEOF\n") + args = append(args, "--append-system-prompt-file", promptFile) + } + + bash.WriteString("cat << 'INPUTEOF' | claude -p") + for _, arg := range args { + bash.WriteString(fmt.Sprintf(" %q", arg)) + } + bash.WriteString(" 2>&1\n") + bash.WriteString(inputJSONL) + bash.WriteString("\nINPUTEOF") + + return []string{"bash", "-c", bash.String()}, env +} + +// buildProxyConfig creates the claude-proxy configuration JSON. +func buildProxyConfig(host, key, model string, setting map[string]any) []byte { + backendURL := connector.BuildAPIURL(host, "/chat/completions") + config := map[string]any{ + "backend": backendURL, + "api_key": key, + "model": model, + } + opts := make(map[string]any) + for k, v := range setting { + switch k { + case "host", "key", "model", "azure", "capabilities": + continue + default: + opts[k] = v + } + } + if len(opts) > 0 { + config["options"] = opts + } + data, _ := json.MarshalIndent(config, "", " ") + return data +} + +// buildMCPConfig creates the .mcp.json for Claude CLI based on declared servers. +// Each server delegates to "tai call" which bridges stdio JSON-RPC to Yao gRPC. +// Connection is configured via env vars (YAO_GRPC_ADDR, YAO_TOKEN, etc.) +// injected by the sandbox infrastructure at container start. +func buildMCPConfig(servers []types.MCPServer) []byte { + mcpServers := make(map[string]any, len(servers)) + for _, s := range servers { + name := s.ServerID + if name == "" { + continue + } + mcpServers[name] = map[string]any{ + "command": "tai", + "args": []string{"call"}, + } + } + if len(mcpServers) == 0 { + mcpServers["yao"] = map[string]any{ + "command": "tai", + "args": []string{"call"}, + } + } + config := map[string]any{"mcpServers": mcpServers} + data, _ := json.Marshal(config) + return data +} + +// buildMCPAllowedTools generates the --allowedTools pattern from server IDs. +func buildMCPAllowedTools(servers []types.MCPServer) string { + patterns := make([]string, 0, len(servers)) + for _, s := range servers { + if s.ServerID != "" { + patterns = append(patterns, fmt.Sprintf("mcp__%s__*", s.ServerID)) + } + } + if len(patterns) == 0 { + return "mcp__yao__*" + } + return strings.Join(patterns, ",") +} + +// buildSandboxEnvPrompt generates the sandbox environment prompt with the actual working directory. +func buildSandboxEnvPrompt(workDir string) string { + return fmt.Sprintf(`## Sandbox Environment + +You are running in a sandboxed environment with the following setup: + +- **Working Directory**: %[1]s +- **Project Structure**: If this is a new project, create a dedicated project folder (e.g., %[1]s/my-project/) and work inside it +- **File Access**: You have full read/write access to %[1]s +- **Output Files**: Save all output files to the working directory + +When creating new projects: +1. Create a project directory with a descriptive name +2. Initialize the project structure inside that directory +3. Keep all related files organized within the project folder + +## IMPORTANT: Restricted Tools + +The following tools are NOT available in this environment and you must NOT use them: +- EnterPlanMode, ExitPlanMode (use regular text to explain plans instead) +- Task, TaskOutput, TaskStop (complete tasks directly without delegation) +- AskUserQuestion (make reasonable assumptions instead of asking) +- Skill, ToolSearch (not supported) + +Focus on using the core tools: Bash, Read, Write, Edit, Glob, Grep, WebSearch, WebFetch. + +## User Attachments + +User-uploaded files (images, documents, code files, etc.) are placed in %[1]s/.attachments/{chatID}/ +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. + +## GitHub CLI (gh) Usage + +When working with GitHub and a token is provided: +1. First authenticate gh CLI using the token: echo "TOKEN" | gh auth login --with-token +2. Then use gh commands normally (gh repo create, gh pr create, etc.) +3. Do NOT use curl to call GitHub API directly - always prefer gh CLI +`, workDir) +} + +// resolveWorkDir returns the configured working directory, falling back to default. +func resolveWorkDir(cfg *types.SandboxConfig) string { + if cfg != nil && cfg.Computer.WorkDir != "" { + return cfg.Computer.WorkDir + } + return defaultWorkDir +} + +// resolveUserHome returns the home directory for the container user. +// Returns empty string if no user is configured (root or unspecified). +func resolveUserHome(cfg *types.SandboxConfig) string { + if cfg == nil { + return defaultUserHome + } + user := cfg.Computer.User + if user == "" { + user = defaultUser + } + if user == "root" { + return "/root" + } + return fmt.Sprintf("/home/%s", user) +} + +var claudeArgWhitelist = map[string]string{ + "max_turns": "--max-turns", + "disallowed_tools": "--disallowed-tools", + "allowed_tools": "--allowedTools", +} diff --git a/agent/sandbox/v2/claude/runner_test.go b/agent/sandbox/v2/claude/runner_test.go new file mode 100644 index 00000000..97afa979 --- /dev/null +++ b/agent/sandbox/v2/claude/runner_test.go @@ -0,0 +1,279 @@ +package claude_test + +import ( + "bytes" + "context" + "fmt" + "mime/multipart" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/yao/agent/caller" + agentcontext "github.com/yaoapp/yao/agent/context" + sandboxtestutils "github.com/yaoapp/yao/agent/sandbox/v2/testutils" + "github.com/yaoapp/yao/attachment" + oauthtypes "github.com/yaoapp/yao/openapi/oauth/types" +) + +type e2eCase struct { + ID string + Prompt string + Timeout time.Duration +} + +var cases = []e2eCase{ + { + ID: "tests.sandbox-v2.oneshot-cli", + Prompt: "Reply exactly with: hello sandbox v2", + Timeout: 3 * time.Minute, + }, +} + +func TestSandboxV2_Claude_E2E(t *testing.T) { + sandboxtestutils.Prepare(t) + defer sandboxtestutils.Clean(t) + + require.NotNil(t, caller.AgentGetterFunc, "AgentGetterFunc should be registered after Prepare") + + for _, tc := range cases { + tc := tc + t.Run(tc.ID, func(t *testing.T) { + agent, err := caller.AgentGetterFunc(tc.ID) + require.NoError(t, err, "should load assistant %s", tc.ID) + + timeout := tc.Timeout + if timeout == 0 { + timeout = 3 * time.Minute + } + + chatID := fmt.Sprintf("e2e-%s-%d", tc.ID, time.Now().UnixMilli()) + ctx := agentcontext.New( + context.Background(), + &oauthtypes.AuthorizedInfo{ + TeamID: "test-team-e2e", + UserID: "test-user-e2e", + }, + chatID, + ) + + messages := []agentcontext.Message{ + {Role: "user", Content: tc.Prompt}, + } + + done := make(chan struct{}) + var resp *agentcontext.Response + var streamErr error + + go func() { + defer close(done) + resp, streamErr = agent.Stream(ctx, messages) + }() + + select { + case <-done: + case <-time.After(timeout): + t.Fatalf("timeout after %v", timeout) + } + + require.NoError(t, streamErr, "Stream should not return error") + require.NotNil(t, resp, "response should not be nil") + + // ── 1. CompletionResponse should behave like the LLM path ── + require.NotNil(t, resp.Completion, "completion should not be nil") + assert.Equal(t, "assistant", resp.Completion.Role, "role should be assistant") + assert.Equal(t, agentcontext.FinishReasonStop, resp.Completion.FinishReason, "finish_reason should be stop") + assert.NotNil(t, resp.Completion.Content, "Content should be populated (same as LLM path)") + + contentStr, ok := resp.Completion.Content.(string) + require.True(t, ok, "Content should be a string, got %T", resp.Completion.Content) + t.Logf("CompletionResponse.Content (%d chars): %s", len(contentStr), contentStr) + assert.Contains(t, contentStr, "hello sandbox v2", "Content should contain expected text") + + // ── 2. Buffer: frame sequence handled correctly ── + require.NotNil(t, ctx.Buffer, "ctx.Buffer should not be nil") + + msgs := ctx.Buffer.GetMessages() + t.Logf("buffer message count: %d", len(msgs)) + for _, m := range msgs { + t.Logf(" seq=%d role=%s type=%s streaming=%v props_keys=%v", + m.Sequence, m.Role, m.Type, m.IsStreaming, mapKeys(m.Props)) + } + + var userInputCount, assistantTextCount, loadingCount int + var bufferTextContent string + for _, m := range msgs { + switch { + case m.Role == "user" && m.Type == "user_input": + userInputCount++ + case m.Role == "assistant" && m.Type == "loading": + loadingCount++ + case m.Role == "assistant" && m.Type == "text": + assistantTextCount++ + assert.False(t, m.IsStreaming, "text message should not be streaming (handleMessageEnd should have finalized it)") + require.NotNil(t, m.Props, "text message props should not be nil") + if c, ok := m.Props["content"].(string); ok { + bufferTextContent += c + } + } + } + + assert.Equal(t, 1, userInputCount, "should have exactly 1 user_input message") + assert.GreaterOrEqual(t, loadingCount, 1, "should have at least 1 loading message") + assert.Equal(t, 1, assistantTextCount, "should have exactly 1 assistant text message (from handleMessageEnd)") + assert.Contains(t, bufferTextContent, "hello sandbox v2", "buffer text should contain expected content") + + // ── 3. Buffer content matches CompletionResponse.Content ── + assert.Equal(t, contentStr, bufferTextContent, + "CompletionResponse.Content and Buffer text should match") + }) + } +} + +func TestSandboxV2_Claude_Attachments(t *testing.T) { + sandboxtestutils.Prepare(t) + defer sandboxtestutils.Clean(t) + + require.NotNil(t, caller.AgentGetterFunc, "AgentGetterFunc should be registered after Prepare") + + agent, err := caller.AgentGetterFunc("tests.sandbox-v2.oneshot-cli") + require.NoError(t, err) + + // ── 1. Locate testdata via runtime.Caller ── + _, thisFile, _, ok := runtime.Caller(0) + require.True(t, ok) + testdataDir := filepath.Join(filepath.Dir(thisFile), "testdata") + + // ── 2. Create attachment manager and upload test files ── + const uploaderName = "__yao.attachment" + manager, err := attachment.New(attachment.ManagerOption{ + Driver: "local", + MaxSize: "50M", + AllowedTypes: []string{"image/*", "text/*", "application/*", "video/*", ".ts", ".js", ".tsx", ".jsx"}, + Options: map[string]interface{}{"path": filepath.Join(os.TempDir(), "test_sandbox_v2_attach")}, + }) + require.NoError(t, err) + manager.Name = uploaderName + attachment.Managers[uploaderName] = manager + t.Cleanup(func() { delete(attachment.Managers, uploaderName) }) + + imgFile := uploadTestFile(t, manager, testdataDir, "test-image.png", "image/png") + codeFile := uploadTestFile(t, manager, testdataDir, "code.ts", "text/plain") + + imgWrapper := fmt.Sprintf("%s://%s", uploaderName, imgFile.ID) + codeWrapper := fmt.Sprintf("%s://%s", uploaderName, codeFile.ID) + t.Logf("image wrapper: %s", imgWrapper) + t.Logf("code wrapper: %s", codeWrapper) + + // ── 3. Build multimodal messages (same as CUI InputArea) ── + chatID := fmt.Sprintf("e2e-attach-%d", time.Now().UnixMilli()) + ctx := agentcontext.New( + context.Background(), + &oauthtypes.AuthorizedInfo{TeamID: "test-team-e2e", UserID: "test-user-e2e"}, + chatID, + ) + + messages := []agentcontext.Message{ + { + Role: "user", + Content: []interface{}{ + map[string]interface{}{"type": "text", "text": "Describe the attached image and summarize the attached code file. Reply in English."}, + map[string]interface{}{ + "type": "image_url", + "image_url": map[string]interface{}{"url": imgWrapper, "detail": "auto"}, + }, + map[string]interface{}{ + "type": "file", + "file": map[string]interface{}{"url": codeWrapper, "filename": "code.ts"}, + }, + }, + }, + } + + // ── 4. Run E2E stream ── + done := make(chan struct{}) + var resp *agentcontext.Response + var streamErr error + + go func() { + defer close(done) + resp, streamErr = agent.Stream(ctx, messages) + }() + + select { + case <-done: + case <-time.After(5 * time.Minute): + t.Fatalf("timeout after 5m") + } + + require.NoError(t, streamErr, "Stream should not return error") + require.NotNil(t, resp) + require.NotNil(t, resp.Completion) + + contentStr, ok := resp.Completion.Content.(string) + require.True(t, ok, "Content should be a string, got %T", resp.Completion.Content) + t.Logf("Response (%d chars): %s", len(contentStr), contentStr) + + lower := strings.ToLower(contentStr) + + // ── 5. Verify Claude actually read the image ── + imageKeywords := []string{"hello", "utf", "chinese", "text", "emoji"} + imgHit := false + for _, kw := range imageKeywords { + if strings.Contains(lower, kw) { + imgHit = true + break + } + } + assert.True(t, imgHit, "response should mention image content (tried: %v)", imageKeywords) + + // ── 6. Verify Claude actually read the code ── + codeKeywords := []string{"excel", "typescript", "class", "volcengine"} + codeHit := false + for _, kw := range codeKeywords { + if strings.Contains(lower, kw) { + codeHit = true + break + } + } + assert.True(t, codeHit, "response should mention code content (tried: %v)", codeKeywords) +} + +func uploadTestFile(t *testing.T, manager *attachment.Manager, testdataDir, filename, contentType string) *attachment.File { + t.Helper() + path := filepath.Join(testdataDir, filename) + data, err := os.ReadFile(path) + require.NoError(t, err, "read testdata/%s", filename) + + fh := &attachment.FileHeader{ + FileHeader: &multipart.FileHeader{ + Filename: filename, + Size: int64(len(data)), + Header: make(map[string][]string), + }, + } + fh.Header.Set("Content-Type", contentType) + + file, err := manager.Upload(context.Background(), fh, bytes.NewReader(data), attachment.UploadOption{ + Groups: []string{"e2e-sandbox-v2"}, + }) + require.NoError(t, err, "upload testdata/%s", filename) + t.Logf("uploaded %s => ID=%s, Path=%s", filename, file.ID, file.Path) + return file +} + +func mapKeys(m map[string]interface{}) []string { + if m == nil { + return nil + } + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + return keys +} diff --git a/agent/sandbox/v2/claude/testdata/code.ts b/agent/sandbox/v2/claude/testdata/code.ts new file mode 100644 index 00000000..458e8de1 --- /dev/null +++ b/agent/sandbox/v2/claude/testdata/code.ts @@ -0,0 +1,904 @@ +import { Process } from "@yao/runtime"; + +/** + * Excel class for manipulating Excel files via Yao's Excel Module + */ +export class Excel { + private handle: string | null = null; + + /** + * Creates a new Excel instance + * @param file Path to the Excel file + */ + constructor(private file: string, writable: boolean = false) { + this.file = file; + this.Open(writable); + } + + /** + * Read each sheet top n rows + * @param file Path to the Excel file + * @param n number of rows to read + * @returns Object with sheet names as keys and arrays of row values as values + */ + static Heads( + file: string, + n: number = 5, + filters?: string[] + ): Record { + const excel = new Excel(file); + const heads = excel.Heads(n, filters); + excel.Close(); + return heads; + } + + /** + * Read each sheet top n rows + * @param n number of rows to read + * @returns Object with sheet names as keys and arrays of row values as values + * @throws Error if file not opened + */ + Heads(n: number = 5, filters: string[] = []): Record { + if (!this.handle) throw new Error("Excel file not opened"); + + const sheets = this.Sheets(); + const result: Record = {}; + + for (const sheet of sheets) { + if (filters.length > 0 && !filters.includes(sheet)) { + continue; + } + + // Open row iterator for the sheet + const iterator = this.each.OpenRow(sheet); + const rows: any[][] = []; + + // Read n rows + let row; + let count = 0; + while ( + count < n && + (row = Process(`excel.each.NextRow`, iterator)) !== null + ) { + // Add column headers (A, B, C, ...) for the first row + if (count === 0) { + const headerRow = []; + for (let i = 0; i < row.length; i++) { + headerRow.push(this.convert.ColumnNumberToName(i + 1)); + } + rows.push(headerRow); + } + + // Trim Each cell's value + row = row.map((cell) => cell?.trim?.()); + rows.push(row); + count++; + } + + // Close the row iterator + this.each.CloseRow(iterator); + + // Find the max length of each row, and pad the column headers(A, B, C, ...) to the same length + const maxLength = Math.max(...rows.map((row) => row.length)); + const start = rows[0].length; + const neededLength = maxLength - rows[0].length; + for (let i = 0; i < neededLength; i++) { + rows[0].push(this.convert.ColumnNumberToName(start + i + 1)); + } + + // Add the sheet's rows to the result + result[sheet] = rows; + } + + return result; + } + + /** + * Check if a sheet exists in the Excel file + * @param file Path to the Excel file + * @param sheet Sheet name to check + * @returns boolean - true if sheet exists, false otherwise + */ + static Exists(file: string, sheet: string) { + const excel = new Excel(file); + const exists = excel.sheet.Exists(sheet); + excel.Close(); + return exists; + } + + /** + * Opens an Excel file for reading or writing + * @param writable Whether to open in writable mode (true) or read-only mode (false) + * @returns Handle ID used for subsequent operations + */ + Open(writable: boolean = false) { + this.handle = Process(`excel.Open`, this.file, writable); + return this.handle; + } + + /** + * Closes the Excel file and releases resources + * IMPORTANT: Always call this method when done to prevent memory leaks + */ + Close() { + if (this.handle) { + Process(`excel.Close`, this.handle); + this.handle = null; + } + } + + /** + * Saves changes to the Excel file + * @throws Error if file not opened + */ + Save() { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.Save`, this.handle); + } + + /** + * Gets all sheet names in the workbook + * @returns Array of sheet names + * @throws Error if file not opened + */ + Sheets() { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.Sheets`, this.handle); + } + + // Sheet operations + sheet = { + /** + * Creates a new sheet in the workbook + * @param name Name for the new sheet + * @returns number Index of the new sheet + * @throws Error if file not opened + */ + Create: (name: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.sheet.create`, this.handle, name); + }, + + /** + * Lists all sheets in the workbook + * @returns string[] Array of sheet names + * @throws Error if file not opened + */ + List: () => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.sheet.list`, this.handle); + }, + + /** + * Checks if a sheet exists in the workbook + * @param name Sheet name to check + * @returns boolean - true if sheet exists, false otherwise + * @throws Error if file not opened + */ + Exists: (name: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.sheet.exists`, this.handle, name); + }, + + /** + * Reads all data from a sheet + * @param name Sheet name + * @returns any[][] Two-dimensional array of cell values + * @throws Error if file not opened + */ + Read: (name: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.sheet.read`, this.handle, name); + }, + + /** + * Reads all data from a sheet with pagination support + * @param name Sheet name + * @param from Starting row index (0-based) + * @param chunk_size Number of rows to read + * @returns any[][] Two-dimensional array of cell values + * @throws Error if file not opened + */ + Rows: (name: string, from: number, chunk_size: number) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.sheet.rows`, this.handle, name, from, chunk_size); + }, + + /** + * Updates data in a sheet. Creates the sheet if it doesn't exist. + * @param name Sheet name + * @param data Two-dimensional array of values to write + * @throws Error if file not opened + */ + Update: (name: string, data: any[][]) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.sheet.update`, this.handle, name, data); + }, + + /** + * Copies a sheet with all its content and formatting + * @param source Source sheet name + * @param target Target sheet name (must not exist) + * @throws Error if file not opened + */ + Copy: (source: string, target: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.sheet.copy`, this.handle, source, target); + }, + + /** + * Deletes a sheet from the workbook + * @param name Sheet name to delete + * @throws Error if file not opened + */ + Delete: (name: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.sheet.delete`, this.handle, name); + }, + + /** + * Gets the dimensions (number of rows and columns) of a sheet + * @param name Sheet name + * @returns {rows: number, cols: number} - Object containing row and column counts + * @throws Error if file not opened + */ + Dimension: (name: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.sheet.dimension`, this.handle, name); + }, + }; + + // Reading operations + read = { + /** + * Reads a cell's value + * @param sheet Sheet name + * @param cell Cell reference (e.g. "A1") + * @returns Cell value + * @throws Error if file not opened + */ + Cell: (sheet: string, cell: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.read.Cell`, this.handle, sheet, cell); + }, + + /** + * Reads all rows in a sheet + * @param sheet Sheet name + * @returns Two-dimensional array of cell values + * @throws Error if file not opened + */ + Row: (sheet: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.read.Row`, this.handle, sheet); + }, + + /** + * Reads all columns in a sheet + * @param sheet Sheet name + * @returns Two-dimensional array of cell values + * @throws Error if file not opened + */ + Column: (sheet: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.read.Column`, this.handle, sheet); + }, + }; + + // Writing operations + write = { + /** + * Writes a value to a cell + * @param sheet Sheet name + * @param cell Cell reference (e.g. "A1") + * @param value Value to write (string, number, boolean, etc.) + * @throws Error if file not opened + */ + Cell: (sheet: string, cell: string, value: any) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.write.Cell`, this.handle, sheet, cell, value); + }, + + /** + * Writes values to a row starting at the specified cell + * @param sheet Sheet name + * @param startCell Starting cell reference (e.g. "A1") + * @param values Array of values to write + * @throws Error if file not opened + */ + Row: (sheet: string, startCell: string, values: any[]) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.write.Row`, this.handle, sheet, startCell, values); + }, + + /** + * Writes values to a column starting at the specified cell + * @param sheet Sheet name + * @param startCell Starting cell reference (e.g. "A1") + * @param values Array of values to write + * @throws Error if file not opened + */ + Column: (sheet: string, startCell: string, values: any[]) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process( + `excel.write.Column`, + this.handle, + sheet, + startCell, + values + ); + }, + + /** + * Writes a two-dimensional array of values starting at the specified cell + * @param sheet Sheet name + * @param startCell Starting cell reference (e.g. "A1") + * @param values Two-dimensional array of values to write + * @throws Error if file not opened + */ + All: (sheet: string, startCell: string, values: any[][]) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.write.All`, this.handle, sheet, startCell, values); + }, + }; + + // Setting properties + set = { + /** + * Applies a style to a cell + * @param sheet Sheet name + * @param cell Cell reference (e.g. "A1") + * @param styleID Style ID to apply + * @throws Error if file not opened + */ + Style: (sheet: string, cell: string, styleID: number) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.set.Style`, this.handle, sheet, cell, styleID); + }, + + /** + * Sets a row's height + * @param sheet Sheet name + * @param row Row number + * @param height Height in points + * @throws Error if file not opened + */ + RowHeight: (sheet: string, row: number, height: number) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.set.RowHeight`, this.handle, sheet, row, height); + }, + + /** + * Sets column width for a range of columns + * @param sheet Sheet name + * @param startCol Starting column letter + * @param endCol Ending column letter + * @param width Width in points + * @throws Error if file not opened + */ + ColumnWidth: ( + sheet: string, + startCol: string, + endCol: string, + width: number + ) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process( + `excel.set.ColumnWidth`, + this.handle, + sheet, + startCol, + endCol, + width + ); + }, + + /** + * Merges cells in a range + * @param sheet Sheet name + * @param startCell Starting cell reference (e.g. "A1") + * @param endCell Ending cell reference (e.g. "B2") + * @throws Error if file not opened + */ + MergeCell: (sheet: string, startCell: string, endCell: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process( + `excel.set.MergeCell`, + this.handle, + sheet, + startCell, + endCell + ); + }, + + /** + * Unmerges previously merged cells + * @param sheet Sheet name + * @param startCell Starting cell reference (e.g. "A1") + * @param endCell Ending cell reference (e.g. "B2") + * @throws Error if file not opened + */ + UnmergeCell: (sheet: string, startCell: string, endCell: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process( + `excel.set.UnmergeCell`, + this.handle, + sheet, + startCell, + endCell + ); + }, + + /** + * Sets a formula in a cell + * @param sheet Sheet name + * @param cell Cell reference (e.g. "C1") + * @param formula Excel formula without the leading equals sign + * @throws Error if file not opened + */ + Formula: (sheet: string, cell: string, formula: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.set.Formula`, this.handle, sheet, cell, formula); + }, + + /** + * Adds a hyperlink to a cell + * @param sheet Sheet name + * @param cell Cell reference (e.g. "A1") + * @param url URL for the hyperlink + * @param text Display text for the hyperlink + * @throws Error if file not opened + */ + Link: (sheet: string, cell: string, url: string, text: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.set.Link`, this.handle, sheet, cell, url, text); + }, + }; + + // Iteration methods + each = { + /** + * Opens a row iterator + * @param sheet Sheet name + * @returns Row iterator ID + * @throws Error if file not opened + */ + OpenRow: (sheet: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.each.OpenRow`, this.handle, sheet); + }, + + /** + * Gets the next row from the iterator + * @param rowID Row iterator ID from excel.each.OpenRow + * @returns Array of cell values or null if no more rows + */ + NextRow: (rowID: string) => { + return Process(`excel.each.NextRow`, rowID); + }, + + /** + * Closes the row iterator + * @param rowID Row iterator ID from excel.each.OpenRow + */ + CloseRow: (rowID: string) => { + return Process(`excel.each.CloseRow`, rowID); + }, + + /** + * Opens a column iterator + * @param sheet Sheet name + * @returns Column iterator ID + * @throws Error if file not opened + */ + OpenColumn: (sheet: string) => { + if (!this.handle) throw new Error("Excel file not opened"); + return Process(`excel.each.OpenColumn`, this.handle, sheet); + }, + + /** + * Gets the next column from the iterator + * @param colID Column iterator ID from excel.each.OpenColumn + * @returns Array of cell values or null if no more columns + */ + NextColumn: (colID: string) => { + return Process(`excel.each.NextColumn`, colID); + }, + + /** + * Closes the column iterator + * @param colID Column iterator ID from excel.each.OpenColumn + */ + CloseColumn: (colID: string) => { + return Process(`excel.each.CloseColumn`, colID); + }, + }; + + // Conversion utilities + convert = { + /** + * Converts a column name to a column number + * @param colName Column name (e.g. "A", "AB") + * @returns Column number (1-based) + */ + ColumnNameToNumber: (colName: string) => { + return Process(`excel.convert.ColumnNameToNumber`, colName); + }, + + /** + * Converts a column number to a column name + * @param colNum Column number (1-based) + * @returns Column name + */ + ColumnNumberToName: (colNum: number) => { + return Process(`excel.convert.ColumnNumberToName`, colNum); + }, + + /** + * Converts a cell reference to coordinates + * @param cell Cell reference (e.g. "A1") + * @returns Array with [columnNumber, rowNumber] (1-based) + */ + CellNameToCoordinates: (cell: string) => { + return Process(`excel.convert.CellNameToCoordinates`, cell); + }, + + /** + * Converts coordinates to a cell reference + * @param col Column number (1-based) + * @param row Row number (1-based) + * @returns Cell reference + */ + CoordinatesToCellName: (col: number, row: number) => { + return Process(`excel.convert.CoordinatesToCellName`, col, row); + }, + }; +} +/** + * Volcengine OpenAPI SDK + */ +import { Exception, http, Process } from "@yao/runtime"; + +export class Volcengine { + private AccessKeyId: string; + private SecretAccessKey: string; + private Region: string; + private Service: string; + private Endpoint: string; + constructor(option: Option) { + this.AccessKeyId = option.AccessKeyId; + this.SecretAccessKey = option.SecretAccessKey; + this.Region = option.Region; + this.Service = option.Service; + this.Endpoint = option.Endpoint + ? `https://${option.Endpoint}` + : `https://${this.Service}.${this.Region}.volcengineapi.com`; + } + + public Get(query: Record) { + const url = this.Endpoint; + const host = url.split("://")[1].split("/")[0]; + const headers = { host: host }; + const request: Request = { + Method: "GET", + URI: "/", + Query: query, + Headers: headers, + Payload: null, + }; + + const auth = this.getAuthorization(request); + + // Add authorization header + headers["Authorization"] = auth; + headers["Content-Type"] = "application/json"; + + const resp = http.Get(url, query, headers); + if (resp.code > 299 || resp.code < 200) { + const { ResponseMetadata } = resp.data || {}; + const { Error } = ResponseMetadata || {}; + const message = + Error?.Message || (resp.code === 0 ? resp.message : "Unknown error"); + throw new Exception(message, resp.code); + } + + return resp.data; + } + + /** + * Post request + * @param query Query parameters + * @param payload Payload + * @returns Response + */ + public Post(query: Record, payload: Record) { + const url = this.Endpoint; + const host = url.split("://")[1].split("/")[0]; + const headers = { host: host }; + const body = JSON.stringify(payload); + const request: Request = { + Method: "POST", + URI: "/", + Query: query, + Headers: headers, + Payload: body, + }; + + const auth = this.getAuthorization(request); + headers["Authorization"] = auth; + headers["Content-Type"] = "application/json"; + + const resp = http.Post(url, body, null, query, headers); + if (resp.code > 299 || resp.code < 200) { + const { ResponseMetadata } = resp.data || {}; + const { Error } = ResponseMetadata || {}; + const message = + Error?.Message || (resp.code === 0 ? resp.message : "Unknown error"); + throw new Exception(message, resp.code); + } + return resp.data; + } + + /** + * Create a canonical request + * @param request Request object + * @returns Canonical request string + */ + private canonicalRequest(request: Request): string { + const xDate = this.formatDate(new Date()); + + // 1. HTTP Method + const method = request.Method; + + // 2. URI (default to '/' if null) + const uri = request.URI || "/"; + + // 3. Query String + let queryString = ""; + if (request.Query) { + if (Array.isArray(request.Query)) { + // Handle array of query parameters + const queryParams = request.Query.reduce((acc: string[], curr) => { + Object.entries(curr).forEach(([key, value]) => { + if (value !== null && value !== undefined && value !== "") { + acc.push( + `${encodeURIComponent(key)}=${encodeURIComponent(value)}` + ); + } + }); + return acc; + }, []); + queryString = queryParams.sort().join("&"); + } else { + // Handle single query object + const queryParams = Object.entries(request.Query) + .filter( + ([_, value]) => + value !== null && value !== undefined && value !== "" + ) + .map( + ([key, value]) => + `${encodeURIComponent(key)}=${encodeURIComponent(value)}` + ) + .sort(); + queryString = queryParams.join("&"); + } + } + + // 4. Headers + // First, collect all headers in a normalized format + const headers: Record = { "x-date": xDate }; + if (request.Headers) { + if (Array.isArray(request.Headers)) { + request.Headers.forEach((headerObj) => { + Object.entries(headerObj).forEach(([key, value]) => { + if (value !== null && value !== undefined && value.trim() !== "") { + headers[key.toLowerCase()] = value.trim(); + } + }); + }); + } else { + Object.entries(request.Headers).forEach(([key, value]) => { + if (value !== null && value !== undefined && value.trim() !== "") { + headers[key.toLowerCase()] = value.trim(); + } + }); + } + } + + // Get required headers if they exist + const signedHeaderKeys: string[] = []; + const requiredHeaders = ["host", "x-date"]; + + // Add required headers first if they exist + requiredHeaders.forEach((key) => { + if (headers[key]) { + signedHeaderKeys.push(key); + } + }); + + // Add any additional headers + // const additionalHeaders = Object.keys(headers) + // .filter((key) => !requiredHeaders.includes(key)) + // .sort(); + // signedHeaderKeys.push(...additionalHeaders); + + // Build canonical headers string + const canonicalHeaders = signedHeaderKeys + .map((key) => `${key}:${headers[key]}`) + .join("\n"); + + // Build signed headers string + const signedHeaders = signedHeaderKeys.join(";"); + + // 5. Payload/Body + let hashedPayload = Process("crypto.Hash", "SHA256", ""); + if (request.Payload !== null && request.Payload !== undefined) { + if (typeof request.Payload === "string") { + if (request.Payload !== "") { + hashedPayload = Process("crypto.Hash", "SHA256", request.Payload); + } + } else { + const payload = JSON.stringify(request.Payload); + if (payload !== "{}" && payload !== "[]") { + hashedPayload = Process("crypto.Hash", "SHA256", payload); + } + } + } + + // Combine all components + const parts = [ + method, + uri, + queryString, + canonicalHeaders, + "", // Empty line after headers + signedHeaders, + hashedPayload, + ]; + + return parts.join("\n"); + } + + /** + * Format date to YYYYMMDDTHHMMSSZ + * @param date Date object + * @returns Formatted date string + */ + private formatDate(date: Date): string { + const year = date.getUTCFullYear(); + const month = String(date.getUTCMonth() + 1).padStart(2, "0"); + const day = String(date.getUTCDate()).padStart(2, "0"); + const hours = String(date.getUTCHours()).padStart(2, "0"); + const minutes = String(date.getUTCMinutes()).padStart(2, "0"); + const seconds = String(date.getUTCSeconds()).padStart(2, "0"); + return `${year}${month}${day}T${hours}${minutes}${seconds}Z`; + } + + /** + * Create string to sign + * @param canonicalRequest Canonical request string + * @returns String to sign + */ + private stringToSign(canonicalRequest: string): string { + const algorithm = "HMAC-SHA256"; + const requestDateTime = this.formatDate(new Date()); + const requestDate = requestDateTime.slice(0, 8); + const credentialScope = `${requestDate}/${this.Region}/${this.Service}/request`; // YYYYMMDD + + const hashedCanonicalRequest = Process( + "crypto.Hash", + "SHA256", + canonicalRequest + ); + + return `${algorithm}\n${requestDateTime}\n${credentialScope}\n${hashedCanonicalRequest}`; + } + + /** + * Derive signing key + * @param date Date in YYYY/MM/DD format + * @returns Signing key + */ + private getSigningKey(date: string): string { + const kDate = Process("crypto.HMAC", "SHA256", date, this.SecretAccessKey); + const kRegion = Process( + "crypto.HMACWith", + { key: "hex" }, + this.Region, + kDate + ); + const kService = Process( + "crypto.HMACWith", + { key: "hex" }, + this.Service, + kRegion + ); + + const kSigning = Process( + "crypto.HMACWith", + { key: "hex" }, + "request", + kService + ); + return kSigning; + } + + /** + * Calculate signature + * @param stringToSign String to sign + * @param signingKey Signing key + * @returns Signature + */ + private signature(stringToSign: string, signingKey: string): string { + return Process("crypto.HMACWith", { key: "hex" }, stringToSign, signingKey); + } + + /** + * Build authorization header + * @param request Request object + * @returns Authorization header value + */ + public getAuthorization(request: Request): string { + const xDate = this.formatDate(new Date()); + if (request.Headers) { + if (typeof request.Headers === "object") { + request.Headers["x-date"] = request.Headers["x-date"] + ? request.Headers["x-date"] + : xDate; + } + } + + // 1. Create canonical request + const canonicalReq = this.canonicalRequest(request); + + // 2. Create string to sign + const stringToSign = this.stringToSign(canonicalReq); + + // 3. Get date from string to sign + const [algorithm, requestDateTime, credentialScope] = + stringToSign.split("\n"); + const date = requestDateTime.slice(0, 8); + + // 4. Derive signing key + const signingKey = this.getSigningKey(date); + // 5. Calculate signature + const signature = this.signature(stringToSign, signingKey); + + // 6. Build authorization header + let signedHeaders = ""; + if (request.Headers) { + const headers: Record = {}; + if (Array.isArray(request.Headers)) { + request.Headers.forEach((headerObj) => { + Object.entries(headerObj).forEach(([key, value]) => { + headers[key.toLowerCase()] = value.trim(); + }); + }); + } else { + Object.entries(request.Headers).forEach(([key, value]) => { + headers[key.toLowerCase()] = value.trim(); + }); + } + signedHeaders = Object.keys(headers).sort().join(";"); + } + + return `${algorithm} Credential=${this.AccessKeyId}/${credentialScope}, SignedHeaders=${signedHeaders}, Signature=${signature}`; + } +} + +export interface Option { + AccessKeyId: string; + SecretAccessKey: string; + Endpoint?: string; + Region: string; + Service: string; +} + +export interface Request { + Method: "GET" | "POST"; + URI: string | null; // Default / + Query: Record | Record[] | null; + Headers: Record | Record[] | null; + Payload: string | Record | any[] | null; +} diff --git a/agent/sandbox/v2/claude/testdata/test-image.png b/agent/sandbox/v2/claude/testdata/test-image.png new file mode 100644 index 00000000..7c354fd1 Binary files /dev/null and b/agent/sandbox/v2/claude/testdata/test-image.png differ diff --git a/agent/sandbox/v2/init.go b/agent/sandbox/v2/init.go new file mode 100644 index 00000000..cd13d530 --- /dev/null +++ b/agent/sandbox/v2/init.go @@ -0,0 +1,13 @@ +package sandboxv2 + +import ( + "github.com/yaoapp/yao/agent/sandbox/v2/claude" + "github.com/yaoapp/yao/agent/sandbox/v2/types" + yaorunner "github.com/yaoapp/yao/agent/sandbox/v2/yao" +) + +func init() { + Register("claude", func() types.Runner { return claude.New() }) + Register("claude/cli", func() types.Runner { return claude.New() }) + Register("yao", func() types.Runner { return yaorunner.New() }) +} diff --git a/agent/sandbox/v2/lifecycle.go b/agent/sandbox/v2/lifecycle.go new file mode 100644 index 00000000..84680a1c --- /dev/null +++ b/agent/sandbox/v2/lifecycle.go @@ -0,0 +1,150 @@ +package sandboxv2 + +import ( + "context" + "crypto/rand" + "encoding/hex" + "fmt" + "log" + + agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/sandbox/v2/types" + infra "github.com/yaoapp/yao/sandbox/v2" +) + +// 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 string, metadata map[string]any) string { + if cfg.Lifecycle == "oneshot" { + return "" + } + + // Custom identifier from metadata takes precedence. + if metadata != nil { + if cid, ok := metadata["computer_id"].(string); ok && cid != "" { + return fmt.Sprintf("%s-%s", ownerID, cid) + } + } + + switch cfg.Lifecycle { + case "session": + return fmt.Sprintf("%s-%s", ownerID, chatID) + case "longrunning", "persistent": + return fmt.Sprintf("%s-%s", ownerID, assistantID) + default: + return "" + } +} + +// GetComputer obtains or creates a Computer for the current request. +// 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) + identifier := BuildIdentifier(cfg, ownerID, ctx.ChatID, ctx.AssistantID, ctx.Metadata) + + // Fill runtime fields. + cfg.Owner = ownerID + cfg.ID = identifier + + workspaceID := "" + if ctx.Metadata != nil { + if ws, ok := ctx.Metadata["workspace_id"].(string); ok && ws != "" { + workspaceID = ws + } + } + if workspaceID == "" { + workspaceID = ownerID + } + cfg.WorkspaceID = workspaceID + + // Host mode: no image → host computer. + if cfg.Computer.Image == "" { + cfg.Kind = "host" + nodeID := cfg.NodeID + if nodeID == "" { + return nil, identifier, fmt.Errorf("host mode requires a nodeID (set in sandbox.yao or workspace)") + } + host, err := manager.Host(context.Background(), nodeID) + if err != nil { + return nil, identifier, fmt.Errorf("get host computer: %w", err) + } + host.BindWorkplace(workspaceID) + return host, identifier, nil + } + + cfg.Kind = "box" + + // Reuse: non-empty identifier → try Get first. + if identifier != "" { + box, err := manager.Get(context.Background(), identifier) + if err == nil && box != nil { + 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) + } + + // 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.Printf("[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" +} + +func randomID() string { + b := make([]byte, 8) + _, _ = rand.Read(b) + return hex.EncodeToString(b) +} diff --git a/agent/sandbox/v2/lifecycle_test.go b/agent/sandbox/v2/lifecycle_test.go new file mode 100644 index 00000000..64cf2a81 --- /dev/null +++ b/agent/sandbox/v2/lifecycle_test.go @@ -0,0 +1,547 @@ +package sandboxv2_test + +import ( + "context" + "fmt" + "strings" + "testing" + "time" + + agentContext "github.com/yaoapp/yao/agent/context" + sandboxv2 "github.com/yaoapp/yao/agent/sandbox/v2" + "github.com/yaoapp/yao/agent/sandbox/v2/types" + oauthTypes "github.com/yaoapp/yao/openapi/oauth/types" + infra "github.com/yaoapp/yao/sandbox/v2" +) + +// =========================================================================== +// BuildIdentifier — pure-function tests (no infra needed) +// =========================================================================== + +func TestBuildIdentifier_Oneshot(t *testing.T) { + cfg := &types.SandboxConfig{Lifecycle: "oneshot"} + id := sandboxv2.BuildIdentifier(cfg, "owner1", "chat1", "ast1", nil) + if id != "" { + t.Errorf("oneshot should return empty, got %q", id) + } +} + +func TestBuildIdentifier_Session(t *testing.T) { + cfg := &types.SandboxConfig{Lifecycle: "session"} + id := sandboxv2.BuildIdentifier(cfg, "owner1", "chat42", "ast1", nil) + if id != "owner1-chat42" { + t.Errorf("session: got %q, want %q", id, "owner1-chat42") + } +} + +func TestBuildIdentifier_Longrunning(t *testing.T) { + cfg := &types.SandboxConfig{Lifecycle: "longrunning"} + id := sandboxv2.BuildIdentifier(cfg, "owner1", "chat1", "ast99", nil) + if id != "owner1-ast99" { + t.Errorf("longrunning: got %q, want %q", id, "owner1-ast99") + } +} + +func TestBuildIdentifier_Persistent(t *testing.T) { + cfg := &types.SandboxConfig{Lifecycle: "persistent"} + id := sandboxv2.BuildIdentifier(cfg, "owner1", "chat1", "ast99", nil) + if id != "owner1-ast99" { + t.Errorf("persistent: got %q, want %q", id, "owner1-ast99") + } +} + +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", meta) + if id != "owner1-custom-box" { + t.Errorf("metadata override: got %q, want %q", id, "owner1-custom-box") + } +} + +func TestBuildIdentifier_MetadataEmptyIgnored(t *testing.T) { + cfg := &types.SandboxConfig{Lifecycle: "session"} + meta := map[string]any{"computer_id": ""} + id := sandboxv2.BuildIdentifier(cfg, "owner1", "chat42", "ast1", meta) + if id != "owner1-chat42" { + t.Errorf("empty metadata should fall through to session, got %q", id) + } +} + +func TestBuildIdentifier_UnknownLifecycle(t *testing.T) { + cfg := &types.SandboxConfig{Lifecycle: "unknown"} + id := sandboxv2.BuildIdentifier(cfg, "owner1", "chat1", "ast1", nil) + if id != "" { + t.Errorf("unknown lifecycle should return empty, got %q", id) + } +} + +// =========================================================================== +// GetComputer — real container tests +// =========================================================================== + +func makeAgentCtx(teamID, userID, chatID, assistantID string, metadata map[string]any) *agentContext.Context { + var auth *oauthTypes.AuthorizedInfo + if teamID != "" || userID != "" { + auth = &oauthTypes.AuthorizedInfo{TeamID: teamID, UserID: userID} + } + return &agentContext.Context{ + Context: context.Background(), + Authorized: auth, + ChatID: chatID, + AssistantID: assistantID, + Metadata: metadata, + } +} + +func TestGetComputer_BoxCreate(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + ensureImage(t, m, nc) + + wsID := fmt.Sprintf("lc-create-%d", time.Now().UnixNano()) + createTestWorkspace(t, nc.TaiID, wsID) + + cfg := &types.SandboxConfig{ + Version: "2.0", + Lifecycle: "oneshot", + Computer: types.ComputerConfig{Image: testImage()}, + NodeID: nc.TaiID, + } + meta := map[string]any{"workspace_id": wsID} + ctx := makeAgentCtx("team-t1", "", "chat-1", "ast-1", meta) + + computer, identifier, err := sandboxv2.GetComputer(ctx, cfg, m) + if err != nil { + t.Fatalf("GetComputer: %v", err) + } + defer cleanupComputer(t, m, cfg) + + if identifier == "" { + t.Fatal("oneshot should get a random identifier, got empty") + } + + info := computer.ComputerInfo() + if info.Kind != "box" { + t.Errorf("kind = %q, want %q", info.Kind, "box") + } + if cfg.Owner != "team-t1" { + t.Errorf("cfg.Owner = %q, want %q", cfg.Owner, "team-t1") + } + if cfg.Kind != "box" { + t.Errorf("cfg.Kind = %q, want %q", cfg.Kind, "box") + } + }) + } +} + +func TestGetComputer_BoxReuse(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + ensureImage(t, m, nc) + + wsID := fmt.Sprintf("lc-reuse-%d", time.Now().UnixNano()) + createTestWorkspace(t, nc.TaiID, wsID) + + cfg := &types.SandboxConfig{ + Version: "2.0", + Lifecycle: "session", + Computer: types.ComputerConfig{Image: testImage()}, + NodeID: nc.TaiID, + } + meta := map[string]any{"workspace_id": wsID} + ctx := makeAgentCtx("team-reuse", "", "chat-reuse", "ast-1", meta) + + computer1, id1, err := sandboxv2.GetComputer(ctx, cfg, m) + if err != nil { + t.Fatalf("first GetComputer: %v", err) + } + defer cleanupComputer(t, m, cfg) + + cfg2 := &types.SandboxConfig{ + Version: "2.0", + Lifecycle: "session", + Computer: types.ComputerConfig{Image: testImage()}, + NodeID: nc.TaiID, + } + computer2, id2, err := sandboxv2.GetComputer(ctx, cfg2, m) + if err != nil { + t.Fatalf("second GetComputer: %v", err) + } + + if id1 != id2 { + t.Errorf("identifiers differ: %q vs %q", id1, id2) + } + + info1 := computer1.ComputerInfo() + info2 := computer2.ComputerInfo() + if info1.ContainerID != info2.ContainerID { + t.Errorf("container IDs differ: %q vs %q (should reuse)", info1.ContainerID, info2.ContainerID) + } + }) + } +} + +func TestGetComputer_WorkspaceBindAlways(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + ensureImage(t, m, nc) + + wsID := fmt.Sprintf("lc-ws-%d", time.Now().UnixNano()) + createTestWorkspace(t, nc.TaiID, wsID) + + cfg := &types.SandboxConfig{ + Version: "2.0", + Lifecycle: "oneshot", + Computer: types.ComputerConfig{Image: testImage()}, + NodeID: nc.TaiID, + } + meta := map[string]any{"workspace_id": wsID} + ctx := makeAgentCtx("team-ws", "", "chat-ws", "ast-ws", meta) + + computer, _, err := sandboxv2.GetComputer(ctx, cfg, m) + if err != nil { + t.Fatalf("GetComputer: %v", err) + } + defer cleanupComputer(t, m, cfg) + + if cfg.WorkspaceID != wsID { + t.Errorf("WorkspaceID = %q, want %q", cfg.WorkspaceID, wsID) + } + + ws := computer.Workplace() + if ws == nil { + t.Fatal("Workplace() returned nil, workspace should always be bound") + } + }) + } +} + +func TestGetComputer_WorkspaceFallbackOwner(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + ensureImage(t, m, nc) + + ownerID := fmt.Sprintf("lc-owner-%d", time.Now().UnixNano()) + createTestWorkspace(t, nc.TaiID, ownerID) + + cfg := &types.SandboxConfig{ + Version: "2.0", + Lifecycle: "oneshot", + Computer: types.ComputerConfig{Image: testImage()}, + NodeID: nc.TaiID, + } + ctx := makeAgentCtx(ownerID, "", "chat-fb", "ast-fb", nil) + + computer, _, err := sandboxv2.GetComputer(ctx, cfg, m) + if err != nil { + t.Fatalf("GetComputer: %v", err) + } + defer cleanupComputer(t, m, cfg) + + if cfg.WorkspaceID != ownerID { + t.Errorf("WorkspaceID = %q, want %q (should fallback to ownerID)", cfg.WorkspaceID, ownerID) + } + + ws := computer.Workplace() + if ws == nil { + t.Fatal("Workplace() returned nil") + } + }) + } +} + +func TestGetComputer_OwnerPriority(t *testing.T) { + skipIfNoDocker(t) + + nc := boxNodes()[0] + m := setupManager(t, &nc) + ensureImage(t, m, nc) + + t.Run("teamID", func(t *testing.T) { + wsID := fmt.Sprintf("lc-ownp-team-%d", time.Now().UnixNano()) + createTestWorkspace(t, nc.TaiID, wsID) + + cfg := &types.SandboxConfig{ + Version: "2.0", Lifecycle: "oneshot", + Computer: types.ComputerConfig{Image: testImage()}, + NodeID: nc.TaiID, + } + ctx := makeAgentCtx("my-team", "my-user", "c", "a", map[string]any{"workspace_id": wsID}) + _, _, err := sandboxv2.GetComputer(ctx, cfg, m) + if err != nil { + t.Fatalf("GetComputer: %v", err) + } + defer cleanupComputer(t, m, cfg) + if cfg.Owner != "my-team" { + t.Errorf("Owner = %q, want %q (teamID takes precedence)", cfg.Owner, "my-team") + } + }) + + t.Run("userID", func(t *testing.T) { + wsID := fmt.Sprintf("lc-ownp-user-%d", time.Now().UnixNano()) + createTestWorkspace(t, nc.TaiID, wsID) + + cfg := &types.SandboxConfig{ + Version: "2.0", Lifecycle: "oneshot", + Computer: types.ComputerConfig{Image: testImage()}, + NodeID: nc.TaiID, + } + ctx := makeAgentCtx("", "my-user", "c", "a", map[string]any{"workspace_id": wsID}) + _, _, err := sandboxv2.GetComputer(ctx, cfg, m) + if err != nil { + t.Fatalf("GetComputer: %v", err) + } + defer cleanupComputer(t, m, cfg) + if cfg.Owner != "my-user" { + t.Errorf("Owner = %q, want %q", cfg.Owner, "my-user") + } + }) + + t.Run("anonymous", func(t *testing.T) { + wsID := fmt.Sprintf("lc-ownp-anon-%d", time.Now().UnixNano()) + createTestWorkspace(t, nc.TaiID, wsID) + + cfg := &types.SandboxConfig{ + Version: "2.0", Lifecycle: "oneshot", + Computer: types.ComputerConfig{Image: testImage()}, + NodeID: nc.TaiID, + } + ctx := makeAgentCtx("", "", "c", "a", map[string]any{"workspace_id": wsID}) + _, _, err := sandboxv2.GetComputer(ctx, cfg, m) + if err != nil { + t.Fatalf("GetComputer: %v", err) + } + defer cleanupComputer(t, m, cfg) + if cfg.Owner != "anonymous" { + t.Errorf("Owner = %q, want %q", cfg.Owner, "anonymous") + } + }) +} + +func TestGetComputer_HostMode(t *testing.T) { + skipIfNoHostExec(t) + + for _, tgt := range hostTargets() { + tgt := tgt + t.Run(tgt.Name, func(t *testing.T) { + m := setupHostManager(t, &tgt) + + cfg := &types.SandboxConfig{ + Version: "2.0", + Lifecycle: "session", + Computer: types.ComputerConfig{}, + NodeID: tgt.TaiID, + } + ctx := makeAgentCtx("team-host", "", "chat-host", "ast-host", nil) + + computer, _, err := sandboxv2.GetComputer(ctx, cfg, m) + if err != nil { + t.Fatalf("GetComputer host: %v", err) + } + + if cfg.Kind != "host" { + t.Errorf("Kind = %q, want %q", cfg.Kind, "host") + } + + info := computer.ComputerInfo() + if info.Kind != "host" { + t.Errorf("ComputerInfo.Kind = %q, want %q", info.Kind, "host") + } + + ws := computer.Workplace() + if ws == nil { + t.Fatal("Workplace() returned nil on host mode") + } + }) + } +} + +func TestGetComputer_HostMissingNodeID(t *testing.T) { + skipIfNoDocker(t) + + nc := boxNodes()[0] + m := setupManager(t, &nc) + + cfg := &types.SandboxConfig{ + Version: "2.0", + Lifecycle: "session", + Computer: types.ComputerConfig{}, + NodeID: "", + } + ctx := makeAgentCtx("team-err", "", "c", "a", nil) + + _, _, err := sandboxv2.GetComputer(ctx, cfg, m) + if err == nil { + t.Fatal("expected error for host mode without nodeID") + } + if !strings.Contains(err.Error(), "nodeID") { + t.Errorf("error should mention nodeID, got: %v", err) + } +} + +// =========================================================================== +// LifecycleAction — behavior tests +// =========================================================================== + +func TestLifecycleAction_Oneshot(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + ensureImage(t, m, nc) + + wsID := fmt.Sprintf("lc-oneshot-%d", time.Now().UnixNano()) + createTestWorkspace(t, nc.TaiID, wsID) + + cfg := &types.SandboxConfig{ + Version: "2.0", + Lifecycle: "oneshot", + Computer: types.ComputerConfig{Image: testImage()}, + NodeID: nc.TaiID, + } + ctx := makeAgentCtx("team-oneshot", "", "c", "a", map[string]any{"workspace_id": wsID}) + + computer, _, err := sandboxv2.GetComputer(ctx, cfg, m) + if err != nil { + t.Fatalf("GetComputer: %v", err) + } + + boxID := cfg.ID + + sandboxv2.LifecycleAction(context.Background(), cfg, computer, m) + + _, getErr := m.Get(context.Background(), boxID) + if getErr == nil { + t.Error("box should be removed after oneshot LifecycleAction") + } + }) + } +} + +func TestLifecycleAction_Session(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + ensureImage(t, m, nc) + + wsID := fmt.Sprintf("lc-sess-%d", time.Now().UnixNano()) + createTestWorkspace(t, nc.TaiID, wsID) + + cfg := &types.SandboxConfig{ + Version: "2.0", + Lifecycle: "session", + Computer: types.ComputerConfig{Image: testImage()}, + NodeID: nc.TaiID, + } + ctx := makeAgentCtx("team-sess", "", "chat-sess", "ast-sess", map[string]any{"workspace_id": wsID}) + + computer, _, err := sandboxv2.GetComputer(ctx, cfg, m) + if err != nil { + t.Fatalf("GetComputer: %v", err) + } + defer cleanupComputer(t, m, cfg) + + sandboxv2.LifecycleAction(context.Background(), cfg, computer, m) + + box, err := m.Get(context.Background(), cfg.ID) + if err != nil { + t.Fatalf("box should still exist after session LifecycleAction: %v", err) + } + if box == nil { + t.Fatal("box is nil after session LifecycleAction") + } + }) + } +} + +func TestLifecycleAction_Persistent(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + ensureImage(t, m, nc) + + wsID := fmt.Sprintf("lc-pers-%d", time.Now().UnixNano()) + createTestWorkspace(t, nc.TaiID, wsID) + + cfg := &types.SandboxConfig{ + Version: "2.0", + Lifecycle: "persistent", + Computer: types.ComputerConfig{Image: testImage()}, + NodeID: nc.TaiID, + } + ctx := makeAgentCtx("team-pers", "", "chat-pers", "ast-pers", map[string]any{"workspace_id": wsID}) + + computer, _, err := sandboxv2.GetComputer(ctx, cfg, m) + if err != nil { + t.Fatalf("GetComputer: %v", err) + } + defer cleanupComputer(t, m, cfg) + + sandboxv2.LifecycleAction(context.Background(), cfg, computer, m) + + box, err := m.Get(context.Background(), cfg.ID) + if err != nil { + t.Fatalf("box should still exist after persistent LifecycleAction: %v", err) + } + if box == nil { + t.Fatal("box is nil after persistent LifecycleAction") + } + }) + } +} + +func TestLifecycleAction_NilSafe(t *testing.T) { + cfg := &types.SandboxConfig{Lifecycle: "oneshot"} + sandboxv2.LifecycleAction(context.Background(), cfg, nil, nil) + sandboxv2.LifecycleAction(context.Background(), nil, nil, nil) +} + +// =========================================================================== +// helpers +// =========================================================================== + +func ensureImage(t *testing.T, m *infra.Manager, nc nodeConfig) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + if err := m.EnsureImage(ctx, nc.TaiID, testImage(), infra.ImagePullOptions{}); err != nil { + t.Fatalf("EnsureImage: %v", err) + } +} + +func cleanupComputer(t *testing.T, m *infra.Manager, cfg *types.SandboxConfig) { + t.Helper() + if cfg.ID == "" { + return + } + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + if err := m.Remove(ctx, cfg.ID); err != nil { + t.Logf("cleanup Remove(%s): %v", cfg.ID, err) + } +} diff --git a/agent/sandbox/v2/options.go b/agent/sandbox/v2/options.go new file mode 100644 index 00000000..df313a24 --- /dev/null +++ b/agent/sandbox/v2/options.go @@ -0,0 +1,168 @@ +package sandboxv2 + +import ( + "fmt" + "os" + "strings" + "time" + + "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. Pure runtime mapping — no file-system or DSL access. +func BuildCreateOptions(cfg *types.SandboxConfig, identifier, ownerID, workspaceID string) (infra.CreateOptions, error) { + opts := infra.CreateOptions{ + ID: identifier, + Owner: ownerID, + Image: cfg.Computer.Image, + WorkDir: cfg.Computer.WorkDir, + User: cfg.Computer.User, + MountPath: cfg.Computer.MountPath, + MountMode: cfg.Computer.MountMode, + WorkspaceID: workspaceID, + Labels: cfg.Labels, + } + + if opts.Labels == nil { + opts.Labels = make(map[string]string) + } + + // Lifecycle policy + switch cfg.Lifecycle { + case "oneshot": + opts.Policy = infra.OneShot + case "session": + opts.Policy = infra.Session + case "longrunning": + opts.Policy = infra.LongRunning + case "persistent": + opts.Policy = infra.Persistent + default: + opts.Policy = infra.OneShot + } + + // Timeouts + if cfg.IdleTimeout != "" { + d, err := time.ParseDuration(cfg.IdleTimeout) + if err != nil { + return opts, fmt.Errorf("idle_timeout: %w", err) + } + opts.IdleTimeout = d + } + if cfg.MaxLifetime != "" { + d, err := time.ParseDuration(cfg.MaxLifetime) + if err != nil { + return opts, fmt.Errorf("max_lifetime: %w", err) + } + opts.MaxLifetime = d + } + if cfg.StopTimeout != "" { + d, err := time.ParseDuration(cfg.StopTimeout) + if err != nil { + return opts, fmt.Errorf("stop_timeout: %w", err) + } + opts.StopTimeout = d + } + + // Memory (string like "4g" → bytes) + if cfg.Computer.Memory != "" { + mem, err := parseMemory(cfg.Computer.Memory) + if err != nil { + return opts, fmt.Errorf("memory: %w", err) + } + opts.Memory = mem + } + + opts.CPUs = cfg.Computer.CPUs + + // VNC + opts.VNC = cfg.Computer.VNC.Enabled + + // Ports + for _, p := range cfg.Computer.Ports { + opts.Ports = append(opts.Ports, infra.PortMapping{ + ContainerPort: p.Port, + HostPort: p.HostPort, + Protocol: p.Protocol, + }) + } + + // NodeID (host mode pre-selection) + if cfg.NodeID != "" { + opts.NodeID = cfg.NodeID + } + + // Merge environment + secrets into CreateOptions.Env. + // Secrets override environment for same-name keys. + // $ENV.XXX references are resolved at runtime. + envSize := len(cfg.Environment) + len(cfg.Secrets) + if envSize > 0 { + opts.Env = make(map[string]string, envSize) + for k, v := range cfg.Environment { + opts.Env[k] = resolveEnvRef(v) + } + for k, v := range cfg.Secrets { + opts.Env[k] = resolveEnvRef(v) + } + } + + return opts, nil +} + +// parseMemory converts a human-readable memory string to bytes. +// Supported formats: "4GB", "4G", "4g", "512MB", "512M", "512m", "1024KB", "1024K", "1024". +func parseMemory(s string) (int64, error) { + if len(s) == 0 { + return 0, nil + } + + upper := strings.ToUpper(s) + var num string + var multiplier int64 + + switch { + case strings.HasSuffix(upper, "GB"): + num = s[:len(s)-2] + multiplier = 1 << 30 + case strings.HasSuffix(upper, "MB"): + num = s[:len(s)-2] + multiplier = 1 << 20 + case strings.HasSuffix(upper, "KB"): + num = s[:len(s)-2] + multiplier = 1 << 10 + case strings.HasSuffix(upper, "TB"): + num = s[:len(s)-2] + multiplier = 1 << 40 + case strings.HasSuffix(upper, "G"): + num = s[:len(s)-1] + multiplier = 1 << 30 + case strings.HasSuffix(upper, "M"): + num = s[:len(s)-1] + multiplier = 1 << 20 + case strings.HasSuffix(upper, "K"): + num = s[:len(s)-1] + multiplier = 1 << 10 + case strings.HasSuffix(upper, "T"): + num = s[:len(s)-1] + multiplier = 1 << 40 + default: + num = s + multiplier = 1 + } + + var val float64 + if _, err := fmt.Sscanf(num, "%f", &val); err != nil { + return 0, fmt.Errorf("invalid memory value %q", s) + } + return int64(val * float64(multiplier)), nil +} diff --git a/agent/sandbox/v2/prepare.go b/agent/sandbox/v2/prepare.go new file mode 100644 index 00000000..26e50843 --- /dev/null +++ b/agent/sandbox/v2/prepare.go @@ -0,0 +1,171 @@ +package sandboxv2 + +import ( + "context" + "fmt" + "log" + "path" + "strings" + + "github.com/yaoapp/yao/agent/sandbox/v2/types" + infra "github.com/yaoapp/yao/sandbox/v2" + "github.com/yaoapp/yao/tai/workspace" +) + +const onceMarkerDir = ".yao/prepare" + +// RunPrepareSteps executes a list of PrepareStep actions on the given Computer. +// file/copy/marker operations use computer.Workplace() (gRPC volume, cross-platform). +// exec operations use shell via Computer.Exec. +func RunPrepareSteps(ctx context.Context, steps []types.PrepareStep, computer infra.Computer, assistantID, configHash string) error { + if len(steps) == 0 { + return nil + } + + var ws workspace.FS + if computer != nil { + ws = computer.Workplace() + } + + markerDir := onceMarkerDir + if assistantID != "" { + markerDir = onceMarkerDir + "/" + assistantID + } + markerPath := markerDir + "/done" + + skipOnce := false + if configHash != "" && ws != nil { + if data, err := ws.ReadFile(markerPath); err == nil { + if strings.TrimSpace(string(data)) == configHash { + skipOnce = true + } + } + } + + for i, step := range steps { + if step.Once && skipOnce { + continue + } + + var err error + switch step.Action { + case "file": + err = runFileStep(ws, step) + case "copy": + err = runCopyStep(ws, step) + case "exec": + err = runExecStep(ctx, computer, step) + case "process": + log.Printf("[sandbox/v2] prepare step %d: action=process (reserved, skipping)", i) + default: + err = fmt.Errorf("unknown prepare action %q", step.Action) + } + + if err != nil { + if step.IgnoreError { + log.Printf("[sandbox/v2] prepare step %d (%s): ignored error: %v", i, step.Action, err) + continue + } + return fmt.Errorf("prepare step %d (%s): %w", i, step.Action, err) + } + } + + if configHash != "" && ws != nil { + ws.MkdirAll(markerDir, 0755) + ws.WriteFile(markerPath, []byte(configHash), 0644) + } + + return nil +} + +// --------------------------------------------------------------------------- +// Step runners +// --------------------------------------------------------------------------- + +func runFileStep(ws workspace.FS, step types.PrepareStep) error { + if step.Path == "" { + return fmt.Errorf("file step requires path") + } + if ws == nil { + return fmt.Errorf("file step requires workspace") + } + + dir := path.Dir(step.Path) + if dir != "." && dir != "/" { + if err := ws.MkdirAll(dir, 0755); err != nil { + return fmt.Errorf("mkdir %s: %w", dir, err) + } + } + + if err := ws.WriteFile(step.Path, step.Content, 0644); err != nil { + return fmt.Errorf("write file %s: %w", step.Path, err) + } + return nil +} + +func runCopyStep(ws workspace.FS, step types.PrepareStep) error { + if step.Src == "" || step.Dst == "" { + return fmt.Errorf("copy step requires src and dst") + } + if ws == nil { + return fmt.Errorf("copy step requires workspace") + } + + data, err := ws.ReadFile(step.Src) + if err != nil { + return fmt.Errorf("read src %s: %w", step.Src, err) + } + + dir := path.Dir(step.Dst) + if dir != "." && dir != "/" { + if err := ws.MkdirAll(dir, 0755); err != nil { + return fmt.Errorf("mkdir %s: %w", dir, err) + } + } + + if err := ws.WriteFile(step.Dst, data, 0644); err != nil { + return fmt.Errorf("write dst %s: %w", step.Dst, err) + } + return nil +} + +func runExecStep(ctx context.Context, computer infra.Computer, step types.PrepareStep) error { + if step.Cmd == "" { + return fmt.Errorf("exec step requires cmd") + } + + kind := shellFromSystem(computer) + script := step.Cmd + if step.Background { + if kind == shellSh { + script = fmt.Sprintf("nohup %s > /dev/null 2>&1 &", step.Cmd) + } else { + script = fmt.Sprintf("Start-Process -NoNewWindow -FilePath 'cmd.exe' -ArgumentList '/C %s'", step.Cmd) + } + } + + result, err := computer.Exec(ctx, shellWrap(kind, script), infra.WithWorkDir("/")) + if err != nil { + return err + } + label := "exec" + if step.Background { + label = "exec(background)" + } + return checkResult(result, label) +} + +// checkResult inspects ExecResult for errors. +func checkResult(result *infra.ExecResult, label string) error { + if result.Error != "" { + return fmt.Errorf("%s: %s", label, result.Error) + } + if result.ExitCode != 0 { + stderr := result.Stderr + if len(stderr) > 200 { + stderr = stderr[:200] + "..." + } + return fmt.Errorf("%s: exit %d: %s", label, result.ExitCode, stderr) + } + return nil +} diff --git a/agent/sandbox/v2/prepare_test.go b/agent/sandbox/v2/prepare_test.go new file mode 100644 index 00000000..b607a83e --- /dev/null +++ b/agent/sandbox/v2/prepare_test.go @@ -0,0 +1,549 @@ +package sandboxv2_test + +import ( + "context" + "fmt" + "strings" + "testing" + "time" + + sandboxv2 "github.com/yaoapp/yao/agent/sandbox/v2" + "github.com/yaoapp/yao/agent/sandbox/v2/types" +) + +// --------------------------------------------------------------------------- +// Box tests (local + remote) +// --------------------------------------------------------------------------- + +func TestRunPrepareSteps_Exec(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + box := createBox(t, m, nc) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + steps := []types.PrepareStep{ + {Action: "exec", Cmd: "echo hello > /tmp/prep-test"}, + {Action: "exec", Cmd: "echo world >> /tmp/prep-test"}, + } + + err := sandboxv2.RunPrepareSteps(ctx, steps, box, "test-assistant", "") + if err != nil { + t.Fatalf("RunPrepareSteps: %v", err) + } + + result, err := box.Exec(ctx, []string{"cat", "/tmp/prep-test"}) + if err != nil { + t.Fatalf("cat: %v", err) + } + got := strings.TrimSpace(result.Stdout) + if got != "hello\nworld" { + t.Errorf("content = %q, want %q", got, "hello\nworld") + } + }) + } +} + +func TestRunPrepareSteps_File(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + box := createBox(t, m, nc) + wsID := fmt.Sprintf("test-file-%d", time.Now().UnixNano()) + box.BindWorkplace(wsID) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + steps := []types.PrepareStep{ + {Action: "file", Path: "config/test.txt", Content: []byte("file-content-v2")}, + } + + err := sandboxv2.RunPrepareSteps(ctx, steps, box, "test-assistant", "") + if err != nil { + t.Fatalf("RunPrepareSteps: %v", err) + } + + ws := box.Workplace() + data, err := ws.ReadFile("config/test.txt") + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + if string(data) != "file-content-v2" { + t.Errorf("content = %q, want %q", string(data), "file-content-v2") + } + }) + } +} + +func TestRunPrepareSteps_Copy(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + box := createBox(t, m, nc) + wsID := fmt.Sprintf("test-copy-%d", time.Now().UnixNano()) + box.BindWorkplace(wsID) + + ws := box.Workplace() + ws.WriteFile("src.txt", []byte("copy-src"), 0644) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + steps := []types.PrepareStep{ + {Action: "copy", Src: "src.txt", Dst: "dst.txt"}, + } + + err := sandboxv2.RunPrepareSteps(ctx, steps, box, "test-assistant", "") + if err != nil { + t.Fatalf("RunPrepareSteps: %v", err) + } + + data, err := ws.ReadFile("dst.txt") + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + if string(data) != "copy-src" { + t.Errorf("content = %q, want %q", string(data), "copy-src") + } + }) + } +} + +func TestRunPrepareSteps_OnceMarker(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + box := createBox(t, m, nc) + wsID := fmt.Sprintf("test-once-%d", time.Now().UnixNano()) + box.BindWorkplace(wsID) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + counter := "/tmp/once-counter" + steps := []types.PrepareStep{ + {Action: "exec", Cmd: "echo -n x >> " + counter, Once: true}, + } + + hash := "abc123" + assistantID := "test-once" + + if err := sandboxv2.RunPrepareSteps(ctx, steps, box, assistantID, hash); err != nil { + t.Fatalf("first run: %v", err) + } + r1, _ := box.Exec(ctx, []string{"cat", counter}) + if r1.Stdout != "x" { + t.Fatalf("first run: got %q, want %q", r1.Stdout, "x") + } + + if err := sandboxv2.RunPrepareSteps(ctx, steps, box, assistantID, hash); err != nil { + t.Fatalf("second run: %v", err) + } + r2, _ := box.Exec(ctx, []string{"cat", counter}) + if r2.Stdout != "x" { + t.Errorf("second run: got %q, want %q (once step should be skipped)", r2.Stdout, "x") + } + + if err := sandboxv2.RunPrepareSteps(ctx, steps, box, assistantID, "new-hash"); err != nil { + t.Fatalf("third run: %v", err) + } + r3, _ := box.Exec(ctx, []string{"cat", counter}) + if r3.Stdout != "xx" { + t.Errorf("third run: got %q, want %q (hash changed, should re-execute)", r3.Stdout, "xx") + } + }) + } +} + +func TestRunPrepareSteps_OnceIsolation(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + box := createBox(t, m, nc) + wsID := fmt.Sprintf("test-iso-%d", time.Now().UnixNano()) + box.BindWorkplace(wsID) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + stepsA := []types.PrepareStep{ + {Action: "exec", Cmd: "echo -n A >> /tmp/iso-a", Once: true}, + } + stepsB := []types.PrepareStep{ + {Action: "exec", Cmd: "echo -n B >> /tmp/iso-b", Once: true}, + } + + hash := "same-hash" + + if err := sandboxv2.RunPrepareSteps(ctx, stepsA, box, "assistant-a", hash); err != nil { + t.Fatalf("assistant-a: %v", err) + } + if err := sandboxv2.RunPrepareSteps(ctx, stepsB, box, "assistant-b", hash); err != nil { + t.Fatalf("assistant-b: %v", err) + } + + rA, _ := box.Exec(ctx, []string{"cat", "/tmp/iso-a"}) + rB, _ := box.Exec(ctx, []string{"cat", "/tmp/iso-b"}) + if rA.Stdout != "A" { + t.Errorf("assistant-a: got %q, want %q", rA.Stdout, "A") + } + if rB.Stdout != "B" { + t.Errorf("assistant-b: got %q, want %q", rB.Stdout, "B") + } + + if err := sandboxv2.RunPrepareSteps(ctx, stepsA, box, "assistant-a", hash); err != nil { + t.Fatalf("assistant-a re-run: %v", err) + } + if err := sandboxv2.RunPrepareSteps(ctx, stepsB, box, "assistant-b", hash); err != nil { + t.Fatalf("assistant-b re-run: %v", err) + } + rA2, _ := box.Exec(ctx, []string{"cat", "/tmp/iso-a"}) + rB2, _ := box.Exec(ctx, []string{"cat", "/tmp/iso-b"}) + if rA2.Stdout != "A" { + t.Errorf("assistant-a re-run: got %q, want %q (should be skipped)", rA2.Stdout, "A") + } + if rB2.Stdout != "B" { + t.Errorf("assistant-b re-run: got %q, want %q (should be skipped)", rB2.Stdout, "B") + } + }) + } +} + +func TestRunPrepareSteps_IgnoreError(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + box := createBox(t, m, nc) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + steps := []types.PrepareStep{ + {Action: "exec", Cmd: "false", IgnoreError: true}, + {Action: "exec", Cmd: "echo survived > /tmp/survived"}, + } + + err := sandboxv2.RunPrepareSteps(ctx, steps, box, "test-assistant", "") + if err != nil { + t.Fatalf("RunPrepareSteps: %v (ignore_error should have prevented failure)", err) + } + + result, _ := box.Exec(ctx, []string{"cat", "/tmp/survived"}) + if strings.TrimSpace(result.Stdout) != "survived" { + t.Errorf("second step should have executed, got %q", result.Stdout) + } + }) + } +} + +func TestRunPrepareSteps_FailOnError(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + box := createBox(t, m, nc) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + steps := []types.PrepareStep{ + {Action: "exec", Cmd: "false"}, + {Action: "exec", Cmd: "echo should-not-reach > /tmp/unreachable"}, + } + + err := sandboxv2.RunPrepareSteps(ctx, steps, box, "test-assistant", "") + if err == nil { + t.Fatal("expected error from failing step without ignore_error") + } + + result, _ := box.Exec(ctx, []string{"cat", "/tmp/unreachable"}) + if result.ExitCode == 0 { + t.Error("second step should not have executed") + } + }) + } +} + +func TestRunPrepareSteps_UnknownAction(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + _ = createBox(t, m, nc) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + steps := []types.PrepareStep{ + {Action: "unknown_action"}, + } + + err := sandboxv2.RunPrepareSteps(ctx, steps, nil, "test-assistant", "") + if err == nil { + t.Fatal("expected error for unknown action") + } + if !strings.Contains(err.Error(), "unknown_action") { + t.Errorf("error should mention action name, got: %v", err) + } + }) + } +} + +func TestRunPrepareSteps_EmptySteps(t *testing.T) { + err := sandboxv2.RunPrepareSteps(context.Background(), nil, nil, "test-assistant", "hash") + if err != nil { + t.Fatalf("empty steps should succeed: %v", err) + } +} + +func TestRunPrepareSteps_Background(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + box := createBox(t, m, nc) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + steps := []types.PrepareStep{ + {Action: "exec", Cmd: "sleep 30", Background: true}, + {Action: "exec", Cmd: "echo after-bg > /tmp/after-bg"}, + } + + err := sandboxv2.RunPrepareSteps(ctx, steps, box, "test-assistant", "") + if err != nil { + t.Fatalf("RunPrepareSteps: %v", err) + } + + result, _ := box.Exec(ctx, []string{"cat", "/tmp/after-bg"}) + if strings.TrimSpace(result.Stdout) != "after-bg" { + t.Errorf("background step blocked execution, got %q", result.Stdout) + } + }) + } +} + +func TestRunPrepareSteps_MixedActions(t *testing.T) { + skipIfNoDocker(t) + + for _, nc := range boxNodes() { + nc := nc + t.Run(nc.Name, func(t *testing.T) { + m := setupManager(t, &nc) + box := createBox(t, m, nc) + wsID := fmt.Sprintf("test-mixed-%d", time.Now().UnixNano()) + box.BindWorkplace(wsID) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + steps := []types.PrepareStep{ + {Action: "file", Path: "mixed.conf", Content: []byte("key=value")}, + {Action: "exec", Cmd: "echo exec-ok > /tmp/mixed-exec"}, + {Action: "copy", Src: "mixed.conf", Dst: "mixed-copy.conf"}, + } + + err := sandboxv2.RunPrepareSteps(ctx, steps, box, "test-assistant", "") + if err != nil { + t.Fatalf("RunPrepareSteps: %v", err) + } + + ws := box.Workplace() + data, err := ws.ReadFile("mixed-copy.conf") + if err != nil { + t.Fatalf("ReadFile mixed-copy.conf: %v", err) + } + if string(data) != "key=value" { + t.Errorf("copy result: got %q, want %q", string(data), "key=value") + } + + result, _ := box.Exec(ctx, []string{"cat", "/tmp/mixed-exec"}) + if strings.TrimSpace(result.Stdout) != "exec-ok" { + t.Errorf("exec result: got %q, want %q", result.Stdout, "exec-ok") + } + }) + } +} + +// --------------------------------------------------------------------------- +// HostExec tests +// --------------------------------------------------------------------------- + +func TestRunPrepareSteps_HostExec(t *testing.T) { + skipIfNoHostExec(t) + + for _, tgt := range hostTargets() { + tgt := tgt + t.Run(tgt.Name, func(t *testing.T) { + m := setupHostManager(t, &tgt) + host := createHost(t, m, tgt) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + t.Logf("SystemInfo: OS=%q Shell=%q TempDir=%q", + host.ComputerInfo().System.OS, + host.ComputerInfo().System.Shell, + host.ComputerInfo().System.TempDir) + + isWin := tgt.Name == "win-native" + var cmd string + if isWin { + cmd = `Write-Output 'host-ok'` + } else { + cmd = "echo host-ok" + } + steps := []types.PrepareStep{ + {Action: "exec", Cmd: cmd}, + } + + err := sandboxv2.RunPrepareSteps(ctx, steps, host, "test-host", "") + if err != nil { + t.Fatalf("RunPrepareSteps on host: %v", err) + } + }) + } +} + +func TestRunPrepareSteps_HostExecFile(t *testing.T) { + skipIfNoHostExec(t) + + for _, tgt := range hostTargets() { + tgt := tgt + t.Run(tgt.Name, func(t *testing.T) { + m := setupHostManager(t, &tgt) + host := createHost(t, m, tgt) + wsID := fmt.Sprintf("test-hostfile-%d", time.Now().UnixNano()) + host.BindWorkplace(wsID) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + steps := []types.PrepareStep{ + {Action: "file", Path: "host-test.txt", Content: []byte("host-file-data")}, + } + + err := sandboxv2.RunPrepareSteps(ctx, steps, host, "test-host", "") + if err != nil { + t.Fatalf("RunPrepareSteps file: %v", err) + } + + ws := host.Workplace() + data, err := ws.ReadFile("host-test.txt") + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + if string(data) != "host-file-data" { + t.Errorf("content = %q, want %q", string(data), "host-file-data") + } + }) + } +} + +func TestRunPrepareSteps_HostExecCopy(t *testing.T) { + skipIfNoHostExec(t) + + for _, tgt := range hostTargets() { + tgt := tgt + t.Run(tgt.Name, func(t *testing.T) { + m := setupHostManager(t, &tgt) + host := createHost(t, m, tgt) + wsID := fmt.Sprintf("test-hostcopy-%d", time.Now().UnixNano()) + host.BindWorkplace(wsID) + + ws := host.Workplace() + ws.WriteFile("copy-src.txt", []byte("copy-data"), 0644) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + steps := []types.PrepareStep{ + {Action: "copy", Src: "copy-src.txt", Dst: "copy-dst.txt"}, + } + + err := sandboxv2.RunPrepareSteps(ctx, steps, host, "test-host", "") + if err != nil { + t.Fatalf("RunPrepareSteps copy: %v", err) + } + + data, err := ws.ReadFile("copy-dst.txt") + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + if string(data) != "copy-data" { + t.Errorf("content = %q, want %q", string(data), "copy-data") + } + }) + } +} + +func TestRunPrepareSteps_HostExecOnce(t *testing.T) { + skipIfNoHostExec(t) + + for _, tgt := range hostTargets() { + tgt := tgt + t.Run(tgt.Name, func(t *testing.T) { + m := setupHostManager(t, &tgt) + host := createHost(t, m, tgt) + wsID := fmt.Sprintf("test-hostonce-%d", time.Now().UnixNano()) + host.BindWorkplace(wsID) + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + isWin := tgt.Name == "win-native" + var cmd string + if isWin { + cmd = `Write-Output 'once-ok'` + } else { + cmd = "echo once-ok" + } + steps := []types.PrepareStep{ + {Action: "exec", Cmd: cmd, Once: true}, + } + hash := "host-once-hash" + aid := "host-once-aid" + + if err := sandboxv2.RunPrepareSteps(ctx, steps, host, aid, hash); err != nil { + t.Fatalf("first run: %v", err) + } + + ws := host.Workplace() + markerData, err := ws.ReadFile(".yao/prepare/" + aid + "/done") + if err != nil { + t.Fatalf("marker not written: %v", err) + } + if string(markerData) != hash { + t.Errorf("marker = %q, want %q", string(markerData), hash) + } + }) + } +} diff --git a/agent/sandbox/v2/runner.go b/agent/sandbox/v2/runner.go new file mode 100644 index 00000000..b2132607 --- /dev/null +++ b/agent/sandbox/v2/runner.go @@ -0,0 +1,32 @@ +package sandboxv2 + +import ( + "fmt" + "sync" + + "github.com/yaoapp/yao/agent/sandbox/v2/types" +) + +var ( + mu sync.RWMutex + runners = map[string]func() types.Runner{} +) + +// Register adds a runner factory to the global registry. +// Typically called from init() in the runner's package. +func Register(name string, factory func() types.Runner) { + mu.Lock() + defer mu.Unlock() + runners[name] = factory +} + +// Get creates a new Runner instance from the registry. +func Get(name string) (types.Runner, error) { + mu.RLock() + defer mu.RUnlock() + factory, ok := runners[name] + if !ok { + return nil, fmt.Errorf("sandbox runner %q not registered", name) + } + return factory(), nil +} diff --git a/agent/sandbox/v2/shell.go b/agent/sandbox/v2/shell.go new file mode 100644 index 00000000..565ca09b --- /dev/null +++ b/agent/sandbox/v2/shell.go @@ -0,0 +1,47 @@ +package sandboxv2 + +import ( + "strings" + + infra "github.com/yaoapp/yao/sandbox/v2" +) + +// shellKind identifies which shell to use for command execution. +type shellKind int + +const ( + shellSh shellKind = iota // Unix: sh -c + shellPwsh // Windows: pwsh -NoProfile -Command + shellPS // Windows: powershell -NoProfile -Command + shellCmd // Windows: cmd.exe /C (last-resort fallback) +) + +// shellWrap returns the Exec command slice to run a script string. +func shellWrap(kind shellKind, script string) []string { + switch kind { + case shellPwsh: + return []string{"pwsh", "-NoProfile", "-Command", script} + case shellPS: + return []string{"powershell", "-NoProfile", "-Command", script} + case shellCmd: + return []string{"cmd.exe", "/C", script} + default: + return []string{"sh", "-c", script} + } +} + +// shellFromSystem resolves shellKind from ComputerInfo().System.Shell +// reported by the Tai node at registration time. +func shellFromSystem(computer infra.Computer) shellKind { + shell := strings.ToLower(computer.ComputerInfo().System.Shell) + switch shell { + case "pwsh": + return shellPwsh + case "powershell": + return shellPS + case "cmd.exe", "cmd": + return shellCmd + default: + return shellSh + } +} diff --git a/agent/sandbox/v2/stream.go b/agent/sandbox/v2/stream.go new file mode 100644 index 00000000..2066cac5 --- /dev/null +++ b/agent/sandbox/v2/stream.go @@ -0,0 +1,138 @@ +package sandboxv2 + +import ( + "context" + "errors" + "fmt" + "log" + "time" + + agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/output/message" + "github.com/yaoapp/yao/agent/sandbox/v2/types" + infra "github.com/yaoapp/yao/sandbox/v2" +) + +// ExecuteRequest consolidates all parameters for ExecuteSandboxStream. +type ExecuteRequest struct { + Computer infra.Computer + Runner types.Runner + Config *types.SandboxConfig + StreamReq *types.StreamRequest + Manager *infra.Manager +} + +// ExecuteSandboxStream is the V2 replacement for executeSandboxStream. +// It calls runner.Stream, handles interrupts, and performs cleanup/lifecycle +// in defer. +func ExecuteSandboxStream( + ctx *agentContext.Context, + req *ExecuteRequest, + handler message.StreamFunc, +) (*agentContext.CompletionResponse, error) { + + if req.Runner == nil || req.Computer == nil { + return nil, fmt.Errorf("runner and computer are required") + } + + 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) + defer cancelRunner() // Prevent goroutine leak. + + done := make(chan struct{}) + defer close(done) + + go func() { + ticker := time.NewTicker(500 * time.Millisecond) + defer ticker.Stop() + for { + select { + case <-done: + return + case <-ticker.C: + if ctx.Interrupt != nil { + if sig := ctx.Interrupt.Peek(); sig != nil { + cancelRunner() + return + } + if ctx.Interrupt.IsInterrupted() { + cancelRunner() + return + } + } + case <-stdCtx.Done(): + cancelRunner() + return + } + } + }() + + var textContent []byte + wrappedHandler := func(chunkType message.StreamChunkType, data []byte) int { + if chunkType == message.ChunkText { + textContent = append(textContent, data...) + } + if handler != nil { + return handler(chunkType, data) + } + return 0 + } + + err := req.Runner.Stream(runnerCtx, req.StreamReq, wrappedHandler) + + panicked = false // Normal exit reached. + + if err != nil { + if errors.Is(err, context.Canceled) { + return nil, err + } + return nil, fmt.Errorf("runner.Stream: %w", err) + } + + resp := &agentContext.CompletionResponse{ + Role: "assistant", + FinishReason: agentContext.FinishReasonStop, + } + if len(textContent) > 0 { + resp.Content = string(textContent) + } + return resp, nil +} diff --git a/agent/sandbox/v2/testutils/testutils.go b/agent/sandbox/v2/testutils/testutils.go new file mode 100644 index 00000000..7e51a105 --- /dev/null +++ b/agent/sandbox/v2/testutils/testutils.go @@ -0,0 +1,47 @@ +package testutils + +import ( + "context" + "os" + "path/filepath" + "testing" + + agenttestutils "github.com/yaoapp/yao/agent/testutils" + "github.com/yaoapp/yao/config" + sandboxv2 "github.com/yaoapp/yao/sandbox/v2" + "github.com/yaoapp/yao/tai" + "github.com/yaoapp/yao/tai/registry" +) + +// Prepare initializes the full environment required for sandbox V2 E2E tests: +// - agent layer (assistants, LLM, caller) +// - tai registry + local node +// - sandbox V2 manager +func Prepare(t *testing.T) { + t.Helper() + + agenttestutils.Prepare(t) + + if registry.Global() == nil { + registry.Init(nil) + } + + dataDir := filepath.Join(config.Conf.DataRoot, "workspaces") + os.MkdirAll(dataDir, 0755) + tai.RegisterLocal(tai.WithDataDir(dataDir)) + + sandboxv2.Init() + if err := sandboxv2.M().Start(context.Background()); err != nil { + t.Fatalf("sandbox v2 manager start: %v", err) + } + + t.Cleanup(func() { + sandboxv2.M().Close() + }) +} + +// Clean tears down the test environment. +func Clean(t *testing.T) { + t.Helper() + agenttestutils.Clean(t) +} diff --git a/agent/sandbox/v2/testutils_test.go b/agent/sandbox/v2/testutils_test.go new file mode 100644 index 00000000..1ee2c02b --- /dev/null +++ b/agent/sandbox/v2/testutils_test.go @@ -0,0 +1,228 @@ +package sandboxv2_test + +import ( + "context" + "fmt" + "log" + "os" + "strconv" + "strings" + "testing" + "time" + + sandbox "github.com/yaoapp/yao/sandbox/v2" + "github.com/yaoapp/yao/tai" + "github.com/yaoapp/yao/tai/registry" + taisandbox "github.com/yaoapp/yao/tai/sandbox" + "github.com/yaoapp/yao/workspace" +) + +// --------------------------------------------------------------------------- +// node configuration — mirrors sandbox/v2 testutils but scoped to prepare tests +// --------------------------------------------------------------------------- + +type nodeConfig struct { + Name string + Addr string + TaiID string + Options []tai.Option +} + +type hostTarget struct { + Name string + Addr string + TaiID string +} + +// --------------------------------------------------------------------------- +// environment helpers (same conventions as sandbox/v2 + env.local.sh) +// --------------------------------------------------------------------------- + +func testLocalAddr() string { + if addr := os.Getenv("SANDBOX_TEST_LOCAL_ADDR"); addr != "" { + return addr + } + return "local" +} + +func testImage() string { + if img := os.Getenv("SANDBOX_TEST_IMAGE"); img != "" { + return img + } + return "alpine:latest" +} + +func envPort(key string, fallback int) int { + if v := os.Getenv(key); v != "" { + if p, err := strconv.Atoi(v); err == nil { + return p + } + } + return fallback +} + +// --------------------------------------------------------------------------- +// node discovery +// --------------------------------------------------------------------------- + +func boxNodes() []nodeConfig { + nodes := []nodeConfig{ + {Name: "local", Addr: testLocalAddr()}, + } + if addr := os.Getenv("SANDBOX_TEST_REMOTE_ADDR"); addr != "" { + nodes = append(nodes, nodeConfig{Name: "remote", Addr: addr}) + } + return nodes +} + +func hostTargets() []hostTarget { + var targets []hostTarget + if addr := os.Getenv("TAI_TEST_WIN_HOSTEXEC_LINUX"); addr != "" { + targets = append(targets, hostTarget{Name: "win-linux", Addr: addr}) + } + if addr := os.Getenv("TAI_TEST_WIN_HOSTEXEC_NATIVE"); addr != "" { + targets = append(targets, hostTarget{Name: "win-native", Addr: addr}) + } + return targets +} + +// --------------------------------------------------------------------------- +// TestMain — purge stale containers from previous runs +// --------------------------------------------------------------------------- + +func TestMain(m *testing.M) { + purgeStale() + os.Exit(m.Run()) +} + +func purgeStale() { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + for _, nc := range boxNodes() { + client, err := tai.New(nc.Addr, nc.Options...) + if err != nil { + continue + } + sb := client.Sandbox() + if sb == nil { + client.Close() + continue + } + containers, _ := sb.List(ctx, taisandbox.ListOptions{All: true}) + for _, c := range containers { + id := c.Name + if id == "" { + id = c.ID + } + if strings.HasPrefix(id, "sb-prep-") || strings.HasPrefix(id, "sb-lc-") { + sb.Remove(ctx, id, true) + log.Printf("[purge] %s: removed %s", nc.Name, id) + } + } + client.Close() + } +} + +// --------------------------------------------------------------------------- +// Manager + Box helpers +// --------------------------------------------------------------------------- + +func setupManager(t *testing.T, nc *nodeConfig) *sandbox.Manager { + t.Helper() + if registry.Global() == nil { + registry.Init(nil) + } + client, err := tai.New(nc.Addr, nc.Options...) + if err != nil { + t.Fatalf("tai.New(%s): %v", nc.Addr, err) + } + nc.TaiID = client.TaiID() + + sandbox.Init() + m := sandbox.M() + t.Cleanup(func() { m.Close() }) + return m +} + +func createBox(t *testing.T, m *sandbox.Manager, nc nodeConfig) *sandbox.Box { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + + if err := m.EnsureImage(ctx, nc.TaiID, testImage(), sandbox.ImagePullOptions{}); err != nil { + t.Fatalf("EnsureImage: %v", err) + } + + box, err := m.Create(ctx, sandbox.CreateOptions{ + ID: fmt.Sprintf("sb-prep-%d", time.Now().UnixNano()), + Image: testImage(), + Owner: "test-prepare", + NodeID: nc.TaiID, + }) + if err != nil { + t.Fatalf("Create: %v", err) + } + t.Cleanup(func() { + cCtx, cCancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cCancel() + if err := m.Remove(cCtx, box.ID()); err != nil { + t.Logf("cleanup Remove(%s): %v", box.ID(), err) + } + }) + return box +} + +func createHost(t *testing.T, m *sandbox.Manager, tgt hostTarget) *sandbox.Host { + t.Helper() + host, err := m.Host(context.Background(), tgt.TaiID) + if err != nil { + t.Skipf("Host(%s): %v", tgt.Name, err) + } + return host +} + +func setupHostManager(t *testing.T, tgt *hostTarget) *sandbox.Manager { + t.Helper() + nc := nodeConfig{Name: tgt.Name, Addr: fmt.Sprintf("tai://%s", tgt.Addr)} + m := setupManager(t, &nc) + tgt.TaiID = nc.TaiID + return m +} + +// --------------------------------------------------------------------------- +// skip helpers +// --------------------------------------------------------------------------- + +func skipIfNoDocker(t *testing.T) { + t.Helper() + if testLocalAddr() == "" { + t.Skip("SANDBOX_TEST_LOCAL_ADDR not set") + } +} + +func skipIfNoHostExec(t *testing.T) { + t.Helper() + if len(hostTargets()) == 0 { + t.Skip("no HostExec targets configured") + } +} + +func createTestWorkspace(t *testing.T, taiID, wsID string) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + _, err := workspace.M().Create(ctx, workspace.CreateOptions{ + ID: wsID, + Owner: "test", + Node: taiID, + }) + if err != nil && !strings.Contains(err.Error(), "exists") { + t.Fatalf("create workspace %q: %v", wsID, err) + } + t.Cleanup(func() { + cCtx, cCancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cCancel() + workspace.M().Delete(cCtx, wsID, true) + }) +} diff --git a/agent/sandbox/v2/types/config.go b/agent/sandbox/v2/types/config.go new file mode 100644 index 00000000..2f5a57af --- /dev/null +++ b/agent/sandbox/v2/types/config.go @@ -0,0 +1,137 @@ +package types + +import ( + "encoding/json" + "fmt" +) + +const ( + SandboxVersionV1 = "1.0" + SandboxVersionV2 = "2.0" +) + +// SandboxConfig is the V2 sandbox configuration loaded from sandbox.yao or +// the package.yao "sandbox" block when version == "2.0". +type SandboxConfig struct { + Version string `json:"version" yaml:"version"` + Computer ComputerConfig `json:"computer" yaml:"computer"` + Runner RunnerConfig `json:"runner" yaml:"runner"` + Lifecycle string `json:"lifecycle,omitempty" yaml:"lifecycle,omitempty"` + IdleTimeout string `json:"idle_timeout,omitempty" yaml:"idle_timeout,omitempty"` + MaxLifetime string `json:"max_lifetime,omitempty" yaml:"max_lifetime,omitempty"` + StopTimeout string `json:"stop_timeout,omitempty" yaml:"stop_timeout,omitempty"` + Prepare []PrepareStep `json:"prepare,omitempty" yaml:"prepare,omitempty"` + Environment map[string]string `json:"environment,omitempty" yaml:"environment,omitempty"` + Secrets map[string]string `json:"secrets,omitempty" yaml:"secrets,omitempty"` + + // Populated by the framework at runtime (never serialized). + Owner string `json:"-" yaml:"-"` + ID string `json:"-" yaml:"-"` + Labels map[string]string `json:"-" yaml:"-"` + NodeID string `json:"-" yaml:"-"` + Kind string `json:"-" yaml:"-"` + WorkspaceID string `json:"-" yaml:"-"` +} + +// ComputerConfig describes the execution environment (container or host). +type ComputerConfig struct { + Image string `json:"image,omitempty" yaml:"image,omitempty"` + VNC VNCConfig `json:"vnc,omitempty" yaml:"vnc,omitempty"` + Memory string `json:"memory,omitempty" yaml:"memory,omitempty"` + CPUs float64 `json:"cpus,omitempty" yaml:"cpus,omitempty"` + Ports PortList `json:"ports,omitempty" yaml:"ports,omitempty"` + User string `json:"user,omitempty" yaml:"user,omitempty"` + WorkDir string `json:"work_dir,omitempty" yaml:"work_dir,omitempty"` + MountPath string `json:"mount_path,omitempty" yaml:"mount_path,omitempty"` + MountMode string `json:"mount_mode,omitempty" yaml:"mount_mode,omitempty"` +} + +// RunnerConfig identifies which Runner to use and how. +type RunnerConfig struct { + Name string `json:"name" yaml:"name"` + Mode string `json:"mode,omitempty" yaml:"mode,omitempty"` + Options map[string]any `json:"options,omitempty" yaml:"options,omitempty"` +} + +// PrepareStep is a single action executed during Runner.Prepare. +type PrepareStep struct { + Action string `json:"action" yaml:"action"` + Once bool `json:"once,omitempty" yaml:"once,omitempty"` + IgnoreError bool `json:"ignore_error,omitempty" yaml:"ignore_error,omitempty"` + + // action=copy + Src string `json:"src,omitempty" yaml:"src,omitempty"` + Dst string `json:"dst,omitempty" yaml:"dst,omitempty"` + + // action=exec + Cmd string `json:"cmd,omitempty" yaml:"cmd,omitempty"` + Background bool `json:"background,omitempty" yaml:"background,omitempty"` + + // action=file (internal use by Runner.Prepare) + Path string `json:"path,omitempty" yaml:"path,omitempty"` + Content []byte `json:"-" yaml:"-"` + + // action=process (reserved) + Name string `json:"name,omitempty" yaml:"name,omitempty"` + Args []any `json:"args,omitempty" yaml:"args,omitempty"` +} + +// --------------------------------------------------------------------------- +// VNCConfig — supports both bool and object in JSON/YAML: +// true → VNCConfig{Enabled: true} +// {"enabled": true, "password": "xxx"} → full struct +// --------------------------------------------------------------------------- + +type VNCConfig struct { + Enabled bool `json:"enabled,omitempty" yaml:"enabled,omitempty"` + ViewOnly bool `json:"view_only,omitempty" yaml:"view_only,omitempty"` + Password string `json:"password,omitempty" yaml:"password,omitempty"` + Resolution string `json:"resolution,omitempty" yaml:"resolution,omitempty"` +} + +func (v *VNCConfig) UnmarshalJSON(data []byte) error { + var b bool + if err := json.Unmarshal(data, &b); err == nil { + v.Enabled = b + return nil + } + type alias VNCConfig + var a alias + if err := json.Unmarshal(data, &a); err != nil { + return err + } + *v = VNCConfig(a) + return nil +} + +// --------------------------------------------------------------------------- +// PortList — supports both int array and object array in JSON: +// [3000, 8080] → []PortMapping{{Port: 3000}, {Port: 8080}} +// [{"port": 3000, "host_port": 9000}] → full structs +// --------------------------------------------------------------------------- + +type PortList []PortMapping + +type PortMapping struct { + Port int `json:"port" yaml:"port"` + HostPort int `json:"host_port,omitempty" yaml:"host_port,omitempty"` + Protocol string `json:"protocol,omitempty" yaml:"protocol,omitempty"` +} + +func (p *PortList) UnmarshalJSON(data []byte) error { + var ints []int + if err := json.Unmarshal(data, &ints); err == nil { + out := make(PortList, len(ints)) + for i, port := range ints { + out[i] = PortMapping{Port: port} + } + *p = out + return nil + } + var objs []PortMapping + if err := json.Unmarshal(data, &objs); err != nil { + return fmt.Errorf("ports: expected int array or object array: %w", err) + } + *p = objs + return nil +} diff --git a/agent/sandbox/v2/types/runner.go b/agent/sandbox/v2/types/runner.go new file mode 100644 index 00000000..5a90f80f --- /dev/null +++ b/agent/sandbox/v2/types/runner.go @@ -0,0 +1,53 @@ +package types + +import ( + "context" + + "github.com/yaoapp/gou/connector" + agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/agent/output/message" + infra "github.com/yaoapp/yao/sandbox/v2" +) + +// Runner is the interface that all sandbox runners must implement. +// A Runner replaces the LLM invocation layer (executeLLMStream) when a +// sandbox is configured. +type Runner interface { + Name() string + Prepare(ctx context.Context, req *PrepareRequest) error + Stream(ctx context.Context, req *StreamRequest, handler message.StreamFunc) error + Cleanup(ctx context.Context, computer infra.Computer) error +} + +// MCPServer mirrors store/types.MCPServerConfig to avoid a cyclic import +// between this leaf package and agent/store/types. +type MCPServer struct { + ServerID string `json:"server_id,omitempty"` + Resources []string `json:"resources,omitempty"` + Tools []string `json:"tools,omitempty"` +} + +// RunStepsFunc is the signature of RunPrepareSteps. Workspace is obtained +// internally via computer.Workplace(). +type RunStepsFunc func(ctx context.Context, steps []PrepareStep, computer infra.Computer, assistantID, configHash string) error + +// PrepareRequest carries everything needed by Runner.Prepare. +type PrepareRequest struct { + Computer infra.Computer + Config *SandboxConfig + Connector connector.Connector + SkillsDir string + MCPServers []MCPServer + ConfigHash string + RunSteps RunStepsFunc +} + +// StreamRequest carries everything needed by Runner.Stream. +type StreamRequest struct { + Computer infra.Computer + Config *SandboxConfig + Connector connector.Connector + Messages []agentContext.Message + SystemPrompt string + ChatID string +} diff --git a/agent/sandbox/v2/types/token.go b/agent/sandbox/v2/types/token.go new file mode 100644 index 00000000..7f1cf525 --- /dev/null +++ b/agent/sandbox/v2/types/token.go @@ -0,0 +1,9 @@ +package types + +import "time" + +// SandboxToken is a short-lived JWT issued for a sandbox computer. +type SandboxToken struct { + Token string + ExpiresAt time.Time +} diff --git a/agent/sandbox/v2/yao/runner.go b/agent/sandbox/v2/yao/runner.go new file mode 100644 index 00000000..7bbd7109 --- /dev/null +++ b/agent/sandbox/v2/yao/runner.go @@ -0,0 +1,38 @@ +package yao + +import ( + "context" + + "github.com/yaoapp/yao/agent/output/message" + "github.com/yaoapp/yao/agent/sandbox/v2/types" + infra "github.com/yaoapp/yao/sandbox/v2" +) + +// YaoRunner is a no-op Runner for pure Hook-driven sandbox interactions. +// When runner.name == "yao", the assistant relies entirely on Create/Next +// hooks for logic; no external CLI is invoked. +type YaoRunner struct{} + +func New() *YaoRunner { return &YaoRunner{} } + +func (r *YaoRunner) Name() string { return "yao" } + +// Prepare runs user-defined prepare steps (copy, exec, file) but adds +// no runner-specific steps. Connector is not required. +func (r *YaoRunner) Prepare(ctx context.Context, req *types.PrepareRequest) error { + if req.RunSteps != nil && len(req.Config.Prepare) > 0 { + return req.RunSteps(ctx, req.Config.Prepare, req.Computer, req.Config.ID, req.ConfigHash) + } + return nil +} + +// Stream is a no-op — hooks handle all interaction. Returns immediately +// so the assistant framework proceeds to the Next hook. +func (r *YaoRunner) Stream(_ context.Context, _ *types.StreamRequest, _ message.StreamFunc) error { + return nil +} + +// Cleanup is a no-op for the yao runner. +func (r *YaoRunner) Cleanup(_ context.Context, _ infra.Computer) error { + return nil +} diff --git a/agent/sandbox/v2/yao/runner_test.go b/agent/sandbox/v2/yao/runner_test.go new file mode 100644 index 00000000..2f7e4201 --- /dev/null +++ b/agent/sandbox/v2/yao/runner_test.go @@ -0,0 +1,137 @@ +package yao_test + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/yaoapp/yao/agent/caller" + agentcontext "github.com/yaoapp/yao/agent/context" + sandboxtestutils "github.com/yaoapp/yao/agent/sandbox/v2/testutils" + oauthtypes "github.com/yaoapp/yao/openapi/oauth/types" +) + +func TestSandboxV2_Yao_JSAPI(t *testing.T) { + sandboxtestutils.Prepare(t) + defer sandboxtestutils.Clean(t) + + require.NotNil(t, caller.AgentGetterFunc, "AgentGetterFunc should be registered after Prepare") + + agent, err := caller.AgentGetterFunc("tests.sandbox-v2.jsapi-v2") + require.NoError(t, err, "should load assistant tests.sandbox-v2.jsapi-v2") + + chatID := fmt.Sprintf("e2e-jsapi-%d", time.Now().UnixMilli()) + ctx := agentcontext.New( + context.Background(), + &oauthtypes.AuthorizedInfo{ + TeamID: "test-team-jsapi", + UserID: "test-user-jsapi", + }, + chatID, + ) + + messages := []agentcontext.Message{ + {Role: "user", Content: "test jsapi"}, + } + + done := make(chan struct{}) + var resp *agentcontext.Response + var streamErr error + + go func() { + defer close(done) + resp, streamErr = agent.Stream(ctx, messages) + }() + + select { + case <-done: + case <-time.After(3 * time.Minute): + t.Fatalf("timeout after 3m") + } + + require.NoError(t, streamErr, "Stream should not return error") + require.NotNil(t, resp, "response should not be nil") + + // runner=yao goes through executeLLMStream, then Next hook returns { data: results } + // The Next hook result should appear in resp.Next + require.NotNil(t, resp.Next, "resp.Next should not be nil (Next hook returned data)") + t.Logf("resp.Next: %+v", resp.Next) + + nextData, ok := resp.Next.(map[string]interface{}) + if !ok { + t.Fatalf("resp.Next should be a map, got %T: %+v", resp.Next, resp.Next) + } + + // The Next hook returns { data: results }, the framework unwraps .data + data, hasData := nextData["data"] + if hasData { + nextData, ok = data.(map[string]interface{}) + require.True(t, ok, "data should be a map") + } + + t.Logf("JSAPI test results: %+v", nextData) + + // ── Verify ctx.computer was available ── + assert.Equal(t, true, nextData["has_computer"], "ctx.computer should be available") + assert.Equal(t, true, nextData["has_workspace"], "ctx.workspace should be available") + + // ── Verify ctx.computer.Info() ── + if infoRaw, ok := nextData["computer_info"]; ok { + info, ok := infoRaw.(map[string]interface{}) + require.True(t, ok, "computer_info should be a map") + assert.NotEmpty(t, info["kind"], "computer_info.kind should not be empty") + t.Logf("computer info: kind=%v os=%v", info["kind"], info["os"]) + } else { + assert.Nil(t, nextData["computer_info_error"], "computer.Info() should not error") + } + + // ── Verify ctx.computer.Exec() ── + assert.Equal(t, "jsapi-v2-test", nextData["exec_stdout"], "Exec should return expected stdout") + assert.Nil(t, nextData["exec_error"], "Exec should not error") + if exitCode, ok := nextData["exec_exit_code"]; ok { + // JS numbers come back as float64 through JSON + switch v := exitCode.(type) { + case float64: + assert.Equal(t, float64(0), v, "exit_code should be 0") + case int: + assert.Equal(t, 0, v, "exit_code should be 0") + } + } + + // ── Verify ctx.workspace write/read ── + assert.Equal(t, true, nextData["write_read_ok"], "workspace WriteFile+ReadFile round-trip should work") + assert.Equal(t, "hello from jsapi v2", nextData["read_content"], "read content should match") + assert.Nil(t, nextData["write_read_error"], "write/read should not error") + + // ── Verify ctx.workspace MkdirAll + Exists ── + assert.Equal(t, true, nextData["mkdir_exists_ok"], "MkdirAll + Exists should work") + assert.Nil(t, nextData["mkdir_exists_error"], "mkdir/exists should not error") + + // ── Verify ctx.workspace ReadDir ── + assert.Nil(t, nextData["readdir_error"], "ReadDir should not error") + if count, ok := nextData["readdir_count"]; ok { + switch v := count.(type) { + case float64: + assert.Greater(t, v, float64(0), "ReadDir should return entries") + } + } + + // ── Verify ctx.workspace Stat ── + assert.Equal(t, true, nextData["stat_ok"], "Stat should return correct info") + assert.Nil(t, nextData["stat_error"], "Stat should not error") + + // ── Verify ctx.workspace Copy ── + assert.Equal(t, true, nextData["copy_ok"], "Copy should work") + assert.Nil(t, nextData["copy_error"], "Copy should not error") + + // ── Verify ctx.workspace Rename ── + assert.Equal(t, true, nextData["rename_ok"], "Rename should work") + assert.Nil(t, nextData["rename_error"], "Rename should not error") + + // ── Verify ctx.workspace Remove ── + assert.Equal(t, true, nextData["remove_ok"], "Remove should work") + assert.Nil(t, nextData["remove_error"], "Remove should not error") +} diff --git a/agent/store/types/sandbox_v2.go b/agent/store/types/sandbox_v2.go new file mode 100644 index 00000000..88485574 --- /dev/null +++ b/agent/store/types/sandbox_v2.go @@ -0,0 +1,99 @@ +package types + +import ( + "crypto/sha256" + "fmt" + "os" + "path/filepath" + "sort" + "strings" + + jsoniter "github.com/json-iterator/go" + sandboxTypes "github.com/yaoapp/yao/agent/sandbox/v2/types" +) + +// LoadSandboxConfig reads a sandbox.yao file (JSON or YAML) and returns +// the V2 SandboxConfig. Called during Assistant.Load(). +func LoadSandboxConfig(filePath string) (*sandboxTypes.SandboxConfig, error) { + data, err := os.ReadFile(filePath) + if err != nil { + return nil, fmt.Errorf("read sandbox config %s: %w", filePath, err) + } + + ext := strings.ToLower(filepath.Ext(filePath)) + var cfg sandboxTypes.SandboxConfig + + switch ext { + case ".json", ".yao": + if err := jsoniter.Unmarshal(data, &cfg); err != nil { + return nil, fmt.Errorf("parse sandbox config (json): %w", err) + } + default: + if err := jsoniter.Unmarshal(data, &cfg); err != nil { + return nil, fmt.Errorf("parse sandbox config: %w", err) + } + } + + if cfg.Version != sandboxTypes.SandboxVersionV2 { + return nil, fmt.Errorf("sandbox.yao version must be %q, got %q", sandboxTypes.SandboxVersionV2, cfg.Version) + } + + return &cfg, nil +} + +// ToSandboxV2 converts a generic value (typically map[string]any from DSL +// parsing) into a V2 SandboxConfig. +func ToSandboxV2(v any) (*sandboxTypes.SandboxConfig, error) { + if v == nil { + return nil, nil + } + + switch sb := v.(type) { + case *sandboxTypes.SandboxConfig: + return sb, nil + case sandboxTypes.SandboxConfig: + return &sb, nil + default: + raw, err := jsoniter.Marshal(v) + if err != nil { + return nil, fmt.Errorf("sandbox v2 format error: %w", err) + } + var cfg sandboxTypes.SandboxConfig + if err := jsoniter.Unmarshal(raw, &cfg); err != nil { + return nil, fmt.Errorf("sandbox v2 format error: %w", err) + } + return &cfg, nil + } +} + +// ComputeConfigHash computes a SHA-256 fingerprint of the sandbox configuration, +// MCP servers, and skills directory. Used for hot-reload detection in prepare +// step "once" logic. +func ComputeConfigHash(cfg *sandboxTypes.SandboxConfig, mcpServers []MCPServerConfig, skillsDir string) string { + h := sha256.New() + + raw, _ := jsoniter.Marshal(cfg) + h.Write(raw) + + if len(mcpServers) > 0 { + mcpRaw, _ := jsoniter.Marshal(mcpServers) + h.Write(mcpRaw) + } + + if skillsDir != "" { + h.Write([]byte(skillsDir)) + entries, err := os.ReadDir(skillsDir) + if err == nil { + names := make([]string, 0, len(entries)) + for _, e := range entries { + names = append(names, e.Name()) + } + sort.Strings(names) + for _, n := range names { + h.Write([]byte(n)) + } + } + } + + return fmt.Sprintf("%x", h.Sum(nil)) +} diff --git a/agent/store/types/types.go b/agent/store/types/types.go index ea4ae8c1..4063b285 100644 --- a/agent/store/types/types.go +++ b/agent/store/types/types.go @@ -9,6 +9,7 @@ import ( "github.com/yaoapp/xun/dbal/query" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/i18n" + sandboxTypes "github.com/yaoapp/yao/agent/sandbox/v2/types" searchTypes "github.com/yaoapp/yao/agent/search/types" ) @@ -421,42 +422,44 @@ type ConnectorOptions struct { // AssistantModel the assistant database model type AssistantModel struct { - ID string `json:"assistant_id"` // Assistant ID - Type string `json:"type,omitempty"` // Assistant Type, default is assistant - Name string `json:"name,omitempty"` // Assistant Name - Avatar string `json:"avatar,omitempty"` // Assistant Avatar - Connector string `json:"connector"` // AI Connector (default connector) - ConnectorOptions *ConnectorOptions `json:"connector_options,omitempty"` // Connector selection options for user to choose from - Path string `json:"path,omitempty"` // Assistant Path - BuiltIn bool `json:"built_in,omitempty"` // Whether this is a built-in assistant - Sort int `json:"sort,omitempty"` // Assistant Sort - Description string `json:"description,omitempty"` // Assistant Description - Capabilities string `json:"capabilities,omitempty"` // Assistant capabilities description (useful for Robot orchestration) - Tags []string `json:"tags,omitempty"` // Assistant Tags - Modes []string `json:"modes,omitempty"` // Supported modes (e.g., ["task", "chat"]), null means all modes are supported - DefaultMode string `json:"default_mode,omitempty"` // Default mode, can be empty - Readonly bool `json:"readonly,omitempty"` // Whether this assistant is readonly - Public bool `json:"public,omitempty"` // Whether this assistant is shared across all teams in the platform - Share string `json:"share,omitempty"` // Assistant sharing scope (private/team) - Mentionable bool `json:"mentionable,omitempty"` // Whether this assistant is mentionable - Automated bool `json:"automated,omitempty"` // Whether this assistant is automated - Options map[string]interface{} `json:"options,omitempty"` // AI Options - Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts (default prompts) - PromptPresets map[string][]Prompt `json:"prompt_presets,omitempty"` // Prompt presets organized by mode (e.g., "chat", "task", etc.) - DisableGlobalPrompts bool `json:"disable_global_prompts,omitempty"` // Whether to disable global prompts, default is false - KB *KnowledgeBase `json:"kb,omitempty"` // Knowledge base configuration - DB *Database `json:"db,omitempty"` // Database configuration - MCP *MCPServers `json:"mcp,omitempty"` // MCP servers configuration - Workflow *Workflow `json:"workflow,omitempty"` // Workflow configuration - Sandbox *Sandbox `json:"sandbox,omitempty"` // Sandbox configuration for coding agents - Placeholder *Placeholder `json:"placeholder,omitempty"` // Assistant Placeholder - Source string `json:"source,omitempty"` // Hook script source code - Locales i18n.Map `json:"locales,omitempty"` // Assistant Locales - Uses *context.Uses `json:"uses,omitempty"` // Assistant-specific wrapper configurations for vision, audio, etc. If not set, use global settings - Search *searchTypes.Config `json:"search,omitempty"` // Search configuration (web, kb, db, citation, weights, etc.) - Dependencies map[string]string `json:"dependencies,omitempty"` // Dependencies on other MCP Clients (name -> version constraint) - CreatedAt int64 `json:"created_at"` // Creation timestamp - UpdatedAt int64 `json:"updated_at"` // Last update timestamp + ID string `json:"assistant_id"` // Assistant ID + Type string `json:"type,omitempty"` // Assistant Type, default is assistant + Name string `json:"name,omitempty"` // Assistant Name + Avatar string `json:"avatar,omitempty"` // Assistant Avatar + Connector string `json:"connector"` // AI Connector (default connector) + ConnectorOptions *ConnectorOptions `json:"connector_options,omitempty"` // Connector selection options for user to choose from + Path string `json:"path,omitempty"` // Assistant Path + BuiltIn bool `json:"built_in,omitempty"` // Whether this is a built-in assistant + Sort int `json:"sort,omitempty"` // Assistant Sort + Description string `json:"description,omitempty"` // Assistant Description + Capabilities string `json:"capabilities,omitempty"` // Assistant capabilities description (useful for Robot orchestration) + Tags []string `json:"tags,omitempty"` // Assistant Tags + Modes []string `json:"modes,omitempty"` // Supported modes (e.g., ["task", "chat"]), null means all modes are supported + DefaultMode string `json:"default_mode,omitempty"` // Default mode, can be empty + Readonly bool `json:"readonly,omitempty"` // Whether this assistant is readonly + Public bool `json:"public,omitempty"` // Whether this assistant is shared across all teams in the platform + Share string `json:"share,omitempty"` // Assistant sharing scope (private/team) + Mentionable bool `json:"mentionable,omitempty"` // Whether this assistant is mentionable + Automated bool `json:"automated,omitempty"` // Whether this assistant is automated + Options map[string]interface{} `json:"options,omitempty"` // AI Options + Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts (default prompts) + PromptPresets map[string][]Prompt `json:"prompt_presets,omitempty"` // Prompt presets organized by mode (e.g., "chat", "task", etc.) + DisableGlobalPrompts bool `json:"disable_global_prompts,omitempty"` // Whether to disable global prompts, default is false + KB *KnowledgeBase `json:"kb,omitempty"` // Knowledge base configuration + DB *Database `json:"db,omitempty"` // Database configuration + MCP *MCPServers `json:"mcp,omitempty"` // MCP servers configuration + Workflow *Workflow `json:"workflow,omitempty"` // Workflow configuration + Sandbox *Sandbox `json:"sandbox,omitempty"` // Sandbox configuration for coding agents (V1) + SandboxV2 *sandboxTypes.SandboxConfig `json:"-"` // V2 sandbox configuration (runtime only, not persisted in DB) + ConfigHash string `json:"-"` // V2 sandbox config fingerprint for hot-reload + Placeholder *Placeholder `json:"placeholder,omitempty"` // Assistant Placeholder + Source string `json:"source,omitempty"` // Hook script source code + Locales i18n.Map `json:"locales,omitempty"` // Assistant Locales + Uses *context.Uses `json:"uses,omitempty"` // Assistant-specific wrapper configurations for vision, audio, etc. If not set, use global settings + Search *searchTypes.Config `json:"search,omitempty"` // Search configuration (web, kb, db, citation, weights, etc.) + Dependencies map[string]string `json:"dependencies,omitempty"` // Dependencies on other MCP Clients (name -> version constraint) + CreatedAt int64 `json:"created_at"` // Creation timestamp + UpdatedAt int64 `json:"updated_at"` // Last update timestamp // Permission management fields (not exposed in JSON API responses) YaoCreatedBy string `json:"-"` // User who created the assistant (not exposed in JSON) diff --git a/sandbox/v2/box.go b/sandbox/v2/box.go index 67d9999d..5cf4223e 100644 --- a/sandbox/v2/box.go +++ b/sandbox/v2/box.go @@ -29,6 +29,7 @@ type Box struct { vnc bool image string workspaceID string + system SystemInfo ws workspace.FS manager *Manager } @@ -46,6 +47,7 @@ func (b *Box) ComputerInfo() ComputerInfo { return ComputerInfo{ Kind: "box", NodeID: b.nodeID, + System: b.system, Status: "online", BoxID: b.id, ContainerID: b.containerID, @@ -98,10 +100,6 @@ func (b *Box) Exec(ctx context.Context, cmd []string, opts ...ExecOption) (*Exec Stderr: result.Stderr, } - if b.policy == OneShot { - b.manager.Remove(ctx, b.id) - } - return r, nil } diff --git a/sandbox/v2/host.go b/sandbox/v2/host.go index 7a8a6df3..125388c3 100644 --- a/sandbox/v2/host.go +++ b/sandbox/v2/host.go @@ -18,6 +18,7 @@ import ( type Host struct { nodeID string workplaceID string + system SystemInfo manager *Manager } @@ -31,6 +32,7 @@ func (h *Host) ComputerInfo() ComputerInfo { return ComputerInfo{ Kind: "host", NodeID: h.nodeID, + System: h.system, Status: "online", } } diff --git a/sandbox/v2/manager.go b/sandbox/v2/manager.go index 1cdfc2cc..1c9659eb 100644 --- a/sandbox/v2/manager.go +++ b/sandbox/v2/manager.go @@ -98,7 +98,20 @@ func (m *Manager) Host(_ context.Context, nodeID string) (*Host, error) { return nil, fmt.Errorf("sandbox: node %q has no host_exec capability", nodeID) } - return &Host{nodeID: nodeID, manager: m}, nil + var sys SystemInfo + if snap, ok := tai.GetNodeSnapshot(nodeID); ok { + sys = SystemInfo{ + OS: snap.System.OS, + Arch: snap.System.Arch, + Hostname: snap.System.Hostname, + NumCPU: snap.System.NumCPU, + TotalMem: snap.System.TotalMem, + Shell: snap.System.Shell, + TempDir: snap.System.TempDir, + } + } + + return &Host{nodeID: nodeID, system: sys, manager: m}, nil } // Create creates and starts a new sandbox. @@ -113,9 +126,33 @@ func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Box, error) if wsm := workspace.M(); wsm != nil { node, err := wsm.NodeForWorkspace(ctx, opts.WorkspaceID) if err != nil { - return nil, fmt.Errorf("sandbox: resolve workspace %q: %w", opts.WorkspaceID, err) + targetNode := nodeID + if targetNode == "" { + if nodes := wsm.Nodes(); len(nodes) > 0 { + for _, n := range nodes { + if n.Online { + targetNode = n.Name + break + } + } + } + } + if targetNode == "" { + return nil, fmt.Errorf("sandbox: resolve workspace %q: no available node", opts.WorkspaceID) + } + _, err = wsm.Create(ctx, workspace.CreateOptions{ + ID: opts.WorkspaceID, + Name: opts.WorkspaceID, + Owner: opts.Owner, + Node: targetNode, + }) + if err != nil { + return nil, fmt.Errorf("sandbox: auto-create workspace %q: %w", opts.WorkspaceID, err) + } + nodeID = targetNode + } else { + nodeID = node } - nodeID = node } } @@ -154,6 +191,19 @@ func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Box, error) policy = Session } + var sys SystemInfo + if snap, ok := tai.GetNodeSnapshot(nodeID); ok { + sys = SystemInfo{ + OS: snap.System.OS, + Arch: snap.System.Arch, + Hostname: snap.System.Hostname, + NumCPU: snap.System.NumCPU, + TotalMem: snap.System.TotalMem, + Shell: snap.System.Shell, + TempDir: snap.System.TempDir, + } + } + box := &Box{ id: id, containerID: containerID, @@ -169,6 +219,7 @@ func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Box, error) vnc: opts.VNC, image: opts.Image, workspaceID: opts.WorkspaceID, + system: sys, } box.lastCall.Store(time.Now().UnixMilli()) diff --git a/sandbox/v2/types.go b/sandbox/v2/types.go index e5c5bca1..1d4a26d4 100644 --- a/sandbox/v2/types.go +++ b/sandbox/v2/types.go @@ -45,13 +45,15 @@ type ComputerInfo struct { Labels map[string]string } -// SystemInfo describes the hardware of a Tai node. +// SystemInfo describes the hardware and environment of a Tai node. type SystemInfo struct { OS string Arch string Hostname string NumCPU int TotalMem int64 + Shell string // preferred shell: "sh", "pwsh", "powershell", "cmd.exe" + TempDir string // system temp directory } // --------------------------------------------------------------------------- diff --git a/tai/registry/registry.go b/tai/registry/registry.go index 857db710..52115787 100644 --- a/tai/registry/registry.go +++ b/tai/registry/registry.go @@ -22,6 +22,8 @@ type SystemInfo struct { Hostname string `json:"hostname"` NumCPU int `json:"num_cpu"` TotalMem int64 `json:"total_mem,omitempty"` + Shell string `json:"shell,omitempty"` + TempDir string `json:"temp_dir,omitempty"` } // TaiNode represents a registered Tai instance (direct or tunnel). diff --git a/tai/serverinfo/pb/serverinfo.pb.go b/tai/serverinfo/pb/serverinfo.pb.go index 5d85f88a..ac804da4 100644 --- a/tai/serverinfo/pb/serverinfo.pb.go +++ b/tai/serverinfo/pb/serverinfo.pb.go @@ -2,7 +2,7 @@ // versions: // protoc-gen-go v1.36.11 // protoc v4.25.0 -// source: tai/serverinfo/pb/serverinfo.proto +// source: serverinfo.proto package pb @@ -29,7 +29,7 @@ type GetInfoRequest struct { func (x *GetInfoRequest) Reset() { *x = GetInfoRequest{} - mi := &file_tai_serverinfo_pb_serverinfo_proto_msgTypes[0] + mi := &file_serverinfo_proto_msgTypes[0] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -41,7 +41,7 @@ func (x *GetInfoRequest) String() string { func (*GetInfoRequest) ProtoMessage() {} func (x *GetInfoRequest) ProtoReflect() protoreflect.Message { - mi := &file_tai_serverinfo_pb_serverinfo_proto_msgTypes[0] + mi := &file_serverinfo_proto_msgTypes[0] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -54,7 +54,99 @@ func (x *GetInfoRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use GetInfoRequest.ProtoReflect.Descriptor instead. func (*GetInfoRequest) Descriptor() ([]byte, []int) { - return file_tai_serverinfo_pb_serverinfo_proto_rawDescGZIP(), []int{0} + return file_serverinfo_proto_rawDescGZIP(), []int{0} +} + +type SystemInfo struct { + state protoimpl.MessageState `protogen:"open.v1"` + Os string `protobuf:"bytes,1,opt,name=os,proto3" json:"os,omitempty"` + Arch string `protobuf:"bytes,2,opt,name=arch,proto3" json:"arch,omitempty"` + Hostname string `protobuf:"bytes,3,opt,name=hostname,proto3" json:"hostname,omitempty"` + NumCpu int32 `protobuf:"varint,4,opt,name=num_cpu,json=numCpu,proto3" json:"num_cpu,omitempty"` + TotalMem int64 `protobuf:"varint,5,opt,name=total_mem,json=totalMem,proto3" json:"total_mem,omitempty"` + Shell string `protobuf:"bytes,6,opt,name=shell,proto3" json:"shell,omitempty"` // preferred shell: "sh", "pwsh", "powershell", "cmd.exe" + TempDir string `protobuf:"bytes,7,opt,name=temp_dir,json=tempDir,proto3" json:"temp_dir,omitempty"` // system temp directory + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *SystemInfo) Reset() { + *x = SystemInfo{} + mi := &file_serverinfo_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *SystemInfo) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SystemInfo) ProtoMessage() {} + +func (x *SystemInfo) ProtoReflect() protoreflect.Message { + mi := &file_serverinfo_proto_msgTypes[1] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SystemInfo.ProtoReflect.Descriptor instead. +func (*SystemInfo) Descriptor() ([]byte, []int) { + return file_serverinfo_proto_rawDescGZIP(), []int{1} +} + +func (x *SystemInfo) GetOs() string { + if x != nil { + return x.Os + } + return "" +} + +func (x *SystemInfo) GetArch() string { + if x != nil { + return x.Arch + } + return "" +} + +func (x *SystemInfo) GetHostname() string { + if x != nil { + return x.Hostname + } + return "" +} + +func (x *SystemInfo) GetNumCpu() int32 { + if x != nil { + return x.NumCpu + } + return 0 +} + +func (x *SystemInfo) GetTotalMem() int64 { + if x != nil { + return x.TotalMem + } + return 0 +} + +func (x *SystemInfo) GetShell() string { + if x != nil { + return x.Shell + } + return "" +} + +func (x *SystemInfo) GetTempDir() string { + if x != nil { + return x.TempDir + } + return "" } type GetInfoResponse struct { @@ -62,13 +154,14 @@ type GetInfoResponse struct { Version string `protobuf:"bytes,1,opt,name=version,proto3" json:"version,omitempty"` Ports map[string]int32 `protobuf:"bytes,2,rep,name=ports,proto3" json:"ports,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"varint,2,opt,name=value"` // "grpc", "http", "vnc", "docker", "k8s" Capabilities map[string]bool `protobuf:"bytes,3,rep,name=capabilities,proto3" json:"capabilities,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"varint,2,opt,name=value"` // "docker", "k8s" + System *SystemInfo `protobuf:"bytes,4,opt,name=system,proto3" json:"system,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } func (x *GetInfoResponse) Reset() { *x = GetInfoResponse{} - mi := &file_tai_serverinfo_pb_serverinfo_proto_msgTypes[1] + mi := &file_serverinfo_proto_msgTypes[2] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -80,7 +173,7 @@ func (x *GetInfoResponse) String() string { func (*GetInfoResponse) ProtoMessage() {} func (x *GetInfoResponse) ProtoReflect() protoreflect.Message { - mi := &file_tai_serverinfo_pb_serverinfo_proto_msgTypes[1] + mi := &file_serverinfo_proto_msgTypes[2] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -93,7 +186,7 @@ func (x *GetInfoResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use GetInfoResponse.ProtoReflect.Descriptor instead. func (*GetInfoResponse) Descriptor() ([]byte, []int) { - return file_tai_serverinfo_pb_serverinfo_proto_rawDescGZIP(), []int{1} + return file_serverinfo_proto_rawDescGZIP(), []int{2} } func (x *GetInfoResponse) GetVersion() string { @@ -117,17 +210,34 @@ func (x *GetInfoResponse) GetCapabilities() map[string]bool { return nil } -var File_tai_serverinfo_pb_serverinfo_proto protoreflect.FileDescriptor +func (x *GetInfoResponse) GetSystem() *SystemInfo { + if x != nil { + return x.System + } + return nil +} -const file_tai_serverinfo_pb_serverinfo_proto_rawDesc = "" + +var File_serverinfo_proto protoreflect.FileDescriptor + +const file_serverinfo_proto_rawDesc = "" + "\n" + - "\"tai/serverinfo/pb/serverinfo.proto\x12\n" + + "\x10serverinfo.proto\x12\n" + "serverinfo\"\x10\n" + - "\x0eGetInfoRequest\"\xb7\x02\n" + + "\x0eGetInfoRequest\"\xb3\x01\n" + + "\n" + + "SystemInfo\x12\x0e\n" + + "\x02os\x18\x01 \x01(\tR\x02os\x12\x12\n" + + "\x04arch\x18\x02 \x01(\tR\x04arch\x12\x1a\n" + + "\bhostname\x18\x03 \x01(\tR\bhostname\x12\x17\n" + + "\anum_cpu\x18\x04 \x01(\x05R\x06numCpu\x12\x1b\n" + + "\ttotal_mem\x18\x05 \x01(\x03R\btotalMem\x12\x14\n" + + "\x05shell\x18\x06 \x01(\tR\x05shell\x12\x19\n" + + "\btemp_dir\x18\a \x01(\tR\atempDir\"\xe7\x02\n" + "\x0fGetInfoResponse\x12\x18\n" + "\aversion\x18\x01 \x01(\tR\aversion\x12<\n" + "\x05ports\x18\x02 \x03(\v2&.serverinfo.GetInfoResponse.PortsEntryR\x05ports\x12Q\n" + - "\fcapabilities\x18\x03 \x03(\v2-.serverinfo.GetInfoResponse.CapabilitiesEntryR\fcapabilities\x1a8\n" + + "\fcapabilities\x18\x03 \x03(\v2-.serverinfo.GetInfoResponse.CapabilitiesEntryR\fcapabilities\x12.\n" + + "\x06system\x18\x04 \x01(\v2\x16.serverinfo.SystemInfoR\x06system\x1a8\n" + "\n" + "PortsEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + @@ -140,56 +250,58 @@ const file_tai_serverinfo_pb_serverinfo_proto_rawDesc = "" + "\aGetInfo\x12\x1a.serverinfo.GetInfoRequest\x1a\x1b.serverinfo.GetInfoResponseB%Z#github.com/yaoapp/tai/serverinfo/pbb\x06proto3" var ( - file_tai_serverinfo_pb_serverinfo_proto_rawDescOnce sync.Once - file_tai_serverinfo_pb_serverinfo_proto_rawDescData []byte + file_serverinfo_proto_rawDescOnce sync.Once + file_serverinfo_proto_rawDescData []byte ) -func file_tai_serverinfo_pb_serverinfo_proto_rawDescGZIP() []byte { - file_tai_serverinfo_pb_serverinfo_proto_rawDescOnce.Do(func() { - file_tai_serverinfo_pb_serverinfo_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_tai_serverinfo_pb_serverinfo_proto_rawDesc), len(file_tai_serverinfo_pb_serverinfo_proto_rawDesc))) +func file_serverinfo_proto_rawDescGZIP() []byte { + file_serverinfo_proto_rawDescOnce.Do(func() { + file_serverinfo_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_serverinfo_proto_rawDesc), len(file_serverinfo_proto_rawDesc))) }) - return file_tai_serverinfo_pb_serverinfo_proto_rawDescData + return file_serverinfo_proto_rawDescData } -var file_tai_serverinfo_pb_serverinfo_proto_msgTypes = make([]protoimpl.MessageInfo, 4) -var file_tai_serverinfo_pb_serverinfo_proto_goTypes = []any{ +var file_serverinfo_proto_msgTypes = make([]protoimpl.MessageInfo, 5) +var file_serverinfo_proto_goTypes = []any{ (*GetInfoRequest)(nil), // 0: serverinfo.GetInfoRequest - (*GetInfoResponse)(nil), // 1: serverinfo.GetInfoResponse - nil, // 2: serverinfo.GetInfoResponse.PortsEntry - nil, // 3: serverinfo.GetInfoResponse.CapabilitiesEntry + (*SystemInfo)(nil), // 1: serverinfo.SystemInfo + (*GetInfoResponse)(nil), // 2: serverinfo.GetInfoResponse + nil, // 3: serverinfo.GetInfoResponse.PortsEntry + nil, // 4: serverinfo.GetInfoResponse.CapabilitiesEntry } -var file_tai_serverinfo_pb_serverinfo_proto_depIdxs = []int32{ - 2, // 0: serverinfo.GetInfoResponse.ports:type_name -> serverinfo.GetInfoResponse.PortsEntry - 3, // 1: serverinfo.GetInfoResponse.capabilities:type_name -> serverinfo.GetInfoResponse.CapabilitiesEntry - 0, // 2: serverinfo.ServerInfo.GetInfo:input_type -> serverinfo.GetInfoRequest - 1, // 3: serverinfo.ServerInfo.GetInfo:output_type -> serverinfo.GetInfoResponse - 3, // [3:4] is the sub-list for method output_type - 2, // [2:3] is the sub-list for method input_type - 2, // [2:2] is the sub-list for extension type_name - 2, // [2:2] is the sub-list for extension extendee - 0, // [0:2] is the sub-list for field type_name +var file_serverinfo_proto_depIdxs = []int32{ + 3, // 0: serverinfo.GetInfoResponse.ports:type_name -> serverinfo.GetInfoResponse.PortsEntry + 4, // 1: serverinfo.GetInfoResponse.capabilities:type_name -> serverinfo.GetInfoResponse.CapabilitiesEntry + 1, // 2: serverinfo.GetInfoResponse.system:type_name -> serverinfo.SystemInfo + 0, // 3: serverinfo.ServerInfo.GetInfo:input_type -> serverinfo.GetInfoRequest + 2, // 4: serverinfo.ServerInfo.GetInfo:output_type -> serverinfo.GetInfoResponse + 4, // [4:5] is the sub-list for method output_type + 3, // [3:4] is the sub-list for method input_type + 3, // [3:3] is the sub-list for extension type_name + 3, // [3:3] is the sub-list for extension extendee + 0, // [0:3] is the sub-list for field type_name } -func init() { file_tai_serverinfo_pb_serverinfo_proto_init() } -func file_tai_serverinfo_pb_serverinfo_proto_init() { - if File_tai_serverinfo_pb_serverinfo_proto != nil { +func init() { file_serverinfo_proto_init() } +func file_serverinfo_proto_init() { + if File_serverinfo_proto != nil { return } type x struct{} out := protoimpl.TypeBuilder{ File: protoimpl.DescBuilder{ GoPackagePath: reflect.TypeOf(x{}).PkgPath(), - RawDescriptor: unsafe.Slice(unsafe.StringData(file_tai_serverinfo_pb_serverinfo_proto_rawDesc), len(file_tai_serverinfo_pb_serverinfo_proto_rawDesc)), + RawDescriptor: unsafe.Slice(unsafe.StringData(file_serverinfo_proto_rawDesc), len(file_serverinfo_proto_rawDesc)), NumEnums: 0, - NumMessages: 4, + NumMessages: 5, NumExtensions: 0, NumServices: 1, }, - GoTypes: file_tai_serverinfo_pb_serverinfo_proto_goTypes, - DependencyIndexes: file_tai_serverinfo_pb_serverinfo_proto_depIdxs, - MessageInfos: file_tai_serverinfo_pb_serverinfo_proto_msgTypes, + GoTypes: file_serverinfo_proto_goTypes, + DependencyIndexes: file_serverinfo_proto_depIdxs, + MessageInfos: file_serverinfo_proto_msgTypes, }.Build() - File_tai_serverinfo_pb_serverinfo_proto = out.File - file_tai_serverinfo_pb_serverinfo_proto_goTypes = nil - file_tai_serverinfo_pb_serverinfo_proto_depIdxs = nil + File_serverinfo_proto = out.File + file_serverinfo_proto_goTypes = nil + file_serverinfo_proto_depIdxs = nil } diff --git a/tai/serverinfo/pb/serverinfo.proto b/tai/serverinfo/pb/serverinfo.proto index 4eece2ee..844f1881 100644 --- a/tai/serverinfo/pb/serverinfo.proto +++ b/tai/serverinfo/pb/serverinfo.proto @@ -8,8 +8,19 @@ service ServerInfo { message GetInfoRequest {} +message SystemInfo { + string os = 1; + string arch = 2; + string hostname = 3; + int32 num_cpu = 4; + int64 total_mem = 5; + string shell = 6; // preferred shell: "sh", "pwsh", "powershell", "cmd.exe" + string temp_dir = 7; // system temp directory +} + message GetInfoResponse { string version = 1; - map ports = 2; // "grpc", "http", "vnc", "docker", "k8s" + map ports = 2; // "grpc", "http", "vnc", "docker", "k8s" map capabilities = 3; // "docker", "k8s" + SystemInfo system = 4; } diff --git a/tai/serverinfo/pb/serverinfo_grpc.pb.go b/tai/serverinfo/pb/serverinfo_grpc.pb.go index 07e690ec..173e1fa0 100644 --- a/tai/serverinfo/pb/serverinfo_grpc.pb.go +++ b/tai/serverinfo/pb/serverinfo_grpc.pb.go @@ -2,7 +2,7 @@ // versions: // - protoc-gen-go-grpc v1.6.1 // - protoc v4.25.0 -// source: tai/serverinfo/pb/serverinfo.proto +// source: serverinfo.proto package pb @@ -117,5 +117,5 @@ var ServerInfo_ServiceDesc = grpc.ServiceDesc{ }, }, Streams: []grpc.StreamDesc{}, - Metadata: "tai/serverinfo/pb/serverinfo.proto", + Metadata: "serverinfo.proto", } diff --git a/tai/tai.go b/tai/tai.go index 7a1d3bd9..a2663a13 100644 --- a/tai/tai.go +++ b/tai/tai.go @@ -239,15 +239,14 @@ func (c *Client) initRemote(cfg *config) (*Client, error) { c.grpcConn = conn c.he = hepb.NewHostExecClient(conn) - caps, err := c.discoverServerInfo(conn, cfg) + info, err := c.discoverServerInfo(conn, cfg) if err != nil { - // Old Tai without ServerInfo — fall back to legacy behaviour (try Docker). - caps = map[string]bool{"docker": true} + info = &discoveredInfo{Capabilities: map[string]bool{"docker": true}} } - hasDocker := caps["docker"] - hasK8s := caps["k8s"] - hasHostExec := caps["host_exec"] + hasDocker := info.Capabilities["docker"] + hasK8s := info.Capabilities["k8s"] + hasHostExec := info.Capabilities["host_exec"] if !hasDocker && !hasK8s && !hasHostExec { conn.Close() @@ -297,9 +296,12 @@ func (c *Client) initRemote(cfg *config) (*Client, error) { id := fmt.Sprintf("%s-%d", c.host, c.ports.GRPC) c.taiID = id reg.Register(®istry.TaiNode{ - TaiID: id, - Mode: "direct", - Addr: fmt.Sprintf("tai://%s:%d", c.host, c.ports.GRPC), + TaiID: id, + Mode: "direct", + Version: info.Version, + System: info.System, + Capabilities: info.Capabilities, + Addr: fmt.Sprintf("tai://%s:%d", c.host, c.ports.GRPC), Ports: map[string]int{ "grpc": c.ports.GRPC, "http": c.ports.HTTP, @@ -351,13 +353,13 @@ func (c *Client) initTunnel(cfg *config) (*Client, error) { c.he = hepb.NewHostExecClient(conn) c.vol = volume.NewRemote(conn) - caps, err := c.discoverServerInfo(conn, cfg) + info, err := c.discoverServerInfo(conn, cfg) if err != nil { - caps = map[string]bool{"docker": true} + info = &discoveredInfo{Capabilities: map[string]bool{"docker": true}} } - hasDocker := caps["docker"] - hasHostExec := caps["host_exec"] + hasDocker := info.Capabilities["docker"] + hasHostExec := info.Capabilities["host_exec"] if !hasDocker && !hasHostExec { c.closeTunnelListeners() @@ -544,10 +546,16 @@ func isLocalHost(h string) bool { return h == "127.0.0.1" || h == "localhost" || h == "::1" } +type discoveredInfo struct { + Capabilities map[string]bool + System registry.SystemInfo + Version string +} + // discoverServerInfo calls ServerInfo.GetInfo on the remote Tai server, merges -// discovered ports into c.ports, and returns the server's capabilities map. +// discovered ports into c.ports, and returns capabilities + system info. // Ports explicitly set via WithPorts take precedence over server-reported values. -func (c *Client) discoverServerInfo(conn *grpc.ClientConn, cfg *config) (map[string]bool, error) { +func (c *Client) discoverServerInfo(conn *grpc.ClientConn, cfg *config) (*discoveredInfo, error) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() @@ -576,7 +584,25 @@ func (c *Client) discoverServerInfo(conn *grpc.ClientConn, cfg *config) (map[str if caps == nil { caps = make(map[string]bool) } - return caps, nil + + var sys registry.SystemInfo + if s := resp.System; s != nil { + sys = registry.SystemInfo{ + OS: s.Os, + Arch: s.Arch, + Hostname: s.Hostname, + NumCPU: int(s.NumCpu), + TotalMem: s.TotalMem, + Shell: s.Shell, + TempDir: s.TempDir, + } + } + + return &discoveredInfo{ + Capabilities: caps, + System: sys, + Version: resp.Version, + }, nil } // RegisterLocal probes the local Docker environment and, if reachable, @@ -616,3 +642,13 @@ func GetClient(taiID string) (*Client, bool) { } return c, true } + +// GetNodeSnapshot returns the registry snapshot for a Tai node by ID. +// Callers can inspect System, Capabilities, Mode and other registry-level fields. +func GetNodeSnapshot(taiID string) (*registry.NodeSnapshot, bool) { + reg := registry.Global() + if reg == nil { + return nil, false + } + return reg.Get(taiID) +} diff --git a/tai/volume/local.go b/tai/volume/local.go index 53410dec..11ca1e11 100644 --- a/tai/volume/local.go +++ b/tai/volume/local.go @@ -138,14 +138,125 @@ func (l *localStorage) MkdirAll(_ context.Context, sessionID, path string) error return os.MkdirAll(abs, 0o755) } +// Copy duplicates src to dst within the same workspace session. +// Supports single files and directories (recursive). Uses excludes from SyncOption +// and forceFull to overwrite even when mtime+size match. +func (l *localStorage) Copy(_ context.Context, sessionID, src, dst string, opts ...SyncOption) (*SyncResult, error) { + start := time.Now() + cfg := ApplySyncOpts(opts) + + srcAbs, err := l.abs(sessionID, src) + if err != nil { + return nil, err + } + dstAbs, err := l.abs(sessionID, dst) + if err != nil { + return nil, err + } + + srcInfo, err := os.Stat(srcAbs) + if err != nil { + return nil, err + } + + if !srcInfo.IsDir() { + n, err := l.copyFile(srcAbs, dstAbs, srcInfo, cfg.ForceFull) + if err != nil { + return nil, err + } + synced := 0 + if n > 0 { + synced = 1 + } + return &SyncResult{ + FilesSynced: synced, + BytesTransferred: n, + Duration: time.Since(start), + }, nil + } + + var synced int + var transferred int64 + err = filepath.WalkDir(srcAbs, func(abs string, d fs.DirEntry, walkErr error) error { + if walkErr != nil { + if os.IsNotExist(walkErr) { + return nil + } + return walkErr + } + rel, _ := filepath.Rel(srcAbs, abs) + if rel == "." { + return os.MkdirAll(dstAbs, 0o755) + } + + if isExcluded(rel, d.IsDir(), cfg.Excludes) { + if d.IsDir() { + return filepath.SkipDir + } + return nil + } + + target := filepath.Join(dstAbs, rel) + if d.IsDir() { + return os.MkdirAll(target, 0o755) + } + + info, err := d.Info() + if err != nil { + return nil + } + n, err := l.copyFile(abs, target, info, cfg.ForceFull) + if err != nil { + return err + } + if n > 0 { + synced++ + transferred += n + } + return nil + }) + + return &SyncResult{ + FilesSynced: synced, + BytesTransferred: transferred, + Duration: time.Since(start), + }, err +} + +func (l *localStorage) copyFile(srcAbs, dstAbs string, srcInfo os.FileInfo, force bool) (int64, error) { + if !force { + if dstInfo, e := os.Stat(dstAbs); e == nil { + if dstInfo.Size() == srcInfo.Size() && dstInfo.ModTime().Equal(srcInfo.ModTime()) { + return 0, nil + } + } + } + + data, err := os.ReadFile(srcAbs) + if err != nil { + if os.IsNotExist(err) { + return 0, nil + } + return 0, err + } + if err := os.MkdirAll(filepath.Dir(dstAbs), 0o755); err != nil { + return 0, err + } + if err := os.WriteFile(dstAbs, data, srcInfo.Mode()); err != nil { + return 0, err + } + _ = os.Chtimes(dstAbs, srcInfo.ModTime(), srcInfo.ModTime()) + return int64(len(data)), nil +} + // SyncPush copies changed files from localDir to dataDir/{sessionID}/. // Uses mtime+size to detect changes. Files that vanish during sync are skipped. func (l *localStorage) SyncPush(_ context.Context, sessionID, localDir string, opts ...SyncOption) (*SyncResult, error) { start := time.Now() - cfg := applySyncOpts(opts) + cfg := ApplySyncOpts(opts) dst := l.root(sessionID) - if cfg.remotePath != "" { - dst = filepath.Join(dst, filepath.Clean(cfg.remotePath)) + if cfg.RemotePath != "" { + dst = filepath.Join(dst, filepath.Clean(cfg.RemotePath)) } if err := os.MkdirAll(dst, 0o755); err != nil { return nil, err @@ -167,7 +278,7 @@ func (l *localStorage) SyncPush(_ context.Context, sessionID, localDir string, o } rel = filepath.ToSlash(rel) - if isExcluded(rel, d.IsDir(), cfg.excludes) { + if isExcluded(rel, d.IsDir(), cfg.Excludes) { if d.IsDir() { return filepath.SkipDir } @@ -184,7 +295,7 @@ func (l *localStorage) SyncPush(_ context.Context, sessionID, localDir string, o return nil // file vanished between readdir and stat; skip } - if !cfg.forceFull { + if !cfg.ForceFull { if dstInfo, e := os.Stat(target); e == nil { if dstInfo.Size() == srcInfo.Size() && dstInfo.ModTime().Equal(srcInfo.ModTime()) { return nil @@ -222,10 +333,10 @@ func (l *localStorage) SyncPush(_ context.Context, sessionID, localDir string, o // Files that vanish during sync are skipped. func (l *localStorage) SyncPull(_ context.Context, sessionID, localDir string, opts ...SyncOption) (*SyncResult, error) { start := time.Now() - cfg := applySyncOpts(opts) + cfg := ApplySyncOpts(opts) src := l.root(sessionID) - if cfg.remotePath != "" { - src = filepath.Join(src, filepath.Clean(cfg.remotePath)) + if cfg.RemotePath != "" { + src = filepath.Join(src, filepath.Clean(cfg.RemotePath)) } if err := os.MkdirAll(localDir, 0o755); err != nil { return nil, err @@ -247,7 +358,7 @@ func (l *localStorage) SyncPull(_ context.Context, sessionID, localDir string, o } rel = filepath.ToSlash(rel) - if isExcluded(rel, d.IsDir(), cfg.excludes) { + if isExcluded(rel, d.IsDir(), cfg.Excludes) { if d.IsDir() { return filepath.SkipDir } @@ -264,7 +375,7 @@ func (l *localStorage) SyncPull(_ context.Context, sessionID, localDir string, o return nil // file vanished between readdir and stat; skip } - if !cfg.forceFull { + if !cfg.ForceFull { if dstInfo, e := os.Stat(target); e == nil { if dstInfo.Size() == srcInfo.Size() && dstInfo.ModTime().Equal(srcInfo.ModTime()) { return nil diff --git a/tai/volume/mock_test.go b/tai/volume/mock_test.go index 59f0ce06..bdee3352 100644 --- a/tai/volume/mock_test.go +++ b/tai/volume/mock_test.go @@ -197,6 +197,13 @@ func (m *mockVolumeServer) ListDir(_ context.Context, req *pb.FSRequest) (*pb.FS }}, nil } +func (m *mockVolumeServer) Copy(_ context.Context, req *pb.FSCopyRequest) (*pb.SyncResult, error) { + return &pb.SyncResult{ + FilesSynced: 1, + BytesTransferred: 42, + }, nil +} + func startMockServer(t *testing.T, mock *mockVolumeServer) (*grpc.ClientConn, func()) { t.Helper() lis, err := net.Listen("tcp", "127.0.0.1:0") @@ -549,6 +556,10 @@ func (m *errMockVolumeServer) MkdirAll(_ context.Context, _ *pb.FSRequest) (*pb. return nil, fmt.Errorf("injected mkdir error") } +func (m *errMockVolumeServer) Copy(_ context.Context, _ *pb.FSCopyRequest) (*pb.SyncResult, error) { + return nil, fmt.Errorf("injected copy error") +} + func startErrMockServer(t *testing.T) (*grpc.ClientConn, func()) { t.Helper() lis, err := net.Listen("tcp", "127.0.0.1:0") @@ -694,3 +705,46 @@ func TestPbToFileInfo(t *testing.T) { t.Errorf("mode = %v", fi.Mode) } } + +func TestMockRemoteCopy(t *testing.T) { + conn, cleanup := startMockServer(t, &mockVolumeServer{}) + defer cleanup() + + vol := NewRemote(conn) + result, err := vol.Copy(context.Background(), "s1", "src.txt", "dst.txt") + if err != nil { + t.Fatalf("Copy: %v", err) + } + if result.FilesSynced != 1 { + t.Errorf("synced = %d, want 1", result.FilesSynced) + } + if result.BytesTransferred != 42 { + t.Errorf("bytes = %d, want 42", result.BytesTransferred) + } +} + +func TestMockRemoteCopyWithOpts(t *testing.T) { + conn, cleanup := startMockServer(t, &mockVolumeServer{}) + defer cleanup() + + vol := NewRemote(conn) + result, err := vol.Copy(context.Background(), "s1", "src", "dst", + WithExcludes("*.log"), WithForceFull()) + if err != nil { + t.Fatalf("Copy: %v", err) + } + if result.FilesSynced != 1 { + t.Errorf("synced = %d", result.FilesSynced) + } +} + +func TestErrRemoteCopy(t *testing.T) { + conn, cleanup := startErrMockServer(t) + defer cleanup() + + vol := NewRemote(conn) + _, err := vol.Copy(context.Background(), "s1", "a", "b") + if err == nil { + t.Error("expected error") + } +} diff --git a/tai/volume/pb/volume.pb.go b/tai/volume/pb/volume.pb.go index 1437faf6..96e97bd8 100644 --- a/tai/volume/pb/volume.pb.go +++ b/tai/volume/pb/volume.pb.go @@ -1035,6 +1035,82 @@ func (x *FSRenameRequest) GetNewPath() string { return "" } +type FSCopyRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` + SrcPath string `protobuf:"bytes,2,opt,name=src_path,json=srcPath,proto3" json:"src_path,omitempty"` + DstPath string `protobuf:"bytes,3,opt,name=dst_path,json=dstPath,proto3" json:"dst_path,omitempty"` + Excludes []string `protobuf:"bytes,4,rep,name=excludes,proto3" json:"excludes,omitempty"` // glob patterns + Force bool `protobuf:"varint,5,opt,name=force,proto3" json:"force,omitempty"` // overwrite even if mtime/size match + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *FSCopyRequest) Reset() { + *x = FSCopyRequest{} + mi := &file_tai_volume_pb_volume_proto_msgTypes[15] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *FSCopyRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FSCopyRequest) ProtoMessage() {} + +func (x *FSCopyRequest) ProtoReflect() protoreflect.Message { + mi := &file_tai_volume_pb_volume_proto_msgTypes[15] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use FSCopyRequest.ProtoReflect.Descriptor instead. +func (*FSCopyRequest) Descriptor() ([]byte, []int) { + return file_tai_volume_pb_volume_proto_rawDescGZIP(), []int{15} +} + +func (x *FSCopyRequest) GetSessionId() string { + if x != nil { + return x.SessionId + } + return "" +} + +func (x *FSCopyRequest) GetSrcPath() string { + if x != nil { + return x.SrcPath + } + return "" +} + +func (x *FSCopyRequest) GetDstPath() string { + if x != nil { + return x.DstPath + } + return "" +} + +func (x *FSCopyRequest) GetExcludes() []string { + if x != nil { + return x.Excludes + } + return nil +} + +func (x *FSCopyRequest) GetForce() bool { + if x != nil { + return x.Force + } + return false +} + type ArchiveRequest struct { state protoimpl.MessageState `protogen:"open.v1"` SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` @@ -1047,7 +1123,7 @@ type ArchiveRequest struct { func (x *ArchiveRequest) Reset() { *x = ArchiveRequest{} - mi := &file_tai_volume_pb_volume_proto_msgTypes[15] + mi := &file_tai_volume_pb_volume_proto_msgTypes[16] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1059,7 +1135,7 @@ func (x *ArchiveRequest) String() string { func (*ArchiveRequest) ProtoMessage() {} func (x *ArchiveRequest) ProtoReflect() protoreflect.Message { - mi := &file_tai_volume_pb_volume_proto_msgTypes[15] + mi := &file_tai_volume_pb_volume_proto_msgTypes[16] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1072,7 +1148,7 @@ func (x *ArchiveRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use ArchiveRequest.ProtoReflect.Descriptor instead. func (*ArchiveRequest) Descriptor() ([]byte, []int) { - return file_tai_volume_pb_volume_proto_rawDescGZIP(), []int{15} + return file_tai_volume_pb_volume_proto_rawDescGZIP(), []int{16} } func (x *ArchiveRequest) GetSessionId() string { @@ -1113,7 +1189,7 @@ type ArchiveResponse struct { func (x *ArchiveResponse) Reset() { *x = ArchiveResponse{} - mi := &file_tai_volume_pb_volume_proto_msgTypes[16] + mi := &file_tai_volume_pb_volume_proto_msgTypes[17] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1125,7 +1201,7 @@ func (x *ArchiveResponse) String() string { func (*ArchiveResponse) ProtoMessage() {} func (x *ArchiveResponse) ProtoReflect() protoreflect.Message { - mi := &file_tai_volume_pb_volume_proto_msgTypes[16] + mi := &file_tai_volume_pb_volume_proto_msgTypes[17] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1138,7 +1214,7 @@ func (x *ArchiveResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use ArchiveResponse.ProtoReflect.Descriptor instead. func (*ArchiveResponse) Descriptor() ([]byte, []int) { - return file_tai_volume_pb_volume_proto_rawDescGZIP(), []int{16} + return file_tai_volume_pb_volume_proto_rawDescGZIP(), []int{17} } func (x *ArchiveResponse) GetSizeBytes() int64 { @@ -1240,7 +1316,14 @@ const file_tai_volume_pb_volume_proto_rawDesc = "" + "\n" + "session_id\x18\x01 \x01(\tR\tsessionId\x12\x19\n" + "\bold_path\x18\x02 \x01(\tR\aoldPath\x12\x19\n" + - "\bnew_path\x18\x03 \x01(\tR\anewPath\"\x81\x01\n" + + "\bnew_path\x18\x03 \x01(\tR\anewPath\"\x96\x01\n" + + "\rFSCopyRequest\x12\x1d\n" + + "\n" + + "session_id\x18\x01 \x01(\tR\tsessionId\x12\x19\n" + + "\bsrc_path\x18\x02 \x01(\tR\asrcPath\x12\x19\n" + + "\bdst_path\x18\x03 \x01(\tR\adstPath\x12\x1a\n" + + "\bexcludes\x18\x04 \x03(\tR\bexcludes\x12\x14\n" + + "\x05force\x18\x05 \x01(\bR\x05force\"\x81\x01\n" + "\x0eArchiveRequest\x12\x1d\n" + "\n" + "session_id\x18\x01 \x01(\tR\tsessionId\x12\x19\n" + @@ -1251,7 +1334,7 @@ const file_tai_volume_pb_volume_proto_rawDesc = "" + "\n" + "size_bytes\x18\x01 \x01(\x03R\tsizeBytes\x12\x1f\n" + "\vfiles_count\x18\x02 \x01(\x05R\n" + - "filesCount2\xc7\a\n" + + "filesCount2\xfa\a\n" + "\x06Volume\x128\n" + "\bSyncPush\x12\x13.volume.SyncMessage\x1a\x13.volume.SyncMessage(\x010\x01\x127\n" + "\bSyncPull\x12\x14.volume.SyncManifest\x1a\x13.volume.SyncMessage0\x01\x128\n" + @@ -1261,7 +1344,8 @@ const file_tai_volume_pb_volume_proto_rawDesc = "" + "\aListDir\x12\x11.volume.FSRequest\x1a\x16.volume.FSListResponse\x127\n" + "\x06Remove\x12\x17.volume.FSRemoveRequest\x1a\x14.volume.FSOpResponse\x127\n" + "\x06Rename\x12\x17.volume.FSRenameRequest\x1a\x14.volume.FSOpResponse\x123\n" + - "\bMkdirAll\x12\x11.volume.FSRequest\x1a\x14.volume.FSOpResponse\x126\n" + + "\bMkdirAll\x12\x11.volume.FSRequest\x1a\x14.volume.FSOpResponse\x121\n" + + "\x04Copy\x12\x15.volume.FSCopyRequest\x1a\x12.volume.SyncResult\x126\n" + "\x03Zip\x12\x16.volume.ArchiveRequest\x1a\x17.volume.ArchiveResponse\x128\n" + "\x05Unzip\x12\x16.volume.ArchiveRequest\x1a\x17.volume.ArchiveResponse\x127\n" + "\x04Gzip\x12\x16.volume.ArchiveRequest\x1a\x17.volume.ArchiveResponse\x129\n" + @@ -1284,7 +1368,7 @@ func file_tai_volume_pb_volume_proto_rawDescGZIP() []byte { } var file_tai_volume_pb_volume_proto_enumTypes = make([]protoimpl.EnumInfo, 1) -var file_tai_volume_pb_volume_proto_msgTypes = make([]protoimpl.MessageInfo, 17) +var file_tai_volume_pb_volume_proto_msgTypes = make([]protoimpl.MessageInfo, 18) var file_tai_volume_pb_volume_proto_goTypes = []any{ (FileChunk_ChunkType)(0), // 0: volume.FileChunk.ChunkType (*FileInfo)(nil), // 1: volume.FileInfo @@ -1302,8 +1386,9 @@ var file_tai_volume_pb_volume_proto_goTypes = []any{ (*FSListResponse)(nil), // 13: volume.FSListResponse (*FSRemoveRequest)(nil), // 14: volume.FSRemoveRequest (*FSRenameRequest)(nil), // 15: volume.FSRenameRequest - (*ArchiveRequest)(nil), // 16: volume.ArchiveRequest - (*ArchiveResponse)(nil), // 17: volume.ArchiveResponse + (*FSCopyRequest)(nil), // 16: volume.FSCopyRequest + (*ArchiveRequest)(nil), // 17: volume.ArchiveRequest + (*ArchiveResponse)(nil), // 18: volume.ArchiveResponse } var file_tai_volume_pb_volume_proto_depIdxs = []int32{ 1, // 0: volume.SyncManifest.files:type_name -> volume.FileInfo @@ -1322,33 +1407,35 @@ var file_tai_volume_pb_volume_proto_depIdxs = []int32{ 14, // 13: volume.Volume.Remove:input_type -> volume.FSRemoveRequest 15, // 14: volume.Volume.Rename:input_type -> volume.FSRenameRequest 7, // 15: volume.Volume.MkdirAll:input_type -> volume.FSRequest - 16, // 16: volume.Volume.Zip:input_type -> volume.ArchiveRequest - 16, // 17: volume.Volume.Unzip:input_type -> volume.ArchiveRequest - 16, // 18: volume.Volume.Gzip:input_type -> volume.ArchiveRequest - 16, // 19: volume.Volume.Gunzip:input_type -> volume.ArchiveRequest - 16, // 20: volume.Volume.Tar:input_type -> volume.ArchiveRequest - 16, // 21: volume.Volume.Untar:input_type -> volume.ArchiveRequest - 16, // 22: volume.Volume.Tgz:input_type -> volume.ArchiveRequest - 16, // 23: volume.Volume.Untgz:input_type -> volume.ArchiveRequest - 3, // 24: volume.Volume.SyncPush:output_type -> volume.SyncMessage - 3, // 25: volume.Volume.SyncPull:output_type -> volume.SyncMessage - 10, // 26: volume.Volume.ReadFile:output_type -> volume.FSDataChunk - 12, // 27: volume.Volume.WriteFile:output_type -> volume.FSWriteResponse - 1, // 28: volume.Volume.Stat:output_type -> volume.FileInfo - 13, // 29: volume.Volume.ListDir:output_type -> volume.FSListResponse - 8, // 30: volume.Volume.Remove:output_type -> volume.FSOpResponse - 8, // 31: volume.Volume.Rename:output_type -> volume.FSOpResponse - 8, // 32: volume.Volume.MkdirAll:output_type -> volume.FSOpResponse - 17, // 33: volume.Volume.Zip:output_type -> volume.ArchiveResponse - 17, // 34: volume.Volume.Unzip:output_type -> volume.ArchiveResponse - 17, // 35: volume.Volume.Gzip:output_type -> volume.ArchiveResponse - 17, // 36: volume.Volume.Gunzip:output_type -> volume.ArchiveResponse - 17, // 37: volume.Volume.Tar:output_type -> volume.ArchiveResponse - 17, // 38: volume.Volume.Untar:output_type -> volume.ArchiveResponse - 17, // 39: volume.Volume.Tgz:output_type -> volume.ArchiveResponse - 17, // 40: volume.Volume.Untgz:output_type -> volume.ArchiveResponse - 24, // [24:41] is the sub-list for method output_type - 7, // [7:24] is the sub-list for method input_type + 16, // 16: volume.Volume.Copy:input_type -> volume.FSCopyRequest + 17, // 17: volume.Volume.Zip:input_type -> volume.ArchiveRequest + 17, // 18: volume.Volume.Unzip:input_type -> volume.ArchiveRequest + 17, // 19: volume.Volume.Gzip:input_type -> volume.ArchiveRequest + 17, // 20: volume.Volume.Gunzip:input_type -> volume.ArchiveRequest + 17, // 21: volume.Volume.Tar:input_type -> volume.ArchiveRequest + 17, // 22: volume.Volume.Untar:input_type -> volume.ArchiveRequest + 17, // 23: volume.Volume.Tgz:input_type -> volume.ArchiveRequest + 17, // 24: volume.Volume.Untgz:input_type -> volume.ArchiveRequest + 3, // 25: volume.Volume.SyncPush:output_type -> volume.SyncMessage + 3, // 26: volume.Volume.SyncPull:output_type -> volume.SyncMessage + 10, // 27: volume.Volume.ReadFile:output_type -> volume.FSDataChunk + 12, // 28: volume.Volume.WriteFile:output_type -> volume.FSWriteResponse + 1, // 29: volume.Volume.Stat:output_type -> volume.FileInfo + 13, // 30: volume.Volume.ListDir:output_type -> volume.FSListResponse + 8, // 31: volume.Volume.Remove:output_type -> volume.FSOpResponse + 8, // 32: volume.Volume.Rename:output_type -> volume.FSOpResponse + 8, // 33: volume.Volume.MkdirAll:output_type -> volume.FSOpResponse + 6, // 34: volume.Volume.Copy:output_type -> volume.SyncResult + 18, // 35: volume.Volume.Zip:output_type -> volume.ArchiveResponse + 18, // 36: volume.Volume.Unzip:output_type -> volume.ArchiveResponse + 18, // 37: volume.Volume.Gzip:output_type -> volume.ArchiveResponse + 18, // 38: volume.Volume.Gunzip:output_type -> volume.ArchiveResponse + 18, // 39: volume.Volume.Tar:output_type -> volume.ArchiveResponse + 18, // 40: volume.Volume.Untar:output_type -> volume.ArchiveResponse + 18, // 41: volume.Volume.Tgz:output_type -> volume.ArchiveResponse + 18, // 42: volume.Volume.Untgz:output_type -> volume.ArchiveResponse + 25, // [25:43] is the sub-list for method output_type + 7, // [7:25] is the sub-list for method input_type 7, // [7:7] is the sub-list for extension type_name 7, // [7:7] is the sub-list for extension extendee 0, // [0:7] is the sub-list for field type_name @@ -1371,7 +1458,7 @@ func file_tai_volume_pb_volume_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_tai_volume_pb_volume_proto_rawDesc), len(file_tai_volume_pb_volume_proto_rawDesc)), NumEnums: 1, - NumMessages: 17, + NumMessages: 18, NumExtensions: 0, NumServices: 1, }, diff --git a/tai/volume/pb/volume.proto b/tai/volume/pb/volume.proto index 464b4d1b..5fa38279 100644 --- a/tai/volume/pb/volume.proto +++ b/tai/volume/pb/volume.proto @@ -30,6 +30,8 @@ service Volume { rpc Remove(FSRemoveRequest) returns (FSOpResponse); rpc Rename(FSRenameRequest) returns (FSOpResponse); rpc MkdirAll(FSRequest) returns (FSOpResponse); + // Copy: copy src to dst within the same workspace (server-side when remote). + rpc Copy(FSCopyRequest) returns (SyncResult); // --- Archive / Compression --- @@ -150,6 +152,14 @@ message FSRenameRequest { string new_path = 3; } +message FSCopyRequest { + string session_id = 1; + string src_path = 2; + string dst_path = 3; + repeated string excludes = 4; // glob patterns + bool force = 5; // overwrite even if mtime/size match +} + // --- Archive / Compression Messages --- message ArchiveRequest { diff --git a/tai/volume/pb/volume_grpc.pb.go b/tai/volume/pb/volume_grpc.pb.go index b2706be5..a5370fe3 100644 --- a/tai/volume/pb/volume_grpc.pb.go +++ b/tai/volume/pb/volume_grpc.pb.go @@ -28,6 +28,7 @@ const ( Volume_Remove_FullMethodName = "/volume.Volume/Remove" Volume_Rename_FullMethodName = "/volume.Volume/Rename" Volume_MkdirAll_FullMethodName = "/volume.Volume/MkdirAll" + Volume_Copy_FullMethodName = "/volume.Volume/Copy" Volume_Zip_FullMethodName = "/volume.Volume/Zip" Volume_Unzip_FullMethodName = "/volume.Volume/Unzip" Volume_Gzip_FullMethodName = "/volume.Volume/Gzip" @@ -63,6 +64,8 @@ type VolumeClient interface { Remove(ctx context.Context, in *FSRemoveRequest, opts ...grpc.CallOption) (*FSOpResponse, error) Rename(ctx context.Context, in *FSRenameRequest, opts ...grpc.CallOption) (*FSOpResponse, error) MkdirAll(ctx context.Context, in *FSRequest, opts ...grpc.CallOption) (*FSOpResponse, error) + // Copy: copy src to dst within the same workspace (server-side when remote). + Copy(ctx context.Context, in *FSCopyRequest, opts ...grpc.CallOption) (*SyncResult, error) Zip(ctx context.Context, in *ArchiveRequest, opts ...grpc.CallOption) (*ArchiveResponse, error) Unzip(ctx context.Context, in *ArchiveRequest, opts ...grpc.CallOption) (*ArchiveResponse, error) Gzip(ctx context.Context, in *ArchiveRequest, opts ...grpc.CallOption) (*ArchiveResponse, error) @@ -195,6 +198,16 @@ func (c *volumeClient) MkdirAll(ctx context.Context, in *FSRequest, opts ...grpc return out, nil } +func (c *volumeClient) Copy(ctx context.Context, in *FSCopyRequest, opts ...grpc.CallOption) (*SyncResult, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(SyncResult) + err := c.cc.Invoke(ctx, Volume_Copy_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + func (c *volumeClient) Zip(ctx context.Context, in *ArchiveRequest, opts ...grpc.CallOption) (*ArchiveResponse, error) { cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) out := new(ArchiveResponse) @@ -300,6 +313,8 @@ type VolumeServer interface { Remove(context.Context, *FSRemoveRequest) (*FSOpResponse, error) Rename(context.Context, *FSRenameRequest) (*FSOpResponse, error) MkdirAll(context.Context, *FSRequest) (*FSOpResponse, error) + // Copy: copy src to dst within the same workspace (server-side when remote). + Copy(context.Context, *FSCopyRequest) (*SyncResult, error) Zip(context.Context, *ArchiveRequest) (*ArchiveResponse, error) Unzip(context.Context, *ArchiveRequest) (*ArchiveResponse, error) Gzip(context.Context, *ArchiveRequest) (*ArchiveResponse, error) @@ -345,6 +360,9 @@ func (UnimplementedVolumeServer) Rename(context.Context, *FSRenameRequest) (*FSO func (UnimplementedVolumeServer) MkdirAll(context.Context, *FSRequest) (*FSOpResponse, error) { return nil, status.Error(codes.Unimplemented, "method MkdirAll not implemented") } +func (UnimplementedVolumeServer) Copy(context.Context, *FSCopyRequest) (*SyncResult, error) { + return nil, status.Error(codes.Unimplemented, "method Copy not implemented") +} func (UnimplementedVolumeServer) Zip(context.Context, *ArchiveRequest) (*ArchiveResponse, error) { return nil, status.Error(codes.Unimplemented, "method Zip not implemented") } @@ -516,6 +534,24 @@ func _Volume_MkdirAll_Handler(srv interface{}, ctx context.Context, dec func(int return interceptor(ctx, in, info, handler) } +func _Volume_Copy_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(FSCopyRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(VolumeServer).Copy(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Volume_Copy_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(VolumeServer).Copy(ctx, req.(*FSCopyRequest)) + } + return interceptor(ctx, in, info, handler) +} + func _Volume_Zip_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { in := new(ArchiveRequest) if err := dec(in); err != nil { @@ -687,6 +723,10 @@ var Volume_ServiceDesc = grpc.ServiceDesc{ MethodName: "MkdirAll", Handler: _Volume_MkdirAll_Handler, }, + { + MethodName: "Copy", + Handler: _Volume_Copy_Handler, + }, { MethodName: "Zip", Handler: _Volume_Zip_Handler, diff --git a/tai/volume/remote.go b/tai/volume/remote.go index cc41f805..e7019e66 100644 --- a/tai/volume/remote.go +++ b/tai/volume/remote.go @@ -170,7 +170,7 @@ func (r *remoteStorage) MkdirAll(ctx context.Context, sessionID, path string) er // SyncPush sends local files to Tai using the manifest-first bidi streaming protocol. func (r *remoteStorage) SyncPush(ctx context.Context, sessionID, localDir string, opts ...SyncOption) (*SyncResult, error) { start := time.Now() - cfg := applySyncOpts(opts) + cfg := ApplySyncOpts(opts) // Scan local directory var manifest []*pb.FileInfo @@ -183,7 +183,7 @@ func (r *remoteStorage) SyncPush(ctx context.Context, sessionID, localDir string return nil } rel = filepath.ToSlash(rel) - if isExcluded(rel, d.IsDir(), cfg.excludes) { + if isExcluded(rel, d.IsDir(), cfg.Excludes) { if d.IsDir() { return filepath.SkipDir } @@ -217,8 +217,8 @@ func (r *remoteStorage) SyncPush(ctx context.Context, sessionID, localDir string Manifest: &pb.SyncManifest{ SessionId: sessionID, Files: manifest, - ForceFull: cfg.forceFull, - RemotePath: cfg.remotePath, + ForceFull: cfg.ForceFull, + RemotePath: cfg.RemotePath, }, }, }); err != nil { @@ -310,7 +310,7 @@ func (r *remoteStorage) SyncPush(ctx context.Context, sessionID, localDir string // SyncPull receives changed files from Tai. func (r *remoteStorage) SyncPull(ctx context.Context, sessionID, localDir string, opts ...SyncOption) (*SyncResult, error) { start := time.Now() - cfg := applySyncOpts(opts) + cfg := ApplySyncOpts(opts) // Build local manifest var manifest []*pb.FileInfo @@ -323,7 +323,7 @@ func (r *remoteStorage) SyncPull(ctx context.Context, sessionID, localDir string return nil } rel = filepath.ToSlash(rel) - if isExcluded(rel, d.IsDir(), cfg.excludes) { + if isExcluded(rel, d.IsDir(), cfg.Excludes) { if d.IsDir() { return filepath.SkipDir } @@ -346,8 +346,8 @@ func (r *remoteStorage) SyncPull(ctx context.Context, sessionID, localDir string stream, err := r.client.SyncPull(ctx, &pb.SyncManifest{ SessionId: sessionID, Files: manifest, - ForceFull: cfg.forceFull, - RemotePath: cfg.remotePath, + ForceFull: cfg.ForceFull, + RemotePath: cfg.RemotePath, }) if err != nil { return nil, err @@ -514,6 +514,27 @@ func (r *remoteStorage) Untgz(ctx context.Context, sessionID, src, dst string) ( return &ArchiveResult{SizeBytes: resp.SizeBytes, FilesCount: int(resp.FilesCount)}, nil } +func (r *remoteStorage) Copy(ctx context.Context, sessionID, src, dst string, opts ...SyncOption) (*SyncResult, error) { + start := time.Now() + cfg := ApplySyncOpts(opts) + + resp, err := r.client.Copy(ctx, &pb.FSCopyRequest{ + SessionId: sessionID, + SrcPath: src, + DstPath: dst, + Excludes: cfg.Excludes, + Force: cfg.ForceFull, + }) + if err != nil { + return nil, err + } + return &SyncResult{ + FilesSynced: int(resp.FilesSynced), + BytesTransferred: resp.BytesTransferred, + Duration: time.Since(start), + }, nil +} + func (r *remoteStorage) Close() error { return nil } diff --git a/tai/volume/volume.go b/tai/volume/volume.go index be74a0c0..e8498c33 100644 --- a/tai/volume/volume.go +++ b/tai/volume/volume.go @@ -20,6 +20,7 @@ type Volume interface { SyncPush(ctx context.Context, sessionID, localDir string, opts ...SyncOption) (*SyncResult, error) SyncPull(ctx context.Context, sessionID, localDir string, opts ...SyncOption) (*SyncResult, error) + Copy(ctx context.Context, sessionID, src, dst string, opts ...SyncOption) (*SyncResult, error) Zip(ctx context.Context, sessionID, src, dst string, excludes []string) (*ArchiveResult, error) Unzip(ctx context.Context, sessionID, src, dst string) (*ArchiveResult, error) @@ -56,31 +57,33 @@ type ArchiveResult struct { } // SyncOption configures sync behavior. -type SyncOption func(*syncConfig) +type SyncOption func(*SyncConfig) -type syncConfig struct { - forceFull bool - excludes []string - remotePath string +// SyncConfig holds resolved sync options. +type SyncConfig struct { + ForceFull bool + Excludes []string + RemotePath string } // WithForceFull skips snapshot caches and diffs against actual disk. func WithForceFull() SyncOption { - return func(c *syncConfig) { c.forceFull = true } + return func(c *SyncConfig) { c.ForceFull = true } } // WithExcludes adds glob patterns to exclude from sync. func WithExcludes(patterns ...string) SyncOption { - return func(c *syncConfig) { c.excludes = append(c.excludes, patterns...) } + return func(c *SyncConfig) { c.Excludes = append(c.Excludes, patterns...) } } // WithRemotePath sets a sub-path within the workspace root for sync operations. func WithRemotePath(path string) SyncOption { - return func(c *syncConfig) { c.remotePath = path } + return func(c *SyncConfig) { c.RemotePath = path } } -func applySyncOpts(opts []SyncOption) syncConfig { - var cfg syncConfig +// ApplySyncOpts resolves a slice of SyncOption into a SyncConfig. +func ApplySyncOpts(opts []SyncOption) SyncConfig { + var cfg SyncConfig for _, o := range opts { o(&cfg) } diff --git a/tai/volume/volume_test.go b/tai/volume/volume_test.go index 79ef4b9b..eae954a8 100644 --- a/tai/volume/volume_test.go +++ b/tai/volume/volume_test.go @@ -1286,3 +1286,244 @@ func TestLocalSyncExcludes(t *testing.T) { t.Error("excluded file should not exist") } } + +func TestLocalCopyFile(t *testing.T) { + dir := t.TempDir() + vol := NewLocal(dir) + defer vol.Close() + ctx := context.Background() + sid := "copy-file" + + _ = vol.WriteFile(ctx, sid, "src.txt", []byte("hello copy"), 0o644) + + result, err := vol.Copy(ctx, sid, "src.txt", "dst.txt") + if err != nil { + t.Fatalf("Copy file: %v", err) + } + if result.FilesSynced != 1 { + t.Errorf("synced = %d, want 1", result.FilesSynced) + } + if result.BytesTransferred != 10 { + t.Errorf("bytes = %d, want 10", result.BytesTransferred) + } + + data, _, err := vol.ReadFile(ctx, sid, "dst.txt") + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + if string(data) != "hello copy" { + t.Errorf("content = %q", data) + } +} + +func TestLocalCopyDir(t *testing.T) { + dir := t.TempDir() + vol := NewLocal(dir) + defer vol.Close() + ctx := context.Background() + sid := "copy-dir" + + _ = vol.MkdirAll(ctx, sid, "src/sub") + _ = vol.WriteFile(ctx, sid, "src/a.txt", []byte("aaa"), 0o644) + _ = vol.WriteFile(ctx, sid, "src/sub/b.txt", []byte("bbb"), 0o644) + + result, err := vol.Copy(ctx, sid, "src", "dst") + if err != nil { + t.Fatalf("Copy dir: %v", err) + } + if result.FilesSynced != 2 { + t.Errorf("synced = %d, want 2", result.FilesSynced) + } + + data, _, err := vol.ReadFile(ctx, sid, "dst/a.txt") + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + if string(data) != "aaa" { + t.Errorf("content = %q", data) + } + + data, _, err = vol.ReadFile(ctx, sid, "dst/sub/b.txt") + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + if string(data) != "bbb" { + t.Errorf("content = %q", data) + } +} + +func TestLocalCopyExcludes(t *testing.T) { + dir := t.TempDir() + vol := NewLocal(dir) + defer vol.Close() + ctx := context.Background() + sid := "copy-excl" + + _ = vol.MkdirAll(ctx, sid, "src") + _ = vol.WriteFile(ctx, sid, "src/keep.txt", []byte("keep"), 0o644) + _ = vol.WriteFile(ctx, sid, "src/skip.log", []byte("skip"), 0o644) + + result, err := vol.Copy(ctx, sid, "src", "dst", WithExcludes("*.log")) + if err != nil { + t.Fatalf("Copy: %v", err) + } + if result.FilesSynced != 1 { + t.Errorf("synced = %d, want 1", result.FilesSynced) + } + + _, err = vol.Stat(ctx, sid, "dst/keep.txt") + if err != nil { + t.Error("keep.txt should exist") + } + _, err = vol.Stat(ctx, sid, "dst/skip.log") + if !os.IsNotExist(err) { + t.Error("skip.log should not exist") + } +} + +func TestLocalCopySkipsUnchanged(t *testing.T) { + dir := t.TempDir() + vol := NewLocal(dir) + defer vol.Close() + ctx := context.Background() + sid := "copy-skip" + + _ = vol.WriteFile(ctx, sid, "src.txt", []byte("data"), 0o644) + + result1, err := vol.Copy(ctx, sid, "src.txt", "dst.txt", WithForceFull()) + if err != nil { + t.Fatalf("Copy 1: %v", err) + } + if result1.FilesSynced != 1 { + t.Errorf("first copy synced = %d, want 1", result1.FilesSynced) + } + + result2, err := vol.Copy(ctx, sid, "src.txt", "dst.txt") + if err != nil { + t.Fatalf("Copy 2: %v", err) + } + if result2.FilesSynced != 0 { + t.Errorf("second copy synced = %d, want 0 (unchanged)", result2.FilesSynced) + } +} + +func TestLocalCopyForceFull(t *testing.T) { + dir := t.TempDir() + vol := NewLocal(dir) + defer vol.Close() + ctx := context.Background() + sid := "copy-force" + + _ = vol.WriteFile(ctx, sid, "src.txt", []byte("data"), 0o644) + + _, _ = vol.Copy(ctx, sid, "src.txt", "dst.txt", WithForceFull()) + + result, err := vol.Copy(ctx, sid, "src.txt", "dst.txt", WithForceFull()) + if err != nil { + t.Fatalf("Copy: %v", err) + } + if result.FilesSynced != 1 { + t.Errorf("force copy synced = %d, want 1", result.FilesSynced) + } +} + +func TestLocalCopyNotExist(t *testing.T) { + dir := t.TempDir() + vol := NewLocal(dir) + ctx := context.Background() + + _, err := vol.Copy(ctx, "test", "nonexistent", "dst") + if err == nil { + t.Error("expected error for copy nonexistent source") + } +} + +func TestLocalCopyPathTraversal(t *testing.T) { + dir := t.TempDir() + vol := NewLocal(dir) + ctx := context.Background() + + _, err := vol.Copy(ctx, "test", "../../etc/passwd", "dst") + if err == nil { + t.Error("expected error for path traversal in src") + } + + _, err = vol.Copy(ctx, "test", "src", "../../etc/evil") + if err == nil { + t.Error("expected error for path traversal in dst") + } +} + +func TestRemoteCopy(t *testing.T) { + addr := taiTestGRPC() + conn, err := grpc.NewClient(addr, + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + if err != nil { + t.Skipf("gRPC dial %s: %v", addr, err) + } + defer conn.Close() + + vol := NewRemote(conn) + defer vol.Close() + ctx := context.Background() + sid := "copy-remote-test" + + _ = vol.WriteFile(ctx, sid, "src.txt", []byte("remote copy"), 0o644) + + result, err := vol.Copy(ctx, sid, "src.txt", "dst.txt", WithForceFull()) + if err != nil { + t.Fatalf("Copy: %v", err) + } + if result.FilesSynced < 1 { + t.Errorf("synced = %d", result.FilesSynced) + } + + data, _, err := vol.ReadFile(ctx, sid, "dst.txt") + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + if string(data) != "remote copy" { + t.Errorf("content = %q", data) + } + + _ = vol.Remove(ctx, sid, ".", true) +} + +func TestRemoteCopyDir(t *testing.T) { + addr := taiTestGRPC() + conn, err := grpc.NewClient(addr, + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + if err != nil { + t.Skipf("gRPC dial %s: %v", addr, err) + } + defer conn.Close() + + vol := NewRemote(conn) + defer vol.Close() + ctx := context.Background() + sid := "copy-remote-dir" + + _ = vol.MkdirAll(ctx, sid, "src/sub") + _ = vol.WriteFile(ctx, sid, "src/a.txt", []byte("aaa"), 0o644) + _ = vol.WriteFile(ctx, sid, "src/sub/b.txt", []byte("bbb"), 0o644) + + result, err := vol.Copy(ctx, sid, "src", "dst", WithForceFull()) + if err != nil { + t.Fatalf("Copy: %v", err) + } + if result.FilesSynced < 2 { + t.Errorf("synced = %d", result.FilesSynced) + } + + data, _, err := vol.ReadFile(ctx, sid, "dst/sub/b.txt") + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + if string(data) != "bbb" { + t.Errorf("content = %q", data) + } + + _ = vol.Remove(ctx, sid, ".", true) +} diff --git a/tai/workspace/copy.go b/tai/workspace/copy.go new file mode 100644 index 00000000..516e2cbc --- /dev/null +++ b/tai/workspace/copy.go @@ -0,0 +1,158 @@ +package workspace + +import ( + "context" + "io/fs" + "os" + "path/filepath" + + "github.com/yaoapp/yao/tai/volume" +) + +// Copy implements the FS.Copy method with 4-way dispatch: +// +// ws -> ws : Volume.Copy (server-side for remote, local copy for local) +// host -> ws : Volume.SyncPush +// ws -> host : Volume.SyncPull +// host -> host : os-level recursive copy +func (w *workspaceFS) Copy(src, dst string, opts ...volume.SyncOption) (*volume.SyncResult, error) { + srcURI := parseHostURI(src) + dstURI := parseHostURI(dst) + ctx := context.Background() + + switch { + case !srcURI.IsHost && !dstURI.IsHost: + return w.vol.Copy(ctx, w.session, srcURI.Path, dstURI.Path, opts...) + + case srcURI.IsHost && !dstURI.IsHost: + hostPath, err := resolveAbsHostPath(srcURI) + if err != nil { + return nil, err + } + info, err := os.Stat(hostPath) + if err != nil { + return nil, err + } + if !info.IsDir() { + return nil, w.pushSingleFile(ctx, hostPath, dstURI.Path, info) + } + pushOpts := append(sliceClone(opts), volume.WithRemotePath(dstURI.Path)) + return w.vol.SyncPush(ctx, w.session, hostPath, pushOpts...) + + case !srcURI.IsHost && dstURI.IsHost: + hostPath, err := resolveAbsHostPath(dstURI) + if err != nil { + return nil, err + } + srcInfo, statErr := w.vol.Stat(ctx, w.session, srcURI.Path) + if statErr != nil { + return nil, statErr + } + if !srcInfo.IsDir { + return nil, w.pullSingleFile(ctx, srcURI.Path, hostPath) + } + pullOpts := append(sliceClone(opts), volume.WithRemotePath(srcURI.Path)) + return w.vol.SyncPull(ctx, w.session, hostPath, pullOpts...) + + default: + srcPath, err := resolveAbsHostPath(srcURI) + if err != nil { + return nil, err + } + dstPath, err := resolveAbsHostPath(dstURI) + if err != nil { + return nil, err + } + cfg := volume.ApplySyncOpts(opts) + return nil, copyLocalToLocal(srcPath, dstPath, cfg.Excludes) + } +} + +// pushSingleFile reads a host file and writes it into the workspace at dstPath. +func (w *workspaceFS) pushSingleFile(ctx context.Context, hostPath, dstPath string, info os.FileInfo) error { + data, err := os.ReadFile(hostPath) + if err != nil { + return err + } + dir := filepath.Dir(dstPath) + if dir != "" && dir != "." { + if err := w.vol.MkdirAll(ctx, w.session, dir); err != nil { + return err + } + } + perm := info.Mode() + if perm == 0 { + perm = 0o644 + } + return w.vol.WriteFile(ctx, w.session, dstPath, data, perm) +} + +// pullSingleFile reads a workspace file and writes it to hostPath. +func (w *workspaceFS) pullSingleFile(ctx context.Context, srcPath, hostPath string) error { + data, perm, err := w.vol.ReadFile(ctx, w.session, srcPath) + if err != nil { + return err + } + if perm == 0 { + perm = 0o644 + } + if err := os.MkdirAll(filepath.Dir(hostPath), 0o755); err != nil { + return err + } + return os.WriteFile(hostPath, data, perm) +} + +func copyLocalToLocal(src, dst string, excludes []string) error { + info, err := os.Stat(src) + if err != nil { + return err + } + if !info.IsDir() { + data, err := os.ReadFile(src) + if err != nil { + return err + } + if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil { + return err + } + return os.WriteFile(dst, data, info.Mode()) + } + + return filepath.WalkDir(src, func(abs string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + rel, _ := filepath.Rel(src, abs) + if rel == "." { + return os.MkdirAll(dst, 0o755) + } + for _, p := range excludes { + if matched, _ := filepath.Match(p, filepath.Base(rel)); matched { + if d.IsDir() { + return filepath.SkipDir + } + return nil + } + } + target := filepath.Join(dst, rel) + if d.IsDir() { + return os.MkdirAll(target, 0o755) + } + data, err := os.ReadFile(abs) + if err != nil { + return err + } + fi, _ := d.Info() + perm := os.FileMode(0o644) + if fi != nil { + perm = fi.Mode() + } + return os.WriteFile(target, data, perm) + }) +} + +func sliceClone(opts []volume.SyncOption) []volume.SyncOption { + cp := make([]volume.SyncOption, len(opts)) + copy(cp, opts) + return cp +} diff --git a/tai/workspace/uri.go b/tai/workspace/uri.go new file mode 100644 index 00000000..e4ed9056 --- /dev/null +++ b/tai/workspace/uri.go @@ -0,0 +1,49 @@ +package workspace + +import ( + "fmt" + "os" + "path/filepath" + "strings" +) + +// hostURI holds the parsed result of a host URI. +type hostURI struct { + Scheme string // "local" or "tmp"; empty for workspace paths + Path string // resolved absolute path (host) or relative path (workspace) + IsHost bool +} + +// parseHostURI extracts scheme and path from a host URI string. +// +// "local:///abs/path" -> {Scheme:"local", Path:"/abs/path", IsHost:true} +// "tmp:///rel/path" -> {Scheme:"tmp", Path:"rel/path", IsHost:true} +// "some/workspace/path" -> {Scheme:"", Path:"some/workspace/path", IsHost:false} +func parseHostURI(raw string) hostURI { + switch { + case strings.HasPrefix(raw, "local:///"): + return hostURI{Scheme: "local", Path: strings.TrimPrefix(raw, "local:///"), IsHost: true} + case strings.HasPrefix(raw, "tmp:///"): + return hostURI{Scheme: "tmp", Path: strings.TrimPrefix(raw, "tmp:///"), IsHost: true} + default: + return hostURI{Path: raw} + } +} + +// resolveAbsHostPath converts a parsed hostURI into an absolute filesystem path. +// For "local" scheme, Path is already absolute (rooted at /). +// For "tmp" scheme, Path is relative to os.TempDir(). +func resolveAbsHostPath(u hostURI) (string, error) { + switch u.Scheme { + case "local": + abs := filepath.Clean("/" + u.Path) + return abs, nil + case "tmp": + if strings.Contains(u.Path, "..") { + return "", fmt.Errorf("path traversal not allowed in tmp:// URI") + } + return filepath.Join(os.TempDir(), u.Path), nil + default: + return "", fmt.Errorf("not a host URI: %q", u.Path) + } +} diff --git a/tai/workspace/workspace.go b/tai/workspace/workspace.go index 6af4e69d..5ff5b518 100644 --- a/tai/workspace/workspace.go +++ b/tai/workspace/workspace.go @@ -25,6 +25,12 @@ type FS interface { RemoveAll(name string) error Rename(oldname, newname string) error MkdirAll(name string, perm os.FileMode) error + + // Copy copies files between workspace paths and/or host paths. + // Host paths use "local:///" (absolute system path) or "tmp:///" (os.TempDir-relative). + // ws↔ws uses Volume.Copy (server-side for remote volumes, avoiding 2N network round-trips). + // Returns non-nil *SyncResult for host↔workspace and ws↔ws transfers; nil for host↔host. + Copy(src, dst string, opts ...volume.SyncOption) (*volume.SyncResult, error) } // New creates an FS backed by the given Volume for the specified session. diff --git a/workspace/jsapi/fs.go b/workspace/jsapi/fs.go index 227f21fe..c82444b9 100644 --- a/workspace/jsapi/fs.go +++ b/workspace/jsapi/fs.go @@ -402,8 +402,6 @@ func copyHandler(info *v8go.FunctionCallbackInfo, wsID string) *v8go.Value { src := args[0].String() dst := args[1].String() - srcIsHost := isHostURI(src) - dstIsHost := isHostURI(dst) var excludes []string force := false @@ -419,12 +417,7 @@ func copyHandler(info *v8go.FunctionCallbackInfo, wsID string) *v8go.Value { } } - vol, sid, err := workspace.M().Volume(ctx, wsID) - if err != nil { - return throwError(info, err.Error()) - } - - opts := []volume.SyncOption{} + var opts []volume.SyncOption if len(excludes) > 0 { opts = append(opts, volume.WithExcludes(excludes...)) } @@ -432,52 +425,22 @@ func copyHandler(info *v8go.FunctionCallbackInfo, wsID string) *v8go.Value { opts = append(opts, volume.WithForceFull()) } - switch { - case !srcIsHost && !dstIsHost: - if err := copyWithinWorkspace(ctx, wsID, src, dst); err != nil { - return throwError(info, err.Error()) - } - return v8go.Undefined(iso) + // Map JSAPI local:// (AppRoot-relative) to Go-layer local:/// (absolute) + src = mapHostURI(src) + dst = mapHostURI(dst) - case srcIsHost && !dstIsHost: - hostPath, err := resolveHostPath(src) - if err != nil { - return throwError(info, err.Error()) - } - opts = append(opts, volume.WithRemotePath(dst)) - result, err := vol.SyncPush(ctx, sid, hostPath, opts...) - if err != nil { - return throwError(info, err.Error()) - } - return syncResultToJS(info, result) - - case !srcIsHost && dstIsHost: - hostPath, err := resolveHostPath(dst) - if err != nil { - return throwError(info, err.Error()) - } - opts = append(opts, volume.WithRemotePath(src)) - result, err := vol.SyncPull(ctx, sid, hostPath, opts...) - if err != nil { - return throwError(info, err.Error()) - } - return syncResultToJS(info, result) - - case srcIsHost && dstIsHost: - srcPath, err := resolveHostPath(src) - if err != nil { - return throwError(info, err.Error()) - } - dstPath, err := resolveHostPath(dst) - if err != nil { - return throwError(info, err.Error()) - } - if err := copyLocalToLocal(srcPath, dstPath, excludes); err != nil { - return throwError(info, err.Error()) - } - return v8go.Undefined(iso) + wsFS, err := workspace.M().FS(ctx, wsID) + if err != nil { + return throwError(info, err.Error()) } + result, err := wsFS.Copy(src, dst, opts...) + if err != nil { + return throwError(info, err.Error()) + } + if result != nil { + return syncResultToJS(info, result) + } return v8go.Undefined(iso) } @@ -491,123 +454,38 @@ func syncResultToJS(info *v8go.FunctionCallbackInfo, r *volume.SyncResult) *v8go return val } -func copyWithinWorkspace(ctx context.Context, wsID, src, dst string) error { - fsys, err := workspace.M().FS(ctx, wsID) - if err != nil { - return err - } - - fi, err := fsys.Stat(src) - if err != nil { - return err - } - - if !fi.IsDir() { - data, err := fsys.ReadFile(src) - if err != nil { - return err - } - return fsys.WriteFile(dst, data, fi.Mode()) - } - - return fs.WalkDir(fsys, src, func(p string, d fs.DirEntry, err error) error { - if err != nil { - return err - } - rel, _ := filepath.Rel(src, p) - target := filepath.Join(dst, rel) - if d.IsDir() { - return fsys.MkdirAll(target, 0o755) - } - data, err := fsys.ReadFile(p) - if err != nil { - return err - } - info, _ := d.Info() - perm := os.FileMode(0o644) - if info != nil { - perm = info.Mode() - } - return fsys.WriteFile(target, data, perm) - }) -} - -func isHostURI(path string) bool { - return strings.HasPrefix(path, "local://") || strings.HasPrefix(path, "tmp://") -} - -func resolveHostPath(rawPath string) (string, error) { - if strings.HasPrefix(rawPath, "tmp://") { - rel := strings.TrimPrefix(rawPath, "tmp://") +// mapHostURI converts JSAPI host URIs to Go-layer absolute URIs. +// local://relative -> local:///{AppSource}/relative (with security checks) +// tmp://relative -> tmp:///relative (Go layer resolves os.TempDir) +// other -> unchanged (workspace-relative path) +func mapHostURI(raw string) string { + switch { + case strings.HasPrefix(raw, "local://"): + rel := strings.TrimPrefix(raw, "local://") if strings.Contains(rel, "..") { - return "", fmt.Errorf("path traversal not allowed") + return raw } - return filepath.Join(os.TempDir(), rel), nil - } - - appRoot := config.Conf.AppSource - rel := strings.TrimPrefix(rawPath, "local://") - if strings.Contains(rel, "..") { - return "", fmt.Errorf("path traversal not allowed") - } - abs := filepath.Join(appRoot, rel) - resolved, err := filepath.EvalSymlinks(abs) - if err != nil { - resolved = abs - } - if !strings.HasPrefix(resolved, appRoot) { - return "", fmt.Errorf("path escapes app root") - } - return resolved, nil -} - -func copyLocalToLocal(src, dst string, excludes []string) error { - info, err := os.Stat(src) - if err != nil { - return err - } - if !info.IsDir() { - data, err := os.ReadFile(src) + appRoot := config.Conf.AppSource + abs := filepath.Join(appRoot, rel) + resolved, err := filepath.EvalSymlinks(abs) if err != nil { - return err + resolved = abs } - if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil { - return err + if !strings.HasPrefix(resolved, appRoot) { + return raw } - return os.WriteFile(dst, data, info.Mode()) - } + return "local:///" + resolved - return filepath.WalkDir(src, func(abs string, d fs.DirEntry, err error) error { - if err != nil { - return err + case strings.HasPrefix(raw, "tmp://"): + rel := strings.TrimPrefix(raw, "tmp://") + if strings.Contains(rel, "..") { + return raw } - rel, _ := filepath.Rel(src, abs) - if rel == "." { - return os.MkdirAll(dst, 0o755) - } - for _, p := range excludes { - if matched, _ := filepath.Match(p, filepath.Base(rel)); matched { - if d.IsDir() { - return filepath.SkipDir - } - return nil - } - } - target := filepath.Join(dst, rel) - if d.IsDir() { - return os.MkdirAll(target, 0o755) - } - data, err := os.ReadFile(abs) - if err != nil { - return err - } - fi, _ := d.Info() - perm := os.FileMode(0o644) - if fi != nil { - perm = fi.Mode() - } - return os.WriteFile(target, data, perm) - }) + return "tmp:///" + rel + + default: + return raw + } } func parseStringArray(val *v8go.Value) []string {