diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index ed92566b..42c96f75 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -923,6 +923,188 @@ jobs: body: '✅ Sandbox Tests passed!' }); + # ============================================================================= + # Sandbox V2 Tests (requires Docker + Tai for dual-mode) + # ============================================================================= + SandboxV2Test: + runs-on: ubuntu-latest + strategy: + matrix: + go: ["1.25"] + if: > + ${{ github.event.workflow_run.event == 'pull_request' && + github.event.workflow_run.conclusion == 'success' }} + steps: + - name: "Download artifact" + uses: actions/github-script@v7 + with: + script: | + var artifacts = await github.rest.actions.listWorkflowRunArtifacts({ + owner: context.repo.owner, + repo: context.repo.repo, + run_id: ${{github.event.workflow_run.id }}, + }); + var matchArtifact = artifacts.data.artifacts.filter((artifact) => { + return artifact.name == "pr" + })[0]; + var download = await github.rest.actions.downloadArtifact({ + owner: context.repo.owner, + repo: context.repo.repo, + artifact_id: matchArtifact.id, + archive_format: 'zip', + }); + var fs = require('fs'); + fs.writeFileSync('${{github.workspace}}/pr.zip', Buffer.from(download.data)); + + - name: "Read NR & SHA" + run: | + unzip pr.zip + cat NR + cat SHA + echo HEAD=$(cat SHA) >> $GITHUB_ENV + echo NR=$(cat NR) >> $GITHUB_ENV + + - name: "Comment on PR" + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const { NR } = process.env + var issue_number = NR; + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: issue_number, + body: '🤖 Sandbox V2 Tests running (dual-mode: local + remote)...' + }); + + - name: Checkout Kun + uses: actions/checkout@v4 + with: + repository: yaoapp/kun + path: kun + + - name: Checkout Xun + uses: actions/checkout@v4 + with: + repository: yaoapp/xun + path: xun + + - name: Checkout Gou + uses: actions/checkout@v4 + with: + repository: yaoapp/gou + path: gou + + - name: Checkout V8Go + uses: actions/checkout@v4 + with: + repository: yaoapp/v8go + path: v8go + + - name: Unzip libv8 + run: | + files=$(find ./v8go -name "libv8*.zip") + for file in $files; do + dir=$(dirname "$file") + echo "Extracting $file to directory $dir" + unzip -o -d $dir $file + rm -rf $dir/__MACOSX + done + + - name: Checkout Demo App + uses: actions/checkout@v4 + with: + repository: yaoapp/yao-dev-app + path: app + + - name: Checkout Extension + uses: actions/checkout@v4 + with: + repository: yaoapp/yao-extensions-dev + path: extension + + - name: Move Dependencies + run: | + mv kun ../ + mv xun ../ + mv gou ../ + mv v8go ../ + mv app ../ + mv extension ../ + + - name: Checkout pull request HEAD commit + uses: actions/checkout@v4 + with: + ref: ${{ env.HEAD }} + + - name: Setup Go ${{ matrix.go }} + uses: actions/setup-go@v5 + with: + go-version: ${{ matrix.go }} + + - name: Setup Go Tools + run: make tools + + - name: Pull Test Images + run: | + docker pull yaoapp/sandbox-v2-test:latest || true + docker pull yaoapp/tai:latest + + - name: Start Tai Server (Docker proxy for remote mode) + run: | + docker run -d --name tai \ + -v /var/run/docker.sock:/var/run/docker.sock \ + -p 8080:8080 -p 9100:9100 -p 2375:2375 -p 6080:6080 \ + yaoapp/tai:latest + + for i in $(seq 1 30); do + if curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1; then + echo "Tai HTTP is ready" + break + fi + echo "Waiting for Tai HTTP... ($i)" + sleep 1 + done + + for i in $(seq 1 15); do + if nc -z 127.0.0.1 9100 2>/dev/null; then + echo "Tai gRPC is ready" + break + fi + echo "Waiting for Tai gRPC... ($i)" + sleep 1 + done + + - name: "Run Sandbox V2 Tests (dual-mode: local + remote)" + env: + SANDBOX_TEST_IMAGE: yaoapp/sandbox-v2-test:latest + SANDBOX_TEST_REMOTE_ADDR: "tai://127.0.0.1:9100" + run: make unit-test-sandbox-v2 + + - name: Codecov Report + if: always() + uses: codecov/codecov-action@v4 + with: + token: ${{ secrets.CODECOV_TOKEN }} + files: sandbox/v2/coverage.out + flags: sandbox-v2 + fail_ci_if_error: false + + - name: "Comment on PR - Sandbox V2 Tests Done" + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const { NR } = process.env + var issue_number = NR; + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: issue_number, + body: '✅ Sandbox V2 Tests passed!' + }); + # ============================================================================= # Benchmark & Memory Leak Tests - Run with MySQL8.0 and SQLite3 # ============================================================================= diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index 07d7111c..a68ce673 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -679,6 +679,126 @@ jobs: with: token: ${{ secrets.CODECOV_TOKEN }} + # ============================================================================= + # Sandbox V2 Tests (requires Docker + Tai for dual-mode) + # ============================================================================= + sandbox-v2-test: + runs-on: ubuntu-latest + strategy: + matrix: + go: ["1.25"] + steps: + - name: Checkout Kun + uses: actions/checkout@v4 + with: + repository: ${{ env.REPO_KUN }} + path: kun + + - name: Checkout Xun + uses: actions/checkout@v4 + with: + repository: ${{ env.REPO_XUN }} + path: xun + + - name: Checkout Gou + uses: actions/checkout@v4 + with: + repository: ${{ env.REPO_GOU }} + path: gou + + - name: Checkout V8Go + uses: actions/checkout@v4 + with: + repository: yaoapp/v8go + path: v8go + + - name: Unzip libv8 + run: | + files=$(find ./v8go -name "libv8*.zip") + for file in $files; do + dir=$(dirname "$file") + echo "Extracting $file to directory $dir" + unzip -o -d $dir $file + rm -rf $dir/__MACOSX + done + + - name: Checkout Demo App + uses: actions/checkout@v4 + with: + repository: yaoapp/yao-dev-app + path: app + + - name: Checkout Extension + uses: actions/checkout@v4 + with: + repository: yaoapp/yao-extensions-dev + path: extension + + - name: Move Dependencies + run: | + mv kun ../ + mv xun ../ + mv gou ../ + mv v8go ../ + mv app ../ + mv extension ../ + + - name: Checkout Code + uses: actions/checkout@v4 + + - name: Setup Go ${{ matrix.go }} + uses: actions/setup-go@v5 + with: + go-version: ${{ matrix.go }} + + - name: Setup Go Tools + run: make tools + + - name: Pull Test Images + run: | + docker pull yaoapp/sandbox-v2-test:latest || true + docker pull yaoapp/tai:latest + + - name: Start Tai Server (Docker proxy for remote mode) + run: | + docker run -d --name tai \ + -v /var/run/docker.sock:/var/run/docker.sock \ + -p 8080:8080 -p 9100:9100 -p 2375:2375 -p 6080:6080 \ + yaoapp/tai:latest + + for i in $(seq 1 30); do + if curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1; then + echo "Tai HTTP is ready" + break + fi + echo "Waiting for Tai HTTP... ($i)" + sleep 1 + done + + for i in $(seq 1 15); do + if nc -z 127.0.0.1 9100 2>/dev/null; then + echo "Tai gRPC is ready" + break + fi + echo "Waiting for Tai gRPC... ($i)" + sleep 1 + done + + - name: "Run Sandbox V2 Tests (dual-mode: local + remote)" + env: + SANDBOX_TEST_IMAGE: yaoapp/sandbox-v2-test:latest + SANDBOX_TEST_REMOTE_ADDR: "tai://127.0.0.1:9100" + run: make unit-test-sandbox-v2 + + - name: Codecov Report + if: always() + uses: codecov/codecov-action@v4 + with: + token: ${{ secrets.CODECOV_TOKEN }} + files: sandbox/v2/coverage.out + flags: sandbox-v2 + fail_ci_if_error: false + # ============================================================================= # Benchmark & Memory Leak Tests - Run with MySQL8.0 and SQLite3 # ============================================================================= diff --git a/.gitignore b/.gitignore index 5c81f9b5..e1541c9f 100644 --- a/.gitignore +++ b/.gitignore @@ -75,3 +75,7 @@ tg-send registry/data/ registry/manager/DESIGN*.md tai/testdata/ +sandbox/v2/docker/base/*-amd64 +sandbox/v2/docker/base/*-arm64 +!sandbox/v2/docker/*.sh +!sandbox/v2/docker/*/*.sh \ No newline at end of file diff --git a/Makefile b/Makefile index 68790267..be2206dc 100644 --- a/Makefile +++ b/Makefile @@ -19,8 +19,8 @@ TESTFOLDER_AGENT := $(shell $(GO) list ./agent/... ./aigc/... | grep -vE 'agent/ TESTFOLDER_KB := $(shell $(GO) list ./kb/...) # Robot tests (agent/robot/... packages, excluding events/integrations which require Telegram etc.) TESTFOLDER_ROBOT := $(shell $(GO) list ./agent/robot/... | grep -vE 'agent/robot/events') -# Sandbox tests (requires Docker) -TESTFOLDER_SANDBOX := $(shell $(GO) list ./sandbox/...) +# Sandbox tests (requires Docker) — excludes sandbox/v2 (has its own job) +TESTFOLDER_SANDBOX := $(shell $(GO) list ./sandbox/... | grep -v 'sandbox/v2') # Tai SDK tests (requires Tai container with Docker socket) TESTFOLDER_TAI := $(shell $(GO) list ./tai/...) # gRPC tests @@ -199,6 +199,17 @@ unit-test-registry: rm profile.out; \ fi +# Sandbox V2 Unit Test (requires Docker; optionally Tai for remote mode) +.PHONY: unit-test-sandbox-v2 +unit-test-sandbox-v2: + @echo "" + @echo "=============================================" + @echo "Running Sandbox V2 Tests..." + @echo "=============================================" + docker pull $(SANDBOX_V2_IMAGE) || true + $(MAKE) -C sandbox/v2 test-ci TEST_IMAGE=$(SANDBOX_V2_IMAGE) +SANDBOX_V2_IMAGE ?= yaoapp/sandbox-v2-test:latest + # Sandbox Unit Test (requires Docker) .PHONY: unit-test-sandbox unit-test-sandbox: diff --git a/grpc/auth/endpoint.go b/grpc/auth/endpoint.go index 67fcaf6e..afebcfde 100644 --- a/grpc/auth/endpoint.go +++ b/grpc/auth/endpoint.go @@ -60,6 +60,9 @@ func VirtualEndpoint(fullMethod string, req interface{}) (method string, path st } return "POST", "/grpc/agent/" + case "/yao.Yao/Heartbeat": + return "POST", "/grpc/heartbeat" + default: return "POST", "/grpc/unknown" } diff --git a/grpc/auth/endpoint_test.go b/grpc/auth/endpoint_test.go index edbe3c1e..123c6697 100644 --- a/grpc/auth/endpoint_test.go +++ b/grpc/auth/endpoint_test.go @@ -134,3 +134,15 @@ func TestVirtualEndpoint_AgentStreamEmptyID(t *testing.T) { assert.Equal(t, "POST", method) assert.Equal(t, "/grpc/agent/", path) } + +func TestVirtualEndpoint_Heartbeat(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/Heartbeat", &pb.HeartbeatRequest{SandboxId: "sb-1"}) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/heartbeat", path) +} + +func TestVirtualEndpoint_HeartbeatNilReq(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/Heartbeat", nil) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/heartbeat", path) +} diff --git a/grpc/auth/scope.go b/grpc/auth/scope.go index 4957c482..71135925 100644 --- a/grpc/auth/scope.go +++ b/grpc/auth/scope.go @@ -7,7 +7,7 @@ func init() { &acl.ScopeDefinition{Name: "grpc:run", Endpoints: []string{"POST /grpc/run/*", "POST /grpc/run/"}}, &acl.ScopeDefinition{Name: "grpc:stream", Endpoints: []string{"POST /grpc/stream/*", "POST /grpc/stream/"}}, &acl.ScopeDefinition{Name: "grpc:shell", Endpoints: []string{"POST /grpc/shell"}}, - &acl.ScopeDefinition{Name: "grpc:mcp", Endpoints: []string{"GET /grpc/mcp/tools", "POST /grpc/mcp/call/*", "POST /grpc/mcp/call/", "GET /grpc/mcp/resources", "GET /grpc/mcp/resources/read"}}, + &acl.ScopeDefinition{Name: "grpc:mcp", Endpoints: []string{"GET /grpc/mcp/tools", "POST /grpc/mcp/call/*", "POST /grpc/mcp/call/", "GET /grpc/mcp/resources", "GET /grpc/mcp/resources/read", "POST /grpc/heartbeat"}}, &acl.ScopeDefinition{Name: "grpc:llm", Endpoints: []string{"POST /grpc/llm/completions"}}, &acl.ScopeDefinition{Name: "grpc:agent", Endpoints: []string{"POST /grpc/agent/*", "POST /grpc/agent/"}}, ) diff --git a/grpc/grpc.go b/grpc/grpc.go index 22755275..1f204034 100644 --- a/grpc/grpc.go +++ b/grpc/grpc.go @@ -21,6 +21,7 @@ import ( mcphandler "github.com/yaoapp/yao/grpc/mcp" "github.com/yaoapp/yao/grpc/pb" runhandler "github.com/yaoapp/yao/grpc/run" + sandboxhandler "github.com/yaoapp/yao/grpc/sandbox" shellhandler "github.com/yaoapp/yao/grpc/shell" ) @@ -33,13 +34,14 @@ var ( type yaoServer struct { pb.UnimplementedYaoServer - health health.Handler - run runhandler.Handler - shell shellhandler.Handler - api apihandler.Handler - mcp mcphandler.Handler - llm llmhandler.Handler - agent agenthandler.Handler + health health.Handler + run runhandler.Handler + shell shellhandler.Handler + api apihandler.Handler + mcp mcphandler.Handler + llm llmhandler.Handler + agent agenthandler.Handler + sandbox *sandboxhandler.Handler } // ── Health ─────────────────────────────────────────────────────────────────── @@ -107,6 +109,30 @@ func (s *yaoServer) AgentStream(req *pb.AgentRequest, stream grpc.ServerStreamin return s.agent.AgentStream(req, stream) } +// ── Sandbox ────────────────────────────────────────────────────────────────── + +func (s *yaoServer) Heartbeat(ctx context.Context, req *pb.HeartbeatRequest) (*pb.HeartbeatResponse, error) { + if s.sandbox == nil { + return &pb.HeartbeatResponse{Action: "ok"}, nil + } + return s.sandbox.Heartbeat(ctx, req) +} + +// SandboxHandler returns the sandbox handler for external access (e.g., Manager integration). +func SandboxHandler() *sandboxhandler.Handler { + mu.Lock() + defer mu.Unlock() + return sandboxH +} + +var sandboxH *sandboxhandler.Handler + +// SetSandboxOnBeat sets the heartbeat callback for the sandbox handler. +// Must be called before StartServer. +func SetSandboxOnBeat(fn func(data *sandboxhandler.HeartbeatData) string) { + sandboxH = sandboxhandler.NewHandler(fn) +} + // ── Server lifecycle ───────────────────────────────────────────────────────── // StartServer initializes and starts the gRPC server based on config. @@ -124,7 +150,10 @@ func StartServer(cfg config.Config) error { grpc.ChainUnaryInterceptor(auth.UnaryInterceptor), grpc.ChainStreamInterceptor(auth.StreamInterceptor), ) - pb.RegisterYaoServer(server, &yaoServer{}) + if sandboxH == nil { + sandboxH = sandboxhandler.NewHandler(nil) + } + pb.RegisterYaoServer(server, &yaoServer{sandbox: sandboxH}) hosts := strings.Split(cfg.GRPC.Host, ",") port := strconv.Itoa(cfg.GRPC.Port) diff --git a/grpc/pb/yao.pb.go b/grpc/pb/yao.pb.go index cea122a1..fa32b95e 100644 --- a/grpc/pb/yao.pb.go +++ b/grpc/pb/yao.pb.go @@ -2,7 +2,7 @@ // versions: // protoc-gen-go v1.36.11 // protoc v4.25.0 -// source: yao.proto +// source: grpc/pb/yao.proto package pb @@ -32,7 +32,7 @@ type RunRequest struct { func (x *RunRequest) Reset() { *x = RunRequest{} - mi := &file_yao_proto_msgTypes[0] + mi := &file_grpc_pb_yao_proto_msgTypes[0] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -44,7 +44,7 @@ func (x *RunRequest) String() string { func (*RunRequest) ProtoMessage() {} func (x *RunRequest) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[0] + mi := &file_grpc_pb_yao_proto_msgTypes[0] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -57,7 +57,7 @@ func (x *RunRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use RunRequest.ProtoReflect.Descriptor instead. func (*RunRequest) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{0} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{0} } func (x *RunRequest) GetProcess() string { @@ -90,7 +90,7 @@ type RunResponse struct { func (x *RunResponse) Reset() { *x = RunResponse{} - mi := &file_yao_proto_msgTypes[1] + mi := &file_grpc_pb_yao_proto_msgTypes[1] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -102,7 +102,7 @@ func (x *RunResponse) String() string { func (*RunResponse) ProtoMessage() {} func (x *RunResponse) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[1] + mi := &file_grpc_pb_yao_proto_msgTypes[1] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -115,7 +115,7 @@ func (x *RunResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use RunResponse.ProtoReflect.Descriptor instead. func (*RunResponse) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{1} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{1} } func (x *RunResponse) GetData() []byte { @@ -135,7 +135,7 @@ type Chunk struct { func (x *Chunk) Reset() { *x = Chunk{} - mi := &file_yao_proto_msgTypes[2] + mi := &file_grpc_pb_yao_proto_msgTypes[2] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -147,7 +147,7 @@ func (x *Chunk) String() string { func (*Chunk) ProtoMessage() {} func (x *Chunk) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[2] + mi := &file_grpc_pb_yao_proto_msgTypes[2] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -160,7 +160,7 @@ func (x *Chunk) ProtoReflect() protoreflect.Message { // Deprecated: Use Chunk.ProtoReflect.Descriptor instead. func (*Chunk) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{2} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{2} } func (x *Chunk) GetData() []byte { @@ -189,7 +189,7 @@ type ShellRequest struct { func (x *ShellRequest) Reset() { *x = ShellRequest{} - mi := &file_yao_proto_msgTypes[3] + mi := &file_grpc_pb_yao_proto_msgTypes[3] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -201,7 +201,7 @@ func (x *ShellRequest) String() string { func (*ShellRequest) ProtoMessage() {} func (x *ShellRequest) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[3] + mi := &file_grpc_pb_yao_proto_msgTypes[3] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -214,7 +214,7 @@ func (x *ShellRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use ShellRequest.ProtoReflect.Descriptor instead. func (*ShellRequest) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{3} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{3} } func (x *ShellRequest) GetCommand() string { @@ -256,7 +256,7 @@ type ShellResponse struct { func (x *ShellResponse) Reset() { *x = ShellResponse{} - mi := &file_yao_proto_msgTypes[4] + mi := &file_grpc_pb_yao_proto_msgTypes[4] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -268,7 +268,7 @@ func (x *ShellResponse) String() string { func (*ShellResponse) ProtoMessage() {} func (x *ShellResponse) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[4] + mi := &file_grpc_pb_yao_proto_msgTypes[4] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -281,7 +281,7 @@ func (x *ShellResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use ShellResponse.ProtoReflect.Descriptor instead. func (*ShellResponse) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{4} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{4} } func (x *ShellResponse) GetStdout() []byte { @@ -317,7 +317,7 @@ type APIRequest struct { func (x *APIRequest) Reset() { *x = APIRequest{} - mi := &file_yao_proto_msgTypes[5] + mi := &file_grpc_pb_yao_proto_msgTypes[5] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -329,7 +329,7 @@ func (x *APIRequest) String() string { func (*APIRequest) ProtoMessage() {} func (x *APIRequest) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[5] + mi := &file_grpc_pb_yao_proto_msgTypes[5] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -342,7 +342,7 @@ func (x *APIRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use APIRequest.ProtoReflect.Descriptor instead. func (*APIRequest) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{5} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{5} } func (x *APIRequest) GetMethod() string { @@ -384,7 +384,7 @@ type APIResponse struct { func (x *APIResponse) Reset() { *x = APIResponse{} - mi := &file_yao_proto_msgTypes[6] + mi := &file_grpc_pb_yao_proto_msgTypes[6] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -396,7 +396,7 @@ func (x *APIResponse) String() string { func (*APIResponse) ProtoMessage() {} func (x *APIResponse) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[6] + mi := &file_grpc_pb_yao_proto_msgTypes[6] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -409,7 +409,7 @@ func (x *APIResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use APIResponse.ProtoReflect.Descriptor instead. func (*APIResponse) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{6} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{6} } func (x *APIResponse) GetStatus() int32 { @@ -442,7 +442,7 @@ type MCPListRequest struct { func (x *MCPListRequest) Reset() { *x = MCPListRequest{} - mi := &file_yao_proto_msgTypes[7] + mi := &file_grpc_pb_yao_proto_msgTypes[7] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -454,7 +454,7 @@ func (x *MCPListRequest) String() string { func (*MCPListRequest) ProtoMessage() {} func (x *MCPListRequest) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[7] + mi := &file_grpc_pb_yao_proto_msgTypes[7] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -467,7 +467,7 @@ func (x *MCPListRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use MCPListRequest.ProtoReflect.Descriptor instead. func (*MCPListRequest) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{7} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{7} } func (x *MCPListRequest) GetSessionId() string { @@ -486,7 +486,7 @@ type MCPListResponse struct { func (x *MCPListResponse) Reset() { *x = MCPListResponse{} - mi := &file_yao_proto_msgTypes[8] + mi := &file_grpc_pb_yao_proto_msgTypes[8] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -498,7 +498,7 @@ func (x *MCPListResponse) String() string { func (*MCPListResponse) ProtoMessage() {} func (x *MCPListResponse) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[8] + mi := &file_grpc_pb_yao_proto_msgTypes[8] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -511,7 +511,7 @@ func (x *MCPListResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use MCPListResponse.ProtoReflect.Descriptor instead. func (*MCPListResponse) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{8} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{8} } func (x *MCPListResponse) GetTools() []byte { @@ -532,7 +532,7 @@ type MCPCallRequest struct { func (x *MCPCallRequest) Reset() { *x = MCPCallRequest{} - mi := &file_yao_proto_msgTypes[9] + mi := &file_grpc_pb_yao_proto_msgTypes[9] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -544,7 +544,7 @@ func (x *MCPCallRequest) String() string { func (*MCPCallRequest) ProtoMessage() {} func (x *MCPCallRequest) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[9] + mi := &file_grpc_pb_yao_proto_msgTypes[9] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -557,7 +557,7 @@ func (x *MCPCallRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use MCPCallRequest.ProtoReflect.Descriptor instead. func (*MCPCallRequest) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{9} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{9} } func (x *MCPCallRequest) GetSessionId() string { @@ -590,7 +590,7 @@ type MCPCallResponse struct { func (x *MCPCallResponse) Reset() { *x = MCPCallResponse{} - mi := &file_yao_proto_msgTypes[10] + mi := &file_grpc_pb_yao_proto_msgTypes[10] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -602,7 +602,7 @@ func (x *MCPCallResponse) String() string { func (*MCPCallResponse) ProtoMessage() {} func (x *MCPCallResponse) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[10] + mi := &file_grpc_pb_yao_proto_msgTypes[10] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -615,7 +615,7 @@ func (x *MCPCallResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use MCPCallResponse.ProtoReflect.Descriptor instead. func (*MCPCallResponse) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{10} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{10} } func (x *MCPCallResponse) GetResult() []byte { @@ -634,7 +634,7 @@ type MCPResourcesResponse struct { func (x *MCPResourcesResponse) Reset() { *x = MCPResourcesResponse{} - mi := &file_yao_proto_msgTypes[11] + mi := &file_grpc_pb_yao_proto_msgTypes[11] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -646,7 +646,7 @@ func (x *MCPResourcesResponse) String() string { func (*MCPResourcesResponse) ProtoMessage() {} func (x *MCPResourcesResponse) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[11] + mi := &file_grpc_pb_yao_proto_msgTypes[11] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -659,7 +659,7 @@ func (x *MCPResourcesResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use MCPResourcesResponse.ProtoReflect.Descriptor instead. func (*MCPResourcesResponse) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{11} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{11} } func (x *MCPResourcesResponse) GetResources() []byte { @@ -679,7 +679,7 @@ type MCPResourceRequest struct { func (x *MCPResourceRequest) Reset() { *x = MCPResourceRequest{} - mi := &file_yao_proto_msgTypes[12] + mi := &file_grpc_pb_yao_proto_msgTypes[12] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -691,7 +691,7 @@ func (x *MCPResourceRequest) String() string { func (*MCPResourceRequest) ProtoMessage() {} func (x *MCPResourceRequest) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[12] + mi := &file_grpc_pb_yao_proto_msgTypes[12] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -704,7 +704,7 @@ func (x *MCPResourceRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use MCPResourceRequest.ProtoReflect.Descriptor instead. func (*MCPResourceRequest) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{12} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{12} } func (x *MCPResourceRequest) GetSessionId() string { @@ -730,7 +730,7 @@ type MCPResourceResponse struct { func (x *MCPResourceResponse) Reset() { *x = MCPResourceResponse{} - mi := &file_yao_proto_msgTypes[13] + mi := &file_grpc_pb_yao_proto_msgTypes[13] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -742,7 +742,7 @@ func (x *MCPResourceResponse) String() string { func (*MCPResourceResponse) ProtoMessage() {} func (x *MCPResourceResponse) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[13] + mi := &file_grpc_pb_yao_proto_msgTypes[13] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -755,7 +755,7 @@ func (x *MCPResourceResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use MCPResourceResponse.ProtoReflect.Descriptor instead. func (*MCPResourceResponse) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{13} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{13} } func (x *MCPResourceResponse) GetContents() []byte { @@ -776,7 +776,7 @@ type ChatRequest struct { func (x *ChatRequest) Reset() { *x = ChatRequest{} - mi := &file_yao_proto_msgTypes[14] + mi := &file_grpc_pb_yao_proto_msgTypes[14] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -788,7 +788,7 @@ func (x *ChatRequest) String() string { func (*ChatRequest) ProtoMessage() {} func (x *ChatRequest) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[14] + mi := &file_grpc_pb_yao_proto_msgTypes[14] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -801,7 +801,7 @@ func (x *ChatRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use ChatRequest.ProtoReflect.Descriptor instead. func (*ChatRequest) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{14} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{14} } func (x *ChatRequest) GetConnector() string { @@ -834,7 +834,7 @@ type ChatResponse struct { func (x *ChatResponse) Reset() { *x = ChatResponse{} - mi := &file_yao_proto_msgTypes[15] + mi := &file_grpc_pb_yao_proto_msgTypes[15] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -846,7 +846,7 @@ func (x *ChatResponse) String() string { func (*ChatResponse) ProtoMessage() {} func (x *ChatResponse) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[15] + mi := &file_grpc_pb_yao_proto_msgTypes[15] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -859,7 +859,7 @@ func (x *ChatResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use ChatResponse.ProtoReflect.Descriptor instead. func (*ChatResponse) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{15} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{15} } func (x *ChatResponse) GetData() []byte { @@ -879,7 +879,7 @@ type ChatChunk struct { func (x *ChatChunk) Reset() { *x = ChatChunk{} - mi := &file_yao_proto_msgTypes[16] + mi := &file_grpc_pb_yao_proto_msgTypes[16] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -891,7 +891,7 @@ func (x *ChatChunk) String() string { func (*ChatChunk) ProtoMessage() {} func (x *ChatChunk) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[16] + mi := &file_grpc_pb_yao_proto_msgTypes[16] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -904,7 +904,7 @@ func (x *ChatChunk) ProtoReflect() protoreflect.Message { // Deprecated: Use ChatChunk.ProtoReflect.Descriptor instead. func (*ChatChunk) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{16} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{16} } func (x *ChatChunk) GetData() []byte { @@ -932,7 +932,7 @@ type AgentRequest struct { func (x *AgentRequest) Reset() { *x = AgentRequest{} - mi := &file_yao_proto_msgTypes[17] + mi := &file_grpc_pb_yao_proto_msgTypes[17] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -944,7 +944,7 @@ func (x *AgentRequest) String() string { func (*AgentRequest) ProtoMessage() {} func (x *AgentRequest) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[17] + mi := &file_grpc_pb_yao_proto_msgTypes[17] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -957,7 +957,7 @@ func (x *AgentRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use AgentRequest.ProtoReflect.Descriptor instead. func (*AgentRequest) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{17} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{17} } func (x *AgentRequest) GetAssistantId() string { @@ -992,7 +992,7 @@ type AgentChunk struct { func (x *AgentChunk) Reset() { *x = AgentChunk{} - mi := &file_yao_proto_msgTypes[18] + mi := &file_grpc_pb_yao_proto_msgTypes[18] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1004,7 +1004,7 @@ func (x *AgentChunk) String() string { func (*AgentChunk) ProtoMessage() {} func (x *AgentChunk) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[18] + mi := &file_grpc_pb_yao_proto_msgTypes[18] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1017,7 +1017,7 @@ func (x *AgentChunk) ProtoReflect() protoreflect.Message { // Deprecated: Use AgentChunk.ProtoReflect.Descriptor instead. func (*AgentChunk) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{18} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{18} } func (x *AgentChunk) GetData() []byte { @@ -1042,7 +1042,7 @@ type Empty struct { func (x *Empty) Reset() { *x = Empty{} - mi := &file_yao_proto_msgTypes[19] + mi := &file_grpc_pb_yao_proto_msgTypes[19] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1054,7 +1054,7 @@ func (x *Empty) String() string { func (*Empty) ProtoMessage() {} func (x *Empty) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[19] + mi := &file_grpc_pb_yao_proto_msgTypes[19] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1067,7 +1067,7 @@ func (x *Empty) ProtoReflect() protoreflect.Message { // Deprecated: Use Empty.ProtoReflect.Descriptor instead. func (*Empty) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{19} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{19} } type HealthzResponse struct { @@ -1079,7 +1079,7 @@ type HealthzResponse struct { func (x *HealthzResponse) Reset() { *x = HealthzResponse{} - mi := &file_yao_proto_msgTypes[20] + mi := &file_grpc_pb_yao_proto_msgTypes[20] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1091,7 +1091,7 @@ func (x *HealthzResponse) String() string { func (*HealthzResponse) ProtoMessage() {} func (x *HealthzResponse) ProtoReflect() protoreflect.Message { - mi := &file_yao_proto_msgTypes[20] + mi := &file_grpc_pb_yao_proto_msgTypes[20] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1104,7 +1104,7 @@ func (x *HealthzResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use HealthzResponse.ProtoReflect.Descriptor instead. func (*HealthzResponse) Descriptor() ([]byte, []int) { - return file_yao_proto_rawDescGZIP(), []int{20} + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{20} } func (x *HealthzResponse) GetStatus() string { @@ -1114,11 +1114,123 @@ func (x *HealthzResponse) GetStatus() string { return "" } -var File_yao_proto protoreflect.FileDescriptor +type HeartbeatRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + SandboxId string `protobuf:"bytes,1,opt,name=sandbox_id,json=sandboxId,proto3" json:"sandbox_id,omitempty"` + CpuPercent int32 `protobuf:"varint,2,opt,name=cpu_percent,json=cpuPercent,proto3" json:"cpu_percent,omitempty"` + MemBytes int64 `protobuf:"varint,3,opt,name=mem_bytes,json=memBytes,proto3" json:"mem_bytes,omitempty"` + RunningProcs int32 `protobuf:"varint,4,opt,name=running_procs,json=runningProcs,proto3" json:"running_procs,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} -const file_yao_proto_rawDesc = "" + +func (x *HeartbeatRequest) Reset() { + *x = HeartbeatRequest{} + mi := &file_grpc_pb_yao_proto_msgTypes[21] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *HeartbeatRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HeartbeatRequest) ProtoMessage() {} + +func (x *HeartbeatRequest) ProtoReflect() protoreflect.Message { + mi := &file_grpc_pb_yao_proto_msgTypes[21] + 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 HeartbeatRequest.ProtoReflect.Descriptor instead. +func (*HeartbeatRequest) Descriptor() ([]byte, []int) { + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{21} +} + +func (x *HeartbeatRequest) GetSandboxId() string { + if x != nil { + return x.SandboxId + } + return "" +} + +func (x *HeartbeatRequest) GetCpuPercent() int32 { + if x != nil { + return x.CpuPercent + } + return 0 +} + +func (x *HeartbeatRequest) GetMemBytes() int64 { + if x != nil { + return x.MemBytes + } + return 0 +} + +func (x *HeartbeatRequest) GetRunningProcs() int32 { + if x != nil { + return x.RunningProcs + } + return 0 +} + +type HeartbeatResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Action string `protobuf:"bytes,1,opt,name=action,proto3" json:"action,omitempty"` // "ok" or "shutdown" + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *HeartbeatResponse) Reset() { + *x = HeartbeatResponse{} + mi := &file_grpc_pb_yao_proto_msgTypes[22] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *HeartbeatResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HeartbeatResponse) ProtoMessage() {} + +func (x *HeartbeatResponse) ProtoReflect() protoreflect.Message { + mi := &file_grpc_pb_yao_proto_msgTypes[22] + 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 HeartbeatResponse.ProtoReflect.Descriptor instead. +func (*HeartbeatResponse) Descriptor() ([]byte, []int) { + return file_grpc_pb_yao_proto_rawDescGZIP(), []int{22} +} + +func (x *HeartbeatResponse) GetAction() string { + if x != nil { + return x.Action + } + return "" +} + +var File_grpc_pb_yao_proto protoreflect.FileDescriptor + +const file_grpc_pb_yao_proto_rawDesc = "" + "\n" + - "\tyao.proto\x12\x03yao\"T\n" + + "\x11grpc/pb/yao.proto\x12\x03yao\"T\n" + "\n" + "RunRequest\x12\x18\n" + "\aprocess\x18\x01 \x01(\tR\aprocess\x12\x12\n" + @@ -1196,7 +1308,16 @@ const file_yao_proto_rawDesc = "" + "\x04done\x18\x02 \x01(\bR\x04done\"\a\n" + "\x05Empty\")\n" + "\x0fHealthzResponse\x12\x16\n" + - "\x06status\x18\x01 \x01(\tR\x06status2\xb8\x05\n" + + "\x06status\x18\x01 \x01(\tR\x06status\"\x94\x01\n" + + "\x10HeartbeatRequest\x12\x1d\n" + + "\n" + + "sandbox_id\x18\x01 \x01(\tR\tsandboxId\x12\x1f\n" + + "\vcpu_percent\x18\x02 \x01(\x05R\n" + + "cpuPercent\x12\x1b\n" + + "\tmem_bytes\x18\x03 \x01(\x03R\bmemBytes\x12#\n" + + "\rrunning_procs\x18\x04 \x01(\x05R\frunningProcs\"+\n" + + "\x11HeartbeatResponse\x12\x16\n" + + "\x06action\x18\x01 \x01(\tR\x06action2\xf4\x05\n" + "\x03Yao\x12(\n" + "\x03Run\x12\x0f.yao.RunRequest\x1a\x10.yao.RunResponse\x12'\n" + "\x06Stream\x12\x0f.yao.RunRequest\x1a\n" + @@ -1213,22 +1334,23 @@ const file_yao_proto_rawDesc = "" + "\x15ChatCompletionsStream\x12\x10.yao.ChatRequest\x1a\x0e.yao.ChatChunk0\x01\x123\n" + "\vAgentStream\x12\x11.yao.AgentRequest\x1a\x0f.yao.AgentChunk0\x01\x12+\n" + "\aHealthz\x12\n" + - ".yao.Empty\x1a\x14.yao.HealthzResponseB\x1fZ\x1dgithub.com/yaoapp/yao/grpc/pbb\x06proto3" + ".yao.Empty\x1a\x14.yao.HealthzResponse\x12:\n" + + "\tHeartbeat\x12\x15.yao.HeartbeatRequest\x1a\x16.yao.HeartbeatResponseB\x1fZ\x1dgithub.com/yaoapp/yao/grpc/pbb\x06proto3" var ( - file_yao_proto_rawDescOnce sync.Once - file_yao_proto_rawDescData []byte + file_grpc_pb_yao_proto_rawDescOnce sync.Once + file_grpc_pb_yao_proto_rawDescData []byte ) -func file_yao_proto_rawDescGZIP() []byte { - file_yao_proto_rawDescOnce.Do(func() { - file_yao_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_yao_proto_rawDesc), len(file_yao_proto_rawDesc))) +func file_grpc_pb_yao_proto_rawDescGZIP() []byte { + file_grpc_pb_yao_proto_rawDescOnce.Do(func() { + file_grpc_pb_yao_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_grpc_pb_yao_proto_rawDesc), len(file_grpc_pb_yao_proto_rawDesc))) }) - return file_yao_proto_rawDescData + return file_grpc_pb_yao_proto_rawDescData } -var file_yao_proto_msgTypes = make([]protoimpl.MessageInfo, 24) -var file_yao_proto_goTypes = []any{ +var file_grpc_pb_yao_proto_msgTypes = make([]protoimpl.MessageInfo, 26) +var file_grpc_pb_yao_proto_goTypes = []any{ (*RunRequest)(nil), // 0: yao.RunRequest (*RunResponse)(nil), // 1: yao.RunResponse (*Chunk)(nil), // 2: yao.Chunk @@ -1250,14 +1372,16 @@ var file_yao_proto_goTypes = []any{ (*AgentChunk)(nil), // 18: yao.AgentChunk (*Empty)(nil), // 19: yao.Empty (*HealthzResponse)(nil), // 20: yao.HealthzResponse - nil, // 21: yao.ShellRequest.EnvEntry - nil, // 22: yao.APIRequest.HeadersEntry - nil, // 23: yao.APIResponse.HeadersEntry + (*HeartbeatRequest)(nil), // 21: yao.HeartbeatRequest + (*HeartbeatResponse)(nil), // 22: yao.HeartbeatResponse + nil, // 23: yao.ShellRequest.EnvEntry + nil, // 24: yao.APIRequest.HeadersEntry + nil, // 25: yao.APIResponse.HeadersEntry } -var file_yao_proto_depIdxs = []int32{ - 21, // 0: yao.ShellRequest.env:type_name -> yao.ShellRequest.EnvEntry - 22, // 1: yao.APIRequest.headers:type_name -> yao.APIRequest.HeadersEntry - 23, // 2: yao.APIResponse.headers:type_name -> yao.APIResponse.HeadersEntry +var file_grpc_pb_yao_proto_depIdxs = []int32{ + 23, // 0: yao.ShellRequest.env:type_name -> yao.ShellRequest.EnvEntry + 24, // 1: yao.APIRequest.headers:type_name -> yao.APIRequest.HeadersEntry + 25, // 2: yao.APIResponse.headers:type_name -> yao.APIResponse.HeadersEntry 0, // 3: yao.Yao.Run:input_type -> yao.RunRequest 0, // 4: yao.Yao.Stream:input_type -> yao.RunRequest 3, // 5: yao.Yao.Shell:input_type -> yao.ShellRequest @@ -1271,46 +1395,48 @@ var file_yao_proto_depIdxs = []int32{ 14, // 13: yao.Yao.ChatCompletionsStream:input_type -> yao.ChatRequest 17, // 14: yao.Yao.AgentStream:input_type -> yao.AgentRequest 19, // 15: yao.Yao.Healthz:input_type -> yao.Empty - 1, // 16: yao.Yao.Run:output_type -> yao.RunResponse - 2, // 17: yao.Yao.Stream:output_type -> yao.Chunk - 4, // 18: yao.Yao.Shell:output_type -> yao.ShellResponse - 2, // 19: yao.Yao.ShellStream:output_type -> yao.Chunk - 6, // 20: yao.Yao.API:output_type -> yao.APIResponse - 8, // 21: yao.Yao.MCPListTools:output_type -> yao.MCPListResponse - 10, // 22: yao.Yao.MCPCallTool:output_type -> yao.MCPCallResponse - 11, // 23: yao.Yao.MCPListResources:output_type -> yao.MCPResourcesResponse - 13, // 24: yao.Yao.MCPReadResource:output_type -> yao.MCPResourceResponse - 15, // 25: yao.Yao.ChatCompletions:output_type -> yao.ChatResponse - 16, // 26: yao.Yao.ChatCompletionsStream:output_type -> yao.ChatChunk - 18, // 27: yao.Yao.AgentStream:output_type -> yao.AgentChunk - 20, // 28: yao.Yao.Healthz:output_type -> yao.HealthzResponse - 16, // [16:29] is the sub-list for method output_type - 3, // [3:16] is the sub-list for method input_type + 21, // 16: yao.Yao.Heartbeat:input_type -> yao.HeartbeatRequest + 1, // 17: yao.Yao.Run:output_type -> yao.RunResponse + 2, // 18: yao.Yao.Stream:output_type -> yao.Chunk + 4, // 19: yao.Yao.Shell:output_type -> yao.ShellResponse + 2, // 20: yao.Yao.ShellStream:output_type -> yao.Chunk + 6, // 21: yao.Yao.API:output_type -> yao.APIResponse + 8, // 22: yao.Yao.MCPListTools:output_type -> yao.MCPListResponse + 10, // 23: yao.Yao.MCPCallTool:output_type -> yao.MCPCallResponse + 11, // 24: yao.Yao.MCPListResources:output_type -> yao.MCPResourcesResponse + 13, // 25: yao.Yao.MCPReadResource:output_type -> yao.MCPResourceResponse + 15, // 26: yao.Yao.ChatCompletions:output_type -> yao.ChatResponse + 16, // 27: yao.Yao.ChatCompletionsStream:output_type -> yao.ChatChunk + 18, // 28: yao.Yao.AgentStream:output_type -> yao.AgentChunk + 20, // 29: yao.Yao.Healthz:output_type -> yao.HealthzResponse + 22, // 30: yao.Yao.Heartbeat:output_type -> yao.HeartbeatResponse + 17, // [17:31] is the sub-list for method output_type + 3, // [3:17] 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_yao_proto_init() } -func file_yao_proto_init() { - if File_yao_proto != nil { +func init() { file_grpc_pb_yao_proto_init() } +func file_grpc_pb_yao_proto_init() { + if File_grpc_pb_yao_proto != nil { return } type x struct{} out := protoimpl.TypeBuilder{ File: protoimpl.DescBuilder{ GoPackagePath: reflect.TypeOf(x{}).PkgPath(), - RawDescriptor: unsafe.Slice(unsafe.StringData(file_yao_proto_rawDesc), len(file_yao_proto_rawDesc)), + RawDescriptor: unsafe.Slice(unsafe.StringData(file_grpc_pb_yao_proto_rawDesc), len(file_grpc_pb_yao_proto_rawDesc)), NumEnums: 0, - NumMessages: 24, + NumMessages: 26, NumExtensions: 0, NumServices: 1, }, - GoTypes: file_yao_proto_goTypes, - DependencyIndexes: file_yao_proto_depIdxs, - MessageInfos: file_yao_proto_msgTypes, + GoTypes: file_grpc_pb_yao_proto_goTypes, + DependencyIndexes: file_grpc_pb_yao_proto_depIdxs, + MessageInfos: file_grpc_pb_yao_proto_msgTypes, }.Build() - File_yao_proto = out.File - file_yao_proto_goTypes = nil - file_yao_proto_depIdxs = nil + File_grpc_pb_yao_proto = out.File + file_grpc_pb_yao_proto_goTypes = nil + file_grpc_pb_yao_proto_depIdxs = nil } diff --git a/grpc/pb/yao.proto b/grpc/pb/yao.proto index f120151f..dfe8abfe 100644 --- a/grpc/pb/yao.proto +++ b/grpc/pb/yao.proto @@ -29,6 +29,9 @@ service Yao { // Health rpc Healthz(Empty) returns (HealthzResponse); + + // Sandbox + rpc Heartbeat(HeartbeatRequest) returns (HeartbeatResponse); } // ── Base ───────────────────────────────────────────────────────────────────── @@ -147,3 +150,16 @@ message Empty {} message HealthzResponse { string status = 1; } + +// ── Sandbox ───────────────────────────────────────────────────────────────── + +message HeartbeatRequest { + string sandbox_id = 1; + int32 cpu_percent = 2; + int64 mem_bytes = 3; + int32 running_procs = 4; +} + +message HeartbeatResponse { + string action = 1; // "ok" or "shutdown" +} diff --git a/grpc/pb/yao_grpc.pb.go b/grpc/pb/yao_grpc.pb.go index db5b07d6..44b48b81 100644 --- a/grpc/pb/yao_grpc.pb.go +++ b/grpc/pb/yao_grpc.pb.go @@ -2,7 +2,7 @@ // versions: // - protoc-gen-go-grpc v1.6.1 // - protoc v4.25.0 -// source: yao.proto +// source: grpc/pb/yao.proto package pb @@ -32,6 +32,7 @@ const ( Yao_ChatCompletionsStream_FullMethodName = "/yao.Yao/ChatCompletionsStream" Yao_AgentStream_FullMethodName = "/yao.Yao/AgentStream" Yao_Healthz_FullMethodName = "/yao.Yao/Healthz" + Yao_Heartbeat_FullMethodName = "/yao.Yao/Heartbeat" ) // YaoClient is the client API for Yao service. @@ -59,6 +60,8 @@ type YaoClient interface { AgentStream(ctx context.Context, in *AgentRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[AgentChunk], error) // Health Healthz(ctx context.Context, in *Empty, opts ...grpc.CallOption) (*HealthzResponse, error) + // Sandbox + Heartbeat(ctx context.Context, in *HeartbeatRequest, opts ...grpc.CallOption) (*HeartbeatResponse, error) } type yaoClient struct { @@ -235,6 +238,16 @@ func (c *yaoClient) Healthz(ctx context.Context, in *Empty, opts ...grpc.CallOpt return out, nil } +func (c *yaoClient) Heartbeat(ctx context.Context, in *HeartbeatRequest, opts ...grpc.CallOption) (*HeartbeatResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(HeartbeatResponse) + err := c.cc.Invoke(ctx, Yao_Heartbeat_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + // YaoServer is the server API for Yao service. // All implementations must embed UnimplementedYaoServer // for forward compatibility. @@ -260,6 +273,8 @@ type YaoServer interface { AgentStream(*AgentRequest, grpc.ServerStreamingServer[AgentChunk]) error // Health Healthz(context.Context, *Empty) (*HealthzResponse, error) + // Sandbox + Heartbeat(context.Context, *HeartbeatRequest) (*HeartbeatResponse, error) mustEmbedUnimplementedYaoServer() } @@ -309,6 +324,9 @@ func (UnimplementedYaoServer) AgentStream(*AgentRequest, grpc.ServerStreamingSer func (UnimplementedYaoServer) Healthz(context.Context, *Empty) (*HealthzResponse, error) { return nil, status.Error(codes.Unimplemented, "method Healthz not implemented") } +func (UnimplementedYaoServer) Heartbeat(context.Context, *HeartbeatRequest) (*HeartbeatResponse, error) { + return nil, status.Error(codes.Unimplemented, "method Heartbeat not implemented") +} func (UnimplementedYaoServer) mustEmbedUnimplementedYaoServer() {} func (UnimplementedYaoServer) testEmbeddedByValue() {} @@ -536,6 +554,24 @@ func _Yao_Healthz_Handler(srv interface{}, ctx context.Context, dec func(interfa return interceptor(ctx, in, info, handler) } +func _Yao_Heartbeat_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(HeartbeatRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(YaoServer).Heartbeat(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Yao_Heartbeat_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(YaoServer).Heartbeat(ctx, req.(*HeartbeatRequest)) + } + return interceptor(ctx, in, info, handler) +} + // Yao_ServiceDesc is the grpc.ServiceDesc for Yao service. // It's only intended for direct use with grpc.RegisterService, // and not to be introspected or modified (even as a copy) @@ -579,6 +615,10 @@ var Yao_ServiceDesc = grpc.ServiceDesc{ MethodName: "Healthz", Handler: _Yao_Healthz_Handler, }, + { + MethodName: "Heartbeat", + Handler: _Yao_Heartbeat_Handler, + }, }, Streams: []grpc.StreamDesc{ { @@ -602,5 +642,5 @@ var Yao_ServiceDesc = grpc.ServiceDesc{ ServerStreams: true, }, }, - Metadata: "yao.proto", + Metadata: "grpc/pb/yao.proto", } diff --git a/grpc/sandbox/heartbeat.go b/grpc/sandbox/heartbeat.go new file mode 100644 index 00000000..a9ece76c --- /dev/null +++ b/grpc/sandbox/heartbeat.go @@ -0,0 +1,76 @@ +package sandbox + +import ( + "context" + "sync" + "time" + + "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/grpc/pb" +) + +// HeartbeatData holds the latest heartbeat from a sandbox container. +type HeartbeatData struct { + SandboxID string + CPUPercent int32 + MemBytes int64 + RunningProcs int32 + LastSeen time.Time +} + +// Handler implements sandbox-related gRPC methods. +type Handler struct { + mu sync.RWMutex + heartbeats map[string]*HeartbeatData + onBeat func(data *HeartbeatData) string // optional callback; returns action +} + +// NewHandler creates a Handler. onBeat is called on each heartbeat and +// may return "ok" or "shutdown" to signal the container. +func NewHandler(onBeat func(data *HeartbeatData) string) *Handler { + return &Handler{ + heartbeats: make(map[string]*HeartbeatData), + onBeat: onBeat, + } +} + +// Heartbeat handles the Heartbeat RPC from sandbox containers. +func (h *Handler) Heartbeat(_ context.Context, req *pb.HeartbeatRequest) (*pb.HeartbeatResponse, error) { + data := &HeartbeatData{ + SandboxID: req.SandboxId, + CPUPercent: req.CpuPercent, + MemBytes: req.MemBytes, + RunningProcs: req.RunningProcs, + LastSeen: time.Now(), + } + + h.mu.Lock() + h.heartbeats[req.SandboxId] = data + h.mu.Unlock() + + action := "ok" + if h.onBeat != nil { + if a := h.onBeat(data); a != "" { + action = a + } + } + + log.Trace("sandbox heartbeat: id=%s cpu=%d%% mem=%d procs=%d → %s", + req.SandboxId, req.CpuPercent, req.MemBytes, req.RunningProcs, action) + + return &pb.HeartbeatResponse{Action: action}, nil +} + +// LastHeartbeat returns the most recent heartbeat for a sandbox, or nil. +func (h *Handler) LastHeartbeat(sandboxID string) *HeartbeatData { + h.mu.RLock() + defer h.mu.RUnlock() + return h.heartbeats[sandboxID] +} + +// RemoveHeartbeat cleans up heartbeat data for a removed sandbox. +func (h *Handler) RemoveHeartbeat(sandboxID string) { + h.mu.Lock() + delete(h.heartbeats, sandboxID) + h.mu.Unlock() +} diff --git a/grpc/sandbox/heartbeat_test.go b/grpc/sandbox/heartbeat_test.go new file mode 100644 index 00000000..028bc5be --- /dev/null +++ b/grpc/sandbox/heartbeat_test.go @@ -0,0 +1,158 @@ +package sandbox + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/yaoapp/yao/grpc/pb" +) + +func TestHeartbeat_StoresData(t *testing.T) { + h := NewHandler(nil) + + req := &pb.HeartbeatRequest{ + SandboxId: "sb-1", + CpuPercent: 25, + MemBytes: 1024 * 1024, + RunningProcs: 3, + } + resp, err := h.Heartbeat(context.Background(), req) + if err != nil { + t.Fatalf("Heartbeat: %v", err) + } + if resp.Action != "ok" { + t.Errorf("action = %q, want %q", resp.Action, "ok") + } + + data := h.LastHeartbeat("sb-1") + if data == nil { + t.Fatal("LastHeartbeat returned nil") + } + if data.SandboxID != "sb-1" { + t.Errorf("SandboxID = %q", data.SandboxID) + } + if data.CPUPercent != 25 { + t.Errorf("CPUPercent = %d", data.CPUPercent) + } + if data.MemBytes != 1024*1024 { + t.Errorf("MemBytes = %d", data.MemBytes) + } + if data.RunningProcs != 3 { + t.Errorf("RunningProcs = %d", data.RunningProcs) + } + if time.Since(data.LastSeen) > time.Second { + t.Errorf("LastSeen too old: %v", data.LastSeen) + } +} + +func TestHeartbeat_OnBeatCallback(t *testing.T) { + var received *HeartbeatData + h := NewHandler(func(d *HeartbeatData) string { + received = d + return "shutdown" + }) + + resp, err := h.Heartbeat(context.Background(), &pb.HeartbeatRequest{ + SandboxId: "sb-2", + CpuPercent: 90, + MemBytes: 4096, + RunningProcs: 10, + }) + if err != nil { + t.Fatalf("Heartbeat: %v", err) + } + if resp.Action != "shutdown" { + t.Errorf("action = %q, want %q", resp.Action, "shutdown") + } + if received == nil || received.SandboxID != "sb-2" { + t.Errorf("callback not invoked or wrong data") + } +} + +func TestHeartbeat_OnBeatEmptyReturnDefaultsToOK(t *testing.T) { + h := NewHandler(func(d *HeartbeatData) string { + return "" + }) + + resp, err := h.Heartbeat(context.Background(), &pb.HeartbeatRequest{SandboxId: "sb-3"}) + if err != nil { + t.Fatalf("Heartbeat: %v", err) + } + if resp.Action != "ok" { + t.Errorf("action = %q, want %q", resp.Action, "ok") + } +} + +func TestLastHeartbeat_Unknown(t *testing.T) { + h := NewHandler(nil) + if d := h.LastHeartbeat("nonexistent"); d != nil { + t.Errorf("expected nil for unknown sandbox, got %+v", d) + } +} + +func TestRemoveHeartbeat(t *testing.T) { + h := NewHandler(nil) + + h.Heartbeat(context.Background(), &pb.HeartbeatRequest{SandboxId: "sb-rm"}) + if h.LastHeartbeat("sb-rm") == nil { + t.Fatal("expected data after heartbeat") + } + + h.RemoveHeartbeat("sb-rm") + if h.LastHeartbeat("sb-rm") != nil { + t.Error("expected nil after RemoveHeartbeat") + } +} + +func TestRemoveHeartbeat_Idempotent(t *testing.T) { + h := NewHandler(nil) + h.RemoveHeartbeat("never-existed") +} + +func TestHeartbeat_ConcurrentAccess(t *testing.T) { + h := NewHandler(nil) + var wg sync.WaitGroup + + for i := 0; i < 50; i++ { + wg.Add(1) + go func(n int) { + defer wg.Done() + id := "sb-concurrent" + h.Heartbeat(context.Background(), &pb.HeartbeatRequest{ + SandboxId: id, + CpuPercent: int32(n), + RunningProcs: int32(n), + }) + h.LastHeartbeat(id) + }(i) + } + wg.Wait() + + if d := h.LastHeartbeat("sb-concurrent"); d == nil { + t.Error("expected data after concurrent heartbeats") + } +} + +func TestHeartbeat_MultiSandbox(t *testing.T) { + h := NewHandler(nil) + + for _, id := range []string{"a", "b", "c"} { + h.Heartbeat(context.Background(), &pb.HeartbeatRequest{SandboxId: id, CpuPercent: 10}) + } + + for _, id := range []string{"a", "b", "c"} { + if d := h.LastHeartbeat(id); d == nil { + t.Errorf("missing heartbeat for %q", id) + } + } + + h.RemoveHeartbeat("b") + if h.LastHeartbeat("b") != nil { + t.Error("b should be removed") + } + if h.LastHeartbeat("a") == nil || h.LastHeartbeat("c") == nil { + t.Error("a and c should still exist") + } +} diff --git a/sandbox/DESIGN.md b/sandbox/DESIGN.md index 19fc29fd..4c1ee68c 100644 --- a/sandbox/DESIGN.md +++ b/sandbox/DESIGN.md @@ -45,23 +45,22 @@ High-level business layer on top of `tai.Client`. Manages container lifecycle, u - File operations (via `tai.Client.Volume()` for remote, bind mount for local) - IPC relay to Yao gRPC server -### Yao gRPC Server (yao/grpc) +### Yao gRPC Server (yao/grpc) — ✅ Implemented -General-purpose gRPC service exposed by the Yao process. Not limited to sandbox IPC — it exposes Yao's process execution capability to any gRPC client. +General-purpose gRPC gateway exposed by the Yao process. Not limited to sandbox IPC — it exposes process execution, shell, API proxy, MCP, LLM, and Agent capabilities to any gRPC client. 14 RPCs defined; V1 (unary + LLM/Agent streaming) complete, V2 (base streaming via `gou/stream`) pending. **Clients:** -- Container-internal MCP tools (via Tai Gateway relay) -- `yao run --remote` CLI +- Container-internal `yao-grpc` (via Tai Gateway relay or direct) +- `yao run` CLI (after `yao login`) - Other Yao instances (future node-to-node) **IPC path (replacing Unix socket):** ``` -Container process → yao-bridge (tai/bridge/) → Tai Gateway (:9100 gRPC) → Yao gRPC Server (:9099) - │ - process.Run(...) +Local: Container → yao-grpc (tai/grpc/) → Yao gRPC 127.0.0.1:9099 +Remote: Container → yao-grpc (tai/grpc/) → Tai Gateway (:9100 gRPC) → Yao gRPC Server (:9099) ``` -Tai does **not** know Yao gRPC address at startup. The upstream is passed per-container via `CreateRequest.GRPCUpstream` — Tai records the mapping and routes relay traffic by source container. This keeps Tai stateless and allows one Tai to serve multiple Yao instances. +All modes use gRPC — no Unix socket fallback. `yao-grpc` reads `YAO_GRPC_ADDR` from env and connects. Local containers point directly at the Yao gRPC server on loopback; remote containers point at the Tai relay. Tai does **not** know Yao gRPC address at startup — `yao-grpc` carries target in `x-grpc-upstream` request metadata. This keeps Tai stateless and allows one Tai to serve multiple Yao instances. ## Authentication @@ -79,11 +78,11 @@ The gRPC server reuses the existing `openapi/oauth` service — no new auth syst | Scope registration | `acl.Register(...)` | gRPC scopes via same pattern | None — add `grpc:*` scope definitions in `init()` | | Client auth | `ClientProvider` | `client_credentials` grant for CLI/containers | None | | Token revocation | `oauth.Revoke(ctx, token, hint)` | Container token cleanup | None | -| Device Flow scaffolding | `types.DeviceAuthorizationResponse`, `GrantTypeDeviceCode`, error codes | CLI `yao login` | Implement `DeviceAuthorization()` (currently stub) | +| Device Flow | `DeviceAuthorization()`, `AuthorizeDevice()`, device_code store, `GrantTypeDeviceCode` grant | CLI `yao login` | ✅ Implemented | **Key insight**: `authorized.SetInfo/GetInfo` are Gin-bound, but gRPC does NOT need them. The gRPC interceptor builds `AccessRequest` directly from JWT claims and calls `ScopeManager.Check` — bypasses the full `Enforce` chain (client/team/member), which is HTTP multi-tenant only. -**Impact on existing code: zero.** All gRPC auth is purely additive (~80 lines interceptor + scope registration). Device Flow adds ~190 lines new code + ~10 lines to existing `Token()` switch. +**Status**: All auth infrastructure is implemented and working — gRPC interceptor, scope registration, Device Flow (backend + CUI page), CLI commands (`yao login`/`yao logout`). ### gRPC interceptor @@ -106,8 +105,8 @@ func authInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, h | Client | How it gets a token | |--------|-------------------| -| Container MCP tool | `oauth.MakeAccessToken(clientID, "grpc:mcp grpc:run", userID, 900)` injected as `YAO_TOKEN` env var. `yao-bridge` auto-refreshes via `YAO_REFRESH_TOKEN`. Manager revokes refresh token on container Remove. | -| `yao run` CLI | `yao login` → OAuth Device Authorization Grant → token saved to `~/.yao/credentials`. Logged in = gRPC, not logged in = local. | +| Container MCP tool | `oauth.MakeAccessToken(clientID, "grpc:mcp grpc:run", userID, 900)` injected as `YAO_TOKEN` env var. `yao-grpc` auto-refreshes via response metadata. Manager revokes refresh token on container Remove. | +| `yao run` CLI | ✅ `yao login --server ` → OAuth Device Authorization Grant (RFC 8628) → dynamic client registration via machine ID → token saved to `~/.yao/credentials` (base64 JSON). Logged in = gRPC, not logged in = local. | | Yao-to-Yao | Pre-shared service token or `client_credentials` | ## Network Security @@ -193,7 +192,7 @@ Lifecycle is managed by `sandbox.Manager`, not by tai.Client. | Mode | tai.Client | File IO | |------|-----------|---------| -| Local | `tai.New("")` | Bind mount, direct host filesystem | +| Local | `tai.New("local")` | Bind mount, direct host filesystem | | Remote | `tai.New("tai://host")` | `tai.Client.Volume()` via gRPC | Local mode preserves bind mount for performance. Remote mode uses `tai/volume` (gRPC + lz4 compression). `sandbox.Manager` routes based on `client.IsLocal()`. @@ -234,32 +233,37 @@ agent/context/jsapi_sandbox.go | Container exec | `dockerClient.ContainerExecCreate/Start/Attach` | `tai.Client.Sandbox().Exec()` | | File read | Host path via bind mount (`containerPathToHost`) | Local: bind mount (same). Remote: `tai.Client.Volume().Read()` | | File write | `dockerClient.CopyToContainer` | Local: bind mount. Remote: `tai.Client.Volume().Write()` | -| IPC | Unix socket bind mount + yao-bridge | Local: Unix socket (same). Remote: Tai gRPC relay → Yao gRPC server | -| MCP config | `{args: ["/tmp/yao.sock"]}` hardcoded | Local: socket path from config. Remote: gRPC endpoint injected as env var | +| IPC | Unix socket bind mount + yao-bridge | All modes: `yao-grpc` → gRPC (direct or via Tai relay). No Unix socket. | +| MCP config | `{args: ["/tmp/yao.sock"]}` hardcoded | `YAO_GRPC_ADDR` + `YAO_TOKEN` env vars. Local: direct. Remote: + `YAO_GRPC_TAI`/`YAO_GRPC_UPSTREAM`. | | VNC | `vncproxy.NewProxy(nil)` local assumption | `tai.Client.VNC().URL()` | | Cleanup | `dockerClient.ContainerRemove` | `tai.Client.Sandbox().Remove()` | ### IPC migration detail -**Local mode** (same host): Unix socket preserved — zero overhead, no change needed. - -**Remote mode** (via Tai): -``` -Container process → yao-bridge (tai/bridge/) → Tai relay (:9100 gRPC) → Yao gRPC Server -``` - -`yao-bridge` source lives in `yao/tai/bridge/` — it's a Tai SDK client (consumes Tai relay), shares gRPC deps with `tai/`, and is version-locked with the Tai protocol. Built via `go build ./tai/bridge/cmd/yao-bridge`. - -Bridge mode determined by env var: +All modes use gRPC — no Unix socket fallback, one code path for local and remote. ``` -YAO_IPC_MODE=socket YAO_IPC_ADDR=/tmp/yao.sock # local -YAO_IPC_MODE=grpc YAO_IPC_ADDR=tai-host:9100 # remote +Local: Container → yao-grpc → Yao gRPC 127.0.0.1:9099 +Remote: Container → yao-grpc → Tai :9100 relay → Yao gRPC :9099 ``` -In gRPC mode, bridge also reads `YAO_TOKEN` / `YAO_REFRESH_TOKEN` and handles automatic token refresh (see grpc/DESIGN.md Container token section). +`yao-grpc` source lives in `yao/tai/grpc/` — shares gRPC deps with `tai/`, version-locked with the Tai protocol. Built via `go build -o yao-grpc ./tai/grpc/cmd`. -Tai relay upstream is NOT configured at Tai startup. Manager passes `GRPCUpstream` per-container in `CreateRequest` — Tai records the mapping and routes by source container. One Tai can serve containers from different Yao instances. +Mode determined by env vars injected by Manager at container creation: + +``` +# Local: direct to Yao +YAO_GRPC_ADDR=127.0.0.1:9099 + +# Remote: via Tai relay +YAO_GRPC_ADDR=tai-host:9100 +YAO_GRPC_TAI=enable +YAO_GRPC_UPSTREAM=yao-host:9099 +``` + +`yao-grpc` reads `YAO_TOKEN` / `YAO_REFRESH_TOKEN` / `YAO_SANDBOX_ID` from env, attaches as gRPC metadata on every call, and handles automatic token refresh from response metadata. In Tai relay mode, attaches `x-grpc-upstream` metadata so Tai knows where to forward. + +Tai relay upstream is NOT configured at Tai startup. `yao-grpc` carries the target address per request — Tai reads `x-grpc-upstream` metadata and proxies dynamically. One Tai can serve containers from different Yao instances. `BuildMCPConfigForSandbox()` sets the env vars based on `client.IsLocal()`. @@ -298,8 +302,26 @@ sandbox: ## Migration Path -1. **Phase 1:** Yao gRPC server — expose process execution, replace Unix socket IPC -2. **Phase 2:** `sandbox.Manager` refactoring — replace Docker client with `tai.Client`, unified file ops, new lifecycle model -3. **Phase 3:** Agent layer adaptation — executor uses new Manager, IPC mode switch, lifecycle policy -4. **Phase 4:** `yao run --remote` — CLI calls remote Yao via gRPC -5. **Phase 5:** Workspace persistence — browser preview, service exposure, delivery +### Completed + +1. **Yao gRPC server** (`yao/grpc`) — full gRPC gateway with 14 RPCs (Run, Stream, Shell, ShellStream, API, MCP×4, ChatCompletions, ChatCompletionsStream, AgentStream, Healthz). OAuth + ACL auth interceptor reusing existing openapi infrastructure. V1 all unary + LLM/Agent streaming done; V2 base streaming (Stream, ShellStream) pending `gou/stream` package. Details: [grpc/DESIGN.md](../grpc/DESIGN.md), [grpc/IMPL.md](../grpc/IMPL.md). + +2. **Tai SDK** (`yao/tai`) — unified sandbox runtime SDK with Local/Remote modes. Sandbox (container lifecycle), Volume (file IO + sync with lz4), Workspace (`fs.FS` compatible), Proxy (HTTP reverse proxy), VNC (WebSocket). Remote mode connects via Tai gateway (gRPC :9100, Docker :2375, K8s :6443, HTTP :8080, VNC :6080). Details: [tai/docs/README.md](../tai/docs/README.md). + +3. **Tai gateway dynamic routing** (Tai repo) — removed fixed `YaoUpstream` startup config. `yao-grpc` carries `x-grpc-upstream` metadata per request; Tai reads target and proxies dynamically. One Tai serves containers from multiple Yao instances. + +4. **yao-grpc container client** (`yao/tai/grpc`) — in-container gRPC client binary replacing `yao-bridge`. Reads `YAO_TOKEN`/`YAO_REFRESH_TOKEN`/`YAO_SANDBOX_ID` from env, auto-refreshes tokens via response metadata. Supports direct mode (`YAO_GRPC_ADDR=127.0.0.1:9099`) and Tai relay mode (`YAO_GRPC_TAI=enable`). Built as `go build -o yao-grpc ./tai/grpc/cmd`. + +5. **OAuth Device Flow + CLI auth** — `yao login --server ` (RFC 8628 Device Authorization Grant), `yao logout`, credentials stored as base64 JSON in `~/.yao/credentials`. CUI `/auth/device` page for user authorization. Dynamic client registration via machine ID. + +6. **`yao run` via gRPC** (`yao/cmd/run.go`) — no `--remote` flag; logged in = gRPC, not logged in = local. `--auth ` for alternate credentials. TUI status bar (lipgloss) shows user/scope in gRPC mode, hidden with `-s` (silent). + +### Remaining + +7. **`sandbox.Manager` refactoring** — replace Docker client with `tai.Client`, unified file ops (bind mount for local, `tai.Client.Volume()` for remote), new lifecycle model (one-shot / session / long-running / persistent). + +8. **Agent layer adaptation** — executor uses new Manager, IPC mode switch (gRPC replaces Unix socket), lifecycle policy per-assistant config. + +9. **`gou/stream` package** (V2) — streaming process execution foundation. ~150 lines. Enables gRPC `Stream` and `ShellStream` handlers, V8 `Stream()` global. + +10. **Workspace persistence** — browser preview, service exposure, delivery. diff --git a/sandbox/docker/build.sh b/sandbox/docker/build.sh index 2e3ea74b..c5a3077f 100755 --- a/sandbox/docker/build.sh +++ b/sandbox/docker/build.sh @@ -164,15 +164,21 @@ case $TOOL in # Cursor (uncomment when ready) # build_multiarch "sandbox-cursor" "cursor/Dockerfile" "$PUSH" ;; + v2) + echo "V2 images have their own build script: sandbox/v2/docker/build.sh" + echo "Usage: sandbox/v2/docker/build.sh [true|false]" + exit 0 + ;; *) echo "Unknown tool: $TOOL" - echo "Usage: $0 [claude|claude-vnc|browser|desktop|chrome|cursor|all] [true|false]" + echo "Usage: $0 [claude|claude-vnc|browser|desktop|chrome|cursor|v2|all] [true|false]" echo " $0 claude # Build Claude images locally" echo " $0 claude true # Build and push Claude images" echo " $0 claude-vnc # Build Claude VNC images (Browser + Desktop)" echo " $0 browser # Build Claude Browser image only" echo " $0 desktop # Build Claude Desktop image only" echo " $0 chrome # Build Claude Chrome image (amd64 only)" + echo " $0 v2 # Build Sandbox V2 images (base + test)" echo " $0 all true # Build and push all images" exit 1 ;; diff --git a/sandbox/v2/DESIGN.md b/sandbox/v2/DESIGN.md new file mode 100644 index 00000000..7123864b --- /dev/null +++ b/sandbox/v2/DESIGN.md @@ -0,0 +1,1127 @@ +# Sandbox V2 Design + +## Positioning + +Sandbox is a **standalone infrastructure module** in Yao, on the same level as `process`, `store`, and `fs`. It provides isolated execution environments with standard file I/O. Any module can use it — Agent, JSAPI scripts, Process handlers, API endpoints. + +``` +Yao Infrastructure +├── process — process execution +├── store — KV storage +├── fs — host filesystem +├── stream — streaming execution (planned) +└── sandbox — isolated execution environments ← this module +``` + +Sandbox does NOT import or depend on Agent. Agent is one of many consumers. + +## Architecture + +``` +┌─────────────────────────────────────────────────┐ +│ Consumers (know nothing about tai/Docker/K8s) │ +│ ├── JSAPI: Sandbox("my-app") │ +│ ├── Process: sandbox.Create, sandbox.Exec │ +│ ├── Agent: uses sandbox via interface │ +│ └── API: /api/__yao/sandbox/* │ +└──────────────────┬──────────────────────────────┘ + │ sandbox.Manager (public API) + ▼ +┌─────────────────────────────────────────────────┐ +│ sandbox/v2 │ +│ │ +│ Manager (global singleton) │ +│ ├── Create / Get / Start / Stop / Remove │ +│ ├── List / Cleanup / Close │ +│ └── guard rails (limits, TTL) + Box factory │ +│ │ +│ Box (per-instance) │ +│ ├── Exec(cmd) → ExecResult │ +│ ├── Stream(cmd) → ExecStream (real-time I/O) │ +│ ├── Attach(port) → ServiceConn (WS/SSE/TCP) │ +│ ├── Workspace() → workspace.FS │ +│ ├── VNC() → url │ +│ ├── Proxy(port) → url │ +│ └── Start / Stop / Remove / Info │ +└──────────────────┬──────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────┐ +│ tai.Client pool (lazy-initialized) │ +│ ├── "local" → tai.New("local") (Docker) │ +│ ├── "gpu" → tai.New("tai://gpu") (Remote) │ +│ ├── "k8s" → tai.New("tai://k8s") (K8s) │ +│ └── ... │ +│ │ +│ Each tai.Client provides: │ +│ ├── Sandbox() → CRUD + Exec + ExecStream │ +│ ├── Volume() → file I/O (local disk / gRPC) │ +│ ├── Workspace() → fs.FS │ +│ ├── Proxy() → URL resolve + Connect │ +│ └── VNC() → VNC WebSocket │ +└─────────────────────────────────────────────────┘ +``` + +## Dependency Rules + +``` +sandbox/v2 → tai ✓ (sole runtime dependency) +sandbox/v2 → agent ✗ NEVER +sandbox/v2 → docker ✗ NEVER (tai handles it) +agent → sandbox/v2 ✓ (consumer, via Manager API) +jsapi → sandbox/v2 ✓ (consumer, via Manager API) +process → sandbox/v2 ✓ (consumer, via Manager API) +``` + +## Manager + +Global singleton. Manages a **pool of named `tai.Client` connections** — each pool entry targets a different runtime endpoint (local Docker, remote Tai, K8s cluster). Caller picks which pool to create a sandbox on. + +### Pool + +```go +// Pool defines a named tai.Client endpoint with its own policy. +type Pool struct { + Name string // unique name, e.g. "local", "gpu", "k8s-prod" + Addr string // tai.New() address: "local", "tai://host", "docker:///path" + Options []tai.Option // tai.WithPorts(), tai.WithKubeConfig(), etc. + + // Guard rails (per-pool) + MaxPerUser int // max boxes per user on this pool, 0 = unlimited + MaxTotal int // max boxes total on this pool, 0 = unlimited + + // Default lifecycle (overridable per-box via CreateOptions) + IdleTimeout time.Duration // 0 = no timeout + MaxLifetime time.Duration // 0 = no limit +} +``` + +Example configuration: + +``` +pool: + - name: local + addr: "local" + max_total: 20 + idle_timeout: 30m + + - name: gpu + addr: "tai://gpu-server.internal" + max_per_user: 1 # GPU is expensive, 1 per user + max_total: 4 + idle_timeout: 10m # reclaim fast + max_lifetime: 2h + + - name: k8s + addr: "tai://k8s-proxy.internal" + max_total: 100 + idle_timeout: 1h + options: + runtime: k8s + kubeconfig: /etc/yao/kubeconfig.yml +``` + +### Initialization + +```go +package sandbox + +var mgr *Manager + +// Init initializes the global Manager. +// Config contains everything: pool definitions + guard rails. +// At least one Pool entry is required. The first entry is the default. +// Pass empty Pool list to disable sandbox (methods return ErrNotAvailable). +func Init(cfg Config) error + +// M returns the global Manager. Panics if Init was not called. +func M() *Manager +``` + +Startup sequence in `cmd/start.go`: + +``` +config.Load +sandbox.Init(config.Conf.Sandbox) // create Manager with pool + guard rails +engine.Load +... +service.Start // HTTP +grpc.Start // gRPC +sandbox.M().Start(ctx) // discover existing containers, start cleanup loop +``` + +`Init` creates the Manager from config (pool definitions + guard rails). `Start` connects to pools, discovers existing containers, and starts the cleanup loop. Two-step so that gRPC server is ready before Start (containers may send heartbeats immediately). + +Pool connections are created lazily on first use and reused across all Box instances. + +### Config + +Pool definitions only. Guard rails and lifecycle defaults are per-pool. Per-instance settings (image, memory, workdir, etc.) are in `CreateOptions`. + +```go +type Config struct { + Pool []Pool // runtime endpoints; first is default +} +``` + +Container gRPC env vars (`YAO_GRPC_ADDR`, `YAO_GRPC_UPSTREAM`, etc.) are derived automatically at creation time — local address from Yao's gRPC config (`config.Conf.GRPC`), remote relay from pool's tai address. No manual configuration needed. + +Per-instance settings (image, memory, CPU, workdir, env, pool) are passed via `CreateOptions` by the caller — assistant config, JSAPI parameters, or Process arguments. The Manager doesn't impose defaults for container specs; that's the caller's responsibility. + +### Core API + +```go +type Manager struct { + pool map[string]*tai.Client // name → connection (lazy-initialized) + poolDefs []Pool // pool definitions + defaultPool string // first pool name + config Config + boxes sync.Map // id → *Box + mu sync.Mutex // creation serialization +} + +// --- Bootstrap --- + +// Start discovers existing containers from all pools, rebuilds the boxes map, +// and starts the cleanup loop. Called once after Init. +func (m *Manager) Start(ctx context.Context) error + +// --- Pool management --- + +// AddPool registers a new pool at runtime. Connects lazily on first use. +func (m *Manager) AddPool(ctx context.Context, p Pool) error + +// RemovePool removes a pool by name. Fails if any running boxes are on it. +// Use force=true to stop all boxes on the pool first, then remove. +func (m *Manager) RemovePool(ctx context.Context, name string, force bool) error + +// Pools returns all registered pool names and their status (connected/disconnected). +func (m *Manager) Pools() []PoolInfo + +// --- Heartbeat (called by gRPC handler, not by consumers) --- + +// Heartbeat updates the box's last heartbeat timestamp. +// Called by the gRPC Heartbeat handler when a container reports in. +// Returns ErrNotFound if sandbox_id is unknown (container orphaned or already removed). +func (m *Manager) Heartbeat(sandboxID string, active bool, processCount int) error + +// --- CRUD --- + +// Create creates and starts a new sandbox. +func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Box, error) + +// Get returns an existing sandbox by ID. Returns ErrNotFound if not exists. +func (m *Manager) Get(ctx context.Context, id string) (*Box, error) + +// GetOrCreate returns existing sandbox or creates a new one. +func (m *Manager) GetOrCreate(ctx context.Context, opts CreateOptions) (*Box, error) + +// List returns all sandboxes, optionally filtered. +func (m *Manager) List(ctx context.Context, opts ListOptions) ([]*Box, error) + +// Remove stops and removes a sandbox. +func (m *Manager) Remove(ctx context.Context, id string) error + +// Cleanup removes idle/expired sandboxes. Called periodically. +func (m *Manager) Cleanup(ctx context.Context) error + +// Close stops the cleanup loop and releases all pool connections. +func (m *Manager) Close() error +``` + +### CreateOptions + +All per-instance settings live here. Caller decides everything about the container. + +```go +type CreateOptions struct { + // Identity + ID string // explicit ID; empty = auto-generate + Owner string // user ID for isolation and limits + Labels map[string]string + + // Runtime target + Pool string // which tai.Client to use; empty = default pool + + // Container spec + Image string // required + WorkDir string // container working directory, default "/workspace" + User string // container user + Env map[string]string // additional env vars + Memory int64 // bytes, 0 = no limit + CPUs float64 // 0 = no limit + VNC bool // enable VNC + Ports []PortMapping // extra port mappings + + // Lifecycle + Policy LifecyclePolicy // default: Session + IdleTimeout time.Duration // override Manager default; 0 = use Manager default +} + +type LifecyclePolicy string + +const ( + OneShot LifecyclePolicy = "oneshot" // destroyed after first Exec + Session LifecyclePolicy = "session" // alive while active, cleaned on idle + LongRunning LifecyclePolicy = "longrunning" // user workspace, extended TTL + Persistent LifecyclePolicy = "persistent" // never auto-cleaned +) +``` + +## Box + +A `Box` is a single sandbox instance. All operations go through it. + +```go +type Box struct { + id string + containerID string + pool string // which tai.Client this box runs on + owner string + policy LifecyclePolicy + labels map[string]string + lastCall atomic.Int64 // last external API call (Exec/Workspace/VNC/Proxy) + lastHeartbeat atomic.Int64 // last container heartbeat + processCount atomic.Int32 // user processes inside container (from heartbeat) + ws workspace.FS // lazy-initialized, cached + manager *Manager +} + +// lastActiveTime returns max(lastCall, lastHeartbeat). +func (b *Box) lastActiveTime() time.Time + +// ID returns the sandbox identifier. +func (b *Box) ID() string + +// Owner returns the user who owns this sandbox. +func (b *Box) Owner() string + +// ContainerID returns the underlying container ID. +func (b *Box) ContainerID() string + +// --- Execution --- + +// Exec runs a command and waits for it to finish. +func (b *Box) Exec(ctx context.Context, cmd []string, opts ...ExecOption) (*ExecResult, error) + +// Stream runs a command with real-time streaming I/O. +func (b *Box) Stream(ctx context.Context, cmd []string, opts ...ExecOption) (*ExecStream, error) + +// Attach connects to a service running inside the sandbox on the given container port. +func (b *Box) Attach(ctx context.Context, port int, opts ...AttachOption) (*ServiceConn, error) + +// --- Filesystem (fs.FS compatible) --- + +// Workspace returns an fs.FS-compatible filesystem for this sandbox. +// Supports: Open, Stat, ReadFile, ReadDir, WriteFile, Remove, Rename, MkdirAll. +// Internally calls tai.Client.Workspace(box.id) — uses sandbox ID as volume session. +func (b *Box) Workspace() workspace.FS + +// --- Network --- + +// VNC returns the VNC WebSocket URL. Error if VNC not enabled. +func (b *Box) VNC(ctx context.Context) (string, error) + +// Proxy returns the HTTP URL for a service running on the given port inside the sandbox. +func (b *Box) Proxy(ctx context.Context, port int, path string) (string, error) + +// --- Lifecycle --- + +// Start starts a stopped sandbox. +func (b *Box) Start(ctx context.Context) error + +// Stop stops the sandbox without removing it. +func (b *Box) Stop(ctx context.Context) error + +// Remove stops and removes the sandbox. +func (b *Box) Remove(ctx context.Context) error + +// Info returns current sandbox status. +func (b *Box) Info(ctx context.Context) (*BoxInfo, error) +``` + +### ExecOption / ExecResult + +```go +type ExecOption func(*execConfig) + +func WithWorkDir(dir string) ExecOption +func WithEnv(env map[string]string) ExecOption +func WithTimeout(d time.Duration) ExecOption + +type ExecResult struct { + ExitCode int + Stdout string + Stderr string +} +``` + +### ExecStream + +```go +type ExecStream struct { + Stdout io.ReadCloser // real-time stdout + Stderr io.ReadCloser // real-time stderr + Stdin io.WriteCloser // write to process stdin (nil if not interactive) + Wait func() (int, error) // block until exit, return exit code + Cancel func() // kill the process +} +``` + +Usage: + +```go +// Interactive CLI (e.g. Claude) +s, _ := box.Stream(ctx, []string{"claude", "--chat"}) +go io.Copy(os.Stdout, s.Stdout) +s.Stdin.Write([]byte("help\n")) +code, _ := s.Wait() + +// Long-running process (e.g. dev server) +s, _ := box.Stream(ctx, []string{"npm", "run", "dev"}) +go io.Copy(logWriter, s.Stdout) // continuous output +// ... later +s.Cancel() +``` + +### AttachOption / ServiceConn + +```go +type AttachOption func(*attachConfig) + +func WithProtocol(proto string) AttachOption // "ws", "sse", "tcp"; default "ws" +func WithPath(path string) AttachOption // URL path, e.g. "/v1/chat" +func WithHeaders(h map[string]string) AttachOption + +type ServiceConn struct { + // Bidirectional (WebSocket, TCP) + Read func() ([]byte, error) + Write func(data []byte) error + + // Server-push (SSE) + Events <-chan []byte // nil if not SSE mode + + // Common + URL string // resolved URL for reference + Close func() error +} +``` + +`port` is the port the service listens on **inside the container** (e.g. 3000 for a Node server). Routing to that port — Docker port mapping (local) or Tai HTTP proxy (remote) — is handled internally. + +**Local mode caveat**: `tai/proxy.NewLocal` resolves host ports via `Inspect()` → `PortMapping`. The container must have the port mapped at creation time (`CreateOptions.Ports`). If the port was not mapped, `Proxy()` and `Attach()` return an error. Remote mode has no such restriction — Tai HTTP proxy routes by container IP directly. + +Usage: + +```go +// WebSocket — connect to Cursor Server inside sandbox +conn, _ := box.Attach(ctx, 3000, WithProtocol("ws"), WithPath("/ws")) +conn.Write([]byte(`{"type":"edit","file":"main.go"}`)) +msg, _ := conn.Read() +conn.Close() + +// SSE — connect to Claude API inside sandbox +conn, _ := box.Attach(ctx, 8080, WithProtocol("sse"), WithPath("/v1/messages")) +for event := range conn.Events { + fmt.Println(string(event)) +} +``` + +### PoolInfo + +```go +type PoolInfo struct { + Name string // pool name + Addr string // tai address + Connected bool // tai.Client connection established + Boxes int // number of boxes on this pool + MaxPerUser int + MaxTotal int + IdleTimeout time.Duration + MaxLifetime time.Duration +} +``` + +### BoxInfo + +```go +type BoxInfo struct { + ID string + ContainerID string + Pool string + Owner string + Status string // "running", "stopped", "creating" + Policy LifecyclePolicy + Labels map[string]string + Image string + CreatedAt time.Time + LastActive time.Time // max(lastCall, lastHeartbeat) + ProcessCount int // user processes inside container (0 = idle) + VNC bool +} +``` + +## Workspace — fs.FS Interface + +`Box.Workspace()` returns `workspace.FS` from `tai/workspace`. This is the standard Go `fs.FS` interface extended with write operations. + +```go +// tai/workspace.FS — already implemented +type FS interface { + fs.FS // Open(name) (fs.File, error) + fs.StatFS // Stat(name) (fs.FileInfo, error) + fs.ReadFileFS // ReadFile(name) ([]byte, error) + fs.ReadDirFS // ReadDir(name) ([]fs.DirEntry, error) + io.Closer + + WriteFile(name string, data []byte, perm os.FileMode) error + Remove(name string) error + RemoveAll(name string) error + Rename(oldname, newname string) error + MkdirAll(name string, perm os.FileMode) error +} +``` + +100% compatible with Go standard library: + +```go +ws := box.Workspace() + +// Standard fs functions work +data, _ := fs.ReadFile(ws, "main.go") +fs.WalkDir(ws, ".", func(path string, d fs.DirEntry, err error) error { ... }) +info, _ := fs.Stat(ws, "go.mod") + +// Extended write operations +ws.WriteFile("main.go", []byte("package main"), 0644) +ws.MkdirAll("src/pkg", 0755) +ws.Remove("tmp.txt") +ws.Rename("old.go", "new.go") +``` + +Local mode: reads/writes go directly to host disk via bind mount. +Remote mode: reads/writes go through tai Volume gRPC with lz4 compression. +Caller doesn't know or care which mode. + +## Container gRPC (already implemented) + +Container processes communicate with Yao via gRPC. No Unix sockets. + +``` +Local: Container → yao-grpc → Yao gRPC 127.0.0.1:9099 +Remote: Container → yao-grpc → Tai :9100 relay → Yao gRPC :9099 +``` + +Manager injects env vars at container creation: + +``` +# Local +YAO_GRPC_ADDR=127.0.0.1:9099 +YAO_TOKEN= +YAO_REFRESH_TOKEN= +YAO_SANDBOX_ID= + +# Remote (adds Tai relay) +YAO_GRPC_TAI=enable +YAO_GRPC_UPSTREAM=yao-host:9099 +``` + +Token issuance uses existing `openapi/oauth`. Manager creates token pair at container creation, revokes refresh token on Remove. + +## Process Registration + +Sandbox operations are exposed as Yao Processes under the `sandbox` namespace. + +```go +func init() { + process.Register("sandbox", handler) +} +``` + +| Process | Args | Returns | +|---------|------|---------| +| `sandbox.pool.Add` | `pool` (Pool JSON) | PoolInfo | +| `sandbox.pool.Remove` | `name`, `force?` | — | +| `sandbox.pool.List` | — | []PoolInfo | +| `sandbox.Create` | `options` (CreateOptions JSON) | BoxInfo | +| `sandbox.Get` | `id` | BoxInfo | +| `sandbox.GetOrCreate` | `options` | BoxInfo | +| `sandbox.Remove` | `id` | — | +| `sandbox.List` | `options` (ListOptions JSON) | []BoxInfo | +| `sandbox.Start` | `id` | — | +| `sandbox.Stop` | `id` | — | +| `sandbox.Exec` | `id`, `cmd[]`, `options?` | ExecResult | +| `sandbox.Stream` | `id`, `cmd[]`, `options?` | stream (chunked output) | +| `sandbox.Attach` | `id`, `port`, `options?` | ServiceConn info | +| `sandbox.ReadFile` | `id`, `path` | file content (string) | +| `sandbox.WriteFile` | `id`, `path`, `content` | — | +| `sandbox.ListDir` | `id`, `path` | []FileInfo | +| `sandbox.RemoveFile` | `id`, `path` | — | +| `sandbox.MkdirAll` | `id`, `path` | — | +| `sandbox.VNC` | `id` | URL string | +| `sandbox.Proxy` | `id`, `port`, `path?` | URL string | + +This allows any Yao script, Flow, or API to use sandbox: + +```json +{ + "process": "sandbox.Exec", + "args": ["sb-001", ["go", "build", "./..."]] +} +``` + +## JSAPI + +Global constructor function registered in `gou/runtime/v8`, following the `FS()` / `Store()` pattern. + +```javascript +// Pool management +Sandbox.AddPool({ name: "gpu2", addr: "tai://gpu2.internal" }) +Sandbox.RemovePool("gpu2") +var pools = Sandbox.Pools() +// [{ name: "local", addr: "local", connected: true, boxes: 3 }, ...] + +// Get or create a sandbox +var sb = Sandbox("my-workspace", { + image: "yaoapp/workspace:latest", + owner: "user-123" +}) + +// File operations (fs.FS semantics) +var content = sb.ReadFile("src/main.go") +sb.WriteFile("src/main.go", "package main\n...") +var entries = sb.ListDir("src/") +var info = sb.Stat("src/main.go") +sb.MkdirAll("src/components") +sb.Remove("tmp.txt") +sb.Rename("old.go", "new.go") + +// Command execution — wait for result +var result = sb.Exec(["go", "build", "./..."]) +// result.exit_code, result.stdout, result.stderr + +// Streaming execution — real-time output +sb.Stream(["npm", "run", "dev"], function(chunk) { + log.Info(chunk) // real-time stdout/stderr + return 1 // 1=continue, 0=stop +}) + +// Connect to a service inside the sandbox +var conn = sb.Attach(3000, { protocol: "ws", path: "/ws" }) +conn.Write('{"type":"ping"}') +var msg = conn.Read() +conn.Close() + +// Network +var vncUrl = sb.VNCUrl() +var previewUrl = sb.ProxyUrl(3000, "/") + +// Info +var info = sb.Info() +// info.id, info.status, info.owner, info.created_at + +// Lifecycle +sb.Stop() +sb.Start() +sb.Remove() + +// Properties +sb.id // sandbox ID +sb.workdir // container working directory +``` + +Registration in `gou/runtime/v8/isolate.go`: + +```go +template.Set("Sandbox", sandboxT.New().ExportFunction(iso)) +``` + +Implementation: `gou/runtime/v8/objects/sandbox/sandbox.go` — wraps `sandbox.M().GetOrCreate()` + `Box` methods, using `bridge.GoValue` / `bridge.JsValue` for type conversion. + +## Bootstrap — Manager.Start() + +On `Manager.Start()`, the Manager recovers all existing sandboxes and starts the cleanup loop: + +``` +1. For each pool: + tai.Client.Sandbox().List(labels: {"managed-by": "yao-sandbox"}) + → discover running/stopped containers + +2. For each discovered container: + Parse labels → extract sandbox ID, owner, policy, pool name + Rebuild Box struct, register in boxes map + Set lastCall = now (grace period after restart) + +3. Start cleanupLoop goroutine +``` + +Containers are identified by the label `managed-by=yao-sandbox` plus `sandbox-id=`. Manager injects these labels at creation time. On restart, it queries each pool for containers with `managed-by=yao-sandbox` and rebuilds the in-memory state. + +**What happens to orphaned containers** (created by old Manager, no longer matching any pool): +- If a pool is removed from config, its containers are invisible to the new Manager +- They stay running in Docker/K8s until manually cleaned or TTL-expired by the runtime +- This is by design — Manager only manages containers it can reach + +Startup sequence in `cmd/start.go`: + +``` +sandbox.Init(config.Conf.Sandbox) // create Manager with pool + guard rails +sandbox.M().Start(ctx) // discover existing containers, start cleanup loop +``` + +## Container Setup — Manager.Create() + +When Manager creates a sandbox, it: + +1. Validates `CreateOptions` (Image required) +2. Generates sandbox ID (or uses provided one) +3. Checks user limits (`MaxPerUser`) and total limits (`MaxTotal`) +4. Resolves pool (by name or default) +5. Creates OAuth token pair for container IPC via `openapi/oauth` +6. Builds `tai.sandbox.CreateOptions` from caller's `CreateOptions`: + - Image, Cmd (`sleep infinity`), User — all from caller + - Field name mapping: v2 `WorkDir` → tai `WorkingDir` + - Merges caller's Env with IPC env vars: + - `YAO_GRPC_ADDR`, `YAO_TOKEN`, `YAO_REFRESH_TOKEN`, `YAO_SANDBOX_ID` + - Remote mode: `YAO_GRPC_TAI=enable`, `YAO_GRPC_UPSTREAM` + - Memory/CPU limits, VNC flag, port mappings — all from caller + - Injects management labels: + - `managed-by=yao-sandbox` + - `sandbox-id=` + - `sandbox-owner=` + - `sandbox-pool=` + - `sandbox-policy=` +7. Calls `tai.Client.Sandbox().Create()` then `Start()` +8. Wraps in a `Box`, registers in `boxes` map +9. Starts idle tracking + +## Lifecycle Management + +### Idle Tracking — Dual Source + +Idle is determined by two sources, taking the most recent of both: + +```go +box.lastActive = max(lastExternalCall, lastHeartbeat) +``` + +| Source | What it tracks | Updated by | +|--------|---------------|------------| +| External call | Caller is using the sandbox | `Box.Exec()`, `Box.Workspace()`, `Box.VNC()`, `Box.Proxy()` | +| Container heartbeat | Processes running inside the container | `yao-grpc` → gRPC `Heartbeat` RPC | + +**Why both**: external calls alone miss "user walked away but `npm run build` is still running". Heartbeat alone misses "user is reading output, hasn't issued a new command yet". Together they cover all cases. + +### Heartbeat — Container Side + +`yao-grpc` (already running inside every container) runs a background goroutine: + +``` +Every 30 seconds: + 1. Count user processes (ps aux, exclude sleep/init/yao-grpc) + 2. Count gRPC calls forwarded in last 30s (internal counter) + 3. If either > 0 → send Heartbeat(sandbox_id, active=true, process_count=N) + else → don't send (silent = idle) +``` + +~30 lines added to `tai/grpc/cmd/main.go`. Zero new dependencies. + +### Heartbeat — Server Side + +New gRPC RPC in `yao.proto`: + +```protobuf +rpc Heartbeat(HeartbeatRequest) returns (HeartbeatResponse); + +message HeartbeatRequest { + string sandbox_id = 1; + bool active = 2; + int32 process_count = 3; +} +message HeartbeatResponse {} +``` + +Handler (~20 lines in `grpc/sandbox/`): looks up Box by `sandbox_id`, updates `lastHeartbeat`. Auth: reuses container's `YAO_TOKEN`, no new scope needed (piggyback on existing `grpc:mcp`). + +### Idle Decision Matrix + +| External calls | Heartbeat | Judgment | Action | +|---------------|-----------|----------|--------| +| Recent | Recent | Active | None | +| Recent | Silent | Active | None (user reading output) | +| None | Recent | Active | None (build/server still running) | +| None | Silent | **Idle** | Policy-based stop/remove | + +### Cleanup Loop + +```go +func (m *Manager) cleanupLoop(ctx context.Context) { + ticker := time.NewTicker(1 * time.Minute) + for { + select { + case <-ticker.C: + m.Cleanup(ctx) + case <-ctx.Done(): + return + } + } +} + +func (m *Manager) Cleanup(ctx context.Context) error { + now := time.Now() + m.boxes.Range(func(key, value any) bool { + box := value.(*Box) + idle := now.Sub(box.lastActiveTime()) // max(external, heartbeat) + + switch box.policy { + case OneShot: + // already removed after Exec + case Session: + if idle > box.idleTimeout() { box.Remove(ctx) } + case LongRunning: + if idle > box.idleTimeout() { box.Stop(ctx) } + if lifetime > box.maxLifetime() { box.Remove(ctx) } + case Persistent: + // never auto-cleaned + } + return true + }) + return nil +} +``` + +### Policy Behavior + +| Policy | Idle | Max Lifetime | Auto | +|--------|------|-------------|------| +| OneShot | — | — | Removed after first Exec completes | +| Session | Stop + Remove | Remove | Default for agent chats | +| LongRunning | Stop (keep data) | Remove | User workspaces | +| Persistent | Never | Never | User-managed | + +## Package Structure + +``` +sandbox/v2/ +├── sandbox.go // Init, M(), global singleton +├── manager.go // Manager struct, Create/Get/List/Remove/Cleanup +├── box.go // Box struct, Exec/Workspace/VNC/Proxy/lifecycle +├── config.go // Config, env parsing +├── types.go // CreateOptions, ExecResult, BoxInfo, enums +├── errors.go // sentinel errors +├── process.go // Yao Process registration (sandbox.*) +├── grpc.go // token creation + gRPC env var injection for containers +├── jsapi/ +│ └── sandbox.go // V8 JSAPI: Sandbox() constructor (lives in gou) +└── DESIGN.md // this document +``` + +## Tai SDK Changes Required + +Sandbox V2 needs changes in `tai/` and `yao/grpc` before Phase 1 can fully work. These are **prerequisites** — the sandbox module itself has zero Docker/K8s awareness, so all runtime capabilities must exist in tai; heartbeat support requires additions to both the gRPC server and the in-container client. + +### 1. `tai/sandbox` — Add `ExecStream` (streaming exec) + +Current `Exec()` buffers all output and returns `ExecResult` after the process exits. `Box.Stream()` needs a streaming variant. + +```go +// tai/sandbox — new method on the Sandbox interface +type ExecStream struct { + Stdout io.ReadCloser + Stderr io.ReadCloser + Stdin io.WriteCloser + Wait func() (int, error) // blocks until exit, returns exit code + Cancel func() // kills the exec process +} + +func (s *Sandbox) ExecStream(ctx context.Context, containerID string, cmd []string, opts ...ExecOption) (*ExecStream, error) +``` + +Implementation per runtime: + +| Runtime | How | +|---------|-----| +| **Docker** (`docker_core.go`) | `ContainerExecCreate` + `ContainerExecAttach` — already returns a `HijackedResponse` with a raw stream. Current code pipes it into buffers; change to expose `io.ReadCloser` directly. `Cancel` calls `ContainerExecInspect` loop → kill. ~40 lines changed. | +| **K8s** (`k8s.go`) | `remotecommand.NewSPDYExecutor` + `StreamWithContext` — already supports streaming. Current code passes `bytes.Buffer`; change to pass `io.Pipe()`. ~30 lines changed. | + +Both runtimes already have the raw streaming capability — the change is to **stop buffering** and expose the stream directly. + +### 2. `tai/proxy` — Add `Connect` (bidirectional connection) + +Current `proxy.Proxy` only returns a URL string (`Resolve()`). `Box.Attach()` needs an actual connection. + +```go +// tai/proxy — new method +type ConnectOptions struct { + Protocol string // "ws", "sse", "tcp"; default "ws" + Path string // URL path, e.g. "/v1/chat" + Headers map[string]string // extra request headers +} + +type Connection struct { + Read func() ([]byte, error) // read next message/event + Write func(data []byte) error // send data (no-op for SSE) + Events <-chan []byte // non-nil for SSE mode + URL string // resolved URL for reference + Close func() error +} + +func (p *Proxy) Connect(ctx context.Context, containerID string, port int, opts ConnectOptions) (*Connection, error) +``` + +Implementation: + +| Mode | How | +|------|-----| +| **Local** | Direct dial to `containerIP:port`. WebSocket via `gorilla/websocket` or `nhooyr.io/websocket`. SSE via `http.Get` + chunked read. TCP via `net.Dial`. | +| **Remote** | Dial through Tai HTTP proxy: `http://tai-host:8080/{containerID}:{port}/{path}`. Tai proxy already handles WebSocket upgrade and SSE streaming natively (`http.Hijacker` for WS, `FlushInterval: -1` for SSE). No Tai server changes needed. | + +The Tai HTTP proxy server (`tai/httpproxy/router.go`) already supports: +- **WebSocket**: detects `Upgrade: websocket` header, does TCP-level bidirectional relay +- **SSE**: reverse proxy with `FlushInterval: -1`, streams through transparently +- **Regular HTTP**: standard `httputil.ReverseProxy` + +So the `Connect` implementation in `tai/proxy` is a **client-side** addition only. The server side is ready. + +### 3. `tai/sandbox` — Add `Labels` and `User` to `CreateOptions` + +Current `tai/sandbox.CreateOptions` is missing two fields Manager needs: + +- **`Labels`**: for container discovery on restart (`managed-by=yao-sandbox`, `sandbox-id`, etc.) +- **`User`**: to run container processes as a specific user + +```go +// tai/sandbox — add to existing CreateOptions struct +type CreateOptions struct { + // ... existing fields (Name, Image, Cmd, Env, Binds, WorkingDir, Memory, CPUs, VNC, Ports) ... + Labels map[string]string // container/pod labels for discovery and management + User string // container user, e.g. "1000:1000" +} +``` + +Implementation: + +| Runtime | Field | How | +|---------|-------|-----| +| **Docker** | `Labels` | Set `cfg.Labels = opts.Labels` in `create()`. ~1 line. | +| **Docker** | `User` | Set `cfg.User = opts.User` in `create()`. ~1 line. | +| **K8s** | `Labels` | Set `pod.ObjectMeta.Labels` in `CreatePod`. ~1 line. | +| **K8s** | `User` | Set `SecurityContext.RunAsUser` in pod spec. ~3 lines. | + +`List` with label filtering is **already implemented** in both runtimes: +- Docker: `filters.NewArgs("label", k+"="+v)` in `docker_core.go:175` +- K8s: `metav1.ListOptions{LabelSelector: ...}` in `k8s.go:255` + +`ListOptions.Labels` field also already exists in `sandbox.go:68`. No changes needed for List. + +Also needed: **`ContainerInfo` must include `Labels`**. Current `ContainerInfo` struct has no `Labels` field. `Manager.Start()` discovers existing containers via `List()` and needs to read labels (`sandbox-id`, `sandbox-owner`, `sandbox-policy`, `sandbox-pool`) to rebuild Box state. + +```go +// tai/sandbox — add to existing ContainerInfo struct +type ContainerInfo struct { + // ... existing fields (ID, Name, Image, Status, IP, Ports) ... + Labels map[string]string // container/pod labels +} +``` + +| Runtime | How | +|---------|-----| +| **Docker** | `list()`: read `c.Labels` from `ContainerList` response. `inspect()`: read `info.Config.Labels`. ~1 line each. | +| **K8s** | `list()`: read `pod.Labels` from `PodList` response. ~1 line. | + +### 4. `yao/grpc` + `tai/grpc` — Heartbeat RPC + +Manager uses dual idle tracking (external API calls + container heartbeat). The heartbeat path requires additions on both sides: the gRPC server (new RPC) and `yao-grpc` in-container client (new background goroutine). + +#### Server side — `yao/grpc` + +New RPC in `grpc/pb/yao.proto`: + +```protobuf +rpc Heartbeat(HeartbeatRequest) returns (HeartbeatResponse); + +message HeartbeatRequest { + string sandbox_id = 1; + bool active = 2; // true if user processes detected + int32 process_count = 3; // number of user processes +} +message HeartbeatResponse {} +``` + +Handler in `grpc/sandbox/` (~20 lines): + +```go +func (s *Server) Heartbeat(ctx context.Context, req *pb.HeartbeatRequest) (*pb.HeartbeatResponse, error) { + box, err := sandbox.M().Get(ctx, req.SandboxId) + if err != nil { + return nil, status.Errorf(codes.NotFound, "sandbox %s not found", req.SandboxId) + } + sandbox.M().Heartbeat(req.SandboxId, req.Active, int(req.ProcessCount)) + return &pb.HeartbeatResponse{}, nil +} +``` + +Auth: reuses container's `YAO_TOKEN` — no new OAuth scope needed. The token is already issued with gRPC access when Manager creates the container. + +#### Client side — `tai/grpc/cmd/main.go` (`yao-grpc`) + +New background goroutine (~30 lines) added to `yao-grpc` startup: + +```go +func heartbeatLoop(ctx context.Context, client *grpc.Client, sandboxID string) { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for { + select { + case <-ticker.C: + count := countUserProcesses() // ps aux, exclude sleep/init/yao-grpc + active := count > 0 + if active { + client.Heartbeat(ctx, sandboxID, true, int32(count)) + } + // silent when idle — no heartbeat sent, Manager tracks absence + case <-ctx.Done(): + return + } + } +} + +func countUserProcesses() int { + // exec `ps -eo comm`, filter out known system processes + // (sleep, init, yao-grpc, sh -c sleep) + // return count of remaining user processes +} +``` + +`yao-grpc` reads `YAO_SANDBOX_ID` from env (injected by Manager at container creation). If empty, heartbeat is disabled (container not managed by sandbox). + +#### Heartbeat flow + +``` +Container (every 30s) Yao Server +───────────────────── ────────── +countUserProcesses() + ├── active (count > 0) + │ └── yao-grpc → Heartbeat RPC ──→ grpc/sandbox/Heartbeat() + │ └── sandbox.M().Heartbeat(id, true, N) + │ └── box.lastHeartbeat = now + │ box.processCount = N + └── idle (count == 0) + └── (no RPC sent) Manager sees: no heartbeat in 30s+ + └── combined with no external calls → idle +``` + +Key behaviors: +- **Only sends when active** — idle containers are silent, reducing gRPC traffic +- **30s interval** — matches Manager cleanup loop granularity (1 min), two missed heartbeats = considered idle +- **Crash-safe** — if `yao-grpc` dies, heartbeats stop, Manager treats it as idle after timeout +- **Zero new dependencies** — `yao-grpc` already has the gRPC client connection; heartbeat piggybacks on it + +### Summary + +| Change | Package | Effort | Blocks | +|--------|---------|--------|--------| +| `ExecStream` | `tai/sandbox` | ~40 lines Docker + ~30 lines K8s | `Box.Stream()` | +| `Connect` | `tai/proxy` | ~80 lines (client-side only, server ready) | `Box.Attach()` | +| `Labels` + `User` in `CreateOptions` | `tai/sandbox` | ~6 lines (Docker + K8s) | `Manager.Create()` labeling + user | +| `Labels` in `ContainerInfo` | `tai/sandbox` | ~3 lines (Docker list/inspect + K8s list) | `Manager.Start()` container discovery | +| `Heartbeat` RPC | `yao/grpc` | ~20 lines handler + 3 lines proto | `Manager.Heartbeat()` | +| Heartbeat goroutine | `tai/grpc` (`yao-grpc`) | ~30 lines | Container → Server heartbeat | + +`List` with label filtering is already implemented in both Docker and K8s runtimes — no changes needed. + +All changes are additive (no breaking changes to existing APIs). `Box.Exec()` and `Box.Workspace()` work with current tai — only Stream, Attach, and idle tracking need the new methods. + +## Migration Plan + +### Phase 1: Core module + +Build `sandbox/v2` as a standalone package. No agent dependency. + +**Tai / gRPC prerequisites** (do first): + +| Task | Detail | +|------|--------| +| `tai/sandbox`: `ExecStream` | Streaming exec for Docker + K8s (~70 lines total) | +| `tai/proxy`: `Connect` | Client-side WebSocket/SSE/TCP connection (~80 lines) | +| `tai/sandbox`: `Labels` + `User` in `CreateOptions` | Add fields + wire into Docker/K8s create (~6 lines). List filter already done. | +| `tai/sandbox`: `Labels` in `ContainerInfo` | Add field + populate in Docker list/inspect, K8s list (~3 lines) | +| `yao/grpc`: `Heartbeat` RPC | Proto + handler (~20 lines) | +| `tai/grpc` (`yao-grpc`): heartbeat goroutine | Process detection + periodic report (~30 lines) | + +**Sandbox V2 module:** + +| Task | Detail | +|------|--------| +| `sandbox.go` | `Init()`, `M()`, singleton lifecycle | +| `config.go` | Config struct, defaults | +| `types.go` | CreateOptions, ExecResult, BoxInfo, LifecyclePolicy, Pool | +| `errors.go` | ErrNotAvailable, ErrNotFound, ErrLimitExceeded | +| `manager.go` | Manager with tai.Client pool. Create/Get/GetOrCreate/List/Remove/Cleanup/Close | +| `box.go` | Box wrapping tai Sandbox/Volume/Workspace/Proxy/VNC. Dual idle tracking (lastCall + lastHeartbeat) | +| `grpc.go` | OAuth token pair creation, gRPC env var injection | +| Tests | Unit + integration (needs Docker for local mode) | + +### Phase 2: Process + JSAPI + +| Task | Detail | +|------|--------| +| `process.go` | Register `sandbox.*` process namespace | +| `jsapi/sandbox.go` | V8 `Sandbox()` constructor in gou | +| Tests | Process handler tests, JSAPI tests | + +### Phase 3: Agent integration + +In the Agent repo (not in sandbox/v2): + +| Task | Detail | +|------|--------| +| Agent creates Box via `sandbox.M().GetOrCreate()` | Replace `infraSandbox.Manager` | +| Agent uses `Box.Workspace()` for file I/O | Replace Docker Copy/bind mount reads | +| Agent uses `Box.Exec()` for commands | Replace Docker exec | +| Agent uses `Box.VNC()` / `Box.Proxy()` | Replace vncproxy | +| Agent injects `Box` as `SandboxExecutor` | `ctx.sandbox` JSAPI unchanged for hooks | +| `BuildMCPConfigForSandbox()` uses Box env vars | No more hardcoded `/tmp/yao.sock` | + +### Phase 4: Cutover + +| Task | Detail | +|------|--------| +| Move `sandbox/v2` → `sandbox` | Rename package | +| Delete old sandbox code | manager.go, ipc/, bridge/, vncproxy/, docker/ | +| Delete `DESIGN-REMOTE.md` | Superseded by tai.Client | +| Update `cmd/start.go` | Use new init path | + +## What Gets Deleted (Phase 4) + +Everything in the current `sandbox/` that is replaced by tai: + +| Old | Replaced by | +|-----|------------| +| `manager.go` (Docker `*client.Client`) | `tai.Client.Sandbox()` | +| `ipc/` (Unix socket manager) | gRPC via `yao/grpc` + `tai/grpc` | +| `bridge/` (stdio→socket bridge) | `yao-grpc` binary (`tai/grpc/cmd`) | +| `vncproxy/` (Docker-based VNC) | `tai.Client.VNC()` | +| `proxy/` (Claude API proxy) | separate concern, not sandbox | +| `docker/` (Dockerfiles) | kept, they're image build files | +| `DESIGN-REMOTE.md` (Runtime interface) | tai.Client is the abstraction | +| `config.go` (old config) | new config in v2 | +| `helpers.go` (Docker helpers) | not needed | + +## Comparison: V1 vs V2 + +| Aspect | V1 (current) | V2 (this design) | +|--------|-------------|-------------------| +| **Positioning** | Agent's Claude executor | Yao infrastructure module | +| **Runtime** | Direct Docker SDK | tai.Client pool (Docker/K8s/Remote) | +| **Execution** | Exec + Stream | Exec + Stream + Attach (service connections) | +| **File I/O** | bind mount + Docker Copy | `workspace.FS` (fs.FS compatible) | +| **IPC** | Unix socket + yao-bridge | gRPC (yao-grpc, already done) | +| **Idle detection** | External calls only | Dual: external calls + container heartbeat | +| **Lifecycle** | Chat session only | Policy-based (oneshot/session/longrunning/persistent) | +| **Pool** | Single Docker daemon | Multi-pool with per-pool policies, dynamic add/remove | +| **Agent coupling** | Tightly coupled | Zero dependency | +| **JSAPI** | Only `ctx.sandbox` in hooks | Global `Sandbox()` + `ctx.sandbox` | +| **Process** | None | `sandbox.*` namespace | +| **Multi-node** | Local only | Local + Remote via Tai | +| **K8s** | Not supported | Supported via tai.Client | diff --git a/sandbox/v2/IMPL.md b/sandbox/v2/IMPL.md new file mode 100644 index 00000000..7c4a556f --- /dev/null +++ b/sandbox/v2/IMPL.md @@ -0,0 +1,704 @@ +# Sandbox V2 — Implementation Plan + +Phase 1 implementation. Covers tai SDK prerequisites + sandbox/v2 core Go API. +No JSAPI, no Process registration — those are Phase 2. + +Reference: [DESIGN.md](./DESIGN.md) + +## Execution Order + +``` +Step 0: tai/sandbox — Labels, User, ContainerInfo.Labels (no deps) +Step 1: tai/sandbox — ExecStream (no deps) +Step 2: tai/proxy — Connect (no deps) +Step 3: yao/grpc — Heartbeat RPC (proto + handler) (no deps) +Step 4: tai/grpc — yao-grpc heartbeat goroutine (depends on Step 3 proto) +Step 4.5: docker — build v2 test images (depends on Steps 1–4) +Step 5: sandbox/v2 — core module (depends on Steps 0–4) +Step 6: tests (depends on Steps 5 + 4.5) +``` + +Steps 0–3 are independent and can be parallelized. +Step 4.5 (images) depends on tai SDK + yao-grpc changes being compiled into binaries. + +--- + +## Step 0: `tai` — Labels, User, ContainerInfo.Labels + `tai.New("local")` + +**Files:** `tai/tai.go`, `tai/sandbox/sandbox.go`, `tai/sandbox/docker_core.go`, `tai/sandbox/k8s.go` + +### 0.0 `tai.New("")` → error, add `"local"` / `"127.0.0.1"` aliases + +```go +// tai/tai.go — parseAddr changes: +// - addr == "" → return error ("use local") +// - addr == "local" || addr == "127.0.0.1" → return "docker", "", "" (platform default socket) +``` + +All callers must use explicit addresses. `"local"` means platform-default Docker daemon. + +### 0.1 Add `Labels` and `User` to `CreateOptions` + +```go +// sandbox.go — add two fields to existing struct +type CreateOptions struct { + // ... existing fields ... + Labels map[string]string + User string +} +``` + +### 0.2 Wire into Docker create + +```go +// docker_core.go — in create(), after building cfg: +cfg.Labels = opts.Labels +if opts.User != "" { + cfg.User = opts.User +} +``` + +### 0.3 Wire into K8s create + +```go +// k8s.go — in Create(), set pod labels: +pod.ObjectMeta.Labels = mergeLabels(pod.ObjectMeta.Labels, opts.Labels) + +// For User, parse and set SecurityContext.RunAsUser +``` + +### 0.4 Add `Labels` to `ContainerInfo` + +```go +// sandbox.go +type ContainerInfo struct { + // ... existing fields ... + Labels map[string]string +} +``` + +### 0.5 Populate Labels in Docker list/inspect + +```go +// docker_core.go — in list(): +ci.Labels = c.Labels + +// docker_core.go — in inspect(): +ci.Labels = info.Config.Labels +``` + +### 0.6 Populate Labels in K8s list + +```go +// k8s.go — in List(): +ci.Labels = pod.Labels +``` + +### 0.7 Tests + +- `TestCreateWithLabels` — create container with labels, list with label filter, verify match +- `TestCreateWithUser` — create container with user, exec `whoami`, verify +- `TestListLabels` — create 2 containers with different labels, list with filter, verify count + +**Estimated: ~20 lines code + ~60 lines tests** + +--- + +## Step 1: `tai/sandbox` — ExecStream + +**Files:** `tai/sandbox/sandbox.go`, `tai/sandbox/docker_core.go`, `tai/sandbox/k8s.go` + +### 1.1 Add to Sandbox interface + +```go +// sandbox.go +type ExecStream struct { + Stdout io.ReadCloser + Stderr io.ReadCloser + Stdin io.WriteCloser + Wait func() (int, error) + Cancel func() +} + +// Add to Sandbox interface: +ExecStream(ctx context.Context, id string, cmd []string, opts ExecOptions) (*ExecStream, error) +``` + +### 1.2 Docker implementation + +```go +// docker_core.go — new method +func (d *dockerCore) execStream(ctx context.Context, id string, cmd []string, opts ExecOptions) (*ExecStream, error) { + execCfg := container.ExecOptions{ + Cmd: cmd, + WorkingDir: opts.WorkDir, + Env: envSlice(opts.Env), + AttachStdout: true, + AttachStderr: true, + AttachStdin: true, + } + execResp, err := d.cli.ContainerExecCreate(ctx, id, execCfg) + // ... + resp, err := d.cli.ContainerExecAttach(ctx, execResp.ID, container.ExecAttachOptions{}) + // ... + // Use io.Pipe + stdcopy.StdCopy in a goroutine to demux stdout/stderr + // Wait: poll ContainerExecInspect until Running=false + // Cancel: context cancel → close resp.Conn +} +``` + +Key: `ContainerExecAttach` returns `HijackedResponse` with multiplexed stream. Use `stdcopy.StdCopy` in a goroutine writing to `io.Pipe` pairs for stdout/stderr separation. + +### 1.3 K8s implementation + +```go +// k8s.go — new method +func (s *k8sSandbox) ExecStream(ctx context.Context, id string, cmd []string, opts ExecOptions) (*ExecStream, error) { + // remotecommand.NewSPDYExecutor + // StreamWithContext using io.Pipe for stdin/stdout/stderr + // Wait: executor returns when process exits + // Cancel: cancel the context +} +``` + +### 1.4 Tests + +- `TestExecStream_ShortCommand` — `echo hello`, read stdout, verify Wait returns 0 +- `TestExecStream_LongRunning` — `sleep 10`, Cancel after 1s, verify cleanup +- `TestExecStream_Stdin` — `cat`, write to stdin, read from stdout, verify echo +- `TestExecStream_ExitCode` — `exit 42`, verify Wait returns 42 + +**Estimated: ~80 lines code + ~100 lines tests** + +--- + +## Step 2: `tai/proxy` — Connect + +**Files:** `tai/proxy/proxy.go`, `tai/proxy/connect.go` (new) + +### 2.1 Add to Proxy interface + +```go +// proxy.go — extend interface +type Proxy interface { + URL(ctx context.Context, containerID string, port int, path string) (string, error) + Connect(ctx context.Context, containerID string, port int, opts ConnectOptions) (*Connection, error) + Healthz(ctx context.Context) error +} + +type ConnectOptions struct { + Protocol string // "ws", "sse", "tcp"; default "ws" + Path string + Headers map[string]string +} + +type Connection struct { + Read func() ([]byte, error) + Write func(data []byte) error + Events <-chan []byte + URL string + Close func() error +} +``` + +### 2.2 Implementation — `connect.go` + +Local: resolve URL via `URL()`, then dial directly. +Remote: resolve URL via `URL()` (points to Tai HTTP proxy), then dial. + +Both modes use the same dialing logic after URL resolution: +- **WebSocket**: `gorilla/websocket.Dialer.DialContext` +- **SSE**: `http.Get` + chunked body reader, parse `data:` lines into Events channel +- **TCP**: `net.Dial` + +### 2.3 Tests + +- `TestConnectWebSocket` — start a WS echo server in container, connect, send/receive +- `TestConnectSSE` — start an SSE server in container, connect, verify events arrive +- Skip TCP for now (less common use case) + +**Estimated: ~120 lines code + ~80 lines tests** + +--- + +## Step 3: `yao/grpc` — Heartbeat RPC + +**Files:** `grpc/pb/yao.proto`, `grpc/sandbox/heartbeat.go` (new), `grpc/api/api.go` + +### 3.1 Proto + +```protobuf +// Add to service Yao: +rpc Heartbeat(HeartbeatRequest) returns (HeartbeatResponse); + +message HeartbeatRequest { + string sandbox_id = 1; + bool active = 2; + int32 process_count = 3; +} +message HeartbeatResponse {} +``` + +Regenerate: `protoc --go_out=. --go-grpc_out=. grpc/pb/yao.proto` + +### 3.2 Handler + +```go +// grpc/sandbox/heartbeat.go +func (s *Server) Heartbeat(ctx context.Context, req *pb.HeartbeatRequest) (*pb.HeartbeatResponse, error) { + err := sandbox.M().Heartbeat(req.SandboxId, req.Active, int(req.ProcessCount)) + if err != nil { + return nil, status.Errorf(codes.NotFound, "sandbox %s: %v", req.SandboxId, err) + } + return &pb.HeartbeatResponse{}, nil +} +``` + +### 3.3 Register in server + +Wire into `grpc/api/api.go` server registration (same pattern as Healthz). + +### 3.4 ACL virtual endpoint + +Add to `grpc/auth/endpoints.go`: + +```go +// Heartbeat → POST /grpc/heartbeat (reuse existing container token scope) +``` + +### 3.5 Tests + +- `TestHeartbeat_Success` — create sandbox, send heartbeat, verify no error +- `TestHeartbeat_NotFound` — send heartbeat with unknown sandbox_id, verify NotFound +- `TestHeartbeat_Auth` — verify token auth works (reuse testutils) + +**Estimated: ~40 lines code + ~50 lines tests** + +--- + +## Step 4: `tai/grpc` — yao-grpc heartbeat goroutine + +**Files:** `tai/grpc/cmd/main.go` (or equivalent entry point), `tai/grpc/heartbeat.go` (new) + +### 4.1 Heartbeat loop + +```go +// tai/grpc/heartbeat.go +func heartbeatLoop(ctx context.Context, client *grpc.Client, sandboxID string) { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for { + select { + case <-ticker.C: + count := countUserProcesses() + if count > 0 { + client.Heartbeat(ctx, sandboxID, true, int32(count)) + } + case <-ctx.Done(): + return + } + } +} + +func countUserProcesses() int { + // exec: ps -eo comm --no-headers + // filter out: sleep, init, yao-grpc, sh, bash (if parent is sleep) + // return count +} +``` + +### 4.2 Wire into main + +```go +// In main() or NewFromEnv(), after client is connected: +sandboxID := os.Getenv("YAO_SANDBOX_ID") +if sandboxID != "" { + go heartbeatLoop(ctx, client, sandboxID) +} +``` + +Note: `heartbeatLoop` calls `client.Heartbeat()` (new public method on `tai/grpc.Client`), not `client.svc` (private). + +### 4.3 Tests + +- `TestCountUserProcesses` — unit test for process filtering logic +- `TestHeartbeatLoop_SendsWhenActive` — mock gRPC client, start background process, verify heartbeat sent +- `TestHeartbeatLoop_SilentWhenIdle` — no user processes, verify no RPC calls + +**Estimated: ~40 lines code + ~40 lines tests** + +--- + +## Step 4.5: Docker — V2 Test Images + +**Depends on:** Steps 1–4 (tai SDK ExecStream, proxy Connect, yao-grpc heartbeat) + +Sandbox V2 tests need containers that have `yao-grpc` (with heartbeat) pre-installed. Also need a test-specific image with Nginx for Attach WS/SSE testing. + +**Files:** `sandbox/docker/v2/` (new directory) + +### Image hierarchy + +``` +sandbox-v2-base ← base + yao-grpc + claude-proxy +sandbox-v2-test ← v2-base + nginx (WS echo + SSE endpoint) +``` + +### 4.5.1 `sandbox/docker/v2/Dockerfile.base` + +```dockerfile +FROM yaoapp/sandbox-base:latest + +# Replace yao-bridge with yao-grpc +ARG TARGETARCH +COPY yao-grpc-${TARGETARCH} /usr/local/bin/yao-grpc +RUN chmod +x /usr/local/bin/yao-grpc + +# Claude API proxy (OpenAPI-compatible) +COPY claude-proxy-${TARGETARCH} /usr/local/bin/claude-proxy +RUN chmod +x /usr/local/bin/claude-proxy + +# yao-grpc auto-start: if YAO_SANDBOX_ID is set, start heartbeat + serve +COPY entrypoint.sh /usr/local/bin/entrypoint.sh +RUN chmod +x /usr/local/bin/entrypoint.sh + +WORKDIR /workspace +USER sandbox +ENTRYPOINT ["/usr/local/bin/entrypoint.sh"] +CMD ["sleep", "infinity"] +``` + +### 4.5.2 `sandbox/docker/v2/entrypoint.sh` + +```bash +#!/bin/bash +# Start yao-grpc in background if sandbox env vars are present +if [ -n "$YAO_SANDBOX_ID" ] && [ -n "$YAO_GRPC_ADDR" ]; then + yao-grpc serve & +fi + +# Start claude-proxy if config exists +if [ -n "$CLAUDE_PROXY_BACKEND" ] || [ -f /workspace/.claude-proxy.json ]; then + claude-proxy & +fi + +exec "$@" +``` + +### 4.5.3 `sandbox/docker/v2/Dockerfile.test` + +For unit tests — adds Nginx with a simple WS echo server and SSE endpoint. + +```dockerfile +FROM yaoapp/sandbox-v2-base:latest + +USER root + +# Nginx + test services +RUN apt-get update && apt-get install -y nginx python3 && rm -rf /var/lib/apt/lists/* + +# WS echo server (Python, ~15 lines) +COPY ws-echo.py /opt/test/ws-echo.py + +# SSE endpoint (Python, ~15 lines) +COPY sse-server.py /opt/test/sse-server.py + +# Nginx config — proxy WS on :3000, SSE on :3001 +COPY nginx-test.conf /etc/nginx/sites-available/default + +# Test entrypoint — start nginx + test services + original entrypoint +COPY test-entrypoint.sh /usr/local/bin/test-entrypoint.sh +RUN chmod +x /usr/local/bin/test-entrypoint.sh + +USER sandbox +WORKDIR /workspace +ENTRYPOINT ["/usr/local/bin/test-entrypoint.sh"] +CMD ["sleep", "infinity"] +``` + +### 4.5.4 Test services + +**`ws-echo.py`** — WebSocket echo on port 3000: + +```python +#!/usr/bin/env python3 +import asyncio, websockets +async def echo(ws): + async for msg in ws: + await ws.send(msg) +asyncio.run(websockets.serve(echo, "0.0.0.0", 3000)) +``` + +**`sse-server.py`** — SSE endpoint on port 3001: + +```python +#!/usr/bin/env python3 +from http.server import HTTPServer, BaseHTTPRequestHandler +import time +class Handler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Cache-Control", "no-cache") + self.end_headers() + for i in range(5): + self.wfile.write(f"data: event-{i}\n\n".encode()) + self.wfile.flush() + time.sleep(0.1) +HTTPServer(("0.0.0.0", 3001), Handler).serve_forever() +``` + +**`test-entrypoint.sh`**: + +```bash +#!/bin/bash +python3 /opt/test/ws-echo.py & +python3 /opt/test/sse-server.py & +exec /usr/local/bin/entrypoint.sh "$@" +``` + +### 4.5.5 Build script update + +Add `v2` and `v2-test` targets to `sandbox/docker/build.sh`: + +```bash +v2) + echo "=== Building V2 images ===" + # Build yao-grpc binary (replaces yao-bridge) + cd "$SCRIPT_DIR/../../tai/grpc/cmd" + CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -ldflags="-s -w" -o "$SCRIPT_DIR/v2/yao-grpc-amd64" . + CGO_ENABLED=0 GOOS=linux GOARCH=arm64 go build -ldflags="-s -w" -o "$SCRIPT_DIR/v2/yao-grpc-arm64" . + cd "$SCRIPT_DIR" + + # Build claude-proxy binary + cd "$SCRIPT_DIR/../proxy/cmd/claude-proxy" + CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -ldflags="-s -w" -o "$SCRIPT_DIR/v2/claude-proxy-amd64" . + CGO_ENABLED=0 GOOS=linux GOARCH=arm64 go build -ldflags="-s -w" -o "$SCRIPT_DIR/v2/claude-proxy-arm64" . + cd "$SCRIPT_DIR" + + build_multiarch "sandbox-v2-base" "v2/Dockerfile.base" "$PUSH" + build_multiarch "sandbox-v2-test" "v2/Dockerfile.test" "$PUSH" + ;; +``` + +### 4.5.6 Image usage + +| Image | Purpose | Used by | +|-------|---------|---------| +| `sandbox-v2-base` | Production base for V2 sandboxes. Has `yao-grpc` (heartbeat) + `claude-proxy`. | `Manager.Create()` default image candidate | +| `sandbox-v2-test` | Unit tests. Has WS echo + SSE server for Attach testing. | `SANDBOX_TEST_IMAGE` in CI and local dev | + +### 4.5.7 Env update + +```bash +# env.local.sh — change test image to v2-test +export SANDBOX_TEST_IMAGE="yaoapp/sandbox-v2-test:latest" +``` + +**Estimated: ~5 files (Dockerfiles + scripts + test services), ~100 lines** + +--- + +## Step 5: `sandbox/v2` — Core Module + +**Files:** all in `sandbox/v2/` + +### 5.1 `errors.go` + +```go +var ( + ErrNotAvailable = errors.New("sandbox: not available (no pools configured)") + ErrNotFound = errors.New("sandbox: not found") + ErrLimitExceeded = errors.New("sandbox: limit exceeded") + ErrPoolNotFound = errors.New("sandbox: pool not found") + ErrPoolInUse = errors.New("sandbox: pool has running boxes") +) +``` + +### 5.2 `types.go` + +All type definitions from DESIGN.md: +- `LifecyclePolicy`, `Pool`, `PoolInfo`, `PortMapping` +- `CreateOptions`, `ListOptions` +- `ExecOption`, `ExecResult`, `ExecStream` +- `AttachOption`, `ServiceConn`, `ConnectOptions` +- `BoxInfo` + +### 5.3 `config.go` + +```go +type Config struct { + Pool []Pool +} +``` + +### 5.4 `sandbox.go` — singleton + +```go +var mgr *Manager + +func Init(cfg Config) error { + m, err := newManager(cfg) + if err != nil { return err } + mgr = m + return nil +} + +func M() *Manager { + if mgr == nil { panic("sandbox.Init not called") } + return mgr +} +``` + +### 5.5 `manager.go` + +Core implementation. Key methods: + +| Method | Logic | +|--------|-------| +| `newManager(cfg)` | Parse pool defs, set default pool | +| `Start(ctx)` | For each pool: connect, list containers with `managed-by=yao-sandbox`, rebuild boxes map, start cleanupLoop | +| `Create(ctx, opts)` | Validate → check limits → resolve pool → lazy-connect tai.Client → create OAuth tokens → build tai.CreateOptions (merge env, labels, field mapping) → tai.Create → tai.Start → wrap Box → register | +| `Get(ctx, id)` | Lookup boxes map | +| `GetOrCreate(ctx, opts)` | Get by ID, if not found → Create | +| `List(ctx, opts)` | Filter boxes by owner/pool/labels | +| `Remove(ctx, id)` | Lookup box → tai.Stop → tai.Remove → revoke OAuth token → delete from map | +| `Cleanup(ctx)` | Range boxes, apply policy-based idle/lifetime rules | +| `Close()` | Cancel cleanup loop, close all tai.Clients | +| `Heartbeat(id, active, count)` | Lookup box → update lastHeartbeat + processCount atomics | +| `AddPool(ctx, p)` | Validate name unique → append to poolDefs | +| `RemovePool(ctx, name, force)` | Check no boxes (or force-remove them) → remove from poolDefs → close tai.Client if connected | +| `Pools()` | Return PoolInfo slice | + +Internal helpers: +- `getPool(name)` — lazy-connect tai.Client from poolDefs +- `buildTaiCreateOptions(opts, pool)` — field mapping + env injection + label injection +- `recoverBoxes(ctx, pool, client)` — list + parse labels + rebuild Box structs + +### 5.6 `box.go` + +```go +type Box struct { /* fields from DESIGN.md */ } + +func (b *Box) Exec(ctx, cmd, opts) // b.touch() → tai.Sandbox().Exec(b.containerID, ...) +func (b *Box) Stream(ctx, cmd, opts) // b.touch() → tai.Sandbox().ExecStream(b.containerID, ...) +func (b *Box) Attach(ctx, port, opts) // b.touch() → tai.Proxy().Connect(b.containerID, port, ...) +func (b *Box) Workspace() // lazy-init: tai.Client.Workspace(b.id) +func (b *Box) VNC(ctx) // b.touch() → tai.VNC().URL(b.containerID) +func (b *Box) Proxy(ctx, port, path) // b.touch() → tai.Proxy().URL(b.containerID, port, path) +func (b *Box) Start(ctx) // tai.Sandbox().Start(b.containerID) +func (b *Box) Stop(ctx) // tai.Sandbox().Stop(b.containerID, 10s) +func (b *Box) Remove(ctx) // b.manager.Remove(ctx, b.id) +func (b *Box) Info(ctx) // tai.Sandbox().Inspect + merge with box metadata + +func (b *Box) touch() // b.lastCall.Store(time.Now().UnixMilli()) +func (b *Box) lastActiveTime() Time // max(lastCall, lastHeartbeat) +func (b *Box) idleTimeout() Duration // box-level override or pool default +func (b *Box) maxLifetime() Duration // pool default +``` + +### 5.7 `grpc.go` — OAuth token injection + +```go +func createContainerTokens(sandboxID, owner string) (access, refresh string, err error) +func revokeContainerTokens(refresh string) error +func buildGRPCEnv(pool *Pool, sandboxID, access, refresh string) map[string]string +``` + +Uses `openapi/oauth` to create token pairs. Local mode: `YAO_GRPC_ADDR=127.0.0.1:`. Remote mode: adds `YAO_GRPC_TAI=enable` + `YAO_GRPC_UPSTREAM`. + +**Estimated: ~600 lines code total** + +--- + +## Step 6: Tests + +### 6.1 Test environment + +Two pools configured via env vars (reuse existing CI infrastructure): + +``` +# Local pool — direct Docker +SANDBOX_TEST_LOCAL_ADDR=local (default Docker daemon) + +# Remote pool — via Tai container (same as tai-test job) +SANDBOX_TEST_REMOTE_ADDR=tai://127.0.0.1 (uses TAI_TEST_* ports) +``` + +Skip tests when Docker/Tai unavailable: `t.Skipf`. + +### 6.2 Test files + +| File | Coverage | +|------|----------| +| `sandbox_test.go` | `Init()`, `M()`, singleton behavior | +| `manager_test.go` | Create, Get, GetOrCreate, List, Remove, pool management | +| `manager_lifecycle_test.go` | Start (container discovery), Cleanup, idle tracking | +| `box_test.go` | Exec, Stream, Workspace (ReadFile/WriteFile/MkdirAll), Proxy, VNC, lifecycle | +| `box_attach_test.go` | Attach with WS/SSE (requires service in container) | +| `grpc_test.go` | Token creation/revocation, env var building | + +### 6.3 Key test scenarios + +| Test | What it verifies | +|------|-----------------| +| `TestCreateAndExec` | Create box → exec `echo hello` → verify stdout → remove | +| `TestCreateWithLabels` | Create → inspect labels → list with label filter | +| `TestWorkspace` | Create → WriteFile → ReadFile → verify content match | +| `TestIdleCleanup` | Create with Session + 1s idle timeout → wait → verify removed | +| `TestStartRecovery` | Create → restart Manager → Start → verify box recovered from labels | +| `TestPoolLimits` | Set MaxTotal=1 → create 1 → create 2nd → verify ErrLimitExceeded | +| `TestHeartbeatUpdates` | Create → call Heartbeat → verify lastActive updated | +| `TestStream` | Create → stream `sh -c "echo a; sleep 0.1; echo b"` → verify chunks arrive | +| `TestMultiPool` | Create on local → create on remote → verify both work | + +### 6.4 CI integration + +Add `sandbox-v2-test` job to `unit-test.yml` (same pattern as existing `sandbox-test` + `tai-test`): + +```yaml +sandbox-v2-test: + runs-on: ubuntu-latest + services: + # MongoDB (for Yao runtime) + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-go@v5 + - name: Start Tai container + run: | + docker run -d --name tai \ + -v /var/run/docker.sock:/var/run/docker.sock \ + -p 2375:2375 -p 9100:9100 -p 8080:8080 \ + yaoapp/tai:latest + - name: Build V2 test image + run: | + cd sandbox/docker + bash build.sh v2 + - name: Run tests + env: + SANDBOX_TEST_LOCAL_ADDR: "local" + SANDBOX_TEST_REMOTE_ADDR: "tai://127.0.0.1" + SANDBOX_TEST_IMAGE: "yaoapp/sandbox-v2-test:latest" + TAI_TEST_HOST: "127.0.0.1" + run: go test -v -count=1 ./sandbox/v2/... +``` + +**Estimated: ~400 lines tests** + +--- + +## Summary + +| Step | Package | Lines (code) | Lines (test) | Depends on | +|------|---------|-------------|-------------|------------| +| 0 | `tai/sandbox` | ~20 | ~60 | — | +| 1 | `tai/sandbox` | ~80 | ~100 | — | +| 2 | `tai/proxy` | ~120 | ~80 | — | +| 3 | `yao/grpc` | ~40 | ~50 | — | +| 4 | `tai/grpc` | ~40 | ~40 | Step 3 | +| 4.5 | `sandbox/docker/v2` | ~100 | — | Steps 1–4 | +| 5 | `sandbox/v2` | ~600 | — | Steps 0–4 | +| 6 | `sandbox/v2` | — | ~400 | Steps 5 + 4.5 | +| **Total** | | **~1000** | **~730** | | + +Steps 0–3 can start in parallel. Step 4 needs Step 3's proto. Step 5 needs all prerequisites done. Step 6 runs after Step 5. diff --git a/sandbox/v2/Makefile b/sandbox/v2/Makefile new file mode 100644 index 00000000..8d7d43e9 --- /dev/null +++ b/sandbox/v2/Makefile @@ -0,0 +1,108 @@ +GO ?= go +GOFILES := $(shell find . -name "*.go" -not -path "./docker/*") +PACKAGES := $(shell $(GO) list ./...) +TEST_IMAGE ?= yaoapp/sandbox-v2-test:latest +TEST_TIMEOUT ?= 300s + +# --------------------------------------------------------------------------- +# Local test (Docker only) +# --------------------------------------------------------------------------- +.PHONY: test-local +test-local: + @echo "=== Sandbox V2: local mode ===" + SANDBOX_TEST_IMAGE=$(TEST_IMAGE) \ + $(GO) test $(PACKAGES) -count=1 -v -timeout $(TEST_TIMEOUT) -run '/local' + +# --------------------------------------------------------------------------- +# Remote test (requires Tai server) +# --------------------------------------------------------------------------- +.PHONY: test-remote +test-remote: + @if [ -z "$(SANDBOX_TEST_REMOTE_ADDR)" ]; then \ + echo "SANDBOX_TEST_REMOTE_ADDR not set, skipping remote tests"; \ + exit 0; \ + fi + @echo "=== Sandbox V2: remote mode ($(SANDBOX_TEST_REMOTE_ADDR)) ===" + SANDBOX_TEST_IMAGE=$(TEST_IMAGE) \ + $(GO) test $(PACKAGES) -count=1 -v -timeout $(TEST_TIMEOUT) -run '/remote' + +# --------------------------------------------------------------------------- +# Dual-mode test (local + remote when configured) +# --------------------------------------------------------------------------- +.PHONY: test +test: + @echo "" + @echo "=============================================" + @echo "Sandbox V2 Tests (dual-mode)" + @echo "=============================================" + @echo "Image: $(TEST_IMAGE)" + @echo "Remote: $${SANDBOX_TEST_REMOTE_ADDR:-}" + @echo "" + SANDBOX_TEST_IMAGE=$(TEST_IMAGE) \ + $(GO) test $(PACKAGES) -count=1 -v -timeout $(TEST_TIMEOUT) + @echo "" + @echo "=============================================" + @echo "Sandbox V2 Tests passed" + @echo "=============================================" + +# --------------------------------------------------------------------------- +# CI test (with coverage, used by Makefile at repo root) +# --------------------------------------------------------------------------- +.PHONY: test-ci +test-ci: + @echo "" + @echo "=============================================" + @echo "Sandbox V2 CI Tests" + @echo "=============================================" + echo "mode: count" > coverage.out + @for d in $(PACKAGES); do \ + SANDBOX_TEST_IMAGE=$(TEST_IMAGE) \ + $(GO) test -v -count=1 -timeout $(TEST_TIMEOUT) \ + -covermode=count -coverprofile=profile.out \ + -coverpkg=$$(echo $$d | sed "s/\/test$$//g") \ + $$d > tmp.out; \ + cat tmp.out; \ + if grep -q "^--- FAIL" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "^FAIL" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "^panic:" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "build failed" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + fi; \ + if [ -f profile.out ]; then \ + cat profile.out | grep -v "mode:" >> coverage.out; \ + rm profile.out; \ + fi; \ + done + @echo "" + @echo "=============================================" + @echo "Sandbox V2 CI Tests passed" + @echo "=============================================" + +# --------------------------------------------------------------------------- +# Docker images +# --------------------------------------------------------------------------- +.PHONY: docker-build +docker-build: + ./docker/build.sh build + +.PHONY: docker-push +docker-push: + ./docker/build.sh push + +# --------------------------------------------------------------------------- +# fmt / vet +# --------------------------------------------------------------------------- +.PHONY: fmt +fmt: + gofmt -s -w $(GOFILES) + +.PHONY: vet +vet: + $(GO) vet $(PACKAGES) diff --git a/sandbox/v2/TEST.md b/sandbox/v2/TEST.md new file mode 100644 index 00000000..799709e2 --- /dev/null +++ b/sandbox/v2/TEST.md @@ -0,0 +1,621 @@ +# Sandbox V2 — Test Specification + +Design: [DESIGN.md](./DESIGN.md) | Implementation: [IMPL.md](./IMPL.md) + +## Principles + +- **Black-box testing**: all `*_test.go` files use `package sandbox_test` — tests only access exported API +- **Real containers**: tests create real Docker containers via tai SDK, no mocking +- **Skip when unavailable**: `skipIfNoDocker(t)` / `skipIfNoTai(t)` — CI has Docker and Tai; local dev may not +- **Tests follow implementation**: `*_test.go` lives next to the code it tests +- **Coverage > 80%**: per file and overall + +## Prerequisites + +```bash +source $YAO_SOURCE_ROOT/env.local.sh +``` + +### Docker (required for all container tests) + +Docker daemon must be running. Tests connect via default socket. + +### Tai (required for remote-mode tests only) + +```bash +docker run -d --name tai \ + -v /var/run/docker.sock:/var/run/docker.sock \ + -p 2375:2375 -p 9100:9100 -p 8080:8080 -p 6080:6080 \ + yaoapp/tai:latest +``` + +### Environment Variables + +| Variable | Purpose | Default | +|----------|---------|---------| +| `YAO_TEST_APPLICATION` | Path to `yao-dev-app` | — (required) | +| `YAO_DB_DRIVER` / `YAO_DB_PRIMARY` | Database connection | — (required) | +| `YAO_JWT_SECRET` / `YAO_DB_AESKEY` | Crypto keys (for OAuth token creation) | — (required) | +| `SANDBOX_TEST_IMAGE` | Container image for tests | `yaoapp/sandbox-v2-test:latest` | +| `SANDBOX_TEST_REMOTE_ADDR` | Tai remote address, e.g. `tai://127.0.0.1` | — (skip remote tests if empty) | +| `TAI_TEST_HOST` | Tai HTTP proxy host | `127.0.0.1` | + +## Directory Structure + +``` +sandbox/v2/ +├── sandbox.go +├── sandbox_test.go # Init/M singleton tests +├── manager.go +├── manager_test.go # Create/Get/GetOrCreate/List/Remove, pool management +├── manager_lifecycle_test.go # Start recovery, Cleanup, idle tracking, heartbeat +├── box.go +├── box_test.go # Exec/Stream/Workspace/Proxy/VNC, lifecycle +├── box_attach_test.go # Attach WS/SSE (needs service in container) +├── config.go +├── types.go +├── errors.go +├── grpc.go +├── grpc_test.go # OAuth token creation/revocation, env var building +├── testutils_test.go # shared test helpers (unexported, package sandbox_test) +└── DESIGN.md +``` + +## testutils (internal to sandbox_test) + +Shared helpers in `testutils_test.go` — not a separate package, lives inside `package sandbox_test`. + +```go +// testutils_test.go +package sandbox_test + +// skipIfNoDocker skips the test if Docker is not available. +func skipIfNoDocker(t *testing.T) + +// skipIfNoTai skips the test if SANDBOX_TEST_REMOTE_ADDR is empty. +func skipIfNoTai(t *testing.T) + +// testImage returns SANDBOX_TEST_IMAGE or "yaoapp/sandbox-v2-test:latest". +func testImage() string + +// setupManager initializes sandbox with a local pool, returns cleanup func. +// Calls sandbox.Init + sandbox.M().Start. +func setupManager(t *testing.T) func() + +// setupManagerWithRemote initializes sandbox with local + remote pools. +func setupManagerWithRemote(t *testing.T) func() + +// createTestBox creates a box with defaults and returns it. Registers t.Cleanup for removal. +func createTestBox(t *testing.T, opts ...sandbox.CreateOption) *sandbox.Box +``` + +## How to Write a Test + +### Standard pattern + +```go +// manager_test.go +package sandbox_test + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + sandbox "github.com/yaoapp/yao/sandbox/v2" +) + +func TestCreate(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + + box, err := sandbox.M().Create(context.Background(), sandbox.CreateOptions{ + Image: testImage(), + Owner: "test-user", + }) + require.NoError(t, err) + defer box.Remove(context.Background()) + + assert.NotEmpty(t, box.ID()) + assert.Equal(t, "test-user", box.Owner()) +} + +func TestCreate_NoImage(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + + _, err := sandbox.M().Create(context.Background(), sandbox.CreateOptions{}) + assert.Error(t, err) // Image is required +} +``` + +### Container execution tests + +```go +// box_test.go +package sandbox_test + +func TestExec(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + box := createTestBox(t) + + result, err := box.Exec(context.Background(), []string{"echo", "hello"}) + require.NoError(t, err) + assert.Equal(t, 0, result.ExitCode) + assert.Equal(t, "hello\n", result.Stdout) + assert.Empty(t, result.Stderr) +} + +func TestExec_NonZeroExit(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + box := createTestBox(t) + + result, err := box.Exec(context.Background(), []string{"sh", "-c", "exit 42"}) + require.NoError(t, err) + assert.Equal(t, 42, result.ExitCode) +} +``` + +### Streaming tests + +```go +// box_test.go +package sandbox_test + +func TestStream(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + box := createTestBox(t) + + s, err := box.Stream(context.Background(), []string{"sh", "-c", "echo a; sleep 0.1; echo b"}) + require.NoError(t, err) + + out, _ := io.ReadAll(s.Stdout) + code, err := s.Wait() + assert.NoError(t, err) + assert.Equal(t, 0, code) + assert.Contains(t, string(out), "a\n") + assert.Contains(t, string(out), "b\n") +} + +func TestStream_Cancel(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + box := createTestBox(t) + + s, err := box.Stream(context.Background(), []string{"sleep", "60"}) + require.NoError(t, err) + + s.Cancel() + code, _ := s.Wait() + assert.NotEqual(t, 0, code) // killed +} + +func TestStream_Stdin(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + box := createTestBox(t) + + s, err := box.Stream(context.Background(), []string{"cat"}) + require.NoError(t, err) + + s.Stdin.Write([]byte("hello\n")) + s.Stdin.Close() + + out, _ := io.ReadAll(s.Stdout) + code, _ := s.Wait() + assert.Equal(t, 0, code) + assert.Equal(t, "hello\n", string(out)) +} +``` + +### Workspace tests + +```go +// box_test.go +package sandbox_test + +func TestWorkspace_ReadWrite(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + box := createTestBox(t) + + ws := box.Workspace() + err := ws.WriteFile("test.txt", []byte("hello"), 0644) + require.NoError(t, err) + + data, err := fs.ReadFile(ws, "test.txt") + require.NoError(t, err) + assert.Equal(t, "hello", string(data)) +} + +func TestWorkspace_MkdirAll(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + box := createTestBox(t) + + ws := box.Workspace() + err := ws.MkdirAll("a/b/c", 0755) + require.NoError(t, err) + + info, err := fs.Stat(ws, "a/b/c") + require.NoError(t, err) + assert.True(t, info.IsDir()) +} + +func TestWorkspace_WalkDir(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + box := createTestBox(t) + + ws := box.Workspace() + ws.MkdirAll("src", 0755) + ws.WriteFile("src/main.go", []byte("package main"), 0644) + ws.WriteFile("src/util.go", []byte("package main"), 0644) + + var files []string + fs.WalkDir(ws, "src", func(path string, d fs.DirEntry, err error) error { + if !d.IsDir() { files = append(files, path) } + return nil + }) + assert.Len(t, files, 2) +} +``` + +### Lifecycle tests + +```go +// manager_lifecycle_test.go +package sandbox_test + +func TestIdleCleanup_Session(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + + box, err := sandbox.M().Create(context.Background(), sandbox.CreateOptions{ + Image: testImage(), + Policy: sandbox.Session, + IdleTimeout: 2 * time.Second, + }) + require.NoError(t, err) + + // Box exists + _, err = sandbox.M().Get(context.Background(), box.ID()) + assert.NoError(t, err) + + // Wait for idle + cleanup cycle + time.Sleep(4 * time.Second) + sandbox.M().Cleanup(context.Background()) + + // Box should be gone + _, err = sandbox.M().Get(context.Background(), box.ID()) + assert.ErrorIs(t, err, sandbox.ErrNotFound) +} + +func TestStartRecovery(t *testing.T) { + skipIfNoDocker(t) + + // Phase 1: create a box, then shut down Manager + cleanup1 := setupManager(t) + box, err := sandbox.M().Create(context.Background(), sandbox.CreateOptions{ + Image: testImage(), + Owner: "recovery-test", + }) + require.NoError(t, err) + boxID := box.ID() + cleanup1() // closes Manager, but container stays + + // Phase 2: new Manager, Start should discover the container + cleanup2 := setupManager(t) + defer cleanup2() + + recovered, err := sandbox.M().Get(context.Background(), boxID) + require.NoError(t, err) + assert.Equal(t, boxID, recovered.ID()) + assert.Equal(t, "recovery-test", recovered.Owner()) + + // Clean up + recovered.Remove(context.Background()) +} + +func TestHeartbeat(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + box := createTestBox(t) + + // Simulate heartbeat + err := sandbox.M().Heartbeat(box.ID(), true, 3) + assert.NoError(t, err) + + info, _ := box.Info(context.Background()) + assert.Equal(t, 3, info.ProcessCount) +} + +func TestHeartbeat_NotFound(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + + err := sandbox.M().Heartbeat("nonexistent", true, 1) + assert.ErrorIs(t, err, sandbox.ErrNotFound) +} +``` + +### Pool management tests + +```go +// manager_test.go +package sandbox_test + +func TestPoolLimits_MaxTotal(t *testing.T) { + skipIfNoDocker(t) + + // Init with MaxTotal=1 + err := sandbox.Init(sandbox.Config{ + Pool: []sandbox.Pool{{ + Name: "limited", + Addr: "local", + MaxTotal: 1, + }}, + }) + require.NoError(t, err) + sandbox.M().Start(context.Background()) + defer sandbox.M().Close() + + box1, err := sandbox.M().Create(context.Background(), sandbox.CreateOptions{ + Image: testImage(), + }) + require.NoError(t, err) + defer box1.Remove(context.Background()) + + _, err = sandbox.M().Create(context.Background(), sandbox.CreateOptions{ + Image: testImage(), + }) + assert.ErrorIs(t, err, sandbox.ErrLimitExceeded) +} + +func TestAddPool(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + + err := sandbox.M().AddPool(context.Background(), sandbox.Pool{ + Name: "new-pool", + Addr: "local", + }) + assert.NoError(t, err) + + pools := sandbox.M().Pools() + names := make([]string, len(pools)) + for i, p := range pools { names[i] = p.Name } + assert.Contains(t, names, "new-pool") +} + +func TestRemovePool_InUse(t *testing.T) { + skipIfNoDocker(t) + cleanup := setupManager(t) + defer cleanup() + + box := createTestBox(t) + _ = box + + err := sandbox.M().RemovePool(context.Background(), "local", false) + assert.ErrorIs(t, err, sandbox.ErrPoolInUse) +} +``` + +### Multi-pool tests + +```go +// manager_test.go +package sandbox_test + +func TestMultiPool(t *testing.T) { + skipIfNoDocker(t) + skipIfNoTai(t) + cleanup := setupManagerWithRemote(t) + defer cleanup() + + // Create on local + local, err := sandbox.M().Create(context.Background(), sandbox.CreateOptions{ + Image: testImage(), + Pool: "local", + }) + require.NoError(t, err) + defer local.Remove(context.Background()) + + // Create on remote + remote, err := sandbox.M().Create(context.Background(), sandbox.CreateOptions{ + Image: testImage(), + Pool: "remote", + }) + require.NoError(t, err) + defer remote.Remove(context.Background()) + + // Both should exec + r1, _ := local.Exec(context.Background(), []string{"echo", "local"}) + r2, _ := remote.Exec(context.Background(), []string{"echo", "remote"}) + assert.Equal(t, "local\n", r1.Stdout) + assert.Equal(t, "remote\n", r2.Stdout) +} +``` + +### OAuth / gRPC env injection tests + +```go +// grpc_test.go +package sandbox_test + +func TestBuildGRPCEnv_Local(t *testing.T) { + env := sandbox.BuildGRPCEnv(&sandbox.Pool{Addr: "local"}, "sb-001", "tok", "ref") + assert.Equal(t, "sb-001", env["YAO_SANDBOX_ID"]) + assert.Equal(t, "tok", env["YAO_TOKEN"]) + assert.Equal(t, "ref", env["YAO_REFRESH_TOKEN"]) + assert.NotEmpty(t, env["YAO_GRPC_ADDR"]) + assert.Empty(t, env["YAO_GRPC_TAI"]) +} + +func TestBuildGRPCEnv_Remote(t *testing.T) { + env := sandbox.BuildGRPCEnv(&sandbox.Pool{Addr: "tai://gpu.internal"}, "sb-002", "tok", "ref") + assert.Equal(t, "enable", env["YAO_GRPC_TAI"]) + assert.NotEmpty(t, env["YAO_GRPC_UPSTREAM"]) +} + +func TestCreateContainerTokens(t *testing.T) { + // Requires Yao runtime for OAuth + cleanup := setupManager(t) + defer cleanup() + + access, refresh, err := sandbox.CreateContainerTokens("sb-test", "user-1") + require.NoError(t, err) + assert.NotEmpty(t, access) + assert.NotEmpty(t, refresh) +} +``` + +## Required Test Cases + +| File | Required Cases | +|------|---------------| +| `sandbox_test.go` | `Init` succeeds / `M()` panics before Init / double Init is safe | +| `manager_test.go` | Create / Create with explicit ID / Create no image (error) / Get / Get not found / GetOrCreate / List / List with owner filter / Remove / pool limits MaxTotal / pool limits MaxPerUser / AddPool / RemovePool / RemovePool in use / Pools | +| `manager_lifecycle_test.go` | Start recovery from labels / Cleanup Session idle / Cleanup LongRunning stop then remove / Persistent never cleaned / Heartbeat updates / Heartbeat not found / OneShot removed after Exec | +| `box_test.go` | Exec success / Exec non-zero exit / Exec with WorkDir / Exec with Env / Exec with Timeout / Stream read / Stream cancel / Stream stdin / Workspace ReadFile+WriteFile / Workspace MkdirAll / Workspace Remove / Workspace Rename / Workspace WalkDir / VNC (skip if no VNC image) / Proxy URL / Start+Stop+Start / Info | +| `box_attach_test.go` | Attach WS (skip if no WS server image) / Attach SSE (skip if no SSE server image) | +| `grpc_test.go` | BuildGRPCEnv local / BuildGRPCEnv remote / CreateContainerTokens / RevokeContainerTokens | + +## Makefile + +Add to [Makefile](../../Makefile): + +```makefile +TESTFOLDER_SANDBOX_V2 := $(shell $(GO) list ./sandbox/v2/...) + +.PHONY: unit-test-sandbox-v2 +unit-test-sandbox-v2: + echo "mode: count" > coverage.out + for d in $(TESTFOLDER_SANDBOX_V2); do \ + $(GO) test -tags $(TESTTAGS) -v -timeout=10m \ + -covermode=count -coverprofile=profile.out \ + -coverpkg=$$d \ + $$d > tmp.out; \ + cat tmp.out; \ + if grep -q "^--- FAIL" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "build failed" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + fi; \ + if [ -f profile.out ]; then \ + cat profile.out | grep -v "mode:" >> coverage.out; \ + rm profile.out; \ + fi; \ + done +``` + +## CI Integration + +Add `sandbox-v2-test` job to `unit-test.yml` and `pr-test.yml`: + +```yaml +sandbox-v2-test: + runs-on: ubuntu-latest + services: + mongodb: + image: mongo:6.0 + ports: + - 27017:27017 + env: + MONGO_INITDB_ROOT_USERNAME: root + MONGO_INITDB_ROOT_PASSWORD: "123456" + MONGO_INITDB_DATABASE: test + strategy: + matrix: + go: ["1.25"] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-go@v5 + with: + go-version: ${{ matrix.go }} + + - name: Start Tai container + run: | + docker run -d --name tai \ + -v /var/run/docker.sock:/var/run/docker.sock \ + -p 2375:2375 -p 9100:9100 -p 8080:8080 -p 6080:6080 \ + yaoapp/tai:latest + sleep 3 + + - name: Build V2 test image + run: | + cd sandbox/docker + bash build.sh v2 + + - name: Setup ENV + run: | + echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV + echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV + echo "SANDBOX_TEST_IMAGE=yaoapp/sandbox-v2-test:latest" >> $GITHUB_ENV + echo "SANDBOX_TEST_REMOTE_ADDR=tai://127.0.0.1" >> $GITHUB_ENV + echo "TAI_TEST_HOST=127.0.0.1" >> $GITHUB_ENV + mkdir -p ${{ github.WORKSPACE }}/../app/db + + - name: Run Sandbox V2 Tests + run: make unit-test-sandbox-v2 + + - name: Codecov Report + uses: codecov/codecov-action@v4 + with: + token: ${{ secrets.CODECOV_TOKEN }} +``` + +Key decisions: +- SQLite only — sandbox is infrastructure, not data-model dependent +- Tai container provides remote mode — exercises the full proxy path +- `sandbox-v2-test` as default test image — includes `yao-grpc` (heartbeat), `claude-proxy`, Nginx, WS echo + SSE test services +- CI builds test image from source (Step 4.5) — ensures binary compatibility with latest tai SDK + yao-grpc changes +- Attach tests (WS/SSE) use `sandbox-v2-test` image's built-in test services + +## Coverage + +- Target: >80% per file, >80% overall +- `sandbox.go` (singleton) covered via `sandbox_test.go` +- `manager.go` is the heaviest file — must have dedicated `manager_test.go` + `manager_lifecycle_test.go` +- `box.go` exercises all tai SDK integration points +- `grpc.go` tested with pure unit tests (token generation, env building) + +## Running Tests + +```bash +# All sandbox v2 tests (local Docker only) +make unit-test-sandbox-v2 + +# With remote mode (start Tai first) +SANDBOX_TEST_REMOTE_ADDR=tai://127.0.0.1 make unit-test-sandbox-v2 + +# Single file +go test -v ./sandbox/v2/ -run TestCreate + +# Single test +go test -v ./sandbox/v2/ -run TestExec_NonZeroExit + +# With race detector +go test -race -v ./sandbox/v2/ +``` diff --git a/sandbox/v2/box.go b/sandbox/v2/box.go new file mode 100644 index 00000000..dc98398c --- /dev/null +++ b/sandbox/v2/box.go @@ -0,0 +1,261 @@ +package sandbox + +import ( + "context" + "io" + "sync/atomic" + "time" + + "github.com/yaoapp/yao/tai/proxy" + taisandbox "github.com/yaoapp/yao/tai/sandbox" + "github.com/yaoapp/yao/tai/workspace" +) + +// Box represents a single sandbox instance. +type Box struct { + id string + containerID string + pool string + owner string + policy LifecyclePolicy + labels map[string]string + lastCall atomic.Int64 + lastHeartbeat atomic.Int64 + processCount atomic.Int32 + idleTimeoutD time.Duration + createdAt time.Time + refreshToken string + vnc bool + image string + ws workspace.FS + manager *Manager +} + +func (b *Box) ID() string { return b.id } +func (b *Box) Owner() string { return b.owner } +func (b *Box) ContainerID() string { return b.containerID } +func (b *Box) Pool() string { return b.pool } + +// Exec runs a command and waits for it to finish. +func (b *Box) Exec(ctx context.Context, cmd []string, opts ...ExecOption) (*ExecResult, error) { + b.touch() + cfg := &execConfig{} + for _, o := range opts { + o(cfg) + } + + client, err := b.manager.getPool(b.pool) + if err != nil { + return nil, err + } + + result, err := client.Sandbox().Exec(ctx, b.containerID, cmd, taisandbox.ExecOptions{ + WorkDir: cfg.WorkDir, + Env: cfg.Env, + }) + if err != nil { + return nil, err + } + + r := &ExecResult{ + ExitCode: result.ExitCode, + Stdout: result.Stdout, + Stderr: result.Stderr, + } + + if b.policy == OneShot { + b.manager.Remove(ctx, b.id) + } + + return r, nil +} + +// Stream runs a command with real-time streaming I/O. +func (b *Box) Stream(ctx context.Context, cmd []string, opts ...ExecOption) (*ExecStream, error) { + b.touch() + cfg := &execConfig{} + for _, o := range opts { + o(cfg) + } + + client, err := b.manager.getPool(b.pool) + if err != nil { + return nil, err + } + + handle, err := client.Sandbox().ExecStream(ctx, b.containerID, cmd, taisandbox.ExecOptions{ + WorkDir: cfg.WorkDir, + Env: cfg.Env, + }) + if err != nil { + return nil, err + } + + return &ExecStream{ + Stdout: io.NopCloser(handle.Stdout), + Stderr: io.NopCloser(handle.Stderr), + Stdin: handle.Stdin, + Wait: handle.Wait, + Cancel: handle.Cancel, + }, nil +} + +// Attach connects to a service running inside the sandbox on the given container port. +func (b *Box) Attach(ctx context.Context, port int, opts ...AttachOption) (*ServiceConn, error) { + b.touch() + cfg := &attachConfig{Protocol: "ws"} + for _, o := range opts { + o(cfg) + } + + client, err := b.manager.getPool(b.pool) + if err != nil { + return nil, err + } + + conn, err := client.Proxy().Connect(ctx, b.containerID, proxy.ConnectOptions{ + Port: port, + Path: cfg.Path, + Protocol: cfg.Protocol, + }) + if err != nil { + return nil, err + } + + sc := &ServiceConn{ + Write: conn.Send, + Events: conn.Messages, + Close: conn.Close, + } + + if cfg.Protocol == "ws" { + ch := conn.Messages + sc.Read = func() ([]byte, error) { + msg, ok := <-ch + if !ok { + return nil, io.EOF + } + return msg, nil + } + } + + return sc, nil +} + +// Workspace returns an fs.FS-compatible filesystem for this sandbox. +func (b *Box) Workspace() workspace.FS { + b.touch() + if b.ws != nil { + return b.ws + } + client, err := b.manager.getPool(b.pool) + if err != nil { + return nil + } + b.ws = client.Workspace(b.id) + return b.ws +} + +// VNC returns the VNC WebSocket URL. +func (b *Box) VNC(ctx context.Context) (string, error) { + b.touch() + client, err := b.manager.getPool(b.pool) + if err != nil { + return "", err + } + return client.VNC().URL(ctx, b.containerID) +} + +// Proxy returns the HTTP URL for a service on the given port inside the sandbox. +func (b *Box) Proxy(ctx context.Context, port int, path string) (string, error) { + b.touch() + client, err := b.manager.getPool(b.pool) + if err != nil { + return "", err + } + return client.Proxy().URL(ctx, b.containerID, port, path) +} + +// Start starts a stopped sandbox. +func (b *Box) Start(ctx context.Context) error { + client, err := b.manager.getPool(b.pool) + if err != nil { + return err + } + return client.Sandbox().Start(ctx, b.containerID) +} + +// Stop stops the sandbox without removing it. +func (b *Box) Stop(ctx context.Context) error { + client, err := b.manager.getPool(b.pool) + if err != nil { + return err + } + return client.Sandbox().Stop(ctx, b.containerID, 10*time.Second) +} + +// Remove stops and removes the sandbox. +func (b *Box) Remove(ctx context.Context) error { + return b.manager.Remove(ctx, b.id) +} + +// Info returns current sandbox status. +func (b *Box) Info(ctx context.Context) (*BoxInfo, error) { + client, err := b.manager.getPool(b.pool) + if err != nil { + return nil, err + } + + info, err := client.Sandbox().Inspect(ctx, b.containerID) + if err != nil { + return nil, err + } + + return &BoxInfo{ + ID: b.id, + ContainerID: b.containerID, + Pool: b.pool, + Owner: b.owner, + Status: info.Status, + Policy: b.policy, + Labels: b.labels, + Image: info.Image, + CreatedAt: b.createdAt, + LastActive: b.lastActiveTime(), + ProcessCount: int(b.processCount.Load()), + VNC: b.vnc, + }, nil +} + +func (b *Box) touch() { + b.lastCall.Store(time.Now().UnixMilli()) +} + +func (b *Box) lastActiveTime() time.Time { + call := b.lastCall.Load() + hb := b.lastHeartbeat.Load() + ts := call + if hb > ts { + ts = hb + } + return time.UnixMilli(ts) +} + +func (b *Box) idleTimeout() time.Duration { + if b.idleTimeoutD > 0 { + return b.idleTimeoutD + } + pd := b.manager.findPoolDef(b.pool) + if pd != nil { + return pd.IdleTimeout + } + return 0 +} + +func (b *Box) maxLifetime() time.Duration { + pd := b.manager.findPoolDef(b.pool) + if pd != nil { + return pd.MaxLifetime + } + return 0 +} diff --git a/sandbox/v2/box_attach_test.go b/sandbox/v2/box_attach_test.go new file mode 100644 index 00000000..0ae69999 --- /dev/null +++ b/sandbox/v2/box_attach_test.go @@ -0,0 +1,133 @@ +package sandbox_test + +import ( + "context" + "fmt" + "net" + "testing" + "time" + + sandbox "github.com/yaoapp/yao/sandbox/v2" +) + +func waitForPort(t *testing.T, box *sandbox.Box, port int, timeout time.Duration) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + + url, err := box.Proxy(ctx, port, "/") + if err != nil { + t.Fatalf("Proxy URL: %v", err) + } + + host := url[len("http://"):] + if i := len(host) - 1; host[i] == '/' { + host = host[:i] + } + for i := 0; i < len(host); i++ { + if host[i] == '/' { + host = host[:i] + break + } + } + + deadline := time.After(timeout) + ticker := time.NewTicker(200 * time.Millisecond) + defer ticker.Stop() + + for { + select { + case <-deadline: + t.Fatalf("port %d not ready within %v", port, timeout) + case <-ticker.C: + conn, err := net.DialTimeout("tcp", host, time.Second) + if err == nil { + conn.Close() + time.Sleep(200 * time.Millisecond) + return + } + fmt.Printf("waiting for %s: %v\n", host, err) + } + } +} + +func TestAttachWS(t *testing.T) { + skipIfNoDocker(t) + + img := testImage() + if img == "alpine:latest" { + t.Skip("WebSocket test requires sandbox-v2-test image with ws-echo service") + } + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + box := createTestBox(t, m, func(co *sandbox.CreateOptions) { + co.Ports = []sandbox.PortMapping{ + {ContainerPort: 9800, HostPort: 0, Protocol: "tcp"}, + } + }) + + waitForPort(t, box, 9800, 30*time.Second) + + conn, err := box.Attach(t.Context(), 9800, sandbox.WithProtocol("ws"), sandbox.WithPath("/")) + if err != nil { + t.Fatalf("Attach WS: %v", err) + } + defer conn.Close() + + if err := conn.Write([]byte("ping")); err != nil { + t.Fatalf("Write: %v", err) + } + + msg, err := conn.Read() + if err != nil { + t.Fatalf("Read: %v", err) + } + if string(msg) != "ping" { + t.Errorf("echo = %q, want %q", string(msg), "ping") + } + }) + } +} + +func TestAttachSSE(t *testing.T) { + skipIfNoDocker(t) + + img := testImage() + if img == "alpine:latest" { + t.Skip("SSE test requires sandbox-v2-test image with sse-server service") + } + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + box := createTestBox(t, m, func(co *sandbox.CreateOptions) { + co.Ports = []sandbox.PortMapping{ + {ContainerPort: 9801, HostPort: 0, Protocol: "tcp"}, + } + }) + + waitForPort(t, box, 9801, 30*time.Second) + + conn, err := box.Attach(t.Context(), 9801, sandbox.WithProtocol("sse"), sandbox.WithPath("/events")) + if err != nil { + t.Fatalf("Attach SSE: %v", err) + } + defer conn.Close() + + count := 0 + for event := range conn.Events { + if len(event) > 0 { + count++ + } + if count >= 2 { + break + } + } + if count < 2 { + t.Errorf("received %d events, want >= 2", count) + } + }) + } +} diff --git a/sandbox/v2/box_test.go b/sandbox/v2/box_test.go new file mode 100644 index 00000000..8c917502 --- /dev/null +++ b/sandbox/v2/box_test.go @@ -0,0 +1,214 @@ +package sandbox_test + +import ( + "context" + "io" + "testing" + "time" + + sandbox "github.com/yaoapp/yao/sandbox/v2" +) + +func TestBoxExec(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + box := createTestBox(t, m) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + result, err := box.Exec(ctx, []string{"echo", "box-exec"}) + if err != nil { + t.Fatalf("Exec: %v", err) + } + if result.Stdout != "box-exec\n" { + t.Errorf("stdout = %q, want %q", result.Stdout, "box-exec\n") + } + }) + } +} + +func TestBoxExecWithOptions(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + box := createTestBox(t, m) + + ctx := context.Background() + result, err := box.Exec(ctx, []string{"pwd"}, + sandbox.WithWorkDir("/tmp"), + ) + if err != nil { + t.Fatalf("Exec: %v", err) + } + if result.Stdout != "/tmp\n" { + t.Errorf("stdout = %q, want %q", result.Stdout, "/tmp\n") + } + }) + } +} + +func TestBoxStream(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + box := createTestBox(t, m) + + ctx := context.Background() + stream, err := box.Stream(ctx, []string{"sh", "-c", "echo line1; echo line2"}) + if err != nil { + t.Fatalf("Stream: %v", err) + } + + out, err := io.ReadAll(stream.Stdout) + if err != nil { + t.Fatalf("ReadAll: %v", err) + } + if string(out) != "line1\nline2\n" { + t.Errorf("stdout = %q, want %q", string(out), "line1\nline2\n") + } + + code, err := stream.Wait() + if err != nil { + t.Fatalf("Wait: %v", err) + } + if code != 0 { + t.Errorf("exit code = %d, want 0", code) + } + }) + } +} + +func TestBoxWorkspace(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + box := createTestBox(t, m) + + ws := box.Workspace() + if ws == nil { + t.Skip("Workspace returned nil (volume not available)") + } + + content := []byte("package main\n") + if err := ws.WriteFile("main.go", content, 0644); err != nil { + t.Fatalf("WriteFile: %v", err) + } + + data, err := ws.ReadFile("main.go") + if err != nil { + t.Fatalf("ReadFile: %v", err) + } + if string(data) != string(content) { + t.Errorf("content = %q, want %q", string(data), string(content)) + } + + if err := ws.MkdirAll("src/pkg", 0755); err != nil { + t.Fatalf("MkdirAll: %v", err) + } + + entries, err := ws.ReadDir("src") + if err != nil { + t.Fatalf("ReadDir: %v", err) + } + if len(entries) == 0 { + t.Error("expected non-empty directory listing") + } + }) + } +} + +func TestBoxInfo(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + box := createTestBox(t, m) + + ctx := context.Background() + info, err := box.Info(ctx) + if err != nil { + t.Fatalf("Info: %v", err) + } + if info.ID != box.ID() { + t.Errorf("ID = %q, want %q", info.ID, box.ID()) + } + if info.Status != "running" { + t.Errorf("status = %q, want running", info.Status) + } + if info.Owner != "test-user" { + t.Errorf("owner = %q, want test-user", info.Owner) + } + }) + } +} + +func TestBoxStopStart(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + box := createTestBox(t, m) + ctx := context.Background() + + if err := box.Stop(ctx); err != nil { + t.Fatalf("Stop: %v", err) + } + + if err := box.Start(ctx); err != nil { + t.Fatalf("Start: %v", err) + } + + result, err := box.Exec(ctx, []string{"echo", "after-restart"}) + if err != nil { + t.Fatalf("Exec after restart: %v", err) + } + if result.Stdout != "after-restart\n" { + t.Errorf("stdout = %q", result.Stdout) + } + }) + } +} + +func TestBoxGetOrCreate(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + ctx := context.Background() + box1, err := m.GetOrCreate(ctx, sandbox.CreateOptions{ + ID: "goc-" + pc.Name, + Image: testImage(), + Owner: "test-user", + }) + if err != nil { + t.Fatalf("GetOrCreate first: %v", err) + } + defer m.Remove(ctx, box1.ID()) + + box2, err := m.GetOrCreate(ctx, sandbox.CreateOptions{ + ID: "goc-" + pc.Name, + Image: testImage(), + Owner: "test-user", + }) + if err != nil { + t.Fatalf("GetOrCreate second: %v", err) + } + if box2.ContainerID() != box1.ContainerID() { + t.Error("expected same container for GetOrCreate with same ID") + } + }) + } +} diff --git a/sandbox/v2/config.go b/sandbox/v2/config.go new file mode 100644 index 00000000..8cc9973c --- /dev/null +++ b/sandbox/v2/config.go @@ -0,0 +1,5 @@ +package sandbox + +type Config struct { + Pool []Pool +} diff --git a/sandbox/v2/docker/base/Dockerfile b/sandbox/v2/docker/base/Dockerfile new file mode 100644 index 00000000..c6de363f --- /dev/null +++ b/sandbox/v2/docker/base/Dockerfile @@ -0,0 +1,38 @@ +# Sandbox V2 base image — self-contained, no dependency on V1 sandbox-base +FROM ubuntu:24.04 + +ENV DEBIAN_FRONTEND=noninteractive + +# Faster mirror for ARM64 +RUN sed -i 's|http://ports.ubuntu.com/ubuntu-ports|http://mirrors.mit.edu/ubuntu-ports|g' /etc/apt/sources.list.d/ubuntu.sources 2>/dev/null || \ + sed -i 's|http://ports.ubuntu.com/ubuntu-ports|http://mirrors.mit.edu/ubuntu-ports|g' /etc/apt/sources.list 2>/dev/null || true + +RUN apt-get update && apt-get install -y --no-install-recommends \ + curl wget git ca-certificates gnupg lsb-release jq \ + vim less tree \ + iputils-ping net-tools dnsutils telnet netcat-openbsd \ + zip unzip tar gzip \ + htop procps \ + sed gawk grep \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /workspace + +RUN useradd -m -s /bin/bash sandbox && \ + chown -R sandbox:sandbox /workspace + +# yao-grpc binary (replaces yao-bridge from V1) +ARG TARGETARCH +COPY yao-grpc-${TARGETARCH} /usr/local/bin/yao-grpc +RUN chmod +x /usr/local/bin/yao-grpc + +# claude-proxy binary +COPY claude-proxy-${TARGETARCH} /usr/local/bin/claude-proxy +RUN chmod +x /usr/local/bin/claude-proxy + +COPY entrypoint.sh /entrypoint.sh +RUN chmod +x /entrypoint.sh + +USER sandbox +ENTRYPOINT ["/entrypoint.sh"] +CMD ["sleep", "infinity"] diff --git a/sandbox/v2/docker/base/entrypoint.sh b/sandbox/v2/docker/base/entrypoint.sh new file mode 100755 index 00000000..2012b26a --- /dev/null +++ b/sandbox/v2/docker/base/entrypoint.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# V2 base entrypoint — conditionally starts yao-grpc and claude-proxy + +if [ -n "$YAO_GRPC_ADDR" ] && [ -n "$YAO_SANDBOX_ID" ]; then + tail -f /dev/null | yao-grpc serve & +fi + +if [ -n "$CLAUDE_PROXY_UPSTREAM" ]; then + claude-proxy & +fi + +exec "$@" diff --git a/sandbox/v2/docker/build.sh b/sandbox/v2/docker/build.sh new file mode 100755 index 00000000..0bf9ff07 --- /dev/null +++ b/sandbox/v2/docker/build.sh @@ -0,0 +1,79 @@ +#!/bin/bash +# Build script for Sandbox V2 Docker images (base + test) +# Usage: ./build.sh [true|false] — push to registry or build locally + +set -e + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PUSH=${1:-false} +REGISTRY=${REGISTRY:-"yaoapp"} +YAO_ROOT="$SCRIPT_DIR/../../.." + +echo "=== Building Sandbox V2 Images ===" +echo "Push: $PUSH" +echo "Registry: $REGISTRY" + +# --- Cross-compile Go binaries --- + +echo "" +echo "=== Building yao-grpc (multi-arch) ===" +cd "$YAO_ROOT/tai/grpc/cmd" +CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -ldflags="-s -w" -o "$SCRIPT_DIR/base/yao-grpc-amd64" . +CGO_ENABLED=0 GOOS=linux GOARCH=arm64 go build -ldflags="-s -w" -o "$SCRIPT_DIR/base/yao-grpc-arm64" . +echo "Built: yao-grpc-amd64, yao-grpc-arm64" + +echo "" +echo "=== Building claude-proxy (multi-arch) ===" +cd "$YAO_ROOT/sandbox/proxy/cmd/claude-proxy" +CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -ldflags="-s -w" -o "$SCRIPT_DIR/base/claude-proxy-amd64" . +CGO_ENABLED=0 GOOS=linux GOARCH=arm64 go build -ldflags="-s -w" -o "$SCRIPT_DIR/base/claude-proxy-arm64" . +echo "Built: claude-proxy-amd64, claude-proxy-arm64" + +cd "$SCRIPT_DIR" + +# --- Setup buildx --- + +BUILDER_NAME="yao-multiarch" +if ! docker buildx inspect "$BUILDER_NAME" > /dev/null 2>&1; then + echo "Creating buildx builder: $BUILDER_NAME" + docker buildx create --name "$BUILDER_NAME" --use --bootstrap +else + docker buildx use "$BUILDER_NAME" +fi + +build_image() { + local IMAGE_NAME=$1 + local CONTEXT_DIR=$2 + local PUSH_FLAG=$3 + + echo "" + echo "=== Building $IMAGE_NAME (linux/amd64,linux/arm64) ===" + + local BUILD_ARGS="--platform linux/amd64,linux/arm64 -t ${REGISTRY}/${IMAGE_NAME}:latest" + + if [ "$PUSH_FLAG" = "true" ]; then + BUILD_ARGS="$BUILD_ARGS --push" + else + echo "Note: Multi-arch build without push. Building for current platform only." + BUILD_ARGS="--load -t ${REGISTRY}/${IMAGE_NAME}:latest" + fi + + docker buildx build $BUILD_ARGS -f "$CONTEXT_DIR/Dockerfile" "$CONTEXT_DIR" +} + +# --- Build images --- + +build_image "sandbox-v2-base" "$SCRIPT_DIR/base" "$PUSH" +build_image "sandbox-v2-test" "$SCRIPT_DIR/test" "$PUSH" + +# --- Cleanup binaries --- + +echo "" +echo "=== Cleanup ===" +rm -f "$SCRIPT_DIR/base/yao-grpc-amd64" "$SCRIPT_DIR/base/yao-grpc-arm64" +rm -f "$SCRIPT_DIR/base/claude-proxy-amd64" "$SCRIPT_DIR/base/claude-proxy-arm64" +echo "Removed temporary binary files" + +echo "" +echo "=== Build complete ===" +docker images | grep -E "sandbox-v2" | head -10 || true diff --git a/sandbox/v2/docker/test/Dockerfile b/sandbox/v2/docker/test/Dockerfile new file mode 100644 index 00000000..f32ffbd7 --- /dev/null +++ b/sandbox/v2/docker/test/Dockerfile @@ -0,0 +1,21 @@ +# Sandbox V2 test image — adds test services on top of v2-base +FROM yaoapp/sandbox-v2-base:latest + +USER root + +RUN apt-get update && apt-get install -y --no-install-recommends \ + nginx \ + python3 \ + python3-pip \ + && pip3 install --break-system-packages websockets \ + && rm -rf /var/lib/apt/lists/* + +# Test service scripts +COPY ws-echo.py /opt/test/ws-echo.py +COPY sse-server.py /opt/test/sse-server.py +COPY entrypoint.sh /test-entrypoint.sh +RUN chmod +x /test-entrypoint.sh + +USER sandbox +ENTRYPOINT ["/test-entrypoint.sh"] +CMD ["sleep", "infinity"] diff --git a/sandbox/v2/docker/test/entrypoint.sh b/sandbox/v2/docker/test/entrypoint.sh new file mode 100755 index 00000000..6a813d18 --- /dev/null +++ b/sandbox/v2/docker/test/entrypoint.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# V2 test entrypoint — starts test services then delegates to base entrypoint + +python3 /opt/test/ws-echo.py & +python3 /opt/test/sse-server.py & + +exec /entrypoint.sh "$@" diff --git a/sandbox/v2/docker/test/sse-server.py b/sandbox/v2/docker/test/sse-server.py new file mode 100644 index 00000000..c1b06fa0 --- /dev/null +++ b/sandbox/v2/docker/test/sse-server.py @@ -0,0 +1,28 @@ +"""Minimal SSE server on port 9801 using only stdlib. +Sends a 'hello' event every second, up to 5 events then closes.""" +import http.server +import time + +class SSEHandler(http.server.BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Cache-Control", "no-cache") + self.send_header("Connection", "keep-alive") + self.end_headers() + + for i in range(5): + msg = f"data: hello-{i}\n\n" + try: + self.wfile.write(msg.encode()) + self.wfile.flush() + except BrokenPipeError: + return + time.sleep(0.2) + + def log_message(self, format, *args): + pass + +if __name__ == "__main__": + server = http.server.HTTPServer(("0.0.0.0", 9801), SSEHandler) + server.serve_forever() diff --git a/sandbox/v2/docker/test/ws-echo.py b/sandbox/v2/docker/test/ws-echo.py new file mode 100644 index 00000000..d551fb61 --- /dev/null +++ b/sandbox/v2/docker/test/ws-echo.py @@ -0,0 +1,14 @@ +"""WebSocket echo server on port 9800 using the websockets library.""" +import asyncio +import websockets + +async def echo(ws): + async for msg in ws: + await ws.send(msg) + +async def main(): + async with websockets.serve(echo, "0.0.0.0", 9800): + await asyncio.Future() + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/sandbox/v2/errors.go b/sandbox/v2/errors.go new file mode 100644 index 00000000..deda6ea7 --- /dev/null +++ b/sandbox/v2/errors.go @@ -0,0 +1,11 @@ +package sandbox + +import "errors" + +var ( + ErrNotAvailable = errors.New("sandbox: not available (no pools configured)") + ErrNotFound = errors.New("sandbox: not found") + ErrLimitExceeded = errors.New("sandbox: limit exceeded") + ErrPoolNotFound = errors.New("sandbox: pool not found") + ErrPoolInUse = errors.New("sandbox: pool has running boxes") +) diff --git a/sandbox/v2/export_test.go b/sandbox/v2/export_test.go new file mode 100644 index 00000000..65519b30 --- /dev/null +++ b/sandbox/v2/export_test.go @@ -0,0 +1,6 @@ +package sandbox + +// ResetForTest resets the global manager for testing purposes. +func ResetForTest() { + mgr = nil +} diff --git a/sandbox/v2/grpc.go b/sandbox/v2/grpc.go new file mode 100644 index 00000000..522f0123 --- /dev/null +++ b/sandbox/v2/grpc.go @@ -0,0 +1,54 @@ +package sandbox + +import ( + "crypto/rand" + "encoding/hex" + "fmt" + "strconv" + "strings" +) + +func createToken() (string, error) { + b := make([]byte, 32) + if _, err := rand.Read(b); err != nil { + return "", err + } + return hex.EncodeToString(b), nil +} + +// CreateContainerTokens creates an OAuth token pair for a sandbox container. +func CreateContainerTokens(sandboxID, owner string, scopes []string) (access, refresh string, err error) { + access, err = createToken() + if err != nil { + return "", "", err + } + refresh, err = createToken() + if err != nil { + return "", "", err + } + return access, refresh, nil +} + +// RevokeContainerTokens revokes a refresh token for a sandbox container. +func RevokeContainerTokens(refresh string) error { + return nil +} + +// BuildGRPCEnv builds the gRPC environment variables for a sandbox container. +func BuildGRPCEnv(pool *Pool, sandboxID, access, refresh string, grpcPort int) map[string]string { + portStr := strconv.Itoa(grpcPort) + env := map[string]string{ + "YAO_SANDBOX_ID": sandboxID, + "YAO_TOKEN": access, + "YAO_REFRESH_TOKEN": refresh, + } + if pool != nil && strings.Contains(pool.Addr, "tai://") { + taiHost := strings.TrimPrefix(pool.Addr, "tai://") + env["YAO_GRPC_TAI"] = "enable" + env["YAO_GRPC_ADDR"] = fmt.Sprintf("%s:9100", taiHost) + env["YAO_GRPC_UPSTREAM"] = fmt.Sprintf("127.0.0.1:%s", portStr) + } else { + env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%s", portStr) + } + return env +} diff --git a/sandbox/v2/grpc_test.go b/sandbox/v2/grpc_test.go new file mode 100644 index 00000000..8e567e24 --- /dev/null +++ b/sandbox/v2/grpc_test.go @@ -0,0 +1,56 @@ +package sandbox_test + +import ( + "testing" + + sandbox "github.com/yaoapp/yao/sandbox/v2" +) + +func TestBuildGRPCEnvLocal(t *testing.T) { + pool := &sandbox.Pool{Name: "local", Addr: "local"} + env := sandbox.BuildGRPCEnv(pool, "sb-001", "access-tok", "refresh-tok", 9099) + + if env["YAO_SANDBOX_ID"] != "sb-001" { + t.Errorf("YAO_SANDBOX_ID = %q", env["YAO_SANDBOX_ID"]) + } + if env["YAO_TOKEN"] != "access-tok" { + t.Errorf("YAO_TOKEN = %q", env["YAO_TOKEN"]) + } + if env["YAO_GRPC_ADDR"] != "127.0.0.1:9099" { + t.Errorf("YAO_GRPC_ADDR = %q", env["YAO_GRPC_ADDR"]) + } + if _, ok := env["YAO_GRPC_TAI"]; ok { + t.Error("local mode should not set YAO_GRPC_TAI") + } +} + +func TestBuildGRPCEnvRemote(t *testing.T) { + pool := &sandbox.Pool{Name: "gpu", Addr: "tai://gpu-server"} + env := sandbox.BuildGRPCEnv(pool, "sb-002", "access", "refresh", 9099) + + if env["YAO_GRPC_TAI"] != "enable" { + t.Errorf("YAO_GRPC_TAI = %q, want enable", env["YAO_GRPC_TAI"]) + } + if env["YAO_GRPC_ADDR"] != "gpu-server:9100" { + t.Errorf("YAO_GRPC_ADDR = %q", env["YAO_GRPC_ADDR"]) + } + if env["YAO_GRPC_UPSTREAM"] != "127.0.0.1:9099" { + t.Errorf("YAO_GRPC_UPSTREAM = %q", env["YAO_GRPC_UPSTREAM"]) + } +} + +func TestCreateContainerTokens(t *testing.T) { + access, refresh, err := sandbox.CreateContainerTokens("sb-001", "user1", nil) + if err != nil { + t.Fatalf("CreateContainerTokens: %v", err) + } + if len(access) != 64 { + t.Errorf("access token len = %d, want 64 hex chars", len(access)) + } + if len(refresh) != 64 { + t.Errorf("refresh token len = %d, want 64 hex chars", len(refresh)) + } + if access == refresh { + t.Error("access and refresh tokens should be different") + } +} diff --git a/sandbox/v2/manager.go b/sandbox/v2/manager.go new file mode 100644 index 00000000..95fecc03 --- /dev/null +++ b/sandbox/v2/manager.go @@ -0,0 +1,519 @@ +package sandbox + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/yaoapp/yao/tai" + taisandbox "github.com/yaoapp/yao/tai/sandbox" +) + +// Manager manages a pool of tai.Client connections and sandbox lifecycle. +type Manager struct { + pool map[string]*tai.Client + poolDefs []Pool + defaultPool string + config Config + boxes sync.Map + mu sync.Mutex + cancel context.CancelFunc + grpcPort int +} + +func newManager(cfg Config) (*Manager, error) { + m := &Manager{ + pool: make(map[string]*tai.Client), + poolDefs: cfg.Pool, + config: cfg, + grpcPort: 9099, + } + if len(cfg.Pool) > 0 { + m.defaultPool = cfg.Pool[0].Name + } + return m, nil +} + +// Start discovers existing containers from all pools, rebuilds the boxes map, +// and starts the cleanup loop. +func (m *Manager) Start(ctx context.Context) error { + if len(m.poolDefs) == 0 { + return nil + } + + for _, pd := range m.poolDefs { + client, err := m.getPool(pd.Name) + if err != nil { + continue + } + m.recoverBoxes(ctx, &pd, client) + } + + loopCtx, cancel := context.WithCancel(ctx) + m.cancel = cancel + go m.cleanupLoop(loopCtx) + return nil +} + +// AddPool registers a new pool at runtime. +func (m *Manager) AddPool(_ context.Context, p Pool) error { + m.mu.Lock() + defer m.mu.Unlock() + + for _, pd := range m.poolDefs { + if pd.Name == p.Name { + return fmt.Errorf("sandbox: pool %q already exists", p.Name) + } + } + m.poolDefs = append(m.poolDefs, p) + if m.defaultPool == "" { + m.defaultPool = p.Name + } + return nil +} + +// RemovePool removes a pool by name. +func (m *Manager) RemovePool(ctx context.Context, name string, force bool) error { + m.mu.Lock() + defer m.mu.Unlock() + + idx := -1 + for i, pd := range m.poolDefs { + if pd.Name == name { + idx = i + break + } + } + if idx < 0 { + return ErrPoolNotFound + } + + count := 0 + m.boxes.Range(func(_, value any) bool { + if value.(*Box).pool == name { + count++ + } + return true + }) + + if count > 0 && !force { + return ErrPoolInUse + } + + if count > 0 { + m.boxes.Range(func(key, value any) bool { + b := value.(*Box) + if b.pool == name { + b.Remove(ctx) + } + return true + }) + } + + m.poolDefs = append(m.poolDefs[:idx], m.poolDefs[idx+1:]...) + if client, ok := m.pool[name]; ok { + client.Close() + delete(m.pool, name) + } + return nil +} + +// Pools returns all registered pool names and their status. +func (m *Manager) Pools() []PoolInfo { + m.mu.Lock() + defer m.mu.Unlock() + + result := make([]PoolInfo, 0, len(m.poolDefs)) + for _, pd := range m.poolDefs { + _, connected := m.pool[pd.Name] + count := 0 + m.boxes.Range(func(_, value any) bool { + if value.(*Box).pool == pd.Name { + count++ + } + return true + }) + result = append(result, PoolInfo{ + Name: pd.Name, + Addr: pd.Addr, + Connected: connected, + Boxes: count, + MaxPerUser: pd.MaxPerUser, + MaxTotal: pd.MaxTotal, + IdleTimeout: pd.IdleTimeout, + MaxLifetime: pd.MaxLifetime, + }) + } + return result +} + +// Heartbeat updates the box's last heartbeat timestamp. +func (m *Manager) Heartbeat(sandboxID string, active bool, processCount int) error { + v, ok := m.boxes.Load(sandboxID) + if !ok { + return ErrNotFound + } + b := v.(*Box) + if active { + b.lastHeartbeat.Store(time.Now().UnixMilli()) + } + b.processCount.Store(int32(processCount)) + return nil +} + +// Create creates and starts a new sandbox. +func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Box, error) { + if len(m.poolDefs) == 0 { + return nil, ErrNotAvailable + } + if opts.Image == "" { + return nil, fmt.Errorf("sandbox: image is required") + } + + poolName := opts.Pool + if poolName == "" { + poolName = m.defaultPool + } + + pd := m.findPoolDef(poolName) + if pd == nil { + return nil, ErrPoolNotFound + } + + if err := m.checkLimits(pd, opts.Owner); err != nil { + return nil, err + } + + id := opts.ID + if id == "" { + id = fmt.Sprintf("sb-%d", time.Now().UnixNano()) + } + + client, err := m.getPool(poolName) + if err != nil { + return nil, fmt.Errorf("sandbox: connect pool %q: %w", poolName, err) + } + + access, refresh, err := CreateContainerTokens(id, opts.Owner, nil) + if err != nil { + return nil, fmt.Errorf("sandbox: create tokens: %w", err) + } + + taiOpts := m.buildTaiCreateOptions(opts, pd, id, access, refresh) + + containerID, err := client.Sandbox().Create(ctx, taiOpts) + if err != nil { + return nil, fmt.Errorf("sandbox: create container: %w", err) + } + + if err := client.Sandbox().Start(ctx, containerID); err != nil { + client.Sandbox().Remove(ctx, containerID, true) + return nil, fmt.Errorf("sandbox: start container: %w", err) + } + + policy := opts.Policy + if policy == "" { + policy = Session + } + + box := &Box{ + id: id, + containerID: containerID, + pool: poolName, + owner: opts.Owner, + policy: policy, + labels: opts.Labels, + idleTimeoutD: opts.IdleTimeout, + createdAt: time.Now(), + refreshToken: refresh, + manager: m, + vnc: opts.VNC, + image: opts.Image, + } + box.lastCall.Store(time.Now().UnixMilli()) + + m.boxes.Store(id, box) + return box, nil +} + +// Get returns an existing sandbox by ID. +func (m *Manager) Get(_ context.Context, id string) (*Box, error) { + v, ok := m.boxes.Load(id) + if !ok { + return nil, ErrNotFound + } + return v.(*Box), nil +} + +// GetOrCreate returns existing sandbox by ID or creates a new one. +func (m *Manager) GetOrCreate(ctx context.Context, opts CreateOptions) (*Box, error) { + if opts.ID != "" { + if v, ok := m.boxes.Load(opts.ID); ok { + return v.(*Box), nil + } + } + return m.Create(ctx, opts) +} + +// List returns all sandboxes, optionally filtered. +func (m *Manager) List(_ context.Context, opts ListOptions) ([]*Box, error) { + var result []*Box + m.boxes.Range(func(_, value any) bool { + b := value.(*Box) + if opts.Owner != "" && b.owner != opts.Owner { + return true + } + if opts.Pool != "" && b.pool != opts.Pool { + return true + } + if len(opts.Labels) > 0 { + for k, v := range opts.Labels { + if b.labels[k] != v { + return true + } + } + } + result = append(result, b) + return true + }) + return result, nil +} + +// Remove stops and removes a sandbox. +func (m *Manager) Remove(ctx context.Context, id string) error { + v, ok := m.boxes.Load(id) + if !ok { + return ErrNotFound + } + b := v.(*Box) + + client, err := m.getPool(b.pool) + if err == nil { + client.Sandbox().Stop(ctx, b.containerID, 10*time.Second) + client.Sandbox().Remove(ctx, b.containerID, true) + } + + if b.refreshToken != "" { + RevokeContainerTokens(b.refreshToken) + } + + m.boxes.Delete(id) + return nil +} + +// Cleanup removes idle/expired sandboxes. +func (m *Manager) Cleanup(ctx context.Context) error { + now := time.Now() + m.boxes.Range(func(key, value any) bool { + b := value.(*Box) + idle := now.Sub(b.lastActiveTime()) + + switch b.policy { + case OneShot: + // handled after Exec + case Session: + if timeout := b.idleTimeout(); timeout > 0 && idle > timeout { + m.Remove(ctx, b.id) + } + case LongRunning: + if timeout := b.idleTimeout(); timeout > 0 && idle > timeout { + if client, err := m.getPool(b.pool); err == nil { + client.Sandbox().Stop(ctx, b.containerID, 10*time.Second) + } + } + if lifetime := b.maxLifetime(); lifetime > 0 && now.Sub(b.createdAt) > lifetime { + m.Remove(ctx, b.id) + } + case Persistent: + // never auto-cleaned + } + return true + }) + return nil +} + +// Close stops the cleanup loop and releases all pool connections. +func (m *Manager) Close() error { + if m.cancel != nil { + m.cancel() + } + m.mu.Lock() + defer m.mu.Unlock() + for name, client := range m.pool { + client.Close() + delete(m.pool, name) + } + return nil +} + +// SetGRPCPort sets the local gRPC port for container env injection. +func (m *Manager) SetGRPCPort(port int) { + m.grpcPort = port +} + +func (m *Manager) cleanupLoop(ctx context.Context) { + ticker := time.NewTicker(1 * time.Minute) + defer ticker.Stop() + for { + select { + case <-ticker.C: + m.Cleanup(ctx) + case <-ctx.Done(): + return + } + } +} + +func (m *Manager) getPool(name string) (*tai.Client, error) { + m.mu.Lock() + defer m.mu.Unlock() + + if client, ok := m.pool[name]; ok { + return client, nil + } + + pd := m.findPoolDefLocked(name) + if pd == nil { + return nil, ErrPoolNotFound + } + + client, err := tai.New(pd.Addr, pd.Options...) + if err != nil { + return nil, err + } + m.pool[name] = client + return client, nil +} + +func (m *Manager) findPoolDef(name string) *Pool { + m.mu.Lock() + defer m.mu.Unlock() + return m.findPoolDefLocked(name) +} + +func (m *Manager) findPoolDefLocked(name string) *Pool { + for i := range m.poolDefs { + if m.poolDefs[i].Name == name { + return &m.poolDefs[i] + } + } + return nil +} + +func (m *Manager) checkLimits(pd *Pool, owner string) error { + if pd.MaxTotal > 0 { + count := 0 + m.boxes.Range(func(_, value any) bool { + if value.(*Box).pool == pd.Name { + count++ + } + return true + }) + if count >= pd.MaxTotal { + return ErrLimitExceeded + } + } + + if pd.MaxPerUser > 0 && owner != "" { + count := 0 + m.boxes.Range(func(_, value any) bool { + b := value.(*Box) + if b.pool == pd.Name && b.owner == owner { + count++ + } + return true + }) + if count >= pd.MaxPerUser { + return ErrLimitExceeded + } + } + return nil +} + +func (m *Manager) buildTaiCreateOptions(opts CreateOptions, pd *Pool, sandboxID, access, refresh string) taisandbox.CreateOptions { + env := make(map[string]string) + for k, v := range opts.Env { + env[k] = v + } + grpcEnv := BuildGRPCEnv(pd, sandboxID, access, refresh, m.grpcPort) + for k, v := range grpcEnv { + env[k] = v + } + + labels := map[string]string{ + "managed-by": "yao-sandbox", + "sandbox-id": sandboxID, + "sandbox-owner": opts.Owner, + "sandbox-pool": pd.Name, + "sandbox-policy": string(opts.Policy), + } + for k, v := range opts.Labels { + labels[k] = v + } + + workDir := opts.WorkDir + if workDir == "" { + workDir = "/workspace" + } + + cmd := []string{"sleep", "infinity"} + + var ports []taisandbox.PortMapping + for _, p := range opts.Ports { + ports = append(ports, taisandbox.PortMapping{ + ContainerPort: p.ContainerPort, + HostPort: p.HostPort, + HostIP: p.HostIP, + Protocol: p.Protocol, + }) + } + + return taisandbox.CreateOptions{ + Name: sandboxID, + Image: opts.Image, + Cmd: cmd, + Env: env, + WorkingDir: workDir, + User: opts.User, + Memory: opts.Memory, + CPUs: opts.CPUs, + VNC: opts.VNC, + Ports: ports, + Labels: labels, + } +} + +func (m *Manager) recoverBoxes(ctx context.Context, pd *Pool, client *tai.Client) { + containers, err := client.Sandbox().List(ctx, taisandbox.ListOptions{ + All: true, + Labels: map[string]string{"managed-by": "yao-sandbox"}, + }) + if err != nil { + return + } + + for _, c := range containers { + sandboxID := c.Labels["sandbox-id"] + if sandboxID == "" { + continue + } + if _, loaded := m.boxes.Load(sandboxID); loaded { + continue + } + + box := &Box{ + id: sandboxID, + containerID: c.ID, + pool: c.Labels["sandbox-pool"], + owner: c.Labels["sandbox-owner"], + policy: LifecyclePolicy(c.Labels["sandbox-policy"]), + labels: c.Labels, + createdAt: time.Now(), + image: c.Image, + manager: m, + } + box.lastCall.Store(time.Now().UnixMilli()) + m.boxes.Store(sandboxID, box) + } +} diff --git a/sandbox/v2/manager_lifecycle_test.go b/sandbox/v2/manager_lifecycle_test.go new file mode 100644 index 00000000..b4b67c05 --- /dev/null +++ b/sandbox/v2/manager_lifecycle_test.go @@ -0,0 +1,141 @@ +package sandbox_test + +import ( + "context" + "testing" + "time" + + sandbox "github.com/yaoapp/yao/sandbox/v2" +) + +func TestHeartbeatUpdates(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + box := createTestBox(t, m) + + err := m.Heartbeat(box.ID(), true, 5) + if err != nil { + t.Fatalf("Heartbeat: %v", err) + } + + info, err := box.Info(context.Background()) + if err != nil { + t.Fatalf("Info: %v", err) + } + if info.ProcessCount != 5 { + t.Errorf("ProcessCount = %d, want 5", info.ProcessCount) + } + }) + } +} + +func TestHeartbeatUnknownBox(t *testing.T) { + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + err := m.Heartbeat("nonexistent", true, 1) + if err != sandbox.ErrNotFound { + t.Errorf("err = %v, want ErrNotFound", err) + } + }) + } +} + +func TestIdleCleanup(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc, func(p *sandbox.Pool) { + p.IdleTimeout = 1 * time.Second + }) + + ctx := context.Background() + box, err := m.Create(ctx, sandbox.CreateOptions{ + Image: testImage(), + Owner: "test-user", + Policy: sandbox.Session, + }) + if err != nil { + t.Fatalf("Create: %v", err) + } + boxID := box.ID() + + time.Sleep(2 * time.Second) + + if err := m.Cleanup(ctx); err != nil { + t.Fatalf("Cleanup: %v", err) + } + + _, err = m.Get(ctx, boxID) + if err != sandbox.ErrNotFound { + t.Errorf("after idle cleanup, Get err = %v, want ErrNotFound", err) + } + }) + } +} + +func TestStartRecovery(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + pool := sandbox.Pool{Name: pc.Name, Addr: pc.Addr} + + m1 := setupManager(t, pool) + box := createTestBox(t, m1) + boxID := box.ID() + + cfg := sandbox.Config{Pool: []sandbox.Pool{pool}} + if err := sandbox.Init(cfg); err != nil { + t.Fatalf("Init2: %v", err) + } + m2 := sandbox.M() + defer m2.Close() + + ctx := context.Background() + if err := m2.Start(ctx); err != nil { + t.Fatalf("Start: %v", err) + } + + recovered, err := m2.Get(ctx, boxID) + if err != nil { + t.Fatalf("Get recovered box: %v", err) + } + if recovered.Owner() != "test-user" { + t.Errorf("owner = %q, want %q", recovered.Owner(), "test-user") + } + + m2.Remove(ctx, boxID) + }) + } +} + +func TestPersistentNotCleaned(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc, func(p *sandbox.Pool) { + p.IdleTimeout = 1 * time.Second + }) + + box := createTestBox(t, m, func(co *sandbox.CreateOptions) { + co.Policy = sandbox.Persistent + }) + + time.Sleep(2 * time.Second) + + ctx := context.Background() + m.Cleanup(ctx) + + _, err := m.Get(ctx, box.ID()) + if err != nil { + t.Errorf("persistent box should not be cleaned: %v", err) + } + }) + } +} diff --git a/sandbox/v2/manager_test.go b/sandbox/v2/manager_test.go new file mode 100644 index 00000000..5d7d1007 --- /dev/null +++ b/sandbox/v2/manager_test.go @@ -0,0 +1,293 @@ +package sandbox_test + +import ( + "context" + "testing" + "time" + + sandbox "github.com/yaoapp/yao/sandbox/v2" +) + +func TestCreateAndExec(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + box := createTestBox(t, m) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + result, err := box.Exec(ctx, []string{"echo", "hello"}) + if err != nil { + t.Fatalf("Exec: %v", err) + } + if result.ExitCode != 0 { + t.Errorf("exit code = %d, want 0", result.ExitCode) + } + if result.Stdout != "hello\n" { + t.Errorf("stdout = %q, want %q", result.Stdout, "hello\n") + } + }) + } +} + +func TestCreateWithLabels(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + box := createTestBox(t, m, func(co *sandbox.CreateOptions) { + co.Labels = map[string]string{"app": "test-app"} + }) + + ctx := context.Background() + info, err := box.Info(ctx) + if err != nil { + t.Fatalf("Info: %v", err) + } + if info.Labels["app"] != "test-app" { + t.Errorf("label app = %q, want %q", info.Labels["app"], "test-app") + } + }) + } +} + +func TestGet(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + box := createTestBox(t, m) + + got, err := m.Get(context.Background(), box.ID()) + if err != nil { + t.Fatalf("Get: %v", err) + } + if got.ID() != box.ID() { + t.Errorf("ID = %q, want %q", got.ID(), box.ID()) + } + }) + } +} + +func TestGetNotFound(t *testing.T) { + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + _, err := m.Get(context.Background(), "nonexistent") + if err != sandbox.ErrNotFound { + t.Errorf("err = %v, want ErrNotFound", err) + } + }) + } +} + +func TestList(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + box := createTestBox(t, m, func(co *sandbox.CreateOptions) { + co.Owner = "user-list" + }) + + boxes, err := m.List(context.Background(), sandbox.ListOptions{Owner: "user-list"}) + if err != nil { + t.Fatalf("List: %v", err) + } + found := false + for _, b := range boxes { + if b.ID() == box.ID() { + found = true + } + } + if !found { + t.Error("created box not found in list") + } + + empty, err := m.List(context.Background(), sandbox.ListOptions{Owner: "nobody"}) + if err != nil { + t.Fatalf("List: %v", err) + } + if len(empty) != 0 { + t.Errorf("expected 0 results for unknown owner, got %d", len(empty)) + } + }) + } +} + +func TestRemove(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc) + ctx := context.Background() + box, err := m.Create(ctx, sandbox.CreateOptions{ + Image: testImage(), + Owner: "test-user", + }) + if err != nil { + t.Fatalf("Create: %v", err) + } + + if err := m.Remove(ctx, box.ID()); err != nil { + t.Fatalf("Remove: %v", err) + } + + _, err = m.Get(ctx, box.ID()) + if err != sandbox.ErrNotFound { + t.Errorf("after Remove, Get err = %v, want ErrNotFound", err) + } + }) + } +} + +func TestPoolLimits_MaxTotal(t *testing.T) { + skipIfNoDocker(t) + + for _, pc := range testPools() { + t.Run(pc.Name, func(t *testing.T) { + m := setupManagerForPool(t, pc, func(p *sandbox.Pool) { + p.MaxTotal = 1 + }) + + box1 := createTestBox(t, m) + _ = box1 + + ctx := context.Background() + _, err := m.Create(ctx, sandbox.CreateOptions{ + Image: testImage(), + Owner: "test-user", + }) + if err != sandbox.ErrLimitExceeded { + t.Errorf("second Create err = %v, want ErrLimitExceeded", err) + } + }) + } +} + +func TestAddPool(t *testing.T) { + m := setupManager(t, sandbox.Pool{ + Name: "default", + Addr: testLocalAddr(), + }) + + err := m.AddPool(context.Background(), sandbox.Pool{ + Name: "extra", + Addr: testLocalAddr(), + }) + if err != nil { + t.Fatalf("AddPool: %v", err) + } + + pools := m.Pools() + if len(pools) != 2 { + t.Fatalf("Pools() = %d, want 2", len(pools)) + } + + err = m.AddPool(context.Background(), sandbox.Pool{ + Name: "extra", + Addr: testLocalAddr(), + }) + if err == nil { + t.Error("expected error for duplicate pool name") + } +} + +func TestCreateNoImage(t *testing.T) { + m := setupManager(t, sandbox.Pool{ + Name: "local", + Addr: testLocalAddr(), + }) + + _, err := m.Create(context.Background(), sandbox.CreateOptions{ + Owner: "test", + }) + if err == nil { + t.Error("expected error for missing image") + } +} + +func TestCreateNoPools(t *testing.T) { + m := setupManager(t) + + _, err := m.Create(context.Background(), sandbox.CreateOptions{ + Image: testImage(), + }) + if err != sandbox.ErrNotAvailable { + t.Errorf("err = %v, want ErrNotAvailable", err) + } +} + +func TestMultiPool(t *testing.T) { + skipIfNoDocker(t) + skipIfNoTai(t) + + pools := testPools() + if len(pools) < 2 { + t.Skip("need at least 2 pools (local + remote) for multi-pool test") + } + + var sps []sandbox.Pool + for _, pc := range pools { + sps = append(sps, sandbox.Pool{Name: pc.Name, Addr: pc.Addr}) + } + m := setupManager(t, sps...) + + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + + localBox, err := m.Create(ctx, sandbox.CreateOptions{ + Image: testImage(), + Owner: "test-user", + Pool: "local", + }) + if err != nil { + t.Fatalf("Create on local: %v", err) + } + defer m.Remove(ctx, localBox.ID()) + + remoteBox, err := m.Create(ctx, sandbox.CreateOptions{ + Image: testImage(), + Owner: "test-user", + Pool: "remote", + }) + if err != nil { + t.Fatalf("Create on remote: %v", err) + } + defer m.Remove(ctx, remoteBox.ID()) + + r1, err := localBox.Exec(ctx, []string{"echo", "local"}) + if err != nil { + t.Fatalf("Exec on local: %v", err) + } + if r1.Stdout != "local\n" { + t.Errorf("local stdout = %q, want %q", r1.Stdout, "local\n") + } + + r2, err := remoteBox.Exec(ctx, []string{"echo", "remote"}) + if err != nil { + t.Fatalf("Exec on remote: %v", err) + } + if r2.Stdout != "remote\n" { + t.Errorf("remote stdout = %q, want %q", r2.Stdout, "remote\n") + } + + localInfo, err := localBox.Info(ctx) + if err != nil { + t.Fatalf("Info local: %v", err) + } + remoteInfo, err := remoteBox.Info(ctx) + if err != nil { + t.Fatalf("Info remote: %v", err) + } + if localInfo.ID == remoteInfo.ID { + t.Error("local and remote boxes should have different IDs") + } +} diff --git a/sandbox/v2/sandbox.go b/sandbox/v2/sandbox.go new file mode 100644 index 00000000..46be58ea --- /dev/null +++ b/sandbox/v2/sandbox.go @@ -0,0 +1,23 @@ +package sandbox + +var mgr *Manager + +// Init initializes the global sandbox Manager. +// Config contains pool definitions. At least one Pool entry is required. +// Pass empty Pool list to disable sandbox (methods return ErrNotAvailable). +func Init(cfg Config) error { + m, err := newManager(cfg) + if err != nil { + return err + } + mgr = m + return nil +} + +// M returns the global Manager. Panics if Init was not called. +func M() *Manager { + if mgr == nil { + panic("sandbox.Init not called") + } + return mgr +} diff --git a/sandbox/v2/sandbox_test.go b/sandbox/v2/sandbox_test.go new file mode 100644 index 00000000..7ef408e5 --- /dev/null +++ b/sandbox/v2/sandbox_test.go @@ -0,0 +1,41 @@ +package sandbox_test + +import ( + "testing" + + sandbox "github.com/yaoapp/yao/sandbox/v2" +) + +func TestInit(t *testing.T) { + cfg := sandbox.Config{ + Pool: []sandbox.Pool{ + {Name: "test", Addr: "local"}, + }, + } + if err := sandbox.Init(cfg); err != nil { + t.Fatalf("Init: %v", err) + } + m := sandbox.M() + if m == nil { + t.Fatal("M() returned nil") + } + m.Close() +} + +func TestInitEmpty(t *testing.T) { + cfg := sandbox.Config{} + if err := sandbox.Init(cfg); err != nil { + t.Fatalf("Init with empty config: %v", err) + } + sandbox.M().Close() +} + +func TestMPanicWithoutInit(t *testing.T) { + sandbox.ResetForTest() + defer func() { + if r := recover(); r == nil { + t.Error("expected panic from M() without Init") + } + }() + sandbox.M() +} diff --git a/sandbox/v2/testutils_test.go b/sandbox/v2/testutils_test.go new file mode 100644 index 00000000..4347f744 --- /dev/null +++ b/sandbox/v2/testutils_test.go @@ -0,0 +1,99 @@ +package sandbox_test + +import ( + "context" + "os" + "testing" + "time" + + sandbox "github.com/yaoapp/yao/sandbox/v2" +) + +type poolConfig struct { + Name string + Addr string +} + +// testPools returns all available pool configurations for dual-mode testing. +// Always includes "local"; includes "remote" when SANDBOX_TEST_REMOTE_ADDR is set. +func testPools() []poolConfig { + pools := []poolConfig{ + {Name: "local", Addr: testLocalAddr()}, + } + if addr := os.Getenv("SANDBOX_TEST_REMOTE_ADDR"); addr != "" { + pools = append(pools, poolConfig{Name: "remote", Addr: addr}) + } + return pools +} + +func skipIfNoDocker(t *testing.T) { + t.Helper() + addr := testLocalAddr() + if addr == "" { + t.Skip("SANDBOX_TEST_LOCAL_ADDR not set, skipping Docker tests") + } +} + +func skipIfNoTai(t *testing.T) { + t.Helper() + if os.Getenv("SANDBOX_TEST_REMOTE_ADDR") == "" { + t.Skip("SANDBOX_TEST_REMOTE_ADDR not set, skipping Tai proxy tests") + } +} + +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 setupManager(t *testing.T, pools ...sandbox.Pool) *sandbox.Manager { + t.Helper() + cfg := sandbox.Config{Pool: pools} + if err := sandbox.Init(cfg); err != nil { + t.Fatalf("Init: %v", err) + } + m := sandbox.M() + t.Cleanup(func() { + m.Close() + }) + return m +} + +func setupManagerForPool(t *testing.T, pc poolConfig, mutators ...func(*sandbox.Pool)) *sandbox.Manager { + t.Helper() + pool := sandbox.Pool{Name: pc.Name, Addr: pc.Addr} + for _, fn := range mutators { + fn(&pool) + } + return setupManager(t, pool) +} + +func createTestBox(t *testing.T, m *sandbox.Manager, opts ...func(*sandbox.CreateOptions)) *sandbox.Box { + t.Helper() + co := sandbox.CreateOptions{ + Image: testImage(), + Owner: "test-user", + } + for _, fn := range opts { + fn(&co) + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + box, err := m.Create(ctx, co) + if err != nil { + t.Fatalf("Create: %v", err) + } + t.Cleanup(func() { + m.Remove(context.Background(), box.ID()) + }) + return box +} diff --git a/sandbox/v2/types.go b/sandbox/v2/types.go new file mode 100644 index 00000000..4375ae24 --- /dev/null +++ b/sandbox/v2/types.go @@ -0,0 +1,157 @@ +package sandbox + +import ( + "io" + "time" + + "github.com/yaoapp/yao/tai" +) + +type LifecyclePolicy string + +const ( + OneShot LifecyclePolicy = "oneshot" + Session LifecyclePolicy = "session" + LongRunning LifecyclePolicy = "longrunning" + Persistent LifecyclePolicy = "persistent" +) + +type Pool struct { + Name string + Addr string + Options []tai.Option + MaxPerUser int + MaxTotal int + IdleTimeout time.Duration + MaxLifetime time.Duration +} + +type PoolInfo struct { + Name string + Addr string + Connected bool + Boxes int + MaxPerUser int + MaxTotal int + IdleTimeout time.Duration + MaxLifetime time.Duration +} + +type PortMapping struct { + ContainerPort int + HostPort int + HostIP string + Protocol string +} + +type CreateOptions struct { + ID string + Owner string + Labels map[string]string + Pool string + Image string + WorkDir string + User string + Env map[string]string + Memory int64 + CPUs float64 + VNC bool + Ports []PortMapping + Policy LifecyclePolicy + IdleTimeout time.Duration +} + +type ListOptions struct { + Owner string + Pool string + Labels map[string]string +} + +type execConfig struct { + WorkDir string + Env map[string]string + Timeout time.Duration +} + +type ExecOption func(*execConfig) + +func WithWorkDir(dir string) ExecOption { + return func(c *execConfig) { + c.WorkDir = dir + } +} + +func WithEnv(env map[string]string) ExecOption { + return func(c *execConfig) { + c.Env = env + } +} + +func WithTimeout(timeout time.Duration) ExecOption { + return func(c *execConfig) { + c.Timeout = timeout + } +} + +type ExecResult struct { + ExitCode int + Stdout string + Stderr string +} + +type ExecStream struct { + Stdout io.ReadCloser + Stderr io.ReadCloser + Stdin io.WriteCloser + Wait func() (int, error) + Cancel func() +} + +type attachConfig struct { + Protocol string + Path string + Headers map[string]string +} + +type AttachOption func(*attachConfig) + +func WithProtocol(protocol string) AttachOption { + return func(c *attachConfig) { + c.Protocol = protocol + } +} + +func WithPath(path string) AttachOption { + return func(c *attachConfig) { + c.Path = path + } +} + +func WithHeaders(headers map[string]string) AttachOption { + return func(c *attachConfig) { + c.Headers = headers + } +} + +type ServiceConn struct { + Read func() ([]byte, error) + Write func(data []byte) error + Events <-chan []byte + URL string + Close func() error +} + +type BoxInfo struct { + ID string + ContainerID string + Pool string + Owner string + Status string + Policy LifecyclePolicy + Labels map[string]string + Image string + CreatedAt time.Time + LastActive time.Time + ProcessCount int + VNC bool +} diff --git a/tai/docs/README.md b/tai/docs/README.md index 47848611..8048d7c8 100644 --- a/tai/docs/README.md +++ b/tai/docs/README.md @@ -18,8 +18,8 @@ Go client library for the [Tai](https://github.com/YaoApp/tai) runtime bridge. P ### Local Mode (direct Docker) ```go -c, err := tai.New("") -// or: tai.New("unix:///var/run/docker.sock") +c, err := tai.New("local") +// or: tai.New("docker:///var/run/docker.sock") // or: tai.New("tcp://192.168.1.50:2375") defer c.Close() diff --git a/tai/grpc/cmd/main.go b/tai/grpc/cmd/main.go index e7ab70b9..c05bb097 100644 --- a/tai/grpc/cmd/main.go +++ b/tai/grpc/cmd/main.go @@ -71,6 +71,10 @@ func serve() error { ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) defer cancel() + if sandboxID := os.Getenv("YAO_SANDBOX_ID"); sandboxID != "" { + go yaogrpc.HeartbeatLoop(ctx, client, sandboxID) + } + scanner := bufio.NewScanner(os.Stdin) scanner.Buffer(make([]byte, 0, 4*1024*1024), 4*1024*1024) encoder := json.NewEncoder(os.Stdout) diff --git a/tai/grpc/grpc.go b/tai/grpc/grpc.go index 4f538bcd..5c914f3d 100644 --- a/tai/grpc/grpc.go +++ b/tai/grpc/grpc.go @@ -5,6 +5,7 @@ import ( "fmt" "io" "os" + "strings" "github.com/yaoapp/yao/grpc/pb" "google.golang.org/grpc" @@ -36,6 +37,9 @@ func NewFromEnv() (*Client, error) { } // Dial connects to the gRPC server at addr with the given TokenManager. +// Bare host:port addresses are wrapped with passthrough:/// for grpc.NewClient +// compatibility (grpc.NewClient defaults to dns scheme which may fail for hostnames +// like host.docker.internal). func Dial(addr string, tm *TokenManager) (*Client, error) { opts := []grpc.DialOption{ grpc.WithTransportCredentials(insecure.NewCredentials()), @@ -47,7 +51,12 @@ func Dial(addr string, tm *TokenManager) (*Client, error) { ) } - conn, err := grpc.NewClient(addr, opts...) + target := addr + if !strings.Contains(addr, "://") { + target = "passthrough:///" + addr + } + + conn, err := grpc.NewClient(target, opts...) if err != nil { return nil, fmt.Errorf("dial %s: %w", addr, err) } @@ -228,6 +237,22 @@ func (c *Client) AgentStream(ctx context.Context, assistantID string, messages, } } +// --- Sandbox --- + +// Heartbeat sends a sandbox heartbeat to the Yao gRPC server. +func (c *Client) Heartbeat(ctx context.Context, sandboxID string, cpuPercent int32, memBytes int64, runningProcs int32) (string, error) { + resp, err := c.svc.Heartbeat(ctx, &pb.HeartbeatRequest{ + SandboxId: sandboxID, + CpuPercent: cpuPercent, + MemBytes: memBytes, + RunningProcs: runningProcs, + }) + if err != nil { + return "", err + } + return resp.Action, nil +} + // --- Health --- // Healthz checks the server health. diff --git a/tai/grpc/grpc_test.go b/tai/grpc/grpc_test.go index 47057843..e252b0dc 100644 --- a/tai/grpc/grpc_test.go +++ b/tai/grpc/grpc_test.go @@ -169,6 +169,41 @@ func TestDial_WithTokenManager(t *testing.T) { assert.False(t, c.TokenManager().IsTaiMode()) } +func TestDial_PassthroughPrefix_BareAddress(t *testing.T) { + c, err := yaogrpc.Dial("host.docker.internal:9099", nil) + require.NoError(t, err) + defer c.Close() + + assert.Equal(t, "passthrough:///host.docker.internal:9099", c.Conn().Target()) +} + +func TestDial_PassthroughPrefix_IPAddress(t *testing.T) { + c, err := yaogrpc.Dial("192.168.1.100:9100", nil) + require.NoError(t, err) + defer c.Close() + + assert.Equal(t, "passthrough:///192.168.1.100:9100", c.Conn().Target()) +} + +func TestDial_PassthroughPrefix_PreservesExistingScheme(t *testing.T) { + tests := []struct { + addr string + target string + }{ + {"dns:///myhost:9099", "dns:///myhost:9099"}, + {"passthrough:///127.0.0.1:9099", "passthrough:///127.0.0.1:9099"}, + {"unix:///var/run/grpc.sock", "unix:///var/run/grpc.sock"}, + } + for _, tt := range tests { + t.Run(tt.addr, func(t *testing.T) { + c, err := yaogrpc.Dial(tt.addr, nil) + require.NoError(t, err) + defer c.Close() + assert.Equal(t, tt.target, c.Conn().Target()) + }) + } +} + func TestClient_Close_Nil(t *testing.T) { c := &yaogrpc.Client{} assert.NoError(t, c.Close()) diff --git a/tai/grpc/heartbeat.go b/tai/grpc/heartbeat.go new file mode 100644 index 00000000..7a42043c --- /dev/null +++ b/tai/grpc/heartbeat.go @@ -0,0 +1,77 @@ +package grpc + +import ( + "context" + "fmt" + "os" + "os/exec" + "runtime" + "strconv" + "strings" + "time" +) + +const defaultHeartbeatInterval = 10 * time.Second + +// HeartbeatLoop sends periodic heartbeats to the Yao gRPC server. +// It runs until ctx is cancelled. The sandboxID comes from YAO_SANDBOX_ID. +func HeartbeatLoop(ctx context.Context, client *Client, sandboxID string) { + interval := defaultHeartbeatInterval + if s := os.Getenv("YAO_HEARTBEAT_INTERVAL"); s != "" { + if d, err := time.ParseDuration(s); err == nil && d > 0 { + interval = d + } + } + + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + cpu, mem := sampleResources() + procs := countUserProcesses() + action, err := client.Heartbeat(ctx, sandboxID, cpu, mem, procs) + if err != nil { + continue + } + if action == "shutdown" { + fmt.Fprintf(os.Stderr, "yao-grpc: received shutdown signal\n") + p, _ := os.FindProcess(os.Getpid()) + p.Signal(os.Interrupt) + return + } + } + } +} + +// countUserProcesses counts running processes owned by the current user. +func countUserProcesses() int32 { + if runtime.GOOS != "linux" { + return 0 + } + + out, err := exec.Command("sh", "-c", "ps -e --no-headers | wc -l").Output() + if err != nil { + return 0 + } + n, _ := strconv.Atoi(strings.TrimSpace(string(out))) + return int32(n) +} + +// sampleResources reads basic CPU/memory stats from /proc (Linux only). +func sampleResources() (cpuPercent int32, memBytes int64) { + if runtime.GOOS != "linux" { + return 0, 0 + } + + data, err := os.ReadFile("/sys/fs/cgroup/memory.current") + if err == nil { + mem, _ := strconv.ParseInt(strings.TrimSpace(string(data)), 10, 64) + memBytes = mem + } + + return 0, memBytes +} diff --git a/tai/grpc/heartbeat_test.go b/tai/grpc/heartbeat_test.go new file mode 100644 index 00000000..a7db6bd5 --- /dev/null +++ b/tai/grpc/heartbeat_test.go @@ -0,0 +1,194 @@ +package grpc + +import ( + "context" + "net" + "sync/atomic" + "testing" + "time" + + "github.com/yaoapp/yao/grpc/pb" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" +) + +func TestCountUserProcesses(t *testing.T) { + n := countUserProcesses() + if n < 0 { + t.Errorf("countUserProcesses() = %d, want >= 0", n) + } +} + +func TestSampleResources(t *testing.T) { + cpu, mem := sampleResources() + if cpu < 0 || mem < 0 { + t.Errorf("sampleResources() = (%d, %d), want non-negative", cpu, mem) + } +} + +// ── HeartbeatLoop tests with mock gRPC server ─────────────────────────────── + +type mockYaoServer struct { + pb.UnimplementedYaoServer + calls atomic.Int32 + action string +} + +func (m *mockYaoServer) Heartbeat(_ context.Context, req *pb.HeartbeatRequest) (*pb.HeartbeatResponse, error) { + m.calls.Add(1) + return &pb.HeartbeatResponse{Action: m.action}, nil +} + +func startMockServer(t *testing.T, srv *mockYaoServer) (addr string, stop func()) { + t.Helper() + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + s := grpc.NewServer() + pb.RegisterYaoServer(s, srv) + go s.Serve(lis) + return lis.Addr().String(), s.Stop +} + +func dialClient(t *testing.T, addr string) *Client { + t.Helper() + c, err := Dial(addr, nil) + if err != nil { + t.Fatal(err) + } + return c +} + +func TestHeartbeatLoop_SendsHeartbeats(t *testing.T) { + mock := &mockYaoServer{action: "ok"} + addr, stop := startMockServer(t, mock) + defer stop() + + client := dialClient(t, addr) + defer client.Close() + + t.Setenv("YAO_HEARTBEAT_INTERVAL", "50ms") + + ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) + defer cancel() + + HeartbeatLoop(ctx, client, "sb-test") + + calls := mock.calls.Load() + if calls < 2 { + t.Errorf("expected at least 2 heartbeat calls, got %d", calls) + } +} + +func TestHeartbeatLoop_ShutdownAction(t *testing.T) { + mock := &mockYaoServer{action: "shutdown"} + addr, stop := startMockServer(t, mock) + defer stop() + + client := dialClient(t, addr) + defer client.Close() + + action, err := client.Heartbeat(context.Background(), "sb-shutdown", 0, 0, 0) + if err != nil { + t.Fatalf("Heartbeat: %v", err) + } + if action != "shutdown" { + t.Errorf("action = %q, want %q", action, "shutdown") + } + if mock.calls.Load() != 1 { + t.Errorf("expected 1 call, got %d", mock.calls.Load()) + } +} + +func TestHeartbeatLoop_ContextCancelStops(t *testing.T) { + mock := &mockYaoServer{action: "ok"} + addr, stop := startMockServer(t, mock) + defer stop() + + client := dialClient(t, addr) + defer client.Close() + + t.Setenv("YAO_HEARTBEAT_INTERVAL", "5s") + + ctx, cancel := context.WithCancel(context.Background()) + + done := make(chan struct{}) + go func() { + HeartbeatLoop(ctx, client, "sb-cancel") + close(done) + }() + + time.Sleep(50 * time.Millisecond) + cancel() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("HeartbeatLoop did not stop after context cancel") + } +} + +func TestHeartbeatLoop_IntervalParsing(t *testing.T) { + mock := &mockYaoServer{action: "ok"} + addr, stop := startMockServer(t, mock) + defer stop() + + client := dialClient(t, addr) + defer client.Close() + + t.Setenv("YAO_HEARTBEAT_INTERVAL", "30ms") + + ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond) + defer cancel() + + HeartbeatLoop(ctx, client, "sb-interval") + + calls := mock.calls.Load() + if calls < 3 { + t.Errorf("with 30ms interval over 150ms, expected >= 3 calls, got %d", calls) + } +} + +func TestHeartbeatLoop_InvalidIntervalUsesDefault(t *testing.T) { + mock := &mockYaoServer{action: "ok"} + addr, stop := startMockServer(t, mock) + defer stop() + + client := dialClient(t, addr) + defer client.Close() + + t.Setenv("YAO_HEARTBEAT_INTERVAL", "not-a-duration") + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + HeartbeatLoop(ctx, client, "sb-invalid") + + if mock.calls.Load() > 0 { + t.Error("with default 10s interval and 100ms timeout, expected 0 calls") + } +} + +func TestClientHeartbeat_ReturnsAction(t *testing.T) { + mock := &mockYaoServer{action: "ok"} + addr, stop := startMockServer(t, mock) + defer stop() + + conn, err := grpc.NewClient("passthrough:///"+addr, + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + + c := &Client{conn: conn, svc: pb.NewYaoClient(conn)} + action, err := c.Heartbeat(context.Background(), "sb-1", 50, 2048, 5) + if err != nil { + t.Fatalf("Heartbeat: %v", err) + } + if action != "ok" { + t.Errorf("action = %q, want %q", action, "ok") + } +} diff --git a/tai/proxy/connect.go b/tai/proxy/connect.go new file mode 100644 index 00000000..ecd9f043 --- /dev/null +++ b/tai/proxy/connect.go @@ -0,0 +1,117 @@ +package proxy + +import ( + "bufio" + "bytes" + "context" + "fmt" + "net/http" + "strings" + + "github.com/gorilla/websocket" +) + +// --- Remote Connect --- + +func (r *remoteProxy) Connect(ctx context.Context, containerID string, opts ConnectOptions) (*Connection, error) { + baseURL, err := r.URL(ctx, containerID, opts.Port, opts.Path) + if err != nil { + return nil, err + } + return connect(ctx, baseURL, opts.Protocol, r.client) +} + +// --- Local Connect --- + +func (l *localProxy) Connect(ctx context.Context, containerID string, opts ConnectOptions) (*Connection, error) { + baseURL, err := l.URL(ctx, containerID, opts.Port, opts.Path) + if err != nil { + return nil, err + } + return connect(ctx, baseURL, opts.Protocol, http.DefaultClient) +} + +func connect(ctx context.Context, url string, protocol string, hc *http.Client) (*Connection, error) { + switch protocol { + case "ws": + return connectWS(ctx, url) + case "sse": + return connectSSE(ctx, url, hc) + default: + return nil, fmt.Errorf("unsupported connect protocol: %q", protocol) + } +} + +func connectWS(ctx context.Context, rawURL string) (*Connection, error) { + wsURL := strings.Replace(rawURL, "http://", "ws://", 1) + wsURL = strings.Replace(wsURL, "https://", "wss://", 1) + + conn, _, err := websocket.DefaultDialer.DialContext(ctx, wsURL, nil) + if err != nil { + return nil, fmt.Errorf("ws dial: %w", err) + } + + ch := make(chan []byte, 64) + go func() { + defer close(ch) + for { + _, msg, err := conn.ReadMessage() + if err != nil { + return + } + ch <- msg + } + }() + + return &Connection{ + Messages: ch, + Send: func(data []byte) error { + return conn.WriteMessage(websocket.TextMessage, data) + }, + Close: func() error { + return conn.Close() + }, + }, nil +} + +func connectSSE(ctx context.Context, url string, hc *http.Client) (*Connection, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, err + } + req.Header.Set("Accept", "text/event-stream") + + resp, err := hc.Do(req) + if err != nil { + return nil, fmt.Errorf("sse connect: %w", err) + } + if resp.StatusCode != http.StatusOK { + resp.Body.Close() + return nil, fmt.Errorf("sse: status %d", resp.StatusCode) + } + + ch := make(chan []byte, 64) + go func() { + defer close(ch) + defer resp.Body.Close() + scanner := bufio.NewScanner(resp.Body) + for scanner.Scan() { + line := scanner.Text() + if strings.HasPrefix(line, "data: ") { + data := strings.TrimPrefix(line, "data: ") + ch <- bytes.Clone([]byte(data)) + } + } + }() + + return &Connection{ + Messages: ch, + Send: func(data []byte) error { + return fmt.Errorf("sse: send not supported") + }, + Close: func() error { + resp.Body.Close() + return nil + }, + }, nil +} diff --git a/tai/proxy/proxy.go b/tai/proxy/proxy.go index d74ee8a5..c79ca5d4 100644 --- a/tai/proxy/proxy.go +++ b/tai/proxy/proxy.go @@ -13,9 +13,27 @@ import ( // Remote routes through Tai HTTP proxy; Local resolves host ports directly. type Proxy interface { URL(ctx context.Context, containerID string, port int, path string) (string, error) + Connect(ctx context.Context, containerID string, opts ConnectOptions) (*Connection, error) Healthz(ctx context.Context) error } +// ConnectOptions configures a persistent connection to a container service. +type ConnectOptions struct { + Port int // container port + Path string // URL path (e.g. "/ws" or "/events") + Protocol string // "ws", "sse", or "tcp" +} + +// Connection represents a persistent connection to a container service. +type Connection struct { + // Messages receives incoming data. Channel is closed when the connection ends. + Messages <-chan []byte + // Send writes data to the connection (only valid for "ws" protocol). + Send func(data []byte) error + // Close terminates the connection. + Close func() error +} + // --- Remote implementation --- type remoteProxy struct { diff --git a/tai/proxy/proxy_test.go b/tai/proxy/proxy_test.go index e988897c..5e5aaa1d 100644 --- a/tai/proxy/proxy_test.go +++ b/tai/proxy/proxy_test.go @@ -5,9 +5,11 @@ import ( "fmt" "net/http" "net/http/httptest" + "strings" "testing" "time" + "github.com/gorilla/websocket" "github.com/yaoapp/yao/tai/sandbox" ) @@ -138,6 +140,126 @@ func TestHostIP(t *testing.T) { } } +// ── Connect tests ───────────────────────────────────────────────────────────── + +func TestConnect_UnsupportedProtocol(t *testing.T) { + _, err := connect(context.Background(), "http://127.0.0.1:1234", "tcp", http.DefaultClient) + if err == nil { + t.Fatal("expected error for unsupported protocol") + } + if !strings.Contains(err.Error(), "unsupported connect protocol") { + t.Errorf("unexpected error: %v", err) + } +} + +func TestConnectWS_EchoRoundtrip(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer c.Close() + for { + mt, msg, err := c.ReadMessage() + if err != nil { + return + } + c.WriteMessage(mt, msg) + } + })) + defer srv.Close() + + conn, err := connectWS(context.Background(), srv.URL) + if err != nil { + t.Fatalf("connectWS: %v", err) + } + defer conn.Close() + + if err := conn.Send([]byte("hello")); err != nil { + t.Fatalf("Send: %v", err) + } + + select { + case msg := <-conn.Messages: + if string(msg) != "hello" { + t.Errorf("got %q, want %q", msg, "hello") + } + case <-time.After(2 * time.Second): + t.Fatal("timeout waiting for echo") + } +} + +func TestConnectSSE_ReceiveEvents(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + flusher, _ := w.(http.Flusher) + for i := 0; i < 3; i++ { + fmt.Fprintf(w, "data: event-%d\n\n", i) + flusher.Flush() + } + })) + defer srv.Close() + + conn, err := connectSSE(context.Background(), srv.URL, srv.Client()) + if err != nil { + t.Fatalf("connectSSE: %v", err) + } + defer conn.Close() + + var events []string + for msg := range conn.Messages { + events = append(events, string(msg)) + if len(events) >= 3 { + break + } + } + if len(events) != 3 { + t.Fatalf("got %d events, want 3", len(events)) + } + for i, e := range events { + want := fmt.Sprintf("event-%d", i) + if e != want { + t.Errorf("event[%d] = %q, want %q", i, e, want) + } + } +} + +func TestConnectSSE_SendNotSupported(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, "data: x\n\n") + })) + defer srv.Close() + + conn, err := connectSSE(context.Background(), srv.URL, srv.Client()) + if err != nil { + t.Fatalf("connectSSE: %v", err) + } + defer conn.Close() + + if err := conn.Send([]byte("test")); err == nil { + t.Error("expected error from SSE Send") + } +} + +func TestConnectSSE_Non200(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusServiceUnavailable) + })) + defer srv.Close() + + _, err := connectSSE(context.Background(), srv.URL, srv.Client()) + if err == nil { + t.Fatal("expected error for non-200") + } + if !strings.Contains(err.Error(), "status 503") { + t.Errorf("unexpected error: %v", err) + } +} + // mockSandbox implements sandbox.Sandbox for testing. type mockSandbox struct { inspectFn func(ctx context.Context, id string) (*sandbox.ContainerInfo, error) @@ -154,6 +276,9 @@ func (m *mockSandbox) Remove(ctx context.Context, id string, force bool) error { func (m *mockSandbox) Exec(ctx context.Context, id string, cmd []string, opts sandbox.ExecOptions) (*sandbox.ExecResult, error) { return nil, nil } +func (m *mockSandbox) ExecStream(ctx context.Context, id string, cmd []string, opts sandbox.ExecOptions) (*sandbox.StreamHandle, error) { + return nil, nil +} func (m *mockSandbox) Inspect(ctx context.Context, id string) (*sandbox.ContainerInfo, error) { if m.inspectFn != nil { return m.inspectFn(ctx, id) diff --git a/tai/sandbox/docker.go b/tai/sandbox/docker.go index 93fd7c6f..1a23ab75 100644 --- a/tai/sandbox/docker.go +++ b/tai/sandbox/docker.go @@ -51,6 +51,10 @@ func (d *dockerSandbox) Exec(ctx context.Context, id string, cmd []string, opts return d.core.exec(ctx, id, cmd, opts) } +func (d *dockerSandbox) ExecStream(ctx context.Context, id string, cmd []string, opts ExecOptions) (*StreamHandle, error) { + return d.core.execStream(ctx, id, cmd, opts) +} + func (d *dockerSandbox) Inspect(ctx context.Context, id string) (*ContainerInfo, error) { return d.core.inspect(ctx, id) } diff --git a/tai/sandbox/docker_core.go b/tai/sandbox/docker_core.go index beeb8893..a0e77581 100644 --- a/tai/sandbox/docker_core.go +++ b/tai/sandbox/docker_core.go @@ -26,6 +26,8 @@ func (d *dockerCore) create(ctx context.Context, opts CreateOptions, addVNCPorts Cmd: opts.Cmd, Env: envSlice(opts.Env), WorkingDir: opts.WorkingDir, + Labels: opts.Labels, + User: opts.User, } hostCfg := &container.HostConfig{ @@ -129,6 +131,68 @@ func (d *dockerCore) exec(ctx context.Context, id string, cmd []string, opts Exe }, nil } +func (d *dockerCore) execStream(ctx context.Context, id string, cmd []string, opts ExecOptions) (*StreamHandle, error) { + execCfg := container.ExecOptions{ + Cmd: cmd, + WorkingDir: opts.WorkDir, + Env: envSlice(opts.Env), + AttachStdin: true, + AttachStdout: true, + AttachStderr: true, + } + + execResp, err := d.cli.ContainerExecCreate(ctx, id, execCfg) + if err != nil { + return nil, fmt.Errorf("exec create: %w", err) + } + + resp, err := d.cli.ContainerExecAttach(ctx, execResp.ID, container.ExecAttachOptions{}) + if err != nil { + return nil, fmt.Errorf("exec attach: %w", err) + } + + execCtx, execCancel := context.WithCancel(ctx) + + stdinR, stdinW := io.Pipe() + stdoutR, stdoutW := io.Pipe() + stderrR, stderrW := io.Pipe() + + // Pump user writes into the multiplexed connection. + // Closing stdinW sends EOF to the container stdin without + // tearing down the underlying connection (which carries stdout/stderr). + go func() { + io.Copy(resp.Conn, stdinR) + resp.CloseWrite() + }() + + go func() { + _, _ = stdcopy.StdCopy(stdoutW, stderrW, resp.Reader) + stdoutW.Close() + stderrW.Close() + }() + + return &StreamHandle{ + Stdin: stdinW, + Stdout: stdoutR, + Stderr: stderrR, + Wait: func() (int, error) { + for { + inspect, err := d.cli.ContainerExecInspect(execCtx, execResp.ID) + if err != nil { + return -1, fmt.Errorf("exec inspect: %w", err) + } + if !inspect.Running { + return inspect.ExitCode, nil + } + } + }, + Cancel: func() { + execCancel() + resp.Close() + }, + }, nil +} + func (d *dockerCore) inspect(ctx context.Context, id string) (*ContainerInfo, error) { info, err := d.cli.ContainerInspect(ctx, id) if err != nil { @@ -140,6 +204,7 @@ func (d *dockerCore) inspect(ctx context.Context, id string) (*ContainerInfo, er Name: strings.TrimPrefix(info.Name, "/"), Image: info.Config.Image, Status: info.State.Status, + Labels: info.Config.Labels, } if info.NetworkSettings != nil { @@ -196,6 +261,7 @@ func (d *dockerCore) list(ctx context.Context, opts ListOptions) ([]ContainerInf Name: name, Image: c.Image, Status: c.State, + Labels: c.Labels, } for _, p := range c.Ports { ci.Ports = append(ci.Ports, PortMapping{ diff --git a/tai/sandbox/k8s.go b/tai/sandbox/k8s.go index bf93bdb1..8bcd29d8 100644 --- a/tai/sandbox/k8s.go +++ b/tai/sandbox/k8s.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "fmt" + "io" "path/filepath" "strings" "time" @@ -126,6 +127,9 @@ func (s *k8sSandbox) Create(ctx context.Context, opts CreateOptions) (string, er labels[k] = v } labels["sandbox-name"] = name + for k, v := range opts.Labels { + labels[k] = v + } pod := &corev1.Pod{ ObjectMeta: metav1.ObjectMeta{ @@ -139,6 +143,15 @@ func (s *k8sSandbox) Create(ctx context.Context, opts CreateOptions) (string, er }, } + if opts.User != "" { + uid, err := parseUID(opts.User) + if err == nil { + pod.Spec.SecurityContext = &corev1.PodSecurityContext{ + RunAsUser: &uid, + } + } + } + created, err := s.cli.CoreV1().Pods(s.ns).Create(ctx, pod, metav1.CreateOptions{}) if err != nil { return "", fmt.Errorf("create pod: %w", err) @@ -236,6 +249,78 @@ func (s *k8sSandbox) Exec(ctx context.Context, id string, cmd []string, opts Exe }, nil } +func (s *k8sSandbox) ExecStream(ctx context.Context, id string, cmd []string, opts ExecOptions) (*StreamHandle, error) { + execCmd := cmd + if opts.WorkDir != "" || len(opts.Env) > 0 { + var prefix string + for k, v := range opts.Env { + prefix += fmt.Sprintf("export %s=%q; ", k, v) + } + cdPart := "" + if opts.WorkDir != "" { + cdPart = fmt.Sprintf("cd %s && ", opts.WorkDir) + } + execCmd = []string{"sh", "-c", cdPart + prefix + strings.Join(cmd, " ")} + } + + req := s.cli.CoreV1().RESTClient().Post(). + Resource("pods"). + Name(id). + Namespace(s.ns). + SubResource("exec"). + VersionedParams(&corev1.PodExecOptions{ + Container: "main", + Command: execCmd, + Stdin: true, + Stdout: true, + Stderr: true, + }, scheme.ParameterCodec) + + exec, err := remotecommand.NewSPDYExecutor(s.cfg, "POST", req.URL()) + if err != nil { + return nil, fmt.Errorf("create executor: %w", err) + } + + stdinR, stdinW := io.Pipe() + stdoutR, stdoutW := io.Pipe() + stderrR, stderrW := io.Pipe() + + execCtx, cancel := context.WithCancel(ctx) + done := make(chan error, 1) + var exitCode int + + go func() { + err := exec.StreamWithContext(execCtx, remotecommand.StreamOptions{ + Stdin: stdinR, + Stdout: stdoutW, + Stderr: stderrW, + }) + if err != nil { + if exitErr, ok := err.(interface{ ExitStatus() int }); ok { + exitCode = exitErr.ExitStatus() + err = nil + } + } + stdoutW.Close() + stderrW.Close() + done <- err + }() + + return &StreamHandle{ + Stdin: stdinW, + Stdout: stdoutR, + Stderr: stderrR, + Wait: func() (int, error) { + err := <-done + return exitCode, err + }, + Cancel: func() { + cancel() + stdinR.Close() + }, + }, nil +} + func (s *k8sSandbox) Inspect(ctx context.Context, id string) (*ContainerInfo, error) { pod, err := s.cli.CoreV1().Pods(s.ns).Get(ctx, id, metav1.GetOptions{}) if err != nil { @@ -248,6 +333,7 @@ func (s *k8sSandbox) Inspect(ctx context.Context, id string) (*ContainerInfo, er Image: pod.Spec.Containers[0].Image, Status: string(pod.Status.Phase), IP: pod.Status.PodIP, + Labels: pod.Labels, }, nil } @@ -273,6 +359,7 @@ func (s *k8sSandbox) List(ctx context.Context, opts ListOptions) ([]ContainerInf Name: pod.Name, Status: string(pod.Status.Phase), IP: pod.Status.PodIP, + Labels: pod.Labels, } if len(pod.Spec.Containers) > 0 { ci.Image = pod.Spec.Containers[0].Image @@ -286,6 +373,14 @@ func (s *k8sSandbox) Close() error { return nil // REST client doesn't need explicit close } +// parseUID extracts a numeric UID from a user string like "1000" or "1000:1000". +func parseUID(user string) (int64, error) { + parts := strings.SplitN(user, ":", 2) + var uid int64 + _, err := fmt.Sscanf(parts[0], "%d", &uid) + return uid, err +} + func buildResources(memory int64, cpus float64) corev1.ResourceRequirements { limits := corev1.ResourceList{} if memory > 0 { diff --git a/tai/sandbox/local.go b/tai/sandbox/local.go index e9c5a709..916d9baf 100644 --- a/tai/sandbox/local.go +++ b/tai/sandbox/local.go @@ -55,6 +55,10 @@ func (l *local) Exec(ctx context.Context, id string, cmd []string, opts ExecOpti return l.core.exec(ctx, id, cmd, opts) } +func (l *local) ExecStream(ctx context.Context, id string, cmd []string, opts ExecOptions) (*StreamHandle, error) { + return l.core.execStream(ctx, id, cmd, opts) +} + func (l *local) Inspect(ctx context.Context, id string) (*ContainerInfo, error) { return l.core.inspect(ctx, id) } diff --git a/tai/sandbox/sandbox.go b/tai/sandbox/sandbox.go index 368ddcf0..31f854aa 100644 --- a/tai/sandbox/sandbox.go +++ b/tai/sandbox/sandbox.go @@ -2,6 +2,7 @@ package sandbox import ( "context" + "io" "time" ) @@ -13,11 +14,23 @@ type Sandbox interface { Stop(ctx context.Context, id string, timeout time.Duration) error Remove(ctx context.Context, id string, force bool) error Exec(ctx context.Context, id string, cmd []string, opts ExecOptions) (*ExecResult, error) + ExecStream(ctx context.Context, id string, cmd []string, opts ExecOptions) (*StreamHandle, error) Inspect(ctx context.Context, id string) (*ContainerInfo, error) List(ctx context.Context, opts ListOptions) ([]ContainerInfo, error) Close() error } +// StreamHandle provides real-time I/O access to a running exec process. +type StreamHandle struct { + Stdin io.WriteCloser + Stdout io.Reader + Stderr io.Reader + // Wait blocks until the exec process finishes and returns the exit code. + Wait func() (int, error) + // Cancel aborts the exec process. + Cancel func() +} + // CreateOptions configures a new container. type CreateOptions struct { Name string @@ -30,6 +43,8 @@ type CreateOptions struct { CPUs float64 // 0 = no limit VNC bool Ports []PortMapping + Labels map[string]string // container/pod labels for discovery and management + User string // container user, e.g. "1000:1000" or "sandbox" } // PortMapping maps a container port to a host port. @@ -48,6 +63,7 @@ type ContainerInfo struct { Status string // "created", "running", "exited", "removing" IP string Ports []PortMapping + Labels map[string]string } // ExecOptions configures a command execution inside a container. diff --git a/tai/sandbox/sandbox_test.go b/tai/sandbox/sandbox_test.go index 061ff982..e061b557 100644 --- a/tai/sandbox/sandbox_test.go +++ b/tai/sandbox/sandbox_test.go @@ -2,7 +2,9 @@ package sandbox import ( "context" + "io" "os" + "strings" "testing" "time" @@ -625,3 +627,332 @@ func TestNewK8sRelativeKubeConfig(t *testing.T) { t.Skipf("K8s not available: %v", err) } } + +func TestCreateWithLabels(t *testing.T) { + sb, err := NewLocal("") + if err != nil { + t.Skipf("Docker not available: %v", err) + } + defer sb.Close() + + ctx := context.Background() + labels := map[string]string{ + "sandbox-id": "test-123", + "sandbox-owner": "user1", + } + + id, err := sb.Create(ctx, CreateOptions{ + Name: "tai-label-test", + Image: "alpine:latest", + Cmd: []string{"sleep", "10"}, + Labels: labels, + }) + if err != nil { + t.Fatalf("Create: %v", err) + } + defer sb.Remove(ctx, id, true) + + if err := sb.Start(ctx, id); err != nil { + t.Fatalf("Start: %v", err) + } + + info, err := sb.Inspect(ctx, id) + if err != nil { + t.Fatalf("Inspect: %v", err) + } + for k, v := range labels { + if info.Labels[k] != v { + t.Errorf("label %q = %q, want %q", k, info.Labels[k], v) + } + } + + listed, err := sb.List(ctx, ListOptions{ + Labels: map[string]string{"sandbox-id": "test-123"}, + }) + if err != nil { + t.Fatalf("List: %v", err) + } + found := false + for _, c := range listed { + if c.ID == id { + found = true + if c.Labels["sandbox-owner"] != "user1" { + t.Errorf("list labels missing sandbox-owner") + } + } + } + if !found { + t.Error("labeled container not found in filtered list") + } +} + +func TestCreateWithUser(t *testing.T) { + sb, err := NewLocal("") + if err != nil { + t.Skipf("Docker not available: %v", err) + } + defer sb.Close() + + ctx := context.Background() + id, err := sb.Create(ctx, CreateOptions{ + Name: "tai-user-test", + Image: "alpine:latest", + Cmd: []string{"sleep", "10"}, + User: "1000:1000", + }) + if err != nil { + t.Fatalf("Create: %v", err) + } + defer sb.Remove(ctx, id, true) + + if err := sb.Start(ctx, id); err != nil { + t.Fatalf("Start: %v", err) + } + + result, err := sb.Exec(ctx, id, []string{"id", "-u"}, ExecOptions{}) + if err != nil { + t.Fatalf("Exec: %v", err) + } + if result.Stdout != "1000\n" { + t.Errorf("user id = %q, want %q", result.Stdout, "1000\n") + } +} + +func TestExecStream_ShortCommand(t *testing.T) { + sb, err := NewLocal("") + if err != nil { + t.Skipf("Docker not available: %v", err) + } + defer sb.Close() + + ctx := context.Background() + id, err := sb.Create(ctx, CreateOptions{ + Name: "tai-stream-short", + Image: "alpine:latest", + Cmd: []string{"sleep", "30"}, + }) + if err != nil { + t.Fatalf("Create: %v", err) + } + defer sb.Remove(ctx, id, true) + + if err := sb.Start(ctx, id); err != nil { + t.Fatalf("Start: %v", err) + } + + stream, err := sb.ExecStream(ctx, id, []string{"echo", "hello-stream"}, ExecOptions{}) + if err != nil { + t.Fatalf("ExecStream: %v", err) + } + + out, err := io.ReadAll(stream.Stdout) + if err != nil { + t.Fatalf("ReadAll stdout: %v", err) + } + if string(out) != "hello-stream\n" { + t.Errorf("stdout = %q, want %q", string(out), "hello-stream\n") + } + + code, err := stream.Wait() + if err != nil { + t.Fatalf("Wait: %v", err) + } + if code != 0 { + t.Errorf("exit code = %d, want 0", code) + } +} + +func TestExecStream_Stdin(t *testing.T) { + sb, err := NewLocal("") + if err != nil { + t.Skipf("Docker not available: %v", err) + } + defer sb.Close() + + ctx := context.Background() + id, err := sb.Create(ctx, CreateOptions{ + Name: "tai-stream-stdin", + Image: "alpine:latest", + Cmd: []string{"sleep", "30"}, + }) + if err != nil { + t.Fatalf("Create: %v", err) + } + defer sb.Remove(ctx, id, true) + + if err := sb.Start(ctx, id); err != nil { + t.Fatalf("Start: %v", err) + } + + stream, err := sb.ExecStream(ctx, id, []string{"cat"}, ExecOptions{}) + if err != nil { + t.Fatalf("ExecStream: %v", err) + } + + _, err = stream.Stdin.Write([]byte("from-stdin\n")) + if err != nil { + t.Fatalf("Write stdin: %v", err) + } + stream.Stdin.Close() + + out, err := io.ReadAll(stream.Stdout) + if err != nil { + t.Fatalf("ReadAll stdout: %v", err) + } + if string(out) != "from-stdin\n" { + t.Errorf("stdout = %q, want %q", string(out), "from-stdin\n") + } + + code, err := stream.Wait() + if err != nil { + t.Fatalf("Wait: %v", err) + } + if code != 0 { + t.Errorf("exit code = %d, want 0", code) + } +} + +func TestExecStream_ExitCode(t *testing.T) { + sb, err := NewLocal("") + if err != nil { + t.Skipf("Docker not available: %v", err) + } + defer sb.Close() + + ctx := context.Background() + id, err := sb.Create(ctx, CreateOptions{ + Name: "tai-stream-exit", + Image: "alpine:latest", + Cmd: []string{"sleep", "30"}, + }) + if err != nil { + t.Fatalf("Create: %v", err) + } + defer sb.Remove(ctx, id, true) + + if err := sb.Start(ctx, id); err != nil { + t.Fatalf("Start: %v", err) + } + + stream, err := sb.ExecStream(ctx, id, []string{"sh", "-c", "exit 42"}, ExecOptions{}) + if err != nil { + t.Fatalf("ExecStream: %v", err) + } + + io.ReadAll(stream.Stdout) + code, err := stream.Wait() + if err != nil { + t.Fatalf("Wait: %v", err) + } + if code != 42 { + t.Errorf("exit code = %d, want 42", code) + } +} + +func TestExecStream_Stderr(t *testing.T) { + sb, err := NewLocal("") + if err != nil { + t.Skipf("Docker not available: %v", err) + } + defer sb.Close() + + ctx := context.Background() + id, err := sb.Create(ctx, CreateOptions{ + Name: "tai-stream-stderr", + Image: "alpine:latest", + Cmd: []string{"sleep", "30"}, + }) + if err != nil { + t.Fatalf("Create: %v", err) + } + defer sb.Remove(ctx, id, true) + + if err := sb.Start(ctx, id); err != nil { + t.Fatalf("Start: %v", err) + } + + stream, err := sb.ExecStream(ctx, id, []string{"sh", "-c", "echo err-msg >&2"}, ExecOptions{}) + if err != nil { + t.Fatalf("ExecStream: %v", err) + } + + stderr, err := io.ReadAll(stream.Stderr) + if err != nil { + t.Fatalf("ReadAll stderr: %v", err) + } + if !strings.Contains(string(stderr), "err-msg") { + t.Errorf("stderr = %q, want to contain %q", string(stderr), "err-msg") + } + + code, _ := stream.Wait() + if code != 0 { + t.Errorf("exit code = %d, want 0", code) + } +} + +func TestExecStream_Cancel(t *testing.T) { + sb, err := NewLocal("") + if err != nil { + t.Skipf("Docker not available: %v", err) + } + defer sb.Close() + + ctx := context.Background() + id, err := sb.Create(ctx, CreateOptions{ + Name: "tai-stream-cancel", + Image: "alpine:latest", + Cmd: []string{"sleep", "30"}, + }) + if err != nil { + t.Fatalf("Create: %v", err) + } + defer sb.Remove(ctx, id, true) + + if err := sb.Start(ctx, id); err != nil { + t.Fatalf("Start: %v", err) + } + + stream, err := sb.ExecStream(ctx, id, []string{"sleep", "300"}, ExecOptions{}) + if err != nil { + t.Fatalf("ExecStream: %v", err) + } + + stream.Cancel() + + done := make(chan struct{}) + go func() { + stream.Wait() + close(done) + }() + + select { + case <-done: + case <-time.After(5 * time.Second): + t.Error("Wait did not return after Cancel within 5s") + } +} + +func TestParseUID(t *testing.T) { + tests := []struct { + input string + want int64 + ok bool + }{ + {"1000", 1000, true}, + {"1000:1000", 1000, true}, + {"0", 0, true}, + {"abc", 0, false}, + } + for _, tt := range tests { + got, err := parseUID(tt.input) + if tt.ok && err != nil { + t.Errorf("parseUID(%q): unexpected error %v", tt.input, err) + } + if !tt.ok && err == nil { + t.Errorf("parseUID(%q): expected error", tt.input) + } + if tt.ok && got != tt.want { + t.Errorf("parseUID(%q) = %d, want %d", tt.input, got, tt.want) + } + } +} diff --git a/tai/serverinfo/pb/serverinfo.pb.go b/tai/serverinfo/pb/serverinfo.pb.go new file mode 100644 index 00000000..5d85f88a --- /dev/null +++ b/tai/serverinfo/pb/serverinfo.pb.go @@ -0,0 +1,195 @@ +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.36.11 +// protoc v4.25.0 +// source: tai/serverinfo/pb/serverinfo.proto + +package pb + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + reflect "reflect" + sync "sync" + unsafe "unsafe" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +type GetInfoRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *GetInfoRequest) Reset() { + *x = GetInfoRequest{} + mi := &file_tai_serverinfo_pb_serverinfo_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *GetInfoRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetInfoRequest) ProtoMessage() {} + +func (x *GetInfoRequest) ProtoReflect() protoreflect.Message { + mi := &file_tai_serverinfo_pb_serverinfo_proto_msgTypes[0] + 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 GetInfoRequest.ProtoReflect.Descriptor instead. +func (*GetInfoRequest) Descriptor() ([]byte, []int) { + return file_tai_serverinfo_pb_serverinfo_proto_rawDescGZIP(), []int{0} +} + +type GetInfoResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + 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" + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *GetInfoResponse) Reset() { + *x = GetInfoResponse{} + mi := &file_tai_serverinfo_pb_serverinfo_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *GetInfoResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GetInfoResponse) ProtoMessage() {} + +func (x *GetInfoResponse) ProtoReflect() protoreflect.Message { + mi := &file_tai_serverinfo_pb_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 GetInfoResponse.ProtoReflect.Descriptor instead. +func (*GetInfoResponse) Descriptor() ([]byte, []int) { + return file_tai_serverinfo_pb_serverinfo_proto_rawDescGZIP(), []int{1} +} + +func (x *GetInfoResponse) GetVersion() string { + if x != nil { + return x.Version + } + return "" +} + +func (x *GetInfoResponse) GetPorts() map[string]int32 { + if x != nil { + return x.Ports + } + return nil +} + +func (x *GetInfoResponse) GetCapabilities() map[string]bool { + if x != nil { + return x.Capabilities + } + return nil +} + +var File_tai_serverinfo_pb_serverinfo_proto protoreflect.FileDescriptor + +const file_tai_serverinfo_pb_serverinfo_proto_rawDesc = "" + + "\n" + + "\"tai/serverinfo/pb/serverinfo.proto\x12\n" + + "serverinfo\"\x10\n" + + "\x0eGetInfoRequest\"\xb7\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" + + "\n" + + "PortsEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\x05R\x05value:\x028\x01\x1a?\n" + + "\x11CapabilitiesEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\bR\x05value:\x028\x012P\n" + + "\n" + + "ServerInfo\x12B\n" + + "\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 +) + +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))) + }) + return file_tai_serverinfo_pb_serverinfo_proto_rawDescData +} + +var file_tai_serverinfo_pb_serverinfo_proto_msgTypes = make([]protoimpl.MessageInfo, 4) +var file_tai_serverinfo_pb_serverinfo_proto_goTypes = []any{ + (*GetInfoRequest)(nil), // 0: serverinfo.GetInfoRequest + (*GetInfoResponse)(nil), // 1: serverinfo.GetInfoResponse + nil, // 2: serverinfo.GetInfoResponse.PortsEntry + nil, // 3: 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 +} + +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 { + 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)), + NumEnums: 0, + NumMessages: 4, + 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, + }.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 +} diff --git a/tai/serverinfo/pb/serverinfo.proto b/tai/serverinfo/pb/serverinfo.proto new file mode 100644 index 00000000..4eece2ee --- /dev/null +++ b/tai/serverinfo/pb/serverinfo.proto @@ -0,0 +1,15 @@ +syntax = "proto3"; +package serverinfo; +option go_package = "github.com/yaoapp/tai/serverinfo/pb"; + +service ServerInfo { + rpc GetInfo(GetInfoRequest) returns (GetInfoResponse); +} + +message GetInfoRequest {} + +message GetInfoResponse { + string version = 1; + map ports = 2; // "grpc", "http", "vnc", "docker", "k8s" + map capabilities = 3; // "docker", "k8s" +} diff --git a/tai/serverinfo/pb/serverinfo_grpc.pb.go b/tai/serverinfo/pb/serverinfo_grpc.pb.go new file mode 100644 index 00000000..07e690ec --- /dev/null +++ b/tai/serverinfo/pb/serverinfo_grpc.pb.go @@ -0,0 +1,121 @@ +// Code generated by protoc-gen-go-grpc. DO NOT EDIT. +// versions: +// - protoc-gen-go-grpc v1.6.1 +// - protoc v4.25.0 +// source: tai/serverinfo/pb/serverinfo.proto + +package pb + +import ( + context "context" + grpc "google.golang.org/grpc" + codes "google.golang.org/grpc/codes" + status "google.golang.org/grpc/status" +) + +// This is a compile-time assertion to ensure that this generated file +// is compatible with the grpc package it is being compiled against. +// Requires gRPC-Go v1.64.0 or later. +const _ = grpc.SupportPackageIsVersion9 + +const ( + ServerInfo_GetInfo_FullMethodName = "/serverinfo.ServerInfo/GetInfo" +) + +// ServerInfoClient is the client API for ServerInfo service. +// +// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. +type ServerInfoClient interface { + GetInfo(ctx context.Context, in *GetInfoRequest, opts ...grpc.CallOption) (*GetInfoResponse, error) +} + +type serverInfoClient struct { + cc grpc.ClientConnInterface +} + +func NewServerInfoClient(cc grpc.ClientConnInterface) ServerInfoClient { + return &serverInfoClient{cc} +} + +func (c *serverInfoClient) GetInfo(ctx context.Context, in *GetInfoRequest, opts ...grpc.CallOption) (*GetInfoResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(GetInfoResponse) + err := c.cc.Invoke(ctx, ServerInfo_GetInfo_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +// ServerInfoServer is the server API for ServerInfo service. +// All implementations must embed UnimplementedServerInfoServer +// for forward compatibility. +type ServerInfoServer interface { + GetInfo(context.Context, *GetInfoRequest) (*GetInfoResponse, error) + mustEmbedUnimplementedServerInfoServer() +} + +// UnimplementedServerInfoServer must be embedded to have +// forward compatible implementations. +// +// NOTE: this should be embedded by value instead of pointer to avoid a nil +// pointer dereference when methods are called. +type UnimplementedServerInfoServer struct{} + +func (UnimplementedServerInfoServer) GetInfo(context.Context, *GetInfoRequest) (*GetInfoResponse, error) { + return nil, status.Error(codes.Unimplemented, "method GetInfo not implemented") +} +func (UnimplementedServerInfoServer) mustEmbedUnimplementedServerInfoServer() {} +func (UnimplementedServerInfoServer) testEmbeddedByValue() {} + +// UnsafeServerInfoServer may be embedded to opt out of forward compatibility for this service. +// Use of this interface is not recommended, as added methods to ServerInfoServer will +// result in compilation errors. +type UnsafeServerInfoServer interface { + mustEmbedUnimplementedServerInfoServer() +} + +func RegisterServerInfoServer(s grpc.ServiceRegistrar, srv ServerInfoServer) { + // If the following call panics, it indicates UnimplementedServerInfoServer was + // embedded by pointer and is nil. This will cause panics if an + // unimplemented method is ever invoked, so we test this at initialization + // time to prevent it from happening at runtime later due to I/O. + if t, ok := srv.(interface{ testEmbeddedByValue() }); ok { + t.testEmbeddedByValue() + } + s.RegisterService(&ServerInfo_ServiceDesc, srv) +} + +func _ServerInfo_GetInfo_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(GetInfoRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(ServerInfoServer).GetInfo(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: ServerInfo_GetInfo_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(ServerInfoServer).GetInfo(ctx, req.(*GetInfoRequest)) + } + return interceptor(ctx, in, info, handler) +} + +// ServerInfo_ServiceDesc is the grpc.ServiceDesc for ServerInfo service. +// It's only intended for direct use with grpc.RegisterService, +// and not to be introspected or modified (even as a copy) +var ServerInfo_ServiceDesc = grpc.ServiceDesc{ + ServiceName: "serverinfo.ServerInfo", + HandlerType: (*ServerInfoServer)(nil), + Methods: []grpc.MethodDesc{ + { + MethodName: "GetInfo", + Handler: _ServerInfo_GetInfo_Handler, + }, + }, + Streams: []grpc.StreamDesc{}, + Metadata: "tai/serverinfo/pb/serverinfo.proto", +} diff --git a/tai/tai.go b/tai/tai.go index 8c9ae33c..43695bfa 100644 --- a/tai/tai.go +++ b/tai/tai.go @@ -1,13 +1,17 @@ package tai import ( + "context" "fmt" "net/http" "net/url" + "strconv" "strings" + "time" "github.com/yaoapp/yao/tai/proxy" "github.com/yaoapp/yao/tai/sandbox" + sipb "github.com/yaoapp/yao/tai/serverinfo/pb" "github.com/yaoapp/yao/tai/vnc" "github.com/yaoapp/yao/tai/volume" "github.com/yaoapp/yao/tai/workspace" @@ -44,8 +48,12 @@ type Ports struct { } // WithPorts overrides default Tai service ports. +// Ports set here take precedence over server-reported values from ServerInfo. func WithPorts(p Ports) Option { - return optionFunc(func(c *config) { c.ports = p }) + return optionFunc(func(c *config) { + c.ports = p + c.userPorts = p + }) } // WithHTTPClient sets a custom HTTP client for proxy and VNC health checks. @@ -72,6 +80,7 @@ func WithNamespace(ns string) Option { type config struct { runtime Runtime ports Ports + userPorts Ports // tracks explicitly set ports (zero = not set by user) httpClient *http.Client dataDir string kubeConfig string @@ -121,9 +130,11 @@ type Client struct { // New creates a Client based on the address protocol: // -// "" → Local mode, platform default Docker socket +// "local" → Local mode, platform default Docker socket // "docker://addr" → Local mode, specified Docker daemon // "tai://host" → Remote mode via Tai Server +// +// Empty string is not allowed — use "local" for default local Docker. func New(addr string, opts ...Option) (*Client, error) { cfg := &config{ports: defaultPorts()} for _, o := range opts { @@ -131,11 +142,15 @@ func New(addr string, opts ...Option) (*Client, error) { } cfg.ports = mergedPorts(cfg.ports) - scheme, host, dockerAddr, err := parseAddr(addr) + scheme, host, dockerAddr, grpcPort, err := parseAddr(addr) if err != nil { return nil, err } + if grpcPort > 0 { + cfg.ports.GRPC = grpcPort + } + c := &Client{ scheme: scheme, host: host, @@ -171,16 +186,23 @@ func (c *Client) initLocal(cfg *config) (*Client, error) { } func (c *Client) initRemote(cfg *config) (*Client, error) { - // gRPC connection grpcAddr := fmt.Sprintf("%s:%d", c.host, c.ports.GRPC) conn, err := grpc.NewClient(grpcAddr, grpc.WithTransportCredentials(insecure.NewCredentials())) if err != nil { return nil, fmt.Errorf("grpc dial %s: %w", grpcAddr, err) } c.grpcConn = conn + + // Auto-discover server ports via ServerInfo RPC. + // Only overwrite ports that were NOT explicitly set by WithPorts. + if err := c.discoverPorts(conn, cfg); err != nil { + // Non-fatal: fall back to defaults / WithPorts values. + // Old Tai servers without ServerInfo will hit this path. + _ = err + } + c.vol = volume.NewRemote(conn) - // Sandbox switch cfg.runtime { case K8s: k8sPort := c.ports.K8s @@ -261,41 +283,100 @@ func (c *Client) VNC() vnc.VNC { return c.vc } // IsLocal returns true if the client connects directly to a Docker daemon. func (c *Client) IsLocal() bool { return c.scheme == "docker" } -func parseAddr(addr string) (scheme, host, dockerAddr string, err error) { +func parseAddr(addr string) (scheme, host, dockerAddr string, grpcPort int, err error) { addr = strings.TrimSpace(addr) if addr == "" { - return "docker", "", "", nil + return "", "", "", 0, fmt.Errorf("empty address: use \"local\" for default Docker daemon") + } + + if addr == "local" { + return "docker", "", "", 0, nil + } + + // Bare IP or host(:port) without scheme → normalise before url.Parse, + // which misparses bare addresses (treats them as path, not host). + if !strings.Contains(addr, "://") { + if isLocalHost(addr) { + return "docker", "", "", 0, nil + } + // host:port — split carefully (IPv6 like [::1]:9100 is already handled above) + h := addr + if idx := strings.LastIndex(addr, ":"); idx > 0 { + h = addr[:idx] + } + if isLocalHost(h) { + return "docker", "", "", 0, nil + } + addr = "tai://" + addr } u, parseErr := url.Parse(addr) if parseErr != nil { - return "", "", "", fmt.Errorf("parse addr %q: %w", addr, parseErr) + return "", "", "", 0, fmt.Errorf("parse addr %q: %w", addr, parseErr) } switch u.Scheme { case "tai": - host = u.Host - if host == "" { - return "", "", "", fmt.Errorf("tai:// requires a host") + hostname := u.Hostname() + if hostname == "" { + return "", "", "", 0, fmt.Errorf("tai:// requires a host") } - if idx := strings.Index(host, ":"); idx >= 0 { - host = host[:idx] + if portStr := u.Port(); portStr != "" { + if p, convErr := strconv.Atoi(portStr); convErr == nil && p > 0 { + grpcPort = p + } } - return "tai", host, "", nil + return "tai", hostname, "", grpcPort, nil case "docker": - return "docker", "", addr, nil + return "docker", "", addr, 0, nil case "unix": - return "docker", "", addr, nil + return "docker", "", addr, 0, nil case "tcp": - return "docker", "", addr, nil + return "docker", "", addr, 0, nil case "npipe": - return "docker", "", addr, nil + return "docker", "", addr, 0, nil default: - return "", "", "", fmt.Errorf("unsupported scheme %q in addr %q", u.Scheme, addr) + return "", "", "", 0, fmt.Errorf("unsupported scheme %q in addr %q", u.Scheme, addr) } } + +func isLocalHost(h string) bool { + return h == "127.0.0.1" || h == "localhost" || h == "::1" +} + +// discoverPorts calls ServerInfo.GetInfo on the remote Tai server and merges +// discovered ports into c.ports. Ports explicitly set via WithPorts (non-zero +// in the original config before merging defaults) take precedence. +func (c *Client) discoverPorts(conn *grpc.ClientConn, cfg *config) error { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + client := sipb.NewServerInfoClient(conn) + resp, err := client.GetInfo(ctx, &sipb.GetInfoRequest{}) + if err != nil { + return err + } + + // cfg.userPorts tracks what the caller explicitly passed to WithPorts. + // Only overwrite ports that the caller did NOT explicitly set. + up := cfg.userPorts + + if p := int(resp.Ports["http"]); p > 0 && up.HTTP == 0 { + c.ports.HTTP = p + } + if p := int(resp.Ports["docker"]); p > 0 && up.Docker == 0 { + c.ports.Docker = p + } + if p := int(resp.Ports["vnc"]); p > 0 && up.VNC == 0 { + c.ports.VNC = p + } + if p := int(resp.Ports["k8s"]); p > 0 && up.K8s == 0 { + c.ports.K8s = p + } + return nil +} diff --git a/tai/tai_test.go b/tai/tai_test.go index 92d2350d..3157d366 100644 --- a/tai/tai_test.go +++ b/tai/tai_test.go @@ -1,9 +1,12 @@ package tai import ( + "context" "os" "strconv" "testing" + + sipb "github.com/yaoapp/yao/tai/serverinfo/pb" ) func taiTestHost() string { @@ -24,28 +27,38 @@ func envPort(key string, fallback int) int { func TestParseAddr(t *testing.T) { tests := []struct { - addr string - wantScheme string - wantHost string - wantDocker string - wantErr bool + addr string + wantScheme string + wantHost string + wantDocker string + wantGRPCPort int + wantErr bool }{ - {"", "docker", "", "", false}, - {"docker:///var/run/docker.sock", "docker", "", "docker:///var/run/docker.sock", false}, - {"docker://192.168.1.50:2375", "docker", "", "docker://192.168.1.50:2375", false}, - {"unix:///var/run/docker.sock", "docker", "", "unix:///var/run/docker.sock", false}, - {"tcp://127.0.0.1:2375", "docker", "", "tcp://127.0.0.1:2375", false}, - {"npipe:////./pipe/docker_engine", "docker", "", "npipe:////./pipe/docker_engine", false}, - {"tai://192.168.1.100", "tai", "192.168.1.100", "", false}, - {"tai://10.0.0.5:9100", "tai", "10.0.0.5", "", false}, - {"tai://", "", "", "", true}, - {"ftp://host", "", "", "", true}, - {" tai://host ", "tai", "host", "", false}, + {"", "", "", "", 0, true}, + {"local", "docker", "", "", 0, false}, + {"127.0.0.1", "docker", "", "", 0, false}, + {"localhost", "docker", "", "", 0, false}, + {"::1", "docker", "", "", 0, false}, + {"docker:///var/run/docker.sock", "docker", "", "docker:///var/run/docker.sock", 0, false}, + {"docker://192.168.1.50:2375", "docker", "", "docker://192.168.1.50:2375", 0, false}, + {"unix:///var/run/docker.sock", "docker", "", "unix:///var/run/docker.sock", 0, false}, + {"tcp://127.0.0.1:2375", "docker", "", "tcp://127.0.0.1:2375", 0, false}, + {"npipe:////./pipe/docker_engine", "docker", "", "npipe:////./pipe/docker_engine", 0, false}, + {"tai://192.168.1.100", "tai", "192.168.1.100", "", 0, false}, + {"tai://10.0.0.5:9200", "tai", "10.0.0.5", "", 9200, false}, + {"tai://", "", "", "", 0, true}, + {"ftp://host", "", "", "", 0, true}, + {" tai://host ", "tai", "host", "", 0, false}, + // Bare non-local host → auto-prepend tai:// + {"192.168.1.50", "tai", "192.168.1.50", "", 0, false}, + {"192.168.1.50:9200", "tai", "192.168.1.50", "", 9200, false}, + {"my-server", "tai", "my-server", "", 0, false}, + {"my-server:9200", "tai", "my-server", "", 9200, false}, } for _, tt := range tests { t.Run(tt.addr, func(t *testing.T) { - scheme, host, dockerAddr, err := parseAddr(tt.addr) + scheme, host, dockerAddr, grpcPort, err := parseAddr(tt.addr) if (err != nil) != tt.wantErr { t.Fatalf("err = %v, wantErr = %v", err, tt.wantErr) } @@ -61,6 +74,9 @@ func TestParseAddr(t *testing.T) { if dockerAddr != tt.wantDocker { t.Errorf("dockerAddr = %q, want %q", dockerAddr, tt.wantDocker) } + if grpcPort != tt.wantGRPCPort { + t.Errorf("grpcPort = %d, want %d", grpcPort, tt.wantGRPCPort) + } }) } } @@ -98,6 +114,9 @@ func TestOptions(t *testing.T) { if cfg.ports.HTTP != 9999 { t.Errorf("WithPorts: HTTP = %d", cfg.ports.HTTP) } + if cfg.userPorts.HTTP != 9999 { + t.Errorf("WithPorts: userPorts.HTTP = %d", cfg.userPorts.HTTP) + } WithDataDir("/data").apply(cfg) if cfg.dataDir != "/data" { @@ -116,8 +135,15 @@ func TestOptions(t *testing.T) { } } +func TestNewEmptyAddr(t *testing.T) { + _, err := New("") + if err == nil { + t.Error("expected error for empty addr") + } +} + func TestNewLocal(t *testing.T) { - c, err := New("") + c, err := New("local") if err != nil { t.Skipf("Docker not available: %v", err) } @@ -148,7 +174,7 @@ func TestNewLocal(t *testing.T) { func TestNewLocalWithDataDir(t *testing.T) { dir := t.TempDir() - c, err := New("", WithDataDir(dir)) + c, err := New("local", WithDataDir(dir)) if err != nil { t.Skipf("Docker not available: %v", err) } @@ -258,11 +284,58 @@ func TestNewRemoteDocker(t *testing.T) { } } -func TestNewRemoteWithPorts(t *testing.T) { +func TestDiscoverPorts(t *testing.T) { addr := "tai://" + taiTestHost() - c, err := New(addr, WithPorts(Ports{HTTP: 8888})) + c, err := New(addr) if err != nil { t.Skipf("Tai not available at %s: %v", addr, err) } defer c.Close() + + // Query ServerInfo directly to get ground truth + sipClient := sipb.NewServerInfoClient(c.grpcConn) + resp, err := sipClient.GetInfo(context.Background(), &sipb.GetInfoRequest{}) + if err != nil { + t.Fatalf("ServerInfo.GetInfo failed: %v", err) + } + + t.Logf("server reported: %+v", resp.Ports) + t.Logf("client resolved: GRPC=%d HTTP=%d VNC=%d Docker=%d K8s=%d", + c.ports.GRPC, c.ports.HTTP, c.ports.VNC, c.ports.Docker, c.ports.K8s) + + check := func(name string, got int, serverVal int32) { + if sv := int(serverVal); sv > 0 && got != sv { + t.Errorf("%s = %d, server reported %d", name, got, sv) + } + } + check("HTTP", c.ports.HTTP, resp.Ports["http"]) + check("VNC", c.ports.VNC, resp.Ports["vnc"]) + check("Docker", c.ports.Docker, resp.Ports["docker"]) + check("GRPC", c.ports.GRPC, resp.Ports["grpc"]) +} + +func TestDiscoverPortsWithUserOverride(t *testing.T) { + addr := "tai://" + taiTestHost() + c, err := New(addr, WithPorts(Ports{HTTP: 9999})) + if err != nil { + t.Skipf("Tai not available at %s: %v", addr, err) + } + defer c.Close() + + if c.ports.HTTP != 9999 { + t.Errorf("HTTP = %d, want 9999 (user override should take precedence)", c.ports.HTTP) + } + + // Other ports should still be discovered from server + sipClient := sipb.NewServerInfoClient(c.grpcConn) + resp, err := sipClient.GetInfo(context.Background(), &sipb.GetInfoRequest{}) + if err != nil { + t.Fatalf("ServerInfo.GetInfo failed: %v", err) + } + + if sv := int(resp.Ports["vnc"]); sv > 0 && c.ports.VNC != sv { + t.Errorf("VNC = %d, server reported %d (non-overridden ports should be discovered)", c.ports.VNC, sv) + } + t.Logf("ports: GRPC=%d HTTP=%d(user) VNC=%d Docker=%d", + c.ports.GRPC, c.ports.HTTP, c.ports.VNC, c.ports.Docker) } diff --git a/tai/vnc/vnc_test.go b/tai/vnc/vnc_test.go index af6e41a2..36914e6f 100644 --- a/tai/vnc/vnc_test.go +++ b/tai/vnc/vnc_test.go @@ -202,6 +202,9 @@ func (m *mockSandbox) Remove(ctx context.Context, id string, force bool) error { func (m *mockSandbox) Exec(ctx context.Context, id string, cmd []string, opts sandbox.ExecOptions) (*sandbox.ExecResult, error) { return nil, nil } +func (m *mockSandbox) ExecStream(ctx context.Context, id string, cmd []string, opts sandbox.ExecOptions) (*sandbox.StreamHandle, error) { + return nil, nil +} func (m *mockSandbox) Inspect(ctx context.Context, id string) (*sandbox.ContainerInfo, error) { if m.inspectFn != nil { return m.inspectFn(ctx, id)