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

Merge pull request #1487 from trheyi/main

Browse source 
Implement Sandbox V2 support
This commit is contained in:
Max 2026-03-06 10:16:11 +08:00 • committed by GitHub
commit 4df3700ac5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
87 changed files with 12356 additions and 382 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -923,6 +923,283 @@ jobs:
body: '✅ Sandbox Tests passed!'
});
# =============================================================================
# Sandbox V2 Tests (tai + sandbox/v2 + workspace, Docker + K8s via k3d)
# =============================================================================
SandboxV2Test:
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"]
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 (tai + sandbox-v2 + workspace)...'
});
- 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 Apple Private Key
run: |
mkdir -p ../app/openapi/certs/apple
echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8
- name: Setup Go ${{ matrix.go }}
uses: actions/setup-go@v5
with:
go-version: ${{ matrix.go }}
- name: Start Redis
run: docker run --name redis --publish 6379:6379 --detach redis:6
- name: Setup Go Tools
run: make tools
- name: Setup ENV (SQLite)
run: |
mkdir -p ${{ github.WORKSPACE }}/../app/db
echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV
echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV
- name: Pull Test Images
run: |
docker pull yaoapp/sandbox-v2-test:latest || true
docker pull yaoapp/tai:latest
docker pull alpine:latest
- name: Install k3d
run: curl -s https://raw.githubusercontent.com/k3d-io/k3d/main/install.sh | bash
- name: Create k3d cluster
run: |
k3d cluster create tai-test --no-lb --wait --api-port 16443
kubectl wait --for=condition=Ready node --all --timeout=60s
k3d image import alpine:latest -c tai-test
- name: Start Tai Docker instance
run: |
docker run -d --name tai-docker \
-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 Docker HTTP ready"; break
fi
echo "Waiting for Tai Docker HTTP... ($i)"; sleep 1
done
curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1 || {
echo "::error::Tai Docker HTTP failed"; docker logs tai-docker 2>&1; exit 1
}
for i in $(seq 1 15); do
if nc -z 127.0.0.1 9100 2>/dev/null; then
echo "Tai Docker gRPC ready"; break
fi
echo "Waiting for Tai Docker gRPC... ($i)"; sleep 1
done
nc -z 127.0.0.1 9100 2>/dev/null || {
echo "::error::Tai Docker gRPC failed"; docker logs tai-docker 2>&1; exit 1
}
- name: Generate kubeconfig for Tai K8s
run: |
K3D_IP=$(docker inspect k3d-tai-test-server-0 | jq -r '.[0].NetworkSettings.Networks["k3d-tai-test"].IPAddress')
echo "k3d server IP: ${K3D_IP}"
k3d kubeconfig get tai-test > /tmp/kubeconfig-k3d.yml
# Kubeconfig for tai-k8s container (uses k3d-internal IP)
sed "s|server: .*|server: https://${K3D_IP}:6443|" /tmp/kubeconfig-k3d.yml \
> /tmp/kubeconfig-tai-k8s.yml
echo "Container kubeconfig server:"
grep server: /tmp/kubeconfig-tai-k8s.yml
# Kubeconfig for test runner (uses localhost via port-mapped 6443)
sed 's|server: .*|server: https://127.0.0.1:6443|' /tmp/kubeconfig-k3d.yml \
> ${{ runner.temp }}/kubeconfig-tai.yml
echo "Test runner kubeconfig server:"
grep server: ${{ runner.temp }}/kubeconfig-tai.yml
- name: Start Tai K8s instance
run: |
K3D_IP=$(docker inspect k3d-tai-test-server-0 | jq -r '.[0].NetworkSettings.Networks["k3d-tai-test"].IPAddress')
echo "k3d server IP: ${K3D_IP}"
docker run -d --name tai-k8s \
--network k3d-tai-test \
-p 8081:8080 -p 9101:9100 -p 6443:6443 -p 6081:6080 \
-v /var/run/docker.sock:/var/run/docker.sock:ro \
-v /tmp/kubeconfig-tai-k8s.yml:/etc/tai/kubeconfig.yml:ro \
-e TAI_K8S_UPSTREAM="tcp://${K3D_IP}:6443" \
-e TAI_KUBECONFIG=/etc/tai/kubeconfig.yml \
yaoapp/tai:latest
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:8081/healthz > /dev/null 2>&1; then
echo "Tai K8s HTTP ready"; break
fi
echo "Waiting for Tai K8s HTTP... ($i)"; sleep 1
done
curl -sf http://127.0.0.1:8081/healthz > /dev/null 2>&1 || {
echo "::error::Tai K8s HTTP failed"; docker logs tai-k8s 2>&1; exit 1
}
for i in $(seq 1 15); do
if nc -z 127.0.0.1 9101 2>/dev/null; then
echo "Tai K8s gRPC ready"; break
fi
echo "Waiting for Tai K8s gRPC... ($i)"; sleep 1
done
nc -z 127.0.0.1 9101 2>/dev/null || {
echo "::error::Tai K8s gRPC failed"; docker logs tai-k8s 2>&1; exit 1
}
- name: Run Sandbox V2 Tests (tai + sandbox-v2 + workspace)
env:
TAI_TEST_HOST: "127.0.0.1"
TAI_TEST_DOCKER: "tcp://127.0.0.1:2375"
TAI_TEST_GRPC_PORT: "9100"
TAI_TEST_K8S_HOST: "127.0.0.1"
TAI_TEST_K8S_PORT: "6443"
TAI_TEST_K8S_GRPC_PORT: "9101"
TAI_TEST_KUBECONFIG: "${{ runner.temp }}/kubeconfig-tai.yml"
TAI_TEST_HOST_IP: "172.17.0.1"
SANDBOX_TEST_REMOTE_ADDR: "tai://127.0.0.1:9100"
SANDBOX_TEST_IMAGE: "yaoapp/sandbox-v2-test:latest"
run: make unit-test-sandbox-v2
- name: Codecov Report
if: always()
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
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 (tai + sandbox-v2 + workspace)!'
});
# =============================================================================
# Benchmark & Memory Leak Tests - Run with MySQL8.0 and SQLite3
# =============================================================================
@ -1533,20 +1810,10 @@ jobs:
});
# =============================================================================
# Tai SDK Tests (requires Tai container with Docker socket mount)
# Benchmark: Sandbox V2 + Workspace (parallel with SandboxV2Test, non-blocking)
# =============================================================================
TaiTest:
BenchmarkSandboxV2:
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"]
@ -1583,20 +1850,6 @@ jobs:
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: '🤖 Tai SDK Tests running...'
});
- name: Checkout Kun
uses: actions/checkout@v4
with:
@ -1679,8 +1932,9 @@ jobs:
echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV
echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV
- name: Pull Tai & Test Images
- name: Pull Test Images
run: |
docker pull yaoapp/sandbox-v2-test:latest || true
docker pull yaoapp/tai:latest
docker pull alpine:latest
@ -1693,91 +1947,89 @@ jobs:
kubectl wait --for=condition=Ready node --all --timeout=60s
k3d image import alpine:latest -c tai-test
- name: Start Tai (with Docker socket + K8s proxy)
- name: Start Tai Docker instance (benchmarks)
run: |
docker run -d --name tai-docker \
-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 Docker HTTP ready"; break
fi
echo "Waiting for Tai Docker HTTP... ($i)"; sleep 1
done
curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1 || {
echo "::error::Tai Docker HTTP failed"; docker logs tai-docker 2>&1; exit 1
}
for i in $(seq 1 15); do
if nc -z 127.0.0.1 9100 2>/dev/null; then
echo "Tai Docker gRPC ready"; break
fi
echo "Waiting for Tai Docker gRPC... ($i)"; sleep 1
done
nc -z 127.0.0.1 9100 2>/dev/null || {
echo "::error::Tai Docker gRPC failed"; docker logs tai-docker 2>&1; exit 1
}
- name: Generate kubeconfig for benchmarks
run: |
K3D_IP=$(docker inspect k3d-tai-test-server-0 | jq -r '.[0].NetworkSettings.Networks["k3d-tai-test"].IPAddress')
k3d kubeconfig get tai-test > /tmp/kubeconfig-k3d-bench.yml
sed "s|server: .*|server: https://${K3D_IP}:6443|" /tmp/kubeconfig-k3d-bench.yml \
> /tmp/kubeconfig-tai-k8s-bench.yml
sed 's|server: .*|server: https://127.0.0.1:6443|' /tmp/kubeconfig-k3d-bench.yml \
> ${{ runner.temp }}/kubeconfig-tai.yml
- name: Start Tai K8s instance (benchmarks)
run: |
K3D_IP=$(docker inspect k3d-tai-test-server-0 | jq -r '.[0].NetworkSettings.Networks["k3d-tai-test"].IPAddress')
echo "k3d server IP: ${K3D_IP}"
docker run -d --name tai \
docker run -d --name tai-k8s \
--network k3d-tai-test \
-v /var/run/docker.sock:/var/run/docker.sock \
-p 8080:8080 -p 9100:9100 -p 2375:2375 -p 6080:6080 -p 6443:6443 \
-p 8081:8080 -p 9101:9100 -p 6443:6443 -p 6081:6080 \
-v /var/run/docker.sock:/var/run/docker.sock:ro \
-v /tmp/kubeconfig-tai-k8s-bench.yml:/etc/tai/kubeconfig.yml:ro \
-e TAI_K8S_UPSTREAM="tcp://${K3D_IP}:6443" \
-e TAI_KUBECONFIG=/etc/tai/kubeconfig.yml \
yaoapp/tai:latest
TAI_HTTP_READY=false
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"
TAI_HTTP_READY=true
break
if curl -sf http://127.0.0.1:8081/healthz > /dev/null 2>&1; then
echo "Tai K8s HTTP ready"; break
fi
echo "Waiting for Tai HTTP... ($i)"
sleep 1
echo "Waiting for Tai K8s HTTP... ($i)"; sleep 1
done
if [ "$TAI_HTTP_READY" != "true" ]; then
echo "::error::Tai HTTP failed to become ready within 30s"
echo "--- Tai container logs ---"
docker logs tai 2>&1 || true
echo "--- Tai container status ---"
docker inspect tai --format='{{.State.Status}} exit={{.State.ExitCode}}' || true
exit 1
fi
curl -sf http://127.0.0.1:8081/healthz > /dev/null 2>&1 || {
echo "::error::Tai K8s HTTP failed"; docker logs tai-k8s 2>&1; exit 1
}
TAI_GRPC_READY=false
for i in $(seq 1 15); do
if nc -z 127.0.0.1 9100 2>/dev/null; then
echo "Tai gRPC is ready"
TAI_GRPC_READY=true
break
if nc -z 127.0.0.1 9101 2>/dev/null; then
echo "Tai K8s gRPC ready"; break
fi
echo "Waiting for Tai gRPC... ($i)"
sleep 1
echo "Waiting for Tai K8s gRPC... ($i)"; sleep 1
done
if [ "$TAI_GRPC_READY" != "true" ]; then
echo "::error::Tai gRPC failed to become ready within 15s"
echo "--- Tai container logs ---"
docker logs tai 2>&1 || true
exit 1
fi
nc -z 127.0.0.1 9101 2>/dev/null || {
echo "::error::Tai K8s gRPC failed"; docker logs tai-k8s 2>&1; exit 1
}
- name: Generate kubeconfig for Tai K8s proxy
run: |
k3d kubeconfig get tai-test > /tmp/kubeconfig-k3d.yml
sed 's|server: .*|server: https://127.0.0.1:6443|' /tmp/kubeconfig-k3d.yml \
> ${{ runner.temp }}/kubeconfig-tai.yml
echo "Generated kubeconfig:"
grep server: ${{ runner.temp }}/kubeconfig-tai.yml
- name: Run Tai SDK Tests
- name: Run Benchmarks
env:
TAI_TEST_HOST: "127.0.0.1"
TAI_TEST_GRPC: "127.0.0.1:9100"
TAI_TEST_DOCKER: "tcp://127.0.0.1:2375"
TAI_TEST_GRPC_PORT: "9100"
TAI_TEST_K8S_HOST: "127.0.0.1"
TAI_TEST_K8S_PORT: "6443"
TAI_TEST_K8S_GRPC_PORT: "9101"
TAI_TEST_KUBECONFIG: "${{ runner.temp }}/kubeconfig-tai.yml"
TAI_TEST_HOST_IP: "172.17.0.1"
run: make unit-test-tai
- name: Codecov Report
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
- name: "Comment on PR - Tai 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: '✅ Tai SDK Tests passed!'
});
SANDBOX_TEST_REMOTE_ADDR: "tai://127.0.0.1:9100"
SANDBOX_TEST_IMAGE: "yaoapp/sandbox-v2-test:latest"
run: make benchmark-sandbox-v2
# =============================================================================
# gRPC Tests - Run once with SQLite (transport layer, no DB matrix needed)

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

@ -679,6 +679,221 @@ jobs:
with:
token: ${{ secrets.CODECOV_TOKEN }}
# =============================================================================
# Sandbox V2 Tests (tai + sandbox/v2 + workspace, Docker + K8s via k3d)
# =============================================================================
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:
- 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 Apple Private Key
run: |
mkdir -p ../app/openapi/certs/apple
echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8
- name: Setup Go ${{ matrix.go }}
uses: actions/setup-go@v5
with:
go-version: ${{ matrix.go }}
- name: Start Redis
run: docker run --name redis --publish 6379:6379 --detach redis:6
- name: Setup Go Tools
run: make tools
- name: Setup ENV (SQLite)
run: |
mkdir -p ${{ github.WORKSPACE }}/../app/db
echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV
echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV
- name: Pull Test Images
run: |
docker pull yaoapp/sandbox-v2-test:latest || true
docker pull yaoapp/tai:latest
docker pull alpine:latest
- name: Install k3d
run: curl -s https://raw.githubusercontent.com/k3d-io/k3d/main/install.sh | bash
- name: Create k3d cluster
run: |
k3d cluster create tai-test --no-lb --wait --api-port 16443
kubectl wait --for=condition=Ready node --all --timeout=60s
k3d image import alpine:latest -c tai-test
- name: Start Tai Docker instance
run: |
docker run -d --name tai-docker \
-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 Docker HTTP ready"; break
fi
echo "Waiting for Tai Docker HTTP... ($i)"; sleep 1
done
curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1 || {
echo "::error::Tai Docker HTTP failed"; docker logs tai-docker 2>&1; exit 1
}
for i in $(seq 1 15); do
if nc -z 127.0.0.1 9100 2>/dev/null; then
echo "Tai Docker gRPC ready"; break
fi
echo "Waiting for Tai Docker gRPC... ($i)"; sleep 1
done
nc -z 127.0.0.1 9100 2>/dev/null || {
echo "::error::Tai Docker gRPC failed"; docker logs tai-docker 2>&1; exit 1
}
- name: Generate kubeconfig for Tai K8s
run: |
K3D_IP=$(docker inspect k3d-tai-test-server-0 | jq -r '.[0].NetworkSettings.Networks["k3d-tai-test"].IPAddress')
echo "k3d server IP: ${K3D_IP}"
k3d kubeconfig get tai-test > /tmp/kubeconfig-k3d.yml
# Kubeconfig for tai-k8s container (uses k3d-internal IP)
sed "s|server: .*|server: https://${K3D_IP}:6443|" /tmp/kubeconfig-k3d.yml \
> /tmp/kubeconfig-tai-k8s.yml
echo "Container kubeconfig server:"
grep server: /tmp/kubeconfig-tai-k8s.yml
# Kubeconfig for test runner (uses localhost via port-mapped 6443)
sed 's|server: .*|server: https://127.0.0.1:6443|' /tmp/kubeconfig-k3d.yml \
> ${{ runner.temp }}/kubeconfig-tai.yml
echo "Test runner kubeconfig server:"
grep server: ${{ runner.temp }}/kubeconfig-tai.yml
- name: Start Tai K8s instance
run: |
K3D_IP=$(docker inspect k3d-tai-test-server-0 | jq -r '.[0].NetworkSettings.Networks["k3d-tai-test"].IPAddress')
echo "k3d server IP: ${K3D_IP}"
docker run -d --name tai-k8s \
--network k3d-tai-test \
-p 8081:8080 -p 9101:9100 -p 6443:6443 -p 6081:6080 \
-v /var/run/docker.sock:/var/run/docker.sock:ro \
-v /tmp/kubeconfig-tai-k8s.yml:/etc/tai/kubeconfig.yml:ro \
-e TAI_K8S_UPSTREAM="tcp://${K3D_IP}:6443" \
-e TAI_KUBECONFIG=/etc/tai/kubeconfig.yml \
yaoapp/tai:latest
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:8081/healthz > /dev/null 2>&1; then
echo "Tai K8s HTTP ready"; break
fi
echo "Waiting for Tai K8s HTTP... ($i)"; sleep 1
done
curl -sf http://127.0.0.1:8081/healthz > /dev/null 2>&1 || {
echo "::error::Tai K8s HTTP failed"; docker logs tai-k8s 2>&1; exit 1
}
for i in $(seq 1 15); do
if nc -z 127.0.0.1 9101 2>/dev/null; then
echo "Tai K8s gRPC ready"; break
fi
echo "Waiting for Tai K8s gRPC... ($i)"; sleep 1
done
nc -z 127.0.0.1 9101 2>/dev/null || {
echo "::error::Tai K8s gRPC failed"; docker logs tai-k8s 2>&1; exit 1
}
- name: Run Sandbox V2 Tests (tai + sandbox-v2 + workspace)
env:
TAI_TEST_HOST: "127.0.0.1"
TAI_TEST_DOCKER: "tcp://127.0.0.1:2375"
TAI_TEST_GRPC_PORT: "9100"
TAI_TEST_K8S_HOST: "127.0.0.1"
TAI_TEST_K8S_PORT: "6443"
TAI_TEST_K8S_GRPC_PORT: "9101"
TAI_TEST_KUBECONFIG: "${{ runner.temp }}/kubeconfig-tai.yml"
TAI_TEST_HOST_IP: "172.17.0.1"
SANDBOX_TEST_REMOTE_ADDR: "tai://127.0.0.1:9100"
SANDBOX_TEST_IMAGE: "yaoapp/sandbox-v2-test:latest"
run: make unit-test-sandbox-v2
- name: Codecov Report
if: always()
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
fail_ci_if_error: false
# =============================================================================
# Benchmark & Memory Leak Tests - Run with MySQL8.0 and SQLite3
# =============================================================================
@ -1135,20 +1350,10 @@ jobs:
token: ${{ secrets.CODECOV_TOKEN }}
# =============================================================================
# Tai SDK Tests (requires Tai container with Docker socket mount)
# Benchmark: Sandbox V2 + Workspace (parallel with sandbox-v2-test)
# =============================================================================
tai-test:
benchmark-sandbox-v2:
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"]
@ -1233,8 +1438,9 @@ jobs:
echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV
echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV
- name: Pull Tai & Test Images
- name: Pull Test Images
run: |
docker pull yaoapp/sandbox-v2-test:latest || true
docker pull yaoapp/tai:latest
docker pull alpine:latest
@ -1247,77 +1453,89 @@ jobs:
kubectl wait --for=condition=Ready node --all --timeout=60s
k3d image import alpine:latest -c tai-test
- name: Start Tai (with Docker socket + K8s proxy)
- name: Start Tai Docker instance (benchmarks)
run: |
docker run -d --name tai-docker \
-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 Docker HTTP ready"; break
fi
echo "Waiting for Tai Docker HTTP... ($i)"; sleep 1
done
curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1 || {
echo "::error::Tai Docker HTTP failed"; docker logs tai-docker 2>&1; exit 1
}
for i in $(seq 1 15); do
if nc -z 127.0.0.1 9100 2>/dev/null; then
echo "Tai Docker gRPC ready"; break
fi
echo "Waiting for Tai Docker gRPC... ($i)"; sleep 1
done
nc -z 127.0.0.1 9100 2>/dev/null || {
echo "::error::Tai Docker gRPC failed"; docker logs tai-docker 2>&1; exit 1
}
- name: Generate kubeconfig for benchmarks
run: |
K3D_IP=$(docker inspect k3d-tai-test-server-0 | jq -r '.[0].NetworkSettings.Networks["k3d-tai-test"].IPAddress')
k3d kubeconfig get tai-test > /tmp/kubeconfig-k3d-bench.yml
sed "s|server: .*|server: https://${K3D_IP}:6443|" /tmp/kubeconfig-k3d-bench.yml \
> /tmp/kubeconfig-tai-k8s-bench.yml
sed 's|server: .*|server: https://127.0.0.1:6443|' /tmp/kubeconfig-k3d-bench.yml \
> ${{ runner.temp }}/kubeconfig-tai.yml
- name: Start Tai K8s instance (benchmarks)
run: |
K3D_IP=$(docker inspect k3d-tai-test-server-0 | jq -r '.[0].NetworkSettings.Networks["k3d-tai-test"].IPAddress')
echo "k3d server IP: ${K3D_IP}"
docker run -d --name tai \
docker run -d --name tai-k8s \
--network k3d-tai-test \
-v /var/run/docker.sock:/var/run/docker.sock \
-p 8080:8080 -p 9100:9100 -p 2375:2375 -p 6080:6080 -p 6443:6443 \
-p 8081:8080 -p 9101:9100 -p 6443:6443 -p 6081:6080 \
-v /var/run/docker.sock:/var/run/docker.sock:ro \
-v /tmp/kubeconfig-tai-k8s-bench.yml:/etc/tai/kubeconfig.yml:ro \
-e TAI_K8S_UPSTREAM="tcp://${K3D_IP}:6443" \
-e TAI_KUBECONFIG=/etc/tai/kubeconfig.yml \
yaoapp/tai:latest
TAI_HTTP_READY=false
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"
TAI_HTTP_READY=true
break
if curl -sf http://127.0.0.1:8081/healthz > /dev/null 2>&1; then
echo "Tai K8s HTTP ready"; break
fi
echo "Waiting for Tai HTTP... ($i)"
sleep 1
echo "Waiting for Tai K8s HTTP... ($i)"; sleep 1
done
if [ "$TAI_HTTP_READY" != "true" ]; then
echo "::error::Tai HTTP failed to become ready within 30s"
echo "--- Tai container logs ---"
docker logs tai 2>&1 || true
echo "--- Tai container status ---"
docker inspect tai --format='{{.State.Status}} exit={{.State.ExitCode}}' || true
exit 1
fi
curl -sf http://127.0.0.1:8081/healthz > /dev/null 2>&1 || {
echo "::error::Tai K8s HTTP failed"; docker logs tai-k8s 2>&1; exit 1
}
TAI_GRPC_READY=false
for i in $(seq 1 15); do
if nc -z 127.0.0.1 9100 2>/dev/null; then
echo "Tai gRPC is ready"
TAI_GRPC_READY=true
break
if nc -z 127.0.0.1 9101 2>/dev/null; then
echo "Tai K8s gRPC ready"; break
fi
echo "Waiting for Tai gRPC... ($i)"
sleep 1
echo "Waiting for Tai K8s gRPC... ($i)"; sleep 1
done
if [ "$TAI_GRPC_READY" != "true" ]; then
echo "::error::Tai gRPC failed to become ready within 15s"
echo "--- Tai container logs ---"
docker logs tai 2>&1 || true
exit 1
fi
nc -z 127.0.0.1 9101 2>/dev/null || {
echo "::error::Tai K8s gRPC failed"; docker logs tai-k8s 2>&1; exit 1
}
- name: Generate kubeconfig for Tai K8s proxy
run: |
k3d kubeconfig get tai-test > /tmp/kubeconfig-k3d.yml
sed 's|server: .*|server: https://127.0.0.1:6443|' /tmp/kubeconfig-k3d.yml \
> ${{ runner.temp }}/kubeconfig-tai.yml
echo "Generated kubeconfig:"
grep server: ${{ runner.temp }}/kubeconfig-tai.yml
- name: Run Tai SDK Tests
- name: Run Benchmarks
env:
TAI_TEST_HOST: "127.0.0.1"
TAI_TEST_GRPC: "127.0.0.1:9100"
TAI_TEST_DOCKER: "tcp://127.0.0.1:2375"
TAI_TEST_GRPC_PORT: "9100"
TAI_TEST_K8S_HOST: "127.0.0.1"
TAI_TEST_K8S_PORT: "6443"
TAI_TEST_K8S_GRPC_PORT: "9101"
TAI_TEST_KUBECONFIG: "${{ runner.temp }}/kubeconfig-tai.yml"
TAI_TEST_HOST_IP: "172.17.0.1"
run: make unit-test-tai
- name: Codecov Report
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
SANDBOX_TEST_REMOTE_ADDR: "tai://127.0.0.1:9100"
SANDBOX_TEST_IMAGE: "yaoapp/sandbox-v2-test:latest"
run: make benchmark-sandbox-v2
# =============================================================================
# gRPC Tests - Run once with SQLite (transport layer, no DB matrix needed)

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

