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

Merge pull request #1489 from trheyi/main

Browse source 
PR Title: Refactor server lifecycle, unify gRPC client, and enhance sandbox streaming with Tai 1.2.0
This commit is contained in:
Max 2026-03-08 12:38:21 +08:00 • committed by GitHub
commit 793c9f9b51
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
76 changed files with 4740 additions and 3979 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -924,7 +924,8 @@ jobs:
});
# =============================================================================
# Sandbox V2 Tests (tai + sandbox/v2 + workspace, Docker + K8s via k3d)
# Sandbox V2 Tests (tai SDK + workspace, Docker + K8s via k3d)
# Full sandbox/v2 integration tests are run locally.
# =============================================================================
SandboxV2Test:
runs-on: ubuntu-latest
@ -941,9 +942,6 @@ jobs:
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
@ -985,7 +983,7 @@ jobs:
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: issue_number,
body: '🤖 Sandbox V2 Tests running (tai + sandbox-v2 + workspace)...'
body: '🤖 Sandbox V2 CI Tests running (tai + workspace)...'
});
- name: Checkout Kun
@ -1072,8 +1070,8 @@ jobs:
- name: Pull Test Images
run: |
docker pull yaoapp/sandbox-v2-test:latest || true
docker pull yaoapp/tai:latest
docker pull yaoapp/tai-sandbox-test:latest || true
docker pull yaoapp/tai:1.2.0
docker pull alpine:latest
- name: Install k3d
@ -1089,26 +1087,27 @@ jobs:
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
-p 19100:19100 -p 8099:8099 -p 12375:12375 -p 16080:16080 \
yaoapp/tai:1.2.0 server \
-grpc 0.0.0.0:19100 -http 0.0.0.0:8099 -vnc 0.0.0.0:16080 -docker 0.0.0.0:12375
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1; then
if curl -sf http://127.0.0.1:8099/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 || {
curl -sf http://127.0.0.1:8099/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
if nc -z 127.0.0.1 19100 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 || {
nc -z 127.0.0.1 19100 2>/dev/null || {
echo "::error::Tai Docker gRPC failed"; docker logs tai-docker 2>&1; exit 1
}
@ -1138,45 +1137,49 @@ jobs:
docker run -d --name tai-k8s \
--network k3d-tai-test \
-p 8081:8080 -p 9101:9100 -p 6443:6443 -p 6081:6080 \
-p 19101:19100 -p 8100:8099 -p 6443:16443 -p 16081:16080 \
-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
yaoapp/tai:1.2.0 server \
-grpc 0.0.0.0:19100 -http 0.0.0.0:8099 -vnc 0.0.0.0:16080 -k8s 0.0.0.0:16443
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:8081/healthz > /dev/null 2>&1; then
if curl -sf http://127.0.0.1:8100/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 || {
curl -sf http://127.0.0.1:8100/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
if nc -z 127.0.0.1 19101 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 || {
nc -z 127.0.0.1 19101 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)
- name: Run Sandbox V2 CI Tests (tai + 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_DOCKER: "tcp://127.0.0.1:12375"
TAI_TEST_GRPC_PORT: "19100"
TAI_TEST_HTTP_PORT: "8099"
TAI_TEST_VNC_PORT: "16080"
TAI_TEST_DOCKER_PORT: "12375"
TAI_TEST_K8S_HOST: "127.0.0.1"
TAI_TEST_K8S_PORT: "6443"
TAI_TEST_K8S_GRPC_PORT: "9101"
TAI_TEST_K8S_GRPC_PORT: "19101"
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"
SANDBOX_TEST_REMOTE_ADDR: "tai://127.0.0.1:19100"
SANDBOX_TEST_IMAGE: "yaoapp/tai-sandbox-test:latest"
run: make unit-test-sandbox-v2
- name: Codecov Report
@ -1197,7 +1200,7 @@ jobs:
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: issue_number,
body: '✅ Sandbox V2 Tests passed (tai + sandbox-v2 + workspace)!'
body: '✅ Sandbox V2 CI Tests passed (tai + workspace)!'
});
# =============================================================================
@ -1809,228 +1812,6 @@ jobs:
body: '✅ Registry Client SDK Tests passed!'
});
# =============================================================================
# Benchmark: Sandbox V2 + Workspace (parallel with SandboxV2Test, non-blocking)
# =============================================================================
BenchmarkSandboxV2:
runs-on: ubuntu-latest
strategy:
matrix:
go: ["1.25"]
if: >
${{ github.event.workflow_run.event == 'pull_request' &&
github.event.workflow_run.conclusion == 'success' }}
steps:
- name: "Download artifact"
uses: actions/github-script@v7
with:
script: |
var artifacts = await github.rest.actions.listWorkflowRunArtifacts({
owner: context.repo.owner,
repo: context.repo.repo,
run_id: ${{github.event.workflow_run.id }},
});
var matchArtifact = artifacts.data.artifacts.filter((artifact) => {
return artifact.name == "pr"
})[0];
var download = await github.rest.actions.downloadArtifact({
owner: context.repo.owner,
repo: context.repo.repo,
artifact_id: matchArtifact.id,
archive_format: 'zip',
});
var fs = require('fs');
fs.writeFileSync('${{github.workspace}}/pr.zip', Buffer.from(download.data));
- name: "Read NR & SHA"
run: |
unzip pr.zip
cat NR
cat SHA
echo HEAD=$(cat SHA) >> $GITHUB_ENV
echo NR=$(cat NR) >> $GITHUB_ENV
- name: 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 (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-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-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
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 Benchmarks
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 benchmark-sandbox-v2
# =============================================================================
# gRPC Tests - Run once with SQLite (transport layer, no DB matrix needed)
# =============================================================================

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

@ -680,7 +680,8 @@ jobs:
token: ${{ secrets.CODECOV_TOKEN }}
# =============================================================================
# Sandbox V2 Tests (tai + sandbox/v2 + workspace, Docker + K8s via k3d)
# Sandbox V2 Tests (tai SDK + workspace, Docker + K8s via k3d)
# Full sandbox/v2 integration tests are run locally.
# =============================================================================
sandbox-v2-test:
runs-on: ubuntu-latest
@ -780,8 +781,8 @@ jobs:
- name: Pull Test Images
run: |
docker pull yaoapp/sandbox-v2-test:latest || true
docker pull yaoapp/tai:latest
docker pull yaoapp/tai-sandbox-test:latest || true
docker pull yaoapp/tai:1.2.0
docker pull alpine:latest
- name: Install k3d
@ -797,26 +798,27 @@ jobs:
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
-p 19100:19100 -p 8099:8099 -p 12375:12375 -p 16080:16080 \
yaoapp/tai:1.2.0 server \
-grpc 0.0.0.0:19100 -http 0.0.0.0:8099 -vnc 0.0.0.0:16080 -docker 0.0.0.0:12375
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1; then
if curl -sf http://127.0.0.1:8099/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 || {
curl -sf http://127.0.0.1:8099/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
if nc -z 127.0.0.1 19100 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 || {
nc -z 127.0.0.1 19100 2>/dev/null || {
echo "::error::Tai Docker gRPC failed"; docker logs tai-docker 2>&1; exit 1
}
@ -846,45 +848,49 @@ jobs:
docker run -d --name tai-k8s \
--network k3d-tai-test \
-p 8081:8080 -p 9101:9100 -p 6443:6443 -p 6081:6080 \
-p 19101:19100 -p 8100:8099 -p 6443:16443 -p 16081:16080 \
-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
yaoapp/tai:1.2.0 server \
-grpc 0.0.0.0:19100 -http 0.0.0.0:8099 -vnc 0.0.0.0:16080 -k8s 0.0.0.0:16443
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:8081/healthz > /dev/null 2>&1; then
if curl -sf http://127.0.0.1:8100/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 || {
curl -sf http://127.0.0.1:8100/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
if nc -z 127.0.0.1 19101 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 || {
nc -z 127.0.0.1 19101 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)
- name: Run Sandbox V2 CI Tests (tai + 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_DOCKER: "tcp://127.0.0.1:12375"
TAI_TEST_GRPC_PORT: "19100"
TAI_TEST_HTTP_PORT: "8099"
TAI_TEST_VNC_PORT: "16080"
TAI_TEST_DOCKER_PORT: "12375"
TAI_TEST_K8S_HOST: "127.0.0.1"
TAI_TEST_K8S_PORT: "6443"
TAI_TEST_K8S_GRPC_PORT: "9101"
TAI_TEST_K8S_GRPC_PORT: "19101"
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"
SANDBOX_TEST_REMOTE_ADDR: "tai://127.0.0.1:19100"
SANDBOX_TEST_IMAGE: "yaoapp/tai-sandbox-test:latest"
run: make unit-test-sandbox-v2
- name: Codecov Report
@ -1349,194 +1355,6 @@ jobs:
with:
token: ${{ secrets.CODECOV_TOKEN }}
# =============================================================================
# Benchmark: Sandbox V2 + Workspace (parallel with sandbox-v2-test)
# =============================================================================
benchmark-sandbox-v2:
runs-on: ubuntu-latest
strategy:
matrix:
go: ["1.25"]
steps:
- name: Checkout Kun
uses: actions/checkout@v4
with:
repository: ${{ env.REPO_KUN }}
path: kun
- name: Checkout Xun
uses: actions/checkout@v4
with:
repository: ${{ env.REPO_XUN }}
path: xun
- name: Checkout Gou
uses: actions/checkout@v4
with:
repository: ${{ env.REPO_GOU }}
path: gou
- name: Checkout V8Go
uses: actions/checkout@v4
with:
repository: yaoapp/v8go
path: v8go
- name: Unzip libv8
run: |
files=$(find ./v8go -name "libv8*.zip")
for file in $files; do
dir=$(dirname "$file")
echo "Extracting $file to directory $dir"
unzip -o -d $dir $file
rm -rf $dir/__MACOSX
done
- name: Checkout Demo App
uses: actions/checkout@v4
with:
repository: yaoapp/yao-dev-app
path: app
- name: Checkout Extension
uses: actions/checkout@v4
with:
repository: yaoapp/yao-extensions-dev
path: extension
- name: Move Dependencies
run: |
mv kun ../
mv xun ../
mv gou ../
mv v8go ../
mv app ../
mv extension ../
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup 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 (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-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-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
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 Benchmarks
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 benchmark-sandbox-v2
# =============================================================================
# gRPC Tests - Run once with SQLite (transport layer, no DB matrix needed)
# =============================================================================

6
.gitignore vendored
View file

@ -74,8 +74,4 @@ tg-login
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
tai/testdata/

46
Makefile
View file

@ -202,35 +202,17 @@ unit-test-registry:
fi
# ---------------------------------------------------------------------------
# Sandbox V2 Integration Test (tai + sandbox/v2 + workspace)
# Requires: Docker, Tai container, optionally k3d for K8s mode
# Sandbox V2 CI Test (tai SDK + workspace only)
# Full sandbox/v2 integration tests (multi-pool, K8s, etc.) are run locally.
# ---------------------------------------------------------------------------
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
unit-test-sandbox-v2: unit-test-tai unit-test-workspace
@echo ""
@echo "============================================="
@echo "All Sandbox V2 integration tests passed"
@echo "All Sandbox V2 CI tests passed (tai + workspace)"
@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:
@ -271,26 +253,6 @@ unit-test-workspace:
@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:

878
agent/llm/providers/openai/gpt5_test.go
View file

@ -1,440 +1,442 @@
package openai_test
import (
gocontext "context"
"testing"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/llm"
"github.com/yaoapp/yao/agent/output/message"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/test"
)
// TestGPT5StreamBasic tests basic streaming completion with GPT-5
func TestGPT5StreamBasic(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-5")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Streaming: true,
Reasoning: true, // GPT-5 supports reasoning
ToolCalls: true,
Vision: true,
Multimodal: true,
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "What is 1+1? Reply with just the number.",
},
}
maxTokens := 100
options.MaxCompletionTokens = &maxTokens
ctx := newGPT5TestContext("test-gpt5-basic", "openai.gpt-5")
var chunks []string
handler := func(chunkType message.StreamChunkType, data []byte) int {
chunks = append(chunks, string(data))
t.Logf("Stream chunk [%s]: %s", chunkType, string(data))
return 0
}
response, err := llmInstance.Stream(ctx, messages, options, handler)
if err != nil {
t.Fatalf("Stream failed: %v", err)
}
if response == nil {
t.Fatal("Response is nil")
}
// Basic validation
if response.ID == "" {
t.Error("Response ID is empty")
}
if response.Model == "" {
t.Error("Response Model is empty")
}
// GPT-5 may use all tokens for reasoning, so content could be empty
// Just log the content instead of failing
t.Logf("Response content: %v", response.Content)
t.Logf("Usage: prompt=%d, completion=%d, total=%d",
response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
if response.Usage != nil && response.Usage.CompletionTokensDetails != nil {
t.Logf("Reasoning tokens: %d", response.Usage.CompletionTokensDetails.ReasoningTokens)
}
t.Logf("Final response: %+v", response)
t.Logf("Total chunks received: %d", len(chunks))
}
// TestGPT5ReasoningEffort tests reasoning_effort parameter with different levels
func TestGPT5ReasoningEffort(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-5")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
// Test with different reasoning effort levels
effortLevels := []string{"low", "medium", "high"}
for _, effort := range effortLevels {
t.Run("effort_"+effort, func(t *testing.T) {
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Reasoning: true,
ToolCalls: true,
},
ReasoningEffort: &effort,
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Solve: If all Bloops are Razzies and all Razzies are Lazzies, are all Bloops Lazzies?",
},
}
maxTokens := 1000
options.MaxCompletionTokens = &maxTokens
ctx := newGPT5TestContext("test-gpt5-reasoning-"+effort, "openai.gpt-5")
response, err := llmInstance.Post(ctx, messages, options)
if err != nil {
t.Fatalf("Post failed with effort=%s: %v", effort, err)
}
if response == nil {
t.Fatal("Response is nil")
}
// Check reasoning tokens
var reasoningTokens int
if response.Usage != nil && response.Usage.CompletionTokensDetails != nil {
reasoningTokens = response.Usage.CompletionTokensDetails.ReasoningTokens
}
t.Logf("Reasoning effort: %s", effort)
t.Logf("Reasoning tokens: %d", reasoningTokens)
t.Logf("Total tokens: %d", response.Usage.TotalTokens)
t.Logf("Content: %s", response.Content)
// GPT-5 reasoning is hidden (no reasoning_content field)
// But should have reasoning_tokens in usage
if effort != "low" {
if reasoningTokens == 0 {
t.Logf("Warning: Expected reasoning_tokens > 0 for effort='%s', got 0", effort)
}
}
})
}
}
// TestGPT5PostWithToolCalls tests GPT-5 with tool calls
func TestGPT5PostWithToolCalls(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-5")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Reasoning: true,
ToolCalls: true,
},
}
// Define a calculation tool
calcTool := map[string]interface{}{
"type": "function",
"function": map[string]interface{}{
"name": "calculate",
"description": "Perform a mathematical calculation",
"parameters": map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"expression": map[string]interface{}{
"type": "string",
"description": "The mathematical expression to evaluate",
},
},
"required": []string{"expression"},
},
},
}
options.Tools = []map[string]interface{}{calcTool}
options.ToolChoice = "auto"
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Use the calculate function to compute 2 * 3",
},
}
ctx := newGPT5TestContext("test-gpt5-tools", "openai.gpt-5")
response, err := llmInstance.Post(ctx, messages, options)
if err != nil {
t.Fatalf("Post with tool calls failed: %v", err)
}
if response == nil {
t.Fatal("Response is nil")
}
// GPT-5 reasoning models may not always use tool calls
// Log what we got instead of failing
if len(response.ToolCalls) == 0 {
t.Logf("No tool calls returned. Content: %v", response.Content)
} else {
tc := response.ToolCalls[0]
t.Logf("✓ Tool call: %s(%s)", tc.Function.Name, tc.Function.Arguments)
if tc.Function.Name != "calculate" {
t.Logf("Warning: Expected tool name 'calculate', got '%s'", tc.Function.Name)
}
}
if response.Usage != nil {
t.Logf("Usage: prompt=%d, completion=%d, total=%d",
response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
if response.Usage.CompletionTokensDetails != nil {
t.Logf("Reasoning tokens: %d", response.Usage.CompletionTokensDetails.ReasoningTokens)
}
}
t.Logf("Response: %+v", response)
}
// TestGPT5Vision tests GPT-5 with image input
func TestGPT5Vision(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-5")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Reasoning: true,
Vision: true,
Multimodal: true,
},
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
// Message with image content
messages := []context.Message{
{
Role: context.RoleUser,
Content: []context.ContentPart{
{
Type: context.ContentText,
Text: "What is in this image? Describe briefly.",
},
{
Type: context.ContentImageURL,
ImageURL: &context.ImageURL{
URL: "https://raw.githubusercontent.com/YaoApp/yao/refs/heads/main/yao/data/icons/icon.png",
},
},
},
},
}
maxTokens := 200
options.MaxCompletionTokens = &maxTokens
ctx := newGPT5TestContext("test-gpt5-vision", "openai.gpt-5")
response, err := llmInstance.Post(ctx, messages, options)
if err != nil {
t.Fatalf("Post with vision failed: %v", err)
}
if response == nil {
t.Fatal("Response is nil")
}
// Should have content describing the image
// Content can be string or []ContentPart for multimodal responses
var contentStr string
switch v := response.Content.(type) {
case string:
contentStr = v
case []interface{}:
// Handle []ContentPart serialized as []interface{}
for _, part := range v {
if partMap, ok := part.(map[string]interface{}); ok {
if text, ok := partMap["text"].(string); ok {
contentStr += text
}
}
}
case []context.ContentPart:
for _, part := range v {
if part.Type == context.ContentText {
contentStr += part.Text
}
}
case nil:
// GPT-5 reasoning models may use all tokens for reasoning, leaving no content
t.Log("Content is nil (reasoning model may have used all tokens for reasoning)")
default:
t.Logf("Unexpected content type: %T", response.Content)
}
if contentStr != "" {
t.Logf("Image description: %s", contentStr)
} else if response.Content != nil {
t.Logf("Warning: Expected text content describing the image, got empty or non-text content")
}
if response.Usage != nil {
t.Logf("Usage: prompt=%d, completion=%d, total=%d",
response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
}
}
// TestGPT5ReasoningEffortWithGPT4o tests that GPT-4o ignores reasoning_effort
func TestGPT5ReasoningEffortWithGPT4o(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
// Use GPT-4o which doesn't support reasoning
conn, err := connector.Select("openai.gpt-4o")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
effort := "high"
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Reasoning: false, // GPT-4o doesn't support reasoning
ToolCalls: true,
},
ReasoningEffort: &effort, // Should be ignored by adapter
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Say 'OK'",
},
}
maxTokens := 10
options.MaxCompletionTokens = &maxTokens
ctx := newGPT5TestContext("test-gpt4o-no-reasoning", "openai.gpt-4o")
// Should succeed (adapter removes reasoning_effort parameter)
response, err := llmInstance.Post(ctx, messages, options)
if err != nil {
t.Fatalf("Post failed: %v", err)
}
if response == nil {
t.Fatal("Response is nil")
}
// Should have 0 reasoning tokens (GPT-4o doesn't do reasoning)
if response.Usage != nil && response.Usage.CompletionTokensDetails != nil {
reasoningTokens := response.Usage.CompletionTokensDetails.ReasoningTokens
if reasoningTokens != 0 {
t.Errorf("Expected reasoning_tokens=0 for GPT-4o, got %d", reasoningTokens)
} else {
t.Log("✓ GPT-4o correctly shows reasoning_tokens=0")
}
}
t.Log("✓ ReasoningAdapter correctly removed reasoning_effort parameter for GPT-4o")
}
// ============================================================================
// Helper Functions
// ============================================================================
// newGPT5TestContext creates a real Context for testing GPT-5 provider
func newGPT5TestContext(chatID, connectorID string) *context.Context {
authorized := &types.AuthorizedInfo{
Subject: "test-user",
ClientID: "test-client",
UserID: "test-user-123",
TeamID: "test-team-456",
TenantID: "test-tenant-789",
SessionID: "test-session-id",
Constraints: types.DataConstraints{
TeamOnly: true,
Extra: map[string]interface{}{
"test": "gpt5-provider",
},
},
}
ctx := context.New(gocontext.Background(), authorized, chatID)
ctx.AssistantID = "test-assistant"
ctx.Locale = "en-us"
ctx.Theme = "light"
ctx.Client = context.Client{
Type: "web",
UserAgent: "GPT5ProviderTest/1.0",
IP: "127.0.0.1",
}
ctx.Referer = context.RefererAPI
ctx.Accept = context.AcceptStandard
ctx.Route = "/api/test"
ctx.Metadata = make(map[string]interface{})
return ctx
}
// GPT-5 tests temporarily commented out
//
// import (
// gocontext "context"
// "testing"
//
// "github.com/yaoapp/gou/connector"
// "github.com/yaoapp/gou/connector/openai"
// "github.com/yaoapp/yao/agent/context"
// "github.com/yaoapp/yao/agent/llm"
// "github.com/yaoapp/yao/agent/output/message"
// "github.com/yaoapp/yao/config"
// "github.com/yaoapp/yao/openapi/oauth/types"
// "github.com/yaoapp/yao/test"
// )
//
// // TestGPT5StreamBasic tests basic streaming completion with GPT-5
// func TestGPT5StreamBasic(t *testing.T) {
// test.Prepare(t, config.Conf)
// defer test.Clean()
//
// conn, err := connector.Select("openai.gpt-5")
// if err != nil {
// t.Fatalf("Failed to select connector: %v", err)
// }
//
// options := &context.CompletionOptions{
// Capabilities: &openai.Capabilities{
// Streaming: true,
// Reasoning: true, // GPT-5 supports reasoning
// ToolCalls: true,
// Vision: true,
// Multimodal: true,
// },
// }
//
// llmInstance, err := llm.New(conn, options)
// if err != nil {
// t.Fatalf("Failed to create LLM instance: %v", err)
// }
//
// messages := []context.Message{
// {
// Role: context.RoleUser,
// Content: "What is 1+1? Reply with just the number.",
// },
// }
//
// maxTokens := 100
// options.MaxCompletionTokens = &maxTokens
//
// ctx := newGPT5TestContext("test-gpt5-basic", "openai.gpt-5")
//
// var chunks []string
// handler := func(chunkType message.StreamChunkType, data []byte) int {
// chunks = append(chunks, string(data))
// t.Logf("Stream chunk [%s]: %s", chunkType, string(data))
// return 0
// }
//
// response, err := llmInstance.Stream(ctx, messages, options, handler)
// if err != nil {
// t.Fatalf("Stream failed: %v", err)
// }
//
// if response == nil {
// t.Fatal("Response is nil")
// }
//
// // Basic validation
// if response.ID == "" {
// t.Error("Response ID is empty")
// }
// if response.Model == "" {
// t.Error("Response Model is empty")
// }
//
// // GPT-5 may use all tokens for reasoning, so content could be empty
// // Just log the content instead of failing
// t.Logf("Response content: %v", response.Content)
// t.Logf("Usage: prompt=%d, completion=%d, total=%d",
// response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
//
// if response.Usage != nil && response.Usage.CompletionTokensDetails != nil {
// t.Logf("Reasoning tokens: %d", response.Usage.CompletionTokensDetails.ReasoningTokens)
// }
//
// t.Logf("Final response: %+v", response)
// t.Logf("Total chunks received: %d", len(chunks))
// }
//
// // TestGPT5ReasoningEffort tests reasoning_effort parameter with different levels
// func TestGPT5ReasoningEffort(t *testing.T) {
// test.Prepare(t, config.Conf)
// defer test.Clean()
//
// conn, err := connector.Select("openai.gpt-5")
// if err != nil {
// t.Fatalf("Failed to select connector: %v", err)
// }
//
// // Test with different reasoning effort levels
// effortLevels := []string{"low", "medium", "high"}
//
// for _, effort := range effortLevels {
// t.Run("effort_"+effort, func(t *testing.T) {
// options := &context.CompletionOptions{
// Capabilities: &openai.Capabilities{
// Reasoning: true,
// ToolCalls: true,
// },
// ReasoningEffort: &effort,
// }
//
// llmInstance, err := llm.New(conn, options)
// if err != nil {
// t.Fatalf("Failed to create LLM instance: %v", err)
// }
//
// messages := []context.Message{
// {
// Role: context.RoleUser,
// Content: "Solve: If all Bloops are Razzies and all Razzies are Lazzies, are all Bloops Lazzies?",
// },
// }
//
// maxTokens := 1000
// options.MaxCompletionTokens = &maxTokens
//
// ctx := newGPT5TestContext("test-gpt5-reasoning-"+effort, "openai.gpt-5")
//
// response, err := llmInstance.Post(ctx, messages, options)
// if err != nil {
// t.Fatalf("Post failed with effort=%s: %v", effort, err)
// }
//
// if response == nil {
// t.Fatal("Response is nil")
// }
//
// // Check reasoning tokens
// var reasoningTokens int
// if response.Usage != nil && response.Usage.CompletionTokensDetails != nil {
// reasoningTokens = response.Usage.CompletionTokensDetails.ReasoningTokens
// }
//
// t.Logf("Reasoning effort: %s", effort)
// t.Logf("Reasoning tokens: %d", reasoningTokens)
// t.Logf("Total tokens: %d", response.Usage.TotalTokens)
// t.Logf("Content: %s", response.Content)
//
// // GPT-5 reasoning is hidden (no reasoning_content field)
// // But should have reasoning_tokens in usage
// if effort != "low" {
// if reasoningTokens == 0 {
// t.Logf("Warning: Expected reasoning_tokens > 0 for effort='%s', got 0", effort)
// }
// }
// })
// }
// }
//
// // TestGPT5PostWithToolCalls tests GPT-5 with tool calls
// func TestGPT5PostWithToolCalls(t *testing.T) {
// test.Prepare(t, config.Conf)
// defer test.Clean()
//
// conn, err := connector.Select("openai.gpt-5")
// if err != nil {
// t.Fatalf("Failed to select connector: %v", err)
// }
//
// options := &context.CompletionOptions{
// Capabilities: &openai.Capabilities{
// Reasoning: true,
// ToolCalls: true,
// },
// }
//
// // Define a calculation tool
// calcTool := map[string]interface{}{
// "type": "function",
// "function": map[string]interface{}{
// "name": "calculate",
// "description": "Perform a mathematical calculation",
// "parameters": map[string]interface{}{
// "type": "object",
// "properties": map[string]interface{}{
// "expression": map[string]interface{}{
// "type": "string",
// "description": "The mathematical expression to evaluate",
// },
// },
// "required": []string{"expression"},
// },
// },
// }
//
// options.Tools = []map[string]interface{}{calcTool}
// options.ToolChoice = "auto"
//
// llmInstance, err := llm.New(conn, options)
// if err != nil {
// t.Fatalf("Failed to create LLM instance: %v", err)
// }
//
// messages := []context.Message{
// {
// Role: context.RoleUser,
// Content: "Use the calculate function to compute 2 * 3",
// },
// }
//
// ctx := newGPT5TestContext("test-gpt5-tools", "openai.gpt-5")
//
// response, err := llmInstance.Post(ctx, messages, options)
// if err != nil {
// t.Fatalf("Post with tool calls failed: %v", err)
// }
//
// if response == nil {
// t.Fatal("Response is nil")
// }
//
// // GPT-5 reasoning models may not always use tool calls
// // Log what we got instead of failing
// if len(response.ToolCalls) == 0 {
// t.Logf("No tool calls returned. Content: %v", response.Content)
// } else {
// tc := response.ToolCalls[0]
// t.Logf("✓ Tool call: %s(%s)", tc.Function.Name, tc.Function.Arguments)
//
// if tc.Function.Name != "calculate" {
// t.Logf("Warning: Expected tool name 'calculate', got '%s'", tc.Function.Name)
// }
// }
//
// if response.Usage != nil {
// t.Logf("Usage: prompt=%d, completion=%d, total=%d",
// response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
// if response.Usage.CompletionTokensDetails != nil {
// t.Logf("Reasoning tokens: %d", response.Usage.CompletionTokensDetails.ReasoningTokens)
// }
// }
//
// t.Logf("Response: %+v", response)
// }
//
// // TestGPT5Vision tests GPT-5 with image input
// func TestGPT5Vision(t *testing.T) {
// test.Prepare(t, config.Conf)
// defer test.Clean()
//
// conn, err := connector.Select("openai.gpt-5")
// if err != nil {
// t.Fatalf("Failed to select connector: %v", err)
// }
//
// options := &context.CompletionOptions{
// Capabilities: &openai.Capabilities{
// Reasoning: true,
// Vision: true,
// Multimodal: true,
// },
// }
//
// llmInstance, err := llm.New(conn, options)
// if err != nil {
// t.Fatalf("Failed to create LLM instance: %v", err)
// }
//
// // Message with image content
// messages := []context.Message{
// {
// Role: context.RoleUser,
// Content: []context.ContentPart{
// {
// Type: context.ContentText,
// Text: "What is in this image? Describe briefly.",
// },
// {
// Type: context.ContentImageURL,
// ImageURL: &context.ImageURL{
// URL: "https://raw.githubusercontent.com/YaoApp/yao/refs/heads/main/yao/data/icons/icon.png",
// },
// },
// },
// },
// }
//
// maxTokens := 200
// options.MaxCompletionTokens = &maxTokens
//
// ctx := newGPT5TestContext("test-gpt5-vision", "openai.gpt-5")
//
// response, err := llmInstance.Post(ctx, messages, options)
// if err != nil {
// t.Fatalf("Post with vision failed: %v", err)
// }
//
// if response == nil {
// t.Fatal("Response is nil")
// }
//
// // Should have content describing the image
// // Content can be string or []ContentPart for multimodal responses
// var contentStr string
// switch v := response.Content.(type) {
// case string:
// contentStr = v
// case []interface{}:
// // Handle []ContentPart serialized as []interface{}
// for _, part := range v {
// if partMap, ok := part.(map[string]interface{}); ok {
// if text, ok := partMap["text"].(string); ok {
// contentStr += text
// }
// }
// }
// case []context.ContentPart:
// for _, part := range v {
// if part.Type == context.ContentText {
// contentStr += part.Text
// }
// }
// case nil:
// // GPT-5 reasoning models may use all tokens for reasoning, leaving no content
// t.Log("Content is nil (reasoning model may have used all tokens for reasoning)")
// default:
// t.Logf("Unexpected content type: %T", response.Content)
// }
//
// if contentStr != "" {
// t.Logf("Image description: %s", contentStr)
// } else if response.Content != nil {
// t.Logf("Warning: Expected text content describing the image, got empty or non-text content")
// }
//
// if response.Usage != nil {
// t.Logf("Usage: prompt=%d, completion=%d, total=%d",
// response.Usage.PromptTokens, response.Usage.CompletionTokens, response.Usage.TotalTokens)
// }
// }
//
// // TestGPT5ReasoningEffortWithGPT4o tests that GPT-4o ignores reasoning_effort
// func TestGPT5ReasoningEffortWithGPT4o(t *testing.T) {
// test.Prepare(t, config.Conf)
// defer test.Clean()
//
// // Use GPT-4o which doesn't support reasoning
// conn, err := connector.Select("openai.gpt-4o")
// if err != nil {
// t.Fatalf("Failed to select connector: %v", err)
// }
//
// effort := "high"
// options := &context.CompletionOptions{
// Capabilities: &openai.Capabilities{
// Reasoning: false, // GPT-4o doesn't support reasoning
// ToolCalls: true,
// },
// ReasoningEffort: &effort, // Should be ignored by adapter
// }
//
// llmInstance, err := llm.New(conn, options)
// if err != nil {
// t.Fatalf("Failed to create LLM instance: %v", err)
// }
//
// messages := []context.Message{
// {
// Role: context.RoleUser,
// Content: "Say 'OK'",
// },
// }
//
// maxTokens := 10
// options.MaxCompletionTokens = &maxTokens
//
// ctx := newGPT5TestContext("test-gpt4o-no-reasoning", "openai.gpt-4o")
//
// // Should succeed (adapter removes reasoning_effort parameter)
// response, err := llmInstance.Post(ctx, messages, options)
// if err != nil {
// t.Fatalf("Post failed: %v", err)
// }
//
// if response == nil {
// t.Fatal("Response is nil")
// }
//
// // Should have 0 reasoning tokens (GPT-4o doesn't do reasoning)
// if response.Usage != nil && response.Usage.CompletionTokensDetails != nil {
// reasoningTokens := response.Usage.CompletionTokensDetails.ReasoningTokens
// if reasoningTokens != 0 {
// t.Errorf("Expected reasoning_tokens=0 for GPT-4o, got %d", reasoningTokens)
// } else {
// t.Log("✓ GPT-4o correctly shows reasoning_tokens=0")
// }
// }
//
// t.Log("✓ ReasoningAdapter correctly removed reasoning_effort parameter for GPT-4o")
// }
//
// // ============================================================================
// // Helper Functions
// // ============================================================================
//
// // newGPT5TestContext creates a real Context for testing GPT-5 provider
// func newGPT5TestContext(chatID, connectorID string) *context.Context {
// authorized := &types.AuthorizedInfo{
// Subject: "test-user",
// ClientID: "test-client",
// UserID: "test-user-123",
// TeamID: "test-team-456",
// TenantID: "test-tenant-789",
// SessionID: "test-session-id",
// Constraints: types.DataConstraints{
// TeamOnly: true,
// Extra: map[string]interface{}{
// "test": "gpt5-provider",
// },
// },
// }
//
// ctx := context.New(gocontext.Background(), authorized, chatID)
// ctx.AssistantID = "test-assistant"
// ctx.Locale = "en-us"
// ctx.Theme = "light"
// ctx.Client = context.Client{
// Type: "web",
// UserAgent: "GPT5ProviderTest/1.0",
// IP: "127.0.0.1",
// }
// ctx.Referer = context.RefererAPI
// ctx.Accept = context.AcceptStandard
// ctx.Route = "/api/test"
// ctx.Metadata = make(map[string]interface{})
// return ctx
// }

