opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Implement Sandbox V2 support in the Yao SDK

Browse source 
- Add new gRPC endpoint for Heartbeat in the Yao service, enabling communication with the sandbox.
- Update Makefile to include a dedicated unit test target for Sandbox V2, ensuring proper testing of new features.
- Enhance CI workflows to incorporate Sandbox V2 tests, allowing for dual-mode testing (local and remote) with Docker.
- Modify .gitignore to exclude specific Docker files while allowing shell scripts for Sandbox V2.
- Update documentation in DESIGN.md to reflect the new architecture and capabilities of the Sandbox V2.

These changes enhance the Yao SDK's functionality, providing improved support for sandbox operations and testing.
This commit is contained in:
Max 2026-03-05 13:21:09 +08:00
parent df63584a6f
commit dfb33681f9
62 changed files with 7377 additions and 201 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

182
.github/workflows/pr-test.yml vendored
View file

@ -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
# =============================================================================

120
.github/workflows/unit-test.yml vendored
View file

@ -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
# =============================================================================

4
.gitignore vendored
View file

@ -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

15
Makefile
View file

@ -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:

3
grpc/auth/endpoint.go
View file

@ -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"
}

12
grpc/auth/endpoint_test.go
View file

@ -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)
}

2
grpc/auth/scope.go
View file

@ -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/"}},
)

31
grpc/grpc.go
View file

@ -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"
)
@ -40,6 +41,7 @@ type yaoServer struct {
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)

346
grpc/pb/yao.pb.go
View file

@ -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
}

16
grpc/pb/yao.proto
View file

@ -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"
}

44
grpc/pb/yao_grpc.pb.go
View file

@ -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",
}

76
grpc/sandbox/heartbeat.go Normal file
View file

@ -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()
}

158
grpc/sandbox/heartbeat_test.go Normal file
View file

@ -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")
}
}

90
sandbox/DESIGN.md
View file

@ -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 <url>` → 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 <url>` (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 <path>` 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.

8
sandbox/docker/build.sh
View file

@ -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
;;

1127
sandbox/v2/DESIGN.md Normal file
View file

File diff suppressed because it is too large Load diff

704
sandbox/v2/IMPL.md Normal file
View file

@ -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:<port>`. 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.

108
sandbox/v2/Makefile Normal file
View file

@ -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:-<not set, local only>}"
@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)

621
sandbox/v2/TEST.md Normal file
View file

@ -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/
```

261
sandbox/v2/box.go Normal file
View file

@ -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
}

133
sandbox/v2/box_attach_test.go Normal file
View file

@ -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)
}
})
}
}

214
sandbox/v2/box_test.go Normal file
View file

@ -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")
}
})
}
}

5
sandbox/v2/config.go Normal file
View file

@ -0,0 +1,5 @@
package sandbox
type Config struct {
Pool []Pool
}

38
sandbox/v2/docker/base/Dockerfile Normal file
View file

@ -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"]

12
sandbox/v2/docker/base/entrypoint.sh Executable file
View file

@ -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 "$@"

79
sandbox/v2/docker/build.sh Executable file
View file

@ -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

21
sandbox/v2/docker/test/Dockerfile Normal file
View file

@ -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"]

7
sandbox/v2/docker/test/entrypoint.sh Executable file
View file

@ -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 "$@"

28
sandbox/v2/docker/test/sse-server.py Normal file
View file

@ -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()

14
sandbox/v2/docker/test/ws-echo.py Normal file
View file

@ -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())

11
sandbox/v2/errors.go Normal file
View file

@ -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")
)

6
sandbox/v2/export_test.go Normal file
View file

@ -0,0 +1,6 @@
package sandbox
// ResetForTest resets the global manager for testing purposes.
func ResetForTest() {
mgr = nil
}

54
sandbox/v2/grpc.go Normal file
View file

@ -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
}

56
sandbox/v2/grpc_test.go Normal file
View file

@ -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")
}
}

519
sandbox/v2/manager.go Normal file
View file

@ -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)
}
}

141
sandbox/v2/manager_lifecycle_test.go Normal file
View file

@ -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)
}
})
}
}

293
sandbox/v2/manager_test.go Normal file
View file

@ -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")
}
}

23
sandbox/v2/sandbox.go Normal file
View file

@ -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
}

41
sandbox/v2/sandbox_test.go Normal file
View file

@ -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()
}

99
sandbox/v2/testutils_test.go Normal file
View file

@ -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
}

157
sandbox/v2/types.go Normal file
View file

@ -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
}

4
tai/docs/README.md
View file

@ -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()

4
tai/grpc/cmd/main.go
View file

@ -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)

27
tai/grpc/grpc.go
View file

@ -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.

35
tai/grpc/grpc_test.go
View file

@ -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())

77
tai/grpc/heartbeat.go Normal file
View file

@ -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
}

194
tai/grpc/heartbeat_test.go Normal file
View file

@ -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")
}
}

117
tai/proxy/connect.go Normal file
View file

@ -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
}

18
tai/proxy/proxy.go
View file

@ -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 {

125
tai/proxy/proxy_test.go
View file

@ -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)

4
tai/sandbox/docker.go
View file

@ -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)
}

66
tai/sandbox/docker_core.go
View file

@ -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{

95
tai/sandbox/k8s.go
View file

@ -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 {

4
tai/sandbox/local.go
View file

@ -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)
}

16
tai/sandbox/sandbox.go
View file

@ -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.

331
tai/sandbox/sandbox_test.go
View file

@ -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)
}
}
}

195
tai/serverinfo/pb/serverinfo.pb.go Normal file
View file

@ -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
}

15
tai/serverinfo/pb/serverinfo.proto Normal file
View file

@ -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<string, int32> ports = 2; // "grpc", "http", "vnc", "docker", "k8s"
map<string, bool> capabilities = 3; // "docker", "k8s"
}

121
tai/serverinfo/pb/serverinfo_grpc.pb.go Normal file
View file

@ -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",
}

119
tai/tai.go
View file

@ -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
}

105
tai/tai_test.go
View file

@ -1,9 +1,12 @@
package tai
import (
"context"
"os"
"strconv"
"testing"
sipb "github.com/yaoapp/yao/tai/serverinfo/pb"
)
func taiTestHost() string {
@ -28,24 +31,34 @@ func TestParseAddr(t *testing.T) {
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)
}

3
tai/vnc/vnc_test.go
View file

@ -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)