99
Makefile
View file

@ -19,10 +19,12 @@ 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/...)
# Workspace tests (requires Tai for remote mode)
TESTFOLDER_WORKSPACE := $(shell $(GO) list ./workspace/...)
# gRPC tests
TESTFOLDER_GRPC := $(shell $(GO) list ./grpc/...)
TESTTAGS ?= ""
@ -199,6 +201,96 @@ unit-test-registry:
rm profile.out; \
fi
# ---------------------------------------------------------------------------
# Sandbox V2 Integration Test (tai + sandbox/v2 + workspace)
# Requires: Docker, Tai container, optionally k3d for K8s mode
# ---------------------------------------------------------------------------
SANDBOX_V2_IMAGE ?= yaoapp/sandbox-v2-test:latest
.PHONY: unit-test-sandbox-v2
unit-test-sandbox-v2: unit-test-sandbox-v2-pull unit-test-tai unit-test-sandbox-v2-core unit-test-workspace
@echo ""
@echo "============================================="
@echo "All Sandbox V2 integration tests passed"
@echo "============================================="
.PHONY: unit-test-sandbox-v2-pull
unit-test-sandbox-v2-pull:
@echo ""
@echo "============================================="
@echo "Pulling test images..."
@echo "============================================="
docker pull $(SANDBOX_V2_IMAGE) || true
docker pull alpine:latest || true
.PHONY: unit-test-sandbox-v2-core
unit-test-sandbox-v2-core:
@echo ""
@echo "============================================="
@echo "Running Sandbox V2 Tests..."
@echo "============================================="
$(MAKE) -C sandbox/v2 test-ci TEST_IMAGE=$(SANDBOX_V2_IMAGE)
# Workspace Unit Test (requires Tai for remote mode)
.PHONY: unit-test-workspace
unit-test-workspace:
@echo ""
@echo "============================================="
@echo "Running Workspace Tests..."
@echo "============================================="
echo "mode: count" > coverage.out
for d in $(TESTFOLDER_WORKSPACE); do \
$(GO) test -tags $(TESTTAGS) -v -timeout=10m -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; \
elif grep -q "setup failed" tmp.out; then \
rm tmp.out; \
exit 1; \
elif grep -q "runtime error" 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 "All workspace tests passed"
@echo "============================================="
# Benchmark: Sandbox V2 + Workspace
.PHONY: benchmark-sandbox-v2
benchmark-sandbox-v2:
@echo ""
@echo "============================================="
@echo "Running Sandbox V2 + Workspace Benchmarks..."
@echo "============================================="
@for d in $$($(GO) list ./sandbox/v2/... ./workspace/...); do \
if $(GO) test -list=Benchmark $$d 2>/dev/null | grep -q "^Benchmark"; then \
echo ""; \
echo "Benchmarking: $$d"; \
echo "---------------------------------------------"; \
$(GO) test -bench=. -benchmem -benchtime=1x -run='^$$' -timeout=600s $$d || true; \
fi; \
done
@echo ""
@echo "============================================="
@echo "All benchmarks completed"
@echo "============================================="
# Sandbox Unit Test (requires Docker)
.PHONY: unit-test-sandbox
unit-test-sandbox:
@ -251,9 +343,6 @@ unit-test-tai:
@echo "============================================="
@echo "Running Tai SDK Tests (requires Tai container)..."
@echo "============================================="
@echo "Pulling test images..."
docker pull alpine:latest || true
@echo ""
echo "mode: count" > coverage.out
for d in $(TESTFOLDER_TAI); do \
$(GO) test -tags $(TESTTAGS) -v -timeout=5m -covermode=count -coverprofile=profile.out -coverpkg=$$(echo $$d | sed "s/\/test$$//g") $$d > tmp.out; \

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