192
agent/llm/providers/openai/temperature_test.go
View file

@ -14,53 +14,54 @@ import (
)
// TestTemperatureGPT5AutoReset tests that GPT-5 automatically resets temperature to 1.0
func TestTemperatureGPT5AutoReset(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-5")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
invalidTemp := 0.7 // GPT-5 doesn't support this
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Reasoning: true,
},
Temperature: &invalidTemp, // Should be reset to 1.0
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Say 'OK'",
},
}
maxTokens := 10
options.MaxCompletionTokens = &maxTokens
ctx := newTemperatureTestContext("test-gpt5-temp", "openai.gpt-5")
// Should succeed (temperature automatically reset to 1.0)
response, err := llmInstance.Post(ctx, messages, options)
if err != nil {
t.Fatalf("Post failed: %v", err)
}
if response == nil {
t.Fatal("Response is nil")
}
t.Log("✓ GPT-5 successfully handled invalid temperature by resetting to 1.0")
t.Logf("Response: %v", response.Content)
}
// Temporarily commented out
// func TestTemperatureGPT5AutoReset(t *testing.T) {
// test.Prepare(t, config.Conf)
// defer test.Clean()
//
// conn, err := connector.Select("openai.gpt-5")
// if err != nil {
// t.Fatalf("Failed to select connector: %v", err)
// }
//
// invalidTemp := 0.7 // GPT-5 doesn't support this
// options := &context.CompletionOptions{
// Capabilities: &openai.Capabilities{
// Reasoning: true,
// },
// Temperature: &invalidTemp, // Should be reset to 1.0
// }
//
// llmInstance, err := llm.New(conn, options)
// if err != nil {
// t.Fatalf("Failed to create LLM instance: %v", err)
// }
//
// messages := []context.Message{
// {
// Role: context.RoleUser,
// Content: "Say 'OK'",
// },
// }
//
// maxTokens := 10
// options.MaxCompletionTokens = &maxTokens
//
// ctx := newTemperatureTestContext("test-gpt5-temp", "openai.gpt-5")
//
// // Should succeed (temperature automatically reset to 1.0)
// response, err := llmInstance.Post(ctx, messages, options)
// if err != nil {
// t.Fatalf("Post failed: %v", err)
// }
//
// if response == nil {
// t.Fatal("Response is nil")
// }
//
// t.Log("✓ GPT-5 successfully handled invalid temperature by resetting to 1.0")
// t.Logf("Response: %v", response.Content)
// }
// TestTemperatureDeepSeekR1AutoReset tests that DeepSeek R1 automatically resets temperature to 1.0
func TestTemperatureDeepSeekR1AutoReset(t *testing.T) {
@ -215,53 +216,54 @@ func TestTemperatureDeepSeekV3Preserved(t *testing.T) {
}
// TestTemperatureGPT5Default tests that GPT-5 with temperature=1.0 works fine
func TestTemperatureGPT5Default(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
conn, err := connector.Select("openai.gpt-5")
if err != nil {
t.Fatalf("Failed to select connector: %v", err)
}
defaultTemp := 1.0 // GPT-5's valid temperature
options := &context.CompletionOptions{
Capabilities: &openai.Capabilities{
Reasoning: true,
},
Temperature: &defaultTemp, // Should work fine
}
llmInstance, err := llm.New(conn, options)
if err != nil {
t.Fatalf("Failed to create LLM instance: %v", err)
}
messages := []context.Message{
{
Role: context.RoleUser,
Content: "What is 2+2? Reply with just the number.",
},
}
maxTokens := 10
options.MaxCompletionTokens = &maxTokens
ctx := newTemperatureTestContext("test-gpt5-temp-default", "openai.gpt-5")
// Should succeed with default temperature
response, err := llmInstance.Post(ctx, messages, options)
if err != nil {
t.Fatalf("Post failed: %v", err)
}
if response == nil {
t.Fatal("Response is nil")
}
t.Log("✓ GPT-5 successfully handled default temperature (1.0)")
t.Logf("Response: %v", response.Content)
}
// Temporarily commented out
// func TestTemperatureGPT5Default(t *testing.T) {
// test.Prepare(t, config.Conf)
// defer test.Clean()
//
// conn, err := connector.Select("openai.gpt-5")
// if err != nil {
// t.Fatalf("Failed to select connector: %v", err)
// }
//
// defaultTemp := 1.0 // GPT-5's valid temperature
// options := &context.CompletionOptions{
// Capabilities: &openai.Capabilities{
// Reasoning: true,
// },
// Temperature: &defaultTemp, // Should work fine
// }
//
// llmInstance, err := llm.New(conn, options)
// if err != nil {
// t.Fatalf("Failed to create LLM instance: %v", err)
// }
//
// messages := []context.Message{
// {
// Role: context.RoleUser,
// Content: "What is 2+2? Reply with just the number.",
// },
// }
//
// maxTokens := 10
// options.MaxCompletionTokens = &maxTokens
//
// ctx := newTemperatureTestContext("test-gpt5-temp-default", "openai.gpt-5")
//
// // Should succeed with default temperature
// response, err := llmInstance.Post(ctx, messages, options)
// if err != nil {
// t.Fatalf("Post failed: %v", err)
// }
//
// if response == nil {
// t.Fatal("Response is nil")
// }
//
// t.Log("✓ GPT-5 successfully handled default temperature (1.0)")
// t.Logf("Response: %v", response.Content)
// }
// TestTemperatureNoTemperatureProvided tests that models work when no temperature is provided
func TestTemperatureNoTemperatureProvided(t *testing.T) {
@ -273,7 +275,7 @@ func TestTemperatureNoTemperatureProvided(t *testing.T) {
connector string
reasoning bool
}{
{"GPT-5 No Temp", "openai.gpt-5", true},
// {"GPT-5 No Temp", "openai.gpt-5", true}, // Temporarily commented out
{"GPT-4o No Temp", "openai.gpt-4o", false},
{"DeepSeek R1 No Temp", "deepseek.r1", true},
{"DeepSeek V3 No Temp", "deepseek.v3", false},

6
cmd/run.go
View file

@ -16,9 +16,9 @@ import (
"github.com/yaoapp/kun/exception"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/engine"
grpcclient "github.com/yaoapp/yao/grpc/client"
ischedule "github.com/yaoapp/yao/schedule"
"github.com/yaoapp/yao/share"
taigrpc "github.com/yaoapp/yao/tai/grpc"
itask "github.com/yaoapp/yao/task"
)
@ -93,8 +93,8 @@ func runGRPC(cred *Credential, args []string) {
os.Exit(1)
}
tm := taigrpc.NewTokenManager(cred.AccessToken, cred.RefreshToken, "", "")
client, err := taigrpc.Dial(cred.GRPCAddr, tm)
tm := grpcclient.NewTokenManager(cred.AccessToken, cred.RefreshToken, "")
client, err := grpcclient.Dial(cred.GRPCAddr, tm)
if err != nil {
color.Red(" %s %s\n", L("gRPC connect failed:"), err.Error())
os.Exit(1)

162
cmd/start.go
View file

@ -2,9 +2,11 @@ package cmd
import (
"fmt"
"net"
"os"
"os/signal"
"path/filepath"
"strconv"
"strings"
"syscall"
@ -16,7 +18,6 @@ import (
"github.com/yaoapp/gou/mcp"
"github.com/yaoapp/gou/plugin"
"github.com/yaoapp/gou/schedule"
"github.com/yaoapp/gou/server/http"
"github.com/yaoapp/gou/store"
"github.com/yaoapp/gou/task"
"github.com/yaoapp/gou/websocket"
@ -128,49 +129,24 @@ var startCmd = &cobra.Command{
// print the messages under the development mode
if mode == "development" {
// Start Studio Server
// Yao Studio will be deprecated in the future
// go func() {
// err = studio.Load(config.Conf)
// if err != nil {
// // fmt.Println(color.RedString(L("Studio Load: %s"), err.Error()))
// log.Error("Studio Load: %s", err.Error())
// return
// }
// err := studio.Start(config.Conf)
// if err != nil {
// log.Error("Studio Start: %s", err.Error())
// return
// }
// }()
// defer studio.Stop()
printApis(false)
printTasks(false)
printSchedules(false)
printConnectors(false)
printStores(false)
printMCPs(false)
}
root, _ := adminRoot()
endpoints := []setup.Endpoint{{URL: fmt.Sprintf("http://%s%s", "127.0.0.1", port), Interface: "localhost"}}
switch host {
case "0.0.0.0":
// All interfaces
if values, err := setup.Endpoints(config.Conf); err == nil {
endpoints = append(endpoints, values...)
}
break
case "127.0.0.1":
// Localhost only
break
default:
// Filter by the host IP
matched := false
endpoints = []setup.Endpoint{}
if values, err := setup.Endpoints(config.Conf); err == nil {
@ -187,32 +163,6 @@ var startCmd = &cobra.Command{
}
}
// Print gRPC listen addresses
grpcAddrs := yaogrpc.Addr()
for _, addr := range grpcAddrs {
fmt.Println(color.WhiteString(L("Listening")), color.GreenString(" %s (gRPC)", addr))
}
fmt.Println(color.WhiteString("\n---------------------------------"))
fmt.Println(color.WhiteString(L("Access Points")))
fmt.Println(color.WhiteString("---------------------------------"))
apiRoot := "/api"
if openapi.Server != nil {
apiRoot = openapi.Server.Config.BaseURL
}
for _, endpoint := range endpoints {
fmt.Println(color.CyanString("\n%s", endpoint.Interface))
fmt.Println(color.WhiteString("--------------------------"))
fmt.Println(color.WhiteString(L("Website")), color.GreenString(" %s", endpoint.URL))
fmt.Println(color.WhiteString(L("Dashboard")), color.GreenString(" %s/%s/auth/entry", endpoint.URL, strings.Trim(root, "/")))
if openapi.Server != nil {
fmt.Println(color.WhiteString(L("OpenAPI")), color.GreenString(" %s%s", endpoint.URL, apiRoot))
} else {
fmt.Println(color.WhiteString(L("API")), color.GreenString(" %s%s", endpoint.URL, apiRoot))
}
}
fmt.Println("")
// Print welcome message for the new application
if isnew {
printWelcome()
@ -230,32 +180,66 @@ var startCmd = &cobra.Command{
// (must happen before HTTP/gRPC start so handlers can access it)
tairegistry.Init(nil)
// Start HTTP Server
srv, err := service.Start(config.Conf)
defer func() {
service.Stop(srv)
fmt.Println(color.GreenString(L("✨Exited successfully!")))
}()
// Pre-flight: detect port conflicts before attempting to start servers.
if occupied, proc := portOccupied(config.Conf.Host, config.Conf.Port); occupied {
fmt.Println(color.RedString(L("Fatal: HTTP port %d is already in use%s"), config.Conf.Port, proc))
return
}
if strings.ToLower(config.Conf.GRPC.Enabled) != "off" {
for _, h := range strings.Split(config.Conf.GRPC.Host, ",") {
if occupied, proc := portOccupied(strings.TrimSpace(h), config.Conf.GRPC.Port); occupied {
fmt.Println(color.RedString(L("Fatal: gRPC port %d is already in use%s"), config.Conf.GRPC.Port, proc))
return
}
}
}
// Start all servers (gRPC + HTTP) as a single unit.
// Start() blocks until HTTP port is bound (READY) or returns error.
svc, err := service.Start(config.Conf, service.ServerHooks{
Start: yaogrpc.StartServer,
Stop: yaogrpc.Stop,
Addrs: yaogrpc.Addr,
})
if err != nil {
fmt.Println(color.RedString(L("Fatal: %s"), err.Error()))
os.Exit(1)
return
}
// Start gRPC Server (after HTTP, LIFO shutdown: gRPC stops before HTTP)
if grpcErr := yaogrpc.StartServer(config.Conf); grpcErr != nil {
fmt.Println(color.RedString(L("gRPC: %s"), grpcErr.Error()))
os.Exit(1)
// Access Points (printed after servers are up so addresses are known)
fmt.Println(color.WhiteString("\n---------------------------------"))
fmt.Println(color.WhiteString(L("Access Points")))
fmt.Println(color.WhiteString("---------------------------------"))
if grpcAddrs := svc.HookAddrs(); len(grpcAddrs) > 0 {
fmt.Println(color.CyanString("\ngRPC"))
fmt.Println(color.WhiteString("--------------------------"))
for _, addr := range grpcAddrs {
fmt.Println(color.WhiteString(L("Server")), color.GreenString(" %s", addr))
}
}
defer yaogrpc.Stop()
apiRoot := "/api"
if openapi.Server != nil {
apiRoot = openapi.Server.Config.BaseURL
}
for _, endpoint := range endpoints {
fmt.Println(color.CyanString("\n%s", endpoint.Interface))
fmt.Println(color.WhiteString("--------------------------"))
fmt.Println(color.WhiteString(L("Website")), color.GreenString(" %s", endpoint.URL))
fmt.Println(color.WhiteString(L("Dashboard")), color.GreenString(" %s/%s/auth/entry", endpoint.URL, strings.Trim(root, "/")))
if openapi.Server != nil {
fmt.Println(color.WhiteString(L("OpenAPI")), color.GreenString(" %s%s", endpoint.URL, apiRoot))
} else {
fmt.Println(color.WhiteString(L("API")), color.GreenString(" %s%s", endpoint.URL, apiRoot))
}
}
fmt.Println("")
// Start watching
watchDone := make(chan uint8, 1)
if mode == "development" && !startDisableWatching {
// fmt.Println(color.WhiteString("\n---------------------------------"))
// fmt.Println(color.WhiteString(L("Watching")))
// fmt.Println(color.WhiteString("---------------------------------"))
go service.Watch(srv, watchDone)
go svc.Watch(watchDone)
}
// Print the messages under the production mode
@ -279,31 +263,15 @@ var startCmd = &cobra.Command{
fmt.Printf("\n")
}
fmt.Println(color.GreenString(L("Server is up and running...")))
fmt.Println(color.GreenString("Ctrl+C to stop"))
for {
select {
case v := <-srv.Event():
switch v {
case http.READY:
fmt.Println(color.GreenString(L("Server is up and running...")))
fmt.Println(color.GreenString("Ctrl+C to stop"))
break
case http.CLOSED:
fmt.Println(color.GreenString(L("✨Exited successfully!")))
watchDone <- 1
return
case http.ERROR:
color.Red("Fatal: check the error information in the log")
watchDone <- 1
return
default:
fmt.Println("Signal:", v)
}
case <-interrupt:
fmt.Println(color.WhiteString("\nShutting down..."))
svc.Stop()
fmt.Println(color.GreenString(L("✨Exited successfully!")))
watchDone <- 1
return
}
@ -399,7 +367,7 @@ func printStores(silent bool) {
}
fmt.Println(color.WhiteString("\n---------------------------------"))
fmt.Println(color.WhiteString(L("Stores List (%d)"), len(connector.Connectors)))
fmt.Println(color.WhiteString(L("Stores List (%d)"), len(store.Pools)))
fmt.Println(color.WhiteString("---------------------------------"))
for name := range store.Pools {
fmt.Print(color.CyanString("[Store]"))
@ -647,6 +615,18 @@ func colorMehtod(method string) string {
}
}
// portOccupied probes whether host:port is already bound.
// Returns (true, " (pid XXXX)") when occupied, (false, "") otherwise.
func portOccupied(host string, port int) (bool, string) {
addr := net.JoinHostPort(host, strconv.Itoa(port))
ln, err := net.Listen("tcp", addr)
if err != nil {
return true, fmt.Sprintf(" (%s)", err.Error())
}
ln.Close()
return false, ""
}
func init() {
startCmd.PersistentFlags().BoolVarP(&startDebug, "debug", "", false, L("Development mode"))
startCmd.PersistentFlags().BoolVarP(&startDisableWatching, "disable-watching", "", false, L("Disable watching"))

30
grpc/IMPL.md
View file

@ -40,12 +40,11 @@ grpc/
Container client:
```
tai/grpc/
├── grpc.go // gRPC client, Dial, method wrappers
├── auth.go // read env tokens, attach metadata, handle refresh
├── grpc_test.go
└── cmd/
└── main.go // yao-grpc binary entry
grpc/client/ // gRPC client (moved from tai/grpc/ to grpc/client/)
├── client.go // gRPC client, Dial, method wrappers
└── token.go // read env tokens, attach metadata, handle refresh
tai repo: tai/call/ // container-side binary (replaces yao-grpc)
```
## V1 Phases
@ -163,20 +162,20 @@ Deliverable: LLM (unary + stream) and Agent streaming via gRPC.
### Phase 4: Tai gateway change (Tai repo) ✅
Depends on: Phase 1 (need proto definitions for testing). yao-grpc depends on this.
Depends on: Phase 1 (need proto definitions for testing). `tai call` (tai repo) depends on this.
Tai gateway currently dials a fixed `YaoUpstream` at startup. New behavior: yao-grpc tells Tai where to forward via request metadata (`x-grpc-upstream`). Tai reads the target address and proxies to it — removes `YaoUpstream` startup config.
Tai gateway receives the upstream address during registration (`SetUpstream`). All gRPC requests are forwarded to the configured upstream — no per-request metadata required.
| Task | Detail | Status |
|------|--------|--------|
| Tai `gateway/gateway.go` | Remove fixed `upstream *grpc.ClientConn`. On each request, read `x-grpc-upstream` from metadata → lookup/create conn from `sync.Map` cache (key = address string) → forward. Typical deployment has 1 upstream, cache stays tiny. | ✅ Done |
| Tai `server/server.go` | Remove `YaoUpstream` from `Config`. Gateway init no longer needs an address. | ✅ Done |
| Tai `gateway/gateway.go` | Removed fixed `upstream *grpc.ClientConn`. `SetUpstream` configures the forwarding target. `sync.Map` cache for connections (key = address string). | ✅ Done |
| Tai `server/server.go` | Remove `YaoUpstream` from `Config`. Gateway uses `SetUpstream` after registration. | ✅ Done |
| Tai `main.go` | Remove `--yao` flag, `TAI_YAO_UPSTREAM` env var, YAML `yao` field, and required check. | ✅ Done |
| Tai `gateway/gateway_test.go` | Updated tests: dynamic routing, missing metadata → InvalidArgument, metadata forwarding (x-grpc-upstream stripped), upstream error propagation, multiple upstreams, connection cache. Coverage: 88.8%. | ✅ Done |
| Tai `gateway/gateway_test.go` | Updated tests: SetUpstream routing, no-upstream → Unavailable, metadata forwarding, upstream error propagation, upstream switching, connection cache. | ✅ Done |
Connection cache: `sync.Map[string, *grpc.ClientConn]` — lazy dial on first request per upstream, reuse thereafter. No eviction needed (upstream count ≈ 1 in practice). `Close` closes all cached connections.
Deliverable: Tai starts without Yao address. Forwards based on request metadata.
Deliverable: Tai receives upstream via registration. Forwards all requests to configured upstream.
### Phase 5: yao-grpc container client ✅
@ -184,14 +183,11 @@ Depends on: Phase 1 (server + auth), Phase 4 (Tai gateway accepts `x-grpc-upstre
| Task | Detail | Status |
|------|--------|--------|
| `tai/grpc/auth.go` | `TokenManager`: read `YAO_TOKEN` / `YAO_REFRESH_TOKEN` / `YAO_SANDBOX_ID` / `YAO_GRPC_UPSTREAM` from env. `YAO_GRPC_TAI=enable` triggers Tai relay mode (requires `YAO_GRPC_UPSTREAM`). Attach as gRPC metadata on every call via unary + stream interceptors. Auto-refresh from response headers. | ✅ Done |
| `tai/grpc/grpc.go` | `Client`: `Dial(addr, TokenManager)`, `NewFromEnv()`. Method wrappers for all RPCs: Run, Shell, API, MCP (list/call/resources/read), ChatCompletions, ChatCompletionsStream, AgentStream, Healthz. | ✅ Done |
| `tai/grpc/cmd/main.go` | Stdio MCP server: JSON-RPC → gRPC. `yao-grpc version` prints version/commit/build time (via `-ldflags`). `yao-grpc serve` reads stdin JSON-RPC, dispatches to gRPC client. | ✅ Done |
| `tai/grpc/grpc_test.go` + `integration_test.go` | Black-box tests (package `grpc_test`). Unit: TokenManager metadata attachment, env parsing, refresh handling. Integration: real Yao gRPC server, all method wrappers, token refresh, auth rejection. Coverage: 83.9%. | ✅ Done |
| `tai call` (tai repo) | In-container gRPC bridge. Reads `YAO_TOKEN` / `YAO_REFRESH_TOKEN` / `YAO_SANDBOX_ID` / `YAO_GRPC_ADDR` from env. Attaches auth metadata on every call via unary + stream interceptors. Auto-refresh from response headers. | ✅ Done |
Container token issuance uses existing `oauth.MakeAccessToken` / `oauth.MakeRefreshToken` — called by sandbox Manager at container creation, injected as env vars. Revoke on Remove. No new auth code needed on the issuance side.
Deliverable: `go build -o yao-grpc ./tai/grpc/cmd`.
Deliverable: `tai call` subcommand (part of Tai binary).
### Phase 6: Device Flow + CLI auth ✅

14
tai/grpc/grpc.go → grpc/client/client.go
View file

@ -1,4 +1,4 @@
package grpc
package client
import (
"context"
@ -12,15 +12,14 @@ import (
"google.golang.org/grpc/credentials/insecure"
)
// Client wraps a gRPC connection to a Yao server (direct or via Tai relay).
// TokenManager handles auth metadata attachment and token refresh automatically.
// Client wraps a gRPC connection to a Yao server.
type Client struct {
conn *grpc.ClientConn
svc pb.YaoClient
token *TokenManager
}
// NewFromEnv reads YAO_GRPC_ADDR (required) and token env vars, dials the
// NewFromEnv reads YAO_GRPC_ADDR and token env vars, dials the
// gRPC server, and returns a connected Client.
func NewFromEnv() (*Client, error) {
addr := os.Getenv("YAO_GRPC_ADDR")
@ -36,10 +35,7 @@ func NewFromEnv() (*Client, error) {
return Dial(addr, tm)
}
// 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).
// Dial connects to a Yao gRPC server at addr with the given TokenManager.
func Dial(addr string, tm *TokenManager) (*Client, error) {
opts := []grpc.DialOption{
grpc.WithTransportCredentials(insecure.NewCredentials()),
@ -180,7 +176,6 @@ func (c *Client) ChatCompletions(ctx context.Context, connector string, messages
}
// ChatCompletionsStream sends a streaming chat completion request.
// The callback receives each chunk's data; return a non-nil error to stop.
func (c *Client) ChatCompletionsStream(ctx context.Context, connector string, messages, options []byte, cb func(data []byte, done bool) error) error {
stream, err := c.svc.ChatCompletionsStream(ctx, &pb.ChatRequest{
Connector: connector,
@ -210,7 +205,6 @@ func (c *Client) ChatCompletionsStream(ctx context.Context, connector string, me
// --- Agent ---
// AgentStream calls an agent with streaming response.
// The callback receives each chunk's data; return a non-nil error to stop.
func (c *Client) AgentStream(ctx context.Context, assistantID string, messages, options []byte, cb func(data []byte, done bool) error) error {
stream, err := c.svc.AgentStream(ctx, &pb.AgentRequest{
AssistantId: assistantID,

52
tai/grpc/auth.go → grpc/client/token.go
View file

@ -1,8 +1,7 @@
package grpc
package client
import (
"context"
"fmt"
"os"
"sync"
@ -10,46 +9,30 @@ import (
"google.golang.org/grpc/metadata"
)
// TokenManager reads auth credentials from environment variables and attaches
// them as gRPC metadata on every call. It also handles automatic token refresh
// by reading new tokens from response headers.
// TokenManager attaches auth credentials as gRPC metadata on every call
// and handles automatic token refresh from response headers.
type TokenManager struct {
mu sync.RWMutex
accessToken string
refreshToken string
sandboxID string
upstream string // only set when YAO_GRPC_TAI=enable
taiMode bool
}
// NewTokenManagerFromEnv creates a TokenManager from environment variables.
// Returns an error if required variables are missing.
func NewTokenManagerFromEnv() (*TokenManager, error) {
tm := &TokenManager{
return &TokenManager{
accessToken: os.Getenv("YAO_TOKEN"),
refreshToken: os.Getenv("YAO_REFRESH_TOKEN"),
sandboxID: os.Getenv("YAO_SANDBOX_ID"),
}
if os.Getenv("YAO_GRPC_TAI") == "enable" {
tm.taiMode = true
tm.upstream = os.Getenv("YAO_GRPC_UPSTREAM")
if tm.upstream == "" {
return nil, fmt.Errorf("YAO_GRPC_TAI=enable but YAO_GRPC_UPSTREAM is not set")
}
}
return tm, nil
}, nil
}
// NewTokenManager creates a TokenManager with explicit values (for testing).
func NewTokenManager(accessToken, refreshToken, sandboxID, upstream string) *TokenManager {
// NewTokenManager creates a TokenManager with explicit values.
func NewTokenManager(accessToken, refreshToken, sandboxID string) *TokenManager {
return &TokenManager{
accessToken: accessToken,
refreshToken: refreshToken,
sandboxID: sandboxID,
upstream: upstream,
taiMode: upstream != "",
}
}
@ -58,7 +41,7 @@ func (tm *TokenManager) AttachMetadata(ctx context.Context) context.Context {
tm.mu.RLock()
defer tm.mu.RUnlock()
pairs := []string{}
var pairs []string
if tm.accessToken != "" {
pairs = append(pairs, "authorization", "Bearer "+tm.accessToken)
}
@ -68,9 +51,6 @@ func (tm *TokenManager) AttachMetadata(ctx context.Context) context.Context {
if tm.sandboxID != "" {
pairs = append(pairs, "x-sandbox-id", tm.sandboxID)
}
if tm.taiMode && tm.upstream != "" {
pairs = append(pairs, "x-grpc-upstream", tm.upstream)
}
if len(pairs) == 0 {
return ctx
@ -79,7 +59,7 @@ func (tm *TokenManager) AttachMetadata(ctx context.Context) context.Context {
}
// HandleResponseHeaders reads new tokens from response headers and updates
// the in-memory credentials. Call after each gRPC response.
// the in-memory credentials.
func (tm *TokenManager) HandleResponseHeaders(header metadata.MD) {
if header == nil {
return
@ -103,10 +83,8 @@ func (tm *TokenManager) UnaryInterceptor() grpc.UnaryClientInterceptor {
cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
ctx = tm.AttachMetadata(ctx)
var header metadata.MD
opts = append(opts, grpc.Header(&header))
err := invoker(ctx, method, req, reply, cc, opts...)
tm.HandleResponseHeaders(header)
return err
@ -114,8 +92,7 @@ func (tm *TokenManager) UnaryInterceptor() grpc.UnaryClientInterceptor {
}
// StreamInterceptor returns a gRPC stream client interceptor that attaches
// auth metadata. Token refresh from stream headers is handled by the caller
// via stream.Header().
// auth metadata.
func (tm *TokenManager) StreamInterceptor() grpc.StreamClientInterceptor {
return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn,
method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
@ -125,23 +102,16 @@ func (tm *TokenManager) StreamInterceptor() grpc.StreamClientInterceptor {
if err != nil {
return nil, err
}
if header, hErr := stream.Header(); hErr == nil {
tm.HandleResponseHeaders(header)
}
return stream, nil
}
}
// AccessToken returns the current access token (for testing/debugging).
// AccessToken returns the current access token.
func (tm *TokenManager) AccessToken() string {
tm.mu.RLock()
defer tm.mu.RUnlock()
return tm.accessToken
}
// IsTaiMode returns whether the client is configured for Tai relay mode.
func (tm *TokenManager) IsTaiMode() bool {
return tm.taiMode
}

41
grpc/grpc.go
View file

@ -6,6 +6,7 @@ import (
"strconv"
"strings"
"sync"
"time"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
@ -162,7 +163,7 @@ func StartServer(cfg config.Config) error {
addr := net.JoinHostPort(strings.TrimSpace(h), port)
lis, err := net.Listen("tcp", addr)
if err != nil {
Stop()
stopLocked()
return err
}
listeners = append(listeners, lis)
@ -179,17 +180,39 @@ func StartServer(cfg config.Config) error {
return nil
}
// Stop gracefully stops the gRPC server. Safe to call if server was never started.
// stopLocked performs cleanup while the caller already holds mu.
func stopLocked() {
s := server
server = nil
listeners = nil
addrs = nil
if s == nil {
return
}
done := make(chan struct{})
go func() {
s.GracefulStop()
close(done)
}()
select {
case <-done:
log.Info("gRPC server stopped gracefully")
case <-time.After(5 * time.Second):
log.Warn("gRPC server graceful stop timed out, forcing stop")
s.Stop()
}
}
// Stop gracefully stops the gRPC server with a 5-second timeout.
// If GracefulStop doesn't complete in time (e.g. active streams), it forces Stop.
// Safe to call if server was never started.
func Stop() {
mu.Lock()
defer mu.Unlock()
if server != nil {
server.GracefulStop()
server = nil
}
listeners = nil
addrs = nil
stopLocked()
}
// GRPCServer returns the active gRPC server instance.

6
openapi/openapi.go
View file

@ -28,6 +28,7 @@ import (
"github.com/yaoapp/yao/openapi/team"
openapiTrace "github.com/yaoapp/yao/openapi/trace"
"github.com/yaoapp/yao/openapi/user"
taiapi "github.com/yaoapp/yao/tai/api"
taitunnel "github.com/yaoapp/yao/tai/tunnel"
)
@ -182,6 +183,11 @@ func (openapi *OpenAPI) Attach(router *gin.Engine) {
group.Any("/tai/:taiID/proxy/*path", taitunnel.HandleProxy)
group.GET("/tai/:taiID/vnc/*path", taitunnel.HandleVNC)
// Tai direct registration API (uses /tai-nodes/ prefix to avoid routing conflict with /tai/:taiID/)
group.POST("/tai-nodes/register", taiapi.HandleRegister)
group.POST("/tai-nodes/heartbeat", taiapi.HandleHeartbeat)
group.DELETE("/tai-nodes/register/:tai_id", taiapi.HandleUnregister)
// Custom handlers (Defined by developer)
}

26
sandbox/DESIGN.md
View file

@ -50,14 +50,14 @@ High-level business layer on top of `tai.Client`. Manages container lifecycle, u
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 `yao-grpc` (via Tai Gateway relay or direct)
- Container-internal `tai call` (via Tai Gateway relay or direct)
- `yao run` CLI (after `yao login`)
- Other Yao instances (future node-to-node)
**IPC path (replacing Unix socket):**
```
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)
Local: Container → tai call (tai repo) → Yao gRPC 127.0.0.1:9099
Remote: Container → tai call (tai repo) → Tai Gateway (:9100 gRPC) → Yao gRPC Server (:9099)
```
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.
@ -233,8 +233,8 @@ 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 | 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`. |
| IPC | Unix socket bind mount + yao-bridge | All modes: `tai call` → gRPC (direct or via Tai gateway). No Unix socket. |
| MCP config | `{args: ["/tmp/yao.sock"]}` hardcoded | `YAO_GRPC_ADDR` + `YAO_TOKEN` env vars. Local: direct. Remote: via Tai gateway. |
| VNC | `vncproxy.NewProxy(nil)` local assumption | `tai.Client.VNC().URL()` |
| Cleanup | `dockerClient.ContainerRemove` | `tai.Client.Sandbox().Remove()` |
@ -243,11 +243,11 @@ agent/context/jsapi_sandbox.go
All modes use gRPC — no Unix socket fallback, one code path for local and remote.
```
Local: Container → yao-grpc → Yao gRPC 127.0.0.1:9099
Remote: Container → yao-grpc → Tai :9100 relay → Yao gRPC :9099
Local: Container → tai call → Yao gRPC 127.0.0.1:9099
Remote: Container → tai call → Tai :9100 gateway → Yao gRPC :9099
```
`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 call` is the in-container gRPC bridge (part of the Tai binary). It reads env vars and bridges JSON-RPC/stdio to the Yao gRPC server.
Mode determined by env vars injected by Manager at container creation:
@ -255,13 +255,11 @@ 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
# Remote: via Tai gateway
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 call` 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. The Tai gateway uses the upstream address configured during Tai registration to forward requests to Yao.
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.
@ -308,9 +306,9 @@ sandbox:
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.
3. **Tai gateway dynamic routing** (Tai repo) — removed fixed `YaoUpstream` startup config. Tai receives upstream address during registration (`SetUpstream`) and forwards all gRPC requests to it. 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`.
4. **`tai call` container client** (Tai repo) — in-container gRPC bridge replacing `yao-bridge`. Reads `YAO_TOKEN`/`YAO_REFRESH_TOKEN`/`YAO_SANDBOX_ID` from env, auto-refreshes tokens via response metadata. `YAO_GRPC_ADDR` determines the target (local Yao or remote Tai gateway).
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.

4
sandbox/docker/build.sh
View file

@ -165,8 +165,8 @@ case $TOOL in
# 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]"
echo "V2 images have moved to the tai repo: tai/docker/sandbox/build.sh"
echo "See: https://github.com/yaoapp/tai/tree/main/docker/sandbox"
exit 0
;;
*)

8
sandbox/v2/DESIGN.md
View file

@ -158,7 +158,7 @@ type Config struct {
}
```
Container gRPC env vars (`YAO_GRPC_ADDR`, `YAO_GRPC_UPSTREAM`, etc.) are derived automatically at creation time. Per-instance settings (image, memory, CPU, workdir, env, pool) are passed via `CreateOptions`.
Container gRPC env vars (`YAO_GRPC_ADDR`, etc.) are derived automatically at creation time. Per-instance settings (image, memory, CPU, workdir, env, pool) are passed via `CreateOptions`.
### Core API
@ -548,10 +548,8 @@ YAO_TOKEN=<access_token>
YAO_REFRESH_TOKEN=<refresh_token>
YAO_GRPC_ADDR=127.0.0.1:9099
# Remote mode (tai://) adds:
YAO_GRPC_TAI=enable
# Remote mode (tai://)
YAO_GRPC_ADDR=<tai-host>:9100
YAO_GRPC_UPSTREAM=127.0.0.1:9099
```
## Errors
@ -1121,7 +1119,7 @@ Permission control is the responsibility of the caller (JS scripts, Agent hooks,
| **Runtime** | Direct Docker SDK | tai.Client pool (Docker/K8s/Remote) |
| **Execution** | Exec + Stream | Exec + Stream + Attach (WS/SSE) |
| **File I/O** | bind mount + Docker Copy | `workspace.FS` (fs.FS compatible) |
| **IPC** | Unix socket + yao-bridge | gRPC (yao-grpc) |
| **IPC** | Unix socket + yao-bridge | gRPC (tai call) |
| **Idle detection** | External calls only | Dual: external calls + container heartbeat |
| **Lifecycle** | Chat session only | Policy-based (oneshot/session/longrunning/persistent) |
| **Pool** | Single Docker daemon | Multi-pool with per-pool policies |

2
sandbox/v2/Makefile
View file

@ -1,7 +1,7 @@
GO ?= go
GOFILES := $(shell find . -name "*.go" -not -path "./docker/*")
PACKAGES := $(shell $(GO) list ./...)
TEST_IMAGE ?= yaoapp/sandbox-v2-test:latest
TEST_IMAGE ?= yaoapp/tai-sandbox-test:latest
TEST_TIMEOUT ?= 600s
# ---------------------------------------------------------------------------

8
sandbox/v2/TEST.md
View file

@ -468,13 +468,11 @@ func TestBuildGRPCEnv_Local(t *testing.T) {
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"])
assert.NotEmpty(t, env["YAO_GRPC_ADDR"])
}
func TestCreateContainerTokens(t *testing.T) {
@ -589,8 +587,8 @@ sandbox-v2-test:
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
- `sandbox-v2-test` as default test image — includes `tai` (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 changes
- Attach tests (WS/SSE) use `sandbox-v2-test` image's built-in test services
## Coverage

8
sandbox/v2/box_attach_test.go
View file

@ -62,7 +62,7 @@ func TestAttachWS(t *testing.T) {
img := testImage()
if img == "alpine:latest" {
t.Skip("WebSocket test requires sandbox-v2-test image with ws-echo service")
t.Skip("WebSocket test requires tai-sandbox-test image with ws-echo service")
}
for _, pc := range testPools() {
@ -110,7 +110,7 @@ func TestAttachSSE(t *testing.T) {
img := testImage()
if img == "alpine:latest" {
t.Skip("SSE test requires sandbox-v2-test image with sse-server service")
t.Skip("SSE test requires tai-sandbox-test image with sse-server service")
}
for _, pc := range testPools() {
@ -159,7 +159,7 @@ func TestVNCURL(t *testing.T) {
img := testImage()
if img == "alpine:latest" {
t.Skip("VNC test requires sandbox-v2-test image with VNC desktop")
t.Skip("VNC test requires tai-sandbox-test image with VNC desktop")
}
for _, pc := range testPools() {
@ -189,7 +189,7 @@ func TestVNCConnect(t *testing.T) {
img := testImage()
if img == "alpine:latest" {
t.Skip("VNC test requires sandbox-v2-test image with VNC desktop")
t.Skip("VNC test requires tai-sandbox-test image with VNC desktop")
}
for _, pc := range testPools() {

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

@ -1,39 +0,0 @@
# 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
View file

@ -1,12 +0,0 @@
#!/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
View file

@ -1,7 +0,0 @@
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
View file

@ -1,419 +0,0 @@
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
View file

@ -1,510 +0,0 @@
// 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
View file

@ -1,244 +0,0 @@
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
View file

@ -1,79 +0,0 @@
#!/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
View file

@ -1,28 +0,0 @@
# 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
View file

@ -1,22 +0,0 @@
#!/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
View file

@ -1,28 +0,0 @@
"""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
View file

@ -1,14 +0,0 @@
"""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())

807
sandbox/v2/docs/API.md Normal file
View file

@ -0,0 +1,807 @@
# Sandbox V2 — Go API Reference
Package: `github.com/yaoapp/yao/sandbox/v2`
Sandbox V2 manages sandboxes through a pool of Tai nodes. Two primary abstractions:
- **Box** — a container (Docker or K8s pod). Created via `Manager.Create`.
- **Host** — the Tai host machine itself. Obtained via `Manager.Host` (no Create needed).
Supports workspace mounting, VNC, WebSocket proxying, and HostExec.
---
## Initialization
### Init
```go
func Init(cfg Config) error
```
Initializes the global Manager singleton. Must be called once at startup.
```go
err := sandbox.Init(sandbox.Config{
Pool: []sandbox.Pool{
{
Name: "docker",
Addr: "tai://192.168.1.10:19100",
MaxPerUser: 5,
MaxTotal: 20,
IdleTimeout: 30 * time.Minute,
MaxLifetime: 24 * time.Hour,
StopTimeout: 5 * time.Second,
},
},
})
```
### M
```go
func M() *Manager
```
Returns the global Manager. Panics if `Init` was not called.
```go
mgr := sandbox.M()
```
---
## Config
```go
type Config struct {
Pool []Pool
}
```
### Pool
```go
type Pool struct {
Name string
Addr string // "tai://host:port", "tunnel://host:port", or Docker socket
Options []tai.Option // tai.Client options
MaxPerUser int // 0 = unlimited
MaxTotal int // 0 = unlimited
IdleTimeout time.Duration // 0 = no idle cleanup
MaxLifetime time.Duration // 0 = no max lifetime
StopTimeout time.Duration // SIGTERM grace period; 0 = DefaultStopTimeout (2s)
}
```
---
## Lifecycle Policies
```go
type LifecyclePolicy string
const (
OneShot LifecyclePolicy = "oneshot" // removed after first Exec
Session LifecyclePolicy = "session" // removed after idle timeout
LongRunning LifecyclePolicy = "longrunning" // stopped after idle, removed after max lifetime
Persistent LifecyclePolicy = "persistent" // never auto-cleaned
)
```
---
## Manager
### Start
```go
func (m *Manager) Start(ctx context.Context) error
```
Recovers existing containers from all pools and starts the background cleanup loop (1 min interval).
```go
ctx := context.Background()
err := sandbox.M().Start(ctx)
```
### Close
```go
func (m *Manager) Close() error
```
Stops the cleanup loop and closes all pool connections.
### Create
```go
func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Box, error)
```
Creates and starts a new sandbox container. Returns a `Box` handle.
```go
box, err := sandbox.M().Create(ctx, sandbox.CreateOptions{
Image: "alpine:latest",
Owner: "user-123",
Pool: "docker",
Policy: sandbox.Session,
WorkDir: "/workspace",
Env: map[string]string{"LANG": "en_US.UTF-8"},
Memory: 512 * 1024 * 1024, // 512MB
CPUs: 1.0,
VNC: true,
Labels: map[string]string{"project": "demo"},
Ports: []sandbox.PortMapping{
{ContainerPort: 8080, HostPort: 0, Protocol: "tcp"},
},
IdleTimeout: 15 * time.Minute,
StopTimeout: 3 * time.Second,
WorkspaceID: "ws-abc",
MountMode: "rw",
MountPath: "/workspace",
})
```
### Host
```go
func (m *Manager) Host(ctx context.Context, pool string) (*Host, error)
```
Returns a `Host` handle for the given pool. Unlike `Create`, no container is provisioned —
the Host is available as long as the pool's Tai server reports `host_exec` capability.
Returns `ErrPoolNotFound` if the pool does not exist, or an error if the pool has no `host_exec`.
```go
host, err := sandbox.M().Host(ctx, "remote")
```
### Get
```go
func (m *Manager) Get(ctx context.Context, id string) (*Box, error)
```
Returns an existing sandbox by ID. Returns `ErrNotFound` if absent.
```go
box, err := sandbox.M().Get(ctx, "sb-12345")
```
### GetOrCreate
```go
func (m *Manager) GetOrCreate(ctx context.Context, opts CreateOptions) (*Box, error)
```
Returns existing sandbox by `opts.ID` or creates a new one.
```go
box, err := sandbox.M().GetOrCreate(ctx, sandbox.CreateOptions{
ID: "sb-session-xyz",
Image: "alpine:latest",
Owner: "user-123",
})
```
### List
```go
func (m *Manager) List(ctx context.Context, opts ListOptions) ([]*Box, error)
```
Returns all sandboxes matching the given filters. Empty fields = no filter.
```go
boxes, err := sandbox.M().List(ctx, sandbox.ListOptions{
Owner: "user-123",
Pool: "docker",
Labels: map[string]string{"project": "demo"},
})
```
### Remove
```go
func (m *Manager) Remove(ctx context.Context, id string) error
```
Force-removes a sandbox (SIGKILL + delete). Revokes container tokens.
```go
err := sandbox.M().Remove(ctx, "sb-12345")
```
### Cleanup
```go
func (m *Manager) Cleanup(ctx context.Context) error
```
Removes idle/expired sandboxes based on lifecycle policies. Called automatically by
the cleanup loop, but can also be invoked manually.
### Heartbeat
```go
func (m *Manager) Heartbeat(sandboxID string, active bool, processCount int) error
```
Updates a sandbox's last-active timestamp. Called by the gRPC heartbeat service.
```go
err := sandbox.M().Heartbeat("sb-12345", true, 3)
```
### AddPool
```go
func (m *Manager) AddPool(ctx context.Context, p Pool) error
```
Registers a new pool at runtime.
```go
err := sandbox.M().AddPool(ctx, sandbox.Pool{
Name: "k8s-gpu",
Addr: "tai://10.0.0.5:19100",
MaxTotal: 10,
})
```
### RemovePool
```go
func (m *Manager) RemovePool(ctx context.Context, name string, force bool) error
```
Removes a pool. Returns `ErrPoolInUse` if the pool has running boxes and `force=false`.
With `force=true`, all boxes in the pool are removed first.
### Pools
```go
func (m *Manager) Pools() []PoolInfo
```
Returns all registered pools and their status.
```go
for _, p := range sandbox.M().Pools() {
fmt.Printf("pool=%s addr=%s connected=%v boxes=%d\n",
p.Name, p.Addr, p.Connected, p.Boxes)
}
```
### SetGRPCPort
```go
func (m *Manager) SetGRPCPort(port int)
```
Sets the local gRPC port injected into container env vars (`YAO_GRPC_ADDR`). Default: `9099`.
### SetWorkspaceManager
```go
func (m *Manager) SetWorkspaceManager(wm *workspace.Manager)
```
Links the workspace manager. When `CreateOptions.WorkspaceID` is set, the Manager uses it
to resolve the workspace's bound node and route the container to the correct pool.
### ImageExists
```go
func (m *Manager) ImageExists(ctx context.Context, pool, ref string) (bool, error)
```
Reports whether the given image ref exists on the target pool node.
Returns `(true, nil)` when the pool has no image service (e.g. K8s — kubelet handles pulls).
```go
exists, err := sandbox.M().ImageExists(ctx, "docker", "alpine:latest")
```
### PullImage
```go
func (m *Manager) PullImage(ctx context.Context, pool, ref string, opts ImagePullOptions) (<-chan taisandbox.PullProgress, error)
```
Pulls an image to the target pool node. Returns a channel of `taisandbox.PullProgress`
(from `github.com/yaoapp/yao/tai/sandbox`). Returns `(nil, nil)` when the pool has no image
service (e.g. K8s).
`PullProgress` fields: `Status string`, `Layer string`, `Current int64`, `Total int64`, `Error string`.
```go
ch, err := sandbox.M().PullImage(ctx, "docker", "myapp:v2", sandbox.ImagePullOptions{
Auth: &sandbox.RegistryAuth{
Username: "user",
Password: "pass",
Server: "registry.example.com",
},
})
for p := range ch {
fmt.Printf("pull: %s layer=%s %d/%d\n", p.Status, p.Layer, p.Current, p.Total)
}
```
### EnsureImage
```go
func (m *Manager) EnsureImage(ctx context.Context, pool, ref string, opts ImagePullOptions) error
```
Checks if the image exists; if not, pulls it and blocks until complete.
```go
err := sandbox.M().EnsureImage(ctx, "docker", "alpine:latest", sandbox.ImagePullOptions{})
```
---
## Box
A `Box` is a handle to a running sandbox container.
### Accessors
```go
func (b *Box) ID() string
func (b *Box) Owner() string
func (b *Box) ContainerID() string
func (b *Box) Pool() string
func (b *Box) WorkspaceID() string
```
### Exec
```go
func (b *Box) Exec(ctx context.Context, cmd []string, opts ...ExecOption) (*ExecResult, error)
```
Runs a command and waits for completion. If the box policy is `OneShot`, the box is
auto-removed after execution.
```go
result, err := box.Exec(ctx, []string{"python3", "-c", "print('hello')"},
sandbox.WithWorkDir("/workspace"),
sandbox.WithEnv(map[string]string{"PYTHONPATH": "/lib"}),
sandbox.WithTimeout(30*time.Second),
)
fmt.Printf("exit=%d stdout=%s stderr=%s\n", result.ExitCode, result.Stdout, result.Stderr)
```
### Stream
```go
func (b *Box) Stream(ctx context.Context, cmd []string, opts ...ExecOption) (*ExecStream, error)
```
Runs a command with real-time streaming I/O.
```go
stream, err := box.Stream(ctx, []string{"bash"})
go io.Copy(os.Stdout, stream.Stdout)
go io.Copy(os.Stderr, stream.Stderr)
fmt.Fprintln(stream.Stdin, "echo hello")
stream.Stdin.Close()
exitCode, _ := stream.Wait()
```
### Attach
```go
func (b *Box) Attach(ctx context.Context, port int, opts ...AttachOption) (*ServiceConn, error)
```
Connects to a service running inside the sandbox via WebSocket proxy.
```go
conn, err := box.Attach(ctx, 8080,
sandbox.WithProtocol("ws"),
sandbox.WithPath("/api/stream"),
sandbox.WithHeaders(map[string]string{"Authorization": "Bearer xxx"}),
)
defer conn.Close()
conn.Write([]byte(`{"action":"subscribe"}`))
data, _ := conn.Read()
```
### VNC
```go
func (b *Box) VNC(ctx context.Context) (string, error)
```
Returns the VNC WebSocket URL for the sandbox (requires `VNC: true` at creation).
```go
url, err := box.VNC(ctx)
// url = "ws://tai-host:16080/vnc/xxx/ws"
```
### Proxy
```go
func (b *Box) Proxy(ctx context.Context, port int, path string) (string, error)
```
Returns the HTTP proxy URL for a service on the given port.
```go
url, err := box.Proxy(ctx, 3000, "/api/health")
// url = "http://tai-host:8099/container-id:3000/api/health"
```
### Workspace
```go
func (b *Box) Workspace() workspace.FS
```
Returns a `workspace.FS` interface (`github.com/yaoapp/yao/tai/workspace`) for file
operations on the sandbox's workspace volume. The interface embeds `fs.FS`, `fs.StatFS`,
`fs.ReadFileFS`, `fs.ReadDirFS`, `io.Closer`, and adds write methods (`WriteFile`,
`Remove`, `RemoveAll`, `Rename`, `MkdirAll`).
```go
ws := box.Workspace()
data, _ := ws.ReadFile("main.py")
ws.WriteFile("output.txt", []byte("result"), 0644)
ws.MkdirAll("src/pkg", 0755)
ws.Remove("tmp.log")
```
### Start / Stop / Remove
```go
func (b *Box) Start(ctx context.Context) error
func (b *Box) Stop(ctx context.Context) error
func (b *Box) Remove(ctx context.Context) error
```
```go
box.Stop(ctx) // SIGTERM with grace period, then SIGKILL
box.Start(ctx) // restart a stopped sandbox
box.Remove(ctx) // force remove
```
### Info
```go
func (b *Box) Info(ctx context.Context) (*BoxInfo, error)
```
Returns current sandbox status from the underlying container runtime.
```go
info, err := box.Info(ctx)
fmt.Printf("status=%s processes=%d vnc=%v created=%s\n",
info.Status, info.ProcessCount, info.VNC, info.CreatedAt)
```
---
## Host
A `Host` represents a Tai host machine execution environment, distinct from `Box` (containers).
No `Create` call is needed — a Host is available as long as the pool's Tai server reports `host_exec`.
### Accessors
```go
func (h *Host) Pool() string
```
### Exec
```go
func (h *Host) Exec(ctx context.Context, cmd string, args []string, opts ...HostExecOption) (*HostExecResult, error)
```
Runs a command directly on the Tai host machine via HostExec gRPC.
```go
host, _ := sandbox.M().Host(ctx, "remote")
result, err := host.Exec(ctx, "git", []string{"status"},
sandbox.WithHostWorkDir("/data/repos/project"),
sandbox.WithHostEnv(map[string]string{"GIT_AUTHOR_NAME": "bot"}),
sandbox.WithHostTimeout(10000), // 10s
sandbox.WithHostMaxOutput(1024*1024), // 1MB
)
fmt.Printf("exit=%d stdout=%s duration=%dms\n",
result.ExitCode, string(result.Stdout), result.DurationMs)
```
### Stream
```go
func (h *Host) Stream(ctx context.Context, cmd string, args []string, opts ...HostExecOption) (*HostExecStream, error)
```
Runs a command on the Tai host and streams stdout/stderr in real time via HostExec gRPC
ExecStream. Returns a `HostExecStream` with separate channels for stdout and stderr.
```go
host, _ := sandbox.M().Host(ctx, "remote")
stream, err := host.Stream(ctx, "tail", []string{"-f", "/var/log/app.log"},
sandbox.WithHostWorkDir("/data"),
sandbox.WithHostTimeout(60000),
)
go func() {
for chunk := range stream.Stderr {
fmt.Fprintf(os.Stderr, "%s", chunk)
}
}()
for chunk := range stream.Stdout {
fmt.Printf("%s", chunk)
}
exitCode, err := stream.Wait()
```
To cancel a long-running stream early:
```go
stream.Cancel()
```
### Workspace
```go
func (h *Host) Workspace(sessionID string) workspace.FS
```
Returns a `workspace.FS` for the given session on the host. Files are stored under
`dataDir/{sessionID}/` on the Tai host, accessed via Volume gRPC (independent of container
bind mounts).
```go
ws := host.Workspace("ws-abc")
ws.WriteFile("input.txt", []byte("data"), 0644)
data, _ := ws.ReadFile("output.txt")
entries, _ := ws.ReadDir(".")
```
---
## ExecOption Functions
```go
func WithWorkDir(dir string) ExecOption
func WithEnv(env map[string]string) ExecOption
func WithTimeout(timeout time.Duration) ExecOption
```
## AttachOption Functions
```go
func WithProtocol(protocol string) AttachOption // "ws" (default) or "sse"
func WithPath(path string) AttachOption // URL path on the target service
func WithHeaders(headers map[string]string) AttachOption
```
## HostExecOption Functions
```go
func WithHostWorkDir(dir string) HostExecOption
func WithHostEnv(env map[string]string) HostExecOption
func WithHostStdin(data []byte) HostExecOption
func WithHostTimeout(ms int64) HostExecOption
func WithHostMaxOutput(bytes int64) HostExecOption
```
---
## Types
### CreateOptions
```go
type CreateOptions struct {
ID string
Owner string
Labels map[string]string
Pool string // empty = default pool
Image string // required
WorkDir string // default "/workspace"
User string // container user
Env map[string]string
Memory int64 // bytes; 0 = unlimited
CPUs float64 // 0 = unlimited
VNC bool
Ports []PortMapping
Policy LifecyclePolicy // default Session
IdleTimeout time.Duration // overrides pool default
StopTimeout time.Duration // overrides pool default
WorkspaceID string // workspace to mount; empty = none
MountMode string // "rw" (default) or "ro"
MountPath string // default "/workspace"
}
```
### ListOptions
```go
type ListOptions struct {
Owner string
Pool string
Labels map[string]string
}
```
### PortMapping
```go
type PortMapping struct {
ContainerPort int
HostPort int // 0 = auto-assign
HostIP string
Protocol string // "tcp" (default), "udp"
}
```
### ExecResult
```go
type ExecResult struct {
ExitCode int
Stdout string
Stderr string
}
```
### ExecStream
```go
type ExecStream struct {
Stdout io.ReadCloser
Stderr io.ReadCloser
Stdin io.WriteCloser
Wait func() (int, error) // blocks until exit; returns exit code
Cancel func() // kills the process
}
```
### ServiceConn
```go
type ServiceConn struct {
Read func() ([]byte, error)
Write func(data []byte) error
Events <-chan []byte
URL string
Close func() error
}
```
### BoxInfo
```go
type BoxInfo struct {
ID string
ContainerID string
Pool string
Owner string
Status string // "running", "stopped", etc.
Policy LifecyclePolicy
Labels map[string]string
Image string
CreatedAt time.Time
LastActive time.Time
ProcessCount int
VNC bool
}
```
### PoolInfo
```go
type PoolInfo struct {
Name string
Addr string
Connected bool
Boxes int
MaxPerUser int
MaxTotal int
IdleTimeout time.Duration
MaxLifetime time.Duration
}
```
### ImagePullOptions / RegistryAuth
```go
type ImagePullOptions struct {
Auth *RegistryAuth // nil = anonymous
}
type RegistryAuth struct {
Username string
Password string
Server string
}
```
### HostExecResult
```go
type HostExecResult struct {
ExitCode int
Stdout []byte
Stderr []byte
DurationMs int64
Error string
Truncated bool
}
```
### HostExecStream
```go
type HostExecStream struct {
Stdout <-chan []byte
Stderr <-chan []byte
Wait func() (int, error) // blocks until exit; returns exit code
Cancel func() // cancels the stream context
}
```
---
## Errors
```go
var (
ErrNotAvailable = errors.New("sandbox: not available (no pools configured)")
ErrNotFound = errors.New("sandbox: not found")
ErrLimitExceeded = errors.New("sandbox: limit exceeded")
ErrPoolNotFound = errors.New("sandbox: pool not found")
ErrPoolInUse = errors.New("sandbox: pool has running boxes")
)
```
---
## Helper Functions
### CreateContainerTokens
```go
func CreateContainerTokens(sandboxID, owner string, scopes []string) (access, refresh string, err error)
```
Creates an OAuth token pair for a sandbox container.
### RevokeContainerTokens
```go
func RevokeContainerTokens(refresh string) error
```
Revokes a container refresh token.
### BuildGRPCEnv
```go
func BuildGRPCEnv(pool *Pool, sandboxID, access, refresh string, grpcPort int) map[string]string
```
Builds environment variables injected into sandbox containers:
| Variable | Description |
|--------------------|--------------------------------------|
| `YAO_SANDBOX_ID` | Sandbox identifier |
| `YAO_TOKEN` | Access token for gRPC auth |
| `YAO_REFRESH_TOKEN` | Refresh token for token rotation |
| `YAO_GRPC_ADDR` | gRPC server address (auto-derived) |
Address derivation logic:
- `tai://host:port` → `host:port` (default port 19100 when omitted)
- `tunnel://...` → `127.0.0.1:<grpcPort>`
- Local/default → `127.0.0.1:<grpcPort>`

6
sandbox/v2/grpc.go
View file

@ -52,9 +52,7 @@ func BuildGRPCEnv(pool *Pool, sandboxID, access, refresh string, grpcPort int) m
switch {
case strings.HasPrefix(pool.Addr, "tunnel://"):
env["YAO_GRPC_TAI"] = "enable"
env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%d", grpcPort)
env["YAO_GRPC_UPSTREAM"] = fmt.Sprintf("127.0.0.1:%s", portStr)
case strings.HasPrefix(pool.Addr, "tai://"):
u, err := url.Parse(pool.Addr)
@ -65,11 +63,9 @@ func BuildGRPCEnv(pool *Pool, sandboxID, access, refresh string, grpcPort int) m
taiHost := u.Hostname()
taiPort := u.Port()
if taiPort == "" {
taiPort = "9100"
taiPort = "19100"
}
env["YAO_GRPC_TAI"] = "enable"
env["YAO_GRPC_ADDR"] = fmt.Sprintf("%s:%s", taiHost, taiPort)
env["YAO_GRPC_UPSTREAM"] = fmt.Sprintf("127.0.0.1:%s", portStr)
default:
env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%s", portStr)

20
sandbox/v2/grpc_test.go
View file

@ -19,23 +19,23 @@ func TestBuildGRPCEnvLocal(t *testing.T) {
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:19100" {
t.Errorf("YAO_GRPC_ADDR = %q, want gpu-server:19100", env["YAO_GRPC_ADDR"])
}
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 TestBuildGRPCEnvTunnel(t *testing.T) {
pool := &sandbox.Pool{Name: "tunnel", Addr: "tunnel://relay.example.com"}
env := sandbox.BuildGRPCEnv(pool, "sb-003", "access", "refresh", 9099)
if env["YAO_GRPC_ADDR"] != "127.0.0.1:9099" {
t.Errorf("YAO_GRPC_ADDR = %q, want 127.0.0.1:9099", env["YAO_GRPC_ADDR"])
}
}

161
sandbox/v2/host.go Normal file
View file

@ -0,0 +1,161 @@
package sandbox
import (
"context"
"fmt"
hepb "github.com/yaoapp/yao/tai/hostexec/pb"
"github.com/yaoapp/yao/tai/workspace"
)
// Host represents a Tai host machine execution environment.
// Unlike Box (which wraps a container), Host executes commands directly on
// the Tai server's OS via HostExec gRPC and accesses files via Volume gRPC.
//
// A Host is bound to a pool and does not require Create — it is available as
// long as the pool's Tai server reports host_exec capability.
type Host struct {
pool string
manager *Manager
}
// Pool returns the pool name this Host belongs to.
func (h *Host) Pool() string { return h.pool }
// Exec runs a command on the Tai host machine via HostExec gRPC.
func (h *Host) Exec(ctx context.Context, cmd string, args []string, opts ...HostExecOption) (*HostExecResult, error) {
client, err := h.manager.getPool(h.pool)
if err != nil {
return nil, err
}
he := client.HostExec()
if he == nil {
return nil, fmt.Errorf("sandbox: host_exec not available on pool %q", h.pool)
}
cfg := &hostExecConfig{}
for _, o := range opts {
o(cfg)
}
req := &hepb.ExecRequest{
Command: cmd,
Args: args,
WorkingDir: cfg.WorkDir,
Stdin: cfg.Stdin,
TimeoutMs: cfg.TimeoutMs,
MaxOutputBytes: cfg.MaxOutputBytes,
}
if cfg.Env != nil {
req.Env = cfg.Env
}
resp, err := he.Exec(ctx, req)
if err != nil {
return nil, fmt.Errorf("hostexec rpc: %w", err)
}
return &HostExecResult{
ExitCode: int(resp.ExitCode),
Stdout: resp.Stdout,
Stderr: resp.Stderr,
DurationMs: resp.DurationMs,
Error: resp.Error,
Truncated: resp.Truncated,
}, nil
}
// Stream runs a command on the Tai host and streams stdout/stderr in real time
// via HostExec gRPC ExecStream. Returns a HostExecStream with separate channels
// for stdout and stderr, plus Wait (blocks until exit) and Cancel.
func (h *Host) Stream(ctx context.Context, cmd string, args []string, opts ...HostExecOption) (*HostExecStream, error) {
client, err := h.manager.getPool(h.pool)
if err != nil {
return nil, err
}
he := client.HostExec()
if he == nil {
return nil, fmt.Errorf("sandbox: host_exec not available on pool %q", h.pool)
}
cfg := &hostExecConfig{}
for _, o := range opts {
o(cfg)
}
req := &hepb.ExecRequest{
Command: cmd,
Args: args,
WorkingDir: cfg.WorkDir,
Stdin: cfg.Stdin,
TimeoutMs: cfg.TimeoutMs,
MaxOutputBytes: cfg.MaxOutputBytes,
}
if cfg.Env != nil {
req.Env = cfg.Env
}
streamCtx, cancel := context.WithCancel(ctx)
rpcStream, err := he.ExecStream(streamCtx, req)
if err != nil {
cancel()
return nil, fmt.Errorf("hostexec stream rpc: %w", err)
}
stdoutCh := make(chan []byte, 64)
stderrCh := make(chan []byte, 64)
doneCh := make(chan struct{})
var exitCode int
var exitErr error
go func() {
defer close(stdoutCh)
defer close(stderrCh)
defer close(doneCh)
for {
msg, err := rpcStream.Recv()
if err != nil {
exitErr = fmt.Errorf("hostexec stream recv: %w", err)
return
}
if len(msg.Data) > 0 {
switch msg.Stream {
case hepb.ExecOutput_STDOUT:
stdoutCh <- msg.Data
case hepb.ExecOutput_STDERR:
stderrCh <- msg.Data
}
}
if msg.Done {
exitCode = int(msg.ExitCode)
if msg.Error != "" {
exitErr = fmt.Errorf("hostexec: %s", msg.Error)
}
return
}
}
}()
return &HostExecStream{
Stdout: stdoutCh,
Stderr: stderrCh,
Wait: func() (int, error) {
<-doneCh
return exitCode, exitErr
},
Cancel: cancel,
}, nil
}
// Workspace returns a filesystem interface for the given session on the host.
// The sessionID typically corresponds to a workspace ID; files are stored
// under dataDir/{sessionID}/ on the Tai host, accessed via Volume gRPC.
func (h *Host) Workspace(sessionID string) workspace.FS {
client, err := h.manager.getPool(h.pool)
if err != nil {
return nil
}
return client.Workspace(sessionID)
}

420
sandbox/v2/host_test.go Normal file
View file

@ -0,0 +1,420 @@
package sandbox_test
import (
"context"
"fmt"
"strings"
"testing"
"time"
sandbox "github.com/yaoapp/yao/sandbox/v2"
"github.com/yaoapp/yao/tai"
)
func setupHostManager(t *testing.T, tgt hostExecTarget) *sandbox.Manager {
t.Helper()
addr := fmt.Sprintf("tai://%s", tgt.Addr)
pool := sandbox.Pool{Name: tgt.Name, Addr: addr}
cfg := sandbox.Config{Pool: []sandbox.Pool{pool}}
if err := sandbox.Init(cfg); err != nil {
t.Fatalf("Init: %v", err)
}
m := sandbox.M()
t.Cleanup(func() { m.Close() })
return m
}
func TestHost_Exec_Echo(t *testing.T) {
skipIfNoHostExec(t)
for _, tgt := range hostExecTargets() {
t.Run(tgt.Name, func(t *testing.T) {
m := setupHostManager(t, tgt)
host, err := m.Host(context.Background(), tgt.Name)
if err != nil {
t.Skipf("Host(%s): %v", tgt.Name, err)
}
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
cmd, args := linuxCmd(tgt, "echo", "hello", "from", "host")
result, err := host.Exec(ctx, cmd, args)
if err != nil {
t.Fatalf("Exec: %v", err)
}
if result.Error != "" {
if strings.Contains(result.Error, "not in the allowed list") {
t.Skipf("command not allowed on %s", tgt.Name)
}
t.Fatalf("error: %s", result.Error)
}
if result.ExitCode != 0 {
t.Errorf("exit_code = %d, want 0", result.ExitCode)
}
got := strings.TrimSpace(string(result.Stdout))
if !strings.Contains(got, "hello") {
t.Errorf("stdout = %q, want contains 'hello'", got)
}
})
}
}
func TestHost_Exec_Env(t *testing.T) {
skipIfNoHostExec(t)
for _, tgt := range hostExecTargets() {
t.Run(tgt.Name, func(t *testing.T) {
m := setupHostManager(t, tgt)
host, err := m.Host(context.Background(), tgt.Name)
if err != nil {
t.Skipf("Host(%s): %v", tgt.Name, err)
}
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
var cmd string
var args []string
if tgt.IsWinNative {
cmd = "cmd.exe"
args = []string{"/c", "echo", "%MY_VAR%"}
} else {
cmd = "sh"
args = []string{"-c", "echo $MY_VAR"}
}
result, err := host.Exec(ctx, cmd, args, sandbox.WithHostEnv(map[string]string{"MY_VAR": "host_test_value"}))
if err != nil {
t.Fatalf("Exec: %v", err)
}
if result.Error != "" {
if strings.Contains(result.Error, "not in the allowed list") {
t.Skipf("command not allowed on %s", tgt.Name)
}
t.Fatalf("error: %s", result.Error)
}
got := strings.TrimSpace(string(result.Stdout))
if !strings.Contains(got, "host_test_value") {
t.Errorf("stdout = %q, want contains 'host_test_value'", got)
}
})
}
}
func TestHost_Workspace(t *testing.T) {
skipIfNoHostExec(t)
for _, tgt := range hostExecTargets() {
t.Run(tgt.Name, func(t *testing.T) {
m := setupHostManager(t, tgt)
host, err := m.Host(context.Background(), tgt.Name)
if err != nil {
t.Skipf("Host(%s): %v", tgt.Name, err)
}
sessionID := fmt.Sprintf("host-test-%d", time.Now().UnixNano())
ws := host.Workspace(sessionID)
if ws == nil {
t.Fatal("Workspace returned nil")
}
content := []byte("hello from host workspace test")
if err := ws.WriteFile("test.txt", content, 0644); err != nil {
t.Fatalf("WriteFile: %v", err)
}
got, err := ws.ReadFile("test.txt")
if err != nil {
t.Fatalf("ReadFile: %v", err)
}
if string(got) != string(content) {
t.Errorf("ReadFile = %q, want %q", got, content)
}
if err := ws.MkdirAll("sub/dir", 0755); err != nil {
t.Fatalf("MkdirAll: %v", err)
}
if err := ws.WriteFile("sub/dir/nested.txt", []byte("nested"), 0644); err != nil {
t.Fatalf("WriteFile nested: %v", err)
}
entries, err := ws.ReadDir("sub/dir")
if err != nil {
t.Fatalf("ReadDir: %v", err)
}
if len(entries) != 1 {
t.Errorf("ReadDir len = %d, want 1", len(entries))
}
if err := ws.RemoveAll(sessionID); err != nil && !strings.Contains(err.Error(), "not found") {
t.Logf("cleanup RemoveAll: %v", err)
}
})
}
}
func TestHost_Stream_Incremental(t *testing.T) {
skipIfNoHostExec(t)
for _, tgt := range hostExecTargets() {
if tgt.IsWinNative {
continue
}
t.Run(tgt.Name, func(t *testing.T) {
m := setupHostManager(t, tgt)
host, err := m.Host(context.Background(), tgt.Name)
if err != nil {
t.Skipf("Host(%s): %v", tgt.Name, err)
}
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
stream, err := host.Stream(ctx, "sh", []string{"-c",
"for i in 1 2 3 4 5; do echo chunk$i; sleep 0.2; done"})
if err != nil {
t.Fatalf("Stream: %v", err)
}
var chunks []string
for chunk := range stream.Stdout {
chunks = append(chunks, string(chunk))
}
exitCode, err := stream.Wait()
if err != nil && !strings.Contains(err.Error(), "EOF") {
if strings.Contains(err.Error(), "not in the allowed list") {
t.Skipf("command not allowed on %s", tgt.Name)
}
t.Fatalf("Wait: %v", err)
}
if exitCode != 0 {
t.Errorf("exit_code = %d, want 0", exitCode)
}
combined := strings.Join(chunks, "")
for _, expect := range []string{"chunk1", "chunk3", "chunk5"} {
if !strings.Contains(combined, expect) {
t.Errorf("output = %q, want contains %q", combined, expect)
}
}
if len(chunks) < 2 {
t.Errorf("received %d chunks, want >= 2 (proves streaming, not buffered)", len(chunks))
}
t.Logf("received %d chunks over stream", len(chunks))
})
}
}
func TestHost_Stream_MultiLine(t *testing.T) {
skipIfNoHostExec(t)
for _, tgt := range hostExecTargets() {
if tgt.IsWinNative {
continue
}
t.Run(tgt.Name, func(t *testing.T) {
m := setupHostManager(t, tgt)
host, err := m.Host(context.Background(), tgt.Name)
if err != nil {
t.Skipf("Host(%s): %v", tgt.Name, err)
}
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
stream, err := host.Stream(ctx, "sh", []string{"-c", "for i in 1 2 3; do echo line$i; done"})
if err != nil {
t.Fatalf("Stream: %v", err)
}
var stdout []byte
for chunk := range stream.Stdout {
stdout = append(stdout, chunk...)
}
exitCode, err := stream.Wait()
if err != nil && !strings.Contains(err.Error(), "EOF") {
if strings.Contains(err.Error(), "not in the allowed list") {
t.Skipf("command not allowed on %s", tgt.Name)
}
t.Fatalf("Wait: %v", err)
}
if exitCode != 0 {
t.Errorf("exit_code = %d, want 0", exitCode)
}
got := strings.TrimSpace(string(stdout))
for _, expect := range []string{"line1", "line2", "line3"} {
if !strings.Contains(got, expect) {
t.Errorf("stdout = %q, want contains %q", got, expect)
}
}
})
}
}
func TestHost_Stream_Stderr(t *testing.T) {
skipIfNoHostExec(t)
for _, tgt := range hostExecTargets() {
if tgt.IsWinNative {
continue
}
t.Run(tgt.Name, func(t *testing.T) {
m := setupHostManager(t, tgt)
host, err := m.Host(context.Background(), tgt.Name)
if err != nil {
t.Skipf("Host(%s): %v", tgt.Name, err)
}
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
stream, err := host.Stream(ctx, "sh", []string{"-c", "echo err-msg >&2"})
if err != nil {
t.Fatalf("Stream: %v", err)
}
var stderr []byte
done := make(chan struct{})
go func() {
for chunk := range stream.Stdout {
_ = chunk
}
close(done)
}()
for chunk := range stream.Stderr {
stderr = append(stderr, chunk...)
}
<-done
exitCode, err := stream.Wait()
if err != nil && !strings.Contains(err.Error(), "EOF") {
if strings.Contains(err.Error(), "not in the allowed list") {
t.Skipf("command not allowed on %s", tgt.Name)
}
t.Fatalf("Wait: %v", err)
}
if exitCode != 0 {
t.Errorf("exit_code = %d, want 0", exitCode)
}
got := strings.TrimSpace(string(stderr))
if !strings.Contains(got, "err-msg") {
t.Errorf("stderr = %q, want contains 'err-msg'", got)
}
})
}
}
func TestHost_Stream_Cancel(t *testing.T) {
skipIfNoHostExec(t)
for _, tgt := range hostExecTargets() {
if tgt.IsWinNative {
continue
}
t.Run(tgt.Name, func(t *testing.T) {
m := setupHostManager(t, tgt)
host, err := m.Host(context.Background(), tgt.Name)
if err != nil {
t.Skipf("Host(%s): %v", tgt.Name, err)
}
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
stream, err := host.Stream(ctx, "sh", []string{"-c", "while true; do echo tick; sleep 0.1; done"})
if err != nil {
if strings.Contains(err.Error(), "not in the allowed list") {
t.Skipf("command not allowed on %s", tgt.Name)
}
t.Fatalf("Stream: %v", err)
}
received := 0
for chunk := range stream.Stdout {
_ = chunk
received++
if received >= 3 {
stream.Cancel()
break
}
}
_, waitErr := stream.Wait()
if waitErr != nil && strings.Contains(waitErr.Error(), "not in the allowed list") {
t.Skipf("command not allowed on %s", tgt.Name)
}
if received < 3 && waitErr == nil {
t.Errorf("received %d chunks before cancel, want >= 3", received)
}
})
}
}
func TestHost_CreateRejectsNoContainerPool(t *testing.T) {
// Use the Windows native HostExec target which has no Docker.
tgt := findHostExecOnly(t)
if tgt == nil {
t.Skip("no host-exec-only target available")
}
m := setupHostManager(t, *tgt)
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cancel()
_, err := m.Create(ctx, sandbox.CreateOptions{
Image: "alpine:latest",
Owner: "test",
Pool: tgt.Name,
})
if err == nil {
t.Fatal("expected error for Create on host-exec-only pool, got nil")
}
if !strings.Contains(err.Error(), "no container runtime") {
t.Errorf("error = %q, want contains 'no container runtime'", err.Error())
}
}
func TestHost_PoolNotFound(t *testing.T) {
skipIfNoHostExec(t)
tgt := hostExecTargets()[0]
m := setupHostManager(t, tgt)
_, err := m.Host(context.Background(), "nonexistent-pool")
if err == nil {
t.Fatal("expected error, got nil")
}
}
// findHostExecOnly returns a hostExecTarget that is likely host-exec-only
// (Windows native Tai without Docker).
func findHostExecOnly(t *testing.T) *hostExecTarget {
t.Helper()
for _, tgt := range hostExecTargets() {
if tgt.IsWinNative {
// Windows native Tai typically has no Docker
addr := fmt.Sprintf("tai://%s", tgt.Addr)
client, err := tai.New(addr)
if err != nil {
continue
}
hasNoSandbox := client.Sandbox() == nil
client.Close()
if hasNoSandbox {
return &tgt
}
}
}
return nil
}

57
sandbox/v2/manager.go
View file

@ -164,6 +164,32 @@ func (m *Manager) Heartbeat(sandboxID string, active bool, processCount int) err
return nil
}
// Host returns a Host handle for executing commands on the Tai host machine.
// The pool must be connected to a Tai server with host_exec capability.
// Unlike Create/Box, Host does not create a container — it is available
// immediately as long as the pool is reachable.
func (m *Manager) Host(_ context.Context, pool string) (*Host, error) {
if pool == "" {
pool = m.defaultPool
}
pd := m.findPoolDef(pool)
if pd == nil {
return nil, ErrPoolNotFound
}
client, err := m.getPool(pool)
if err != nil {
return nil, fmt.Errorf("sandbox: connect pool %q: %w", pool, err)
}
if client.HostExec() == nil {
return nil, fmt.Errorf("sandbox: pool %q has no host_exec capability", pool)
}
return &Host{pool: pool, manager: m}, nil
}
// Create creates and starts a new sandbox.
func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Box, error) {
if len(m.poolDefs) == 0 {
@ -207,6 +233,10 @@ func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Box, error)
return nil, fmt.Errorf("sandbox: connect pool %q: %w", poolName, err)
}
if client.Sandbox() == nil {
return nil, fmt.Errorf("sandbox: pool %q has no container runtime", poolName)
}
access, refresh, err := CreateContainerTokens(id, opts.Owner, nil)
if err != nil {
return nil, fmt.Errorf("sandbox: create tokens: %w", err)
@ -303,7 +333,7 @@ func (m *Manager) Remove(ctx context.Context, id string) error {
b := v.(*Box)
client, err := m.getPool(b.pool)
if err == nil {
if err == nil && client.Sandbox() != nil {
client.Sandbox().Remove(ctx, b.containerID, true)
}
@ -331,7 +361,7 @@ func (m *Manager) Cleanup(ctx context.Context) error {
}
case LongRunning:
if timeout := b.idleTimeout(); timeout > 0 && idle > timeout {
if client, err := m.getPool(b.pool); err == nil {
if client, err := m.getPool(b.pool); err == nil && client.Sandbox() != nil {
client.Sandbox().Stop(ctx, b.containerID, b.stopTimeout())
}
}
@ -526,6 +556,9 @@ func (m *Manager) buildTaiCreateOptions(opts CreateOptions, pd *Pool, sandboxID,
}
func (m *Manager) recoverBoxes(ctx context.Context, pd *Pool, client *tai.Client) {
if client.Sandbox() == nil {
return
}
containers, err := client.Sandbox().List(ctx, taisandbox.ListOptions{
All: true,
Labels: map[string]string{"managed-by": "yao-sandbox"},
@ -543,9 +576,13 @@ func (m *Manager) recoverBoxes(ctx context.Context, pd *Pool, client *tai.Client
continue
}
cid := c.ID
if c.Name != "" {
cid = c.Name
}
box := &Box{
id: sandboxID,
containerID: c.ID,
containerID: cid,
pool: c.Labels["sandbox-pool"],
owner: c.Labels["sandbox-owner"],
policy: LifecyclePolicy(c.Labels["sandbox-policy"]),
@ -561,12 +598,18 @@ func (m *Manager) recoverBoxes(ctx context.Context, pd *Pool, client *tai.Client
}
// ImageExists reports whether the given image ref exists on the target pool node.
// Returns (true, nil) when the pool has no image service (e.g. K8s — kubelet
// handles image pulls transparently).
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)
img := client.Image()
if img == nil {
return true, nil
}
return img.Exists(ctx, ref)
}
// PullImage pulls an image to the target pool node, returning a channel of
@ -576,6 +619,10 @@ func (m *Manager) PullImage(ctx context.Context, pool, ref string, opts ImagePul
if err != nil {
return nil, err
}
img := client.Image()
if img == nil {
return nil, nil
}
pullOpts := taisandbox.PullOptions{}
if opts.Auth != nil {
pullOpts.Auth = &taisandbox.RegistryAuth{
@ -584,7 +631,7 @@ func (m *Manager) PullImage(ctx context.Context, pool, ref string, opts ImagePul
Server: opts.Auth.Server,
}
}
return client.Image().Pull(ctx, ref, pullOpts)
return img.Pull(ctx, ref, pullOpts)
}
// EnsureImage checks whether the image exists on the pool node; if not, it

2
sandbox/v2/manager_lifecycle_test.go
View file

@ -84,7 +84,7 @@ func TestStartRecovery(t *testing.T) {
for _, pc := range testPools() {
t.Run(pc.Name, func(t *testing.T) {
pool := sandbox.Pool{Name: pc.Name, Addr: pc.Addr}
pool := sandbox.Pool{Name: pc.Name, Addr: pc.Addr, Options: pc.Options}
m1 := setupManager(t, pool)
box := createTestBox(t, m1)

173
sandbox/v2/testutils_test.go
View file

@ -3,17 +3,101 @@ package sandbox_test
import (
"context"
"fmt"
"log"
"os"
"strconv"
"strings"
"sync"
"testing"
"time"
sandbox "github.com/yaoapp/yao/sandbox/v2"
"github.com/yaoapp/yao/tai"
taisandbox "github.com/yaoapp/yao/tai/sandbox"
"github.com/yaoapp/yao/tai/volume"
"github.com/yaoapp/yao/workspace"
)
// k8sSem limits concurrent K8s pod creation to avoid overwhelming the cluster.
var k8sSem = make(chan struct{}, 2)
// k8sCleanupMu serialises K8s pod cleanup to prevent overlapping API calls
// when many tests finish at once.
var k8sCleanupMu sync.Mutex
func TestMain(m *testing.M) {
purgeStaleContainers()
os.Exit(m.Run())
}
// purgeStaleContainers removes leftover sb-* containers/pods from previous
// test runs across all configured pools (Docker + K8s).
func purgeStaleContainers() {
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
type target struct {
name string
addr string
opts []tai.Option
}
var targets []target
targets = append(targets, target{name: "local", addr: testLocalAddr()})
if addr := os.Getenv("SANDBOX_TEST_REMOTE_ADDR"); addr != "" {
targets = append(targets, target{name: "remote", addr: addr})
}
if host := os.Getenv("TAI_TEST_CONTAINERIZED_HOST"); host != "" {
grpcPort := envPort("TAI_TEST_CONTAINERIZED_GRPC_PORT", 9200)
targets = append(targets, target{name: "containerized", addr: fmt.Sprintf("tai://%s:%d", host, grpcPort)})
}
if host := os.Getenv("TAI_TEST_K8S_HOST"); host != "" {
kubeconfig := os.Getenv("TAI_TEST_KUBECONFIG")
if kubeconfig != "" {
grpcPort := envPort("TAI_TEST_K8S_GRPC_PORT", envPort("TAI_TEST_GRPC_PORT", 19100))
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))
}
targets = append(targets, target{name: "k8s", addr: fmt.Sprintf("tai://%s:%d", host, grpcPort), opts: opts})
}
}
for _, tgt := range targets {
client, err := tai.New(tgt.addr, tgt.opts...)
if err != nil {
continue
}
sb := client.Sandbox()
if sb == nil {
client.Close()
continue
}
containers, err := sb.List(ctx, taisandbox.ListOptions{All: true})
if err != nil {
client.Close()
continue
}
for _, c := range containers {
id := c.Name
if id == "" {
id = c.ID
}
if !strings.HasPrefix(id, "sb-") && !strings.HasPrefix(c.Labels["sandbox-id"], "sb-") {
continue
}
sb.Remove(ctx, id, true)
log.Printf("[purge] %s: removed stale container %s", tgt.name, id)
}
client.Close()
}
}
type poolConfig struct {
Name string
Addr string
@ -44,7 +128,7 @@ func testPools() []poolConfig {
if kubeconfig == "" {
return pools
}
grpcPort := envPort("TAI_TEST_K8S_GRPC_PORT", envPort("TAI_TEST_GRPC_PORT", 9100))
grpcPort := envPort("TAI_TEST_K8S_GRPC_PORT", envPort("TAI_TEST_GRPC_PORT", 19100))
addr := fmt.Sprintf("tai://%s:%d", host, grpcPort)
opts := []tai.Option{
tai.K8s,
@ -77,6 +161,67 @@ func skipIfNoTai(t *testing.T) {
}
}
type hostExecTarget struct {
Name string
Addr string // host:port (without tai:// prefix)
IsWinNative bool
}
// hostExecTargets returns all Tai instances that support HostExec gRPC.
// No container creation needed — these are direct gRPC connections.
func hostExecTargets() []hostExecTarget {
var targets []hostExecTarget
if addr := os.Getenv("SANDBOX_TEST_REMOTE_ADDR"); addr != "" {
addr = strings.TrimPrefix(addr, "tai://")
targets = append(targets, hostExecTarget{Name: "remote", Addr: addr})
}
if host := os.Getenv("TAI_TEST_K8S_HOST"); host != "" {
grpcPort := envPort("TAI_TEST_K8S_GRPC_PORT", envPort("TAI_TEST_GRPC_PORT", 19100))
targets = append(targets, hostExecTarget{Name: "k8s", Addr: fmt.Sprintf("%s:%d", host, grpcPort)})
}
if addr := os.Getenv("TAI_TEST_WIN_HOSTEXEC_LINUX"); addr != "" {
targets = append(targets, hostExecTarget{Name: "win-linux", Addr: addr})
}
if addr := os.Getenv("TAI_TEST_WIN_HOSTEXEC_NATIVE"); addr != "" {
targets = append(targets, hostExecTarget{Name: "win-native", Addr: addr, IsWinNative: true})
}
return targets
}
func skipIfNoHostExec(t *testing.T) {
t.Helper()
if len(hostExecTargets()) == 0 {
t.Skip("no HostExec targets configured")
}
}
// linuxCmd adapts a Linux command to the equivalent Windows command for
// Windows native Tai targets.
func linuxCmd(tgt hostExecTarget, cmd string, args ...string) (string, []string) {
if tgt.IsWinNative {
switch cmd {
case "echo":
return "cmd.exe", append([]string{"/c", "echo"}, args...)
case "pwd":
return "cmd.exe", []string{"/c", "cd"}
case "env":
return "cmd.exe", []string{"/c", "set"}
case "sleep":
return "cmd.exe", []string{"/c", "ping", "-n", "10", "127.0.0.1"}
case "cat":
return "cmd.exe", []string{"/c", "more"}
case "sh":
if len(args) >= 2 && args[0] == "-c" {
return "cmd.exe", []string{"/c", args[1]}
}
return "cmd.exe", append([]string{"/c"}, args...)
default:
return cmd, args
}
}
return cmd, args
}
func testLocalAddr() string {
if addr := os.Getenv("SANDBOX_TEST_LOCAL_ADDR"); addr != "" {
return addr
@ -176,21 +321,43 @@ func createTestBox(t *testing.T, m *sandbox.Manager, opts ...func(*sandbox.Creat
}
}
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
isK8s := pool == "k8s"
if isK8s {
k8sSem <- struct{}{}
}
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
if pool != "" {
if err := m.EnsureImage(ctx, pool, co.Image, sandbox.ImagePullOptions{}); err != nil {
if isK8s {
<-k8sSem
}
t.Fatalf("EnsureImage(%s, %s): %v", pool, co.Image, err)
}
}
box, err := m.Create(ctx, co)
if err != nil {
if isK8s {
<-k8sSem
}
t.Fatalf("Create: %v", err)
}
t.Cleanup(func() {
m.Remove(context.Background(), box.ID())
cleanCtx, cleanCancel := context.WithTimeout(context.Background(), 15*time.Second)
defer cleanCancel()
if isK8s {
k8sCleanupMu.Lock()
defer k8sCleanupMu.Unlock()
}
if err := m.Remove(cleanCtx, box.ID()); err != nil {
t.Logf("cleanup Remove(%s): %v", box.ID(), err)
}
if isK8s {
<-k8sSem
}
})
return box
}

50
sandbox/v2/types.go
View file

@ -176,3 +176,53 @@ type BoxInfo struct {
ProcessCount int
VNC bool
}
// HostExecResult holds the outcome of a command executed on the Tai host.
type HostExecResult struct {
ExitCode int
Stdout []byte
Stderr []byte
DurationMs int64
Error string
Truncated bool
}
// HostExecStream provides real-time streaming output from a command running
// on the Tai host machine via HostExec gRPC ExecStream.
type HostExecStream struct {
Stdout <-chan []byte
Stderr <-chan []byte
Wait func() (int, error) // blocks until exit; returns exit code
Cancel func() // cancels the stream context
}
type hostExecConfig struct {
WorkDir string
Env map[string]string
Stdin []byte
TimeoutMs int64
MaxOutputBytes int64
}
// HostExecOption configures an ExecOnHost call.
type HostExecOption func(*hostExecConfig)
func WithHostWorkDir(dir string) HostExecOption {
return func(c *hostExecConfig) { c.WorkDir = dir }
}
func WithHostEnv(env map[string]string) HostExecOption {
return func(c *hostExecConfig) { c.Env = env }
}
func WithHostStdin(data []byte) HostExecOption {
return func(c *hostExecConfig) { c.Stdin = data }
}
func WithHostTimeout(ms int64) HostExecOption {
return func(c *hostExecConfig) { c.TimeoutMs = ms }
}
func WithHostMaxOutput(bytes int64) HostExecOption {
return func(c *hostExecConfig) { c.MaxOutputBytes = bytes }
}

105
service/service.go
View file

@ -1,6 +1,7 @@
package service
import (
"fmt"
"time"
"github.com/gin-gonic/gin"
@ -15,8 +16,23 @@ import (
// requests internally without an HTTP round-trip.
var Router *gin.Engine
// Start the yao service
func Start(cfg config.Config) (*http.Server, error) {
// ServerHooks allows the caller to inject gRPC (or other) server lifecycle
// without creating import cycles.
type ServerHooks struct {
Start func(cfg config.Config) error // called before HTTP starts; nil = skip
Stop func() // called on shutdown; nil = skip
Addrs func() []string // returns listen addresses; nil = skip
}
// Service manages HTTP and optional gRPC servers as a single unit.
type Service struct {
http *http.Server
hooks ServerHooks
}
// Start launches optional hook servers (e.g. gRPC) and the HTTP server.
// Returns a Service handle for shutdown coordination.
func Start(cfg config.Config, hooks ...ServerHooks) (*Service, error) {
if cfg.AllowFrom == nil {
cfg.AllowFrom = []string{}
@ -27,29 +43,31 @@ func Start(cfg config.Config) (*http.Server, error) {
return nil, err
}
var h ServerHooks
if len(hooks) > 0 {
h = hooks[0]
}
// Start hook server (gRPC, etc.)
if h.Start != nil {
if err := h.Start(cfg); err != nil {
return nil, err
}
}
router := gin.New()
Router = router
router.Use(Middlewares...)
var apiRoot string
if openapi.Server != nil {
// OpenAPI mode: use OAuth guards and dynamic routing
apiRoot = openapi.Server.Config.BaseURL
api.SetGuards(OpenAPIGuards())
// Developer APIs: use dynamic proxy (supports hot-reload)
router.Any(apiRoot+"/api/*path", DynamicAPIHandler)
// Widgets and system APIs: static registration
api.SetRoutes(router, apiRoot, cfg.AllowFrom...)
// Build route table for dynamic lookup
api.BuildRouteTable()
// Attach OpenAPI built-in features
openapi.Server.Attach(router)
} else {
// Traditional mode: unchanged
apiRoot = "/api"
api.SetGuards(Guards)
api.SetRoutes(router, "/api", cfg.AllowFrom...)
@ -63,21 +81,57 @@ func Start(cfg config.Config) (*http.Server, error) {
Timeout: 5 * time.Second,
})
// Start HTTP in background; wait for the first event to confirm
// the port is bound before returning.
go func() {
err = srv.Start()
srv.Start()
}()
return srv, nil
// Block until HTTP reports READY or ERROR
ev := <-srv.Event()
if ev != http.READY {
if h.Stop != nil {
h.Stop()
}
return nil, fmt.Errorf("HTTP server failed to start on %s:%d", cfg.Host, cfg.Port)
}
return &Service{http: srv, hooks: h}, nil
}
// Restart the yao service
func Restart(srv *http.Server, cfg config.Config) error {
// Event returns the HTTP server event channel (READY, CLOSED, ERROR).
func (s *Service) Event() chan uint8 {
return s.http.Event()
}
// Stop shuts down hook servers (gRPC, etc.) then signals the HTTP server to close.
func (s *Service) Stop() {
if s.hooks.Stop != nil {
s.hooks.Stop()
}
s.http.Stop()
}
// HookAddrs returns the hook server listen addresses (e.g. gRPC addresses).
func (s *Service) HookAddrs() []string {
if s.hooks.Addrs != nil {
return s.hooks.Addrs()
}
return nil
}
// Watch starts file watching in development mode. Blocking; run in a goroutine.
func (s *Service) Watch(done chan uint8) {
watch(s, done)
}
// Restart the HTTP server with a fresh router (hook servers stay running).
func Restart(svc *Service, cfg config.Config) error {
router := gin.New()
Router = router
router.Use(Middlewares...)
if openapi.Server != nil {
// OpenAPI mode
baseURL := openapi.Server.Config.BaseURL
api.SetGuards(OpenAPIGuards())
router.Any(baseURL+"/api/*path", DynamicAPIHandler)
@ -85,28 +139,15 @@ func Restart(srv *http.Server, cfg config.Config) error {
api.BuildRouteTable()
openapi.Server.Attach(router)
} else {
// Traditional mode: unchanged
api.SetGuards(Guards)
api.SetRoutes(router, "/api", cfg.AllowFrom...)
}
srv.Reset(router)
return srv.Restart()
}
// Stop the yao service
func Stop(srv *http.Server) error {
err := srv.Stop()
if err != nil {
return err
}
<-srv.Event()
return nil
svc.http.Reset(router)
return svc.http.Restart()
}
func prepare() error {
// Session server
err := share.SessionStart()
if err != nil {
return err

21
service/service_test.go
View file

@ -31,24 +31,12 @@ func TestStartStop(t *testing.T) {
if err != nil {
t.Fatal(err)
}
defer Stop(srv)
defer srv.Stop()
<-srv.Event()
if !srv.Ready() {
t.Fatal("server not ready")
}
port, err := srv.Port()
if err != nil {
t.Fatal(err)
}
if port <= 0 {
t.Fatal("invalid port")
}
// API Server
req := test.NewRequest(port).Route("/api/__yao/app/setting")
req := test.NewRequest(cfg.Port).Route("/api/__yao/app/setting")
res, err := req.Get()
if err != nil {
t.Fatal(err)
@ -58,11 +46,10 @@ func TestStartStop(t *testing.T) {
if err != nil {
t.Fatal(err)
}
// assert.Equal(t, "Demo Application", data["name"])
assert.True(t, len(data["name"].(string)) > 0)
// Public
req = test.NewRequest(port).Route("/")
req = test.NewRequest(cfg.Port).Route("/")
res, err = req.Get()
if err != nil {
t.Fatal(err)
@ -71,7 +58,7 @@ func TestStartStop(t *testing.T) {
assert.Equal(t, "Hello World\n", res.Body())
// XGEN
req = test.NewRequest(port).Route("/admin/")
req = test.NewRequest(cfg.Port).Route("/admin/")
res, err = req.Get()
if err != nil {
t.Fatal(err)

12
service/watch.go
View file

@ -6,14 +6,13 @@ import (
"github.com/fatih/color"
"github.com/yaoapp/gou/application"
"github.com/yaoapp/gou/server/http"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/engine"
"github.com/yaoapp/yao/openapi"
)
// Watch the application code change for hot update
func Watch(srv *http.Server, interrupt chan uint8) (err error) {
func watch(svc *Service, interrupt chan uint8) error {
if application.App == nil {
return fmt.Errorf("Application is not initialized")
@ -24,23 +23,19 @@ func Watch(srv *http.Server, interrupt chan uint8) (err error) {
return
}
// Reload
err = engine.Reload(config.Conf, engine.LoadOption{Action: "watch"})
err := engine.Reload(config.Conf, engine.LoadOption{Action: "watch"})
if err != nil {
fmt.Println(color.RedString("[Watch] Reload: %s", err.Error()))
return
}
fmt.Println(color.GreenString("[Watch] Reload Completed"))
// Model
if strings.HasPrefix(name, "/models") {
fmt.Println(color.GreenString("[Watch] Model: %s changed (Please run yao migrate manually)", name))
}
// API changes: hot reload or restart
if strings.HasPrefix(name, "/apis") {
if openapi.Server != nil {
// OpenAPI mode: hot reload (no server restart needed)
err = ReloadAPIs()
if err != nil {
fmt.Println(color.RedString("[Watch] Reload APIs: %s", err.Error()))
@ -48,8 +43,7 @@ func Watch(srv *http.Server, interrupt chan uint8) (err error) {
}
fmt.Println(color.GreenString("[Watch] APIs Reloaded"))
} else {
// Traditional mode: restart server
err = Restart(srv, config.Conf)
err = Restart(svc, config.Conf)
if err != nil {
fmt.Println(color.RedString("[Watch] Restart: %s", err.Error()))
return

4
service/watch_test.go
View file

@ -18,10 +18,10 @@ func TestWatch(t *testing.T) {
if err != nil {
t.Fatal(err)
}
defer Stop(srv)
defer srv.Stop()
done := make(chan uint8, 1)
go Watch(srv, done)
go srv.Watch(done)
select {
case <-time.After(200 * time.Millisecond):

205
tai/api/register.go Normal file
View file

@ -0,0 +1,205 @@
package api
import (
"fmt"
"log/slog"
"net/http"
"strings"
"github.com/gin-gonic/gin"
"github.com/yaoapp/yao/openapi/oauth"
"github.com/yaoapp/yao/tai/registry"
)
// authenticateBearer validates a Bearer token and returns the caller's identity.
// Package-level var so tests can inject a mock without an OAuth service.
var authenticateBearer = authenticateBearerDefault
func authenticateBearerDefault(token string) (registry.AuthInfo, error) {
svc := oauth.OAuth
if svc == nil {
return registry.AuthInfo{}, fmt.Errorf("oauth service not initialized")
}
result, err := svc.AuthenticateToken(oauth.AuthInput{AccessToken: token})
if err != nil {
return registry.AuthInfo{}, err
}
info := registry.AuthInfo{}
if result.Info != nil {
info.Subject = result.Info.Subject
info.UserID = result.Info.UserID
info.ClientID = result.Info.ClientID
info.Scope = result.Info.Scope
info.TeamID = result.Info.TeamID
info.TenantID = result.Info.TenantID
}
return info, nil
}
func extractBearer(r *http.Request) string {
auth := r.Header.Get("Authorization")
if len(auth) > 7 && strings.EqualFold(auth[:7], "bearer ") {
return auth[7:]
}
return ""
}
// registerRequest is the JSON body for POST /tai-nodes/register.
type registerRequest struct {
TaiID string `json:"tai_id"`
MachineID string `json:"machine_id"`
Version string `json:"version"`
Addr string `json:"addr"`
Ports map[string]int `json:"ports"`
Capabilities map[string]bool `json:"capabilities"`
System registry.SystemInfo `json:"system"`
}
// heartbeatRequest is the JSON body for POST /tai-nodes/heartbeat.
type heartbeatRequest struct {
TaiID string `json:"tai_id"`
}
// HandleRegister handles POST /tai-nodes/register.
// Validates Bearer token, extracts AuthInfo, and writes the node to the Registry.
func HandleRegister(c *gin.Context) {
reg := registry.Global()
if reg == nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "registry not initialized"})
return
}
bearer := extractBearer(c.Request)
if bearer == "" {
c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authorization"})
return
}
authInfo, err := authenticateBearer(bearer)
if err != nil {
slog.Warn("tai register auth failed", "err", err)
c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"})
return
}
var req registerRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
return
}
if req.TaiID == "" {
c.JSON(http.StatusBadRequest, gin.H{"error": "tai_id is required"})
return
}
node := &registry.TaiNode{
TaiID: req.TaiID,
MachineID: req.MachineID,
Version: req.Version,
Auth: authInfo,
System: req.System,
Mode: "direct",
Addr: req.Addr,
Ports: req.Ports,
Capabilities: req.Capabilities,
}
reg.Register(node)
remoteIP := c.ClientIP()
slog.Info("tai node registered via API",
"tai_id", req.TaiID, "remote_ip", remoteIP, "user_id", authInfo.UserID)
c.JSON(http.StatusOK, gin.H{
"status": "registered",
"tai_id": req.TaiID,
"remote_ip": remoteIP,
})
}
// HandleHeartbeat handles POST /tai-nodes/heartbeat.
// Validates Bearer token and updates the node's last ping timestamp.
func HandleHeartbeat(c *gin.Context) {
reg := registry.Global()
if reg == nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "registry not initialized"})
return
}
bearer := extractBearer(c.Request)
if bearer == "" {
c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authorization"})
return
}
authInfo, err := authenticateBearer(bearer)
if err != nil {
c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"})
return
}
var req heartbeatRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request body"})
return
}
if req.TaiID == "" {
c.JSON(http.StatusBadRequest, gin.H{"error": "tai_id is required"})
return
}
snap, ok := reg.Get(req.TaiID)
if !ok {
c.JSON(http.StatusNotFound, gin.H{"error": "tai node not found"})
return
}
if snap.Auth.ClientID != authInfo.ClientID {
c.JSON(http.StatusForbidden, gin.H{"error": "tai_id does not belong to this client"})
return
}
reg.UpdatePing(req.TaiID)
c.JSON(http.StatusOK, gin.H{"status": "ok"})
}
// HandleUnregister handles DELETE /tai-nodes/register/:tai_id.
// Validates Bearer token, checks ownership, and removes the node.
func HandleUnregister(c *gin.Context) {
reg := registry.Global()
if reg == nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "registry not initialized"})
return
}
bearer := extractBearer(c.Request)
if bearer == "" {
c.JSON(http.StatusUnauthorized, gin.H{"error": "missing authorization"})
return
}
authInfo, err := authenticateBearer(bearer)
if err != nil {
c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"})
return
}
taiID := c.Param("tai_id")
if taiID == "" {
c.JSON(http.StatusBadRequest, gin.H{"error": "tai_id is required"})
return
}
snap, ok := reg.Get(taiID)
if !ok {
c.JSON(http.StatusNotFound, gin.H{"error": "tai node not found"})
return
}
if snap.Auth.ClientID != authInfo.ClientID {
c.JSON(http.StatusForbidden, gin.H{"error": "tai_id does not belong to this client"})
return
}
reg.Unregister(taiID)
slog.Info("tai node unregistered via API", "tai_id", taiID, "user_id", authInfo.UserID)
c.JSON(http.StatusOK, gin.H{"status": "unregistered"})
}

267
tai/api/register_test.go Normal file
View file

@ -0,0 +1,267 @@
package api
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/gin-gonic/gin"
"github.com/yaoapp/yao/tai/registry"
)
func init() {
gin.SetMode(gin.TestMode)
}
func setupTest() func() {
r := registry.NewForTest()
registry.SetGlobalForTest(r)
origAuth := authenticateBearer
authenticateBearer = func(token string) (registry.AuthInfo, error) {
return registry.AuthInfo{
Subject: "sub-001",
UserID: "user-alice",
ClientID: "tai-abc123",
Scope: "tai:connect",
TeamID: "team-dev",
}, nil
}
return func() {
authenticateBearer = origAuth
registry.SetGlobalForTest(nil)
}
}
func jsonBody(v interface{}) *bytes.Buffer {
b, _ := json.Marshal(v)
return bytes.NewBuffer(b)
}
func TestHandleRegister_Success(t *testing.T) {
teardown := setupTest()
defer teardown()
body := registerRequest{
TaiID: "tai-abc123",
MachineID: "m-001",
Version: "0.2.0",
Addr: "192.168.1.100",
Ports: map[string]int{"grpc": 19100, "http": 8099},
Capabilities: map[string]bool{"docker": true, "host_exec": false},
System: registry.SystemInfo{
OS: "linux", Arch: "amd64", Hostname: "docker-host-01", NumCPU: 16,
},
}
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("POST", "/tai-nodes/register", jsonBody(body))
c.Request.Header.Set("Authorization", "Bearer test-token")
c.Request.Header.Set("Content-Type", "application/json")
HandleRegister(c)
if w.Code != http.StatusOK {
t.Fatalf("status = %d, want %d; body = %s", w.Code, http.StatusOK, w.Body.String())
}
var resp map[string]interface{}
json.Unmarshal(w.Body.Bytes(), &resp)
if resp["status"] != "registered" {
t.Errorf("status = %v, want registered", resp["status"])
}
if resp["tai_id"] != "tai-abc123" {
t.Errorf("tai_id = %v, want tai-abc123", resp["tai_id"])
}
if _, ok := resp["remote_ip"]; !ok {
t.Error("response missing remote_ip")
}
snap, ok := registry.Global().Get("tai-abc123")
if !ok {
t.Fatal("node not found in registry after register")
}
if snap.Mode != "direct" {
t.Errorf("Mode = %q, want direct", snap.Mode)
}
if snap.System.OS != "linux" {
t.Errorf("System.OS = %q, want linux", snap.System.OS)
}
if snap.Auth.UserID != "user-alice" {
t.Errorf("Auth.UserID = %q, want user-alice", snap.Auth.UserID)
}
}
func TestHandleRegister_MissingAuth(t *testing.T) {
teardown := setupTest()
defer teardown()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("POST", "/tai-nodes/register", jsonBody(registerRequest{TaiID: "x"}))
c.Request.Header.Set("Content-Type", "application/json")
HandleRegister(c)
if w.Code != http.StatusUnauthorized {
t.Errorf("status = %d, want %d", w.Code, http.StatusUnauthorized)
}
}
func TestHandleRegister_MissingTaiID(t *testing.T) {
teardown := setupTest()
defer teardown()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("POST", "/tai-nodes/register", jsonBody(registerRequest{}))
c.Request.Header.Set("Authorization", "Bearer test-token")
c.Request.Header.Set("Content-Type", "application/json")
HandleRegister(c)
if w.Code != http.StatusBadRequest {
t.Errorf("status = %d, want %d", w.Code, http.StatusBadRequest)
}
}
func TestHandleHeartbeat_Success(t *testing.T) {
teardown := setupTest()
defer teardown()
reg := registry.Global()
reg.Register(&registry.TaiNode{
TaiID: "tai-abc123",
Mode: "direct",
Auth: registry.AuthInfo{ClientID: "tai-abc123"},
})
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("POST", "/tai-nodes/heartbeat",
jsonBody(heartbeatRequest{TaiID: "tai-abc123"}))
c.Request.Header.Set("Authorization", "Bearer test-token")
c.Request.Header.Set("Content-Type", "application/json")
HandleHeartbeat(c)
if w.Code != http.StatusOK {
t.Fatalf("status = %d, want %d; body = %s", w.Code, http.StatusOK, w.Body.String())
}
}
func TestHandleHeartbeat_WrongOwner(t *testing.T) {
teardown := setupTest()
defer teardown()
reg := registry.Global()
reg.Register(&registry.TaiNode{
TaiID: "tai-other",
Mode: "direct",
Auth: registry.AuthInfo{ClientID: "different-client"},
})
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("POST", "/tai-nodes/heartbeat",
jsonBody(heartbeatRequest{TaiID: "tai-other"}))
c.Request.Header.Set("Authorization", "Bearer test-token")
c.Request.Header.Set("Content-Type", "application/json")
HandleHeartbeat(c)
if w.Code != http.StatusForbidden {
t.Errorf("status = %d, want %d", w.Code, http.StatusForbidden)
}
}
func TestHandleHeartbeat_NotFound(t *testing.T) {
teardown := setupTest()
defer teardown()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("POST", "/tai-nodes/heartbeat",
jsonBody(heartbeatRequest{TaiID: "ghost"}))
c.Request.Header.Set("Authorization", "Bearer test-token")
c.Request.Header.Set("Content-Type", "application/json")
HandleHeartbeat(c)
if w.Code != http.StatusNotFound {
t.Errorf("status = %d, want %d", w.Code, http.StatusNotFound)
}
}
func TestHandleUnregister_Success(t *testing.T) {
teardown := setupTest()
defer teardown()
reg := registry.Global()
reg.Register(&registry.TaiNode{
TaiID: "tai-abc123",
Mode: "direct",
Auth: registry.AuthInfo{ClientID: "tai-abc123"},
})
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("DELETE", "/tai-nodes/register/tai-abc123", nil)
c.Request.Header.Set("Authorization", "Bearer test-token")
c.Params = gin.Params{{Key: "tai_id", Value: "tai-abc123"}}
HandleUnregister(c)
if w.Code != http.StatusOK {
t.Fatalf("status = %d, want %d; body = %s", w.Code, http.StatusOK, w.Body.String())
}
if _, ok := reg.Get("tai-abc123"); ok {
t.Error("node should be removed after unregister")
}
}
func TestHandleUnregister_WrongOwner(t *testing.T) {
teardown := setupTest()
defer teardown()
reg := registry.Global()
reg.Register(&registry.TaiNode{
TaiID: "tai-other",
Mode: "direct",
Auth: registry.AuthInfo{ClientID: "different-client"},
})
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("DELETE", "/tai-nodes/register/tai-other", nil)
c.Request.Header.Set("Authorization", "Bearer test-token")
c.Params = gin.Params{{Key: "tai_id", Value: "tai-other"}}
HandleUnregister(c)
if w.Code != http.StatusForbidden {
t.Errorf("status = %d, want %d", w.Code, http.StatusForbidden)
}
}
func TestHandleUnregister_NotFound(t *testing.T) {
teardown := setupTest()
defer teardown()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("DELETE", "/tai-nodes/register/ghost", nil)
c.Request.Header.Set("Authorization", "Bearer test-token")
c.Params = gin.Params{{Key: "tai_id", Value: "ghost"}}
HandleUnregister(c)
if w.Code != http.StatusNotFound {
t.Errorf("status = %d, want %d", w.Code, http.StatusNotFound)
}
}

87
tai/docs/README.md
View file

@ -1,6 +1,6 @@
# Tai SDK
Go client library for the [Tai](https://github.com/YaoApp/tai) runtime bridge. Provides unified access to container sandboxes, volume IO, HTTP proxy, and VNC routing — transparently working in both **Local** (direct Docker) and **Remote** (via Tai server) modes.
Go client library for the [Tai](https://github.com/YaoApp/tai) runtime bridge. Provides unified access to container sandboxes, volume IO, HTTP proxy, and VNC routing — transparently working in **Local** (direct Docker), **Remote** (via Tai server), and **Tunnel** (via Yao WebSocket tunnel) modes.
## Package Layout
@ -12,6 +12,11 @@ Go client library for the [Tai](https://github.com/YaoApp/tai) runtime bridge. P
| `workspace` | `github.com/yaoapp/yao/tai/workspace` | `fs.FS`-compatible filesystem over Volume |
| `proxy` | `github.com/yaoapp/yao/tai/proxy` | HTTP reverse proxy URL resolution |
| `vnc` | `github.com/yaoapp/yao/tai/vnc` | VNC WebSocket URL resolution |
| `registry` | `github.com/yaoapp/yao/tai/registry` | In-memory Tai node registry (direct + tunnel) |
| `api` | `github.com/yaoapp/yao/tai/api` | HTTP handlers for node registration/heartbeat |
| `tunnel` | `github.com/yaoapp/yao/tai/tunnel` | WebSocket tunnel server (control + data + proxy) |
| `hostexec/pb` | `github.com/yaoapp/yao/tai/hostexec/pb` | HostExec gRPC client (host command execution) |
| `serverinfo/pb` | `github.com/yaoapp/yao/tai/serverinfo/pb` | ServerInfo gRPC client (port/capability discovery) |
## Quick Start
@ -47,52 +52,71 @@ fmt.Println(result.Stdout) // "hello\n"
c, err := tai.New("tai://192.168.1.100", tai.K8s,
tai.WithKubeConfig("/path/to/kubeconfig.yml"),
tai.WithNamespace("default"),
tai.WithPorts(tai.Ports{K8s: 6443}),
tai.WithPorts(tai.Ports{K8s: 16443}),
)
defer c.Close()
```
### Tunnel Mode (via Yao WebSocket tunnel)
```go
// Requires a running Yao server with the Tai node registered via tunnel.
// The taiID is the node's identifier in the registry.
c, err := tai.New("tunnel://tai-abc123")
defer c.Close()
```
## Address Protocols
| Address | Mode | Description |
|---------|------|-------------|
| `""` | Local | Platform default Docker socket |
| `"local"` | Local | Platform default Docker socket |
| `"127.0.0.1"` / `"localhost"` / `"::1"` | Local | Auto-detected as local Docker |
| `unix:///var/run/docker.sock` | Local | Explicit Unix socket |
| `tcp://host:port` | Local | Explicit TCP Docker daemon |
| `npipe:////./pipe/docker_engine` | Local | Windows named pipe |
| `docker://host:port` | Local | Docker scheme |
| `tai://host` | Remote | Connect via Tai server |
| `tai://host` | Remote | Connect via Tai server (gRPC default 19100) |
| `tai://host:port` | Remote | Connect via Tai server on custom gRPC port |
| `tunnel://tai-id` | Tunnel | Connect via Yao WebSocket tunnel |
| `192.168.x.x` (non-local IP) | Remote | Auto-prepends `tai://` |
## Options
| Option | Description | Default |
|--------|-------------|---------|
| `WithPorts(Ports{...})` | Override Tai service ports | gRPC=9100, HTTP=8080, VNC=6080 |
| `WithPorts(Ports{...})` | Override Tai service ports (takes precedence over ServerInfo) | gRPC=19100, HTTP=8099, VNC=16080 |
| `WithHTTPClient(*http.Client)` | Custom HTTP client for proxy/VNC | `http.DefaultClient` |
| `WithDataDir(path)` | Volume storage root (Local mode) | `/tmp/tai-volumes` |
| `WithKubeConfig(path)` | Kubeconfig file path (K8s mode, **required**) | - |
| `WithNamespace(ns)` | K8s namespace | `"default"` |
| `WithVolume(vol)` | Inject custom Volume implementation (testing) | - |
## Default Ports
| Service | Port | Description |
|---------|------|-------------|
| gRPC | 9100 | Volume IO + Gateway |
| HTTP | 8080 | HTTP reverse proxy |
| VNC | 6080 | VNC WebSocket router |
| Docker | 2375 | Docker API proxy |
| K8s | 6443 | Kubernetes API proxy |
| gRPC | 19100 | Volume IO + Gateway + ServerInfo + HostExec |
| HTTP | 8099 | HTTP reverse proxy |
| VNC | 16080 | VNC WebSocket router |
| Docker | 12375 | Docker API proxy |
| K8s | 16443 | Kubernetes API proxy |
Ports are auto-discovered via Tai's `ServerInfo.GetInfo` gRPC call. Values set via `WithPorts` take precedence over server-reported values.
## Client API
```go
c.Volume() // volume.Volume
c.Workspace(sessionID) // workspace.FS
c.Sandbox() // sandbox.Sandbox
c.Proxy() // proxy.Proxy
c.VNC() // vnc.VNC
c.IsLocal() // bool
c.Close() // error
c.Volume() // volume.Volume — file IO (never nil)
c.Workspace(sessionID) // workspace.FS — fs.FS over Volume
c.DataDir() // string — host-side data directory (Local mode only)
c.Sandbox() // sandbox.Sandbox — container lifecycle (nil if host-exec-only)
c.Image() // sandbox.Image — image management (nil if host-exec-only)
c.Proxy() // proxy.Proxy — HTTP reverse proxy (nil if host-exec-only)
c.VNC() // vnc.VNC — VNC WebSocket (nil if host-exec-only)
c.HostExec() // hepb.HostExecClient — host command execution (nil in local mode)
c.IsLocal() // bool — true for local mode (docker/unix/tcp/npipe/local)
c.Close() // error — releases all resources
```
## Runtime Constants
@ -102,10 +126,37 @@ tai.Docker // default — use Docker runtime via Tai
tai.K8s // use Kubernetes runtime via Tai
```
## Yao gRPC Compatibility
The `tai` package re-exports Yao gRPC helpers for backward compatibility:
```go
tai.NewTokenManagerFromEnv() // *TokenManager from env vars
tai.NewTokenManager(access, refresh, sandboxID)
tai.NewYaoClientFromEnv() // *YaoClient from env vars
tai.DialYao(addr, tm) // connect to Yao gRPC
tai.Run(ctx, client, process, args, timeout) // execute Yao process
tai.Shell(ctx, client, cmd, args, env, timeout) // execute shell command
tai.HeartbeatLoop(ctx, client, sandboxID) // periodic heartbeat (blocks)
```
New code should use `grpc/client` directly. These wrappers exist for sandbox/container code that imports `tai`.
## Capabilities
When connecting to a remote Tai server, the client calls `ServerInfo.GetInfo` to discover:
- **Ports**: actual listening ports (http, docker, vnc, k8s)
- **Capabilities**: `docker`, `k8s`, `host_exec`
If no usable capabilities are found, `New()` returns an error. Remote mode checks `docker`/`k8s`/`host_exec`; Tunnel mode checks `docker`/`host_exec` (K8s is not supported over tunnel).
## Sub-Package Documentation
- [sandbox.md](sandbox.md) — Container lifecycle management
- [sandbox.md](sandbox.md) — Container lifecycle & Image management
- [volume.md](volume.md) — File IO and sync
- [workspace.md](workspace.md) — fs.FS-compatible filesystem
- [proxy.md](proxy.md) — HTTP reverse proxy
- [vnc.md](vnc.md) — VNC WebSocket routing
- [registry.md](registry.md) — Tai node registry (direct + tunnel)
- [api.md](api.md) — HTTP registration API
- [tunnel.md](tunnel.md) — WebSocket tunnel handlers

135
tai/docs/api.md Normal file
View file

@ -0,0 +1,135 @@
# Package `api`
HTTP handlers for Tai node registration, heartbeat, and unregistration. Built on [Gin](https://github.com/gin-gonic/gin), these handlers are mounted on the Yao server to allow remote Tai instances to register themselves.
## Routes
| Method | Path | Handler | Description |
|--------|------|---------|-------------|
| `POST` | `/tai-nodes/register` | `HandleRegister` | Register a Tai node |
| `POST` | `/tai-nodes/heartbeat` | `HandleHeartbeat` | Update heartbeat timestamp |
| `DELETE` | `/tai-nodes/register/:tai_id` | `HandleUnregister` | Remove a Tai node |
All endpoints require a `Bearer` token in the `Authorization` header. Tokens are validated via the Yao OAuth service.
## Authentication
```
Authorization: Bearer <access_token>
```
The token is validated against `oauth.OAuth.AuthenticateToken()`. On success, an `AuthInfo` is extracted containing `Subject`, `UserID`, `ClientID`, `Scope`, `TeamID`, and `TenantID`. The `ClientID` is used for ownership checks on heartbeat and unregister.
## Endpoints
### POST /tai-nodes/register
Registers a new Tai node in the global registry.
**Request Body:**
```json
{
"tai_id": "tai-abc123",
"machine_id": "m-001",
"version": "1.2.0",
"addr": "192.168.1.100",
"ports": {"grpc": 19100, "http": 8099, "vnc": 16080, "docker": 12375},
"capabilities": {"docker": true, "host_exec": true},
"system": {
"os": "linux",
"arch": "amd64",
"hostname": "docker-host-01",
"num_cpu": 16,
"total_mem": 34359738368
}
}
```
| Field | Type | Required | Description |
|-------|------|----------|-------------|
| `tai_id` | string | yes | Unique identifier for this Tai instance |
| `machine_id` | string | no | Host machine identifier |
| `version` | string | no | Tai version string |
| `addr` | string | no | Reachable address of the Tai server |
| `ports` | map[string]int | no | Service ports (grpc, http, vnc, docker, k8s) |
| `capabilities` | map[string]bool | no | Supported features (docker, k8s, host_exec) |
| `system` | object | no | Host system information |
**Response (200):**
```json
{
"status": "registered",
"tai_id": "tai-abc123",
"remote_ip": "203.0.113.50"
}
```
**Errors:**
| Code | Condition |
|------|-----------|
| 400 | Missing `tai_id` or invalid JSON body |
| 401 | Missing or invalid Bearer token |
| 500 | Registry not initialized |
### POST /tai-nodes/heartbeat
Updates the `LastPing` timestamp for a registered node. The node's `ClientID` must match the token's `ClientID`.
**Request Body:**
```json
{
"tai_id": "tai-abc123"
}
```
**Response (200):**
```json
{
"status": "ok"
}
```
**Errors:**
| Code | Condition |
|------|-----------|
| 400 | Missing `tai_id` or invalid JSON body |
| 401 | Missing or invalid Bearer token |
| 403 | `tai_id` belongs to a different client |
| 404 | `tai_id` not found in registry |
| 500 | Registry not initialized |
### DELETE /tai-nodes/register/:tai_id
Removes a registered node. The node's `ClientID` must match the token's `ClientID`.
**Response (200):**
```json
{
"status": "unregistered"
}
```
**Errors:**
| Code | Condition |
|------|-----------|
| 400 | Missing `tai_id` path parameter |
| 401 | Missing or invalid Bearer token |
| 403 | `tai_id` belongs to a different client |
| 404 | `tai_id` not found in registry |
| 500 | Registry not initialized |
## Node Mode
Nodes registered via this HTTP API are marked with `Mode: "direct"`. This means the Yao server can reach the Tai instance directly over the network. For tunnel-mode nodes (registered via WebSocket), see [registry.md](registry.md).
## Health Check
The registry runs a background health checker (started via `Registry.StartHealthCheck`). Direct-mode nodes that miss heartbeats beyond the configured timeout are marked `"offline"`. Nodes that remain offline longer than the cleanup threshold are automatically unregistered.

72
tai/docs/proxy.md
View file

@ -7,6 +7,7 @@ HTTP reverse proxy URL resolution. Resolves service URLs for containers so that
```go
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
}
```
@ -15,8 +16,9 @@ type Proxy interface {
| Implementation | Constructor | Mode | URL Pattern |
|----------------|-------------|------|-------------|
| **Remote** | `NewRemote(host, port, hc)` | Via Tai HTTP proxy | `http://tai-host:8080/{containerID}:{port}/{path}` |
| **Remote** | `NewRemote(host, port, hc)` | Via Tai HTTP proxy | `http://tai-host:8099/{containerID}:{port}/{path}` |
| **Local** | `NewLocal(sb)` | Direct host port lookup | `http://127.0.0.1:{hostPort}/{path}` |
| **Tunnel** | `NewTunnel(taiID, yaoBase)` | Via Yao reverse proxy | `{yaoBase}/tai/{taiID}/proxy/{containerID}:{port}/{path}` |
## Constructors
@ -29,7 +31,7 @@ func NewRemote(host string, port int, hc *http.Client) Proxy
Creates a Proxy that routes through Tai's HTTP reverse proxy. URLs are constructed by combining the Tai server address with the container ID and port.
- `host` — Tai server hostname/IP
- `port` — Tai HTTP proxy port (default 8080)
- `port` — Tai HTTP proxy port (default 8099)
- `hc` — custom HTTP client, `nil` uses `http.DefaultClient`
### NewLocal
@ -42,6 +44,17 @@ Creates a Proxy that resolves URLs by inspecting the container's port mappings v
Returns an error if the requested port is not mapped.
### NewTunnel
```go
func NewTunnel(taiID, yaoBase string) Proxy
```
Creates a Proxy that routes through Yao's HTTP reverse proxy for tunnel-mode connections.
- `taiID` — the Tai node identifier in the registry
- `yaoBase` — the Yao server base URL (e.g. `"http://yao-server:5099"`)
## Methods
### URL
@ -53,11 +66,44 @@ URL(ctx context.Context, containerID string, port int, path string) (string, err
Resolves an HTTP URL to reach a service running on `port` inside the given container.
**Remote example:** container `abc123` port `3000` path `/api/health`
→ `http://tai-host:8080/abc123:3000/api/health`
→ `http://tai-host:8099/abc123:3000/api/health`
**Local example:** container `abc123` port `3000` mapped to host port `32768`
→ `http://127.0.0.1:32768/api/health`
### Connect
```go
Connect(ctx context.Context, containerID string, opts ConnectOptions) (*Connection, error)
```
Establishes a persistent connection to a container service. Supports WebSocket and SSE protocols.
### ConnectOptions
```go
type ConnectOptions struct {
Port int // container port
Path string // URL path (e.g. "/ws" or "/events")
Protocol string // "ws" or "sse"
}
```
### Connection
```go
type Connection struct {
Messages <-chan []byte // incoming data; closed when connection ends
Send func(data []byte) error // write data (only valid for "ws" protocol)
Close func() error // terminate the connection
}
```
| Protocol | Messages | Send | Description |
|----------|----------|------|-------------|
| `"ws"` | WebSocket messages | write to WS | Full-duplex WebSocket |
| `"sse"` | SSE `data:` lines | returns error | Read-only Server-Sent Events |
### Healthz
```go
@ -68,6 +114,7 @@ Checks the health of the proxy backend.
- **Remote**: sends `GET /healthz` to the Tai HTTP proxy server
- **Local**: always returns `nil` (no external dependency)
- **Tunnel**: always returns `nil`
## Example
@ -83,4 +130,23 @@ resp, _ := http.Get(url)
if err := c.Proxy().Healthz(ctx); err != nil {
log.Fatal("Tai HTTP proxy is down:", err)
}
// WebSocket connection to a service
conn, _ := c.Proxy().Connect(ctx, containerID, proxy.ConnectOptions{
Port: 8080, Path: "/ws", Protocol: "ws",
})
defer conn.Close()
conn.Send([]byte(`{"action":"ping"}`))
for msg := range conn.Messages {
fmt.Println(string(msg))
}
// SSE event stream
conn, _ = c.Proxy().Connect(ctx, containerID, proxy.ConnectOptions{
Port: 8080, Path: "/events", Protocol: "sse",
})
defer conn.Close()
for msg := range conn.Messages {
fmt.Println("event:", string(msg))
}
```

178
tai/docs/registry.md Normal file
View file

@ -0,0 +1,178 @@
# Package `registry`
In-memory registry for Tai nodes. Manages both **direct** (network-reachable) and **tunnel** (WebSocket-bridged) connections. Used server-side by Yao to track all connected Tai instances.
## Architecture
```
Direct Mode: Yao ── TCP ──> Tai (gRPC/HTTP/Docker/VNC)
Tunnel Mode: Yao <── WS ── Tai (control channel)
Yao <── WS ── Tai (data channels, on-demand)
```
## Types
### TaiNode
```go
type TaiNode struct {
TaiID string
MachineID string
Version string
Auth AuthInfo
System SystemInfo
Mode string // "direct" | "tunnel"
Addr string // direct: "tai-host"; tunnel: empty
YaoBase string // tunnel: Yao server base URL
Ports map[string]int // {"grpc":19100, "http":8099, ...}
Capabilities map[string]bool // {"docker":true, "host_exec":true}
ControlConn *websocket.Conn // tunnel: WS control channel
Status string // "online" | "offline" | "connecting"
ConnectedAt time.Time
LastPing time.Time
PoolName string
}
```
### NodeSnapshot
Read-only copy of `TaiNode` safe to use outside locks. Returned by `Get()` and `List()`.
```go
type NodeSnapshot struct {
TaiID, MachineID, Version string
Auth AuthInfo
System SystemInfo
Mode, Addr, YaoBase string
Ports map[string]int
Capabilities map[string]bool
Status string
ConnectedAt, LastPing time.Time
PoolName string
}
```
### AuthInfo
```go
type AuthInfo struct {
Subject string
UserID string
ClientID string
Scope string
TeamID string
TenantID string
}
```
### SystemInfo
```go
type SystemInfo struct {
OS string `json:"os"`
Arch string `json:"arch"`
Hostname string `json:"hostname"`
NumCPU int `json:"num_cpu"`
TotalMem int64 `json:"total_mem,omitempty"`
}
```
## Registry API
### Init / Global
```go
func Init(logger *slog.Logger)
func Global() *Registry
```
`Init` creates the global singleton (once). `Global` returns it (nil before Init).
### Register / Unregister
```go
func (r *Registry) Register(node *TaiNode)
func (r *Registry) Unregister(taiID string)
```
`Register` adds or replaces a node, setting `Status="online"` and recording timestamps. `Unregister` closes all tunnel listeners and the control connection, then removes the node.
### Query
```go
func (r *Registry) Get(taiID string) (*NodeSnapshot, bool)
func (r *Registry) List() []NodeSnapshot
func (r *Registry) ListByTeam(teamID string) []NodeSnapshot
```
### Heartbeat
```go
func (r *Registry) UpdatePing(taiID string)
```
### Health Check
```go
func (r *Registry) StartHealthCheck(done <-chan struct{}, interval, timeout, cleanupAfter time.Duration)
```
Runs a background goroutine that:
1. Marks direct-mode nodes as `"offline"` if `LastPing` exceeds `timeout`
2. Auto-unregisters nodes that stay offline longer than `timeout + cleanupAfter`
## Tunnel API
For tunnel-mode nodes, the registry manages on-demand TCP-over-WebSocket channels.
### RequestChannel
```go
func (r *Registry) RequestChannel(taiID string, targetPort int) (channelID string, result chan net.Conn, err error)
```
Sends an `{"type":"open", "channel_id":"...", "target_port":...}` command to the node's control WebSocket. Returns a channel that receives the `net.Conn` when Tai connects back with the data channel. Times out after 30 seconds.
### AcceptDataChannel
```go
func (r *Registry) AcceptDataChannel(channelID, taiID string, conn net.Conn) error
```
Called when Tai establishes a data WebSocket for a pending channel. Validates `taiID` ownership and delivers the connection to the waiting `RequestChannel` caller.
### WriteControlJSON
```go
func (r *Registry) WriteControlJSON(taiID string, v interface{}) error
```
Thread-safe JSON write to a node's control WebSocket.
### OpenLocalListener
```go
func (r *Registry) OpenLocalListener(taiID string, targetPort int) (net.Listener, error)
```
Creates a `127.0.0.1:0` TCP listener. Every accepted connection is automatically bridged through the tunnel to `targetPort` on the Tai node. Returns the listener so the caller can read `ln.Addr()` to get the ephemeral port.
## Connection Flow (Tunnel)
```
1. Tai → Yao: WebSocket upgrade to GET /ws/tai (Bearer auth)
2. Tai → Yao: sends {"type":"register", "tai_id":"xxx", ...} on WS
3. Yao: Register(node) with Mode="tunnel", ControlConn=ws
4. Yao → Tai: sends {"type":"registered", "tai_id":"xxx"}
5. Client → Yao: tai.New("tunnel://tai-abc123")
6. Yao: OpenLocalListener("tai-abc123", 19100) → 127.0.0.1:54321
7. Yao: grpc.Dial("passthrough:///127.0.0.1:54321") → triggers accept
8. Yao: RequestChannel("tai-abc123", 19100) → sends {"type":"open"} on control WS
9. Tai: receives "open", dials localhost:19100, connects data WS to GET /ws/tai/data/:channel_id
10. Yao: AcceptDataChannel(channelID, taiID, conn) → bridges local TCP ↔ data WS
11. gRPC traffic flows transparently through the tunnel
```
### Keep-alive
Tai sends `{"type":"ping"}` periodically on the control channel. Yao replies `{"type":"pong"}` and updates `LastPing`.

118
tai/docs/sandbox.md
View file

@ -12,17 +12,32 @@ Container lifecycle management. Provides a unified `Sandbox` interface with thre
```go
type Sandbox interface {
Create(ctx context.Context, opts CreateOptions) (id string, err error)
Create(ctx context.Context, opts CreateOptions) (string, error)
Start(ctx context.Context, id string) error
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
```go
type StreamHandle struct {
Stdin io.WriteCloser
Stdout io.Reader
Stderr io.Reader
Wait func() (int, error) // blocks until exec finishes, returns exit code
Cancel func() // aborts the exec process
}
```
`ExecStream` provides real-time I/O access to a running exec process. Unlike `Exec` which collects all output, `ExecStream` returns immediately with readers/writers for interactive use.
## Constructors
### NewLocal
@ -44,7 +59,7 @@ Pings the daemon on creation; returns an error if unreachable.
func NewDocker(addr string) (Sandbox, error)
```
Connects to Docker Engine API through Tai's Docker proxy. `addr` should be `"tcp://tai-host:2375"`.
Connects to Docker Engine API through Tai's Docker proxy. `addr` should be `"tcp://tai-host:12375"`.
### NewK8s
@ -77,8 +92,10 @@ type CreateOptions struct {
WorkingDir string // working directory
Memory int64 // memory limit in bytes, 0 = no limit
CPUs float64 // CPU limit, 0 = no limit
VNC bool // enable VNC port mapping (Local only)
VNC bool // enable VNC port mapping (Local and Docker modes)
Ports []PortMapping // port mappings (Docker only)
Labels map[string]string // container/pod labels for discovery and management
User string // container user, e.g. "1000:1000" or "sandbox"
}
```
@ -97,13 +114,14 @@ type PortMapping struct {
```go
type ContainerInfo struct {
ID string // container/pod ID
Name string // container/pod name
Image string // image name
Status string // "created", "running", "exited", "removing" (Docker)
// "Pending", "Running", "Succeeded", "Failed" (K8s)
IP string // container/pod IP address
Ports []PortMapping // mapped ports (Docker only)
ID string // container/pod ID
Name string // container/pod name
Image string // image name
Status string // "created", "running", "exited", "removing" (Docker)
// "Pending", "Running", "Succeeded", "Failed" (K8s)
IP string // container/pod IP address
Ports []PortMapping // mapped ports (Docker only)
Labels map[string]string // container/pod labels
}
```
@ -149,16 +167,88 @@ type K8sOption struct {
| Behavior | Docker (Local/Remote) | K8s |
|----------|----------------------|-----|
| `Create` returns | container ID (hash) | pod name |
| `Start` | starts a stopped container | polls until pod leaves Pending (up to 30s) |
| `Start` | starts a stopped container | polls until pod leaves Pending (up to 60s) |
| `Stop` | stops with timeout, container persists | deletes the pod with grace period |
| `Remove(force=true)` | force-removes | deletes with grace period 0 |
| `Exec` | Docker exec API | `kubectl exec` via SPDY |
| `Inspect.Ports` | populated from Docker | always empty |
| `List` | filters by `tai-sdk=true` label | filters by `managed-by=yao-tai-sdk` label |
| `List` | filters only by `opts.Labels` (no auto label) | auto-merges `managed-by=yao-tai-sdk` + `opts.Labels` |
| `Binds` | supported | not supported |
| `VNC` flag | auto port-maps 6080 on macOS/Windows | not applicable |
| `VNC` flag | auto port-maps 6080 and 5900 (all platforms) | not applicable |
## Example
## Image Interface
```go
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)
}
```
Accessed via `c.Image()` on the top-level client. Nil when the Tai server has no container runtime.
| Implementation | Constructor | Backend | Notes |
|----------------|-------------|---------|-------|
| **Docker** | `NewDockerImage(cli)` | Docker SDK | Shared by Local and Docker-via-Tai modes |
| **K8s** | `NewK8sImage()` | No-op | Image pulling is handled by kubelet |
### DockerCli Helper
```go
func DockerCli(sb Sandbox) *client.Client
```
Extracts the underlying Docker SDK client from a `Sandbox` (Local or Docker). Returns `nil` for K8s sandboxes. Used internally to construct `NewDockerImage(DockerCli(sb))`.
### Types
```go
type PullOptions struct {
Auth *RegistryAuth // nil = anonymous / public
}
type RegistryAuth struct {
Username string
Password string
Server string // e.g. "ghcr.io", "registry.example.com"
}
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
}
type ImageInfo struct {
ID string
Tags []string
Size int64
Created time.Time
}
```
### Image Example
```go
c, _ := tai.New("tai://192.168.1.100")
defer c.Close()
progress, _ := c.Image().Pull(ctx, "alpine:latest", sandbox.PullOptions{})
for p := range progress {
fmt.Printf("%s %s %d/%d\n", p.Status, p.Layer, p.Current, p.Total)
}
images, _ := c.Image().List(ctx)
for _, img := range images {
fmt.Printf("%s %v\n", img.ID[:12], img.Tags)
}
```
## Sandbox Example
```go
sb, _ := sandbox.NewLocal("")

91
tai/docs/tunnel.md Normal file
View file

@ -0,0 +1,91 @@
# Package `tunnel`
WebSocket tunnel handlers for Tai nodes that cannot be reached directly (e.g. behind NAT/firewall). Provides Gin HTTP handlers mounted on the Yao server.
## Handlers
| Method | Route | Handler | Description |
|--------|-------|---------|-------------|
| `GET` | `/ws/tai` | `HandleControl` | Control channel WebSocket |
| `GET` | `/ws/tai/data/:channel_id` | `HandleData` | Data channel WebSocket |
| `ANY` | `/tai/:taiID/proxy/*path` | `HandleProxy` | HTTP reverse proxy via tunnel |
| `GET` | `/tai/:taiID/vnc/*path` | `HandleVNC` | VNC WebSocket proxy via tunnel |
All WebSocket endpoints require `Authorization: Bearer <token>` header.
## HandleControl
Manages the long-lived control WebSocket for a Tai node.
**Flow:**
1. Authenticate Bearer token via OAuth
2. Upgrade to WebSocket
3. Read `register` message (JSON):
```json
{
"type": "register",
"tai_id": "tai-abc123",
"machine_id": "m-001",
"version": "1.2.0",
"server": "http://yao-server:5099",
"ports": {"grpc": 19100, "http": 8099, "vnc": 16080, "docker": 12375},
"capabilities": {"docker": true, "host_exec": true},
"system": {"os": "linux", "arch": "amd64", "hostname": "host-01", "num_cpu": 8}
}
```
4. Register node in global registry with `Mode="tunnel"`
5. Reply `{"type": "registered", "tai_id": "tai-abc123"}`
6. Loop reading messages:
- `{"type": "ping"}` → update heartbeat, reply `{"type": "pong"}`
7. On disconnect → unregister node
## HandleData
Handles on-demand data channel connections from Tai.
**Flow:**
1. Authenticate Bearer token
2. Extract `:channel_id` from URL
3. Upgrade to WebSocket
4. Wrap WebSocket as `net.Conn` (bidirectional byte bridge)
5. Call `registry.AcceptDataChannel(channelID, clientID, conn)`
The `channel_id` must match a pending `RequestChannel` call. The `clientID` (from token) must match the Tai node that owns the channel.
## HandleProxy
HTTP reverse proxy for tunnel-connected Tai nodes.
**Flow:**
1. Look up Tai node from `:taiID` in registry
2. Get the node's HTTP port (from `node.Ports["http"]`, default 8099)
3. Open a tunnel data channel to that port via `RequestChannel`
4. Forward the incoming HTTP request through the tunnel
5. Read the response and stream it back to the client
## HandleVNC
VNC WebSocket proxy for tunnel-connected Tai nodes.
**Flow:**
1. Look up Tai node from `:taiID` in registry
2. Get the node's VNC port (from `node.Ports["vnc"]`, default 16080)
3. Open a tunnel data channel to that port via `RequestChannel`
4. Upgrade the client connection to WebSocket
5. Bridge client WebSocket ↔ tunnel data channel (binary messages)
## Internal Types
### wsConn
`wsConn` wraps `gorilla/websocket.Conn` to implement `net.Conn` for bidirectional byte bridging. This allows tunnel data channels to be treated as standard TCP connections by the registry's bridge logic.
```go
type wsConn struct { ... }
func (c *wsConn) Read(p []byte) (int, error) // reads WS binary messages
func (c *wsConn) Write(p []byte) (int, error) // writes WS binary messages
func (c *wsConn) Close() error
// Also implements: LocalAddr, RemoteAddr, SetDeadline, SetReadDeadline, SetWriteDeadline
```

23
tai/docs/vnc.md
View file

@ -15,8 +15,9 @@ type VNC interface {
| Implementation | Constructor | Mode | URL Pattern |
|----------------|-------------|------|-------------|
| **Remote** | `NewRemote(host, port, hc)` | Via Tai VNC router | `ws://tai-host:6080/vnc/{containerID}/ws` |
| **Remote** | `NewRemote(host, port, hc)` | Via Tai VNC router | `ws://tai-host:16080/vnc/{containerID}/ws` |
| **Local** | `NewLocal(sb)` | Direct host port lookup | `ws://127.0.0.1:{hostPort}/ws` |
| **Tunnel** | `NewTunnel(taiID, yaoBase)` | Via Yao reverse proxy | `ws(s)://{yaoHost}/tai/{taiID}/vnc/{containerID}/ws` |
## Constructors
@ -29,7 +30,7 @@ func NewRemote(host string, port int, hc *http.Client) VNC
Creates a VNC that routes through Tai's VNC WebSocket router.
- `host` — Tai server hostname/IP
- `port` — Tai VNC router port (default 6080)
- `port` — Tai VNC router port (default 16080)
- `hc` — custom HTTP client for Ping, `nil` uses `http.DefaultClient`
### NewLocal
@ -42,6 +43,17 @@ Creates a VNC that resolves URLs by inspecting the container's port mappings. Lo
Returns an error if port 6080 is not mapped. On macOS and Windows (Docker Desktop), the Local sandbox automatically maps port 6080 when `CreateOptions.VNC` is `true`.
### NewTunnel
```go
func NewTunnel(taiID, yaoBase string) VNC
```
Creates a VNC that routes through Yao's HTTP reverse proxy for tunnel-mode connections.
- `taiID` — the Tai node identifier in the registry
- `yaoBase` — the Yao server base URL (e.g. `"http://yao-server:5099"`)
## Methods
### URL
@ -52,7 +64,7 @@ URL(ctx context.Context, containerID string) (string, error)
Returns a WebSocket URL for connecting to the container's VNC session.
**Remote:** `ws://tai-host:6080/vnc/abc123/ws`
**Remote:** `ws://tai-host:16080/vnc/abc123/ws`
**Local:** `ws://127.0.0.1:32769/ws`
### Ping
@ -63,8 +75,9 @@ Ping(ctx context.Context, containerID string) error
Checks if the VNC endpoint is reachable by making an HTTP GET request to the WebSocket URL. Useful for verifying that the VNC server inside the container is ready before connecting a client.
- **Remote**: sends GET to `http://tai-host:6080/vnc/{containerID}/ws`
- **Remote**: sends GET to `http://tai-host:16080/vnc/{containerID}/ws`
- **Local**: resolves the host port via Inspect, then sends GET
- **Tunnel**: always returns `nil` (no direct network path to probe)
## Example
@ -90,5 +103,5 @@ for i := 0; i < 10; i++ {
// Get the WebSocket URL for a noVNC client
url, _ := c.VNC().URL(ctx, id)
fmt.Println(url) // ws://192.168.1.100:6080/vnc/desktop/ws
fmt.Println(url) // ws://192.168.1.100:16080/vnc/desktop/ws
```

8
tai/docs/volume.md
View file

@ -5,7 +5,7 @@ File IO and directory synchronization. Provides a `Volume` interface with two im
| Implementation | Constructor | Backend | Mode |
|----------------|-------------|---------|------|
| **Local** | `NewLocal(root)` | Direct filesystem | Local |
| **Remote** | `NewRemote(conn)` | gRPC to Tai :9100 | Remote |
| **Remote** | `NewRemote(conn)` | gRPC to Tai :19100 | Remote |
## Interface
@ -33,10 +33,10 @@ All paths are **relative** to the session's workspace root. The `sessionID` iden
### NewLocal
```go
func NewLocal(root string) Volume
func NewLocal(dataDir string) Volume
```
Creates a Volume backed by the local filesystem. Files are stored under `<root>/<sessionID>/`.
Creates a Volume backed by the local filesystem. Files are stored under `<dataDir>/<sessionID>/`.
### NewRemote
@ -44,7 +44,7 @@ Creates a Volume backed by the local filesystem. Files are stored under `<root>/
func NewRemote(conn *grpc.ClientConn) Volume
```
Creates a Volume backed by Tai's gRPC Volume service. The connection should target Tai's gRPC port (default 9100). Uses lz4 compression for `SyncPush`/`SyncPull` bulk transfers.
Creates a Volume backed by Tai's gRPC Volume service. The connection should target Tai's gRPC port (default 19100). Uses lz4 compression for `SyncPush`/`SyncPull` bulk transfers.
## Types

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

@ -1,267 +0,0 @@
package main
import (
"bufio"
"context"
"encoding/json"
"fmt"
"io"
"os"
"os/signal"
"syscall"
yaogrpc "github.com/yaoapp/yao/tai/grpc"
)
// Build-time variables set via -ldflags.
var (
Version = "dev"
Commit = "none"
BuildTime = "unknown"
)
func main() {
if len(os.Args) < 2 {
fmt.Fprintln(os.Stderr, "Usage: yao-grpc <version|serve>")
os.Exit(1)
}
switch os.Args[1] {
case "version":
fmt.Printf("yao-grpc %s (commit: %s, built: %s)\n", Version, Commit, BuildTime)
case "serve":
if err := serve(); err != nil {
fmt.Fprintf(os.Stderr, "error: %v\n", err)
os.Exit(1)
}
default:
fmt.Fprintf(os.Stderr, "Unknown command: %s\nUsage: yao-grpc <version|serve>\n", os.Args[1])
os.Exit(1)
}
}
// jsonrpcRequest is a minimal JSON-RPC 2.0 request.
type jsonrpcRequest struct {
JSONRPC string `json:"jsonrpc"`
ID json.RawMessage `json:"id,omitempty"`
Method string `json:"method"`
Params json.RawMessage `json:"params,omitempty"`
}
// jsonrpcResponse is a minimal JSON-RPC 2.0 response.
type jsonrpcResponse struct {
JSONRPC string `json:"jsonrpc"`
ID json.RawMessage `json:"id,omitempty"`
Result json.RawMessage `json:"result,omitempty"`
Error *jsonrpcError `json:"error,omitempty"`
}
type jsonrpcError struct {
Code int `json:"code"`
Message string `json:"message"`
}
func serve() error {
client, err := yaogrpc.NewFromEnv()
if err != nil {
return err
}
defer client.Close()
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)
for scanner.Scan() {
select {
case <-ctx.Done():
return nil
default:
}
line := scanner.Bytes()
if len(line) == 0 {
continue
}
var req jsonrpcRequest
if err := json.Unmarshal(line, &req); err != nil {
encoder.Encode(jsonrpcResponse{
JSONRPC: "2.0",
Error: &jsonrpcError{Code: -32700, Message: "parse error"},
})
continue
}
resp := dispatch(ctx, client, &req)
encoder.Encode(resp)
}
if err := scanner.Err(); err != nil && err != io.EOF {
return fmt.Errorf("stdin read: %w", err)
}
return nil
}
func dispatch(ctx context.Context, client *yaogrpc.Client, req *jsonrpcRequest) jsonrpcResponse {
base := jsonrpcResponse{JSONRPC: "2.0", ID: req.ID}
switch req.Method {
case "run":
return handleRun(ctx, client, req, base)
case "shell":
return handleShell(ctx, client, req, base)
case "mcp/list_tools":
return handleMCPListTools(ctx, client, req, base)
case "mcp/call_tool":
return handleMCPCallTool(ctx, client, req, base)
case "mcp/list_resources":
return handleMCPListResources(ctx, client, req, base)
case "mcp/read_resource":
return handleMCPReadResource(ctx, client, req, base)
case "healthz":
return handleHealthz(ctx, client, base)
default:
base.Error = &jsonrpcError{Code: -32601, Message: "method not found: " + req.Method}
return base
}
}
// --- handlers ---
type runParams struct {
Process string `json:"process"`
Args json.RawMessage `json:"args,omitempty"`
Timeout int32 `json:"timeout,omitempty"`
}
func handleRun(ctx context.Context, c *yaogrpc.Client, req *jsonrpcRequest, base jsonrpcResponse) jsonrpcResponse {
var p runParams
if err := json.Unmarshal(req.Params, &p); err != nil {
base.Error = &jsonrpcError{Code: -32602, Message: "invalid params: " + err.Error()}
return base
}
data, err := c.Run(ctx, p.Process, p.Args, p.Timeout)
if err != nil {
base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
return base
}
base.Result = data
return base
}
type shellParams struct {
Command string `json:"command"`
Args []string `json:"args,omitempty"`
Env map[string]string `json:"env,omitempty"`
Timeout int32 `json:"timeout,omitempty"`
}
func handleShell(ctx context.Context, c *yaogrpc.Client, req *jsonrpcRequest, base jsonrpcResponse) jsonrpcResponse {
var p shellParams
if err := json.Unmarshal(req.Params, &p); err != nil {
base.Error = &jsonrpcError{Code: -32602, Message: "invalid params: " + err.Error()}
return base
}
resp, err := c.Shell(ctx, p.Command, p.Args, p.Env, p.Timeout)
if err != nil {
base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
return base
}
data, _ := json.Marshal(resp)
base.Result = data
return base
}
type mcpSessionParams struct {
SessionID string `json:"session_id"`
}
func handleMCPListTools(ctx context.Context, c *yaogrpc.Client, req *jsonrpcRequest, base jsonrpcResponse) jsonrpcResponse {
var p mcpSessionParams
if err := json.Unmarshal(req.Params, &p); err != nil {
base.Error = &jsonrpcError{Code: -32602, Message: "invalid params: " + err.Error()}
return base
}
data, err := c.MCPListTools(ctx, p.SessionID)
if err != nil {
base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
return base
}
base.Result = data
return base
}
type mcpCallParams struct {
SessionID string `json:"session_id"`
Tool string `json:"tool"`
Arguments json.RawMessage `json:"arguments,omitempty"`
}
func handleMCPCallTool(ctx context.Context, c *yaogrpc.Client, req *jsonrpcRequest, base jsonrpcResponse) jsonrpcResponse {
var p mcpCallParams
if err := json.Unmarshal(req.Params, &p); err != nil {
base.Error = &jsonrpcError{Code: -32602, Message: "invalid params: " + err.Error()}
return base
}
data, err := c.MCPCallTool(ctx, p.SessionID, p.Tool, p.Arguments)
if err != nil {
base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
return base
}
base.Result = data
return base
}
func handleMCPListResources(ctx context.Context, c *yaogrpc.Client, req *jsonrpcRequest, base jsonrpcResponse) jsonrpcResponse {
var p mcpSessionParams
if err := json.Unmarshal(req.Params, &p); err != nil {
base.Error = &jsonrpcError{Code: -32602, Message: "invalid params: " + err.Error()}
return base
}
data, err := c.MCPListResources(ctx, p.SessionID)
if err != nil {
base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
return base
}
base.Result = data
return base
}
type mcpReadParams struct {
SessionID string `json:"session_id"`
URI string `json:"uri"`
}
func handleMCPReadResource(ctx context.Context, c *yaogrpc.Client, req *jsonrpcRequest, base jsonrpcResponse) jsonrpcResponse {
var p mcpReadParams
if err := json.Unmarshal(req.Params, &p); err != nil {
base.Error = &jsonrpcError{Code: -32602, Message: "invalid params: " + err.Error()}
return base
}
data, err := c.MCPReadResource(ctx, p.SessionID, p.URI)
if err != nil {
base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
return base
}
base.Result = data
return base
}
func handleHealthz(ctx context.Context, c *yaogrpc.Client, base jsonrpcResponse) jsonrpcResponse {
status, err := c.Healthz(ctx)
if err != nil {
base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
return base
}
data, _ := json.Marshal(map[string]string{"status": status})
base.Result = data
return base
}

210
tai/grpc/grpc_test.go
View file

@ -1,210 +0,0 @@
package grpc_test
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"google.golang.org/grpc/metadata"
yaogrpc "github.com/yaoapp/yao/tai/grpc"
)
// ── TokenManager unit tests ──────────────────────────────────────────────────
func TestTokenManager_AttachMetadata_WithAllFields(t *testing.T) {
tm := yaogrpc.NewTokenManager("tok", "ref", "sb-1", "yao:9099")
ctx := tm.AttachMetadata(context.Background())
md, ok := metadata.FromOutgoingContext(ctx)
require.True(t, ok)
assert.Equal(t, []string{"Bearer tok"}, md.Get("authorization"))
assert.Equal(t, []string{"ref"}, md.Get("x-refresh-token"))
assert.Equal(t, []string{"sb-1"}, md.Get("x-sandbox-id"))
assert.Equal(t, []string{"yao:9099"}, md.Get("x-grpc-upstream"))
}
func TestTokenManager_AttachMetadata_DirectMode(t *testing.T) {
tm := yaogrpc.NewTokenManager("tok", "ref", "sb-1", "")
ctx := tm.AttachMetadata(context.Background())
md, ok := metadata.FromOutgoingContext(ctx)
require.True(t, ok)
assert.Equal(t, []string{"Bearer tok"}, md.Get("authorization"))
assert.Empty(t, md.Get("x-grpc-upstream"), "direct mode should not set x-grpc-upstream")
}
func TestTokenManager_AttachMetadata_EmptyTokens(t *testing.T) {
tm := yaogrpc.NewTokenManager("", "", "", "")
ctx := tm.AttachMetadata(context.Background())
_, ok := metadata.FromOutgoingContext(ctx)
assert.False(t, ok, "empty tokens should not produce metadata")
}
func TestTokenManager_HandleResponseHeaders(t *testing.T) {
tm := yaogrpc.NewTokenManager("old-tok", "old-ref", "", "")
tm.HandleResponseHeaders(metadata.New(map[string]string{
"x-access-token": "new-tok",
"x-refresh-token": "new-ref",
}))
assert.Equal(t, "new-tok", tm.AccessToken())
ctx := tm.AttachMetadata(context.Background())
md, _ := metadata.FromOutgoingContext(ctx)
assert.Equal(t, []string{"Bearer new-tok"}, md.Get("authorization"))
assert.Equal(t, []string{"new-ref"}, md.Get("x-refresh-token"))
}
func TestTokenManager_HandleResponseHeaders_Nil(t *testing.T) {
tm := yaogrpc.NewTokenManager("tok", "", "", "")
tm.HandleResponseHeaders(nil)
assert.Equal(t, "tok", tm.AccessToken())
}
func TestTokenManager_HandleResponseHeaders_EmptyValues(t *testing.T) {
tm := yaogrpc.NewTokenManager("tok", "ref", "", "")
tm.HandleResponseHeaders(metadata.New(map[string]string{
"x-access-token": "",
}))
assert.Equal(t, "tok", tm.AccessToken(), "empty header should not overwrite")
}
func TestTokenManager_IsTaiMode(t *testing.T) {
tmDirect := yaogrpc.NewTokenManager("tok", "", "", "")
assert.False(t, tmDirect.IsTaiMode())
tmTai := yaogrpc.NewTokenManager("tok", "", "", "tai:9100")
assert.True(t, tmTai.IsTaiMode())
}
func TestTokenManager_NewFromEnv_MissingUpstream(t *testing.T) {
t.Setenv("YAO_GRPC_TAI", "enable")
t.Setenv("YAO_GRPC_UPSTREAM", "")
t.Setenv("YAO_TOKEN", "tok")
_, err := yaogrpc.NewTokenManagerFromEnv()
assert.Error(t, err)
assert.Contains(t, err.Error(), "YAO_GRPC_UPSTREAM")
}
func TestTokenManager_NewFromEnv_TaiEnabled(t *testing.T) {
t.Setenv("YAO_GRPC_TAI", "enable")
t.Setenv("YAO_GRPC_UPSTREAM", "yao:9099")
t.Setenv("YAO_TOKEN", "my-token")
t.Setenv("YAO_REFRESH_TOKEN", "my-refresh")
t.Setenv("YAO_SANDBOX_ID", "sb-42")
tm, err := yaogrpc.NewTokenManagerFromEnv()
require.NoError(t, err)
assert.True(t, tm.IsTaiMode())
assert.Equal(t, "my-token", tm.AccessToken())
}
func TestTokenManager_NewFromEnv_DirectMode(t *testing.T) {
t.Setenv("YAO_GRPC_TAI", "")
t.Setenv("YAO_GRPC_UPSTREAM", "")
t.Setenv("YAO_TOKEN", "tok")
tm, err := yaogrpc.NewTokenManagerFromEnv()
require.NoError(t, err)
assert.False(t, tm.IsTaiMode())
}
// ── Dial tests ───────────────────────────────────────────────────────────────
func TestNewFromEnv_MissingAddr(t *testing.T) {
t.Setenv("YAO_GRPC_ADDR", "")
_, err := yaogrpc.NewFromEnv()
assert.Error(t, err)
assert.Contains(t, err.Error(), "YAO_GRPC_ADDR")
}
func TestNewFromEnv_Success(t *testing.T) {
t.Setenv("YAO_GRPC_ADDR", "127.0.0.1:9099")
t.Setenv("YAO_TOKEN", "test-token")
t.Setenv("YAO_REFRESH_TOKEN", "test-refresh")
t.Setenv("YAO_SANDBOX_ID", "sb-1")
t.Setenv("YAO_GRPC_TAI", "")
c, err := yaogrpc.NewFromEnv()
require.NoError(t, err)
defer c.Close()
assert.NotNil(t, c.Conn())
assert.Equal(t, "test-token", c.TokenManager().AccessToken())
}
func TestNewFromEnv_TaiMode_MissingUpstream(t *testing.T) {
t.Setenv("YAO_GRPC_ADDR", "tai:9100")
t.Setenv("YAO_GRPC_TAI", "enable")
t.Setenv("YAO_GRPC_UPSTREAM", "")
_, err := yaogrpc.NewFromEnv()
assert.Error(t, err)
assert.Contains(t, err.Error(), "YAO_GRPC_UPSTREAM")
}
func TestDial_WithNilTokenManager(t *testing.T) {
c, err := yaogrpc.Dial("127.0.0.1:0", nil)
require.NoError(t, err)
defer c.Close()
assert.NotNil(t, c.Conn())
assert.Nil(t, c.TokenManager())
}
func TestDial_WithTokenManager(t *testing.T) {
tm := yaogrpc.NewTokenManager("tok", "", "", "")
c, err := yaogrpc.Dial("127.0.0.1:0", tm)
require.NoError(t, err)
defer c.Close()
assert.NotNil(t, c.TokenManager())
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())
}

194
tai/grpc/heartbeat_test.go
View file

@ -1,194 +0,0 @@
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")
}
}

420
tai/grpc/integration_test.go
View file

@ -1,420 +0,0 @@
package grpc_test
import (
"context"
"encoding/json"
"os"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/grpc/tests/testutils"
yaogrpc "github.com/yaoapp/yao/tai/grpc"
)
// Integration tests that start a real Yao gRPC server and test the tai/grpc
// client through the full interceptor -> handler chain.
func setupClient(t *testing.T, scopes ...string) *yaogrpc.Client {
t.Helper()
conn := testutils.Prepare(t)
t.Cleanup(func() {
conn.Close()
testutils.Clean()
})
addr := testutils.Addr()
token := testutils.ObtainAccessToken(t, scopes...)
refreshToken := testutils.ObtainRefreshToken(t, scopes...)
tm := yaogrpc.NewTokenManager(token, refreshToken, "test-sandbox", "")
client, err := yaogrpc.Dial(addr, tm)
require.NoError(t, err)
t.Cleanup(func() { client.Close() })
return client
}
// ── Healthz ──────────────────────────────────────────────────────────────────
func TestIntegration_Healthz(t *testing.T) {
conn := testutils.Prepare(t)
defer func() {
conn.Close()
testutils.Clean()
}()
addr := testutils.Addr()
tm := yaogrpc.NewTokenManager("", "", "", "")
client, err := yaogrpc.Dial(addr, tm)
require.NoError(t, err)
defer client.Close()
status, err := client.Healthz(context.Background())
assert.NoError(t, err)
assert.Equal(t, "ok", status)
}
// ── Run ──────────────────────────────────────────────────────────────────────
func TestIntegration_Run_Ping(t *testing.T) {
client := setupClient(t, "grpc:run")
data, err := client.Run(context.Background(), "utils.app.Ping", nil, 0)
assert.NoError(t, err)
assert.NotNil(t, data)
}
func TestIntegration_Run_InvalidProcess(t *testing.T) {
client := setupClient(t, "grpc:run")
_, err := client.Run(context.Background(), "nonexistent.process", nil, 0)
assert.Error(t, err)
}
func TestIntegration_Run_WithArgs(t *testing.T) {
client := setupClient(t, "grpc:run")
args, _ := json.Marshal([]any{"hello", "world"})
data, err := client.Run(context.Background(), "utils.app.Ping", args, 5)
assert.NoError(t, err)
assert.NotNil(t, data)
}
// ── Shell ────────────────────────────────────────────────────────────────────
func TestIntegration_Shell_Echo(t *testing.T) {
client := setupClient(t, "grpc:shell")
resp, err := client.Shell(context.Background(), "echo", []string{"hello"}, nil, 5)
require.NoError(t, err)
require.NotNil(t, resp)
assert.Equal(t, int32(0), resp.ExitCode)
assert.Contains(t, string(resp.Stdout), "hello")
}
func TestIntegration_Shell_NotFound(t *testing.T) {
client := setupClient(t, "grpc:shell")
_, err := client.Shell(context.Background(), "nonexistent-command-xyz", nil, nil, 5)
assert.Error(t, err)
}
// ── MCP ──────────────────────────────────────────────────────────────────────
func TestIntegration_MCPListTools(t *testing.T) {
client := setupClient(t, "grpc:mcp")
data, err := client.MCPListTools(context.Background(), "echo")
assert.NoError(t, err)
assert.NotNil(t, data)
var tools []any
assert.NoError(t, json.Unmarshal(data, &tools))
assert.Greater(t, len(tools), 0)
}
func TestIntegration_MCPCallTool(t *testing.T) {
client := setupClient(t, "grpc:mcp")
args, _ := json.Marshal(map[string]string{"message": "hi"})
data, err := client.MCPCallTool(context.Background(), "echo", "ping", args)
assert.NoError(t, err)
assert.NotNil(t, data)
}
func TestIntegration_MCPListResources(t *testing.T) {
client := setupClient(t, "grpc:mcp")
data, err := client.MCPListResources(context.Background(), "echo")
assert.NoError(t, err)
assert.NotNil(t, data)
}
func TestIntegration_MCPReadResource(t *testing.T) {
client := setupClient(t, "grpc:mcp")
data, err := client.MCPReadResource(context.Background(), "echo", "echo://info")
assert.NoError(t, err)
assert.NotNil(t, data)
}
// ── API ──────────────────────────────────────────────────────────────────────
func TestIntegration_API_Proxy(t *testing.T) {
client := setupClient(t, "grpc:run", "grpc:mcp")
resp, err := client.API(context.Background(), "GET", "/api/__yao/app/setting", nil, nil)
require.NoError(t, err)
require.NotNil(t, resp)
t.Logf("API proxy status: %d", resp.Status)
}
// ── LLM ──────────────────────────────────────────────────────────────────────
func TestIntegration_ChatCompletions_InvalidConnector(t *testing.T) {
client := setupClient(t, "grpc:llm")
messages, _ := json.Marshal([]map[string]string{
{"role": "user", "content": "test"},
})
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
_, err := client.ChatCompletions(ctx, "nonexistent-connector", messages, nil)
assert.Error(t, err)
}
func TestIntegration_ChatCompletionsStream_InvalidConnector(t *testing.T) {
client := setupClient(t, "grpc:llm")
messages, _ := json.Marshal([]map[string]string{
{"role": "user", "content": "test"},
})
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
err := client.ChatCompletionsStream(ctx, "nonexistent-connector", messages, nil,
func(data []byte, done bool) error { return nil })
assert.Error(t, err)
}
func TestIntegration_ChatCompletions_EmptyMessages(t *testing.T) {
client := setupClient(t, "grpc:llm")
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
_, err := client.ChatCompletions(ctx, "default", nil, nil)
assert.Error(t, err)
}
// ── Agent ────────────────────────────────────────────────────────────────────
func TestIntegration_AgentStream_InvalidRobot(t *testing.T) {
client := setupClient(t, "grpc:agent")
messages, _ := json.Marshal([]map[string]string{
{"role": "user", "content": "hello"},
})
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
err := client.AgentStream(ctx, "nonexistent-robot-xyz", messages, nil,
func(data []byte, done bool) error { return nil })
assert.Error(t, err)
}
// ── Unauthenticated ──────────────────────────────────────────────────────────
func TestIntegration_Run_NoToken(t *testing.T) {
conn := testutils.Prepare(t)
defer func() {
conn.Close()
testutils.Clean()
}()
addr := testutils.Addr()
tm := yaogrpc.NewTokenManager("", "", "", "")
client, err := yaogrpc.Dial(addr, tm)
require.NoError(t, err)
defer client.Close()
_, err = client.Run(context.Background(), "utils.app.Ping", nil, 0)
assert.Error(t, err)
assert.Contains(t, err.Error(), "Unauthenticated")
}
// ── Token Refresh via interceptor ────────────────────────────────────────────
func TestIntegration_TokenRefresh(t *testing.T) {
conn := testutils.Prepare(t)
defer func() {
conn.Close()
testutils.Clean()
}()
addr := testutils.Addr()
scopes := []string{"grpc:run"}
expiredToken := testutils.ObtainExpiredAccessToken(t, scopes...)
refreshToken := testutils.ObtainRefreshToken(t, scopes...)
tm := yaogrpc.NewTokenManager(expiredToken, refreshToken, "sb-test", "")
client, err := yaogrpc.Dial(addr, tm)
require.NoError(t, err)
defer client.Close()
data, err := client.Run(context.Background(), "utils.app.Ping", nil, 0)
assert.NoError(t, err)
assert.NotNil(t, data)
newToken := tm.AccessToken()
if newToken != expiredToken {
t.Logf("token was refreshed: old=%s... new=%s...", expiredToken[:20], newToken[:20])
}
}
// ══════════════════════════════════════════════════════════════════════════════
// Relay mode tests — client → Tai (:9100) → x-grpc-upstream → Yao gRPC
// Requires TAI_TEST_GRPC env var (e.g. 127.0.0.1:9100) and a running Tai server.
// ══════════════════════════════════════════════════════════════════════════════
func setupRelayClient(t *testing.T, scopes ...string) *yaogrpc.Client {
t.Helper()
taiAddr := os.Getenv("TAI_TEST_GRPC")
if taiAddr == "" {
t.Skip("TAI_TEST_GRPC not set, skipping relay mode test")
}
conn := testutils.Prepare(t)
t.Cleanup(func() {
conn.Close()
testutils.Clean()
})
yaoAddr := testutils.RelayAddr()
token := testutils.ObtainAccessToken(t, scopes...)
refreshToken := testutils.ObtainRefreshToken(t, scopes...)
// upstream = Yao gRPC address reachable from the Tai container
tm := yaogrpc.NewTokenManager(token, refreshToken, "relay-sandbox", yaoAddr)
client, err := yaogrpc.Dial(taiAddr, tm)
require.NoError(t, err)
t.Cleanup(func() { client.Close() })
return client
}
func TestRelay_Healthz(t *testing.T) {
taiAddr := os.Getenv("TAI_TEST_GRPC")
if taiAddr == "" {
t.Skip("TAI_TEST_GRPC not set")
}
conn := testutils.Prepare(t)
defer func() {
conn.Close()
testutils.Clean()
}()
yaoAddr := testutils.RelayAddr()
tm := yaogrpc.NewTokenManager("", "", "", yaoAddr)
client, err := yaogrpc.Dial(taiAddr, tm)
require.NoError(t, err)
defer client.Close()
status, err := client.Healthz(context.Background())
require.NoError(t, err)
assert.Equal(t, "ok", status)
}
func TestRelay_Run_Ping(t *testing.T) {
client := setupRelayClient(t, "grpc:run")
data, err := client.Run(context.Background(), "utils.app.Ping", nil, 0)
require.NoError(t, err)
require.NotNil(t, data)
t.Logf("relay Run result: %s", string(data))
}
func TestRelay_Run_InvalidProcess(t *testing.T) {
client := setupRelayClient(t, "grpc:run")
_, err := client.Run(context.Background(), "nonexistent.process", nil, 0)
assert.Error(t, err)
}
func TestRelay_Shell_Echo(t *testing.T) {
client := setupRelayClient(t, "grpc:shell")
resp, err := client.Shell(context.Background(), "echo", []string{"relay-test"}, nil, 5)
require.NoError(t, err)
require.NotNil(t, resp)
assert.Equal(t, int32(0), resp.ExitCode)
assert.Contains(t, string(resp.Stdout), "relay-test")
}
func TestRelay_MCPListTools(t *testing.T) {
client := setupRelayClient(t, "grpc:mcp")
data, err := client.MCPListTools(context.Background(), "echo")
assert.NoError(t, err)
assert.NotNil(t, data)
var tools []any
assert.NoError(t, json.Unmarshal(data, &tools))
assert.Greater(t, len(tools), 0)
}
func TestRelay_MCPCallTool(t *testing.T) {
client := setupRelayClient(t, "grpc:mcp")
args, _ := json.Marshal(map[string]string{"message": "relay"})
data, err := client.MCPCallTool(context.Background(), "echo", "ping", args)
assert.NoError(t, err)
assert.NotNil(t, data)
}
func TestRelay_Run_NoToken(t *testing.T) {
taiAddr := os.Getenv("TAI_TEST_GRPC")
if taiAddr == "" {
t.Skip("TAI_TEST_GRPC not set")
}
conn := testutils.Prepare(t)
defer func() {
conn.Close()
testutils.Clean()
}()
yaoAddr := testutils.RelayAddr()
tm := yaogrpc.NewTokenManager("", "", "", yaoAddr)
client, err := yaogrpc.Dial(taiAddr, tm)
require.NoError(t, err)
defer client.Close()
_, err = client.Run(context.Background(), "utils.app.Ping", nil, 0)
assert.Error(t, err)
assert.Contains(t, err.Error(), "Unauthenticated")
}
func TestRelay_TokenRefresh(t *testing.T) {
taiAddr := os.Getenv("TAI_TEST_GRPC")
if taiAddr == "" {
t.Skip("TAI_TEST_GRPC not set")
}
conn := testutils.Prepare(t)
defer func() {
conn.Close()
testutils.Clean()
}()
yaoAddr := testutils.RelayAddr()
scopes := []string{"grpc:run"}
expiredToken := testutils.ObtainExpiredAccessToken(t, scopes...)
refreshToken := testutils.ObtainRefreshToken(t, scopes...)
tm := yaogrpc.NewTokenManager(expiredToken, refreshToken, "relay-sb", yaoAddr)
client, err := yaogrpc.Dial(taiAddr, tm)
require.NoError(t, err)
defer client.Close()
data, err := client.Run(context.Background(), "utils.app.Ping", nil, 0)
assert.NoError(t, err)
assert.NotNil(t, data)
newToken := tm.AccessToken()
if newToken != expiredToken {
t.Logf("relay token refreshed: old=%s... new=%s...", expiredToken[:20], newToken[:20])
}
}

13
tai/grpc/heartbeat.go → tai/heartbeat.go
View file

@ -1,4 +1,4 @@
package grpc
package tai
import (
"context"
@ -14,8 +14,8 @@ import (
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) {
// It runs until ctx is cancelled.
func HeartbeatLoop(ctx context.Context, client *YaoClient, sandboxID string) {
interval := defaultHeartbeatInterval
if s := os.Getenv("YAO_HEARTBEAT_INTERVAL"); s != "" {
if d, err := time.ParseDuration(s); err == nil && d > 0 {
@ -38,7 +38,7 @@ func HeartbeatLoop(ctx context.Context, client *Client, sandboxID string) {
continue
}
if action == "shutdown" {
fmt.Fprintf(os.Stderr, "yao-grpc: received shutdown signal\n")
fmt.Fprintf(os.Stderr, "tai: received shutdown signal\n")
p, _ := os.FindProcess(os.Getpid())
p.Signal(os.Interrupt)
return
@ -47,12 +47,10 @@ func HeartbeatLoop(ctx context.Context, client *Client, sandboxID string) {
}
}
// 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
@ -61,17 +59,14 @@ func countUserProcesses() int32 {
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
}

424
tai/hostexec/pb/hostexec.pb.go Normal file
View file

@ -0,0 +1,424 @@
// Code generated by protoc-gen-go. DO NOT EDIT.
// versions:
// protoc-gen-go v1.36.11
// protoc v4.25.0
// source: hostexec/pb/hostexec.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 ExecOutput_Stream int32
const (
ExecOutput_STDOUT ExecOutput_Stream = 0
ExecOutput_STDERR ExecOutput_Stream = 1
)
// Enum value maps for ExecOutput_Stream.
var (
ExecOutput_Stream_name = map[int32]string{
0: "STDOUT",
1: "STDERR",
}
ExecOutput_Stream_value = map[string]int32{
"STDOUT": 0,
"STDERR": 1,
}
)
func (x ExecOutput_Stream) Enum() *ExecOutput_Stream {
p := new(ExecOutput_Stream)
*p = x
return p
}
func (x ExecOutput_Stream) String() string {
return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x))
}
func (ExecOutput_Stream) Descriptor() protoreflect.EnumDescriptor {
return file_hostexec_pb_hostexec_proto_enumTypes[0].Descriptor()
}
func (ExecOutput_Stream) Type() protoreflect.EnumType {
return &file_hostexec_pb_hostexec_proto_enumTypes[0]
}
func (x ExecOutput_Stream) Number() protoreflect.EnumNumber {
return protoreflect.EnumNumber(x)
}
// Deprecated: Use ExecOutput_Stream.Descriptor instead.
func (ExecOutput_Stream) EnumDescriptor() ([]byte, []int) {
return file_hostexec_pb_hostexec_proto_rawDescGZIP(), []int{2, 0}
}
type ExecRequest struct {
state protoimpl.MessageState `protogen:"open.v1"`
Command string `protobuf:"bytes,1,opt,name=command,proto3" json:"command,omitempty"`
Args []string `protobuf:"bytes,2,rep,name=args,proto3" json:"args,omitempty"`
WorkingDir string `protobuf:"bytes,3,opt,name=working_dir,json=workingDir,proto3" json:"working_dir,omitempty"`
Env map[string]string `protobuf:"bytes,4,rep,name=env,proto3" json:"env,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"`
Stdin []byte `protobuf:"bytes,5,opt,name=stdin,proto3" json:"stdin,omitempty"`
TimeoutMs int64 `protobuf:"varint,6,opt,name=timeout_ms,json=timeoutMs,proto3" json:"timeout_ms,omitempty"`
MaxOutputBytes int64 `protobuf:"varint,7,opt,name=max_output_bytes,json=maxOutputBytes,proto3" json:"max_output_bytes,omitempty"` // max stdout+stderr size (0 = default 10MB), truncate if exceeded
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *ExecRequest) Reset() {
*x = ExecRequest{}
mi := &file_hostexec_pb_hostexec_proto_msgTypes[0]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *ExecRequest) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*ExecRequest) ProtoMessage() {}
func (x *ExecRequest) ProtoReflect() protoreflect.Message {
mi := &file_hostexec_pb_hostexec_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 ExecRequest.ProtoReflect.Descriptor instead.
func (*ExecRequest) Descriptor() ([]byte, []int) {
return file_hostexec_pb_hostexec_proto_rawDescGZIP(), []int{0}
}
func (x *ExecRequest) GetCommand() string {
if x != nil {
return x.Command
}
return ""
}
func (x *ExecRequest) GetArgs() []string {
if x != nil {
return x.Args
}
return nil
}
func (x *ExecRequest) GetWorkingDir() string {
if x != nil {
return x.WorkingDir
}
return ""
}
func (x *ExecRequest) GetEnv() map[string]string {
if x != nil {
return x.Env
}
return nil
}
func (x *ExecRequest) GetStdin() []byte {
if x != nil {
return x.Stdin
}
return nil
}
func (x *ExecRequest) GetTimeoutMs() int64 {
if x != nil {
return x.TimeoutMs
}
return 0
}
func (x *ExecRequest) GetMaxOutputBytes() int64 {
if x != nil {
return x.MaxOutputBytes
}
return 0
}
type ExecResponse struct {
state protoimpl.MessageState `protogen:"open.v1"`
ExitCode int32 `protobuf:"varint,1,opt,name=exit_code,json=exitCode,proto3" json:"exit_code,omitempty"`
Stdout []byte `protobuf:"bytes,2,opt,name=stdout,proto3" json:"stdout,omitempty"`
Stderr []byte `protobuf:"bytes,3,opt,name=stderr,proto3" json:"stderr,omitempty"`
DurationMs int64 `protobuf:"varint,4,opt,name=duration_ms,json=durationMs,proto3" json:"duration_ms,omitempty"`
Error string `protobuf:"bytes,5,opt,name=error,proto3" json:"error,omitempty"` // non-empty if Tai failed to execute (not the command's error)
Truncated bool `protobuf:"varint,6,opt,name=truncated,proto3" json:"truncated,omitempty"` // true if stdout+stderr exceeded max_output_bytes and was truncated
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *ExecResponse) Reset() {
*x = ExecResponse{}
mi := &file_hostexec_pb_hostexec_proto_msgTypes[1]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *ExecResponse) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*ExecResponse) ProtoMessage() {}
func (x *ExecResponse) ProtoReflect() protoreflect.Message {
mi := &file_hostexec_pb_hostexec_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 ExecResponse.ProtoReflect.Descriptor instead.
func (*ExecResponse) Descriptor() ([]byte, []int) {
return file_hostexec_pb_hostexec_proto_rawDescGZIP(), []int{1}
}
func (x *ExecResponse) GetExitCode() int32 {
if x != nil {
return x.ExitCode
}
return 0
}
func (x *ExecResponse) GetStdout() []byte {
if x != nil {
return x.Stdout
}
return nil
}
func (x *ExecResponse) GetStderr() []byte {
if x != nil {
return x.Stderr
}
return nil
}
func (x *ExecResponse) GetDurationMs() int64 {
if x != nil {
return x.DurationMs
}
return 0
}
func (x *ExecResponse) GetError() string {
if x != nil {
return x.Error
}
return ""
}
func (x *ExecResponse) GetTruncated() bool {
if x != nil {
return x.Truncated
}
return false
}
type ExecOutput struct {
state protoimpl.MessageState `protogen:"open.v1"`
Stream ExecOutput_Stream `protobuf:"varint,1,opt,name=stream,proto3,enum=hostexec.ExecOutput_Stream" json:"stream,omitempty"`
Data []byte `protobuf:"bytes,2,opt,name=data,proto3" json:"data,omitempty"`
// Only set in the final message (exit_code is meaningful).
Done bool `protobuf:"varint,3,opt,name=done,proto3" json:"done,omitempty"`
ExitCode int32 `protobuf:"varint,4,opt,name=exit_code,json=exitCode,proto3" json:"exit_code,omitempty"`
Error string `protobuf:"bytes,5,opt,name=error,proto3" json:"error,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *ExecOutput) Reset() {
*x = ExecOutput{}
mi := &file_hostexec_pb_hostexec_proto_msgTypes[2]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *ExecOutput) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*ExecOutput) ProtoMessage() {}
func (x *ExecOutput) ProtoReflect() protoreflect.Message {
mi := &file_hostexec_pb_hostexec_proto_msgTypes[2]
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 ExecOutput.ProtoReflect.Descriptor instead.
func (*ExecOutput) Descriptor() ([]byte, []int) {
return file_hostexec_pb_hostexec_proto_rawDescGZIP(), []int{2}
}
func (x *ExecOutput) GetStream() ExecOutput_Stream {
if x != nil {
return x.Stream
}
return ExecOutput_STDOUT
}
func (x *ExecOutput) GetData() []byte {
if x != nil {
return x.Data
}
return nil
}
func (x *ExecOutput) GetDone() bool {
if x != nil {
return x.Done
}
return false
}
func (x *ExecOutput) GetExitCode() int32 {
if x != nil {
return x.ExitCode
}
return 0
}
func (x *ExecOutput) GetError() string {
if x != nil {
return x.Error
}
return ""
}
var File_hostexec_pb_hostexec_proto protoreflect.FileDescriptor
const file_hostexec_pb_hostexec_proto_rawDesc = "" +
"\n" +
"\x1ahostexec/pb/hostexec.proto\x12\bhostexec\"\xa5\x02\n" +
"\vExecRequest\x12\x18\n" +
"\acommand\x18\x01 \x01(\tR\acommand\x12\x12\n" +
"\x04args\x18\x02 \x03(\tR\x04args\x12\x1f\n" +
"\vworking_dir\x18\x03 \x01(\tR\n" +
"workingDir\x120\n" +
"\x03env\x18\x04 \x03(\v2\x1e.hostexec.ExecRequest.EnvEntryR\x03env\x12\x14\n" +
"\x05stdin\x18\x05 \x01(\fR\x05stdin\x12\x1d\n" +
"\n" +
"timeout_ms\x18\x06 \x01(\x03R\ttimeoutMs\x12(\n" +
"\x10max_output_bytes\x18\a \x01(\x03R\x0emaxOutputBytes\x1a6\n" +
"\bEnvEntry\x12\x10\n" +
"\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" +
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\xb0\x01\n" +
"\fExecResponse\x12\x1b\n" +
"\texit_code\x18\x01 \x01(\x05R\bexitCode\x12\x16\n" +
"\x06stdout\x18\x02 \x01(\fR\x06stdout\x12\x16\n" +
"\x06stderr\x18\x03 \x01(\fR\x06stderr\x12\x1f\n" +
"\vduration_ms\x18\x04 \x01(\x03R\n" +
"durationMs\x12\x14\n" +
"\x05error\x18\x05 \x01(\tR\x05error\x12\x1c\n" +
"\ttruncated\x18\x06 \x01(\bR\ttruncated\"\xbe\x01\n" +
"\n" +
"ExecOutput\x123\n" +
"\x06stream\x18\x01 \x01(\x0e2\x1b.hostexec.ExecOutput.StreamR\x06stream\x12\x12\n" +
"\x04data\x18\x02 \x01(\fR\x04data\x12\x12\n" +
"\x04done\x18\x03 \x01(\bR\x04done\x12\x1b\n" +
"\texit_code\x18\x04 \x01(\x05R\bexitCode\x12\x14\n" +
"\x05error\x18\x05 \x01(\tR\x05error\" \n" +
"\x06Stream\x12\n" +
"\n" +
"\x06STDOUT\x10\x00\x12\n" +
"\n" +
"\x06STDERR\x10\x012~\n" +
"\bHostExec\x125\n" +
"\x04Exec\x12\x15.hostexec.ExecRequest\x1a\x16.hostexec.ExecResponse\x12;\n" +
"\n" +
"ExecStream\x12\x15.hostexec.ExecRequest\x1a\x14.hostexec.ExecOutput0\x01B#Z!github.com/yaoapp/tai/hostexec/pbb\x06proto3"
var (
file_hostexec_pb_hostexec_proto_rawDescOnce sync.Once
file_hostexec_pb_hostexec_proto_rawDescData []byte
)
func file_hostexec_pb_hostexec_proto_rawDescGZIP() []byte {
file_hostexec_pb_hostexec_proto_rawDescOnce.Do(func() {
file_hostexec_pb_hostexec_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_hostexec_pb_hostexec_proto_rawDesc), len(file_hostexec_pb_hostexec_proto_rawDesc)))
})
return file_hostexec_pb_hostexec_proto_rawDescData
}
var file_hostexec_pb_hostexec_proto_enumTypes = make([]protoimpl.EnumInfo, 1)
var file_hostexec_pb_hostexec_proto_msgTypes = make([]protoimpl.MessageInfo, 4)
var file_hostexec_pb_hostexec_proto_goTypes = []any{
(ExecOutput_Stream)(0), // 0: hostexec.ExecOutput.Stream
(*ExecRequest)(nil), // 1: hostexec.ExecRequest
(*ExecResponse)(nil), // 2: hostexec.ExecResponse
(*ExecOutput)(nil), // 3: hostexec.ExecOutput
nil, // 4: hostexec.ExecRequest.EnvEntry
}
var file_hostexec_pb_hostexec_proto_depIdxs = []int32{
4, // 0: hostexec.ExecRequest.env:type_name -> hostexec.ExecRequest.EnvEntry
0, // 1: hostexec.ExecOutput.stream:type_name -> hostexec.ExecOutput.Stream
1, // 2: hostexec.HostExec.Exec:input_type -> hostexec.ExecRequest
1, // 3: hostexec.HostExec.ExecStream:input_type -> hostexec.ExecRequest
2, // 4: hostexec.HostExec.Exec:output_type -> hostexec.ExecResponse
3, // 5: hostexec.HostExec.ExecStream:output_type -> hostexec.ExecOutput
4, // [4:6] is the sub-list for method output_type
2, // [2:4] 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_hostexec_pb_hostexec_proto_init() }
func file_hostexec_pb_hostexec_proto_init() {
if File_hostexec_pb_hostexec_proto != nil {
return
}
type x struct{}
out := protoimpl.TypeBuilder{
File: protoimpl.DescBuilder{
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
RawDescriptor: unsafe.Slice(unsafe.StringData(file_hostexec_pb_hostexec_proto_rawDesc), len(file_hostexec_pb_hostexec_proto_rawDesc)),
NumEnums: 1,
NumMessages: 4,
NumExtensions: 0,
NumServices: 1,
},
GoTypes: file_hostexec_pb_hostexec_proto_goTypes,
DependencyIndexes: file_hostexec_pb_hostexec_proto_depIdxs,
EnumInfos: file_hostexec_pb_hostexec_proto_enumTypes,
MessageInfos: file_hostexec_pb_hostexec_proto_msgTypes,
}.Build()
File_hostexec_pb_hostexec_proto = out.File
file_hostexec_pb_hostexec_proto_goTypes = nil
file_hostexec_pb_hostexec_proto_depIdxs = nil
}

46
tai/hostexec/pb/hostexec.proto Normal file
View file

@ -0,0 +1,46 @@
syntax = "proto3";
package hostexec;
option go_package = "github.com/yaoapp/yao/tai/hostexec/pb";
// HostExec provides remote command execution on the Tai host machine.
// High privilege — enabled only when --host-exec flag is set.
service HostExec {
// Exec runs a command and returns the result when it completes.
rpc Exec(ExecRequest) returns (ExecResponse);
// ExecStream runs a command and streams stdout/stderr in real time.
rpc ExecStream(ExecRequest) returns (stream ExecOutput);
}
message ExecRequest {
string command = 1;
repeated string args = 2;
string working_dir = 3;
map<string, string> env = 4;
bytes stdin = 5;
int64 timeout_ms = 6;
int64 max_output_bytes = 7; // max stdout+stderr size (0 = default 10MB), truncate if exceeded
}
message ExecResponse {
int32 exit_code = 1;
bytes stdout = 2;
bytes stderr = 3;
int64 duration_ms = 4;
string error = 5; // non-empty if Tai failed to execute (not the command's error)
bool truncated = 6; // true if stdout+stderr exceeded max_output_bytes and was truncated
}
message ExecOutput {
enum Stream {
STDOUT = 0;
STDERR = 1;
}
Stream stream = 1;
bytes data = 2;
// Only set in the final message (exit_code is meaningful).
bool done = 3;
int32 exit_code = 4;
string error = 5;
}

173
tai/hostexec/pb/hostexec_grpc.pb.go Normal file
View file

@ -0,0 +1,173 @@
// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
// versions:
// - protoc-gen-go-grpc v1.6.1
// - protoc v4.25.0
// source: hostexec/pb/hostexec.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 (
HostExec_Exec_FullMethodName = "/hostexec.HostExec/Exec"
HostExec_ExecStream_FullMethodName = "/hostexec.HostExec/ExecStream"
)
// HostExecClient is the client API for HostExec 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.
//
// HostExec provides remote command execution on the Tai host machine.
// High privilege — enabled only when --host-exec flag is set.
type HostExecClient interface {
// Exec runs a command and returns the result when it completes.
Exec(ctx context.Context, in *ExecRequest, opts ...grpc.CallOption) (*ExecResponse, error)
// ExecStream runs a command and streams stdout/stderr in real time.
ExecStream(ctx context.Context, in *ExecRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ExecOutput], error)
}
type hostExecClient struct {
cc grpc.ClientConnInterface
}
func NewHostExecClient(cc grpc.ClientConnInterface) HostExecClient {
return &hostExecClient{cc}
}
func (c *hostExecClient) Exec(ctx context.Context, in *ExecRequest, opts ...grpc.CallOption) (*ExecResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(ExecResponse)
err := c.cc.Invoke(ctx, HostExec_Exec_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
func (c *hostExecClient) ExecStream(ctx context.Context, in *ExecRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ExecOutput], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &HostExec_ServiceDesc.Streams[0], HostExec_ExecStream_FullMethodName, cOpts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[ExecRequest, ExecOutput]{ClientStream: stream}
if err := x.ClientStream.SendMsg(in); err != nil {
return nil, err
}
if err := x.ClientStream.CloseSend(); err != nil {
return nil, err
}
return x, nil
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type HostExec_ExecStreamClient = grpc.ServerStreamingClient[ExecOutput]
// HostExecServer is the server API for HostExec service.
// All implementations must embed UnimplementedHostExecServer
// for forward compatibility.
//
// HostExec provides remote command execution on the Tai host machine.
// High privilege — enabled only when --host-exec flag is set.
type HostExecServer interface {
// Exec runs a command and returns the result when it completes.
Exec(context.Context, *ExecRequest) (*ExecResponse, error)
// ExecStream runs a command and streams stdout/stderr in real time.
ExecStream(*ExecRequest, grpc.ServerStreamingServer[ExecOutput]) error
mustEmbedUnimplementedHostExecServer()
}
// UnimplementedHostExecServer 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 UnimplementedHostExecServer struct{}
func (UnimplementedHostExecServer) Exec(context.Context, *ExecRequest) (*ExecResponse, error) {
return nil, status.Error(codes.Unimplemented, "method Exec not implemented")
}
func (UnimplementedHostExecServer) ExecStream(*ExecRequest, grpc.ServerStreamingServer[ExecOutput]) error {
return status.Error(codes.Unimplemented, "method ExecStream not implemented")
}
func (UnimplementedHostExecServer) mustEmbedUnimplementedHostExecServer() {}
func (UnimplementedHostExecServer) testEmbeddedByValue() {}
// UnsafeHostExecServer may be embedded to opt out of forward compatibility for this service.
// Use of this interface is not recommended, as added methods to HostExecServer will
// result in compilation errors.
type UnsafeHostExecServer interface {
mustEmbedUnimplementedHostExecServer()
}
func RegisterHostExecServer(s grpc.ServiceRegistrar, srv HostExecServer) {
// If the following call panics, it indicates UnimplementedHostExecServer 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(&HostExec_ServiceDesc, srv)
}
func _HostExec_Exec_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(ExecRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(HostExecServer).Exec(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: HostExec_Exec_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(HostExecServer).Exec(ctx, req.(*ExecRequest))
}
return interceptor(ctx, in, info, handler)
}
func _HostExec_ExecStream_Handler(srv interface{}, stream grpc.ServerStream) error {
m := new(ExecRequest)
if err := stream.RecvMsg(m); err != nil {
return err
}
return srv.(HostExecServer).ExecStream(m, &grpc.GenericServerStream[ExecRequest, ExecOutput]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type HostExec_ExecStreamServer = grpc.ServerStreamingServer[ExecOutput]
// HostExec_ServiceDesc is the grpc.ServiceDesc for HostExec service.
// It's only intended for direct use with grpc.RegisterService,
// and not to be introspected or modified (even as a copy)
var HostExec_ServiceDesc = grpc.ServiceDesc{
ServiceName: "hostexec.HostExec",
HandlerType: (*HostExecServer)(nil),
Methods: []grpc.MethodDesc{
{
MethodName: "Exec",
Handler: _HostExec_Exec_Handler,
},
},
Streams: []grpc.StreamDesc{
{
StreamName: "ExecStream",
Handler: _HostExec_ExecStream_Handler,
ServerStreams: true,
},
},
Metadata: "hostexec/pb/hostexec.proto",
}

2
tai/proxy/proxy.go
View file

@ -21,7 +21,7 @@ type Proxy interface {
type ConnectOptions struct {
Port int // container port
Path string // URL path (e.g. "/ws" or "/events")
Protocol string // "ws", "sse", or "tcp"
Protocol string // "ws" or "sse"
}
// Connection represents a persistent connection to a container service.

76
tai/registry/registry.go
View file

@ -13,6 +13,15 @@ import (
"github.com/gorilla/websocket"
)
// SystemInfo describes the host machine running Tai.
type SystemInfo struct {
OS string `json:"os"`
Arch string `json:"arch"`
Hostname string `json:"hostname"`
NumCPU int `json:"num_cpu"`
TotalMem int64 `json:"total_mem,omitempty"`
}
// TaiNode represents a registered Tai instance (direct or tunnel).
// Internal use only; external callers receive NodeSnapshot via Get()/List().
type TaiNode struct {
@ -20,10 +29,11 @@ type TaiNode struct {
MachineID string
Version string
Auth AuthInfo
System SystemInfo
Mode string // "direct" | "tunnel"
Addr string // direct mode: "tai-host"; tunnel mode: empty
YaoBase string // Yao server base URL reported by Tai (tunnel mode)
Ports map[string]int // {"grpc":9100, "http":8080, "vnc":6080, "docker":2375}
Ports map[string]int // {"grpc":19100, "http":8099, "vnc":16080, "docker":12375}
Capabilities map[string]bool
ControlConn *websocket.Conn
@ -43,6 +53,7 @@ type NodeSnapshot struct {
MachineID string
Version string
Auth AuthInfo
System SystemInfo
Mode string
Addr string
YaoBase string
@ -65,7 +76,8 @@ func (n *TaiNode) snapshot() NodeSnapshot {
}
return NodeSnapshot{
TaiID: n.TaiID, MachineID: n.MachineID, Version: n.Version,
Auth: n.Auth, Mode: n.Mode, Addr: n.Addr, YaoBase: n.YaoBase,
Auth: n.Auth, System: n.System,
Mode: n.Mode, Addr: n.Addr, YaoBase: n.YaoBase,
Ports: ports, Capabilities: caps,
Status: n.Status, ConnectedAt: n.ConnectedAt, LastPing: n.LastPing,
PoolName: n.PoolName,
@ -222,6 +234,66 @@ func (r *Registry) UpdatePing(taiID string) {
}
}
// ListByTeam returns snapshots of all nodes belonging to the given team.
func (r *Registry) ListByTeam(teamID string) []NodeSnapshot {
r.mu.RLock()
defer r.mu.RUnlock()
var result []NodeSnapshot
for _, n := range r.nodes {
if n.Auth.TeamID == teamID {
result = append(result, n.snapshot())
}
}
return result
}
// StartHealthCheck runs a background goroutine that periodically checks
// direct-mode nodes for heartbeat timeout. Nodes whose LastPing exceeds
// timeout are marked offline. Nodes that remain offline longer than
// cleanupAfter are automatically unregistered.
// The goroutine stops when ctx.Done() is closed.
func (r *Registry) StartHealthCheck(done <-chan struct{}, interval, timeout, cleanupAfter time.Duration) {
go func() {
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-done:
return
case <-ticker.C:
r.checkHealth(timeout, cleanupAfter)
}
}
}()
}
func (r *Registry) checkHealth(timeout, cleanupAfter time.Duration) {
now := time.Now()
var toRemove []string
r.mu.Lock()
for id, n := range r.nodes {
if n.Mode != "direct" {
continue
}
elapsed := now.Sub(n.LastPing)
if n.Status == "online" && elapsed > timeout {
n.Status = "offline"
r.logger.Warn("tai node offline (heartbeat timeout)",
"tai_id", id, "last_ping", n.LastPing)
}
if n.Status == "offline" && elapsed > timeout+cleanupAfter {
toRemove = append(toRemove, id)
}
}
r.mu.Unlock()
for _, id := range toRemove {
r.logger.Info("tai node auto-unregistered (offline too long)", "tai_id", id)
r.Unregister(id)
}
}
// RequestChannel sends an "open" command to a tunnel-connected Tai via its
// control channel. Returns a channel_id that Tai will use to connect back.
// Blocks until the data channel is established or timeout.

155
tai/registry/registry_test.go
View file

@ -29,7 +29,7 @@ func TestRegister_SetsFieldsAndOnline(t *testing.T) {
MachineID: "m-abc",
Version: "1.0.0",
Mode: "tunnel",
Ports: map[string]int{"grpc": 9100},
Ports: map[string]int{"grpc": 19100},
}
r.Register(node)
@ -117,14 +117,14 @@ func TestSnapshot_DeepCopy(t *testing.T) {
r := newTestRegistry()
r.Register(&TaiNode{
TaiID: "tai-001",
Ports: map[string]int{"grpc": 9100, "http": 8080},
Ports: map[string]int{"grpc": 19100, "http": 8099},
})
snap, _ := r.Get("tai-001")
snap.Ports["grpc"] = 0
snap2, _ := r.Get("tai-001")
if snap2.Ports["grpc"] != 9100 {
if snap2.Ports["grpc"] != 19100 {
t.Error("snapshot modification leaked into registry node")
}
}
@ -165,7 +165,7 @@ func TestWriteControlJSON_NilConn(t *testing.T) {
func TestRequestChannel_NotFound(t *testing.T) {
r := newTestRegistry()
_, _, err := r.RequestChannel("ghost", 9100)
_, _, err := r.RequestChannel("ghost", 19100)
if err == nil {
t.Fatal("expected error for missing node")
}
@ -174,7 +174,7 @@ func TestRequestChannel_NotFound(t *testing.T) {
func TestRequestChannel_DirectMode(t *testing.T) {
r := newTestRegistry()
r.Register(&TaiNode{TaiID: "tai-001", Mode: "direct"})
_, _, err := r.RequestChannel("tai-001", 9100)
_, _, err := r.RequestChannel("tai-001", 19100)
if err == nil {
t.Fatal("expected error for direct-mode node")
}
@ -357,7 +357,7 @@ func TestRequestChannel_Success(t *testing.T) {
r := newTestRegistry()
r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn})
channelID, resultCh, err := r.RequestChannel("tai-001", 9100)
channelID, resultCh, err := r.RequestChannel("tai-001", 19100)
if err != nil {
t.Fatalf("RequestChannel: %v", err)
}
@ -379,8 +379,8 @@ func TestRequestChannel_Success(t *testing.T) {
if cmd["channel_id"] != channelID {
t.Errorf("cmd channel_id = %v, want %s", cmd["channel_id"], channelID)
}
if tp, ok := cmd["target_port"].(float64); !ok || int(tp) != 9100 {
t.Errorf("cmd target_port = %v, want 9100", cmd["target_port"])
if tp, ok := cmd["target_port"].(float64); !ok || int(tp) != 19100 {
t.Errorf("cmd target_port = %v, want 19100", cmd["target_port"])
}
case <-time.After(2 * time.Second):
t.Fatal("timeout waiting for open command")
@ -391,7 +391,7 @@ func TestRequestChannel_NoControlConn(t *testing.T) {
r := newTestRegistry()
r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel"})
_, _, err := r.RequestChannel("tai-001", 9100)
_, _, err := r.RequestChannel("tai-001", 19100)
if err == nil {
t.Fatal("expected error for nil ControlConn")
}
@ -420,7 +420,7 @@ func TestOpenLocalListener_Success(t *testing.T) {
r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn})
ln, err := r.OpenLocalListener("tai-001", 9100)
ln, err := r.OpenLocalListener("tai-001", 19100)
if err != nil {
t.Fatalf("OpenLocalListener: %v", err)
}
@ -448,8 +448,8 @@ func TestOpenLocalListener_Success(t *testing.T) {
if _, ok := cmd["channel_id"].(string); !ok {
t.Error("open cmd missing channel_id")
}
if tp, ok := cmd["target_port"].(float64); !ok || int(tp) != 9100 {
t.Errorf("target_port = %v, want 9100", cmd["target_port"])
if tp, ok := cmd["target_port"].(float64); !ok || int(tp) != 19100 {
t.Errorf("target_port = %v, want 19100", cmd["target_port"])
}
case <-time.After(2 * time.Second):
t.Fatal("timeout waiting for open command from local listener")
@ -458,7 +458,7 @@ func TestOpenLocalListener_Success(t *testing.T) {
func TestOpenLocalListener_NodeNotFound(t *testing.T) {
r := newTestRegistry()
_, err := r.OpenLocalListener("ghost", 9100)
_, err := r.OpenLocalListener("ghost", 19100)
if err == nil {
t.Fatal("expected error for missing node")
}
@ -475,6 +475,135 @@ func newWSServer(handler func(*websocket.Conn)) *httptest.Server {
}))
}
func TestRegister_SystemInfo(t *testing.T) {
r := newTestRegistry()
r.Register(&TaiNode{
TaiID: "tai-001",
System: SystemInfo{
OS: "linux",
Arch: "amd64",
Hostname: "docker-host-01",
NumCPU: 16,
},
})
snap, ok := r.Get("tai-001")
if !ok {
t.Fatal("node not found")
}
if snap.System.OS != "linux" {
t.Errorf("System.OS = %q, want linux", snap.System.OS)
}
if snap.System.Arch != "amd64" {
t.Errorf("System.Arch = %q, want amd64", snap.System.Arch)
}
if snap.System.Hostname != "docker-host-01" {
t.Errorf("System.Hostname = %q, want docker-host-01", snap.System.Hostname)
}
if snap.System.NumCPU != 16 {
t.Errorf("System.NumCPU = %d, want 16", snap.System.NumCPU)
}
}
func TestListByTeam(t *testing.T) {
r := newTestRegistry()
r.Register(&TaiNode{TaiID: "tai-a", Auth: AuthInfo{TeamID: "team-dev"}})
r.Register(&TaiNode{TaiID: "tai-b", Auth: AuthInfo{TeamID: "team-dev"}})
r.Register(&TaiNode{TaiID: "tai-c", Auth: AuthInfo{TeamID: "team-ops"}})
devNodes := r.ListByTeam("team-dev")
if len(devNodes) != 2 {
t.Errorf("ListByTeam(team-dev) = %d nodes, want 2", len(devNodes))
}
opsNodes := r.ListByTeam("team-ops")
if len(opsNodes) != 1 {
t.Errorf("ListByTeam(team-ops) = %d nodes, want 1", len(opsNodes))
}
empty := r.ListByTeam("team-ghost")
if len(empty) != 0 {
t.Errorf("ListByTeam(team-ghost) = %d nodes, want 0", len(empty))
}
}
func TestStartHealthCheck_MarkOffline(t *testing.T) {
r := newTestRegistry()
r.Register(&TaiNode{TaiID: "tai-direct", Mode: "direct"})
r.Register(&TaiNode{TaiID: "tai-tunnel", Mode: "tunnel"})
// Manually set LastPing to the past for the direct node.
r.mu.Lock()
r.nodes["tai-direct"].LastPing = time.Now().Add(-5 * time.Second)
r.mu.Unlock()
done := make(chan struct{})
r.StartHealthCheck(done, 50*time.Millisecond, 2*time.Second, 10*time.Minute)
defer close(done)
time.Sleep(200 * time.Millisecond)
snap, ok := r.Get("tai-direct")
if !ok {
t.Fatal("direct node should still exist")
}
if snap.Status != "offline" {
t.Errorf("direct node Status = %q, want offline", snap.Status)
}
// Tunnel nodes should not be affected.
snap2, ok := r.Get("tai-tunnel")
if !ok {
t.Fatal("tunnel node should still exist")
}
if snap2.Status != "online" {
t.Errorf("tunnel node Status = %q, want online", snap2.Status)
}
}
func TestStartHealthCheck_AutoCleanup(t *testing.T) {
r := newTestRegistry()
r.Register(&TaiNode{TaiID: "tai-stale", Mode: "direct"})
// Set LastPing far in the past so it exceeds both timeout and cleanupAfter.
r.mu.Lock()
r.nodes["tai-stale"].LastPing = time.Now().Add(-1 * time.Hour)
r.mu.Unlock()
done := make(chan struct{})
r.StartHealthCheck(done, 50*time.Millisecond, 1*time.Second, 1*time.Second)
defer close(done)
time.Sleep(200 * time.Millisecond)
if _, ok := r.Get("tai-stale"); ok {
t.Error("stale node should have been auto-unregistered")
}
}
func TestStartHealthCheck_PingKeepsAlive(t *testing.T) {
r := newTestRegistry()
r.Register(&TaiNode{TaiID: "tai-alive", Mode: "direct"})
done := make(chan struct{})
r.StartHealthCheck(done, 50*time.Millisecond, 2*time.Second, 10*time.Minute)
defer close(done)
// Continuously ping to keep the node alive.
for i := 0; i < 4; i++ {
time.Sleep(30 * time.Millisecond)
r.UpdatePing("tai-alive")
}
snap, ok := r.Get("tai-alive")
if !ok {
t.Fatal("node should still exist")
}
if snap.Status != "online" {
t.Errorf("Status = %q, want online", snap.Status)
}
}
func TestNodeSnapshot_AuthInfo(t *testing.T) {
r := newTestRegistry()
r.Register(&TaiNode{

2
tai/sandbox/docker.go
View file

@ -13,7 +13,7 @@ type dockerSandbox struct {
}
// NewDocker creates a Sandbox backed by Docker SDK through Tai's Docker API proxy.
// addr should be "tcp://tai-host:2375".
// addr should be "tcp://tai-host:12375".
func NewDocker(addr string) (Sandbox, error) {
cli, err := client.NewClientWithOpts(
client.WithHost(addr),

16
tai/sandbox/k8s.go
View file

@ -349,12 +349,18 @@ func (s *k8sSandbox) Inspect(ctx context.Context, id string) (*ContainerInfo, er
}
func (s *k8sSandbox) List(ctx context.Context, opts ListOptions) ([]ContainerInfo, error) {
labelSelector := "managed-by=yao-tai-sdk"
if len(opts.Labels) > 0 {
for k, v := range opts.Labels {
labelSelector += "," + k + "=" + v
}
merged := make(map[string]string)
for k, v := range s.labels {
merged[k] = v
}
for k, v := range opts.Labels {
merged[k] = v
}
var parts []string
for k, v := range merged {
parts = append(parts, k+"="+v)
}
labelSelector := strings.Join(parts, ",")
pods, err := s.cli.CoreV1().Pods(s.ns).List(ctx, metav1.ListOptions{
LabelSelector: labelSelector,

158
tai/tai.go
View file

@ -10,6 +10,7 @@ import (
"strings"
"time"
hepb "github.com/yaoapp/yao/tai/hostexec/pb"
"github.com/yaoapp/yao/tai/proxy"
"github.com/yaoapp/yao/tai/registry"
"github.com/yaoapp/yao/tai/sandbox"
@ -42,11 +43,11 @@ func (f optionFunc) apply(c *config) { f(c) }
// Ports configures service ports for Tai server.
type Ports struct {
GRPC int // default 9100
HTTP int // default 8080
VNC int // default 6080
Docker int // default 2375
K8s int // default 6443
GRPC int // default 19100
HTTP int // default 8099
VNC int // default 16080
Docker int // default 12375
K8s int // default 16443
}
// WithPorts overrides default Tai service ports.
@ -98,9 +99,9 @@ type config struct {
func defaultPorts() Ports {
return Ports{
GRPC: 9100,
HTTP: 8080,
VNC: 6080,
GRPC: 19100,
HTTP: 8099,
VNC: 16080,
}
}
@ -136,6 +137,7 @@ type Client struct {
img sandbox.Image
prx proxy.Proxy
vc vnc.VNC
he hepb.HostExecClient
grpcConn *grpc.ClientConn
// tunnel mode: local listeners that bridge to Tai via WS
@ -217,52 +219,61 @@ func (c *Client) initRemote(cfg *config) (*Client, error) {
return nil, fmt.Errorf("grpc dial %s: %w", grpcAddr, err)
}
c.grpcConn = conn
c.he = hepb.NewHostExecClient(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
caps, err := c.discoverServerInfo(conn, cfg)
if err != nil {
// Old Tai without ServerInfo — fall back to legacy behaviour (try Docker).
caps = map[string]bool{"docker": true}
}
hasDocker := caps["docker"]
hasK8s := caps["k8s"]
hasHostExec := caps["host_exec"]
if !hasDocker && !hasK8s && !hasHostExec {
conn.Close()
return nil, fmt.Errorf("tai %s: no capabilities available (docker/k8s/host_exec all false)", c.host)
}
c.vol = volume.NewRemote(conn)
switch cfg.runtime {
case K8s:
if cfg.runtime == K8s {
if cfg.kubeConfig == "" {
conn.Close()
return nil, fmt.Errorf("tai %s: K8s runtime requested but no kubeconfig provided", c.host)
}
k8sPort := c.ports.K8s
if k8sPort == 0 {
k8sPort = 6443
k8sPort = 16443
}
sbAddr := fmt.Sprintf("%s:%d", c.host, k8sPort)
sb, err := sandbox.NewK8s(sbAddr, sandbox.K8sOption{
Namespace: cfg.namespace,
KubeConfig: cfg.kubeConfig,
})
if err != nil {
conn.Close()
return nil, err
if err == nil {
c.sb = sb
c.img = sandbox.NewK8sImage()
}
c.sb = sb
c.img = sandbox.NewK8sImage()
default:
} else if hasDocker {
dockerPort := c.ports.Docker
if dockerPort == 0 {
dockerPort = 2375
dockerPort = 12375
}
sbAddr := fmt.Sprintf("tcp://%s:%d", c.host, dockerPort)
sb, err := sandbox.NewDocker(sbAddr)
if err != nil {
conn.Close()
return nil, err
if err == nil {
c.sb = sb
c.img = sandbox.NewDockerImage(sandbox.DockerCli(sb))
}
c.sb = sb
c.img = sandbox.NewDockerImage(sandbox.DockerCli(sb))
}
hc := cfg.httpClient
c.prx = proxy.NewRemote(c.host, c.ports.HTTP, hc)
c.vc = vnc.NewRemote(c.host, c.ports.VNC, hc)
if c.sb != nil {
hc := cfg.httpClient
c.prx = proxy.NewRemote(c.host, c.ports.HTTP, hc)
c.vc = vnc.NewRemote(c.host, c.ports.VNC, hc)
}
if reg := registry.Global(); reg != nil {
reg.Register(&registry.TaiNode{
@ -295,10 +306,10 @@ func (c *Client) initTunnel(cfg *config) (*Client, error) {
}
c.ports = Ports{
GRPC: nodePort(node.Ports, "grpc", 9100),
HTTP: nodePort(node.Ports, "http", 8080),
VNC: nodePort(node.Ports, "vnc", 6080),
Docker: nodePort(node.Ports, "docker", 2375),
GRPC: nodePort(node.Ports, "grpc", 19100),
HTTP: nodePort(node.Ports, "http", 8099),
VNC: nodePort(node.Ports, "vnc", 16080),
Docker: nodePort(node.Ports, "docker", 12375),
}
grpcLn, err := reg.OpenLocalListener(taiID, c.ports.GRPC)
@ -315,28 +326,40 @@ func (c *Client) initTunnel(cfg *config) (*Client, error) {
return nil, fmt.Errorf("grpc dial tunnel %s: %w", grpcAddr, err)
}
c.grpcConn = conn
c.he = hepb.NewHostExecClient(conn)
c.vol = volume.NewRemote(conn)
dockerLn, err := reg.OpenLocalListener(taiID, c.ports.Docker)
caps, err := c.discoverServerInfo(conn, cfg)
if err != nil {
conn.Close()
grpcLn.Close()
return nil, fmt.Errorf("open docker tunnel listener: %w", err)
caps = map[string]bool{"docker": true}
}
c.tunnelListeners = append(c.tunnelListeners, dockerLn)
sbAddr := fmt.Sprintf("tcp://%s", dockerLn.Addr().String())
sb, err := sandbox.NewDocker(sbAddr)
if err != nil {
hasDocker := caps["docker"]
hasHostExec := caps["host_exec"]
if !hasDocker && !hasHostExec {
c.closeTunnelListeners()
conn.Close()
return nil, err
return nil, fmt.Errorf("tai %s: no capabilities available via tunnel", taiID)
}
c.sb = sb
c.img = sandbox.NewDockerImage(sandbox.DockerCli(sb))
c.prx = proxy.NewTunnel(taiID, node.YaoBase)
c.vc = vnc.NewTunnel(taiID, node.YaoBase)
if hasDocker && c.ports.Docker > 0 {
dockerLn, err := reg.OpenLocalListener(taiID, c.ports.Docker)
if err == nil {
c.tunnelListeners = append(c.tunnelListeners, dockerLn)
sbAddr := fmt.Sprintf("tcp://%s", dockerLn.Addr().String())
sb, err := sandbox.NewDocker(sbAddr)
if err == nil {
c.sb = sb
c.img = sandbox.NewDockerImage(sandbox.DockerCli(sb))
}
}
}
if c.sb != nil {
c.prx = proxy.NewTunnel(taiID, node.YaoBase)
c.vc = vnc.NewTunnel(taiID, node.YaoBase)
}
return c, nil
}
@ -396,18 +419,26 @@ func (c *Client) Workspace(sessionID string) workspace.FS {
return workspace.New(c.vol, sessionID)
}
// Sandbox returns the container lifecycle manager. Never nil.
// Sandbox returns the container lifecycle manager.
// Nil when the Tai server has no container runtime (host-exec-only mode).
func (c *Client) Sandbox() sandbox.Sandbox { return c.sb }
// Image returns the container image manager. Never nil.
// Image returns the container image manager.
// Nil when the Tai server has no container runtime.
func (c *Client) Image() sandbox.Image { return c.img }
// Proxy returns the HTTP reverse proxy helper. Never nil.
// Proxy returns the HTTP reverse proxy helper.
// Nil when the Tai server has no container runtime.
func (c *Client) Proxy() proxy.Proxy { return c.prx }
// VNC returns the VNC WebSocket helper. Never nil.
// VNC returns the VNC WebSocket helper.
// Nil when the Tai server has no container runtime.
func (c *Client) VNC() vnc.VNC { return c.vc }
// HostExec returns the HostExec gRPC client for executing commands on the Tai
// host machine. Returns nil in local mode (no Tai server).
func (c *Client) HostExec() hepb.HostExecClient { return c.he }
// IsLocal returns true if the client connects directly to a Docker daemon.
func (c *Client) IsLocal() bool { return c.scheme == "docker" }
@ -427,7 +458,7 @@ func parseAddr(addr string) (scheme, host, dockerAddr string, grpcPort int, err
if isLocalHost(addr) {
return "docker", "", "", 0, nil
}
// host:port — split carefully (IPv6 like [::1]:9100 is already handled above)
// host:port — split carefully (IPv6 like [::1]:19100 is already handled above)
h := addr
if idx := strings.LastIndex(addr, ":"); idx > 0 {
h = addr[:idx]
@ -484,21 +515,19 @@ 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 {
// discoverServerInfo calls ServerInfo.GetInfo on the remote Tai server, merges
// discovered ports into c.ports, and returns the server's capabilities map.
// Ports explicitly set via WithPorts take precedence over server-reported values.
func (c *Client) discoverServerInfo(conn *grpc.ClientConn, cfg *config) (map[string]bool, 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
return nil, 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 {
@ -513,5 +542,10 @@ func (c *Client) discoverPorts(conn *grpc.ClientConn, cfg *config) error {
if p := int(resp.Ports["k8s"]); p > 0 && up.K8s == 0 {
c.ports.K8s = p
}
return nil
caps := resp.Capabilities
if caps == nil {
caps = make(map[string]bool)
}
return caps, nil
}

47
tai/tai_test.go
View file

@ -14,6 +14,26 @@ func taiTestHost() string {
return "127.0.0.1"
}
// taiRemoteAddr returns the tai:// address for remote tests (e.g. TestNewRemoteDocker).
// Uses TAI_TEST_HOST and, when set, TAI_TEST_GRPC_PORT so Tai on non-default port works.
func taiRemoteAddr() string {
host := taiTestHost()
if p := os.Getenv("TAI_TEST_GRPC_PORT"); p != "" {
return "tai://" + host + ":" + p
}
return "tai://" + host
}
// taiTestPorts builds a Ports struct from TAI_TEST_*_PORT env vars.
// Only non-zero fields are set so they override ServerInfo-discovered values.
func taiTestPorts() Ports {
return Ports{
Docker: envPort("TAI_TEST_DOCKER_PORT", 0),
HTTP: envPort("TAI_TEST_HTTP_PORT", 0),
VNC: envPort("TAI_TEST_VNC_PORT", 0),
}
}
func envPort(key string, fallback int) int {
if v := os.Getenv(key); v != "" {
if p, err := strconv.Atoi(v); err == nil {
@ -84,11 +104,11 @@ func TestMergedPorts(t *testing.T) {
if p.HTTP != 8888 {
t.Errorf("HTTP = %d, want 8888", p.HTTP)
}
if p.GRPC != 9100 {
t.Errorf("GRPC = %d, want 9100 (default)", p.GRPC)
if p.GRPC != 19100 {
t.Errorf("GRPC = %d, want 19100 (default)", p.GRPC)
}
if p.VNC != 6080 {
t.Errorf("VNC = %d, want 6080 (default)", p.VNC)
if p.VNC != 16080 {
t.Errorf("VNC = %d, want 16080 (default)", p.VNC)
}
if p.Docker != 0 {
t.Errorf("Docker = %d, want 0 (unset)", p.Docker)
@ -202,12 +222,12 @@ 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))
grpcPort := envPort("TAI_TEST_K8S_GRPC_PORT", envPort("TAI_TEST_GRPC_PORT", 19100))
ports := Ports{
K8s: envPort("TAI_TEST_K8S_PORT", 6443),
K8s: envPort("TAI_TEST_K8S_PORT", 16443),
GRPC: grpcPort,
HTTP: envPort("TAI_TEST_K8S_HTTP_PORT", 8080),
VNC: envPort("TAI_TEST_K8S_VNC_PORT", 6080),
HTTP: envPort("TAI_TEST_K8S_HTTP_PORT", 8099),
VNC: envPort("TAI_TEST_K8S_VNC_PORT", 16080),
}
c, err := New(fmt.Sprintf("tai://%s:%d", host, grpcPort), K8s,
@ -255,13 +275,16 @@ func TestNewInvalidScheme(t *testing.T) {
}
func TestNewRemoteDocker(t *testing.T) {
addr := "tai://" + taiTestHost()
c, err := New(addr)
addr := taiRemoteAddr()
ports := taiTestPorts()
c, err := New(addr, WithPorts(ports))
if err != nil {
t.Skipf("Tai not available at %s: %v", addr, err)
}
defer c.Close()
t.Logf("remote docker: addr=%s ports=%+v", addr, c.ports)
if c.IsLocal() {
t.Error("expected IsLocal = false for tai://")
}
@ -284,7 +307,7 @@ func TestNewRemoteDocker(t *testing.T) {
}
func TestDiscoverPorts(t *testing.T) {
addr := "tai://" + taiTestHost()
addr := taiRemoteAddr()
c, err := New(addr)
if err != nil {
t.Skipf("Tai not available at %s: %v", addr, err)
@ -303,7 +326,7 @@ func TestDiscoverPorts(t *testing.T) {
}
func TestDiscoverPortsWithUserOverride(t *testing.T) {
addr := "tai://" + taiTestHost()
addr := taiRemoteAddr()
c, err := New(addr, WithPorts(Ports{HTTP: 9999}))
if err != nil {
t.Skipf("Tai not available at %s: %v", addr, err)

17
tai/token.go Normal file
View file

@ -0,0 +1,17 @@
package tai
import grpcclient "github.com/yaoapp/yao/grpc/client"
// TokenManager is an alias for grpc/client.TokenManager.
// New code should use grpc/client.TokenManager directly.
type TokenManager = grpcclient.TokenManager
// NewTokenManagerFromEnv creates a TokenManager from environment variables.
func NewTokenManagerFromEnv() (*TokenManager, error) {
return grpcclient.NewTokenManagerFromEnv()
}
// NewTokenManager creates a TokenManager with explicit values.
func NewTokenManager(accessToken, refreshToken, sandboxID string) *TokenManager {
return grpcclient.NewTokenManager(accessToken, refreshToken, sandboxID)
}

4
tai/tunnel/proxy.go
View file

@ -33,7 +33,7 @@ func HandleProxy(c *gin.Context) {
httpPort := node.Ports["http"]
if httpPort == 0 {
httpPort = 8080
httpPort = 8099
}
channelID, resultCh, err := reg.RequestChannel(taiID, httpPort)
@ -104,7 +104,7 @@ func HandleVNC(c *gin.Context) {
vncPort := node.Ports["vnc"]
if vncPort == 0 {
vncPort = 6080
vncPort = 16080
}
channelID, resultCh, err := reg.RequestChannel(taiID, vncPort)

16
tai/tunnel/server.go
View file

@ -73,6 +73,7 @@ func HandleControl(c *gin.Context) {
MachineID: regMsg.MachineID,
Version: regMsg.Version,
Auth: authInfo,
System: regMsg.System,
Mode: "tunnel",
YaoBase: regMsg.Server,
Ports: regMsg.Ports,
@ -159,13 +160,14 @@ func HandleData(c *gin.Context) {
// registerMessage is the JSON structure for Tai's register message.
type registerMessage struct {
Type string `json:"type"`
TaiID string `json:"tai_id"`
MachineID string `json:"machine_id"`
Version string `json:"version"`
Server string `json:"server"`
Ports map[string]int `json:"ports"`
Capabilities map[string]bool `json:"capabilities"`
Type string `json:"type"`
TaiID string `json:"tai_id"`
MachineID string `json:"machine_id"`
Version string `json:"version"`
Server string `json:"server"`
Ports map[string]int `json:"ports"`
Capabilities map[string]bool `json:"capabilities"`
System registry.SystemInfo `json:"system"`
}
// controlMsg is a generic control channel message.

2
tai/volume/pb/volume.proto
View file

@ -3,7 +3,7 @@ package volume;
option go_package = "github.com/yaoapp/tai/volume/pb";
// Volume provides bulk file synchronization and real-time filesystem I/O.
// Shares gRPC port :9100 with Yao Gateway.
// Shares gRPC port :19100 with Yao Gateway.
service Volume {
// --- Bulk Sync ---

4
tai/volume/pb/volume_grpc.pb.go
View file

@ -35,7 +35,7 @@ const (
// 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.
//
// Volume provides bulk file synchronization and real-time filesystem I/O.
// Shares gRPC port :9100 with Yao Gateway.
// Shares gRPC port :19100 with Yao Gateway.
type VolumeClient interface {
// SyncPush: Yao sends code to Tai (before container start).
// Bidirectional stream:
@ -183,7 +183,7 @@ func (c *volumeClient) MkdirAll(ctx context.Context, in *FSRequest, opts ...grpc
// for forward compatibility.
//
// Volume provides bulk file synchronization and real-time filesystem I/O.
// Shares gRPC port :9100 with Yao Gateway.
// Shares gRPC port :19100 with Yao Gateway.
type VolumeServer interface {
// SyncPush: Yao sends code to Tai (before container start).
// Bidirectional stream:

2
tai/volume/volume.go
View file

@ -8,7 +8,7 @@ import (
)
// Volume provides filesystem IO and directory synchronization.
// Remote connects to Tai gRPC :9100; Local operates directly on disk.
// Remote connects to Tai gRPC :19100; Local operates directly on disk.
type Volume interface {
ReadFile(ctx context.Context, sessionID, path string) ([]byte, os.FileMode, error)
WriteFile(ctx context.Context, sessionID, path string, data []byte, perm os.FileMode) error

10
tai/volume/volume_test.go
View file

@ -14,7 +14,15 @@ func taiTestGRPC() string {
if addr := os.Getenv("TAI_TEST_GRPC"); addr != "" {
return addr
}
return "127.0.0.1:9100"
host := os.Getenv("TAI_TEST_HOST")
if host == "" {
host = "127.0.0.1"
}
port := os.Getenv("TAI_TEST_GRPC_PORT")
if port == "" {
port = "19100"
}
return host + ":" + port
}
func TestLocalVolume(t *testing.T) {

35
tai/yao.go Normal file
View file

@ -0,0 +1,35 @@
package tai
import (
"context"
grpcclient "github.com/yaoapp/yao/grpc/client"
"github.com/yaoapp/yao/grpc/pb"
)
// YaoClient wraps grpc/client.Client for backward compatibility.
// New code should use grpc/client.Client directly.
type YaoClient = grpcclient.Client
// NewYaoClientFromEnv reads YAO_GRPC_ADDR and token env vars, dials the
// gRPC server, and returns a connected YaoClient.
func NewYaoClientFromEnv() (*YaoClient, error) {
return grpcclient.NewFromEnv()
}
// DialYao connects to a Yao gRPC server at addr with the given TokenManager.
func DialYao(addr string, tm *TokenManager) (*YaoClient, error) {
return grpcclient.Dial(addr, tm)
}
// --- Convenience wrappers kept for sandbox/container code ---
// Run executes a Yao process via the given client.
func Run(ctx context.Context, c *YaoClient, process string, args []byte, timeout int32) ([]byte, error) {
return c.Run(ctx, process, args, timeout)
}
// Shell executes a system command via the given client.
func Shell(ctx context.Context, c *YaoClient, command string, args []string, env map[string]string, timeout int32) (*pb.ShellResponse, error) {
return c.Shell(ctx, command, args, env, timeout)
}