45
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"
)
@ -33,13 +34,14 @@ var (
type yaoServer struct {
pb.UnimplementedYaoServer
health health.Handler
run runhandler.Handler
shell shellhandler.Handler
api apihandler.Handler
mcp mcphandler.Handler
llm llmhandler.Handler
agent agenthandler.Handler
health health.Handler
run runhandler.Handler
shell shellhandler.Handler
api apihandler.Handler
mcp mcphandler.Handler
llm llmhandler.Handler
agent agenthandler.Handler
sandbox *sandboxhandler.Handler
}
// ── Health ───────────────────────────────────────────────────────────────────
@ -107,6 +109,30 @@ func (s *yaoServer) AgentStream(req *pb.AgentRequest, stream grpc.ServerStreamin
return s.agent.AgentStream(req, stream)
}
// ── Sandbox ──────────────────────────────────────────────────────────────────
func (s *yaoServer) Heartbeat(ctx context.Context, req *pb.HeartbeatRequest) (*pb.HeartbeatResponse, error) {
if s.sandbox == nil {
return &pb.HeartbeatResponse{Action: "ok"}, nil
}
return s.sandbox.Heartbeat(ctx, req)
}
// SandboxHandler returns the sandbox handler for external access (e.g., Manager integration).
func SandboxHandler() *sandboxhandler.Handler {
mu.Lock()
defer mu.Unlock()
return sandboxH
}
var sandboxH *sandboxhandler.Handler
// SetSandboxOnBeat sets the heartbeat callback for the sandbox handler.
// Must be called before StartServer.
func SetSandboxOnBeat(fn func(data *sandboxhandler.HeartbeatData) string) {
sandboxH = sandboxhandler.NewHandler(fn)
}
// ── Server lifecycle ─────────────────────────────────────────────────────────
// StartServer initializes and starts the gRPC server based on config.
@ -124,7 +150,10 @@ func StartServer(cfg config.Config) error {
grpc.ChainUnaryInterceptor(auth.UnaryInterceptor),
grpc.ChainStreamInterceptor(auth.StreamInterceptor),
)
pb.RegisterYaoServer(server, &yaoServer{})
if sandboxH == nil {
sandboxH = sandboxhandler.NewHandler(nil)
}
pb.RegisterYaoServer(server, &yaoServer{sandbox: sandboxH})
hosts := strings.Split(cfg.GRPC.Host, ",")
port := strconv.Itoa(cfg.GRPC.Port)

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

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

File diff suppressed because it is too large Load diff

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

@ -0,0 +1,263 @@
# Sandbox V2 — Implementation Status
Reference: [DESIGN.md](./DESIGN.md)
---
## Phase 1: Core Module — DONE
### tai SDK Prerequisites — DONE
| Step | Package | What | Status |
|------|---------|------|--------|
| 0 | `tai/sandbox` | Labels, User in CreateOptions + ContainerInfo | DONE |
| 1 | `tai/sandbox` | ExecStream (Docker + K8s) | DONE |
| 2 | `tai/proxy` | Connect (WS/SSE, Local + Remote) | DONE |
| 3 | `tai/sandbox` | Image interface (Exists, Pull, Remove, List) | DONE |
| 4 | `tai/tai.go` | Client: Sandbox(), Image(), Proxy(), VNC(), Volume(), Workspace() | DONE |
| 5 | `yao/grpc` | Heartbeat RPC (proto + handler) | DONE |
### sandbox/v2 Core — DONE
| File | What | Status |
|------|------|--------|
| `sandbox.go` | `Init()`, `M()`, global singleton | DONE |
| `manager.go` | Manager: Create/Get/GetOrCreate/List/Remove/Cleanup/Close, Start (container recovery), AddPool/RemovePool/Pools, Heartbeat, SetGRPCPort, SetWorkspaceManager, ImageExists/PullImage/EnsureImage | DONE |
| `box.go` | Box: Exec, Stream, Attach, Workspace, VNC, Proxy, Start/Stop/Remove, Info, touch/lastActiveTime/idleTimeout/maxLifetime/stopTimeout | DONE |
| `types.go` | LifecyclePolicy (OneShot/Session/LongRunning/Persistent), Pool, PoolInfo, PortMapping, CreateOptions (with WorkspaceID/MountMode/MountPath), ListOptions, ExecOption/ExecResult/ExecStream, AttachOption/ServiceConn, ImagePullOptions/RegistryAuth, BoxInfo, DefaultStopTimeout | DONE |
| `config.go` | Config struct | DONE |
| `errors.go` | ErrNotAvailable, ErrNotFound, ErrLimitExceeded, ErrPoolNotFound, ErrPoolInUse | DONE |
| `grpc.go` | CreateContainerTokens, RevokeContainerTokens, BuildGRPCEnv | DONE |
### workspace Module — DONE
| File | What | Status |
|------|------|--------|
| `workspace.go` | Workspace struct, CreateOptions, ListOptions, UpdateOptions, NodeInfo, MountMode, metadata marshal/unmarshal | DONE |
| `manager.go` | Manager: Create/Get/List/Update/Delete, ReadFile/WriteFile/ListDir/Remove/FS, Nodes/AddPool/RemovePool, NodeForWorkspace/MountPath | DONE |
| `errors.go` | ErrNotFound, ErrNodeMissing, ErrNodeOffline, ErrHasMounts | DONE |
### Tests — DONE
| File | Coverage | Status |
|------|----------|--------|
| **sandbox/v2** | | |
| `sandbox_test.go` | Init, M, singleton | DONE |
| `manager_test.go` | Create, Get, GetOrCreate, List, Remove, pool limits (MaxTotal, MaxPerUser), multi-pool | DONE |
| `manager_lifecycle_test.go` | Start (recovery), Cleanup, idle tracking, Heartbeat | DONE |
| `box_test.go` | Exec, Info, Workspace (ReadFile/WriteFile), status | DONE |
| `box_attach_test.go` | Attach WS echo, Attach SSE events, VNC URL, VNC Connect (RFB handshake) | DONE |
| `box_workspace_test.go` | Workspace mount, file I/O through Box, invalid ID | DONE |
| `box_image_test.go` | ImageExists (Docker+K8s), PullImage (progress+K8s no-op), EnsureImage, bad ref | DONE |
| `grpc_test.go` | Token creation/revocation, env var building (local vs remote) | DONE |
| `bench_test.go` | ContainerLifecycle, Create, Exec, ExecHeavy, Remove, Info, StopStart, WorkspaceReadWrite | DONE |
| `testutils_test.go` | testPools (local/remote/k8s), setupManager, createTestBox, ensureTestImage | DONE |
| `export_test.go` | ResetForTest | DONE |
| **workspace** | | |
| `workspace_test.go` | Create (auto/explicit ID, labels, invalid node), Get, List (owner/node filter), Update (name/labels), Delete, Nodes, NodeForWorkspace, AddPool/RemovePool, MountPath | DONE |
| `fileio_test.go` | ReadWriteFile, nested paths, ListDir, Remove, fs.FS (ReadFile, WriteFile, MkdirAll, Rename, WalkDir, Remove, NotFound) | DONE |
| `bench_test.go` | WriteFile, ReadFile, ReadWriteCycle, WriteLargeFile, ListDir, FSWalkDir, CreateDelete | DONE |
| `testutils_test.go` | testPools, setupManagerForPool, clientForPool, localClient, setupManagerMultiNode, createWorkspace | DONE |
### CI — DONE
| Job | Contents | Status |
|-----|----------|--------|
| `SandboxV2Test` | Consolidated: image pre-pull → tai-test → sandbox/v2 (local+remote+k8s) → workspace (local+remote) | DONE |
| `BenchmarkSandboxV2` | Parallel: performance tests for sandbox/v2 + workspace | DONE |
| `GRPCTest` | Independent: gRPC tests (unchanged) | DONE |
### Performance Optimizations — DONE
| Optimization | Before | After | Impact |
|-------------|--------|-------|--------|
| Remove redundant Stop in Manager.Remove() | 2.14s | 177ms | 12x faster Docker remove |
| Container CMD trap SIGTERM | 2s+ stop | near-instant | Graceful shutdown on Stop |
| K8s Start: respect ctx deadline | 30s hardcoded | ctx-aware + 60s default | Proper timeout propagation |
| K8s Pod spec: Args vs Command | CMD overridden | ENTRYPOINT preserved | Correct container behavior |
---
## Phase 2: JSAPI + OAuth — PENDING
| Task | Package | Detail |
|------|---------|--------|
| `jsapi/sandbox.go` | `sandbox/v2/jsapi` | V8 JSAPI `Sandbox()` + `Workspace()` constructors (registered in gou runtime) |
| Wire `openapi/oauth` | `sandbox/v2/grpc.go` | `CreateContainerTokens` currently generates random strings; `RevokeContainerTokens` is a no-op. Replace with real `openapi/oauth` issue/revoke calls |
| `cmd/start.go` integration | `yao` | Call `sandbox.Init(config.Conf.Sandbox)` + `sandbox.M().Start(ctx)` in startup sequence |
| Heartbeat bridge | `yao/grpc` | Wire gRPC Heartbeat handler → `sandbox.M().Heartbeat()` |
### JSAPI (planned)
```javascript
// Sandbox
var sb = Sandbox("my-workspace", {
image: "yaoapp/workspace:latest",
owner: "user-123"
})
sb.Exec(["go", "build", "./..."])
sb.ReadFile("src/main.go")
sb.WriteFile("src/main.go", "package main\n...")
sb.Stream(["npm", "run", "dev"], function(chunk) { ... })
var conn = sb.Attach(3000, { protocol: "ws", path: "/ws" })
sb.Info()
sb.Stop()
sb.Start()
sb.Remove()
// Workspace
var ws = Workspace("my-workspace")
ws.ReadFile("src/main.go")
ws.WriteFile("src/main.go", "package main\n...")
ws.ListDir("src/")
ws.Remove("tmp.txt")
```
---
## Phase 3: Agent Integration — PENDING
| Task | Detail |
|------|--------|
| Agent creates Box via `sandbox.M().GetOrCreate()` | Replace `infraSandbox.Manager` |
| Agent uses `Box.Workspace()` for file I/O | Replace Docker Copy/bind mount reads |
| Agent uses `Box.Exec()` for commands | Replace Docker exec |
| Agent uses `Box.VNC()` / `Box.Proxy()` | Replace vncproxy |
| Agent injects `Box` as `SandboxExecutor` | `ctx.sandbox` JSAPI unchanged for hooks |
---
## Phase 4: Cutover — PENDING
| Task | Detail |
|------|--------|
| Move `sandbox/v2` → `sandbox` | Rename package |
| Delete old sandbox code | manager.go, ipc/, bridge/, vncproxy/, docker/ |
| Delete `DESIGN-REMOTE.md` | Superseded by tai.Client |
| Update `cmd/start.go` | Use new init path |
| `sandbox/process.go` | Register `sandbox.*` process namespace (post-cutover) |
| `workspace/process.go` | Register `workspace.*` process namespace (post-cutover) |
---
## Implementation Details
### Container CMD
All V2 containers use a SIGTERM-aware sleep as PID 1:
```bash
sh -c "trap 'exit 0' TERM; while :; do sleep 86400 & wait $!; done"
```
This ensures:
- Container stays alive indefinitely (no hardcoded `sleep infinity`)
- Exits immediately on SIGTERM (no 2s wait)
- Works on both Docker and K8s
### Container Labels
Manager injects these labels at creation time:
```
managed-by=yao-sandbox
sandbox-id=<id>
sandbox-owner=<owner>
sandbox-pool=<pool>
sandbox-policy=<policy>
workspace-id=<workspace-id> (if WorkspaceID set)
```
Used by `Manager.Start()` to discover and recover existing containers after restart.
### Workspace Bind Mount
When `CreateOptions.WorkspaceID` is set:
```
1. NodeForWorkspace(wsID) → node name
2. Force pool = node name
3. MountPath(wsID) → hostDir
4. Bind: hostDir:/workspace:rw
```
### Multi-Mode Testing
`testPools()` returns all available pool configurations:
```go
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})
}
if host := os.Getenv("TAI_TEST_K8S_HOST"); host != "" {
// ... K8s pool with kubeconfig, namespace, ports
pools = append(pools, poolConfig{Name: "k8s", ...})
}
return pools
}
```
Every test iterates over all available pools:
```go
func TestSomething(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
// test logic
})
}
}
```
### Benchmark Helpers
```go
func setupManagerForBench(b *testing.B, pc poolConfig) *sandbox.Manager
func ensureTestImageBench(b *testing.B, m *sandbox.Manager, pool string)
func createBoxForBench(b *testing.B, m *sandbox.Manager) *sandbox.Box
```
K8s-specific behavior:
- `BenchmarkStopStart`: skipped (K8s Stop deletes Pod)
- Create/Lifecycle benchmarks: 120s timeout for K8s Pod scheduling
---
## File Inventory
### sandbox/v2 (7 source + 10 test = 17 files)
| File | Lines | Purpose |
|------|-------|---------|
| `sandbox.go` | ~25 | Global singleton |
| `manager.go` | ~620 | Manager implementation |
| `box.go` | ~230 | Box implementation |
| `types.go` | ~170 | Type definitions |
| `config.go` | ~5 | Config struct |
| `errors.go` | ~10 | Error definitions |
| `grpc.go` | ~55 | Token/env injection |
| `testutils_test.go` | ~130 | Test helpers |
| `sandbox_test.go` | ~30 | Singleton tests |
| `manager_test.go` | ~250 | CRUD tests |
| `manager_lifecycle_test.go` | ~120 | Lifecycle tests |
| `box_test.go` | ~200 | Box tests |
| `box_attach_test.go` | ~260 | Attach/VNC tests |
| `box_workspace_test.go` | ~285 | Workspace tests |
| `box_image_test.go` | ~120 | Image tests |
| `grpc_test.go` | ~80 | Token tests |
| `bench_test.go` | ~230 | Benchmarks |
### workspace (3 source + 4 test = 7 files)
| File | Lines | Purpose |
|------|-------|---------|
| `workspace.go` | ~80 | Types + metadata |
| `manager.go` | ~320 | Manager implementation |
| `errors.go` | ~10 | Error definitions |
| `testutils_test.go` | ~90 | Test helpers |
| `workspace_test.go` | ~325 | CRUD tests |
| `fileio_test.go` | ~235 | File I/O tests |
| `bench_test.go` | ~150 | Benchmarks |

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 ?= 600s
# ---------------------------------------------------------------------------
# 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), `openai-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/
```

254
sandbox/v2/bench_test.go Normal file
View file

@ -0,0 +1,254 @@
package sandbox_test
import (
"context"
"fmt"
"testing"
"time"
sandbox "github.com/yaoapp/yao/sandbox/v2"
)
// BenchmarkContainerLifecycle measures the full Create → Exec → Remove cycle.
func BenchmarkContainerLifecycle(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForBench(b, pc)
ensureTestImageBench(b, m, pc.Name)
b.ResetTimer()
for i := 0; i < b.N; i++ {
ctx := context.Background()
box, err := m.Create(ctx, sandbox.CreateOptions{
Image: testImage(),
Owner: "bench",
})
if err != nil {
b.Fatalf("Create: %v", err)
}
_, err = box.Exec(ctx, []string{"echo", "ok"})
if err != nil {
b.Fatalf("Exec: %v", err)
}
m.Remove(ctx, box.ID())
}
})
}
}
// BenchmarkCreate measures container creation time only.
func BenchmarkCreate(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForBench(b, pc)
ensureTestImageBench(b, m, pc.Name)
ids := make([]string, 0, b.N)
b.ResetTimer()
for i := 0; i < b.N; i++ {
box, err := m.Create(context.Background(), sandbox.CreateOptions{
Image: testImage(),
Owner: "bench",
})
if err != nil {
b.Fatalf("Create: %v", err)
}
ids = append(ids, box.ID())
}
b.StopTimer()
for _, id := range ids {
m.Remove(context.Background(), id)
}
})
}
}
// BenchmarkExec measures command execution latency on a pre-created container.
func BenchmarkExec(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForBench(b, pc)
box := createBoxForBench(b, m)
b.ResetTimer()
for i := 0; i < b.N; i++ {
result, err := box.Exec(context.Background(), []string{"echo", "bench"})
if err != nil {
b.Fatalf("Exec: %v", err)
}
if result.ExitCode != 0 {
b.Fatalf("exit code = %d", result.ExitCode)
}
}
})
}
}
// BenchmarkExecHeavy measures execution of a heavier command (write + read file).
func BenchmarkExecHeavy(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForBench(b, pc)
box := createBoxForBench(b, m)
b.ResetTimer()
for i := 0; i < b.N; i++ {
cmd := []string{"sh", "-c", fmt.Sprintf("echo bench-%d > /tmp/b.txt && cat /tmp/b.txt", i)}
result, err := box.Exec(context.Background(), cmd)
if err != nil {
b.Fatalf("Exec: %v", err)
}
if result.ExitCode != 0 {
b.Fatalf("exit code = %d", result.ExitCode)
}
}
})
}
}
// BenchmarkRemove measures container removal time.
func BenchmarkRemove(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForBench(b, pc)
ensureTestImageBench(b, m, pc.Name)
boxes := make([]*sandbox.Box, b.N)
for i := 0; i < b.N; i++ {
box, err := m.Create(context.Background(), sandbox.CreateOptions{
Image: testImage(),
Owner: "bench",
})
if err != nil {
b.Fatalf("Create: %v", err)
}
boxes[i] = box
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
if err := m.Remove(context.Background(), boxes[i].ID()); err != nil {
b.Fatalf("Remove: %v", err)
}
}
})
}
}
// BenchmarkInfo measures Info() latency on a running container.
func BenchmarkInfo(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForBench(b, pc)
box := createBoxForBench(b, m)
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, err := box.Info(context.Background())
if err != nil {
b.Fatalf("Info: %v", err)
}
}
})
}
}
// BenchmarkStopStart measures Stop → Start cycle time.
func BenchmarkStopStart(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
if pc.Name == "k8s" {
b.Skip("K8s Stop deletes Pod; Stop→Start cycle not applicable")
}
m := setupManagerForBench(b, pc)
box := createBoxForBench(b, m)
b.ResetTimer()
for i := 0; i < b.N; i++ {
if err := box.Stop(context.Background()); err != nil {
b.Fatalf("Stop: %v", err)
}
if err := box.Start(context.Background()); err != nil {
b.Fatalf("Start: %v", err)
}
}
})
}
}
// BenchmarkWorkspaceReadWrite measures workspace file read/write via container Box.
func BenchmarkWorkspaceReadWrite(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForBench(b, pc)
box := createBoxForBench(b, m)
ws := box.Workspace()
if ws == nil {
b.Skip("workspace not available")
}
payload := []byte("package main\nfunc main() { println(\"hello\") }\n")
b.ResetTimer()
for i := 0; i < b.N; i++ {
name := fmt.Sprintf("f%d.go", i)
if err := ws.WriteFile(name, payload, 0644); err != nil {
b.Fatalf("WriteFile: %v", err)
}
data, err := ws.ReadFile(name)
if err != nil {
b.Fatalf("ReadFile: %v", err)
}
if len(data) != len(payload) {
b.Fatalf("size mismatch: %d vs %d", len(data), len(payload))
}
}
})
}
}
// --- helpers ---
func setupManagerForBench(b *testing.B, pc poolConfig) *sandbox.Manager {
b.Helper()
pool := sandbox.Pool{Name: pc.Name, Addr: pc.Addr, Options: pc.Options}
cfg := sandbox.Config{Pool: []sandbox.Pool{pool}}
if err := sandbox.Init(cfg); err != nil {
b.Fatalf("Init: %v", err)
}
m := sandbox.M()
b.Cleanup(func() { m.Close() })
return m
}
func ensureTestImageBench(b *testing.B, m *sandbox.Manager, pool string) {
b.Helper()
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
defer cancel()
if err := m.EnsureImage(ctx, pool, testImage(), sandbox.ImagePullOptions{}); err != nil {
b.Fatalf("EnsureImage: %v", err)
}
}
func createBoxForBench(b *testing.B, m *sandbox.Manager) *sandbox.Box {
b.Helper()
pools := m.Pools()
if len(pools) > 0 {
ensureTestImageBench(b, m, pools[0].Name)
}
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
defer cancel()
box, err := m.Create(ctx, sandbox.CreateOptions{
Image: testImage(),
Owner: "bench",
})
if err != nil {
b.Fatalf("Create: %v", err)
}
b.Cleanup(func() { m.Remove(context.Background(), box.ID()) })
return box
}

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

@ -0,0 +1,283 @@
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
stopTimeoutD time.Duration
createdAt time.Time
refreshToken string
vnc bool
image string
workspaceID 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.
// If a workspace is mounted (WorkspaceID set), uses the workspace ID as session;
// otherwise falls back to the sandbox ID (backward compatible).
func (b *Box) Workspace() workspace.FS {
b.touch()
if b.ws != nil {
return b.ws
}
sessionID := b.workspaceID
if sessionID == "" {
sessionID = b.id
}
client, err := b.manager.getPool(b.pool)
if err != nil {
return nil
}
b.ws = client.Workspace(sessionID)
return b.ws
}
// WorkspaceID returns the workspace ID mounted to this sandbox, or empty string.
func (b *Box) WorkspaceID() string { return b.workspaceID }
// 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, b.stopTimeout())
}
// 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
}
func (b *Box) stopTimeout() time.Duration {
if b.stopTimeoutD > 0 {
return b.stopTimeoutD
}
pd := b.manager.findPoolDef(b.pool)
if pd != nil && pd.StopTimeout > 0 {
return pd.StopTimeout
}
return DefaultStopTimeout
}

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

@ -0,0 +1,280 @@
package sandbox_test
import (
"context"
"net"
"net/http"
"strings"
"testing"
"time"
"github.com/gorilla/websocket"
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()
proxyURL, err := box.Proxy(ctx, port, "/")
if err != nil {
t.Fatalf("Proxy URL: %v", err)
}
host := proxyURL[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(500 * 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, 2*time.Second)
if err != nil {
continue
}
conn.Close()
// TCP reachable — give the service process time to accept
// application-layer connections (Python ws/sse servers in CI
// may take 1-3s after the port opens before they're ready).
time.Sleep(2 * time.Second)
return
}
}
}
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)
var conn *sandbox.ServiceConn
var err error
for attempt := 0; attempt < 5; attempt++ {
conn, err = box.Attach(t.Context(), 9800, sandbox.WithProtocol("ws"), sandbox.WithPath("/"))
if err == nil {
break
}
time.Sleep(time.Duration(attempt+1) * 500 * time.Millisecond)
}
if err != nil {
t.Fatalf("Attach WS after retries: %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)
var conn *sandbox.ServiceConn
var err error
for attempt := 0; attempt < 5; attempt++ {
conn, err = box.Attach(t.Context(), 9801, sandbox.WithProtocol("sse"), sandbox.WithPath("/events"))
if err == nil {
break
}
time.Sleep(time.Duration(attempt+1) * 500 * time.Millisecond)
}
if err != nil {
t.Fatalf("Attach SSE after retries: %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)
}
})
}
}
func TestVNCURL(t *testing.T) {
skipIfNoDocker(t)
img := testImage()
if img == "alpine:latest" {
t.Skip("VNC test requires sandbox-v2-test image with VNC desktop")
}
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.VNC = true
})
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
url, err := box.VNC(ctx)
if err != nil {
t.Fatalf("VNC URL: %v", err)
}
if !strings.HasPrefix(url, "ws://") {
t.Fatalf("VNC URL = %q, want ws:// prefix", url)
}
t.Logf("VNC URL: %s", url)
})
}
}
func TestVNCConnect(t *testing.T) {
skipIfNoDocker(t)
img := testImage()
if img == "alpine:latest" {
t.Skip("VNC test requires sandbox-v2-test image with VNC desktop")
}
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.VNC = true
})
ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
defer cancel()
vncURL, err := box.VNC(ctx)
if err != nil {
t.Fatalf("VNC URL: %v", err)
}
t.Logf("VNC URL: %s", vncURL)
waitForWSEndpoint(t, vncURL, 30*time.Second)
dialer := websocket.Dialer{
Subprotocols: []string{"binary"},
HandshakeTimeout: 10 * time.Second,
}
var ws *websocket.Conn
for attempt := 0; attempt < 5; attempt++ {
var resp *http.Response
ws, resp, err = dialer.DialContext(ctx, vncURL, http.Header{})
if err == nil {
break
}
if resp != nil {
resp.Body.Close()
}
time.Sleep(time.Duration(attempt+1) * time.Second)
}
if err != nil {
t.Fatalf("VNC dial after retries: %v", err)
}
defer ws.Close()
ws.SetReadDeadline(time.Now().Add(10 * time.Second))
_, msg, err := ws.ReadMessage()
if err != nil {
t.Fatalf("VNC read: %v", err)
}
if !strings.HasPrefix(string(msg), "RFB ") {
t.Fatalf("VNC banner = %q, want RFB prefix", string(msg))
}
t.Logf("VNC banner: %s", strings.TrimSpace(string(msg)))
})
}
}
func waitForWSEndpoint(t *testing.T, wsURL string, timeout time.Duration) {
t.Helper()
httpURL := "http" + strings.TrimPrefix(wsURL, "ws")
if idx := strings.LastIndex(httpURL, "/ws"); idx > 0 {
httpURL = httpURL[:idx]
}
host := strings.TrimPrefix(httpURL, "http://")
if i := strings.Index(host, "/"); i > 0 {
host = host[:i]
}
deadline := time.After(timeout)
ticker := time.NewTicker(500 * time.Millisecond)
defer ticker.Stop()
for {
select {
case <-deadline:
t.Fatalf("VNC endpoint %s not ready within %v", host, timeout)
case <-ticker.C:
conn, err := net.DialTimeout("tcp", host, 2*time.Second)
if err != nil {
continue
}
conn.Close()
// VNC services (Xvfb → fluxbox → x11vnc → websockify) need time
// after the TCP port is reachable. Give the process chain time to
// stabilize before attempting the WebSocket handshake.
time.Sleep(2 * time.Second)
return
}
}
}

121
sandbox/v2/box_image_test.go Normal file
View file

@ -0,0 +1,121 @@
package sandbox_test
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
sandbox "github.com/yaoapp/yao/sandbox/v2"
)
func TestImageExists(t *testing.T) {
for _, pc := range testPools() {
pc := pc
t.Run(pc.Name, func(t *testing.T) {
if pc.Name == "k8s" {
t.Run("always_true", func(t *testing.T) {
m := setupManagerForPool(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
exists, err := m.ImageExists(ctx, pc.Name, "anything:nonexistent")
require.NoError(t, err)
assert.True(t, exists, "k8s mode should always return true")
})
return
}
m := setupManagerForPool(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
t.Run("existing", func(t *testing.T) {
exists, err := m.ImageExists(ctx, pc.Name, "alpine:latest")
require.NoError(t, err)
assert.True(t, exists)
})
t.Run("missing", func(t *testing.T) {
exists, err := m.ImageExists(ctx, pc.Name, "nonexistent/image:no-such-tag-ever-12345")
require.NoError(t, err)
assert.False(t, exists)
})
})
}
}
func TestImagePull(t *testing.T) {
for _, pc := range testPools() {
pc := pc
t.Run(pc.Name, func(t *testing.T) {
if pc.Name == "k8s" {
t.Run("noop", func(t *testing.T) {
m := setupManagerForPool(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
ch, err := m.PullImage(ctx, pc.Name, "alpine:latest", sandbox.ImagePullOptions{})
require.NoError(t, err)
assert.Nil(t, ch, "k8s mode should return nil channel")
})
return
}
m := setupManagerForPool(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
defer cancel()
t.Run("pull_with_progress", func(t *testing.T) {
ch, err := m.PullImage(ctx, pc.Name, "alpine:latest", sandbox.ImagePullOptions{})
require.NoError(t, err)
require.NotNil(t, ch)
var count int
for p := range ch {
if p.Error != "" {
t.Fatalf("pull error: %s", p.Error)
}
count++
}
assert.Greater(t, count, 0, "should receive at least one progress event")
})
})
}
}
func TestEnsureImage(t *testing.T) {
for _, pc := range testPools() {
pc := pc
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
defer cancel()
err := m.EnsureImage(ctx, pc.Name, "alpine:latest", sandbox.ImagePullOptions{})
require.NoError(t, err)
if pc.Name != "k8s" {
exists, err := m.ImageExists(ctx, pc.Name, "alpine:latest")
require.NoError(t, err)
assert.True(t, exists)
}
})
}
}
func TestEnsureImage_BadRef(t *testing.T) {
for _, pc := range testPools() {
pc := pc
if pc.Name == "k8s" {
continue
}
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
err := m.EnsureImage(ctx, pc.Name, "nonexistent/image:no-such-tag-ever-12345", sandbox.ImagePullOptions{})
assert.Error(t, err)
})
}
}

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

@ -0,0 +1,215 @@
package sandbox_test
import (
"context"
"io"
"strings"
"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 s := strings.ToLower(info.Status); s != "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")
}
})
}
}

284
sandbox/v2/box_workspace_test.go Normal file
View file

@ -0,0 +1,284 @@
package sandbox_test
import (
"context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
sandbox "github.com/yaoapp/yao/sandbox/v2"
"github.com/yaoapp/yao/workspace"
)
func TestWorkspaceID_Set(t *testing.T) {
skipIfNoDocker(t)
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
sbm, wsm := setupManagerWithWorkspace(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
ws, err := wsm.Create(ctx, workspace.CreateOptions{
Name: "test-ws", Owner: "user", Node: pc.Name,
})
require.NoError(t, err)
defer wsm.Delete(context.Background(), ws.ID, true)
box := createTestBox(t, sbm, func(co *sandbox.CreateOptions) {
co.WorkspaceID = ws.ID
})
assert.Equal(t, ws.ID, box.WorkspaceID())
})
}
}
func TestWorkspaceID_Empty(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)
assert.Empty(t, box.WorkspaceID())
})
}
}
func TestWorkspace_NodeRouting(t *testing.T) {
skipIfNoDocker(t)
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
sbm, wsm := setupManagerWithWorkspace(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
ws, err := wsm.Create(ctx, workspace.CreateOptions{
Name: "routed-ws", Owner: "user", Node: pc.Name,
})
require.NoError(t, err)
defer wsm.Delete(context.Background(), ws.ID, true)
box := createTestBox(t, sbm, func(co *sandbox.CreateOptions) {
co.WorkspaceID = ws.ID
})
assert.Equal(t, pc.Name, box.Pool())
})
}
}
func TestWorkspace_InvalidID(t *testing.T) {
skipIfNoDocker(t)
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
sbm, _ := setupManagerWithWorkspace(t, pc)
ensureTestImage(t, sbm, pc.Name)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
_, err := sbm.Create(ctx, sandbox.CreateOptions{
Image: testImage(),
Owner: "user",
WorkspaceID: "nonexistent-workspace",
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "resolve workspace")
})
}
}
func TestWorkspace_BindMountLocal(t *testing.T) {
skipIfNoDocker(t)
pc := poolConfig{Name: "local", Addr: testLocalAddr()}
sbm, wsm := setupManagerWithWorkspace(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
ws, err := wsm.Create(ctx, workspace.CreateOptions{
Name: "mount-ws", Owner: "user", Node: pc.Name,
})
require.NoError(t, err)
defer wsm.Delete(context.Background(), ws.ID, true)
require.NoError(t, wsm.WriteFile(ctx, ws.ID, "seed.txt", []byte("hello from workspace"), 0644))
box := createTestBox(t, sbm, func(co *sandbox.CreateOptions) {
co.WorkspaceID = ws.ID
})
result, err := box.Exec(ctx, []string{"cat", "/workspace/seed.txt"})
require.NoError(t, err)
assert.Equal(t, "hello from workspace", result.Stdout)
}
func TestWorkspace_ContainerWriteBack(t *testing.T) {
skipIfNoDocker(t)
pc := poolConfig{Name: "local", Addr: testLocalAddr()}
sbm, wsm := setupManagerWithWorkspace(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
ws, err := wsm.Create(ctx, workspace.CreateOptions{
Name: "writeback-ws", Owner: "user", Node: pc.Name,
})
require.NoError(t, err)
defer wsm.Delete(context.Background(), ws.ID, true)
box := createTestBox(t, sbm, func(co *sandbox.CreateOptions) {
co.WorkspaceID = ws.ID
})
_, err = box.Exec(ctx, []string{"sh", "-c", "echo 'from container' > /workspace/output.txt"})
require.NoError(t, err)
data, err := wsm.ReadFile(ctx, ws.ID, "output.txt")
require.NoError(t, err)
assert.Equal(t, "from container\n", string(data))
}
func TestWorkspace_ReadOnlyMount(t *testing.T) {
skipIfNoDocker(t)
pc := poolConfig{Name: "local", Addr: testLocalAddr()}
sbm, wsm := setupManagerWithWorkspace(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
ws, err := wsm.Create(ctx, workspace.CreateOptions{
Name: "ro-ws", Owner: "user", Node: pc.Name,
})
require.NoError(t, err)
defer wsm.Delete(context.Background(), ws.ID, true)
require.NoError(t, wsm.WriteFile(ctx, ws.ID, "readonly.txt", []byte("immutable"), 0644))
box := createTestBox(t, sbm, func(co *sandbox.CreateOptions) {
co.WorkspaceID = ws.ID
co.MountMode = "ro"
})
result, err := box.Exec(ctx, []string{"cat", "/workspace/readonly.txt"})
require.NoError(t, err)
assert.Equal(t, "immutable", result.Stdout)
result, err = box.Exec(ctx, []string{"sh", "-c", "echo fail > /workspace/nope.txt 2>&1; echo $?"})
require.NoError(t, err)
// Write to read-only mount should fail (non-zero exit or error message)
assert.True(t, result.ExitCode != 0 || result.Stdout != "0\n" || len(result.Stderr) > 0,
"expected write to read-only mount to fail")
}
func TestWorkspace_CustomMountPath(t *testing.T) {
skipIfNoDocker(t)
pc := poolConfig{Name: "local", Addr: testLocalAddr()}
sbm, wsm := setupManagerWithWorkspace(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
ws, err := wsm.Create(ctx, workspace.CreateOptions{
Name: "custom-path-ws", Owner: "user", Node: pc.Name,
})
require.NoError(t, err)
defer wsm.Delete(context.Background(), ws.ID, true)
require.NoError(t, wsm.WriteFile(ctx, ws.ID, "data.json", []byte(`{"ok":true}`), 0644))
box := createTestBox(t, sbm, func(co *sandbox.CreateOptions) {
co.WorkspaceID = ws.ID
co.MountPath = "/data"
})
result, err := box.Exec(ctx, []string{"cat", "/data/data.json"})
require.NoError(t, err)
assert.Equal(t, `{"ok":true}`, result.Stdout)
}
func TestWorkspace_BoxWorkspaceFS(t *testing.T) {
skipIfNoDocker(t)
for _, pc := range testPools() {
if pc.Name == "local" {
// Local mode: sandbox and workspace use separate tai.Clients with
// different dataDirs, so Box.Workspace() writes to the sandbox volume
// while wsm reads from the workspace volume. Bind mount tests cover
// local workspace I/O end-to-end instead.
continue
}
t.Run(pc.Name, func(t *testing.T) {
sbm, wsm := setupManagerWithWorkspace(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
ws, err := wsm.Create(ctx, workspace.CreateOptions{
Name: "fs-ws", Owner: "user", Node: pc.Name,
})
require.NoError(t, err)
defer wsm.Delete(context.Background(), ws.ID, true)
box := createTestBox(t, sbm, func(co *sandbox.CreateOptions) {
co.WorkspaceID = ws.ID
})
wfs := box.Workspace()
if wfs == nil {
t.Skip("Workspace FS not available")
}
require.NoError(t, wfs.WriteFile("via-box.txt", []byte("box wrote this"), 0644))
data, err := wsm.ReadFile(ctx, ws.ID, "via-box.txt")
require.NoError(t, err)
assert.Equal(t, "box wrote this", string(data))
})
}
}
func TestWorkspace_LabelPersistence(t *testing.T) {
skipIfNoDocker(t)
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
sbm, wsm := setupManagerWithWorkspace(t, pc)
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
ws, err := wsm.Create(ctx, workspace.CreateOptions{
Name: "label-ws", Owner: "user", Node: pc.Name,
})
require.NoError(t, err)
defer wsm.Delete(context.Background(), ws.ID, true)
box := createTestBox(t, sbm, func(co *sandbox.CreateOptions) {
co.WorkspaceID = ws.ID
})
// WorkspaceID getter should reflect what was set
assert.Equal(t, ws.ID, box.WorkspaceID())
// Container should also carry the label (verify via exec reading env or
// just trust that buildTaiCreateOptions sets it — the label is tested
// indirectly by TestWorkspace_NodeRouting which relies on correct routing)
info, err := box.Info(ctx)
require.NoError(t, err)
assert.Contains(t, []string{"running", "Running"}, info.Status)
})
}
}

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

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

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

@ -0,0 +1,39 @@
# 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 \
tini \
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
# openai-proxy: Anthropic Messages API → OpenAI Chat Completions API
COPY openai-proxy-${TARGETARCH} /usr/local/bin/openai-proxy
RUN chmod +x /usr/local/bin/openai-proxy
COPY entrypoint.sh /entrypoint.sh
RUN chmod +x /entrypoint.sh
USER sandbox
ENTRYPOINT ["/usr/bin/tini", "--", "/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 openai-proxy
if [ -n "$YAO_GRPC_ADDR" ] && [ -n "$YAO_SANDBOX_ID" ]; then
tail -f /dev/null | yao-grpc serve &
fi
if [ -n "$OPENAI_PROXY_BACKEND" ]; then
openai-proxy &
fi
exec "$@"

7
sandbox/v2/docker/bin/openai-proxy/cmd/openai-proxy/main.go Normal file
View file

@ -0,0 +1,7 @@
package main
import proxy "github.com/yaoapp/yao/sandbox/v2/docker/bin/openai-proxy"
func main() {
proxy.Main()
}

419
sandbox/v2/docker/bin/openai-proxy/convert.go Normal file
View file

@ -0,0 +1,419 @@
package proxy
import (
"encoding/json"
"fmt"
"strings"
)
func (s *Server) convertRequest(req *AnthropicRequest) *OpenAIRequest {
maxTokens := req.MaxTokens
if s.config.Options != nil {
if mt, ok := s.config.Options["max_tokens"]; ok {
switch v := mt.(type) {
case float64:
maxTokens = int(v)
case int:
maxTokens = v
}
}
}
temperature := req.Temperature
if s.config.Options != nil {
if temp, ok := s.config.Options["temperature"]; ok {
if v, ok := temp.(float64); ok {
temperature = &v
}
}
}
openaiReq := &OpenAIRequest{
Model: s.config.Model,
MaxTokens: maxTokens,
Stream: req.Stream,
Temperature: temperature,
TopP: req.TopP,
Stop: req.StopSequences,
}
if s.config.Options != nil {
openaiReq.ExtraOptions = make(map[string]interface{})
for k, v := range s.config.Options {
switch k {
case "max_tokens", "temperature", "model", "key", "proxy":
continue
default:
openaiReq.ExtraOptions[k] = v
}
}
}
openaiReq.Messages = s.convertMessages(req.Messages, req.System)
if len(req.Tools) > 0 {
openaiReq.Tools = s.convertTools(req.Tools)
}
if req.ToolChoice != nil {
openaiReq.ToolChoice = s.convertToolChoice(req.ToolChoice)
}
return openaiReq
}
func (s *Server) convertMessages(msgs []AnthropicMsg, system interface{}) []OpenAIMsg {
var result []OpenAIMsg
if system != nil {
systemText := extractSystemText(system)
if systemText != "" {
result = append(result, OpenAIMsg{
Role: "system",
Content: systemText,
})
}
}
for _, msg := range msgs {
converted := s.convertMessage(msg)
result = append(result, converted...)
}
return result
}
func (s *Server) convertMessage(msg AnthropicMsg) []OpenAIMsg {
var result []OpenAIMsg
switch content := msg.Content.(type) {
case string:
result = append(result, OpenAIMsg{
Role: mapRole(msg.Role),
Content: content,
})
case []interface{}:
var toolResults []ContentBlock
var otherContent []interface{}
for _, item := range content {
block := parseContentBlock(item)
if block.Type == "tool_result" {
toolResults = append(toolResults, block)
} else {
otherContent = append(otherContent, item)
}
}
for _, tr := range toolResults {
toolMsg := OpenAIMsg{
Role: "tool",
ToolCallID: tr.ToolUseID,
Content: extractToolResultContent(tr.Content),
}
result = append(result, toolMsg)
}
if len(otherContent) > 0 {
openaiContent := s.convertContentBlocks(otherContent)
if len(openaiContent) == 1 && openaiContent[0].Type == "text" {
result = append(result, OpenAIMsg{
Role: mapRole(msg.Role),
Content: openaiContent[0].Text,
})
} else if len(openaiContent) > 0 {
result = append(result, OpenAIMsg{
Role: mapRole(msg.Role),
Content: openaiContent,
})
}
}
if msg.Role == "assistant" {
toolCalls := extractToolUseBlocks(content)
if len(toolCalls) > 0 {
found := false
for i := range result {
if result[i].Role == "assistant" {
result[i].ToolCalls = toolCalls
found = true
break
}
}
if !found {
result = append(result, OpenAIMsg{
Role: "assistant",
Content: "",
ToolCalls: toolCalls,
})
}
}
}
}
return result
}
func (s *Server) convertContentBlocks(blocks []interface{}) []OpenAIContent {
var result []OpenAIContent
for _, item := range blocks {
block := parseContentBlock(item)
switch block.Type {
case "text":
result = append(result, OpenAIContent{
Type: "text",
Text: block.Text,
})
case "image":
if block.Source != nil {
imageURL := convertImageSource(block.Source)
result = append(result, OpenAIContent{
Type: "image_url",
ImageURL: imageURL,
})
}
case "tool_use", "tool_result":
continue
}
}
return result
}
func convertImageSource(source *ImageSource) *OpenAIImageURL {
if source == nil {
return nil
}
switch source.Type {
case "base64":
mediaType := source.MediaType
if mediaType == "" {
mediaType = "image/jpeg"
}
return &OpenAIImageURL{
URL: fmt.Sprintf("data:%s;base64,%s", mediaType, source.Data),
}
case "url":
return &OpenAIImageURL{
URL: source.URL,
}
}
return nil
}
func (s *Server) convertTools(tools []AnthropicTool) []OpenAITool {
var result []OpenAITool
for _, tool := range tools {
result = append(result, OpenAITool{
Type: "function",
Function: OpenAIFunction{
Name: tool.Name,
Description: tool.Description,
Parameters: tool.InputSchema,
},
})
}
return result
}
func (s *Server) convertToolChoice(choice *AnthropicToolChoice) interface{} {
if choice == nil {
return nil
}
switch choice.Type {
case "auto":
return "auto"
case "any":
return "required"
case "tool":
return map[string]interface{}{
"type": "function",
"function": map[string]string{
"name": choice.Name,
},
}
case "none":
return "none"
}
return "auto"
}
func (s *Server) convertResponse(resp *OpenAIResponse) *AnthropicResponse {
result := &AnthropicResponse{
ID: generateID("msg_"),
Type: "message",
Role: "assistant",
Content: []ContentBlock{},
Model: s.config.Model,
}
if len(resp.Choices) > 0 {
choice := resp.Choices[0]
if content, ok := choice.Message.Content.(string); ok && content != "" {
result.Content = append(result.Content, ContentBlock{
Type: "text",
Text: content,
})
}
for _, tc := range choice.Message.ToolCalls {
var input interface{}
json.Unmarshal([]byte(tc.Function.Arguments), &input)
result.Content = append(result.Content, ContentBlock{
Type: "tool_use",
ID: tc.ID,
Name: tc.Function.Name,
Input: input,
})
}
stopReason := mapFinishReason(choice.FinishReason)
result.StopReason = &stopReason
}
if resp.Usage != nil {
result.Usage = &Usage{
InputTokens: resp.Usage.PromptTokens,
OutputTokens: resp.Usage.CompletionTokens,
}
} else {
result.Usage = &Usage{InputTokens: 0, OutputTokens: 0}
}
return result
}
func extractSystemText(system interface{}) string {
switch s := system.(type) {
case string:
return s
case []interface{}:
var texts []string
for _, item := range s {
if block, ok := item.(map[string]interface{}); ok {
if text, ok := block["text"].(string); ok {
if strings.HasPrefix(text, "x-anthropic-") {
continue
}
texts = append(texts, text)
}
}
}
if len(texts) > 0 {
return strings.Join(texts, "\n\n")
}
}
return ""
}
func parseContentBlock(item interface{}) ContentBlock {
var block ContentBlock
switch v := item.(type) {
case map[string]interface{}:
if t, ok := v["type"].(string); ok {
block.Type = t
}
if text, ok := v["text"].(string); ok {
block.Text = text
}
if id, ok := v["id"].(string); ok {
block.ID = id
}
if name, ok := v["name"].(string); ok {
block.Name = name
}
if input, ok := v["input"]; ok {
block.Input = input
}
if toolUseID, ok := v["tool_use_id"].(string); ok {
block.ToolUseID = toolUseID
}
if content, ok := v["content"]; ok {
block.Content = content
}
if isError, ok := v["is_error"].(bool); ok {
block.IsError = isError
}
if source, ok := v["source"].(map[string]interface{}); ok {
block.Source = parseImageSource(source)
}
}
return block
}
func parseImageSource(source map[string]interface{}) *ImageSource {
if source == nil {
return nil
}
result := &ImageSource{}
if t, ok := source["type"].(string); ok {
result.Type = t
}
if mediaType, ok := source["media_type"].(string); ok {
result.MediaType = mediaType
}
if data, ok := source["data"].(string); ok {
result.Data = data
}
if url, ok := source["url"].(string); ok {
result.URL = url
}
return result
}
func extractToolUseBlocks(content []interface{}) []OpenAIToolCall {
var result []OpenAIToolCall
for _, item := range content {
block := parseContentBlock(item)
if block.Type == "tool_use" {
args, _ := json.Marshal(block.Input)
result = append(result, OpenAIToolCall{
ID: block.ID,
Type: "function",
Function: OpenAIFunctionCall{
Name: block.Name,
Arguments: string(args),
},
})
}
}
return result
}
func extractToolResultContent(content interface{}) string {
switch c := content.(type) {
case string:
return c
case []interface{}:
for _, item := range c {
if block, ok := item.(map[string]interface{}); ok {
if block["type"] == "text" {
if text, ok := block["text"].(string); ok {
return text
}
}
}
}
}
return ""
}
func mapRole(role string) string {
switch role {
case "user":
return "user"
case "assistant":
return "assistant"
default:
return role
}
}

510
sandbox/v2/docker/bin/openai-proxy/main.go Normal file
View file

@ -0,0 +1,510 @@
// Package proxy provides a lightweight API proxy that translates
// Anthropic Messages API to OpenAI Chat Completions API.
package proxy
import (
"bufio"
"bytes"
"encoding/json"
"flag"
"fmt"
"io"
"log"
"net/http"
"os"
"strconv"
"strings"
"time"
)
// Config holds the proxy server configuration
type Config struct {
Port int
Backend string
Model string
APIKey string
Timeout int
Verbose bool
LogFile string
Options map[string]interface{}
}
// Server is the API proxy server
type Server struct {
config *Config
client *http.Client
}
// Main is the entry point for the proxy server
func Main() {
config := parseFlags()
if err := config.Validate(); err != nil {
log.Fatalf("Configuration error: %v", err)
}
if config.LogFile != "" {
f, err := os.OpenFile(config.LogFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
if err != nil {
log.Fatalf("Failed to open log file: %v", err)
}
mw := io.MultiWriter(os.Stdout, f)
log.SetOutput(mw)
}
server := NewServer(config)
addr := fmt.Sprintf(":%d", config.Port)
log.Printf("OpenAI Proxy starting on %s", addr)
log.Printf("Backend: %s", config.Backend)
log.Printf("Model: %s", config.Model)
if len(config.Options) > 0 {
optBytes, _ := json.Marshal(config.Options)
log.Printf("Options: %s", string(optBytes))
}
http.HandleFunc("/v1/messages", server.handleMessages)
http.HandleFunc("/health", server.handleHealth)
if err := http.ListenAndServe(addr, nil); err != nil {
log.Fatalf("Server failed: %v", err)
}
}
func parseFlags() *Config {
config := &Config{}
flag.IntVar(&config.Port, "p", 0, "Listen port")
flag.IntVar(&config.Port, "port", 0, "Listen port")
flag.StringVar(&config.Backend, "b", "", "Backend API URL")
flag.StringVar(&config.Backend, "backend", "", "Backend API URL")
flag.StringVar(&config.Model, "m", "", "Backend model name")
flag.StringVar(&config.Model, "model", "", "Backend model name")
flag.StringVar(&config.APIKey, "k", "", "Backend API key")
flag.StringVar(&config.APIKey, "api-key", "", "Backend API key")
flag.IntVar(&config.Timeout, "t", 0, "Request timeout in seconds")
flag.IntVar(&config.Timeout, "timeout", 0, "Request timeout in seconds")
flag.BoolVar(&config.Verbose, "v", false, "Verbose logging")
flag.BoolVar(&config.Verbose, "verbose", false, "Verbose logging")
flag.StringVar(&config.LogFile, "l", "", "Log file path")
flag.StringVar(&config.LogFile, "log", "", "Log file path")
flag.Parse()
if config.Port == 0 {
if v := os.Getenv("OPENAI_PROXY_PORT"); v != "" {
config.Port, _ = strconv.Atoi(v)
}
}
if config.Port == 0 {
config.Port = 3456
}
if config.Backend == "" {
config.Backend = os.Getenv("OPENAI_PROXY_BACKEND")
}
if config.Model == "" {
config.Model = os.Getenv("OPENAI_PROXY_MODEL")
}
if config.APIKey == "" {
config.APIKey = os.Getenv("OPENAI_PROXY_API_KEY")
}
if config.Timeout == 0 {
if v := os.Getenv("OPENAI_PROXY_TIMEOUT"); v != "" {
config.Timeout, _ = strconv.Atoi(v)
}
}
if config.Timeout == 0 {
config.Timeout = 300
}
if optionsStr := os.Getenv("OPENAI_PROXY_OPTIONS"); optionsStr != "" {
var options map[string]interface{}
if err := json.Unmarshal([]byte(optionsStr), &options); err != nil {
log.Printf("Warning: failed to parse OPENAI_PROXY_OPTIONS: %v", err)
} else {
config.Options = options
}
}
return config
}
// Validate checks if the configuration is valid
func (c *Config) Validate() error {
if c.Backend == "" {
return fmt.Errorf("backend URL is required (-b or OPENAI_PROXY_BACKEND)")
}
if c.Model == "" {
return fmt.Errorf("model name is required (-m or OPENAI_PROXY_MODEL)")
}
if c.APIKey == "" {
return fmt.Errorf("API key is required (-k or OPENAI_PROXY_API_KEY)")
}
return nil
}
// NewServer creates a new proxy server
func NewServer(config *Config) *Server {
return &Server{
config: config,
client: &http.Client{
Timeout: time.Duration(config.Timeout) * time.Second,
},
}
}
func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
}
func (s *Server) handleMessages(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
body, err := io.ReadAll(r.Body)
if err != nil {
s.errorResponse(w, http.StatusBadRequest, "invalid_request", "Failed to read request body")
return
}
defer r.Body.Close()
if s.config.Verbose {
log.Printf("Received request: %s", string(body))
}
var anthropicReq AnthropicRequest
if err := json.Unmarshal(body, &anthropicReq); err != nil {
s.errorResponse(w, http.StatusBadRequest, "invalid_request", "Invalid JSON")
return
}
openaiReq := s.convertRequest(&anthropicReq)
if anthropicReq.Stream {
s.handleStreamingRequest(w, openaiReq)
} else {
s.handleNonStreamingRequest(w, openaiReq)
}
}
func (s *Server) handleNonStreamingRequest(w http.ResponseWriter, openaiReq *OpenAIRequest) {
openaiReq.Stream = false
resp, err := s.forwardRequest(openaiReq)
if err != nil {
s.errorResponse(w, http.StatusBadGateway, "backend_error", err.Error())
return
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
s.errorResponse(w, http.StatusBadGateway, "backend_error", "Failed to read backend response")
return
}
if s.config.Verbose {
log.Printf("Backend response: %s", string(body))
}
if resp.StatusCode != http.StatusOK {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(resp.StatusCode)
w.Write(body)
return
}
var openaiResp OpenAIResponse
if err := json.Unmarshal(body, &openaiResp); err != nil {
s.errorResponse(w, http.StatusBadGateway, "backend_error", "Invalid backend response")
return
}
anthropicResp := s.convertResponse(&openaiResp)
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(anthropicResp)
}
func (s *Server) handleStreamingRequest(w http.ResponseWriter, openaiReq *OpenAIRequest) {
openaiReq.Stream = true
openaiReq.StreamOptions = &StreamOptions{IncludeUsage: true}
resp, err := s.forwardRequest(openaiReq)
if err != nil {
s.errorResponse(w, http.StatusBadGateway, "backend_error", err.Error())
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(resp.StatusCode)
w.Write(body)
return
}
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.Header().Set("Connection", "keep-alive")
flusher, ok := w.(http.Flusher)
if !ok {
s.errorResponse(w, http.StatusInternalServerError, "server_error", "Streaming not supported")
return
}
msgID := generateID("msg_")
startEvent := AnthropicStreamEvent{
Type: "message_start",
Message: &AnthropicResponse{
ID: msgID,
Type: "message",
Role: "assistant",
Content: []ContentBlock{},
Model: s.config.Model,
StopReason: nil,
StopSequence: nil,
Usage: &Usage{InputTokens: 0, OutputTokens: 0},
},
}
s.writeSSE(w, flusher, startEvent)
s.processStream(w, flusher, resp.Body, msgID)
}
func (s *Server) processStream(w http.ResponseWriter, flusher http.Flusher, body io.Reader, msgID string) {
scanner := bufio.NewScanner(body)
scanner.Buffer(make([]byte, 64*1024), 1024*1024)
var contentBlockStarted bool
var currentToolCall *ToolCallAccumulator
var toolCalls []*ToolCallAccumulator
var contentIndex int
var finishReason string
var lastUsage *Usage
for scanner.Scan() {
line := scanner.Text()
if !strings.HasPrefix(line, "data: ") {
continue
}
data := strings.TrimPrefix(line, "data: ")
if data == "[DONE]" {
break
}
var chunk OpenAIStreamChunk
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
if s.config.Verbose {
log.Printf("Failed to parse chunk: %s", data)
}
continue
}
if len(chunk.Choices) == 0 {
if chunk.Usage != nil {
lastUsage = &Usage{
InputTokens: chunk.Usage.PromptTokens,
OutputTokens: chunk.Usage.CompletionTokens,
}
}
continue
}
choice := chunk.Choices[0]
if choice.FinishReason != "" {
finishReason = mapFinishReason(choice.FinishReason)
}
if len(choice.Delta.ToolCalls) > 0 {
for _, tc := range choice.Delta.ToolCalls {
if tc.Index != nil {
idx := *tc.Index
if idx >= len(toolCalls) {
if contentBlockStarted && currentToolCall == nil {
stopEvent := AnthropicStreamEvent{
Type: "content_block_stop",
Index: contentIndex - 1,
}
s.writeSSE(w, flusher, stopEvent)
}
currentToolCall = &ToolCallAccumulator{
Index: idx,
ID: tc.ID,
Name: tc.Function.Name,
Args: "",
}
toolCalls = append(toolCalls, currentToolCall)
startEvent := AnthropicStreamEvent{
Type: "content_block_start",
Index: contentIndex,
ContentBlock: &ContentBlock{
Type: "tool_use",
ID: tc.ID,
Name: tc.Function.Name,
Input: map[string]interface{}{},
},
}
s.writeSSE(w, flusher, startEvent)
contentIndex++
}
if tc.Function.Arguments != "" {
currentToolCall.Args += tc.Function.Arguments
deltaEvent := AnthropicStreamEvent{
Type: "content_block_delta",
Index: contentIndex - 1,
Delta: &DeltaContent{
Type: "input_json_delta",
PartialJSON: tc.Function.Arguments,
},
}
s.writeSSE(w, flusher, deltaEvent)
}
}
}
continue
}
if choice.Delta.Content != "" {
if !contentBlockStarted {
startEvent := AnthropicStreamEvent{
Type: "content_block_start",
Index: contentIndex,
ContentBlock: &ContentBlock{
Type: "text",
Text: "",
},
}
s.writeSSE(w, flusher, startEvent)
contentBlockStarted = true
contentIndex++
}
deltaEvent := AnthropicStreamEvent{
Type: "content_block_delta",
Index: contentIndex - 1,
Delta: &DeltaContent{
Type: "text_delta",
Text: choice.Delta.Content,
},
}
s.writeSSE(w, flusher, deltaEvent)
}
}
if contentBlockStarted || len(toolCalls) > 0 {
stopEvent := AnthropicStreamEvent{
Type: "content_block_stop",
Index: contentIndex - 1,
}
s.writeSSE(w, flusher, stopEvent)
}
if finishReason == "" {
finishReason = "end_turn"
}
if lastUsage == nil {
lastUsage = &Usage{InputTokens: 0, OutputTokens: 0}
}
deltaEvent := AnthropicStreamEvent{
Type: "message_delta",
Delta: &DeltaContent{
StopReason: &finishReason,
},
Usage: lastUsage,
}
s.writeSSE(w, flusher, deltaEvent)
stopEvent := AnthropicStreamEvent{
Type: "message_stop",
}
s.writeSSE(w, flusher, stopEvent)
}
func (s *Server) writeSSE(w http.ResponseWriter, flusher http.Flusher, event interface{}) {
data, err := json.Marshal(event)
if err != nil {
return
}
eventType := ""
if e, ok := event.(AnthropicStreamEvent); ok {
eventType = e.Type
}
if eventType != "" {
fmt.Fprintf(w, "event: %s\n", eventType)
}
fmt.Fprintf(w, "data: %s\n\n", data)
flusher.Flush()
if s.config.Verbose {
log.Printf("SSE event: %s", string(data))
}
}
func (s *Server) forwardRequest(openaiReq *OpenAIRequest) (*http.Response, error) {
body, err := json.Marshal(openaiReq)
if err != nil {
return nil, err
}
if s.config.Verbose {
log.Printf("Forwarding to backend: %s", string(body))
}
req, err := http.NewRequest(http.MethodPost, s.config.Backend, bytes.NewReader(body))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+s.config.APIKey)
return s.client.Do(req)
}
func (s *Server) errorResponse(w http.ResponseWriter, status int, errType, message string) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
json.NewEncoder(w).Encode(map[string]interface{}{
"type": "error",
"error": map[string]string{
"type": errType,
"message": message,
},
})
}
func generateID(prefix string) string {
return fmt.Sprintf("%s%d", prefix, time.Now().UnixNano())
}
func mapFinishReason(reason string) string {
switch reason {
case "stop":
return "end_turn"
case "length":
return "max_tokens"
case "tool_calls", "function_call":
return "tool_use"
case "content_filter":
return "end_turn"
default:
return "end_turn"
}
}

244
sandbox/v2/docker/bin/openai-proxy/types.go Normal file
View file

@ -0,0 +1,244 @@
package proxy
import "encoding/json"
// ============================================
// Anthropic API Types
// ============================================
type AnthropicRequest struct {
Model string `json:"model"`
Messages []AnthropicMsg `json:"messages"`
System interface{} `json:"system,omitempty"`
MaxTokens int `json:"max_tokens"`
Stream bool `json:"stream,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
TopK *int `json:"top_k,omitempty"`
StopSequences []string `json:"stop_sequences,omitempty"`
Tools []AnthropicTool `json:"tools,omitempty"`
ToolChoice *AnthropicToolChoice `json:"tool_choice,omitempty"`
Metadata map[string]string `json:"metadata,omitempty"`
}
type AnthropicMsg struct {
Role string `json:"role"`
Content interface{} `json:"content"`
}
type ContentBlock struct {
Type string `json:"type"`
Text string `json:"text,omitempty"`
Source *ImageSource `json:"source,omitempty"`
ID string `json:"id,omitempty"`
Name string `json:"name,omitempty"`
Input interface{} `json:"input,omitempty"`
ToolUseID string `json:"tool_use_id,omitempty"`
Content interface{} `json:"content,omitempty"`
IsError bool `json:"is_error,omitempty"`
}
type ImageSource struct {
Type string `json:"type"`
MediaType string `json:"media_type,omitempty"`
Data string `json:"data,omitempty"`
URL string `json:"url,omitempty"`
}
type SystemBlock struct {
Type string `json:"type"`
Text string `json:"text"`
}
type AnthropicTool struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
InputSchema interface{} `json:"input_schema"`
}
type AnthropicToolChoice struct {
Type string `json:"type"`
Name string `json:"name,omitempty"`
}
type AnthropicResponse struct {
ID string `json:"id"`
Type string `json:"type"`
Role string `json:"role"`
Content []ContentBlock `json:"content"`
Model string `json:"model"`
StopReason *string `json:"stop_reason"`
StopSequence *string `json:"stop_sequence,omitempty"`
Usage *Usage `json:"usage"`
}
type Usage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
}
type AnthropicStreamEvent struct {
Type string `json:"type"`
Index int `json:"index,omitempty"`
Message *AnthropicResponse `json:"message,omitempty"`
ContentBlock *ContentBlock `json:"content_block,omitempty"`
Delta *DeltaContent `json:"delta,omitempty"`
Usage *Usage `json:"usage,omitempty"`
}
type DeltaContent struct {
Type string `json:"type,omitempty"`
Text string `json:"text,omitempty"`
PartialJSON string `json:"partial_json,omitempty"`
StopReason *string `json:"stop_reason,omitempty"`
}
// ============================================
// OpenAI API Types
// ============================================
type OpenAIRequest struct {
Model string `json:"model"`
Messages []OpenAIMsg `json:"messages"`
MaxTokens int `json:"max_tokens,omitempty"`
Stream bool `json:"stream,omitempty"`
StreamOptions *StreamOptions `json:"stream_options,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
TopP *float64 `json:"top_p,omitempty"`
Stop []string `json:"stop,omitempty"`
Tools []OpenAITool `json:"tools,omitempty"`
ToolChoice interface{} `json:"tool_choice,omitempty"`
ExtraOptions map[string]interface{} `json:"-"`
}
func (r OpenAIRequest) MarshalJSON() ([]byte, error) {
m := map[string]interface{}{
"model": r.Model,
"messages": r.Messages,
}
if r.MaxTokens > 0 {
m["max_tokens"] = r.MaxTokens
}
if r.Stream {
m["stream"] = r.Stream
}
if r.StreamOptions != nil {
m["stream_options"] = r.StreamOptions
}
if r.Temperature != nil {
m["temperature"] = *r.Temperature
}
if r.TopP != nil {
m["top_p"] = *r.TopP
}
if len(r.Stop) > 0 {
m["stop"] = r.Stop
}
if len(r.Tools) > 0 {
m["tools"] = r.Tools
}
if r.ToolChoice != nil {
m["tool_choice"] = r.ToolChoice
}
for k, v := range r.ExtraOptions {
if _, exists := m[k]; !exists {
m[k] = v
}
}
return json.Marshal(m)
}
type StreamOptions struct {
IncludeUsage bool `json:"include_usage"`
}
type OpenAIMsg struct {
Role string `json:"role"`
Content interface{} `json:"content,omitempty"`
ToolCalls []OpenAIToolCall `json:"tool_calls,omitempty"`
ToolCallID string `json:"tool_call_id,omitempty"`
Name string `json:"name,omitempty"`
}
type OpenAIContent struct {
Type string `json:"type"`
Text string `json:"text,omitempty"`
ImageURL *OpenAIImageURL `json:"image_url,omitempty"`
}
type OpenAIImageURL struct {
URL string `json:"url"`
Detail string `json:"detail,omitempty"`
}
type OpenAITool struct {
Type string `json:"type"`
Function OpenAIFunction `json:"function"`
}
type OpenAIFunction struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
Parameters interface{} `json:"parameters"`
}
type OpenAIToolCall struct {
ID string `json:"id"`
Type string `json:"type"`
Function OpenAIFunctionCall `json:"function"`
Index *int `json:"index,omitempty"`
}
type OpenAIFunctionCall struct {
Name string `json:"name"`
Arguments string `json:"arguments"`
}
type OpenAIResponse struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []OpenAIChoice `json:"choices"`
Usage *OpenAIUsage `json:"usage,omitempty"`
}
type OpenAIChoice struct {
Index int `json:"index"`
Message OpenAIMsg `json:"message"`
FinishReason string `json:"finish_reason"`
}
type OpenAIUsage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
}
type OpenAIStreamChunk struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []OpenAIStreamChoice `json:"choices"`
Usage *OpenAIUsage `json:"usage,omitempty"`
}
type OpenAIStreamChoice struct {
Index int `json:"index"`
Delta OpenAIStreamDelta `json:"delta"`
FinishReason string `json:"finish_reason,omitempty"`
}
type OpenAIStreamDelta struct {
Role string `json:"role,omitempty"`
Content string `json:"content,omitempty"`
ToolCalls []OpenAIToolCall `json:"tool_calls,omitempty"`
}
type ToolCallAccumulator struct {
Index int
ID string
Name string
Args string
}

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 openai-proxy (multi-arch) ==="
cd "$SCRIPT_DIR/bin/openai-proxy/cmd/openai-proxy"
CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -ldflags="-s -w" -o "$SCRIPT_DIR/base/openai-proxy-amd64" .
CGO_ENABLED=0 GOOS=linux GOARCH=arm64 go build -ldflags="-s -w" -o "$SCRIPT_DIR/base/openai-proxy-arm64" .
echo "Built: openai-proxy-amd64, openai-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/openai-proxy-amd64" "$SCRIPT_DIR/base/openai-proxy-arm64"
echo "Removed temporary binary files"
echo ""
echo "=== Build complete ==="
docker images | grep -E "sandbox-v2" | head -10 || true

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

@ -0,0 +1,28 @@
# Sandbox V2 test image — adds test services + VNC desktop 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 \
xvfb \
x11vnc \
fluxbox \
xterm \
&& pip3 install --break-system-packages websockets websockify \
&& 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
ENV DISPLAY=:99
USER sandbox
EXPOSE 5900 6080
ENTRYPOINT ["/usr/bin/tini", "--", "/test-entrypoint.sh"]
CMD ["sleep", "infinity"]

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

@ -0,0 +1,22 @@
#!/bin/bash
# V2 test entrypoint — starts test services + VNC desktop then delegates to base entrypoint
# Start Xvfb (virtual framebuffer)
Xvfb :99 -screen 0 1024x768x24 -ac +extension GLX +render -noreset &
sleep 0.5
# Start fluxbox window manager
fluxbox &
# Start x11vnc (raw RFB on 5900)
x11vnc -display :99 -rfbport 5900 -nopw -shared -forever -xkb -ncache 10 &
sleep 0.3
# Start websockify (WebSocket on 6080 → RFB 5900)
websockify 0.0.0.0:6080 localhost:5900 &
# Test services
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")
}
}

129
sandbox/v2/jsapi/box.go Normal file
View file

@ -0,0 +1,129 @@
package jsapi
import (
"rogchap.com/v8go"
)
// NewBoxObject creates a JS Box object backed by a sandbox ID string.
// All methods delegate to the Go sandbox.M() singleton — no Go object is
// passed to V8, no bridge registration, no Release() needed.
//
// # Properties (read-only)
//
// box.id → string // sandbox ID ← Box.ID()
// box.owner → string // owner user ID ← Box.Owner()
// box.pool → string // pool name ← Box.Pool()
//
// # Methods — Go mapping
//
// box.Exec(cmd, options?) → ExecResult
//
// Go: Box.Exec(ctx, cmd []string, opts ...ExecOption) (*ExecResult, error)
//
// JS args:
// cmd: string[] → cmd []string
// options: { → ExecOption functional options
// workdir: string, → WithWorkDir(dir)
// env: object, → WithEnv(map[string]string)
// timeout: number → WithTimeout(ms → time.Duration)
// }
// JS returns: {
// exit_code: number, ← ExecResult.ExitCode
// stdout: string, ← ExecResult.Stdout
// stderr: string ← ExecResult.Stderr
// }
//
// box.Stream(cmd, options?) → ExecStream
//
// Go: Box.Stream(ctx, cmd []string, opts ...ExecOption) (*ExecStream, error)
//
// JS returns: {
// stdout: ReadableStream, ← ExecStream.Stdout
// stderr: ReadableStream, ← ExecStream.Stderr
// stdin: WritableStream, ← ExecStream.Stdin
// wait: function() → number, ← ExecStream.Wait() (int, error)
// cancel: function() → void ← ExecStream.Cancel()
// }
//
// box.Attach(port, options?) → ServiceConn
//
// Go: Box.Attach(ctx, port int, opts ...AttachOption) (*ServiceConn, error)
//
// JS args:
// port: number → port int
// options: { → AttachOption functional options
// protocol: "ws"|"sse", → WithProtocol(protocol)
// path: string, → WithPath(path)
// headers: object → WithHeaders(map[string]string)
// }
// JS returns: {
// url: string, ← ServiceConn.URL
// read: function() → Uint8Array, ← ServiceConn.Read() ([]byte, error)
// write: function(data) → void, ← ServiceConn.Write(data) error
// events: AsyncIterable<Uint8Array>, ← ServiceConn.Events <-chan []byte
// close: function() → void ← ServiceConn.Close() error
// }
//
// box.VNC() → string
//
// Go: Box.VNC(ctx) (string, error)
// Returns: VNC WebSocket URL
//
// box.Proxy(port, path?) → string
//
// Go: Box.Proxy(ctx, port int, path string) (string, error)
// Returns: HTTP proxy URL
//
// box.Workspace() → WorkspaceFS
//
// Go: Box.Workspace() workspace.FS
// Box.WorkspaceID() string
// Returns: WorkspaceFS object (see workspace/jsapi/fs.go)
// Uses box.WorkspaceID() to create NewFSObject
//
// box.Info() → BoxInfo
//
// Go: Box.Info(ctx) (*BoxInfo, error)
// JS returns: {
// id: string, ← BoxInfo.ID
// container_id: string, ← BoxInfo.ContainerID
// pool: string, ← BoxInfo.Pool
// owner: string, ← BoxInfo.Owner
// status: string, ← BoxInfo.Status
// image: string, ← BoxInfo.Image
// vnc: boolean, ← BoxInfo.VNC
// policy: string, ← BoxInfo.Policy (LifecyclePolicy)
// labels: object, ← BoxInfo.Labels (map[string]string)
// created_at: string, ← BoxInfo.CreatedAt (ISO 8601)
// last_active: string, ← BoxInfo.LastActive (ISO 8601)
// process_count: number ← BoxInfo.ProcessCount
// }
//
// box.Start() → void
//
// Go: Box.Start(ctx) error
//
// box.Stop() → void
//
// Go: Box.Stop(ctx) error
//
// box.Remove() → void
//
// Go: Box.Remove(ctx) error
func NewBoxObject(v8ctx *v8go.Context, boxID string) (*v8go.Value, error) {
// TODO: Phase 2 implementation
// 1. Create JS object via v8go.NewObjectTemplate
// 2. Set read-only properties: id, owner, pool (from sandbox.M().Get(boxID))
// 3. Bind each method as FunctionTemplate:
// - Exec → sandbox.M().Get(id).Exec(ctx, cmd, opts...)
// - Stream → sandbox.M().Get(id).Stream(ctx, cmd, opts...)
// - Attach → sandbox.M().Get(id).Attach(ctx, port, opts...)
// - VNC → sandbox.M().Get(id).VNC(ctx)
// - Proxy → sandbox.M().Get(id).Proxy(ctx, port, path)
// - Workspace → NewFSObject(v8ctx, sandbox.M().Get(id).WorkspaceID())
// - Info → sandbox.M().Get(id).Info(ctx) → JS object
// - Start → sandbox.M().Get(id).Start(ctx)
// - Stop → sandbox.M().Get(id).Stop(ctx)
// - Remove → sandbox.M().Get(id).Remove(ctx)
return nil, nil
}

155
sandbox/v2/jsapi/jsapi.go Normal file
View file

@ -0,0 +1,155 @@
// Package jsapi registers the sandbox namespace into the Yao V8 runtime.
//
// All methods are static on the sandbox object — no constructor.
//
// # JavaScript API
//
// const box = sandbox.Create({ image: "node:20", owner: "user1" })
// const result = box.Exec(["node", "-e", "console.log('hi')"])
// console.log(result.stdout)
//
// const box = sandbox.Get(id) // → Box
// const list = sandbox.List({ owner: "u1" }) // → BoxInfo[]
// sandbox.Delete(id) // → void
//
// # Go mapping
//
// sandbox.Create(opts) → Manager.Create(ctx, CreateOptions) → Box
// sandbox.Create(opts) → Manager.GetOrCreate(ctx, opts) → Box (when opts.id is set)
// sandbox.Get(id) → Manager.Get(ctx, id) → Box
// sandbox.List(filter?) → Manager.List(ctx, ListOptions) → []*Box → BoxInfo[]
// sandbox.Delete(id) → Manager.Remove(ctx, id) → void
//
// Registration happens via init() — import with:
//
// _ "github.com/yaoapp/yao/sandbox/v2/jsapi"
package jsapi
import (
v8 "github.com/yaoapp/gou/runtime/v8"
"rogchap.com/v8go"
)
func init() {
v8.RegisterObject("sandbox", ExportObject)
}
// ExportObject exports the sandbox namespace object to V8.
func ExportObject(iso *v8go.Isolate) *v8go.ObjectTemplate {
obj := v8go.NewObjectTemplate(iso)
obj.Set("Create", v8go.NewFunctionTemplate(iso, sbCreate))
obj.Set("Get", v8go.NewFunctionTemplate(iso, sbGet))
obj.Set("List", v8go.NewFunctionTemplate(iso, sbList))
obj.Set("Delete", v8go.NewFunctionTemplate(iso, sbDelete))
return obj
}
// sbCreate: `sandbox.Create(options)` → Box
//
// Go: Manager.Create(ctx, CreateOptions) (*Box, error)
//
// Manager.GetOrCreate(ctx, CreateOptions) (*Box, error) — when opts.id is set
//
// JS options → Go CreateOptions mapping:
//
// {
// id: string → CreateOptions.ID // optional; triggers GetOrCreate
// owner: string → CreateOptions.Owner // required
// pool: string → CreateOptions.Pool // default: first pool
// image: string → CreateOptions.Image // required
// workdir: string → CreateOptions.WorkDir
// user: string → CreateOptions.User // e.g. "1000:1000"
// env: object → CreateOptions.Env // map[string]string
// memory: number → CreateOptions.Memory // bytes (int64)
// cpus: number → CreateOptions.CPUs // float64 e.g. 1.5
// vnc: boolean → CreateOptions.VNC
// ports: array → CreateOptions.Ports // [{container, host, host_ip, protocol}] → []PortMapping
// policy: string → CreateOptions.Policy // "oneshot"|"session"|"longrunning"|"persistent"
// idle_timeout: number → CreateOptions.IdleTimeout // ms → time.Duration
// stop_timeout: number → CreateOptions.StopTimeout // ms → time.Duration
// workspace_id: string → CreateOptions.WorkspaceID
// mount_mode: string → CreateOptions.MountMode // "rw"|"ro"
// mount_path: string → CreateOptions.MountPath
// labels: object → CreateOptions.Labels // map[string]string
// }
//
// Returns: Box object (see box.go)
func sbCreate(info *v8go.FunctionCallbackInfo) *v8go.Value {
// TODO: Phase 2
// 1. Parse options from info.Args()[0]
// 2. Validate required fields (image, owner)
// 3. If opts.id != "" → sandbox.M().GetOrCreate(ctx, opts)
// else → sandbox.M().Create(ctx, opts)
// 4. Return NewBoxObject(v8ctx, box.ID())
return v8go.Undefined(info.Context().Isolate())
}
// sbGet: `sandbox.Get(id)` → Box | null
//
// Go: Manager.Get(ctx, id) (*Box, error)
//
// Args:
//
// id: string — sandbox ID
//
// Returns: Box object if found, null if not found
func sbGet(info *v8go.FunctionCallbackInfo) *v8go.Value {
// TODO: Phase 2
// 1. id = info.Args()[0].String()
// 2. box, err := sandbox.M().Get(ctx, id)
// 3. Return NewBoxObject(v8ctx, id) or null
return v8go.Undefined(info.Context().Isolate())
}
// sbList: `sandbox.List(filter?)` → BoxInfo[]
//
// Go: Manager.List(ctx, ListOptions) ([]*Box, error)
//
// then Box.Info(ctx) for each → BoxInfo
//
// JS filter → Go ListOptions mapping:
//
// {
// owner: string → ListOptions.Owner // filter by owner; empty = all
// pool: string → ListOptions.Pool // filter by pool; empty = all
// labels: object → ListOptions.Labels // filter by labels
// }
//
// Returns: BoxInfo[] — each element:
//
// {
// id: string ← BoxInfo.ID
// container_id: string ← BoxInfo.ContainerID
// pool: string ← BoxInfo.Pool
// owner: string ← BoxInfo.Owner
// status: string ← BoxInfo.Status
// image: string ← BoxInfo.Image
// vnc: boolean ← BoxInfo.VNC
// policy: string ← BoxInfo.Policy
// labels: object ← BoxInfo.Labels
// created_at: string ← BoxInfo.CreatedAt (ISO 8601)
// last_active: string ← BoxInfo.LastActive (ISO 8601)
// process_count: number ← BoxInfo.ProcessCount
// }
func sbList(info *v8go.FunctionCallbackInfo) *v8go.Value {
// TODO: Phase 2
// 1. Parse optional filter from info.Args()[0]
// 2. boxes := sandbox.M().List(ctx, opts)
// 3. For each box: box.Info(ctx) → BoxInfo → JS object
// 4. Return JS array of BoxInfo objects
return v8go.Undefined(info.Context().Isolate())
}
// sbDelete: `sandbox.Delete(id)` → void
//
// Go: Manager.Remove(ctx, id) error
//
// Args:
//
// id: string — sandbox ID to remove
func sbDelete(info *v8go.FunctionCallbackInfo) *v8go.Value {
// TODO: Phase 2
// 1. id = info.Args()[0].String()
// 2. sandbox.M().Remove(ctx, id)
return v8go.Undefined(info.Context().Isolate())
}

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

@ -0,0 +1,616 @@
package sandbox
import (
"context"
"fmt"
"sync"
"time"
"github.com/yaoapp/yao/tai"
taisandbox "github.com/yaoapp/yao/tai/sandbox"
"github.com/yaoapp/yao/workspace"
)
// 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
wsManager *workspace.Manager
}
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
}
// Workspace node binding: when WorkspaceID is set, resolve the workspace's
// bound node and force the container onto that pool.
if opts.WorkspaceID != "" && m.wsManager != nil {
node, err := m.wsManager.NodeForWorkspace(ctx, opts.WorkspaceID)
if err != nil {
return nil, fmt.Errorf("sandbox: resolve workspace %q: %w", opts.WorkspaceID, err)
}
poolName = node
}
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,
stopTimeoutD: opts.StopTimeout,
createdAt: time.Now(),
refreshToken: refresh,
manager: m,
vnc: opts.VNC,
image: opts.Image,
workspaceID: opts.WorkspaceID,
}
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 force-removes a sandbox (SIGKILL + delete).
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().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, b.stopTimeout())
}
}
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
}
// SetWorkspaceManager links the workspace manager for workspace-aware container creation.
// When CreateOptions.WorkspaceID is set, the sandbox Manager uses the workspace Manager
// to resolve the workspace's bound node and force container routing.
func (m *Manager) SetWorkspaceManager(wm *workspace.Manager) {
m.wsManager = wm
}
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),
}
if opts.WorkspaceID != "" {
labels["workspace-id"] = opts.WorkspaceID
}
for k, v := range opts.Labels {
labels[k] = v
}
workDir := opts.WorkDir
if workDir == "" {
workDir = "/workspace"
}
cmd := []string{"sh", "-c", "trap 'exit 0' TERM; while :; do sleep 86400 & wait $!; done"}
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,
})
}
// Workspace bind mount
var binds []string
if opts.WorkspaceID != "" && m.wsManager != nil {
mountPath := opts.MountPath
if mountPath == "" {
mountPath = "/workspace"
}
mode := opts.MountMode
if mode == "" {
mode = "rw"
}
hostPath, _ := m.wsManager.MountPath(context.Background(), opts.WorkspaceID)
if hostPath != "" {
binds = append(binds, fmt.Sprintf("%s:%s:%s", hostPath, mountPath, mode))
}
}
return taisandbox.CreateOptions{
Name: sandboxID,
Image: opts.Image,
Cmd: cmd,
Env: env,
Binds: binds,
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,
workspaceID: c.Labels["workspace-id"],
manager: m,
}
box.lastCall.Store(time.Now().UnixMilli())
m.boxes.Store(sandboxID, box)
}
}
// ImageExists reports whether the given image ref exists on the target pool node.
func (m *Manager) ImageExists(ctx context.Context, pool, ref string) (bool, error) {
client, err := m.getPool(pool)
if err != nil {
return false, err
}
return client.Image().Exists(ctx, ref)
}
// PullImage pulls an image to the target pool node, returning a channel of
// real-time progress events. The channel is nil when no pull is needed (e.g. K8s mode).
func (m *Manager) PullImage(ctx context.Context, pool, ref string, opts ImagePullOptions) (<-chan taisandbox.PullProgress, error) {
client, err := m.getPool(pool)
if err != nil {
return nil, err
}
pullOpts := taisandbox.PullOptions{}
if opts.Auth != nil {
pullOpts.Auth = &taisandbox.RegistryAuth{
Username: opts.Auth.Username,
Password: opts.Auth.Password,
Server: opts.Auth.Server,
}
}
return client.Image().Pull(ctx, ref, pullOpts)
}
// EnsureImage checks whether the image exists on the pool node; if not, it
// pulls the image and blocks until the pull completes. Returns the first
// error encountered during pull. For K8s pools this is a no-op.
func (m *Manager) EnsureImage(ctx context.Context, pool, ref string, opts ImagePullOptions) error {
exists, err := m.ImageExists(ctx, pool, ref)
if err != nil {
return fmt.Errorf("image exists check: %w", err)
}
if exists {
return nil
}
ch, err := m.PullImage(ctx, pool, ref, opts)
if err != nil {
return fmt.Errorf("image pull: %w", err)
}
if ch == nil {
return nil
}
for p := range ch {
if p.Error != "" {
return fmt.Errorf("image pull %q: %s", ref, p.Error)
}
}
return nil
}

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

@ -0,0 +1,142 @@
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
})
ensureTestImage(t, m, pc.Name)
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)
}
})
}
}

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

@ -0,0 +1,299 @@
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)
ensureTestImage(t, m, pc.Name)
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
})
ensureTestImage(t, m, pc.Name)
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, Options: pc.Options})
}
m := setupManager(t, sps...)
for _, pc := range pools {
ensureTestImage(t, m, pc.Name)
}
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()
}

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

@ -0,0 +1,196 @@
package sandbox_test
import (
"context"
"fmt"
"os"
"strconv"
"testing"
"time"
sandbox "github.com/yaoapp/yao/sandbox/v2"
"github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/volume"
"github.com/yaoapp/yao/workspace"
)
type poolConfig struct {
Name string
Addr string
Options []tai.Option
}
// testPools returns all available pool configurations for multi-mode testing.
// - local: always present (direct Docker daemon)
// - remote: when SANDBOX_TEST_REMOTE_ADDR is set (Tai on host → Docker)
// - containerized: when TAI_TEST_CONTAINERIZED_HOST is set (Tai in container → Docker)
// - k8s: when TAI_TEST_K8S_HOST + TAI_TEST_KUBECONFIG are set (Tai → K8s)
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})
}
if host := os.Getenv("TAI_TEST_CONTAINERIZED_HOST"); host != "" {
grpcPort := envPort("TAI_TEST_CONTAINERIZED_GRPC_PORT", 9200)
addr := fmt.Sprintf("tai://%s:%d", host, grpcPort)
// No WithPorts for HTTP/VNC — Tai self-inspects its container
// and returns host-mapped ports via ServerInfo automatically.
pools = append(pools, poolConfig{Name: "containerized", Addr: addr})
}
if host := os.Getenv("TAI_TEST_K8S_HOST"); host != "" {
kubeconfig := os.Getenv("TAI_TEST_KUBECONFIG")
if kubeconfig == "" {
return pools
}
grpcPort := envPort("TAI_TEST_K8S_GRPC_PORT", envPort("TAI_TEST_GRPC_PORT", 9100))
addr := fmt.Sprintf("tai://%s:%d", host, grpcPort)
opts := []tai.Option{
tai.K8s,
tai.WithKubeConfig(kubeconfig),
tai.WithPorts(tai.Ports{
K8s: envPort("TAI_TEST_K8S_PORT", 6443),
GRPC: grpcPort,
}),
}
if ns := os.Getenv("TAI_TEST_K8S_NAMESPACE"); ns != "" {
opts = append(opts, tai.WithNamespace(ns))
}
pools = append(pools, poolConfig{Name: "k8s", Addr: addr, Options: opts})
}
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 envPort(key string, fallback int) int {
if v := os.Getenv(key); v != "" {
if p, err := strconv.Atoi(v); err == nil {
return p
}
}
return fallback
}
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, Options: pc.Options}
for _, fn := range mutators {
fn(&pool)
}
return setupManager(t, pool)
}
// setupManagerWithWorkspace creates a sandbox Manager with a linked workspace Manager.
// Returns both managers and a helper to create workspaces on the given pool's node.
func setupManagerWithWorkspace(t *testing.T, pc poolConfig) (*sandbox.Manager, *workspace.Manager) {
t.Helper()
sbm := setupManagerForPool(t, pc)
var wsClient *tai.Client
var err error
if pc.Addr == "local" || pc.Addr == "" {
dataDir := t.TempDir()
vol := volume.NewLocal(dataDir)
wsClient, err = tai.New("local", tai.WithVolume(vol), tai.WithDataDir(dataDir))
} else {
wsClient, err = tai.New(pc.Addr, pc.Options...)
}
if err != nil {
t.Fatalf("tai.New for workspace: %v", err)
}
t.Cleanup(func() { wsClient.Close() })
wsm := workspace.NewManager(map[string]*tai.Client{pc.Name: wsClient})
sbm.SetWorkspaceManager(wsm)
return sbm, wsm
}
// ensureTestImage guarantees testImage() is available on the given pool before
// container creation. Safe for all modes (Docker pull; K8s no-op).
func ensureTestImage(t *testing.T, m *sandbox.Manager, pool string) {
t.Helper()
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
defer cancel()
if err := m.EnsureImage(ctx, pool, testImage(), sandbox.ImagePullOptions{}); err != nil {
t.Fatalf("EnsureImage(%s, %s): %v", pool, testImage(), err)
}
}
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)
}
pool := co.Pool
if pool == "" {
pools := m.Pools()
if len(pools) > 0 {
pool = pools[0].Name
}
}
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
defer cancel()
if pool != "" {
if err := m.EnsureImage(ctx, pool, co.Image, sandbox.ImagePullOptions{}); err != nil {
t.Fatalf("EnsureImage(%s, %s): %v", pool, co.Image, err)
}
}
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
}

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

@ -0,0 +1,178 @@
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"
)
const DefaultStopTimeout = 2 * time.Second
type Pool struct {
Name string
Addr string
Options []tai.Option
MaxPerUser int
MaxTotal int
IdleTimeout time.Duration
MaxLifetime time.Duration
StopTimeout time.Duration // SIGTERM grace period before SIGKILL; 0 = DefaultStopTimeout
}
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
StopTimeout time.Duration // SIGTERM grace period; 0 = pool default or DefaultStopTimeout
WorkspaceID string // workspace to mount; empty = no workspace
MountMode string // "rw" (default) or "ro"
MountPath string // container path; default "/workspace"
}
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
}
}
// ImagePullOptions configures an image pull operation.
type ImagePullOptions struct {
Auth *RegistryAuth // nil = anonymous / public
}
// RegistryAuth holds credentials for a private container registry.
type RegistryAuth struct {
Username string
Password string
Server string
}
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)

20
tai/sandbox/client_accessor.go Normal file
View file

@ -0,0 +1,20 @@
package sandbox
import "github.com/docker/docker/client"
// dockerCliAccessor is implemented by sandbox types that hold a Docker client.
type dockerCliAccessor interface {
dockerClient() *client.Client
}
func (l *local) dockerClient() *client.Client { return l.core.cli }
func (d *dockerSandbox) dockerClient() *client.Client { return d.core.cli }
// DockerCli extracts the underlying Docker SDK client from a Sandbox.
// Returns nil if the Sandbox is not Docker-based (e.g. K8s).
func DockerCli(sb Sandbox) *client.Client {
if a, ok := sb.(dockerCliAccessor); ok {
return a.dockerClient()
}
return nil
}

6
tai/sandbox/docker.go
View file

@ -32,7 +32,7 @@ func NewDocker(addr string) (Sandbox, error) {
}
func (d *dockerSandbox) Create(ctx context.Context, opts CreateOptions) (string, error) {
return d.core.create(ctx, opts, false)
return d.core.create(ctx, opts, true)
}
func (d *dockerSandbox) Start(ctx context.Context, id string) error {
@ -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{

43
tai/sandbox/image.go Normal file
View file

@ -0,0 +1,43 @@
package sandbox
import (
"context"
"time"
)
// Image manages container images on a runtime node.
type Image interface {
Exists(ctx context.Context, ref string) (bool, error)
Pull(ctx context.Context, ref string, opts PullOptions) (<-chan PullProgress, error)
Remove(ctx context.Context, ref string, force bool) error
List(ctx context.Context) ([]ImageInfo, error)
}
// PullOptions configures an image pull operation.
type PullOptions struct {
Auth *RegistryAuth // nil = anonymous / public
}
// RegistryAuth holds credentials for a private container registry.
type RegistryAuth struct {
Username string
Password string
Server string // e.g. "ghcr.io", "registry.example.com"
}
// PullProgress reports real-time progress of an image pull.
type PullProgress struct {
Status string // "Pulling fs layer", "Downloading", "Extracting", "Pull complete", etc.
Layer string // layer digest / short ID
Current int64 // bytes completed
Total int64 // bytes total (0 if unknown)
Error string // non-empty on failure
}
// ImageInfo describes a local image.
type ImageInfo struct {
ID string
Tags []string
Size int64
Created time.Time
}

132
tai/sandbox/image_docker.go Normal file
View file

@ -0,0 +1,132 @@
package sandbox
import (
"context"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"time"
"github.com/docker/docker/api/types/image"
"github.com/docker/docker/api/types/registry"
"github.com/docker/docker/client"
)
// dockerImage implements Image using the Docker SDK.
// Shared by both local and dockerSandbox (via Tai proxy) modes.
type dockerImage struct {
cli *client.Client
}
// NewDockerImage creates an Image backed by a Docker client.
func NewDockerImage(cli *client.Client) Image {
return &dockerImage{cli: cli}
}
func (d *dockerImage) Exists(ctx context.Context, ref string) (bool, error) {
_, _, err := d.cli.ImageInspectWithRaw(ctx, ref)
if err != nil {
if client.IsErrNotFound(err) {
return false, nil
}
return false, fmt.Errorf("image inspect %q: %w", ref, err)
}
return true, nil
}
func (d *dockerImage) Pull(ctx context.Context, ref string, opts PullOptions) (<-chan PullProgress, error) {
pullOpts := image.PullOptions{}
if opts.Auth != nil {
encoded, err := encodeAuth(opts.Auth)
if err != nil {
return nil, err
}
pullOpts.RegistryAuth = encoded
}
reader, err := d.cli.ImagePull(ctx, ref, pullOpts)
if err != nil {
return nil, fmt.Errorf("image pull %q: %w", ref, err)
}
ch := make(chan PullProgress, 32)
go func() {
defer close(ch)
defer reader.Close()
decodePullStream(reader, ch)
}()
return ch, nil
}
func (d *dockerImage) Remove(ctx context.Context, ref string, force bool) error {
_, err := d.cli.ImageRemove(ctx, ref, image.RemoveOptions{Force: force, PruneChildren: true})
if err != nil {
return fmt.Errorf("image remove %q: %w", ref, err)
}
return nil
}
func (d *dockerImage) List(ctx context.Context) ([]ImageInfo, error) {
imgs, err := d.cli.ImageList(ctx, image.ListOptions{})
if err != nil {
return nil, fmt.Errorf("image list: %w", err)
}
result := make([]ImageInfo, len(imgs))
for i, img := range imgs {
result[i] = ImageInfo{
ID: img.ID,
Tags: img.RepoTags,
Size: img.Size,
Created: time.Unix(img.Created, 0),
}
}
return result, nil
}
// dockerPullEvent mirrors the JSON lines emitted by Docker's ImagePull stream.
type dockerPullEvent struct {
Status string `json:"status"`
ID string `json:"id"`
ProgressDetail struct {
Current int64 `json:"current"`
Total int64 `json:"total"`
} `json:"progressDetail"`
Error string `json:"error"`
}
func decodePullStream(r io.Reader, ch chan<- PullProgress) {
dec := json.NewDecoder(r)
for {
var ev dockerPullEvent
if err := dec.Decode(&ev); err != nil {
if err != io.EOF {
ch <- PullProgress{Error: err.Error()}
}
return
}
p := PullProgress{
Status: ev.Status,
Layer: ev.ID,
Current: ev.ProgressDetail.Current,
Total: ev.ProgressDetail.Total,
}
if ev.Error != "" {
p.Error = ev.Error
}
ch <- p
}
}
func encodeAuth(auth *RegistryAuth) (string, error) {
cfg := registry.AuthConfig{
Username: auth.Username,
Password: auth.Password,
ServerAddress: auth.Server,
}
data, err := json.Marshal(cfg)
if err != nil {
return "", fmt.Errorf("encode registry auth: %w", err)
}
return base64.URLEncoding.EncodeToString(data), nil
}

25
tai/sandbox/image_k8s.go Normal file
View file

@ -0,0 +1,25 @@
package sandbox
import "context"
// k8sImage is a no-op Image for K8s mode.
// Image pulling is handled by kubelet based on imagePullPolicy and imagePullSecrets.
type k8sImage struct{}
func NewK8sImage() Image { return &k8sImage{} }
func (k *k8sImage) Exists(_ context.Context, _ string) (bool, error) {
return true, nil
}
func (k *k8sImage) Pull(_ context.Context, _ string, _ PullOptions) (<-chan PullProgress, error) {
return nil, nil
}
func (k *k8sImage) Remove(_ context.Context, _ string, _ bool) error {
return nil
}
func (k *k8sImage) List(_ context.Context) ([]ImageInfo, error) {
return nil, nil
}

118
tai/sandbox/k8s.go
View file

@ -4,6 +4,7 @@ import (
"bytes"
"context"
"fmt"
"io"
"path/filepath"
"strings"
"time"
@ -112,7 +113,7 @@ func (s *k8sSandbox) Create(ctx context.Context, opts CreateOptions) (string, er
container := corev1.Container{
Name: "main",
Image: opts.Image,
Command: opts.Cmd,
Args: opts.Cmd,
Env: envVars,
WorkingDir: opts.WorkingDir,
}
@ -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)
@ -147,9 +160,16 @@ func (s *k8sSandbox) Create(ctx context.Context, opts CreateOptions) (string, er
}
func (s *k8sSandbox) Start(ctx context.Context, id string) error {
// K8s pods start automatically after creation.
// Wait briefly for the pod to leave Pending.
for i := 0; i < 30; i++ {
if _, ok := ctx.Deadline(); !ok {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, 60*time.Second)
defer cancel()
}
ticker := time.NewTicker(1 * time.Second)
defer ticker.Stop()
for {
pod, err := s.cli.CoreV1().Pods(s.ns).Get(ctx, id, metav1.GetOptions{})
if err != nil {
return fmt.Errorf("get pod: %w", err)
@ -157,9 +177,13 @@ func (s *k8sSandbox) Start(ctx context.Context, id string) error {
if pod.Status.Phase == corev1.PodRunning || pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed {
return nil
}
time.Sleep(1 * time.Second)
select {
case <-ctx.Done():
return fmt.Errorf("pod %s did not reach Running: %w", id, ctx.Err())
case <-ticker.C:
}
}
return fmt.Errorf("pod %s did not reach Running within 30s", id)
}
func (s *k8sSandbox) Stop(ctx context.Context, id string, timeout time.Duration) error {
@ -236,6 +260,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 +344,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 +370,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 +384,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 {

13
tai/sandbox/local.go
View file

@ -3,7 +3,6 @@ package sandbox
import (
"context"
"fmt"
"runtime"
"time"
"github.com/docker/docker/client"
@ -36,7 +35,7 @@ func NewLocal(addr string) (Sandbox, error) {
}
func (l *local) Create(ctx context.Context, opts CreateOptions) (string, error) {
return l.core.create(ctx, opts, opts.VNC && needsPortMapping())
return l.core.create(ctx, opts, opts.VNC)
}
func (l *local) Start(ctx context.Context, id string) error {
@ -55,6 +54,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)
}
@ -67,12 +70,6 @@ func (l *local) Close() error {
return l.core.cli.Close()
}
// needsPortMapping returns true on platforms where container IPs are not
// directly reachable (macOS Docker Desktop, Windows).
func needsPortMapping() bool {
return runtime.GOOS == "darwin" || runtime.GOOS == "windows"
}
func portStr(p int) string {
if p == 0 {
return ""

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

164
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.
@ -69,13 +77,21 @@ func WithNamespace(ns string) Option {
return optionFunc(func(c *config) { c.namespace = ns })
}
// WithVolume injects a custom Volume implementation.
// Useful for testing workspace operations without Docker.
func WithVolume(vol volume.Volume) Option {
return optionFunc(func(c *config) { c.volume = vol })
}
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
namespace string
volume volume.Volume // override volume (for testing without Docker)
}
func defaultPorts() Ports {
@ -112,8 +128,10 @@ type Client struct {
host string
addr string
ports Ports
dataDir string // host-side data directory for local volume
vol volume.Volume
sb sandbox.Sandbox
img sandbox.Image
prx proxy.Proxy
vc vnc.VNC
grpcConn *grpc.ClientConn
@ -121,9 +139,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 +151,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,
@ -155,32 +179,48 @@ func New(addr string, opts ...Option) (*Client, error) {
func (c *Client) initLocal(cfg *config) (*Client, error) {
sb, err := sandbox.NewLocal(c.addr)
if err != nil {
if err != nil && cfg.volume == nil {
return nil, err
}
c.sb = sb
c.prx = proxy.NewLocal(sb)
c.vc = vnc.NewLocal(sb)
dataDir := cfg.dataDir
if dataDir == "" {
dataDir = "/tmp/tai-volumes"
if sb != nil {
c.sb = sb
c.img = sandbox.NewDockerImage(sandbox.DockerCli(sb))
c.prx = proxy.NewLocal(sb)
c.vc = vnc.NewLocal(sb)
}
if cfg.volume != nil {
c.vol = cfg.volume
c.dataDir = cfg.dataDir
} else {
dataDir := cfg.dataDir
if dataDir == "" {
dataDir = "/tmp/tai-volumes"
}
c.dataDir = dataDir
c.vol = volume.NewLocal(dataDir)
}
c.vol = volume.NewLocal(dataDir)
return c, nil
}
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
@ -197,6 +237,7 @@ func (c *Client) initRemote(cfg *config) (*Client, error) {
return nil, err
}
c.sb = sb
c.img = sandbox.NewK8sImage()
default:
dockerPort := c.ports.Docker
if dockerPort == 0 {
@ -209,6 +250,7 @@ func (c *Client) initRemote(cfg *config) (*Client, error) {
return nil, err
}
c.sb = sb
c.img = sandbox.NewDockerImage(sandbox.DockerCli(sb))
}
hc := cfg.httpClient
@ -244,6 +286,10 @@ func (c *Client) Close() error {
// Volume returns the Volume IO layer. Never nil.
func (c *Client) Volume() volume.Volume { return c.vol }
// DataDir returns the host-side data directory used by the local volume.
// Empty for remote (Tai gRPC) connections — the Tai server manages paths.
func (c *Client) DataDir() string { return c.dataDir }
// Workspace returns an fs.FS-compatible filesystem for the given session.
func (c *Client) Workspace(sessionID string) workspace.FS {
return workspace.New(c.vol, sessionID)
@ -252,6 +298,9 @@ func (c *Client) Workspace(sessionID string) workspace.FS {
// Sandbox returns the container lifecycle manager. Never nil.
func (c *Client) Sandbox() sandbox.Sandbox { return c.sb }
// Image returns the container image manager. Never nil.
func (c *Client) Image() sandbox.Image { return c.img }
// Proxy returns the HTTP reverse proxy helper. Never nil.
func (c *Client) Proxy() proxy.Proxy { return c.prx }
@ -261,41 +310,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
}

103
tai/tai_test.go
View file

@ -1,6 +1,7 @@
package tai
import (
"fmt"
"os"
"strconv"
"testing"
@ -24,28 +25,38 @@ func envPort(key string, fallback int) int {
func TestParseAddr(t *testing.T) {
tests := []struct {
addr string
wantScheme string
wantHost string
wantDocker string
wantErr bool
addr string
wantScheme string
wantHost string
wantDocker string
wantGRPCPort int
wantErr bool
}{
{"", "docker", "", "", false},
{"docker:///var/run/docker.sock", "docker", "", "docker:///var/run/docker.sock", false},
{"docker://192.168.1.50:2375", "docker", "", "docker://192.168.1.50:2375", false},
{"unix:///var/run/docker.sock", "docker", "", "unix:///var/run/docker.sock", false},
{"tcp://127.0.0.1:2375", "docker", "", "tcp://127.0.0.1:2375", false},
{"npipe:////./pipe/docker_engine", "docker", "", "npipe:////./pipe/docker_engine", false},
{"tai://192.168.1.100", "tai", "192.168.1.100", "", false},
{"tai://10.0.0.5:9100", "tai", "10.0.0.5", "", false},
{"tai://", "", "", "", true},
{"ftp://host", "", "", "", true},
{" tai://host ", "tai", "host", "", false},
{"", "", "", "", 0, true},
{"local", "docker", "", "", 0, false},
{"127.0.0.1", "docker", "", "", 0, false},
{"localhost", "docker", "", "", 0, false},
{"::1", "docker", "", "", 0, false},
{"docker:///var/run/docker.sock", "docker", "", "docker:///var/run/docker.sock", 0, false},
{"docker://192.168.1.50:2375", "docker", "", "docker://192.168.1.50:2375", 0, false},
{"unix:///var/run/docker.sock", "docker", "", "unix:///var/run/docker.sock", 0, false},
{"tcp://127.0.0.1:2375", "docker", "", "tcp://127.0.0.1:2375", 0, false},
{"npipe:////./pipe/docker_engine", "docker", "", "npipe:////./pipe/docker_engine", 0, false},
{"tai://192.168.1.100", "tai", "192.168.1.100", "", 0, false},
{"tai://10.0.0.5:9200", "tai", "10.0.0.5", "", 9200, false},
{"tai://", "", "", "", 0, true},
{"ftp://host", "", "", "", 0, true},
{" tai://host ", "tai", "host", "", 0, false},
// Bare non-local host → auto-prepend tai://
{"192.168.1.50", "tai", "192.168.1.50", "", 0, false},
{"192.168.1.50:9200", "tai", "192.168.1.50", "", 9200, false},
{"my-server", "tai", "my-server", "", 0, false},
{"my-server:9200", "tai", "my-server", "", 9200, false},
}
for _, tt := range tests {
t.Run(tt.addr, func(t *testing.T) {
scheme, host, dockerAddr, err := parseAddr(tt.addr)
scheme, host, dockerAddr, grpcPort, err := parseAddr(tt.addr)
if (err != nil) != tt.wantErr {
t.Fatalf("err = %v, wantErr = %v", err, tt.wantErr)
}
@ -61,6 +72,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 +112,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 +133,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 +172,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)
}
@ -178,14 +202,15 @@ func TestNewRemoteK8s(t *testing.T) {
t.Skip("TAI_TEST_K8S_HOST or TAI_TEST_KUBECONFIG not set")
}
grpcPort := envPort("TAI_TEST_K8S_GRPC_PORT", envPort("TAI_TEST_GRPC_PORT", 9100))
ports := Ports{
K8s: envPort("TAI_TEST_K8S_PORT", 6443),
GRPC: envPort("TAI_TEST_GRPC_PORT", 9100),
HTTP: envPort("TAI_TEST_HTTP_PORT", 8080),
VNC: envPort("TAI_TEST_VNC_PORT", 6080),
GRPC: grpcPort,
HTTP: envPort("TAI_TEST_K8S_HTTP_PORT", 8080),
VNC: envPort("TAI_TEST_K8S_VNC_PORT", 6080),
}
c, err := New("tai://"+host, K8s,
c, err := New(fmt.Sprintf("tai://%s:%d", host, grpcPort), K8s,
WithPorts(ports),
WithKubeConfig(kubeconfig),
WithNamespace("default"),
@ -258,11 +283,39 @@ 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()
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)
if c.ports.GRPC == 0 {
t.Error("GRPC port should be discovered (non-zero)")
}
if c.ports.HTTP == 0 {
t.Error("HTTP port should be discovered (non-zero)")
}
}
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)
}
if c.ports.GRPC == 0 {
t.Error("GRPC port should still be discovered (non-zero)")
}
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)

600
workspace/DESIGN.md Normal file
View file

@ -0,0 +1,600 @@
# Workspace Design Document
> **Status**: Draft
> **Module**: `workspace` (top-level, parallel to `sandbox/v2`)
> **Depends on**: `tai` SDK (Volume, VolumeProvider, Sandbox), `sandbox/v2` Manager
---
## Overview
Workspace is a **first-class, persistent storage entity** independent of containers, chat sessions, and user sessions. It represents a user's project files — source code, configs, build artifacts — that can be mounted into any number of ephemeral containers.
Workspace is the **anchor point** for container scheduling: when a Workspace is created on a specific Tai node (host machine), all subsequent containers that reference it are automatically routed to the same node, because bind mounts require co-location on the same physical host.
---
## Problem
Current design: `Box.Workspace()` returns `workspace.FS` keyed by `box.id` — workspace and container are 1:1, same lifecycle. This couples file storage to container lifetime.
Real usage pattern:
```
User creates a project → uploads files → works on it across multiple chat sessions
→ attaches a long-running dev server → destroys/rebuilds containers freely
→ project files must survive all of this
```
Workspace must outlive containers. It is the persistent artifact; containers are disposable compute.
---
## Architecture
```
┌──────────────────────────────────────────────────┐
│ Application Layer │
│ │
│ Workspace Management UI Chat Interface │
│ ┌─────────────────────┐ ┌─────────────────┐ │
│ │ Create / Delete / UI │ │ Select Workspace│ │
│ │ Browse / Upload │ │ Start Chat │ │
│ └─────────┬───────────┘ └────────┬────────┘ │
│ │ │ │
└────────────┼─────────────────────────┼────────────┘
│ │
▼ ▼
┌──────────────────────────────────────────────────┐
│ Yao Engine │
│ │
│ workspace.Manager sandbox.Manager │
│ ┌────────────────┐ ┌─────────────────┐ │
│ │ CRUD │◄────────│ Mount workspace │ │
│ │ File I/O │ │ Route to node │ │
│ │ Node binding │ │ Create container│ │
│ └────────┬───────┘ └────────┬────────┘ │
│ │ │ │
└───────────┼───────────────────────────┼───────────┘
│ │
▼ ▼
┌──────────────────────────────────────────────────┐
│ Tai Node (Host) │
│ │
│ Volume gRPC Container Runtime │
│ ┌──────────────┐ ┌─────────────────────┐ │
│ │ ReadFile │ │ Container A (rw) │ │
│ │ WriteFile │ │ └─ /workspace ─┐ │ │
│ │ ListDir │ │ │ │ │
│ │ SyncPush/Pull │ │ Container B (ro) │ │ │
│ └──────┬───────┘ │ └─ /workspace ─┐│ │ │
│ │ └────────────────┼┼───┘ │
│ │ ││ │
│ ▼ ▼▼ │
│ ┌──────────────────────────────────────────┐ │
│ │ /data/ws/{workspace-id}/ │ │
│ │ ├── .workspace.json (metadata) │ │
│ │ ├── src/ │ │
│ │ ├── package.json │ │
│ │ └── ... │ │
│ └──────────────────────────────────────────┘ │
│ │
│ VolumeProvider │
│ ┌─────────────┬──────────────┬──────────────┐ │
│ │ BindMount │ DockerVolume │ K8s PVC │ │
│ │ (default) │ │ │ │
│ └─────────────┴──────────────┴──────────────┘ │
└───────────────────────────────────────────────────┘
```
---
## Core Design
### Node Binding
Workspace is physically stored on a Tai node's disk. **Bind mount requires Workspace and container to be on the same host.** Therefore:
- **Workspace binds to a specific Tai node at creation time.** This binding is immutable.
- When a container references a Workspace (`CreateOptions.WorkspaceID`), the container is **automatically routed to the same Tai node** — the caller does not (and should not) specify a Pool.
- One Tai node = one Pool = one host machine. These are equivalent in the current architecture.
```
创建 Workspace:
用户选择节点 "gpu-server" → workspace.Create(opts)
→ Tai "gpu-server" 上创建 /data/ws/ws-123/
创建容器（选了 Workspace）:
→ sandbox.Create(opts, WorkspaceID: "ws-123")
→ Manager 查到 ws-123 绑在 "gpu-server"
→ 自动路由到 "gpu-server" Pool
→ bind mount /data/ws/ws-123:/workspace:rw ✓ 同机
创建容器（没选 Workspace）:
→ 按原逻辑选 Pool（用户指定或默认）
```
This makes Workspace the **scheduling anchor**: once a Workspace is chosen, the node is determined.
### Workspace struct
```go
type Workspace struct {
ID string // unique identifier, e.g. "ws-abc123"
Name string // human-readable, e.g. "my-react-app"
Owner string // user ID
Node string // Tai node name (= Pool name); set at creation, immutable
Labels map[string]string // arbitrary metadata
CreatedAt time.Time
UpdatedAt time.Time
}
```
`Node` is the critical field: it pins this Workspace to a specific machine. All container operations referencing this Workspace are routed to this node.
No container references stored here. Workspace is pure storage — it doesn't know or care about containers.
### MountMode
```go
type MountMode string
const (
MountRW MountMode = "rw" // read-write (default)
MountRO MountMode = "ro" // read-only
)
```
Rules:
- A Workspace can be mounted by multiple containers simultaneously
- Each mount independently specifies `rw` or `ro`
- No write-lock enforcement — caller manages concurrency
- Default is `rw`
Rationale: In practice, Chat containers write source code and Runtime containers write build artifacts/logs — different files, no real conflict. Enforcing locks adds complexity without solving a real problem in this use case.
---
## API Design
### workspace.Manager
Workspace has its own manager, separate from `sandbox.Manager`. It owns Workspace CRUD and file I/O.
```go
package workspace
type Manager struct {
pools map[string]*tai.Client // node name → tai client (shared with sandbox.Manager)
}
// NewManager creates a workspace manager with the given pools.
// Pools are shared with sandbox.Manager — both reference the same tai.Client instances.
func NewManager(pools map[string]*tai.Client) *Manager
```
### Workspace CRUD
```go
type CreateOptions struct {
ID string // explicit ID; empty = auto-generate (uuid)
Name string // human-readable name
Owner string // user ID
Node string // target Tai node (required)
Labels map[string]string
}
type ListOptions struct {
Owner string // filter by owner; empty = all
Node string // filter by node; empty = all
}
// Create allocates storage on the target node and persists metadata.
func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Workspace, error)
// Get returns a workspace by ID.
// Checks the metadata file on the bound node.
func (m *Manager) Get(ctx context.Context, id string) (*Workspace, error)
// List returns workspaces, optionally filtered.
func (m *Manager) List(ctx context.Context, opts ListOptions) ([]*Workspace, error)
// Delete removes workspace storage from the node.
// Fails if containers currently mount it (unless force=true).
func (m *Manager) Delete(ctx context.Context, id string, force bool) error
// Update modifies workspace metadata (Name, Labels).
// Node and Owner are immutable after creation.
func (m *Manager) Update(ctx context.Context, id string, opts UpdateOptions) (*Workspace, error)
type UpdateOptions struct {
Name *string // nil = no change
Labels map[string]string // nil = no change; non-nil replaces all
}
```
### File I/O (no container needed)
File operations go through the Tai `Volume` gRPC service, using the Workspace ID as the session identifier. No container is needed.
```go
// FS returns an fs.FS view of the workspace, backed by Tai Volume gRPC.
func (m *Manager) FS(ctx context.Context, id string) (workspace.FS, error)
// ReadFile reads a file from the workspace.
func (m *Manager) ReadFile(ctx context.Context, id string, path string) ([]byte, error)
// WriteFile writes a file to the workspace.
func (m *Manager) WriteFile(ctx context.Context, id string, path string, data []byte, perm os.FileMode) error
// ListDir lists entries in a workspace directory.
func (m *Manager) ListDir(ctx context.Context, id string, path string) ([]DirEntry, error)
// Remove deletes a file or directory from the workspace.
func (m *Manager) Remove(ctx context.Context, id string, path string) error
// SyncPush uploads a local directory tree to the workspace.
func (m *Manager) SyncPush(ctx context.Context, id string, localPath string) error
// SyncPull downloads the workspace to a local directory.
func (m *Manager) SyncPull(ctx context.Context, id string, localPath string) error
```
These are thin wrappers around `tai.Client.Volume().{ReadFile,WriteFile,ListDir,...}` — the Tai SDK already implements all of these.
---
## Integration with Sandbox
### sandbox.CreateOptions changes
```go
type CreateOptions struct {
// ... existing fields ...
WorkspaceID string // workspace to mount; empty = no workspace
MountMode MountMode // "rw" (default) or "ro"
MountPath string // container path; default "/workspace"
}
```
### Container creation flow
When `WorkspaceID` is set in `CreateOptions`, the sandbox Manager:
```
Manager.Create(ctx, CreateOptions{
Image: "yaoapp/workspace:latest",
WorkspaceID: "ws-abc123",
MountMode: MountRW,
})
1. Validate CreateOptions (image required, etc.)
2. If WorkspaceID is set:
a. ws := workspaceManager.Get(ctx, workspaceID)
b. Force Pool = ws.Node (override any user-specified Pool)
c. spec := taiClient.VolumeProvider().MountSpec(workspaceID)
d. Inject mount into container create:
- Docker: opts.Binds = ["/data/ws/ws-abc123:/workspace:rw"]
- K8s: opts.Volumes + opts.VolumeMounts (PVC)
3. Create container via tai.Client.Sandbox().Create()
4. Start container
5. Return Box
```
### Box.Workspace() behavior change
```go
func (b *Box) Workspace() workspace.FS {
sessionID := b.workspaceID
if sessionID == "" {
sessionID = b.id // backward compatible
}
client, _ := b.manager.getPool(b.pool)
return client.Workspace(sessionID)
}
```
Multiple boxes mounting the same workspace -> same `sessionID` -> same files via Volume API.
---
## Metadata Storage
Workspace metadata (ID, Name, Owner, Node, Labels, timestamps) is stored as a JSON file inside the workspace directory.
### Storage path
```
/data/ws/{id}/.workspace.json
```
### Schema
```json
{
"id": "ws-abc123",
"name": "my-react-app",
"owner": "user-001",
"node": "gpu-server",
"labels": {"project": "frontend"},
"created_at": "2026-03-05T10:00:00Z",
"updated_at": "2026-03-05T12:30:00Z"
}
```
### Operations
| Operation | Implementation |
|-----------|---------------|
| Create | `Volume.WriteFile(id, ".workspace.json", json)` + `Volume.ResolvePath(id)` |
| Get | `Volume.ReadFile(id, ".workspace.json")` → unmarshal |
| List | `Volume.ListDir("")` → iterate dirs → read `.workspace.json` each |
| Update | Read → merge → `Volume.WriteFile(id, ".workspace.json", json)` |
| Delete | `Volume.Cleanup(id)` (removes entire dir) |
Phase 1 strategy: simple JSON files, zero external dependencies. Can migrate to SQLite or Yao's built-in DB if query/filter performance becomes a bottleneck.
---
## Node Management
### Listing available nodes
Application layer needs to present available nodes when user creates a Workspace. This comes from the sandbox Manager's pool configuration:
```go
// In workspace.Manager or sandbox.Manager
func (m *Manager) Nodes() []NodeInfo
type NodeInfo struct {
Name string // pool name = node name, e.g. "gpu-server"
Addr string // tai:// address
Online bool // is tai client connected
// Can be extended with capacity info later
}
```
### Dynamic node configuration
Nodes are configured at the application level (Yao settings/config). When a node is added or removed, both `workspace.Manager` and `sandbox.Manager` share the updated pool map. The Pool configuration API (from `sandbox/v2`) handles this — Workspace inherits it.
```
Application Config:
nodes:
- name: "local"
addr: "tai://localhost"
- name: "gpu-server"
addr: "tai://192.168.1.100:9527"
→ Both managers share:
pools["local"] = tai.Client("tai://localhost")
pools["gpu-server"] = tai.Client("tai://192.168.1.100:9527")
```
### Node failure handling
If a Tai node goes offline:
- Workspace CRUD for that node: returns error (node unreachable)
- Container creation referencing a Workspace on that node: returns error
- Workspaces on that node are not lost — data is still on the node's disk, will be available when node comes back online
- No automatic migration (Phase 1). Can add migration (rsync between nodes) later if needed.
---
## User Flows
### Flow 1: Workspace management UI
```
1. User opens Workspace management UI
→ API: workspace.List(owner: "user-001")
→ Returns list of workspaces with metadata
2. User creates workspace
→ UI shows available nodes (from Nodes() API)
→ User selects "gpu-server"
→ API: workspace.Create({ name: "my-project", node: "gpu-server" })
→ Directory /data/ws/ws-123/ created on gpu-server
→ .workspace.json written
3. User uploads files
→ API: workspace.WriteFile("ws-123", "src/main.go", data)
→ File written to /data/ws/ws-123/src/main.go via Volume gRPC
4. User browses files
→ API: workspace.ListDir("ws-123", "src/")
→ Returns file listing
5. User deletes workspace
→ API: workspace.Delete("ws-123")
→ Checks no active mounts → removes /data/ws/ws-123/
```
### Flow 2: Chat with Workspace
```
1. User opens Chat
→ Chat UI shows workspace selector
→ User picks "my-project" (ws-123, on node "gpu-server")
2. Agent needs a container:
→ sandbox.Create({
image: "yaoapp/workspace:latest",
workspace_id: "ws-123",
mount_mode: "rw",
})
→ Manager resolves ws-123.node = "gpu-server"
→ Container created on "gpu-server" Pool
→ -v /data/ws/ws-123:/workspace:rw
→ Agent can exec "ls /workspace/src/" inside container
3. Chat ends, container destroyed
→ Workspace files persist in /data/ws/ws-123/
4. User opens new Chat, selects same workspace
→ New container, same workspace, all files still there
```
### Flow 3: Long-running Runtime + Chat
```
1. User starts Runtime container for workspace:
→ sandbox.Create({
image: "node:20",
workspace_id: "ws-123",
mount_mode: "rw",
policy: "persistent",
ports: [{ container: 3000 }],
})
→ Container starts on "gpu-server"
→ -v /data/ws/ws-123:/workspace:rw
→ Inside: cd /workspace && npm install && npm run dev
2. User accesses dev server via proxy
→ box.Proxy(ctx, 3000, "/")
3. User opens Chat with same workspace:
→ Second container created on "gpu-server"
→ Same workspace mounted
→ Agent modifies source → Runtime hot-reloads
4. Chat ends, chat container destroyed
→ Runtime container keeps running
→ Workspace files persist
```
---
## Process & JSAPI
### Process registration
| Process | Args | Returns |
|---------|------|---------|
| `workspace.Create` | `options` (CreateOptions JSON) | Workspace |
| `workspace.Get` | `id` | Workspace |
| `workspace.List` | `options` (ListOptions JSON) | []Workspace |
| `workspace.Update` | `id`, `options` (UpdateOptions JSON) | Workspace |
| `workspace.Delete` | `id`, `force?` | — |
| `workspace.ReadFile` | `id`, `path` | file content |
| `workspace.WriteFile` | `id`, `path`, `data` | — |
| `workspace.ListDir` | `id`, `path` | []DirEntry |
| `workspace.Remove` | `id`, `path` | — |
| `workspace.Nodes` | — | []NodeInfo |
### JSAPI
```javascript
// Workspace CRUD
var ws = Workspace.Create({ name: "my-project", node: "gpu-server" })
var ws = Workspace.Get("ws-abc123")
var list = Workspace.List({ owner: "user-001" })
Workspace.Update("ws-abc123", { name: "new-name" })
Workspace.Delete("ws-abc123")
// File operations (no container needed)
var data = Workspace.ReadFile("ws-abc123", "src/main.go")
Workspace.WriteFile("ws-abc123", "src/main.go", "package main\n...")
var entries = Workspace.ListDir("ws-abc123", "src/")
Workspace.Remove("ws-abc123", "tmp.txt")
// List available nodes
var nodes = Workspace.Nodes()
// → [{ name: "local", addr: "tai://localhost", online: true },
// { name: "gpu-server", addr: "tai://192.168.1.100:9527", online: true }]
// Create container with workspace (via Sandbox API)
var sb = Sandbox("my-box", {
image: "node:20",
workspace_id: ws.id, // → auto-routes to ws.node
mount_mode: "rw",
})
```
---
## Storage Backend (Tai)
The `storage.VolumeProvider` interface in Tai Server already has three implementations:
```go
// tai/storage/provider.go
type VolumeProvider interface {
ResolvePath(sessionID string) (string, error)
MountSpec(sessionID string) MountConfig
Cleanup(sessionID string) error
}
type MountConfig struct {
Type string // "bind" | "volume" | "pvc"
Source string
Target string // always /workspace
}
```
| Provider | Backend | MountSpec | Status |
|----------|---------|-----------|--------|
| `BindMountProvider` | Host directory (`/data/ws/{id}/`) | `type:"bind"` | Implemented, default |
| `DockerVolumeProvider` | Docker named volume (`tai-{id}`) | `type:"volume"` | Implemented |
| `K8sPVCProvider` | K8s PVC (`tai-{id}-pvc`, 10Gi RWO) | `type:"pvc"` | Implemented |
Default is `BindMountProvider` for Docker environments (direct host path access for file CRUD). K8s environments use `K8sPVCProvider`.
The Tai `Volume` gRPC service (`ReadFile`, `WriteFile`, `ListDir`, etc.) already operates on the same `dataDir/{sessionID}/` paths. No additional work needed — Workspace file operations reuse existing Volume gRPC endpoints.
---
## Comparison: Before vs After
| Aspect | Before | After |
|--------|--------|-------|
| Workspace lifecycle | Tied to Box (same ID, same lifetime) | Independent entity, outlives containers |
| Workspace identity | `sessionID = box.id` | `sessionID = workspace.id` (explicit) |
| Container ↔ Workspace | 1:1, implicit | N:1, explicit via `CreateOptions.WorkspaceID` |
| Container scheduling | User picks Pool | Workspace determines Pool (node binding) |
| File persistence | Lost when container removed | Persists until workspace deleted |
| Multi-container access | Not possible | Multiple containers mount same workspace |
| Storage backend | Volume gRPC only (no mount) | Volume gRPC + bind mount into container |
| CRUD without container | Not possible | Via Volume API directly |
| Module status | Part of sandbox/v2 | Top-level module, parallel to sandbox/v2 |
---
## Implementation Plan
### Phase 1: Core (target: week 1-2)
| Task | Detail |
|------|--------|
| `workspace/workspace.go` | Workspace struct, MountMode, CreateOptions, metadata JSON read/write |
| `workspace/manager.go` | Manager with CRUD + file I/O (thin wrapper over tai Volume) |
| `workspace/manager_test.go` | Unit tests for CRUD and file operations |
| Node binding | `Workspace.Node` field, `Nodes()` API |
| `sandbox/v2` integration | `CreateOptions.WorkspaceID` → resolve node → force Pool → inject mount |
| `Box.Workspace()` update | Use `workspaceID` as sessionID when set |
### Phase 2: Wire into Tai (target: week 2-3)
| Task | Detail |
|------|--------|
| Tai Server: `VolumeProvider.MountSpec()` | Wire into container creation path |
| Tai gRPC: workspace metadata endpoints | Optional — can use Volume gRPC directly for Phase 1 |
| Process + JSAPI registration | `workspace.*` processes, JS bindings |
### Phase 3: Advanced (target: week 3+)
| Task | Detail |
|------|--------|
| Active mount tracking | Track which containers mount which workspaces |
| Delete safety | Refuse delete if active mounts exist |
| Workspace migration | rsync between nodes (stretch goal) |
| Quota / size limits | Per-workspace storage limits |
| Snapshot / backup | Workspace snapshots for rollback |
### Backward Compatibility
No breaking changes. Containers created without `WorkspaceID` work exactly as before:
- `sessionID = box.id`
- No bind mount
- Workspace FS backed by Volume gRPC as today

36
workspace/Makefile Normal file
View file

@ -0,0 +1,36 @@
GO ?= go
TEST_TIMEOUT ?= 120s
.PHONY: test test-v test-cover test-race
test:
$(GO) test -timeout=$(TEST_TIMEOUT) -count=1 ./...
test-v:
$(GO) test -v -timeout=$(TEST_TIMEOUT) -count=1 ./...
test-cover:
$(GO) test -v -timeout=$(TEST_TIMEOUT) -count=1 \
-coverprofile=coverage.out -covermode=count ./...
$(GO) tool cover -func=coverage.out | tail -1
test-race:
$(GO) test -race -v -timeout=$(TEST_TIMEOUT) -count=1 ./...
test-ci:
@echo "mode: count" > coverage.out
@for d in $$($(GO) list ./...); do \
$(GO) test -v -timeout=$(TEST_TIMEOUT) -count=1 \
-covermode=count -coverprofile=profile.out \
-coverpkg=$$d $$d > tmp.out; \
cat tmp.out; \
if grep -q "^--- FAIL" tmp.out; then \
rm -f tmp.out profile.out; \
exit 1; \
fi; \
if [ -f profile.out ]; then \
grep -v "mode:" profile.out >> coverage.out; \
rm profile.out; \
fi; \
rm -f tmp.out; \
done

74
workspace/TEST.md Normal file
View file

@ -0,0 +1,74 @@
# Workspace — Test Specification
Design: [DESIGN.md](./DESIGN.md)
## Principles
- **Black-box testing**: all `*_test.go` files use `package workspace_test` — tests only access exported API
- **No Docker required**: workspace unit tests use `volume.NewLocal(t.TempDir())` via `tai.WithVolume` — no Docker daemon needed
- **Skip when unavailable**: `skipIfNoTai(t)` for remote-mode tests
- **Tests follow implementation**: `*_test.go` lives next to the code it tests
- **Coverage > 80%**: per file and overall
## Prerequisites
No external services required for unit tests. Tests create a temp directory for storage.
### Remote mode (optional)
For remote-mode tests via Tai gRPC:
```bash
SANDBOX_TEST_REMOTE_ADDR=tai://127.0.0.1 go test -v ./workspace/
```
## Directory Structure
```
workspace/
├── workspace.go # Types (Workspace, CreateOptions, MountMode, etc.)
├── errors.go # Error definitions
├── manager.go # Manager (CRUD, file I/O, Nodes)
├── workspace_test.go # CRUD tests
├── fileio_test.go # File I/O + FS tests
├── testutils_test.go # Shared test helpers
├── DESIGN.md # Design document
├── TEST.md # This file
└── Makefile # Test runner
```
## testutils (internal to workspace_test)
```go
// testutils_test.go
package workspace_test
func setupManager(t *testing.T) *workspace.Manager
func setupManagerMultiNode(t *testing.T) *workspace.Manager
func localClient(t *testing.T, dataDir string) *tai.Client
func createTestWorkspace(t *testing.T, m *workspace.Manager, opts ...func(*workspace.CreateOptions)) *workspace.Workspace
func skipIfNoTai(t *testing.T)
```
## Required Test Cases
| File | Required Cases |
|------|---------------|
| `workspace_test.go` | Create / Create auto ID / Create explicit ID / Create with labels / Create invalid node / Create node not found / Get / Get not found / List / List filter owner / List filter node / Update name / Update labels / Update not found / Delete / Delete not found / Nodes / NodeForWorkspace / NodeForWorkspace not found |
| `fileio_test.go` | ReadWriteFile / WriteFile nested path / ListDir / Remove file / FS ReadFile / FS WriteFile / FS MkdirAll / FS Rename / FS WalkDir / FS Remove / FS not found |
## Running Tests
```bash
# All workspace tests (no Docker needed)
make -C workspace test
# Single test
go test -v ./workspace/ -run TestCreate
# With race detector
go test -race -v ./workspace/
# With coverage
go test -v -coverprofile=coverage.out ./workspace/
```

199
workspace/bench_test.go Normal file
View file

@ -0,0 +1,199 @@
package workspace_test
import (
"context"
"fmt"
"io/fs"
"testing"
"github.com/yaoapp/yao/workspace"
)
// BenchmarkWriteFile measures workspace file write latency.
func BenchmarkWriteFile(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForPool(b, pc)
ws := createWorkspace(b, m, pc.Name)
ctx := context.Background()
payload := []byte("package main\nfunc main() { println(\"bench\") }\n")
b.ResetTimer()
for i := 0; i < b.N; i++ {
if err := m.WriteFile(ctx, ws.ID, fmt.Sprintf("f%d.go", i), payload, 0644); err != nil {
b.Fatalf("WriteFile: %v", err)
}
}
})
}
}
// BenchmarkReadFile measures workspace file read latency.
func BenchmarkReadFile(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForPool(b, pc)
ws := createWorkspace(b, m, pc.Name)
ctx := context.Background()
if err := m.WriteFile(ctx, ws.ID, "bench.txt", []byte("benchmark data here"), 0644); err != nil {
b.Fatalf("setup WriteFile: %v", err)
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
data, err := m.ReadFile(ctx, ws.ID, "bench.txt")
if err != nil {
b.Fatalf("ReadFile: %v", err)
}
if len(data) == 0 {
b.Fatal("empty data")
}
}
})
}
}
// BenchmarkReadWriteCycle measures a full write-then-read cycle.
func BenchmarkReadWriteCycle(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForPool(b, pc)
ws := createWorkspace(b, m, pc.Name)
ctx := context.Background()
payload := []byte("package main\nfunc main() { println(\"cycle\") }\n")
b.ResetTimer()
for i := 0; i < b.N; i++ {
name := fmt.Sprintf("c%d.go", i)
if err := m.WriteFile(ctx, ws.ID, name, payload, 0644); err != nil {
b.Fatalf("WriteFile: %v", err)
}
data, err := m.ReadFile(ctx, ws.ID, name)
if err != nil {
b.Fatalf("ReadFile: %v", err)
}
if len(data) != len(payload) {
b.Fatalf("size mismatch: %d vs %d", len(data), len(payload))
}
}
})
}
}
// BenchmarkWriteLargeFile measures write throughput with a 1MB payload.
func BenchmarkWriteLargeFile(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForPool(b, pc)
ws := createWorkspace(b, m, pc.Name)
ctx := context.Background()
payload := make([]byte, 1<<20) // 1 MB
for i := range payload {
payload[i] = byte('A' + i%26)
}
b.SetBytes(int64(len(payload)))
b.ResetTimer()
for i := 0; i < b.N; i++ {
if err := m.WriteFile(ctx, ws.ID, fmt.Sprintf("large%d.bin", i), payload, 0644); err != nil {
b.Fatalf("WriteFile: %v", err)
}
}
})
}
}
// BenchmarkListDir measures directory listing latency (50 files).
func BenchmarkListDir(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForPool(b, pc)
ws := createWorkspace(b, m, pc.Name)
ctx := context.Background()
for i := 0; i < 50; i++ {
m.WriteFile(ctx, ws.ID, fmt.Sprintf("file%d.txt", i), []byte("x"), 0644)
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
entries, err := m.ListDir(ctx, ws.ID, ".")
if err != nil {
b.Fatalf("ListDir: %v", err)
}
if len(entries) < 50 {
b.Fatalf("expected >= 50 entries, got %d", len(entries))
}
}
})
}
}
// BenchmarkFSWalkDir measures fs.WalkDir performance over a directory tree (45+ entries).
func BenchmarkFSWalkDir(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForPool(b, pc)
ws := createWorkspace(b, m, pc.Name)
ctx := context.Background()
wfs, err := m.FS(ctx, ws.ID)
if err != nil {
b.Fatalf("FS: %v", err)
}
for _, dir := range []string{"src", "src/pkg", "src/cmd", "lib"} {
wfs.MkdirAll(dir, 0755)
}
for i := 0; i < 20; i++ {
wfs.WriteFile(fmt.Sprintf("src/f%d.go", i), []byte("package src"), 0644)
}
for i := 0; i < 10; i++ {
wfs.WriteFile(fmt.Sprintf("src/pkg/p%d.go", i), []byte("package pkg"), 0644)
}
for i := 0; i < 10; i++ {
wfs.WriteFile(fmt.Sprintf("lib/l%d.go", i), []byte("package lib"), 0644)
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
count := 0
fs.WalkDir(wfs, ".", func(_ string, _ fs.DirEntry, err error) error {
if err != nil {
return err
}
count++
return nil
})
if count < 40 {
b.Fatalf("walk returned only %d entries", count)
}
}
})
}
}
// BenchmarkCreateDelete measures workspace CRUD cycle.
func BenchmarkCreateDelete(b *testing.B) {
for _, pc := range testPools() {
b.Run(pc.Name, func(b *testing.B) {
m := setupManagerForPool(b, pc)
ctx := context.Background()
b.ResetTimer()
for i := 0; i < b.N; i++ {
ws, err := m.Create(ctx, workspace.CreateOptions{
Name: "bench-workspace",
Owner: "bench-user",
Node: pc.Name,
})
if err != nil {
b.Fatalf("Create: %v", err)
}
if err := m.Delete(ctx, ws.ID, true); err != nil {
b.Fatalf("Delete: %v", err)
}
}
})
}
}

10
workspace/errors.go Normal file
View file

@ -0,0 +1,10 @@
package workspace
import "errors"
var (
ErrNotFound = errors.New("workspace: not found")
ErrNodeMissing = errors.New("workspace: node is required")
ErrNodeOffline = errors.New("workspace: node is offline or not configured")
ErrHasMounts = errors.New("workspace: workspace has active container mounts")
)

232
workspace/fileio_test.go Normal file
View file

@ -0,0 +1,232 @@
package workspace_test
import (
"context"
"io/fs"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/workspace"
)
func TestReadWriteFile(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
ctx := context.Background()
err := m.WriteFile(ctx, ws.ID, "hello.txt", []byte("hello world"), 0644)
require.NoError(t, err)
data, err := m.ReadFile(ctx, ws.ID, "hello.txt")
require.NoError(t, err)
assert.Equal(t, "hello world", string(data))
})
}
}
func TestWriteFile_NestedPath(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
ctx := context.Background()
err := m.WriteFile(ctx, ws.ID, "src/main.go", []byte("package main"), 0644)
require.NoError(t, err)
data, err := m.ReadFile(ctx, ws.ID, "src/main.go")
require.NoError(t, err)
assert.Equal(t, "package main", string(data))
})
}
}
func TestListDir(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
ctx := context.Background()
require.NoError(t, m.WriteFile(ctx, ws.ID, "a.txt", []byte("a"), 0644))
require.NoError(t, m.WriteFile(ctx, ws.ID, "b.txt", []byte("b"), 0644))
entries, err := m.ListDir(ctx, ws.ID, ".")
require.NoError(t, err)
names := make(map[string]bool)
for _, e := range entries {
names[e.Name] = true
}
assert.True(t, names["a.txt"])
assert.True(t, names["b.txt"])
})
}
}
func TestRemoveFile(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
ctx := context.Background()
require.NoError(t, m.WriteFile(ctx, ws.ID, "tmp.txt", []byte("temp"), 0644))
err := m.Remove(ctx, ws.ID, "tmp.txt")
require.NoError(t, err)
_, err = m.ReadFile(ctx, ws.ID, "tmp.txt")
assert.Error(t, err)
})
}
}
func TestFS_ReadFile(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
ctx := context.Background()
require.NoError(t, m.WriteFile(ctx, ws.ID, "test.txt", []byte("via fs"), 0644))
wfs, err := m.FS(ctx, ws.ID)
require.NoError(t, err)
data, err := fs.ReadFile(wfs, "test.txt")
require.NoError(t, err)
assert.Equal(t, "via fs", string(data))
})
}
}
func TestFS_WriteFile(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
ctx := context.Background()
wfs, err := m.FS(ctx, ws.ID)
require.NoError(t, err)
err = wfs.WriteFile("from-fs.txt", []byte("written via fs"), 0644)
require.NoError(t, err)
data, err := m.ReadFile(ctx, ws.ID, "from-fs.txt")
require.NoError(t, err)
assert.Equal(t, "written via fs", string(data))
})
}
}
func TestFS_MkdirAll(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
ctx := context.Background()
wfs, err := m.FS(ctx, ws.ID)
require.NoError(t, err)
err = wfs.MkdirAll("a/b/c", 0755)
require.NoError(t, err)
info, err := fs.Stat(wfs, "a/b/c")
require.NoError(t, err)
assert.True(t, info.IsDir())
})
}
}
func TestFS_Rename(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
ctx := context.Background()
wfs, err := m.FS(ctx, ws.ID)
require.NoError(t, err)
err = wfs.WriteFile("old.txt", []byte("content"), 0644)
require.NoError(t, err)
err = wfs.Rename("old.txt", "new.txt")
require.NoError(t, err)
data, err := fs.ReadFile(wfs, "new.txt")
require.NoError(t, err)
assert.Equal(t, "content", string(data))
_, err = fs.ReadFile(wfs, "old.txt")
assert.Error(t, err)
})
}
}
func TestFS_WalkDir(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
ctx := context.Background()
wfs, err := m.FS(ctx, ws.ID)
require.NoError(t, err)
require.NoError(t, wfs.MkdirAll("src", 0755))
require.NoError(t, wfs.WriteFile("src/main.go", []byte("package main"), 0644))
require.NoError(t, wfs.WriteFile("src/util.go", []byte("package main"), 0644))
var files []string
err = fs.WalkDir(wfs, "src", func(path string, d fs.DirEntry, err error) error {
if err != nil {
return err
}
if !d.IsDir() {
files = append(files, path)
}
return nil
})
require.NoError(t, err)
assert.Len(t, files, 2)
})
}
}
func TestFS_Remove(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
ctx := context.Background()
wfs, err := m.FS(ctx, ws.ID)
require.NoError(t, err)
require.NoError(t, wfs.WriteFile("removeme.txt", []byte("bye"), 0644))
err = wfs.Remove("removeme.txt")
require.NoError(t, err)
_, err = fs.ReadFile(wfs, "removeme.txt")
assert.Error(t, err)
})
}
}
func TestFS_NotFound(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
_, err := m.FS(context.Background(), "nonexistent")
assert.ErrorIs(t, err, workspace.ErrNotFound)
})
}
}

130
workspace/jsapi/fs.go Normal file
View file

@ -0,0 +1,130 @@
package jsapi
import (
"rogchap.com/v8go"
)
// NewFSObject creates a JS WorkspaceFS object backed by a workspace ID string.
// All methods delegate to workspace.M() → FS — no Go object passed to V8.
//
// # Properties (read-only)
//
// ws.id → string // workspace ID ← workspaceID arg
// ws.name → string // workspace name ← Workspace.Name
// ws.node → string // tai node name ← Workspace.Node
//
// # Methods — Go mapping
//
// Each method internally does: fs, _ := workspace.M().FS(ctx, workspaceID)
// then calls the corresponding method on taiworkspace.FS.
//
// ws.ReadFile(path) → string
//
// Go: FS.ReadFile(name string) ([]byte, error)
// — also available via Manager.ReadFile(ctx, id, path)
// JS args: path string
// JS returns: string (UTF-8 content of the file)
//
// ws.WriteFile(path, data, perm?) → void
//
// Go: FS.WriteFile(name string, data []byte, perm os.FileMode) error
// — also available via Manager.WriteFile(ctx, id, path, data, perm)
// JS args: path string, data string|Uint8Array, perm? number (default 0644)
//
// ws.ReadDir(path?) → DirEntry[]
//
// Go: FS.ReadDir(name string) ([]fs.DirEntry, error)
// — also available via Manager.ListDir(ctx, id, path)
// JS args: path string (default ".")
// JS returns: [{
// name: string, ← DirEntry.Name()
// is_dir: boolean, ← DirEntry.IsDir()
// size: number ← DirEntry.Info().Size()
// }]
//
// ws.Stat(path) → FileInfo
//
// Go: FS.Stat(name string) (fs.FileInfo, error)
// JS args: path string
// JS returns: {
// name: string, ← FileInfo.Name()
// size: number, ← FileInfo.Size()
// is_dir: boolean, ← FileInfo.IsDir()
// mod_time: string ← FileInfo.ModTime() (ISO 8601)
// }
//
// ws.MkdirAll(path, perm?) → void
//
// Go: FS.MkdirAll(name string, perm os.FileMode) error
// JS args: path string, perm? number (default 0755)
//
// ws.Remove(path) → void
//
// Go: FS.Remove(name string) error
// — also available via Manager.Remove(ctx, id, path)
// JS args: path string (single file or empty directory)
//
// ws.RemoveAll(path) → void
//
// Go: FS.RemoveAll(name string) error
// JS args: path string (recursive removal)
//
// ws.Rename(from, to) → void
//
// Go: FS.Rename(oldname, newname string) error
// JS args: from string, to string
//
// # Base64 variants (PLANNED — not yet implemented)
//
// Avoids V8↔Go binary bridge overhead for images, archives, etc.
//
// ws.ReadFileBase64(path) → string
//
// Go: FS.ReadFile(name) → base64.StdEncoding.EncodeToString(data)
// JS args: path string
// JS returns: string (base64-encoded content)
//
// ws.WriteFileBase64(path, b64, perm?) → void
//
// Go: base64.StdEncoding.DecodeString(b64) → FS.WriteFile(name, data, perm)
// JS args: path string, b64 string, perm? number (default 0644)
//
// # Host copy (PLANNED — not yet implemented)
//
// Copy files/dirs from Yao host filesystem into the workspace volume.
// Useful for seeding workspaces with templates, config files, assets, etc.
//
// ws.CopyFromHost(hostPath, destPath?) → void
//
// Copies a single file or directory tree from the Yao host into the workspace.
// Go: read host file(s) → FS.WriteFile / FS.MkdirAll for each entry
// JS args: hostPath string (absolute path on Yao host),
// destPath? string (target path inside workspace, default basename of hostPath)
//
// ws.CopyFromHostArchive(hostPath, destPath?) → void
//
// For large directory trees: zip on host → transfer → unzip on Tai node.
// Requires Tai server-side unarchive support.
// Go: zip hostPath → tai Volume upload → tai unarchive at destPath
// JS args: hostPath string, destPath? string (default ".")
func NewFSObject(v8ctx *v8go.Context, workspaceID string) (*v8go.Value, error) {
// TODO: Phase 2 implementation
// 1. Create JS object via v8go.NewObjectTemplate
// 2. Set read-only properties: id, name, node (from workspace.M().Get(workspaceID))
// 3. Bind each method as FunctionTemplate:
// - ReadFile → workspace.M().FS(ctx, id).ReadFile(path)
// - WriteFile → workspace.M().FS(ctx, id).WriteFile(path, data, perm)
// - ReadDir → workspace.M().FS(ctx, id).ReadDir(path)
// - Stat → workspace.M().FS(ctx, id).Stat(path)
// - MkdirAll → workspace.M().FS(ctx, id).MkdirAll(path, perm)
// - Remove → workspace.M().FS(ctx, id).Remove(path)
// - RemoveAll → workspace.M().FS(ctx, id).RemoveAll(path)
// - Rename → workspace.M().FS(ctx, id).Rename(old, new)
//
// PLANNED (not yet implemented):
// - ReadFileBase64 → ReadFile + base64 encode in Go
// - WriteFileBase64 → base64 decode in Go + WriteFile
// - CopyFromHost → host fs.Read → FS.Write (file-by-file)
// - CopyFromHostArchive → zip on host → tai transfer → unzip (needs Tai support)
return nil, nil
}

129
workspace/jsapi/jsapi.go Normal file
View file

@ -0,0 +1,129 @@
// Package jsapi registers the workspace namespace into the Yao V8 runtime.
//
// All methods are static on the workspace object — no constructor.
//
// # JavaScript API
//
// const ws = workspace.Create({ name: "proj", owner: "user1", node: "default" })
// const ws = workspace.Get(id)
// ws.ReadFile("main.go") → string
// ws.WriteFile("out.txt", data) → void
// ws.ReadDir("src/") → [{ name, is_dir, size }]
// workspace.Delete(id) → void
//
// # Go mapping
//
// workspace.Create(opts) → Manager.Create(ctx, CreateOptions) → *Workspace → WorkspaceFS
// workspace.Get(id) → Manager.Get(ctx, id) → *Workspace → WorkspaceFS
// workspace.List(filter?) → Manager.List(ctx, ListOptions) → []*Workspace → WorkspaceInfo[]
// workspace.Delete(id) → Manager.Delete(ctx, id, false) → void
//
// Registration happens via init() — import with:
//
// _ "github.com/yaoapp/yao/workspace/jsapi"
package jsapi
import (
v8 "github.com/yaoapp/gou/runtime/v8"
"rogchap.com/v8go"
)
func init() {
v8.RegisterObject("workspace", ExportObject)
}
// ExportObject exports the workspace namespace object to V8.
func ExportObject(iso *v8go.Isolate) *v8go.ObjectTemplate {
obj := v8go.NewObjectTemplate(iso)
obj.Set("Create", v8go.NewFunctionTemplate(iso, wsCreate))
obj.Set("Get", v8go.NewFunctionTemplate(iso, wsGet))
obj.Set("List", v8go.NewFunctionTemplate(iso, wsList))
obj.Set("Delete", v8go.NewFunctionTemplate(iso, wsDelete))
return obj
}
// wsCreate: `workspace.Create(options)` → WorkspaceFS
//
// Go: Manager.Create(ctx, CreateOptions) (*Workspace, error)
//
// JS options → Go CreateOptions mapping:
//
// {
// id: string → CreateOptions.ID // optional; auto-generated if empty
// name: string → CreateOptions.Name // required, human-readable name
// owner: string → CreateOptions.Owner // required, user ID
// node: string → CreateOptions.Node // required, target Tai node
// labels: object → CreateOptions.Labels // optional, map[string]string
// }
//
// Returns: WorkspaceFS object (see fs.go)
func wsCreate(info *v8go.FunctionCallbackInfo) *v8go.Value {
// TODO: Phase 2
// 1. Parse options from info.Args()[0]
// 2. Validate required fields (name, owner, node)
// 3. ws := workspace.M().Create(ctx, opts)
// 4. Return NewFSObject(v8ctx, ws.ID)
return v8go.Undefined(info.Context().Isolate())
}
// wsGet: `workspace.Get(id)` → WorkspaceFS | null
//
// Go: Manager.Get(ctx, id) (*Workspace, error)
//
// Args:
//
// id: string — workspace ID
//
// Returns: WorkspaceFS object if found, null if not found
func wsGet(info *v8go.FunctionCallbackInfo) *v8go.Value {
// TODO: Phase 2
// 1. id = info.Args()[0].String()
// 2. ws, err := workspace.M().Get(ctx, id)
// 3. Return NewFSObject(v8ctx, id) or null
return v8go.Undefined(info.Context().Isolate())
}
// wsList: `workspace.List(filter?)` → WorkspaceInfo[]
//
// Go: Manager.List(ctx, ListOptions) ([]*Workspace, error)
//
// JS filter → Go ListOptions mapping:
//
// {
// owner: string → ListOptions.Owner // filter by owner; empty = all
// node: string → ListOptions.Node // filter by node; empty = all
// }
//
// Returns: WorkspaceInfo[] — each element:
//
// {
// id: string ← Workspace.ID
// name: string ← Workspace.Name
// owner: string ← Workspace.Owner
// node: string ← Workspace.Node
// labels: object ← Workspace.Labels
// created_at: string ← Workspace.CreatedAt (ISO 8601)
// updated_at: string ← Workspace.UpdatedAt (ISO 8601)
// }
func wsList(info *v8go.FunctionCallbackInfo) *v8go.Value {
// TODO: Phase 2
// 1. Parse optional filter from info.Args()[0]
// 2. list := workspace.M().List(ctx, opts)
// 3. Convert each *Workspace → JS object
// 4. Return JS array
return v8go.Undefined(info.Context().Isolate())
}
// wsDelete: `workspace.Delete(id)` → void
//
// Go: Manager.Delete(ctx, id string, force bool) error
//
// Args:
//
// id: string — workspace ID to remove (force = false)
func wsDelete(info *v8go.FunctionCallbackInfo) *v8go.Value {
// TODO: Phase 2
// 1. id = info.Args()[0].String()
// 2. workspace.M().Delete(ctx, id, false)
return v8go.Undefined(info.Context().Isolate())
}

319
workspace/manager.go Normal file
View file

@ -0,0 +1,319 @@
package workspace
import (
"context"
"fmt"
"os"
"sync"
"time"
"github.com/yaoapp/yao/tai"
taiworkspace "github.com/yaoapp/yao/tai/workspace"
)
// Manager owns workspace CRUD, file I/O, and node management.
// Pools are shared with sandbox.Manager — both reference the same tai.Client instances.
type Manager struct {
pools map[string]*tai.Client
mu sync.RWMutex
}
// NewManager creates a workspace manager with the given pools.
func NewManager(pools map[string]*tai.Client) *Manager {
if pools == nil {
pools = make(map[string]*tai.Client)
}
return &Manager{pools: pools}
}
// Create allocates storage on the target node and persists metadata.
func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Workspace, error) {
if opts.Node == "" {
return nil, ErrNodeMissing
}
client, err := m.getClient(opts.Node)
if err != nil {
return nil, err
}
id := opts.ID
if id == "" {
id = generateID()
}
now := time.Now().UTC()
ws := &Workspace{
ID: id,
Name: opts.Name,
Owner: opts.Owner,
Node: opts.Node,
Labels: opts.Labels,
CreatedAt: now,
UpdatedAt: now,
}
vol := client.Volume()
if err := vol.MkdirAll(ctx, id, "."); err != nil {
return nil, fmt.Errorf("workspace: create directory: %w", err)
}
data, err := marshalMeta(ws)
if err != nil {
return nil, err
}
if err := vol.WriteFile(ctx, id, metadataFile, data, 0644); err != nil {
return nil, fmt.Errorf("workspace: write metadata: %w", err)
}
return ws, nil
}
// Get returns a workspace by ID.
// If the node is unknown, scans all pools.
func (m *Manager) Get(ctx context.Context, id string) (*Workspace, error) {
m.mu.RLock()
defer m.mu.RUnlock()
for nodeName, client := range m.pools {
ws, err := m.readMeta(ctx, client, id)
if err != nil {
continue
}
if ws.Node == "" {
ws.Node = nodeName
}
return ws, nil
}
return nil, ErrNotFound
}
// List returns workspaces, optionally filtered by owner and/or node.
func (m *Manager) List(ctx context.Context, opts ListOptions) ([]*Workspace, error) {
m.mu.RLock()
defer m.mu.RUnlock()
var result []*Workspace
for nodeName, client := range m.pools {
if opts.Node != "" && nodeName != opts.Node {
continue
}
entries, err := client.Volume().ListDir(ctx, "", ".")
if err != nil {
continue
}
for _, e := range entries {
if !e.IsDir {
continue
}
ws, err := m.readMeta(ctx, client, e.Path)
if err != nil {
continue
}
if ws.Node == "" {
ws.Node = nodeName
}
if opts.Owner != "" && ws.Owner != opts.Owner {
continue
}
result = append(result, ws)
}
}
return result, nil
}
// Update modifies workspace metadata (Name, Labels).
// Node and Owner are immutable after creation.
func (m *Manager) Update(ctx context.Context, id string, opts UpdateOptions) (*Workspace, error) {
ws, client, err := m.resolve(ctx, id)
if err != nil {
return nil, err
}
if opts.Name != nil {
ws.Name = *opts.Name
}
if opts.Labels != nil {
ws.Labels = opts.Labels
}
ws.UpdatedAt = time.Now().UTC()
data, err := marshalMeta(ws)
if err != nil {
return nil, err
}
if err := client.Volume().WriteFile(ctx, id, metadataFile, data, 0644); err != nil {
return nil, fmt.Errorf("workspace: write metadata: %w", err)
}
return ws, nil
}
// Delete removes workspace storage from the node.
func (m *Manager) Delete(ctx context.Context, id string, force bool) error {
_, client, err := m.resolve(ctx, id)
if err != nil {
return err
}
vol := client.Volume()
if err := vol.Remove(ctx, id, ".", true); err != nil {
return fmt.Errorf("workspace: remove: %w", err)
}
return nil
}
// Nodes returns all configured Tai nodes with their online status.
func (m *Manager) Nodes() []NodeInfo {
m.mu.RLock()
defer m.mu.RUnlock()
nodes := make([]NodeInfo, 0, len(m.pools))
for name := range m.pools {
nodes = append(nodes, NodeInfo{
Name: name,
Online: true,
})
}
return nodes
}
// FS returns an fs.FS-compatible filesystem for the given workspace.
func (m *Manager) FS(ctx context.Context, id string) (taiworkspace.FS, error) {
ws, client, err := m.resolve(ctx, id)
if err != nil {
return nil, err
}
_ = ws
return client.Workspace(id), nil
}
// ReadFile reads a file from the workspace.
func (m *Manager) ReadFile(ctx context.Context, id string, path string) ([]byte, error) {
_, client, err := m.resolve(ctx, id)
if err != nil {
return nil, err
}
data, _, err := client.Volume().ReadFile(ctx, id, path)
return data, err
}
// WriteFile writes a file to the workspace.
func (m *Manager) WriteFile(ctx context.Context, id string, path string, data []byte, perm os.FileMode) error {
_, client, err := m.resolve(ctx, id)
if err != nil {
return err
}
return client.Volume().WriteFile(ctx, id, path, data, perm)
}
// ListDir lists entries in a workspace directory.
func (m *Manager) ListDir(ctx context.Context, id string, path string) ([]DirEntry, error) {
_, client, err := m.resolve(ctx, id)
if err != nil {
return nil, err
}
entries, err := client.Volume().ListDir(ctx, id, path)
if err != nil {
return nil, err
}
result := make([]DirEntry, len(entries))
for i, e := range entries {
result[i] = DirEntry{
Name: e.Path,
IsDir: e.IsDir,
Size: e.Size,
}
}
return result, nil
}
// Remove deletes a file or directory from the workspace.
func (m *Manager) Remove(ctx context.Context, id string, path string) error {
_, client, err := m.resolve(ctx, id)
if err != nil {
return err
}
return client.Volume().Remove(ctx, id, path, true)
}
// AddPool registers a new Tai node.
func (m *Manager) AddPool(name string, client *tai.Client) {
m.mu.Lock()
defer m.mu.Unlock()
m.pools[name] = client
}
// RemovePool unregisters a Tai node.
func (m *Manager) RemovePool(name string) {
m.mu.Lock()
defer m.mu.Unlock()
delete(m.pools, name)
}
// NodeForWorkspace returns the node name for a given workspace ID.
// Used by sandbox.Manager to route container creation to the correct pool.
func (m *Manager) NodeForWorkspace(ctx context.Context, id string) (string, error) {
ws, _, err := m.resolve(ctx, id)
if err != nil {
return "", err
}
return ws.Node, nil
}
// MountPath returns the host-side directory path for a workspace,
// suitable for use as a Docker bind mount source.
// For local volumes this is dataDir/{id}; for remote (Tai) the server handles mounts.
func (m *Manager) MountPath(ctx context.Context, id string) (string, error) {
_, client, err := m.resolve(ctx, id)
if err != nil {
return "", err
}
dataDir := client.DataDir()
if dataDir == "" {
return "", nil
}
return dataDir + "/" + id, nil
}
// --- internal ---
func (m *Manager) getClient(node string) (*tai.Client, error) {
m.mu.RLock()
defer m.mu.RUnlock()
client, ok := m.pools[node]
if !ok {
return nil, ErrNodeOffline
}
return client, nil
}
// resolve finds the workspace and its tai.Client by scanning pools.
func (m *Manager) resolve(ctx context.Context, id string) (*Workspace, *tai.Client, error) {
m.mu.RLock()
defer m.mu.RUnlock()
for _, client := range m.pools {
ws, err := m.readMeta(ctx, client, id)
if err != nil {
continue
}
return ws, client, nil
}
return nil, nil, ErrNotFound
}
func (m *Manager) readMeta(ctx context.Context, client *tai.Client, id string) (*Workspace, error) {
data, _, err := client.Volume().ReadFile(ctx, id, metadataFile)
if err != nil {
return nil, err
}
return unmarshalMeta(data)
}
// DirEntry represents a file or directory entry in a workspace listing.
type DirEntry struct {
Name string `json:"name"`
IsDir bool `json:"is_dir"`
Size int64 `json:"size"`
}

89
workspace/testutils_test.go Normal file
View file

@ -0,0 +1,89 @@
package workspace_test
import (
"context"
"os"
"testing"
"time"
"github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/volume"
"github.com/yaoapp/yao/workspace"
)
type poolConfig struct {
Name string
Addr string
}
func testPools() []poolConfig {
pools := []poolConfig{
{Name: "local", Addr: "local"},
}
if addr := os.Getenv("SANDBOX_TEST_REMOTE_ADDR"); addr != "" {
pools = append(pools, poolConfig{Name: "remote", Addr: addr})
}
return pools
}
func setupManagerForPool(tb testing.TB, pc poolConfig) *workspace.Manager {
tb.Helper()
client := clientForPool(tb, pc)
pools := map[string]*tai.Client{pc.Name: client}
return workspace.NewManager(pools)
}
func clientForPool(tb testing.TB, pc poolConfig) *tai.Client {
tb.Helper()
if pc.Addr == "local" {
return localClient(tb, tb.TempDir())
}
client, err := tai.New(pc.Addr)
if err != nil {
tb.Fatalf("tai.New(%s): %v", pc.Addr, err)
}
tb.Cleanup(func() { client.Close() })
return client
}
func localClient(tb testing.TB, dataDir string) *tai.Client {
tb.Helper()
vol := volume.NewLocal(dataDir)
client, err := tai.New("local", tai.WithVolume(vol), tai.WithDataDir(dataDir))
if err != nil {
tb.Fatalf("tai.New local: %v", err)
}
tb.Cleanup(func() { client.Close() })
return client
}
func setupManagerMultiNode(t *testing.T) *workspace.Manager {
t.Helper()
pools := map[string]*tai.Client{
"node-a": localClient(t, t.TempDir()),
"node-b": localClient(t, t.TempDir()),
}
return workspace.NewManager(pools)
}
func createWorkspace(tb testing.TB, m *workspace.Manager, node string, opts ...func(*workspace.CreateOptions)) *workspace.Workspace {
tb.Helper()
co := workspace.CreateOptions{
Name: "test-workspace",
Owner: "test-user",
Node: node,
}
for _, fn := range opts {
fn(&co)
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
ws, err := m.Create(ctx, co)
if err != nil {
tb.Fatalf("Create workspace: %v", err)
}
tb.Cleanup(func() {
m.Delete(context.Background(), ws.ID, true)
})
return ws
}

78
workspace/workspace.go Normal file
View file

@ -0,0 +1,78 @@
package workspace
import (
"encoding/json"
"fmt"
"time"
"github.com/google/uuid"
)
// MountMode controls read-write or read-only access when a workspace is
// bind-mounted into a container.
type MountMode string
const (
MountRW MountMode = "rw"
MountRO MountMode = "ro"
)
const metadataFile = ".workspace.json"
// Workspace is a persistent, user-managed storage entity.
// It is pinned to a specific Tai node (host machine) at creation time;
// containers referencing this workspace are automatically routed to that node.
type Workspace struct {
ID string `json:"id"`
Name string `json:"name"`
Owner string `json:"owner"`
Node string `json:"node"`
Labels map[string]string `json:"labels,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// CreateOptions configures a new workspace.
type CreateOptions struct {
ID string // explicit ID; empty = auto-generate (uuid)
Name string // human-readable name
Owner string // user ID
Node string // target Tai node (required)
Labels map[string]string // arbitrary metadata
}
// ListOptions filters workspace listing.
type ListOptions struct {
Owner string // filter by owner; empty = all
Node string // filter by node; empty = all
}
// UpdateOptions specifies which metadata fields to change.
// nil fields are left unchanged. Node and Owner are immutable.
type UpdateOptions struct {
Name *string // nil = no change
Labels map[string]string // nil = no change; non-nil replaces all labels
}
// NodeInfo describes a Tai node available for workspace storage.
type NodeInfo struct {
Name string // pool name = node name
Addr string // tai:// address
Online bool // tai client is connected
}
func generateID() string {
return fmt.Sprintf("ws-%s", uuid.New().String()[:12])
}
func marshalMeta(ws *Workspace) ([]byte, error) {
return json.MarshalIndent(ws, "", " ")
}
func unmarshalMeta(data []byte) (*Workspace, error) {
var ws Workspace
if err := json.Unmarshal(data, &ws); err != nil {
return nil, fmt.Errorf("workspace: invalid metadata: %w", err)
}
return &ws, nil
}

323
workspace/workspace_test.go Normal file
View file

@ -0,0 +1,323 @@
package workspace_test
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/volume"
"github.com/yaoapp/yao/workspace"
)
func TestCreate(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
assert.NotEmpty(t, ws.ID)
assert.Equal(t, "test-workspace", ws.Name)
assert.Equal(t, "test-user", ws.Owner)
assert.Equal(t, pc.Name, ws.Node)
assert.False(t, ws.CreatedAt.IsZero())
assert.False(t, ws.UpdatedAt.IsZero())
})
}
}
func TestCreate_AutoID(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
assert.True(t, len(ws.ID) > 0)
assert.Contains(t, ws.ID, "ws-")
})
}
}
func TestCreate_ExplicitID(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name, func(co *workspace.CreateOptions) {
co.ID = "my-custom-id"
})
assert.Equal(t, "my-custom-id", ws.ID)
})
}
}
func TestCreate_WithLabels(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name, func(co *workspace.CreateOptions) {
co.Labels = map[string]string{"project": "frontend", "env": "dev"}
})
assert.Equal(t, "frontend", ws.Labels["project"])
assert.Equal(t, "dev", ws.Labels["env"])
})
}
}
func TestCreate_InvalidNode(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
_, err := m.Create(context.Background(), workspace.CreateOptions{
Name: "bad",
Owner: "user",
Node: "",
})
assert.ErrorIs(t, err, workspace.ErrNodeMissing)
})
}
}
func TestCreate_NodeNotFound(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
_, err := m.Create(context.Background(), workspace.CreateOptions{
Name: "bad",
Owner: "user",
Node: "nonexistent-node",
})
assert.ErrorIs(t, err, workspace.ErrNodeOffline)
})
}
}
func TestGet(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
got, err := m.Get(context.Background(), ws.ID)
require.NoError(t, err)
assert.Equal(t, ws.ID, got.ID)
assert.Equal(t, ws.Name, got.Name)
assert.Equal(t, ws.Owner, got.Owner)
assert.Equal(t, ws.Node, got.Node)
})
}
}
func TestGet_NotFound(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")
assert.ErrorIs(t, err, workspace.ErrNotFound)
})
}
}
func TestList(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
createWorkspace(t, m, pc.Name, func(co *workspace.CreateOptions) { co.Name = "ws-1" })
createWorkspace(t, m, pc.Name, func(co *workspace.CreateOptions) { co.Name = "ws-2" })
list, err := m.List(context.Background(), workspace.ListOptions{})
require.NoError(t, err)
assert.GreaterOrEqual(t, len(list), 2)
})
}
}
func TestList_FilterOwner(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
createWorkspace(t, m, pc.Name, func(co *workspace.CreateOptions) {
co.Owner = "alice"
co.Name = "alice-ws"
})
createWorkspace(t, m, pc.Name, func(co *workspace.CreateOptions) {
co.Owner = "bob"
co.Name = "bob-ws"
})
list, err := m.List(context.Background(), workspace.ListOptions{Owner: "alice"})
require.NoError(t, err)
assert.Len(t, list, 1)
assert.Equal(t, "alice", list[0].Owner)
})
}
}
func TestList_FilterNode(t *testing.T) {
m := setupManagerMultiNode(t)
ctx := context.Background()
_, err := m.Create(ctx, workspace.CreateOptions{Name: "a", Owner: "u", Node: "node-a"})
require.NoError(t, err)
_, err = m.Create(ctx, workspace.CreateOptions{Name: "b", Owner: "u", Node: "node-b"})
require.NoError(t, err)
list, err := m.List(ctx, workspace.ListOptions{Node: "node-a"})
require.NoError(t, err)
assert.Len(t, list, 1)
assert.Equal(t, "node-a", list[0].Node)
}
func TestUpdate_Name(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
newName := "renamed-workspace"
updated, err := m.Update(context.Background(), ws.ID, workspace.UpdateOptions{
Name: &newName,
})
require.NoError(t, err)
assert.Equal(t, newName, updated.Name)
assert.Equal(t, ws.Owner, updated.Owner)
assert.True(t, updated.UpdatedAt.After(ws.UpdatedAt) || updated.UpdatedAt.Equal(ws.UpdatedAt))
})
}
}
func TestUpdate_Labels(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name, func(co *workspace.CreateOptions) {
co.Labels = map[string]string{"old": "value"}
})
updated, err := m.Update(context.Background(), ws.ID, workspace.UpdateOptions{
Labels: map[string]string{"new": "label"},
})
require.NoError(t, err)
assert.Equal(t, "label", updated.Labels["new"])
assert.Empty(t, updated.Labels["old"])
})
}
}
func TestUpdate_NotFound(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
_, err := m.Update(context.Background(), "nonexistent", workspace.UpdateOptions{})
assert.ErrorIs(t, err, workspace.ErrNotFound)
})
}
}
func TestDelete(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws, err := m.Create(context.Background(), workspace.CreateOptions{
Name: "to-delete", Owner: "user", Node: pc.Name,
})
require.NoError(t, err)
err = m.Delete(context.Background(), ws.ID, false)
require.NoError(t, err)
_, err = m.Get(context.Background(), ws.ID)
assert.ErrorIs(t, err, workspace.ErrNotFound)
})
}
}
func TestDelete_NotFound(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
err := m.Delete(context.Background(), "nonexistent", false)
assert.ErrorIs(t, err, workspace.ErrNotFound)
})
}
}
func TestNodes(t *testing.T) {
m := setupManagerMultiNode(t)
nodes := m.Nodes()
assert.Len(t, nodes, 2)
names := make(map[string]bool)
for _, n := range nodes {
names[n.Name] = true
assert.True(t, n.Online)
}
assert.True(t, names["node-a"])
assert.True(t, names["node-b"])
}
func TestNodeForWorkspace(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
ws := createWorkspace(t, m, pc.Name)
node, err := m.NodeForWorkspace(context.Background(), ws.ID)
require.NoError(t, err)
assert.Equal(t, pc.Name, node)
})
}
}
func TestNodeForWorkspace_NotFound(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
_, err := m.NodeForWorkspace(context.Background(), "nonexistent")
assert.ErrorIs(t, err, workspace.ErrNotFound)
})
}
}
func TestAddPool(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForPool(t, pc)
assert.Len(t, m.Nodes(), 1)
vol := volume.NewLocal(t.TempDir())
client, err := tai.New("local", tai.WithVolume(vol))
require.NoError(t, err)
defer client.Close()
m.AddPool("new-node", client)
assert.Len(t, m.Nodes(), 2)
})
}
}
func TestRemovePool(t *testing.T) {
m := setupManagerMultiNode(t)
assert.Len(t, m.Nodes(), 2)
m.RemovePool("node-b")
assert.Len(t, m.Nodes(), 1)
}
func TestMountPath(t *testing.T) {
m := setupManagerForPool(t, poolConfig{Name: "local", Addr: "local"})
ws := createWorkspace(t, m, "local")
mountPath, err := m.MountPath(context.Background(), ws.ID)
require.NoError(t, err)
assert.Contains(t, mountPath, ws.ID)
}
func TestMountPath_NotFound(t *testing.T) {
m := setupManagerForPool(t, poolConfig{Name: "local", Addr: "local"})
_, err := m.MountPath(context.Background(), "nonexistent")
assert.ErrorIs(t, err, workspace.ErrNotFound)
}