Merge pull request #1491 from trheyi/main

feat(assistant): enhance assistant info structure and sandbox integration
This commit is contained in:
Max 2026-03-12 21:45:02 +08:00 committed by GitHub
commit 3a58b7aef3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
77 changed files with 6147 additions and 2991 deletions

113
.github/actions/setup-yao/action.yml vendored Normal file
View file

@ -0,0 +1,113 @@
name: "Setup Yao Build Environment"
description: "Checkout dependency repos, setup Go toolchain, and install build tools (v1.0.0)"
inputs:
go-version:
description: "Go version to install"
default: "1.25"
repo-kun:
description: "Kun repository (owner/repo)"
required: true
repo-xun:
description: "Xun repository (owner/repo)"
required: true
repo-gou:
description: "Gou repository (owner/repo)"
required: true
checkout-app:
description: "Checkout yao-dev-app (demo application for tests)"
default: "true"
checkout-init:
description: "Checkout yao-init (for Yao server startup in CI)"
default: "false"
apple-private-key:
description: "Apple private key content for OAuth certs (optional)"
default: ""
runs:
using: "composite"
steps:
# -- Dependency repositories --
- name: Checkout Kun
uses: actions/checkout@v4
with:
repository: ${{ inputs.repo-kun }}
path: kun
- name: Checkout Xun
uses: actions/checkout@v4
with:
repository: ${{ inputs.repo-xun }}
path: xun
- name: Checkout Gou
uses: actions/checkout@v4
with:
repository: ${{ inputs.repo-gou }}
path: gou
- name: Checkout V8Go
uses: actions/checkout@v4
with:
repository: yaoapp/v8go
path: v8go
- name: Unzip libv8
shell: bash
run: |
for file in $(find ./v8go -name "libv8*.zip"); do
dir=$(dirname "$file")
echo "Extracting $file to $dir"
unzip -o -d "$dir" "$file"
rm -rf "$dir/__MACOSX"
done
- name: Checkout Demo App
if: ${{ inputs.checkout-app == 'true' }}
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: Checkout yao-init
if: ${{ inputs.checkout-init == 'true' }}
uses: actions/checkout@v4
with:
repository: yaoapp/yao-init
path: yao-init
# -- Move all dependencies to parent directory (Go workspace layout) --
- name: Move Dependencies
shell: bash
run: |
mv kun ../
mv xun ../
mv gou ../
mv v8go ../
[ -d app ] && mv app ../
mv extension ../
[ -d yao-init ] && mv yao-init ../
# -- Setup Apple Private Key (if provided) --
- name: Setup Apple Private Key
if: ${{ inputs.apple-private-key != '' }}
shell: bash
run: |
mkdir -p ../app/openapi/certs/apple
echo "${{ inputs.apple-private-key }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8
# -- Go toolchain --
- name: Setup Go ${{ inputs.go-version }}
uses: actions/setup-go@v5
with:
go-version: ${{ inputs.go-version }}
- name: Setup Go Tools
shell: bash
run: make tools

141
.github/env/sandbox-v2.env vendored Normal file
View file

@ -0,0 +1,141 @@
# ============================================================
# Yao CI Environment — sandbox-v2 (v1.0.0)
# Loaded via: cat .github/env/sandbox-v2.env >> $GITHUB_ENV
# ============================================================
# ========================================
# Yao Runtime (YAO_ prefix, read by Yao)
# ========================================
YAO_HOST=0.0.0.0
YAO_PORT=5099
YAO_GRPC_HOST=0.0.0.0
YAO_GRPC_PORT=9099
YAO_DB_DRIVER=sqlite3
YAO_SESSION=memory
YAO_ENV=development
# ========================================
# CI Test Parameters (YAO_CI_ prefix)
# ========================================
# -- Network --
YAO_CI_BRIDGE_IP=172.17.0.1
# -- Yao service ports (tests read these, not YAO_PORT/YAO_GRPC_PORT) --
YAO_CI_HTTP_PORT=5099
YAO_CI_GRPC_PORT=9099
YAO_CI_URL=http://127.0.0.1:5099
YAO_CI_GRPC=127.0.0.1:9099
# -- OAuth token generation (ci-token tool) --
YAO_CI_OAUTH_SUBJECT=ci-test-user
YAO_CI_OAUTH_USER_ID=ci-test-user
YAO_CI_OAUTH_TEAM_ID=ci-test-team
YAO_CI_OAUTH_SCOPE=tai:tunnel
YAO_CI_OAUTH_TTL=24h
# ========================================
# Tai Instances
#
# Connection modes:
# tai-local → DIRECT (--direct, Yao dials Tai gRPC directly)
# tai-docker → TUNNEL (default, Tai dials Yao gRPC, reverse tunnel)
# tai-k8s → TUNNEL (same as above, with K8s runtime)
# tai-hostexec → TUNNEL (same as above, no container runtime)
# ========================================
# -- tai-local (DIRECT mode, auto-detect Docker via /var/run/docker.sock) --
# Yao dials tai-local gRPC directly, so gRPC port must be reachable.
# Docker proxy on 12376 (not default 12375) to avoid conflict with tai-docker.
YAO_CI_TAI_LOCAL_HOST=127.0.0.1
YAO_CI_TAI_LOCAL_GRPC_PORT=19103
YAO_CI_TAI_LOCAL_HTTP_PORT=8102
YAO_CI_TAI_LOCAL_VNC_PORT=16083
YAO_CI_TAI_LOCAL_DOCKER_PORT=12376
YAO_CI_TAI_LOCAL_GRPC=127.0.0.1:19103
YAO_CI_TAI_LOCAL_DOCKER_API=tcp://127.0.0.1:12376
# -- tai-docker (TUNNEL mode, explicit Docker API proxy) --
# Tunnel: Tai connects to Yao gRPC. Sandbox connects to Yao, traffic forwarded via tunnel.
# gRPC port used only for Tai's own listener; Yao accesses via tunnel, not direct dial.
YAO_CI_TAI_DOCKER_HOST=127.0.0.1
YAO_CI_TAI_DOCKER_GRPC_PORT=19100
YAO_CI_TAI_DOCKER_HTTP_PORT=8099
YAO_CI_TAI_DOCKER_VNC_PORT=16080
YAO_CI_TAI_DOCKER_API_PORT=12375
YAO_CI_TAI_DOCKER_API=tcp://127.0.0.1:12375
# -- tai-k8s (TUNNEL mode, K8s API proxy via k3d) --
# K8s proxy on 16444 (not 16443) because k3d --api-port already binds 16443.
# TAI_K8S_UPSTREAM points to k3d at 127.0.0.1:16443; proxy exposes on 16444.
YAO_CI_TAI_K8S_HOST=127.0.0.1
YAO_CI_TAI_K8S_GRPC_PORT=19101
YAO_CI_TAI_K8S_HTTP_PORT=8100
YAO_CI_TAI_K8S_VNC_PORT=16081
YAO_CI_TAI_K8S_API_PORT=16444
YAO_CI_K3D_API_PORT=16443
# -- tai-hostexec (TUNNEL mode, no container runtime, HostExec only) --
YAO_CI_TAI_HOSTEXEC_HOST=127.0.0.1
YAO_CI_TAI_HOSTEXEC_GRPC_PORT=19102
YAO_CI_TAI_HOSTEXEC_HTTP_PORT=8101
YAO_CI_TAI_HOSTEXEC_VNC_PORT=16082
# ========================================
# Sandbox V2 addresses (used by test code)
# ========================================
YAO_CI_SANDBOX_LOCAL_ADDR=tai://127.0.0.1:19103
YAO_CI_SANDBOX_DOCKER_ADDR=tai://127.0.0.1:19100
YAO_CI_SANDBOX_K8S_ADDR=tai://127.0.0.1:19101
YAO_CI_SANDBOX_IMAGE=yaoapp/tai-sandbox-test:latest
# ========================================
# HostExec addresses
# ========================================
YAO_CI_HOSTEXEC_LOCAL_ADDR=127.0.0.1:19103
YAO_CI_HOSTEXEC_DOCKER_ADDR=127.0.0.1:19100
YAO_CI_HOSTEXEC_K8S_ADDR=127.0.0.1:19101
YAO_CI_HOSTEXEC_ONLY_ADDR=127.0.0.1:19102
# -- Tunnel --
YAO_CI_TUNNEL=true
# ========================================
# Database (MySQL / PostgreSQL / SQLite)
# ========================================
MYSQL_TEST_HOST=127.0.0.1
MYSQL_TEST_PORT=3308
MYSQL_TEST_USER=test
MYSQL_TEST_PASS=123456
PG_TEST_HOST=127.0.0.1
PG_TEST_PORT=5432
PG_TEST_USER=test
PG_TEST_PASS=123456
SQLITE_DB=./app/db/yao.db
# ========================================
# Legacy variable mapping (migrate later)
# ========================================
TAI_TEST_HOST=127.0.0.1
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=16444
TAI_TEST_K8S_GRPC_PORT=19101
TAI_TEST_K8S_HTTP_PORT=8100
TAI_TEST_K8S_VNC_PORT=16081
TAI_TEST_HOST_IP=172.17.0.1
TAI_TEST_TUNNEL=true
TAI_TEST_YAO_URL=http://127.0.0.1:5099
TAI_TEST_YAO_GRPC=127.0.0.1:9099
SANDBOX_TEST_LOCAL_ADDR=tai://127.0.0.1:19103
SANDBOX_TEST_REMOTE_ADDR=tai://127.0.0.1:19100
SANDBOX_TEST_K8S_REMOTE_ADDR=tai://127.0.0.1:19101
SANDBOX_TEST_HOSTEXEC_ADDR=tai://127.0.0.1:19102
SANDBOX_TEST_IMAGE=yaoapp/tai-sandbox-test:latest
DOCKER_BRIDGE_IP=172.17.0.1

534
.github/workflows/unit-test-v1.yml vendored Normal file
View file

@ -0,0 +1,534 @@
name: Unit Test V1
on:
workflow_dispatch:
inputs:
tags:
description: "Version"
env:
CI_VERSION: "1.0.0"
REPO_KUN: ${{ github.repository_owner }}/kun
REPO_XUN: ${{ github.repository_owner }}/xun
REPO_GOU: ${{ github.repository_owner }}/gou
YAO_DEV: ${{ github.WORKSPACE }}
YAO_ENV: development
YAO_ROOT: ${{ github.WORKSPACE }}/../app
YAO_HOST: 0.0.0.0
YAO_PORT: 5099
YAO_SESSION: "memory"
YAO_LOG: "./logs/application.log"
YAO_LOG_MODE: "TEXT"
YAO_JWT_SECRET: "bLp@bi!oqo-2U+hoTRUG"
YAO_DB_AESKEY: "ZLX=T&f6refeCh-ro*r@"
YAO_EXTENSION_ROOT: ${{ github.WORKSPACE }}/../extension
YAO_TEST_APPLICATION: ${{ github.WORKSPACE }}/../app
YAO_RUNTIME_MIN: 3
YAO_RUNTIME_MAX: 6
YAO_RUNTIME_HEAP_LIMIT: 1500000000
YAO_RUNTIME_HEAP_RELEASE: 10000000
YAO_RUNTIME_HEAP_AVAILABLE: 550000000
YAO_RUNTIME_PRECOMPILE: true
MYSQL_TEST_HOST: "127.0.0.1"
MYSQL_TEST_PORT: "3308"
MYSQL_TEST_USER: "test"
MYSQL_TEST_PASS: "123456"
REDIS_TEST_HOST: "127.0.0.1"
REDIS_TEST_PORT: "6379"
REDIS_TEST_DB: "2"
MONGO_TEST_HOST: "127.0.0.1"
MONGO_TEST_PORT: "27017"
MONGO_TEST_USER: "root"
MONGO_TEST_PASS: "123456"
PG_TEST_HOST: "127.0.0.1"
PG_TEST_PORT: "5432"
PG_TEST_USER: "test"
PG_TEST_PASS: "123456"
jobs:
# =============================================================================
# Environment Setup & Verification
# Build Yao, start services, connect Tai via gRPC tunnel, verify everything.
# No tests are run — this job validates the CI environment is healthy.
# =============================================================================
setup-and-verify:
runs-on: ubuntu-latest
services:
mongodb:
image: mongo:6.0
ports:
- 27017:27017
env:
MONGO_INITDB_ROOT_USERNAME: root
MONGO_INITDB_ROOT_PASSWORD: 123456
MONGO_INITDB_DATABASE: test
postgres:
image: postgres:14
ports:
- 5432:5432
env:
POSTGRES_USER: test
POSTGRES_PASSWORD: 123456
POSTGRES_DB: test
options: >-
--health-cmd="pg_isready -U test"
--health-interval=10s
--health-timeout=5s
--health-retries=5
strategy:
matrix:
go: ["1.25"]
steps:
# ==== Phase 1: Checkout & Setup ====
- name: Checkout Yao
uses: actions/checkout@v4
- name: Setup Build Environment
uses: ./.github/actions/setup-yao
with:
repo-kun: ${{ env.REPO_KUN }}
repo-xun: ${{ env.REPO_XUN }}
repo-gou: ${{ env.REPO_GOU }}
checkout-init: "true"
apple-private-key: ${{ secrets.APPLE_PRIVATE_KEY_USER }}
- name: Load sandbox-v2 env
run: grep -vE '^\s*#|^\s*$' .github/env/sandbox-v2.env >> $GITHUB_ENV
- name: Setup SQLite
run: |
mkdir -p ${{ github.WORKSPACE }}/../app/db
echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV
- name: Start MySQL 8.0
run: |
docker run -d --name mysql \
-e MYSQL_RANDOM_ROOT_PASSWORD=true \
-e MYSQL_USER=${MYSQL_TEST_USER} \
-e MYSQL_PASSWORD=${MYSQL_TEST_PASS} \
-e MYSQL_DATABASE=test \
-p ${MYSQL_TEST_PORT}:3306 \
mysql:8.0 --port=3306 --sql-mode='' \
--character-set-server=utf8mb4 --collation-server=utf8mb4_general_ci
for i in $(seq 1 30); do
if docker exec mysql mysqladmin ping -h 127.0.0.1 -u ${MYSQL_TEST_USER} -p${MYSQL_TEST_PASS} > /dev/null 2>&1; then
echo "MySQL ready"
break
fi
echo "Waiting for MySQL... ($i/30)"
sleep 2
done
- name: Start Redis
run: docker run --name redis -d -p 6379:6379 redis:6
# ==== Phase 2: Code Quality & Build ====
- name: Code Quality Check
run: |
make vet
make fmt-check
- name: Build Yao
run: go build -v -o $RUNNER_TEMP/yao .
- name: Build ci-token
run: go build -tags ci -v -o $RUNNER_TEMP/ci-token ./cmd/ci-token
- name: Extract tai binary from image
run: |
CID=$(docker create yaoapp/tai:latest)
docker cp "$CID":/usr/local/bin/tai $RUNNER_TEMP/tai
docker rm "$CID"
chmod +x $RUNNER_TEMP/tai
$RUNNER_TEMP/tai version || echo "tai binary extracted"
# ==== Phase 3: Prepare & Start Yao ====
- name: Prepare test app directory
run: |
cp -r ${{ github.WORKSPACE }}/../yao-init $RUNNER_TEMP/yao-test-app
mkdir -p $RUNNER_TEMP/yao-test-app/db
- name: Start Yao server
run: |
cd $RUNNER_TEMP/yao-test-app
YAO_ROOT=$(pwd) \
YAO_HOST=0.0.0.0 \
YAO_PORT=5099 \
YAO_GRPC_HOST=0.0.0.0 \
YAO_GRPC_PORT=9099 \
YAO_DB_DRIVER=sqlite3 \
YAO_DB_PRIMARY=$(pwd)/db/yao.db \
YAO_SESSION=memory \
YAO_ENV=development \
YAO_JWT_SECRET="${{ env.YAO_JWT_SECRET }}" \
YAO_DB_AESKEY="${{ env.YAO_DB_AESKEY }}" \
$RUNNER_TEMP/yao start &
# Wait for Yao HTTP to be ready (up to 120s)
for i in $(seq 1 60); do
if curl -sf http://127.0.0.1:5099/.well-known/yao > /dev/null 2>&1; then
echo "Yao HTTP ready"
curl -s http://127.0.0.1:5099/.well-known/yao | jq .
break
fi
echo "Waiting for Yao... ($i/60)"
sleep 2
done
curl -sf http://127.0.0.1:5099/.well-known/yao > /dev/null 2>&1 || {
echo "::error::Yao HTTP failed to start"
exit 1
}
# ==== Phase 4: Generate Tai credentials ====
- name: Generate Tai credentials
run: |
gen_cred() {
local CID=$1 TID=$2 OUT=$3
local TOKEN
TOKEN=$($RUNNER_TEMP/ci-token \
--app $RUNNER_TEMP/yao-test-app \
--client-id "$CID" \
--subject "${YAO_CI_OAUTH_SUBJECT:-ci-tai}" \
--user-id "${YAO_CI_OAUTH_USER_ID}" \
--team-id "${YAO_CI_OAUTH_TEAM_ID}" \
--scope "${YAO_CI_OAUTH_SCOPE:-tai:tunnel}" \
--ttl "${YAO_CI_OAUTH_TTL:-24h}" 2>/dev/null | tail -1 | tr -d '[:space:]')
local JSON="{\"client_id\":\"$CID\",\"tai_id\":\"$TID\",\"machine_id\":\"ci-runner\",\"server\":\"http://127.0.0.1:${YAO_CI_HTTP_PORT}\",\"yao_grpc_addr\":\"127.0.0.1:${YAO_CI_GRPC_PORT}\",\"access_token\":\"$TOKEN\",\"scope\":\"${YAO_CI_OAUTH_SCOPE}\",\"expires_at\":\"2099-01-01T00:00:00Z\",\"registered\":true}"
echo -n "$JSON" | base64 -w0 > "$OUT"
echo ""
echo "Credentials JSON (debug): $JSON" | head -c 200
echo "..."
echo "Generated credentials for $CID → $OUT"
}
gen_cred tai-ci-local tai-local-001 $RUNNER_TEMP/tai-local-credentials
gen_cred tai-ci-docker tai-docker-001 $RUNNER_TEMP/tai-docker-credentials
gen_cred tai-ci-k8s tai-k8s-001 $RUNNER_TEMP/tai-k8s-credentials
gen_cred tai-ci-hostexec tai-hostexec-001 $RUNNER_TEMP/tai-hostexec-credentials
# ==== Phase 5: Pull images & Setup K8s ====
- name: Pull test images
run: |
docker pull yaoapp/tai-sandbox-test:latest || true
docker pull yaoapp/tai:latest
docker pull alpine:latest
- name: Install k3d & create cluster
run: |
curl -s https://raw.githubusercontent.com/k3d-io/k3d/main/install.sh | bash
k3d cluster create tai-test --no-lb --wait --api-port ${YAO_CI_K3D_API_PORT}
kubectl wait --for=condition=Ready node --all --timeout=60s
k3d image import alpine:latest -c tai-test
- name: Generate kubeconfig
run: |
K3D_IP=$(docker inspect k3d-tai-test-server-0 | jq -r '.[0].NetworkSettings.Networks["k3d-tai-test"].IPAddress')
echo "k3d server IP: ${K3D_IP}"
k3d kubeconfig get tai-test > /tmp/kubeconfig-k3d.yml
# For tai-k8s container (uses k3d internal IP)
sed "s|server: .*|server: https://${K3D_IP}:6443|" /tmp/kubeconfig-k3d.yml \
> /tmp/kubeconfig-tai-k8s.yml
echo "Container kubeconfig server:"
grep server: /tmp/kubeconfig-tai-k8s.yml
# For test runner (uses localhost via k3d port-mapped API)
sed "s|server: .*|server: https://127.0.0.1:${YAO_CI_K3D_API_PORT}|" /tmp/kubeconfig-k3d.yml \
> $RUNNER_TEMP/kubeconfig-tai.yml
echo "Test runner kubeconfig server:"
grep server: $RUNNER_TEMP/kubeconfig-tai.yml
# Export for later steps
echo "TAI_TEST_KUBECONFIG=$RUNNER_TEMP/kubeconfig-tai.yml" >> $GITHUB_ENV
echo "YAO_CI_TAI_KUBECONFIG=$RUNNER_TEMP/kubeconfig-tai.yml" >> $GITHUB_ENV
# ==== Phase 6: Start Tai instances (host processes) ====
- name: Start tai-local (DIRECT mode, auto-detect Docker)
run: |
mkdir -p $RUNNER_TEMP/tai-local-data
TAI_CREDENTIALS=$RUNNER_TEMP/tai-local-credentials \
TAI_YAO_SERVER=http://127.0.0.1:${YAO_CI_HTTP_PORT} \
TAI_DATA_DIR=$RUNNER_TEMP/tai-local-data \
$RUNNER_TEMP/tai server \
--grpc 127.0.0.1:${YAO_CI_TAI_LOCAL_GRPC_PORT} \
--http 127.0.0.1:${YAO_CI_TAI_LOCAL_HTTP_PORT} \
--vnc 127.0.0.1:${YAO_CI_TAI_LOCAL_VNC_PORT} \
--docker 127.0.0.1:${YAO_CI_TAI_LOCAL_DOCKER_PORT} \
--direct \
--host-exec --host-exec-full-access \
--log-level debug &
echo $! > $RUNNER_TEMP/tai-local.pid
echo "tai-local PID: $(cat $RUNNER_TEMP/tai-local.pid)"
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:${YAO_CI_TAI_LOCAL_HTTP_PORT}/healthz > /dev/null 2>&1; then
echo "tai-local HTTP ready"
break
fi
echo "Waiting for tai-local HTTP... ($i/30)"
sleep 1
done
- name: Start tai-docker (TUNNEL mode, Docker API proxy)
run: |
mkdir -p $RUNNER_TEMP/tai-docker-data
TAI_CREDENTIALS=$RUNNER_TEMP/tai-docker-credentials \
TAI_YAO_SERVER=http://127.0.0.1:${YAO_CI_HTTP_PORT} \
TAI_DATA_DIR=$RUNNER_TEMP/tai-docker-data \
$RUNNER_TEMP/tai server \
--grpc 127.0.0.1:${YAO_CI_TAI_DOCKER_GRPC_PORT} \
--http 127.0.0.1:${YAO_CI_TAI_DOCKER_HTTP_PORT} \
--vnc 127.0.0.1:${YAO_CI_TAI_DOCKER_VNC_PORT} \
--docker 127.0.0.1:${YAO_CI_TAI_DOCKER_API_PORT} \
--host-exec --host-exec-full-access \
--log-level debug &
echo $! > $RUNNER_TEMP/tai-docker.pid
echo "tai-docker PID: $(cat $RUNNER_TEMP/tai-docker.pid)"
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:${YAO_CI_TAI_DOCKER_HTTP_PORT}/healthz > /dev/null 2>&1; then
echo "tai-docker HTTP ready"
break
fi
echo "Waiting for tai-docker HTTP... ($i/30)"
sleep 1
done
- name: Start tai-k8s (TUNNEL mode, K8s API proxy)
run: |
mkdir -p $RUNNER_TEMP/tai-k8s-data
TAI_CREDENTIALS=$RUNNER_TEMP/tai-k8s-credentials \
TAI_YAO_SERVER=http://127.0.0.1:${YAO_CI_HTTP_PORT} \
TAI_DATA_DIR=$RUNNER_TEMP/tai-k8s-data \
TAI_K8S_UPSTREAM="tcp://127.0.0.1:${YAO_CI_K3D_API_PORT}" \
TAI_KUBECONFIG=$RUNNER_TEMP/kubeconfig-tai.yml \
$RUNNER_TEMP/tai server \
--grpc 127.0.0.1:${YAO_CI_TAI_K8S_GRPC_PORT} \
--http 127.0.0.1:${YAO_CI_TAI_K8S_HTTP_PORT} \
--vnc 127.0.0.1:${YAO_CI_TAI_K8S_VNC_PORT} \
--k8s 127.0.0.1:${YAO_CI_TAI_K8S_API_PORT} \
--docker="" \
--host-exec --host-exec-full-access \
--log-level debug &
echo $! > $RUNNER_TEMP/tai-k8s.pid
echo "tai-k8s PID: $(cat $RUNNER_TEMP/tai-k8s.pid)"
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:${YAO_CI_TAI_K8S_HTTP_PORT}/healthz > /dev/null 2>&1; then
echo "tai-k8s HTTP ready"
break
fi
echo "Waiting for tai-k8s HTTP... ($i/30)"
sleep 1
done
- name: Start tai-hostexec (TUNNEL mode, HostExec only, no runtime)
run: |
mkdir -p $RUNNER_TEMP/tai-hostexec-data
TAI_CREDENTIALS=$RUNNER_TEMP/tai-hostexec-credentials \
TAI_YAO_SERVER=http://127.0.0.1:${YAO_CI_HTTP_PORT} \
TAI_DATA_DIR=$RUNNER_TEMP/tai-hostexec-data \
$RUNNER_TEMP/tai server \
--grpc 127.0.0.1:${YAO_CI_TAI_HOSTEXEC_GRPC_PORT} \
--http 127.0.0.1:${YAO_CI_TAI_HOSTEXEC_HTTP_PORT} \
--vnc 127.0.0.1:${YAO_CI_TAI_HOSTEXEC_VNC_PORT} \
--docker="" \
--host-exec --host-exec-full-access \
--log-level debug &
echo $! > $RUNNER_TEMP/tai-hostexec.pid
echo "tai-hostexec PID: $(cat $RUNNER_TEMP/tai-hostexec.pid)"
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:${YAO_CI_TAI_HOSTEXEC_HTTP_PORT}/healthz > /dev/null 2>&1; then
echo "tai-hostexec HTTP ready"
break
fi
echo "Waiting for tai-hostexec HTTP... ($i/30)"
sleep 1
done
# ==== Phase 7: Environment Verification (fail fast) ====
- name: Verify Environment
run: |
echo "CI Environment v${CI_VERSION}"
echo ""
FAILED=0
check() {
local name=$1; shift
if "$@" > /dev/null 2>&1; then
echo " [PASS] $name"
else
echo " [FAIL] $name"
FAILED=$((FAILED + 1))
fi
}
echo "=== Environment Verification ==="
# ── 1. Service Health ──
echo ""
echo "--- 1. Service Health ---"
echo "[Yao]"
check "Yao HTTP (/.well-known/yao)" curl -sf http://127.0.0.1:5099/.well-known/yao
check "Yao gRPC port" nc -z 127.0.0.1 9099
echo "[tai-local (DIRECT)]"
check "tai-local process alive" kill -0 $(cat $RUNNER_TEMP/tai-local.pid 2>/dev/null || echo 0)
check "tai-local HTTP (/healthz)" curl -sf http://127.0.0.1:${YAO_CI_TAI_LOCAL_HTTP_PORT}/healthz
check "tai-local gRPC reachable (direct)" nc -z 127.0.0.1 ${YAO_CI_TAI_LOCAL_GRPC_PORT}
echo "[tai-docker (TUNNEL)]"
check "tai-docker process alive" kill -0 $(cat $RUNNER_TEMP/tai-docker.pid 2>/dev/null || echo 0)
check "tai-docker HTTP (/healthz)" curl -sf http://127.0.0.1:${YAO_CI_TAI_DOCKER_HTTP_PORT}/healthz
check "tai-docker gRPC listener" nc -z 127.0.0.1 ${YAO_CI_TAI_DOCKER_GRPC_PORT}
echo "[tai-k8s (TUNNEL)]"
check "tai-k8s process alive" kill -0 $(cat $RUNNER_TEMP/tai-k8s.pid 2>/dev/null || echo 0)
check "tai-k8s HTTP (/healthz)" curl -sf http://127.0.0.1:${YAO_CI_TAI_K8S_HTTP_PORT}/healthz
check "tai-k8s gRPC listener" nc -z 127.0.0.1 ${YAO_CI_TAI_K8S_GRPC_PORT}
echo "[tai-hostexec (TUNNEL)]"
check "tai-hostexec process alive" kill -0 $(cat $RUNNER_TEMP/tai-hostexec.pid 2>/dev/null || echo 0)
check "tai-hostexec HTTP (/healthz)" curl -sf http://127.0.0.1:${YAO_CI_TAI_HOSTEXEC_HTTP_PORT}/healthz
check "tai-hostexec gRPC listener" nc -z 127.0.0.1 ${YAO_CI_TAI_HOSTEXEC_GRPC_PORT}
echo "[K8s (k3d via direct API)]"
check "kubectl get nodes (k3d direct)" kubectl --kubeconfig=$RUNNER_TEMP/kubeconfig-tai.yml get nodes
echo "[Data Stores]"
MONGO_CID=$(docker ps -qf "ancestor=mongo:6.0" | head -1)
check "MongoDB ping" docker exec "$MONGO_CID" mongosh --quiet \
-u ${MONGO_TEST_USER} -p ${MONGO_TEST_PASS} --authenticationDatabase admin \
--eval "db.runCommand({ping:1})"
check "MySQL ping" docker exec mysql mysqladmin ping -h 127.0.0.1 \
-u ${MYSQL_TEST_USER} -p${MYSQL_TEST_PASS}
PG_CID=$(docker ps -qf "ancestor=postgres:14" | head -1)
check "PostgreSQL ping" docker exec "$PG_CID" pg_isready -U ${PG_TEST_USER}
check "Redis ping" docker exec redis redis-cli ping
# ── 2. Network Topology ──
echo ""
echo "--- 2. Network Topology ---"
BRIDGE_IP=${YAO_CI_BRIDGE_IP}
echo "[Host network basics]"
check "docker0 bridge exists" ip addr show docker0
check "Bridge IP reachable (${BRIDGE_IP})" ping -c1 -W2 ${BRIDGE_IP}
check "Docker socket accessible" test -S /var/run/docker.sock
check "tai binary on runner" test -x $RUNNER_TEMP/tai
echo "[Yao endpoints (all Tai instances need these)]"
check "Yao HTTP :${YAO_CI_HTTP_PORT}" curl -sf http://127.0.0.1:${YAO_CI_HTTP_PORT}/.well-known/yao
check "Yao gRPC :${YAO_CI_GRPC_PORT}" nc -z 127.0.0.1 ${YAO_CI_GRPC_PORT}
echo "[DIRECT path: Yao → tai-local]"
check "Yao→tai-local gRPC :${YAO_CI_TAI_LOCAL_GRPC_PORT}" nc -z 127.0.0.1 ${YAO_CI_TAI_LOCAL_GRPC_PORT}
check "Yao→tai-local HTTP :${YAO_CI_TAI_LOCAL_HTTP_PORT}" curl -sf http://127.0.0.1:${YAO_CI_TAI_LOCAL_HTTP_PORT}/healthz
check "tai-local Docker API proxy :${YAO_CI_TAI_LOCAL_DOCKER_PORT}" nc -z 127.0.0.1 ${YAO_CI_TAI_LOCAL_DOCKER_PORT}
echo "[TUNNEL path: tai-docker → Yao gRPC (reverse tunnel)]"
check "tai-docker→Yao gRPC :${YAO_CI_GRPC_PORT}" nc -z 127.0.0.1 ${YAO_CI_GRPC_PORT}
check "tai-docker Docker API proxy :${YAO_CI_TAI_DOCKER_API_PORT}" nc -z 127.0.0.1 ${YAO_CI_TAI_DOCKER_API_PORT}
check "Docker API via proxy" bash -c "curl -sf http://127.0.0.1:${YAO_CI_TAI_DOCKER_API_PORT}/version | jq -r .ApiVersion"
echo "[TUNNEL path: tai-k8s → Yao gRPC (reverse tunnel)]"
check "tai-k8s→Yao gRPC :${YAO_CI_GRPC_PORT}" nc -z 127.0.0.1 ${YAO_CI_GRPC_PORT}
check "tai-k8s K8s API proxy :${YAO_CI_TAI_K8S_API_PORT}" nc -z 127.0.0.1 ${YAO_CI_TAI_K8S_API_PORT}
sed "s|server: .*|server: https://127.0.0.1:${YAO_CI_TAI_K8S_API_PORT}|" $RUNNER_TEMP/kubeconfig-tai.yml \
> $RUNNER_TEMP/kubeconfig-tai-proxy.yml
check "K8s API via tai-k8s proxy (kubectl)" kubectl --kubeconfig=$RUNNER_TEMP/kubeconfig-tai-proxy.yml --insecure-skip-tls-verify get nodes
echo "[TUNNEL path: tai-hostexec → Yao gRPC (reverse tunnel)]"
check "tai-hostexec→Yao gRPC :${YAO_CI_GRPC_PORT}" nc -z 127.0.0.1 ${YAO_CI_GRPC_PORT}
echo "[Data Store connectivity from runner]"
check "MongoDB :27017" nc -z 127.0.0.1 27017
check "Redis :6379" nc -z 127.0.0.1 6379
check "MySQL :${MYSQL_TEST_PORT}" nc -z 127.0.0.1 ${MYSQL_TEST_PORT}
check "PostgreSQL :${PG_TEST_PORT}" nc -z 127.0.0.1 ${PG_TEST_PORT}
# ── 3. Connection Mode Verification ──
echo ""
echo "--- 3. Connection Modes ---"
sleep 5
echo "[Credentials (4 tokens)]"
check "tai-local credentials exist" test -f $RUNNER_TEMP/tai-local-credentials
check "tai-docker credentials exist" test -f $RUNNER_TEMP/tai-docker-credentials
check "tai-k8s credentials exist" test -f $RUNNER_TEMP/tai-k8s-credentials
check "tai-hostexec credentials exist" test -f $RUNNER_TEMP/tai-hostexec-credentials
echo "[DIRECT: tai-local → Yao HTTP register → Yao dials tai-local gRPC]"
echo " tai-local registers via POST /tai-nodes/register"
echo " Yao dials back tai-local gRPC at 127.0.0.1:${YAO_CI_TAI_LOCAL_GRPC_PORT}"
echo "[TUNNEL: tai-docker → Yao gRPC :${YAO_CI_GRPC_PORT} (Register + Forward)]"
echo " Sandbox connects Yao gRPC → Forward stream → tai-docker :${YAO_CI_TAI_DOCKER_GRPC_PORT}"
echo "[TUNNEL: tai-k8s → Yao gRPC :${YAO_CI_GRPC_PORT} (Register + Forward)]"
echo " Sandbox connects Yao gRPC → Forward stream → tai-k8s :${YAO_CI_TAI_K8S_GRPC_PORT}"
echo "[TUNNEL: tai-hostexec → Yao gRPC :${YAO_CI_GRPC_PORT} (Register + Forward)]"
echo " HostExec only, no container runtime"
WELL_KNOWN=$(curl -sf http://127.0.0.1:5099/.well-known/yao 2>/dev/null || echo "{}")
echo ""
echo " Yao .well-known/yao:"
echo "$WELL_KNOWN" | jq . 2>/dev/null || echo " $WELL_KNOWN"
# ── 4. HostExec Readiness ──
echo ""
echo "--- 4. HostExec ---"
echo "[HostExec gRPC ports (all 4 instances)]"
check "HostExec tai-local (direct, auto-Docker)" nc -z 127.0.0.1 ${YAO_CI_TAI_LOCAL_GRPC_PORT}
check "HostExec tai-docker (tunnel, Docker proxy)" nc -z 127.0.0.1 ${YAO_CI_TAI_DOCKER_GRPC_PORT}
check "HostExec tai-k8s (tunnel, K8s proxy)" nc -z 127.0.0.1 ${YAO_CI_TAI_K8S_GRPC_PORT}
check "HostExec tai-hostexec (tunnel, no runtime)" nc -z 127.0.0.1 ${YAO_CI_TAI_HOSTEXEC_GRPC_PORT}
echo ""
echo "=========================================="
if [ $FAILED -gt 0 ]; then
echo "::error::$FAILED verification check(s) FAILED"
echo ""
echo "=== Diagnostic Info ==="
echo "--- Docker containers ---"
docker ps -a
echo ""
echo "--- Processes (yao + tai) ---"
ps aux | grep -E "yao|tai" | grep -v grep || true
echo ""
echo "--- Listening ports ---"
ss -tlnp | grep -E "5099|9099|19100|19101|19102|19103|8099|8100|8101|8102|12375|12376|16443|16444" || true
exit 1
else
echo "All verification checks PASSED"
fi

3
.gitignore vendored
View file

@ -75,4 +75,5 @@ tg-send
registry/data/ registry/data/
registry/manager/DESIGN*.md registry/manager/DESIGN*.md
tai/testdata/ tai/testdata/
agent/sandbox/docs/*.md agent/sandbox/docs/*.md
tai/docs/refactor-registration.md

View file

@ -479,11 +479,16 @@ func (ast *Assistant) GetInfo(locale ...string) *store.AssistantInfo {
} }
info := &store.AssistantInfo{ info := &store.AssistantInfo{
AssistantID: ast.ID, AssistantID: ast.ID,
Avatar: ast.Avatar, Avatar: ast.Avatar,
Connector: ast.Connector,
ConnectorOptions: ast.ConnectorOptions,
Modes: ast.Modes,
DefaultMode: ast.DefaultMode,
Sandbox: ast.IsSandbox,
ComputerFilter: ast.ComputerFilter,
} }
// Apply i18n translation if locale is provided
if loc != "" { if loc != "" {
info.Name = ast.GetName(loc) info.Name = ast.GetName(loc)
info.Description = ast.GetDescription(loc) info.Description = ast.GetDescription(loc)

View file

@ -402,6 +402,12 @@ func LoadPath(path string) (*Assistant, error) {
ast.SandboxV2 = sbCfg ast.SandboxV2 = sbCfg
} }
// Extract Sandbox flag and ComputerFilter from V2 sandbox config.
if ast.SandboxV2 != nil {
ast.IsSandbox = true
ast.ComputerFilter = ast.SandboxV2.Filter
}
// Compute config hash for V2 sandbox. // Compute config hash for V2 sandbox.
if ast.SandboxV2 != nil { if ast.SandboxV2 != nil {
var mcpServers []store.MCPServerConfig var mcpServers []store.MCPServerConfig
@ -772,6 +778,8 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
return nil, err return nil, err
} }
assistant.SandboxV2 = sb assistant.SandboxV2 = sb
assistant.IsSandbox = true
assistant.ComputerFilter = sb.Filter
} else { } else {
sb, err := store.ToSandbox(sandbox) sb, err := store.ToSandbox(sandbox)
if err != nil { if err != nil {

View file

@ -0,0 +1,17 @@
//go:build remote
package sandboxv2_test
import "os"
func init() {
extraNodeProviders = append(extraNodeProviders, agentRemoteNodes)
}
func agentRemoteNodes() []nodeConfig {
addr := os.Getenv("SANDBOX_TEST_REMOTE_ADDR")
if addr == "" {
return nil
}
return []nodeConfig{{Name: "remote", Addr: addr}}
}

View file

@ -13,19 +13,27 @@ import (
sandbox "github.com/yaoapp/yao/sandbox/v2" sandbox "github.com/yaoapp/yao/sandbox/v2"
"github.com/yaoapp/yao/tai" "github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/registry"
taisandbox "github.com/yaoapp/yao/tai/sandbox" tairuntime "github.com/yaoapp/yao/tai/runtime"
"github.com/yaoapp/yao/workspace" "github.com/yaoapp/yao/workspace"
) )
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
// node configuration — mirrors sandbox/v2 testutils but scoped to prepare tests // Build-tag extension points (same pattern as sandbox/v2).
// ---------------------------------------------------------------------------
var (
extraNodeProviders []func() []nodeConfig
extraHostExecProviders []func() []hostTarget
)
// ---------------------------------------------------------------------------
// Node / host configuration
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
type nodeConfig struct { type nodeConfig struct {
Name string Name string
Addr string Addr string
TaiID string TaiID string
Options []tai.Option DialOps []tai.DialOption
} }
type hostTarget struct { type hostTarget struct {
@ -35,7 +43,7 @@ type hostTarget struct {
} }
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
// environment helpers (same conventions as sandbox/v2 + env.local.sh) // Environment helpers
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
func testLocalAddr() string { func testLocalAddr() string {
@ -62,30 +70,84 @@ func envPort(key string, fallback int) int {
} }
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
// node discovery // Node / host discovery
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
func boxNodes() []nodeConfig { func boxNodes() []nodeConfig {
nodes := []nodeConfig{ nodes := []nodeConfig{
{Name: "local", Addr: testLocalAddr()}, {Name: "local", Addr: testLocalAddr()},
} }
if addr := os.Getenv("SANDBOX_TEST_REMOTE_ADDR"); addr != "" { for _, fn := range extraNodeProviders {
nodes = append(nodes, nodeConfig{Name: "remote", Addr: addr}) nodes = append(nodes, fn()...)
} }
return nodes return nodes
} }
func hostTargets() []hostTarget { func hostTargets() []hostTarget {
var targets []hostTarget var targets []hostTarget
if addr := os.Getenv("TAI_TEST_WIN_HOSTEXEC_LINUX"); addr != "" { for _, fn := range extraHostExecProviders {
targets = append(targets, hostTarget{Name: "win-linux", Addr: addr}) targets = append(targets, fn()...)
}
if addr := os.Getenv("TAI_TEST_WIN_HOSTEXEC_NATIVE"); addr != "" {
targets = append(targets, hostTarget{Name: "win-native", Addr: addr})
} }
return targets return targets
} }
// ---------------------------------------------------------------------------
// Dial + Register helper (replaces old tai.New)
// ---------------------------------------------------------------------------
func dialForTest(addr string, dialOps ...tai.DialOption) (*tai.ConnResources, error) {
if addr == "local" || addr == "" {
return tai.DialLocal("", "", nil)
}
host, grpcPort := parseHostPort(addr)
ports := tai.Ports{GRPC: grpcPort}
return tai.DialRemote(host, ports, dialOps...)
}
func registerForTest(t testing.TB, addr string, dialOps ...tai.DialOption) (string, *tai.ConnResources) {
t.Helper()
if registry.Global() == nil {
registry.Init(nil)
}
res, err := dialForTest(addr, dialOps...)
if err != nil {
t.Fatalf("dialForTest(%s): %v", addr, err)
}
taiID := taiIDFromAddr(addr)
reg := registry.Global()
reg.Register(&registry.TaiNode{TaiID: taiID, Mode: modeForAddr(addr)})
reg.SetResources(taiID, res)
return taiID, res
}
func taiIDFromAddr(addr string) string {
if addr == "local" || addr == "" {
return "local"
}
addr = strings.TrimPrefix(addr, "tai://")
parts := strings.SplitN(addr, ":", 2)
return parts[0]
}
func modeForAddr(addr string) string {
if addr == "local" || addr == "" {
return "local"
}
return "direct"
}
func parseHostPort(addr string) (string, int) {
addr = strings.TrimPrefix(addr, "tai://")
parts := strings.SplitN(addr, ":", 2)
h := parts[0]
if len(parts) == 2 {
if p, err := strconv.Atoi(parts[1]); err == nil {
return h, p
}
}
return h, 19100
}
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
// TestMain — purge stale containers from previous runs // TestMain — purge stale containers from previous runs
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
@ -100,16 +162,16 @@ func purgeStale() {
defer cancel() defer cancel()
for _, nc := range boxNodes() { for _, nc := range boxNodes() {
client, err := tai.New(nc.Addr, nc.Options...) res, err := dialForTest(nc.Addr, nc.DialOps...)
if err != nil { if err != nil {
continue continue
} }
sb := client.Sandbox() sb := res.Runtime
if sb == nil { if sb == nil {
client.Close() res.Close()
continue continue
} }
containers, _ := sb.List(ctx, taisandbox.ListOptions{All: true}) containers, _ := sb.List(ctx, tairuntime.ListOptions{All: true})
for _, c := range containers { for _, c := range containers {
id := c.Name id := c.Name
if id == "" { if id == "" {
@ -120,7 +182,7 @@ func purgeStale() {
log.Printf("[purge] %s: removed %s", nc.Name, id) log.Printf("[purge] %s: removed %s", nc.Name, id)
} }
} }
client.Close() res.Close()
} }
} }
@ -133,11 +195,9 @@ func setupManager(t *testing.T, nc *nodeConfig) *sandbox.Manager {
if registry.Global() == nil { if registry.Global() == nil {
registry.Init(nil) registry.Init(nil)
} }
client, err := tai.New(nc.Addr, nc.Options...) taiID, res := registerForTest(t, nc.Addr, nc.DialOps...)
if err != nil { nc.TaiID = taiID
t.Fatalf("tai.New(%s): %v", nc.Addr, err) t.Cleanup(func() { res.Close() })
}
nc.TaiID = client.TaiID()
sandbox.Init() sandbox.Init()
m := sandbox.M() m := sandbox.M()
@ -191,7 +251,7 @@ func setupHostManager(t *testing.T, tgt *hostTarget) *sandbox.Manager {
} }
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
// skip helpers // Skip helpers
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
func skipIfNoDocker(t *testing.T) { func skipIfNoDocker(t *testing.T) {

View file

@ -0,0 +1,20 @@
//go:build wintest
package sandboxv2_test
import "os"
func init() {
extraHostExecProviders = append(extraHostExecProviders, agentWinHostExec)
}
func agentWinHostExec() []hostTarget {
var targets []hostTarget
if addr := os.Getenv("TAI_TEST_WIN_HOSTEXEC_LINUX"); addr != "" {
targets = append(targets, hostTarget{Name: "win-linux", Addr: addr})
}
if addr := os.Getenv("TAI_TEST_WIN_HOSTEXEC_NATIVE"); addr != "" {
targets = append(targets, hostTarget{Name: "win-native", Addr: addr})
}
return targets
}

View file

@ -23,6 +23,7 @@ type SandboxConfig struct {
Prepare []PrepareStep `json:"prepare,omitempty" yaml:"prepare,omitempty"` Prepare []PrepareStep `json:"prepare,omitempty" yaml:"prepare,omitempty"`
Environment map[string]string `json:"environment,omitempty" yaml:"environment,omitempty"` Environment map[string]string `json:"environment,omitempty" yaml:"environment,omitempty"`
Secrets map[string]string `json:"secrets,omitempty" yaml:"secrets,omitempty"` Secrets map[string]string `json:"secrets,omitempty" yaml:"secrets,omitempty"`
Filter *ComputerFilter `json:"filter,omitempty" yaml:"filter,omitempty"`
// Populated by the framework at runtime (never serialized). // Populated by the framework at runtime (never serialized).
Owner string `json:"-" yaml:"-"` Owner string `json:"-" yaml:"-"`
@ -33,6 +34,19 @@ type SandboxConfig struct {
WorkspaceID string `json:"-" yaml:"-"` WorkspaceID string `json:"-" yaml:"-"`
} }
// ComputerFilter defines the query parameters for GET /computer/options.
// Declared in DSL sandbox.filter; frontend passes it through to the API.
type ComputerFilter struct {
Kind string `json:"kind,omitempty" yaml:"kind,omitempty"`
Image string `json:"image,omitempty" yaml:"image,omitempty"`
VNC *bool `json:"vnc,omitempty" yaml:"vnc,omitempty"`
OS string `json:"os,omitempty" yaml:"os,omitempty"`
Arch string `json:"arch,omitempty" yaml:"arch,omitempty"`
MinCPUs float64 `json:"min_cpus,omitempty" yaml:"min_cpus,omitempty"`
MinMem string `json:"min_mem,omitempty" yaml:"min_mem,omitempty"`
Labels map[string]string `json:"labels,omitempty" yaml:"labels,omitempty"`
}
// ComputerConfig describes the execution environment (container or host). // ComputerConfig describes the execution environment (container or host).
type ComputerConfig struct { type ComputerConfig struct {
Image string `json:"image,omitempty" yaml:"image,omitempty"` Image string `json:"image,omitempty" yaml:"image,omitempty"`

View file

@ -207,10 +207,16 @@ type AssistantList struct {
// AssistantInfo contains basic assistant information for display // AssistantInfo contains basic assistant information for display
// Used in chat history to show assistant details with i18n support // Used in chat history to show assistant details with i18n support
type AssistantInfo struct { type AssistantInfo struct {
AssistantID string `json:"assistant_id"` AssistantID string `json:"assistant_id"`
Name string `json:"name"` Name string `json:"name"`
Avatar string `json:"avatar,omitempty"` Avatar string `json:"avatar,omitempty"`
Description string `json:"description,omitempty"` Description string `json:"description,omitempty"`
Connector string `json:"connector,omitempty"`
ConnectorOptions *ConnectorOptions `json:"connector_options,omitempty"`
Modes []string `json:"modes,omitempty"`
DefaultMode string `json:"default_mode,omitempty"`
Sandbox bool `json:"sandbox,omitempty"`
ComputerFilter *sandboxTypes.ComputerFilter `json:"computer_filter,omitempty"`
} }
// Tag represents a tag // Tag represents a tag
@ -422,44 +428,46 @@ type ConnectorOptions struct {
// AssistantModel the assistant database model // AssistantModel the assistant database model
type AssistantModel struct { type AssistantModel struct {
ID string `json:"assistant_id"` // Assistant ID ID string `json:"assistant_id"` // Assistant ID
Type string `json:"type,omitempty"` // Assistant Type, default is assistant Type string `json:"type,omitempty"` // Assistant Type, default is assistant
Name string `json:"name,omitempty"` // Assistant Name Name string `json:"name,omitempty"` // Assistant Name
Avatar string `json:"avatar,omitempty"` // Assistant Avatar Avatar string `json:"avatar,omitempty"` // Assistant Avatar
Connector string `json:"connector"` // AI Connector (default connector) Connector string `json:"connector"` // AI Connector (default connector)
ConnectorOptions *ConnectorOptions `json:"connector_options,omitempty"` // Connector selection options for user to choose from ConnectorOptions *ConnectorOptions `json:"connector_options,omitempty"` // Connector selection options for user to choose from
Path string `json:"path,omitempty"` // Assistant Path Path string `json:"path,omitempty"` // Assistant Path
BuiltIn bool `json:"built_in,omitempty"` // Whether this is a built-in assistant BuiltIn bool `json:"built_in,omitempty"` // Whether this is a built-in assistant
Sort int `json:"sort,omitempty"` // Assistant Sort Sort int `json:"sort,omitempty"` // Assistant Sort
Description string `json:"description,omitempty"` // Assistant Description Description string `json:"description,omitempty"` // Assistant Description
Capabilities string `json:"capabilities,omitempty"` // Assistant capabilities description (useful for Robot orchestration) Capabilities string `json:"capabilities,omitempty"` // Assistant capabilities description (useful for Robot orchestration)
Tags []string `json:"tags,omitempty"` // Assistant Tags Tags []string `json:"tags,omitempty"` // Assistant Tags
Modes []string `json:"modes,omitempty"` // Supported modes (e.g., ["task", "chat"]), null means all modes are supported Modes []string `json:"modes,omitempty"` // Supported modes (e.g., ["task", "chat"]), null means all modes are supported
DefaultMode string `json:"default_mode,omitempty"` // Default mode, can be empty DefaultMode string `json:"default_mode,omitempty"` // Default mode, can be empty
Readonly bool `json:"readonly,omitempty"` // Whether this assistant is readonly Readonly bool `json:"readonly,omitempty"` // Whether this assistant is readonly
Public bool `json:"public,omitempty"` // Whether this assistant is shared across all teams in the platform Public bool `json:"public,omitempty"` // Whether this assistant is shared across all teams in the platform
Share string `json:"share,omitempty"` // Assistant sharing scope (private/team) Share string `json:"share,omitempty"` // Assistant sharing scope (private/team)
Mentionable bool `json:"mentionable,omitempty"` // Whether this assistant is mentionable Mentionable bool `json:"mentionable,omitempty"` // Whether this assistant is mentionable
Automated bool `json:"automated,omitempty"` // Whether this assistant is automated Automated bool `json:"automated,omitempty"` // Whether this assistant is automated
Options map[string]interface{} `json:"options,omitempty"` // AI Options Options map[string]interface{} `json:"options,omitempty"` // AI Options
Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts (default prompts) Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts (default prompts)
PromptPresets map[string][]Prompt `json:"prompt_presets,omitempty"` // Prompt presets organized by mode (e.g., "chat", "task", etc.) PromptPresets map[string][]Prompt `json:"prompt_presets,omitempty"` // Prompt presets organized by mode (e.g., "chat", "task", etc.)
DisableGlobalPrompts bool `json:"disable_global_prompts,omitempty"` // Whether to disable global prompts, default is false DisableGlobalPrompts bool `json:"disable_global_prompts,omitempty"` // Whether to disable global prompts, default is false
KB *KnowledgeBase `json:"kb,omitempty"` // Knowledge base configuration KB *KnowledgeBase `json:"kb,omitempty"` // Knowledge base configuration
DB *Database `json:"db,omitempty"` // Database configuration DB *Database `json:"db,omitempty"` // Database configuration
MCP *MCPServers `json:"mcp,omitempty"` // MCP servers configuration MCP *MCPServers `json:"mcp,omitempty"` // MCP servers configuration
Workflow *Workflow `json:"workflow,omitempty"` // Workflow configuration Workflow *Workflow `json:"workflow,omitempty"` // Workflow configuration
Sandbox *Sandbox `json:"sandbox,omitempty"` // Sandbox configuration for coding agents (V1) Sandbox *Sandbox `json:"sandbox,omitempty"` // Sandbox configuration for coding agents (V1)
SandboxV2 *sandboxTypes.SandboxConfig `json:"-"` // V2 sandbox configuration (runtime only, not persisted in DB) SandboxV2 *sandboxTypes.SandboxConfig `json:"-"` // V2 sandbox configuration (runtime only, not persisted in DB)
ConfigHash string `json:"-"` // V2 sandbox config fingerprint for hot-reload IsSandbox bool `json:"-"` // Whether this is a Sandbox assistant (derived from SandboxV2 presence)
Placeholder *Placeholder `json:"placeholder,omitempty"` // Assistant Placeholder ComputerFilter *sandboxTypes.ComputerFilter `json:"-"` // Computer filter from DSL sandbox.filter (runtime only)
Source string `json:"source,omitempty"` // Hook script source code ConfigHash string `json:"-"` // V2 sandbox config fingerprint for hot-reload
Locales i18n.Map `json:"locales,omitempty"` // Assistant Locales Placeholder *Placeholder `json:"placeholder,omitempty"` // Assistant Placeholder
Uses *context.Uses `json:"uses,omitempty"` // Assistant-specific wrapper configurations for vision, audio, etc. If not set, use global settings Source string `json:"source,omitempty"` // Hook script source code
Search *searchTypes.Config `json:"search,omitempty"` // Search configuration (web, kb, db, citation, weights, etc.) Locales i18n.Map `json:"locales,omitempty"` // Assistant Locales
Dependencies map[string]string `json:"dependencies,omitempty"` // Dependencies on other MCP Clients (name -> version constraint) Uses *context.Uses `json:"uses,omitempty"` // Assistant-specific wrapper configurations for vision, audio, etc. If not set, use global settings
CreatedAt int64 `json:"created_at"` // Creation timestamp Search *searchTypes.Config `json:"search,omitempty"` // Search configuration (web, kb, db, citation, weights, etc.)
UpdatedAt int64 `json:"updated_at"` // Last update timestamp Dependencies map[string]string `json:"dependencies,omitempty"` // Dependencies on other MCP Clients (name -> version constraint)
CreatedAt int64 `json:"created_at"` // Creation timestamp
UpdatedAt int64 `json:"updated_at"` // Last update timestamp
// Permission management fields (not exposed in JSON API responses) // Permission management fields (not exposed in JSON API responses)
YaoCreatedBy string `json:"-"` // User who created the assistant (not exposed in JSON) YaoCreatedBy string `json:"-"` // User who created the assistant (not exposed in JSON)

91
cmd/ci-token/main.go Normal file
View file

@ -0,0 +1,91 @@
//go:build ci
package main
import (
"flag"
"fmt"
"os"
"path/filepath"
"time"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/engine"
"github.com/yaoapp/yao/openapi/oauth"
)
func main() {
appPath := flag.String("app", envOr("YAO_CI_APP_PATH", "."), "Yao application directory")
clientID := flag.String("client-id", envOr("YAO_CI_OAUTH_CLIENT_ID", "ci-tai"), "OAuth client ID embedded in token")
subject := flag.String("subject", envOr("YAO_CI_OAUTH_SUBJECT", "ci-tai"), "JWT subject claim")
scope := flag.String("scope", envOr("YAO_CI_OAUTH_SCOPE", "tai:tunnel"), "Token scope (space-separated)")
ttl := flag.String("ttl", envOr("YAO_CI_OAUTH_TTL", "24h"), "Token TTL (e.g. 1h, 24h, 168h)")
userID := flag.String("user-id", envOr("YAO_CI_OAUTH_USER_ID", ""), "User ID claim")
teamID := flag.String("team-id", envOr("YAO_CI_OAUTH_TEAM_ID", ""), "Team ID claim")
flag.Parse()
root, err := filepath.Abs(*appPath)
if err != nil {
fmt.Fprintf(os.Stderr, "ci-token: invalid app path: %v\n", err)
os.Exit(1)
}
if err := os.Chdir(root); err != nil {
fmt.Fprintf(os.Stderr, "ci-token: chdir %s: %v\n", root, err)
os.Exit(1)
}
savedStdout := os.Stdout
os.Stdout, _ = os.Open(os.DevNull)
config.Conf = config.LoadFrom(filepath.Join(root, ".env"))
config.Conf.Root = root
cfg := config.Conf
cfg.Session.IsCLI = true
warnings, err := engine.Load(cfg, engine.LoadOption{Action: "run"})
os.Stdout = savedStdout
if err != nil {
fmt.Fprintf(os.Stderr, "ci-token: engine.Load failed: %v\n", err)
os.Exit(1)
}
for _, w := range warnings {
fmt.Fprintf(os.Stderr, "ci-token: warning [%s]: %v\n", w.Widget, w.Error)
}
if oauth.OAuth == nil {
fmt.Fprintln(os.Stderr, "ci-token: oauth service not initialized (openapi.Load may have failed)")
os.Exit(1)
}
dur, err := time.ParseDuration(*ttl)
if err != nil {
fmt.Fprintf(os.Stderr, "ci-token: invalid --ttl %q: %v\n", *ttl, err)
os.Exit(1)
}
expiresIn := int(dur.Seconds())
extraClaims := map[string]interface{}{}
if *userID != "" {
extraClaims["user_id"] = *userID
}
if *teamID != "" {
extraClaims["team_id"] = *teamID
}
token, err := oauth.OAuth.MakeAccessToken(*clientID, *scope, *subject, expiresIn, extraClaims)
if err != nil {
fmt.Fprintf(os.Stderr, "ci-token: MakeAccessToken failed: %v\n", err)
os.Exit(1)
}
fmt.Print(token)
}
func envOr(key, fallback string) string {
if v := os.Getenv(key); v != "" {
return v
}
return fallback
}

View file

@ -184,8 +184,8 @@ var startCmd = &cobra.Command{
return return
} }
if strings.ToLower(config.Conf.GRPC.Enabled) != "off" { if strings.ToLower(config.Conf.GRPC.Enabled) != "off" {
for _, h := range strings.Split(config.Conf.GRPC.Host, ",") { for _, h := range yaogrpc.ExpandHosts(config.Conf.GRPC.Host) {
if occupied, proc := portOccupied(strings.TrimSpace(h), config.Conf.GRPC.Port); occupied { if occupied, proc := portOccupied(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)) fmt.Println(color.RedString(L("Fatal: gRPC port %d is already in use%s"), config.Conf.GRPC.Port, proc))
return return
} }

View file

@ -1,5 +1,17 @@
package config package config
import "strings"
// HostHasInternal reports whether a comma-separated host string contains "internal".
func HostHasInternal(host string) bool {
for _, h := range strings.Split(host, ",") {
if strings.ToLower(strings.TrimSpace(h)) == "internal" {
return true
}
}
return false
}
// Config 象传应用引擎配置 // Config 象传应用引擎配置
type Config struct { type Config struct {
Mode string `json:"mode,omitempty" env:"YAO_ENV" envDefault:"production"` // The start mode production/development Mode string `json:"mode,omitempty" env:"YAO_ENV" envDefault:"production"` // The start mode production/development
@ -30,6 +42,12 @@ type Config struct {
} }
// GRPCConfig gRPC server configuration // GRPCConfig gRPC server configuration
//
// Host accepts comma-separated bind addresses. Special values:
// - "internal" — 127.0.0.1 + auto-detect all private-network interfaces (10.x, 172.16-31.x, 192.168.x)
// - "localhost" — treated as 127.0.0.1
//
// Example: YAO_GRPC_HOST=127.0.0.1,internal
type GRPCConfig struct { type GRPCConfig struct {
Enabled string `json:"enabled,omitempty" env:"YAO_GRPC"` // Set "off" to disable gRPC server Enabled string `json:"enabled,omitempty" env:"YAO_GRPC"` // Set "off" to disable gRPC server
Host string `json:"host,omitempty" env:"YAO_GRPC_HOST" envDefault:"127.0.0.1"` // Comma-separated bind addresses Host string `json:"host,omitempty" env:"YAO_GRPC_HOST" envDefault:"127.0.0.1"` // Comma-separated bind addresses

View file

@ -63,6 +63,12 @@ func VirtualEndpoint(fullMethod string, req interface{}) (method string, path st
case "/yao.Yao/Heartbeat": case "/yao.Yao/Heartbeat":
return "POST", "/grpc/heartbeat" return "POST", "/grpc/heartbeat"
case "/tai.tunnel.TaiTunnel/Register":
return "POST", "/grpc/tai/register"
case "/tai.tunnel.TaiTunnel/Forward":
return "POST", "/grpc/tai/forward"
default: default:
return "POST", "/grpc/unknown" return "POST", "/grpc/unknown"
} }

View file

@ -15,8 +15,10 @@ import (
) )
const ( const (
healthzMethod = "/yao.Yao/Healthz" healthzMethod = "/yao.Yao/Healthz"
apiMethod = "/yao.Yao/API" apiMethod = "/yao.Yao/API"
taiRegisterMethod = "/tai.tunnel.TaiTunnel/Register"
taiForwardMethod = "/tai.tunnel.TaiTunnel/Forward"
metaAuthorization = "authorization" metaAuthorization = "authorization"
metaRefreshToken = "x-refresh-token" metaRefreshToken = "x-refresh-token"
@ -102,8 +104,8 @@ func authenticate(ctx context.Context, fullMethod string, req interface{}) (cont
)) ))
} }
// ACL scope check — skip for API proxy (the openapi router does its own auth). // ACL scope check — skip for API proxy and Tai tunnel (infrastructure services).
if fullMethod != apiMethod { if fullMethod != apiMethod && fullMethod != taiRegisterMethod && fullMethod != taiForwardMethod {
httpMethod, httpPath := VirtualEndpoint(fullMethod, req) httpMethod, httpPath := VirtualEndpoint(fullMethod, req)
scopes := strings.Fields(result.Info.Scope) scopes := strings.Fields(result.Info.Scope)

View file

@ -11,6 +11,7 @@ import (
"github.com/yaoapp/yao/grpc/pb" "github.com/yaoapp/yao/grpc/pb"
"github.com/yaoapp/yao/grpc/tests/testutils" "github.com/yaoapp/yao/grpc/tests/testutils"
"github.com/yaoapp/yao/tai/tunnel/taipb"
) )
func TestAuth_NoToken_Rejected(t *testing.T) { func TestAuth_NoToken_Rejected(t *testing.T) {
@ -167,3 +168,94 @@ func TestAuth_StreamInterceptor_WrongScope(t *testing.T) {
st, _ := status.FromError(err) st, _ := status.FromError(err)
assert.Equal(t, codes.PermissionDenied, st.Code()) assert.Equal(t, codes.PermissionDenied, st.Code())
} }
// ── TaiTunnel auth tests ──────────────────────────────────────────────────
func TestAuth_TaiTunnel_Register_NoToken(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := taipb.NewTaiTunnelClient(conn)
stream, err := client.Register(context.Background())
if err != nil {
st, _ := status.FromError(err)
assert.Equal(t, codes.Unauthenticated, st.Code())
return
}
_ = stream.Send(&taipb.TunnelControl{Type: "register", NodeId: "n", MachineId: "m"})
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.Unauthenticated, st.Code())
}
func TestAuth_TaiTunnel_Forward_NoToken(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := taipb.NewTaiTunnelClient(conn)
ctx := metadata.AppendToOutgoingContext(context.Background(), "channel_id", "test-ch")
stream, err := client.Forward(ctx)
if err != nil {
st, _ := status.FromError(err)
assert.Equal(t, codes.Unauthenticated, st.Code())
return
}
_ = stream.Send(&taipb.ForwardData{Data: []byte("x")})
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.Unauthenticated, st.Code())
}
func TestAuth_TaiTunnel_Register_ValidToken(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := taipb.NewTaiTunnelClient(conn)
token := testutils.ObtainAccessToken(t, "tai:connect")
ctx := testutils.WithToken(context.Background(), token)
stream, err := client.Register(ctx)
if err != nil {
t.Fatal(err)
}
err = stream.Send(&taipb.TunnelControl{
Type: "register", NodeId: "auth-test-node", MachineId: "auth-test-machine",
})
if err != nil {
t.Fatal(err)
}
resp, err := stream.Recv()
if err != nil {
st, ok := status.FromError(err)
if ok && (st.Code() == codes.Unauthenticated || st.Code() == codes.PermissionDenied) {
t.Fatalf("expected auth to pass, got %v: %v", st.Code(), st.Message())
}
t.Fatal(err)
}
assert.Equal(t, "registered", resp.Type)
assert.NotEmpty(t, resp.TaiId)
stream.CloseSend()
}
func TestAuth_TaiTunnel_Register_ExpiredToken(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := taipb.NewTaiTunnelClient(conn)
token := testutils.ObtainExpiredAccessToken(t, "tai:connect")
ctx := testutils.WithToken(context.Background(), token)
stream, err := client.Register(ctx)
if err != nil {
st, _ := status.FromError(err)
assert.Equal(t, codes.Unauthenticated, st.Code())
return
}
_ = stream.Send(&taipb.TunnelControl{Type: "register", NodeId: "n", MachineId: "m"})
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.Unauthenticated, st.Code())
}

View file

@ -10,5 +10,6 @@ func init() {
&acl.ScopeDefinition{Name: "grpc:mcp", Endpoints: []string{"GET /grpc/mcp/tools", "POST /grpc/mcp/call/*", "POST /grpc/mcp/call/", "GET /grpc/mcp/resources", "GET /grpc/mcp/resources/read", "POST /grpc/heartbeat"}}, &acl.ScopeDefinition{Name: "grpc:mcp", Endpoints: []string{"GET /grpc/mcp/tools", "POST /grpc/mcp/call/*", "POST /grpc/mcp/call/", "GET /grpc/mcp/resources", "GET /grpc/mcp/resources/read", "POST /grpc/heartbeat"}},
&acl.ScopeDefinition{Name: "grpc:llm", Endpoints: []string{"POST /grpc/llm/completions"}}, &acl.ScopeDefinition{Name: "grpc:llm", Endpoints: []string{"POST /grpc/llm/completions"}},
&acl.ScopeDefinition{Name: "grpc:agent", Endpoints: []string{"POST /grpc/agent/*", "POST /grpc/agent/"}}, &acl.ScopeDefinition{Name: "grpc:agent", Endpoints: []string{"POST /grpc/agent/*", "POST /grpc/agent/"}},
&acl.ScopeDefinition{Name: "tai:connect", Endpoints: []string{"POST /grpc/tai/register", "POST /grpc/tai/forward"}},
) )
} }

View file

@ -24,6 +24,9 @@ import (
runhandler "github.com/yaoapp/yao/grpc/run" runhandler "github.com/yaoapp/yao/grpc/run"
sandboxhandler "github.com/yaoapp/yao/grpc/sandbox" sandboxhandler "github.com/yaoapp/yao/grpc/sandbox"
shellhandler "github.com/yaoapp/yao/grpc/shell" shellhandler "github.com/yaoapp/yao/grpc/shell"
"github.com/yaoapp/yao/tai/registry"
"github.com/yaoapp/yao/tai/tunnel"
"github.com/yaoapp/yao/tai/tunnel/taipb"
) )
var ( var (
@ -127,6 +130,7 @@ func SandboxHandler() *sandboxhandler.Handler {
} }
var sandboxH *sandboxhandler.Handler var sandboxH *sandboxhandler.Handler
var tunnelH *tunnel.TunnelHandler
// SetSandboxOnBeat sets the heartbeat callback for the sandbox handler. // SetSandboxOnBeat sets the heartbeat callback for the sandbox handler.
// Must be called before StartServer. // Must be called before StartServer.
@ -156,11 +160,16 @@ func StartServer(cfg config.Config) error {
} }
pb.RegisterYaoServer(server, &yaoServer{sandbox: sandboxH}) pb.RegisterYaoServer(server, &yaoServer{sandbox: sandboxH})
hosts := strings.Split(cfg.GRPC.Host, ",") if reg := registry.Global(); reg != nil {
tunnelH = tunnel.NewTunnelHandler(reg)
taipb.RegisterTaiTunnelServer(server, tunnelH)
}
hosts := ExpandHosts(cfg.GRPC.Host)
port := strconv.Itoa(cfg.GRPC.Port) port := strconv.Itoa(cfg.GRPC.Port)
for _, h := range hosts { for _, h := range hosts {
addr := net.JoinHostPort(strings.TrimSpace(h), port) addr := net.JoinHostPort(h, port)
lis, err := net.Listen("tcp", addr) lis, err := net.Listen("tcp", addr)
if err != nil { if err != nil {
stopLocked() stopLocked()
@ -224,6 +233,13 @@ func GRPCServer() *grpc.Server {
return server return server
} }
// TunnelHandler returns the gRPC tunnel handler for forward requests.
func TunnelHandler() *tunnel.TunnelHandler {
mu.Lock()
defer mu.Unlock()
return tunnelH
}
// Addr returns all addresses the gRPC server is listening on. // Addr returns all addresses the gRPC server is listening on.
func Addr() []string { func Addr() []string {
mu.Lock() mu.Lock()
@ -232,3 +248,84 @@ func Addr() []string {
copy(result, addrs) copy(result, addrs)
return result return result
} }
// expandHosts parses comma-separated host entries, expanding special values:
// - "internal" → 127.0.0.1 + all private-network IPv4 addresses (10.x, 172.16-31.x, 192.168.x)
// - "localhost" → 127.0.0.1
//
// Duplicates are removed.
func ExpandHosts(raw string) []string {
seen := map[string]bool{}
var result []string
for _, h := range strings.Split(raw, ",") {
h = strings.TrimSpace(h)
if h == "" {
continue
}
switch strings.ToLower(h) {
case "localhost":
h = "127.0.0.1"
if !seen[h] {
seen[h] = true
result = append(result, h)
}
case "internal":
if !seen["127.0.0.1"] {
seen["127.0.0.1"] = true
result = append(result, "127.0.0.1")
}
for _, ip := range InternalIPs() {
if !seen[ip] {
seen[ip] = true
result = append(result, ip)
}
}
default:
if !seen[h] {
seen[h] = true
result = append(result, h)
}
}
}
return result
}
// InternalIPs returns all IPv4 addresses on private-network interfaces
// (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16).
func InternalIPs() []string {
var ips []string
ifaces, err := net.Interfaces()
if err != nil {
return nil
}
for _, iface := range ifaces {
if iface.Flags&net.FlagUp == 0 || iface.Flags&net.FlagLoopback != 0 {
continue
}
addrs, err := iface.Addrs()
if err != nil {
continue
}
for _, a := range addrs {
ipNet, ok := a.(*net.IPNet)
if !ok {
continue
}
ip := ipNet.IP.To4()
if ip == nil {
continue
}
if isPrivateIP(ip) {
ips = append(ips, ip.String())
}
}
}
return ips
}
func isPrivateIP(ip net.IP) bool {
return ip[0] == 10 ||
(ip[0] == 172 && ip[1] >= 16 && ip[1] <= 31) ||
(ip[0] == 192 && ip[1] == 168)
}

View file

@ -28,6 +28,7 @@ import (
"github.com/yaoapp/yao/openapi" "github.com/yaoapp/yao/openapi"
"github.com/yaoapp/yao/openapi/oauth" "github.com/yaoapp/yao/openapi/oauth"
"github.com/yaoapp/yao/service" "github.com/yaoapp/yao/service"
"github.com/yaoapp/yao/tai/registry"
"github.com/yaoapp/yao/test" "github.com/yaoapp/yao/test"
_ "github.com/yaoapp/gou/encoding" _ "github.com/yaoapp/gou/encoding"
@ -97,6 +98,10 @@ func Prepare(t *testing.T) *grpc.ClientConn {
service.Router = router service.Router = router
} }
if registry.Global() == nil {
registry.SetGlobalForTest(registry.NewForTest())
}
if err := yaogrpc.StartServer(cfg); err != nil { if err := yaogrpc.StartServer(cfg); err != nil {
t.Fatalf("failed to start gRPC server: %v", err) t.Fatalf("failed to start gRPC server: %v", err)
} }

View file

@ -8,7 +8,7 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/yaoapp/kun/log" "github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/agent" "github.com/yaoapp/yao/agent"
"github.com/yaoapp/yao/agent/assistant" assistantPkg "github.com/yaoapp/yao/agent/assistant"
agenttypes "github.com/yaoapp/yao/agent/store/types" agenttypes "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/openapi/oauth/authorized" "github.com/yaoapp/yao/openapi/oauth/authorized"
"github.com/yaoapp/yao/openapi/oauth/types" "github.com/yaoapp/yao/openapi/oauth/types"
@ -415,13 +415,13 @@ func CreateAssistant(c *gin.Context) {
assistantData["assistant_id"] = id assistantData["assistant_id"] = id
// Clear cache and reload assistant to make it effective // Clear cache and reload assistant to make it effective
cache := assistant.GetCache() cache := assistantPkg.GetCache()
if cache != nil { if cache != nil {
cache.Remove(id) cache.Remove(id)
} }
// Reload the assistant to ensure it's available in cache with updated data // Reload the assistant to ensure it's available in cache with updated data
_, err = assistant.Get(id) _, err = assistantPkg.Get(id)
if err != nil { if err != nil {
// Just log the error, don't fail the request // Just log the error, don't fail the request
log.Error("Error reloading assistant %s: %v", id, err) log.Error("Error reloading assistant %s: %v", id, err)
@ -520,13 +520,13 @@ func UpdateAssistant(c *gin.Context) {
} }
// Clear cache and reload assistant to make it effective // Clear cache and reload assistant to make it effective
cache := assistant.GetCache() cache := assistantPkg.GetCache()
if cache != nil { if cache != nil {
cache.Remove(assistantID) cache.Remove(assistantID)
} }
// Reload the assistant to ensure it's available in cache with updated data // Reload the assistant to ensure it's available in cache with updated data
_, err = assistant.Get(assistantID) _, err = assistantPkg.Get(assistantID)
if err != nil { if err != nil {
// Just log the error, don't fail the request // Just log the error, don't fail the request
log.Error("Error reloading assistant %s: %v", assistantID, err) log.Error("Error reloading assistant %s: %v", assistantID, err)
@ -539,24 +539,10 @@ func UpdateAssistant(c *gin.Context) {
} }
// GetAssistantInfo retrieves essential assistant information for InputArea component // GetAssistantInfo retrieves essential assistant information for InputArea component
// Returns only the fields needed for UI display: id, name, avatar, description, connector, connector_options, modes, default_mode
func GetAssistantInfo(c *gin.Context) { func GetAssistantInfo(c *gin.Context) {
// Get authorized information
authInfo := authorized.GetInfo(c) authInfo := authorized.GetInfo(c)
// Get Agent instance from global variable
agentInstance := agent.GetAgent()
if agentInstance == nil || agentInstance.Store == nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Agent store not initialized",
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Get assistant ID from URL parameter
assistantID := c.Param("id") assistantID := c.Param("id")
if assistantID == "" { if assistantID == "" {
errorResp := &response.ErrorResponse{ errorResp := &response.ErrorResponse{
@ -567,37 +553,11 @@ func GetAssistantInfo(c *gin.Context) {
return return
} }
// Parse locale (optional - defaults to "en-us")
locale := "en-us" locale := "en-us"
if loc := c.Query("locale"); loc != "" { if loc := c.Query("locale"); loc != "" {
locale = strings.ToLower(strings.TrimSpace(loc)) locale = strings.ToLower(strings.TrimSpace(loc))
} }
// Define fields needed for InputArea
infoFields := []string{
"assistant_id",
"name",
"avatar",
"description",
"connector",
"connector_options",
"modes",
"default_mode",
}
// Get assistant with specific fields and locale
assistant, err := agentInstance.Store.GetAssistant(assistantID, infoFields, locale)
if err != nil {
log.Error("Failed to get assistant info %s: %v", assistantID, err)
errorResp := &response.ErrorResponse{
Code: response.ErrInvalidRequest.Code,
ErrorDescription: "Assistant not found: " + err.Error(),
}
response.RespondWithError(c, response.StatusNotFound, errorResp)
return
}
// Check read permission (same as GetAssistant)
hasPermission, err := checkAssistantPermission(authInfo, assistantID, true) hasPermission, err := checkAssistantPermission(authInfo, assistantID, true)
if err != nil { if err != nil {
log.Error("Failed to check permission for assistant %s: %v", assistantID, err) log.Error("Failed to check permission for assistant %s: %v", assistantID, err)
@ -618,28 +578,18 @@ func GetAssistantInfo(c *gin.Context) {
return return
} }
// Build response with only the required fields ast, err := assistantPkg.Get(assistantID)
infoResponse := map[string]interface{}{ if err != nil || ast == nil {
"assistant_id": assistant.ID, log.Error("Failed to get assistant info %s: %v", assistantID, err)
"name": assistant.Name, errorResp := &response.ErrorResponse{
"avatar": assistant.Avatar, Code: response.ErrInvalidRequest.Code,
"description": assistant.Description, ErrorDescription: "Assistant not found",
"connector": assistant.Connector, }
response.RespondWithError(c, response.StatusNotFound, errorResp)
return
} }
// Add optional fields if they exist response.RespondWithSuccess(c, response.StatusOK, ast.GetInfo(locale))
if assistant.ConnectorOptions != nil {
infoResponse["connector_options"] = assistant.ConnectorOptions
}
if len(assistant.Modes) > 0 {
infoResponse["modes"] = assistant.Modes
}
if assistant.DefaultMode != "" {
infoResponse["default_mode"] = assistant.DefaultMode
}
// Return the result with standard response format
response.RespondWithSuccess(c, response.StatusOK, infoResponse)
} }
// checkAssistantPermission checks if the user has permission to access the assistant // checkAssistantPermission checks if the user has permission to access the assistant

View file

@ -0,0 +1,345 @@
package computer
import (
"context"
"net/http"
"strconv"
"strings"
"github.com/gin-gonic/gin"
sandboxv2 "github.com/yaoapp/yao/sandbox/v2"
"github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/registry"
taitypes "github.com/yaoapp/yao/tai/types"
"github.com/yaoapp/yao/openapi/oauth/authorized"
oauthTypes "github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/openapi/response"
)
// Attach registers computer option routes on the given group.
// - GET /options — list available computers (filtered by ComputerFilter query params)
func Attach(group *gin.RouterGroup, oauth oauthTypes.OAuth) {
group.Use(oauth.Guard)
group.GET("/options", handleOptions)
}
type computerSystemInfo struct {
OS string `json:"os"`
Arch string `json:"arch"`
Hostname string `json:"hostname"`
NumCPU int `json:"num_cpu"`
TotalMem int64 `json:"total_mem,omitempty"`
}
type computerOption struct {
Kind string `json:"kind"`
ID string `json:"id"`
DisplayName string `json:"display_name"`
NodeID string `json:"node_id"`
Status string `json:"status"`
Mode string `json:"mode,omitempty"`
Addr string `json:"addr,omitempty"`
Image string `json:"image,omitempty"`
Policy string `json:"policy,omitempty"`
VNC bool `json:"vnc"`
Labels map[string]string `json:"labels,omitempty"`
System computerSystemInfo `json:"system"`
}
func handleOptions(c *gin.Context) {
authInfo := authorized.GetInfo(c)
kindFilter := c.Query("kind")
imageFilter := c.Query("image")
osFilter := c.Query("os")
archFilter := c.Query("arch")
var vncFilter *bool
if v := c.Query("vnc"); v != "" {
b, _ := strconv.ParseBool(v)
vncFilter = &b
}
var minCPUs float64
if v := c.Query("min_cpus"); v != "" {
minCPUs, _ = strconv.ParseFloat(v, 64)
}
var minMem int64
if v := c.Query("min_mem"); v != "" {
minMem = parseMemString(v)
}
var result []computerOption
reg := registry.Global()
if reg == nil {
response.RespondWithSuccess(c, http.StatusOK, []computerOption{})
return
}
snaps := reg.List()
// Host entries: nodes with host_exec capability
if kindFilter == "" || kindFilter == "host" {
for i := range snaps {
s := &snaps[i]
if !nodeOwnedBy(s, authInfo) {
continue
}
if !s.Capabilities.HostExec {
continue
}
if !matchNodeFilter(s, osFilter, archFilter, minCPUs, minMem) {
continue
}
result = append(result, nodeToHostOption(*s))
}
}
// Node entries: nodes with container runtime capability
if kindFilter == "" || kindFilter == "node" {
for i := range snaps {
s := &snaps[i]
if !nodeOwnedBy(s, authInfo) {
continue
}
hasRuntime := s.Capabilities.Docker || s.Capabilities.K8s
if !hasRuntime {
continue
}
if !matchNodeFilter(s, osFilter, archFilter, minCPUs, minMem) {
continue
}
result = append(result, nodeToNodeOption(*s))
}
}
// Box entries: persistent/longrunning boxes only
if kindFilter == "" || kindFilter == "box" {
if mgr := getManager(); mgr != nil {
owner := resolveOwner(authInfo)
boxes, err := mgr.List(context.Background(), sandboxv2.ListOptions{})
if err == nil {
for _, b := range boxes {
snap := b.Snapshot()
if snap.Owner != owner {
continue
}
if snap.Policy != sandboxv2.Persistent && snap.Policy != sandboxv2.LongRunning {
continue
}
if imageFilter != "" && snap.Image != imageFilter {
continue
}
if vncFilter != nil && snap.VNC != *vncFilter {
continue
}
result = append(result, boxToOption(b))
}
}
}
}
if result == nil {
result = []computerOption{}
}
response.RespondWithSuccess(c, http.StatusOK, result)
}
func matchNodeFilter(s *taitypes.NodeMeta, osFilter, archFilter string, minCPUs float64, minMem int64) bool {
if osFilter != "" && !strings.EqualFold(s.System.OS, osFilter) {
return false
}
if archFilter != "" && !strings.EqualFold(s.System.Arch, archFilter) {
return false
}
if minCPUs > 0 && float64(s.System.NumCPU) < minCPUs {
return false
}
if minMem > 0 && s.System.TotalMem < minMem {
return false
}
return true
}
func nodeToHostOption(s taitypes.NodeMeta) computerOption {
displayName := s.DisplayName
if displayName == "" {
displayName = s.System.Hostname
}
if displayName == "" {
displayName = s.TaiID
}
status := "stopped"
if s.Status == "online" {
status = "running"
}
addr := s.Addr
if addr == "" {
scheme := s.Mode
if scheme == "" {
scheme = "tai"
}
addr = scheme + "://" + s.TaiID
}
return computerOption{
Kind: "host",
ID: s.TaiID,
DisplayName: displayName,
NodeID: s.TaiID,
Status: status,
Mode: s.Mode,
Addr: addr,
System: computerSystemInfo{
OS: s.System.OS,
Arch: s.System.Arch,
Hostname: s.System.Hostname,
NumCPU: s.System.NumCPU,
TotalMem: s.System.TotalMem,
},
}
}
func nodeToNodeOption(s taitypes.NodeMeta) computerOption {
displayName := s.DisplayName
if displayName == "" {
displayName = s.System.Hostname
}
if displayName == "" {
displayName = s.TaiID
}
status := "stopped"
if s.Status == "online" {
status = "running"
}
addr := s.Addr
if addr == "" {
scheme := s.Mode
if scheme == "" {
scheme = "tai"
}
addr = scheme + "://" + s.TaiID
}
return computerOption{
Kind: "node",
ID: s.TaiID,
DisplayName: displayName,
NodeID: s.TaiID,
Status: status,
Mode: s.Mode,
Addr: addr,
System: computerSystemInfo{
OS: s.System.OS,
Arch: s.System.Arch,
Hostname: s.System.Hostname,
NumCPU: s.System.NumCPU,
TotalMem: s.System.TotalMem,
},
}
}
func boxToOption(b *sandboxv2.Box) computerOption {
snap := b.Snapshot()
info := b.ComputerInfo()
displayName := info.System.Hostname
if displayName == "" {
displayName = snap.ID
}
var mode, addr string
if ns, ok := tai.GetNodeMeta(snap.NodeID); ok {
mode = ns.Mode
addr = ns.Addr
}
if addr == "" && snap.NodeID != "" {
scheme := mode
if scheme == "" {
scheme = "local"
}
addr = scheme + "://" + snap.NodeID
}
return computerOption{
Kind: "box",
ID: snap.ID,
DisplayName: displayName,
NodeID: snap.NodeID,
Status: snap.Status,
Mode: mode,
Addr: addr,
Image: snap.Image,
Policy: string(snap.Policy),
VNC: snap.VNC,
Labels: snap.Labels,
System: computerSystemInfo{
OS: info.System.OS,
Arch: info.System.Arch,
Hostname: info.System.Hostname,
NumCPU: info.System.NumCPU,
TotalMem: info.System.TotalMem,
},
}
}
func nodeOwnedBy(snap *taitypes.NodeMeta, authInfo *oauthTypes.AuthorizedInfo) bool {
if authInfo == nil {
return true
}
if authInfo.TeamID != "" {
return snap.Auth.TeamID == authInfo.TeamID
}
if authInfo.UserID != "" {
return snap.Auth.TeamID == "" && snap.Auth.UserID == authInfo.UserID
}
return true
}
func resolveOwner(authInfo *oauthTypes.AuthorizedInfo) string {
if authInfo != nil && authInfo.TeamID != "" {
return authInfo.TeamID
}
if authInfo != nil {
return authInfo.UserID
}
return ""
}
func getManager() *sandboxv2.Manager {
defer func() { recover() }()
return sandboxv2.M()
}
func parseMemString(s string) int64 {
s = strings.TrimSpace(strings.ToLower(s))
if s == "" {
return 0
}
multiplier := int64(1)
switch {
case strings.HasSuffix(s, "g"):
multiplier = 1024 * 1024 * 1024
s = strings.TrimSuffix(s, "g")
case strings.HasSuffix(s, "m"):
multiplier = 1024 * 1024
s = strings.TrimSuffix(s, "m")
case strings.HasSuffix(s, "k"):
multiplier = 1024
s = strings.TrimSuffix(s, "k")
}
val, err := strconv.ParseFloat(s, 64)
if err != nil {
return 0
}
return int64(val * float64(multiplier))
}

View file

@ -9,6 +9,7 @@ import (
"github.com/yaoapp/yao/openapi/oauth/types" "github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/openapi/response" "github.com/yaoapp/yao/openapi/response"
"github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/registry"
taitypes "github.com/yaoapp/yao/tai/types"
) )
// Attach registers Tai node endpoints on the given group. // Attach registers Tai node endpoints on the given group.
@ -44,7 +45,7 @@ type systemResponse struct {
Shell string `json:"shell,omitempty"` Shell string `json:"shell,omitempty"`
} }
func snapToResponse(s registry.NodeSnapshot) nodeResponse { func snapToResponse(s taitypes.NodeMeta) nodeResponse {
r := nodeResponse{ r := nodeResponse{
TaiID: s.TaiID, TaiID: s.TaiID,
MachineID: s.MachineID, MachineID: s.MachineID,
@ -53,8 +54,8 @@ func snapToResponse(s registry.NodeSnapshot) nodeResponse {
Mode: s.Mode, Mode: s.Mode,
Addr: s.Addr, Addr: s.Addr,
Status: s.Status, Status: s.Status,
Capabilities: s.Capabilities, Capabilities: map[string]bool{"docker": s.Capabilities.Docker, "k8s": s.Capabilities.K8s, "host_exec": s.Capabilities.HostExec},
Ports: s.Ports, Ports: map[string]int{"grpc": s.Ports.GRPC, "http": s.Ports.HTTP, "vnc": s.Ports.VNC, "docker": s.Ports.Docker, "k8s": s.Ports.K8s},
System: systemResponse{ System: systemResponse{
OS: s.System.OS, OS: s.System.OS,
Arch: s.System.Arch, Arch: s.System.Arch,
@ -75,7 +76,7 @@ func snapToResponse(s registry.NodeSnapshot) nodeResponse {
// nodeOwnedBy checks whether a node belongs to the caller. // nodeOwnedBy checks whether a node belongs to the caller.
// TeamID match → true; no team and UserID match → true. // TeamID match → true; no team and UserID match → true.
func nodeOwnedBy(snap *registry.NodeSnapshot, authInfo *types.AuthorizedInfo) bool { func nodeOwnedBy(snap *taitypes.NodeMeta, authInfo *types.AuthorizedInfo) bool {
if authInfo == nil { if authInfo == nil {
return true return true
} }
@ -98,7 +99,7 @@ func handleList(c *gin.Context) {
authInfo := authorized.GetInfo(c) authInfo := authorized.GetInfo(c)
var snaps []registry.NodeSnapshot var snaps []taitypes.NodeMeta
if authInfo != nil && authInfo.TeamID != "" { if authInfo != nil && authInfo.TeamID != "" {
snaps = reg.ListByTeam(authInfo.TeamID) snaps = reg.ListByTeam(authInfo.TeamID)
} else if authInfo != nil && authInfo.UserID != "" { } else if authInfo != nil && authInfo.UserID != "" {

View file

@ -208,10 +208,14 @@ func (openapi *OpenAPI) handleStandardTokenGrant(c *gin.Context, grantType strin
} }
case types.GrantTypeClientCredentials: case types.GrantTypeClientCredentials:
// No code needed for client credentials
code = "" code = ""
// Validate that client supports client credentials grant // RFC 6749 §4.4: client_credentials requires confidential client
if clientInfo.ClientType == types.ClientTypePublic {
response.RespondWithSecureError(c, response.StatusUnauthorized, response.ErrUnauthorizedClient)
return
}
if !openapi.clientSupportsGrantType(clientInfo, types.GrantTypeClientCredentials) { if !openapi.clientSupportsGrantType(clientInfo, types.GrantTypeClientCredentials) {
response.RespondWithSecureError(c, response.StatusUnauthorized, response.ErrUnauthorizedClient) response.RespondWithSecureError(c, response.StatusUnauthorized, response.ErrUnauthorizedClient)
return return

View file

@ -10,6 +10,7 @@ import (
"github.com/yaoapp/yao/openapi/app" "github.com/yaoapp/yao/openapi/app"
"github.com/yaoapp/yao/openapi/captcha" "github.com/yaoapp/yao/openapi/captcha"
"github.com/yaoapp/yao/openapi/chat" "github.com/yaoapp/yao/openapi/chat"
openapiComputer "github.com/yaoapp/yao/openapi/computer"
"github.com/yaoapp/yao/openapi/dsl" "github.com/yaoapp/yao/openapi/dsl"
"github.com/yaoapp/yao/openapi/file" "github.com/yaoapp/yao/openapi/file"
"github.com/yaoapp/yao/openapi/hello" "github.com/yaoapp/yao/openapi/hello"
@ -181,17 +182,18 @@ func (openapi *OpenAPI) Attach(router *gin.Engine) {
sandbox.Attach(sandboxGroup, openapi.OAuth) sandbox.Attach(sandboxGroup, openapi.OAuth)
sandbox.AttachManage(sandboxGroup) sandbox.AttachManage(sandboxGroup)
// Computer option handlers (for InputArea selector)
openapiComputer.Attach(group.Group("/computer"), openapi.OAuth)
// Workspace handlers // Workspace handlers
openapiWorkspace.Attach(group.Group("/workspace"), openapi.OAuth) openapiWorkspace.Attach(group.Group("/workspace"), openapi.OAuth)
// Tai nodes handlers // Tai nodes handlers
nodes.Attach(group.Group("/nodes"), openapi.OAuth) nodes.Attach(group.Group("/nodes"), openapi.OAuth)
// Tai tunnel WebSocket and reverse proxy routes // Tai tunnel: gRPC Forward-based HTTP/VNC transparent proxy
group.GET("/ws/tai", taitunnel.HandleControl) group.Any("/tai/:taiID/proxy/*path", taitunnel.HandleForwardLazy)
group.GET("/ws/tai/data/:channel_id", taitunnel.HandleData) group.Any("/tai/:taiID/vnc/*path", taitunnel.HandleForwardLazy)
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/) // 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/register", taiapi.HandleRegister)

View file

@ -13,6 +13,7 @@ import (
sandboxv2 "github.com/yaoapp/yao/sandbox/v2" sandboxv2 "github.com/yaoapp/yao/sandbox/v2"
"github.com/yaoapp/yao/tai" "github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/registry"
taitypes "github.com/yaoapp/yao/tai/types"
) )
// AttachManage registers sandbox management CRUD routes on the given group. // AttachManage registers sandbox management CRUD routes on the given group.
@ -115,7 +116,7 @@ func boxToResponse(b *sandboxv2.Box) sandboxResponse {
} }
var mode, addr string var mode, addr string
if ns, ok := tai.GetNodeSnapshot(snap.NodeID); ok { if ns, ok := tai.GetNodeMeta(snap.NodeID); ok {
mode = ns.Mode mode = ns.Mode
addr = ns.Addr addr = ns.Addr
} }
@ -157,7 +158,7 @@ func boxToResponse(b *sandboxv2.Box) sandboxResponse {
} }
} }
func hostToResponse(s registry.NodeSnapshot) sandboxResponse { func hostToResponse(s taitypes.NodeMeta) sandboxResponse {
displayName := s.DisplayName displayName := s.DisplayName
if displayName == "" { if displayName == "" {
displayName = s.System.Hostname displayName = s.System.Hostname
@ -209,7 +210,7 @@ func hostToResponse(s registry.NodeSnapshot) sandboxResponse {
} }
} }
func nodeOwnedBy(snap *registry.NodeSnapshot, authInfo *types.AuthorizedInfo) bool { func nodeOwnedBy(snap *taitypes.NodeMeta, authInfo *types.AuthorizedInfo) bool {
if authInfo == nil { if authInfo == nil {
return true return true
} }
@ -257,7 +258,7 @@ func handleList(c *gin.Context) {
if !nodeOwnedBy(s, authInfo) { if !nodeOwnedBy(s, authInfo) {
continue continue
} }
if !s.Capabilities["host_exec"] { if !s.Capabilities.HostExec {
continue continue
} }
if nodeFilter != "" && s.TaiID != nodeFilter { if nodeFilter != "" && s.TaiID != nodeFilter {

View file

@ -104,7 +104,11 @@ func resolveServerURL(issuerURL string) string {
} }
// resolveGRPCAddr returns the gRPC server address for client discovery. // resolveGRPCAddr returns the gRPC server address for client discovery.
// Uses the request Host's IP with the configured gRPC port. //
// When the listen host includes "internal", "0.0.0.0", or multiple addresses,
// the returned address uses the IP from the incoming HTTP request — if the
// client could reach Yao's HTTP port via that IP, gRPC on the same IP should
// also be reachable. "localhost" is treated as "127.0.0.1".
func resolveGRPCAddr(c *gin.Context) string { func resolveGRPCAddr(c *gin.Context) string {
cfg := config.Conf.GRPC cfg := config.Conf.GRPC
if strings.ToLower(cfg.Enabled) == "off" { if strings.ToLower(cfg.Enabled) == "off" {
@ -116,15 +120,22 @@ func resolveGRPCAddr(c *gin.Context) string {
} }
host := cfg.Host host := cfg.Host
if host == "" || host == "0.0.0.0" { useRequestIP := host == "" || host == "0.0.0.0" ||
strings.Contains(host, ",") ||
config.HostHasInternal(host)
if useRequestIP {
reqHost := c.Request.Host reqHost := c.Request.Host
h, _, err := net.SplitHostPort(reqHost) h, _, err := net.SplitHostPort(reqHost)
if err != nil { if err != nil {
h = reqHost h = reqHost
} }
if strings.ToLower(h) == "localhost" {
h = "127.0.0.1"
}
host = h host = h
} else if strings.Contains(host, ",") { } else if strings.ToLower(strings.TrimSpace(host)) == "localhost" {
host = strings.TrimSpace(strings.Split(host, ",")[0]) host = "127.0.0.1"
} }
return fmt.Sprintf("%s:%s", host, strconv.Itoa(port)) return fmt.Sprintf("%s:%s", host, strconv.Itoa(port))

View file

@ -32,6 +32,7 @@ func Attach(group *gin.RouterGroup, oauth types.OAuth) {
group.Use(oauth.Guard) group.Use(oauth.Guard)
group.GET("", handleList) group.GET("", handleList)
group.GET("/options", handleOptions)
group.POST("", handleCreate) group.POST("", handleCreate)
group.GET("/:id", handleGet) group.GET("/:id", handleGet)
group.PUT("/:id", handleUpdate) group.PUT("/:id", handleUpdate)
@ -170,6 +171,35 @@ func handleList(c *gin.Context) {
response.RespondWithSuccess(c, http.StatusOK, result) response.RespondWithSuccess(c, http.StatusOK, result)
} }
// handleOptions returns workspace options for the InputArea selector.
// Reuses the same logic as handleList (Manager.List with owner+node filter).
// Separated as a dedicated endpoint for clear API responsibility boundary.
func handleOptions(c *gin.Context) {
m := mgr()
if m == nil {
response.RespondWithSuccess(c, http.StatusOK, []workspaceResponse{})
return
}
authInfo := authorized.GetInfo(c)
owner := resolveOwner(authInfo)
list, err := m.List(context.Background(), ws.ListOptions{
Owner: owner,
Node: c.Query("node"),
})
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
return
}
result := make([]workspaceResponse, 0, len(list))
for _, w := range list {
result = append(result, toResponse(w))
}
response.RespondWithSuccess(c, http.StatusOK, result)
}
func handleCreate(c *gin.Context) { func handleCreate(c *gin.Context) {
m := mgr() m := mgr()
if m == nil { if m == nil {

View file

@ -7,7 +7,6 @@ import (
"time" "time"
sandbox "github.com/yaoapp/yao/sandbox/v2" sandbox "github.com/yaoapp/yao/sandbox/v2"
"github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/registry"
) )
@ -225,15 +224,12 @@ func BenchmarkWorkspaceReadWrite(b *testing.B) {
func setupManagerForBench(b *testing.B, pc *nodeConfig) *sandbox.Manager { func setupManagerForBench(b *testing.B, pc *nodeConfig) *sandbox.Manager {
b.Helper() b.Helper()
reg := registry.Global() if registry.Global() == nil {
if reg == nil {
registry.Init(nil) registry.Init(nil)
} }
client, err := tai.New(pc.Addr, pc.Options...) taiID, res := registerForTest(b, pc.Addr, pc.DialOps...)
if err != nil { pc.TaiID = taiID
b.Fatalf("tai.New(%s): %v", pc.Addr, err) b.Cleanup(func() { res.Close() })
}
pc.TaiID = client.TaiID()
sandbox.Init() sandbox.Init()
m := sandbox.M() m := sandbox.M()
b.Cleanup(func() { m.Close() }) b.Cleanup(func() { m.Close() })

View file

@ -7,8 +7,8 @@ import (
"time" "time"
"github.com/yaoapp/yao/tai/proxy" "github.com/yaoapp/yao/tai/proxy"
taisandbox "github.com/yaoapp/yao/tai/sandbox" tairuntime "github.com/yaoapp/yao/tai/runtime"
"github.com/yaoapp/yao/tai/workspace" taiworkspace "github.com/yaoapp/yao/tai/workspace"
) )
// Box represents a single sandbox instance. // Box represents a single sandbox instance.
@ -30,7 +30,7 @@ type Box struct {
image string image string
workspaceID string workspaceID string
system SystemInfo system SystemInfo
ws workspace.FS ws taiworkspace.FS
manager *Manager manager *Manager
} }
@ -69,7 +69,7 @@ func (b *Box) BindWorkplace(workspaceID string) {
// Workplace returns the workspace FS bound to this Box. // Workplace returns the workspace FS bound to this Box.
// If a workspace was bound via CreateOptions.WorkspaceID or BindWorkplace(), // If a workspace was bound via CreateOptions.WorkspaceID or BindWorkplace(),
// returns that workspace's FS. Otherwise returns nil. // returns that workspace's FS. Otherwise returns nil.
func (b *Box) Workplace() workspace.FS { func (b *Box) Workplace() taiworkspace.FS {
return b.Workspace() return b.Workspace()
} }
@ -81,12 +81,12 @@ func (b *Box) Exec(ctx context.Context, cmd []string, opts ...ExecOption) (*Exec
o(cfg) o(cfg)
} }
client, err := b.manager.getNode(b.nodeID) res, err := b.manager.getNode(b.nodeID)
if err != nil { if err != nil {
return nil, err return nil, err
} }
result, err := client.Sandbox().Exec(ctx, b.containerID, cmd, taisandbox.ExecOptions{ result, err := res.Runtime.Exec(ctx, b.containerID, cmd, tairuntime.ExecOptions{
WorkDir: cfg.WorkDir, WorkDir: cfg.WorkDir,
Env: cfg.Env, Env: cfg.Env,
}) })
@ -111,12 +111,12 @@ func (b *Box) Stream(ctx context.Context, cmd []string, opts ...ExecOption) (*Ex
o(cfg) o(cfg)
} }
client, err := b.manager.getNode(b.nodeID) res, err := b.manager.getNode(b.nodeID)
if err != nil { if err != nil {
return nil, err return nil, err
} }
handle, err := client.Sandbox().ExecStream(ctx, b.containerID, cmd, taisandbox.ExecOptions{ handle, err := res.Runtime.ExecStream(ctx, b.containerID, cmd, tairuntime.ExecOptions{
WorkDir: cfg.WorkDir, WorkDir: cfg.WorkDir,
Env: cfg.Env, Env: cfg.Env,
}) })
@ -141,12 +141,12 @@ func (b *Box) Attach(ctx context.Context, port int, opts ...AttachOption) (*Serv
o(cfg) o(cfg)
} }
client, err := b.manager.getNode(b.nodeID) res, err := b.manager.getNode(b.nodeID)
if err != nil { if err != nil {
return nil, err return nil, err
} }
conn, err := client.Proxy().Connect(ctx, b.containerID, proxy.ConnectOptions{ conn, err := res.Proxy.Connect(ctx, b.containerID, proxy.ConnectOptions{
Port: port, Port: port,
Path: cfg.Path, Path: cfg.Path,
Protocol: cfg.Protocol, Protocol: cfg.Protocol,
@ -178,7 +178,7 @@ func (b *Box) Attach(ctx context.Context, port int, opts ...AttachOption) (*Serv
// Workspace returns an fs.FS-compatible filesystem for this sandbox. // Workspace returns an fs.FS-compatible filesystem for this sandbox.
// If a workspace is mounted (WorkspaceID set), uses the workspace ID as session; // If a workspace is mounted (WorkspaceID set), uses the workspace ID as session;
// otherwise falls back to the sandbox ID (backward compatible). // otherwise falls back to the sandbox ID (backward compatible).
func (b *Box) Workspace() workspace.FS { func (b *Box) Workspace() taiworkspace.FS {
b.touch() b.touch()
if b.ws != nil { if b.ws != nil {
return b.ws return b.ws
@ -187,11 +187,11 @@ func (b *Box) Workspace() workspace.FS {
if sessionID == "" { if sessionID == "" {
sessionID = b.id sessionID = b.id
} }
client, err := b.manager.getNode(b.nodeID) res, err := b.manager.getNode(b.nodeID)
if err != nil { if err != nil {
return nil return nil
} }
b.ws = client.Workspace(sessionID) b.ws = taiworkspace.New(res.Volume, sessionID)
return b.ws return b.ws
} }
@ -220,39 +220,39 @@ func (b *Box) Snapshot() BoxInfo {
// VNC returns the VNC WebSocket URL. // VNC returns the VNC WebSocket URL.
func (b *Box) VNC(ctx context.Context) (string, error) { func (b *Box) VNC(ctx context.Context) (string, error) {
b.touch() b.touch()
client, err := b.manager.getNode(b.nodeID) res, err := b.manager.getNode(b.nodeID)
if err != nil { if err != nil {
return "", err return "", err
} }
return client.VNC().URL(ctx, b.containerID) return res.VNC.URL(ctx, b.containerID)
} }
// Proxy returns the HTTP URL for a service on the given port inside the sandbox. // Proxy returns the HTTP URL for a service on the given port inside the sandbox.
func (b *Box) Proxy(ctx context.Context, port int, path string) (string, error) { func (b *Box) Proxy(ctx context.Context, port int, path string) (string, error) {
b.touch() b.touch()
client, err := b.manager.getNode(b.nodeID) res, err := b.manager.getNode(b.nodeID)
if err != nil { if err != nil {
return "", err return "", err
} }
return client.Proxy().URL(ctx, b.containerID, port, path) return res.Proxy.URL(ctx, b.containerID, port, path)
} }
// Start starts a stopped sandbox. // Start starts a stopped sandbox.
func (b *Box) Start(ctx context.Context) error { func (b *Box) Start(ctx context.Context) error {
client, err := b.manager.getNode(b.nodeID) res, err := b.manager.getNode(b.nodeID)
if err != nil { if err != nil {
return err return err
} }
return client.Sandbox().Start(ctx, b.containerID) return res.Runtime.Start(ctx, b.containerID)
} }
// Stop stops the sandbox without removing it. // Stop stops the sandbox without removing it.
func (b *Box) Stop(ctx context.Context) error { func (b *Box) Stop(ctx context.Context) error {
client, err := b.manager.getNode(b.nodeID) res, err := b.manager.getNode(b.nodeID)
if err != nil { if err != nil {
return err return err
} }
return client.Sandbox().Stop(ctx, b.containerID, b.stopTimeout()) return res.Runtime.Stop(ctx, b.containerID, b.stopTimeout())
} }
// Remove stops and removes the sandbox. // Remove stops and removes the sandbox.
@ -262,12 +262,12 @@ func (b *Box) Remove(ctx context.Context) error {
// Info returns current sandbox status. // Info returns current sandbox status.
func (b *Box) Info(ctx context.Context) (*BoxInfo, error) { func (b *Box) Info(ctx context.Context) (*BoxInfo, error) {
client, err := b.manager.getNode(b.nodeID) res, err := b.manager.getNode(b.nodeID)
if err != nil { if err != nil {
return nil, err return nil, err
} }
info, err := client.Sandbox().Inspect(ctx, b.containerID) info, err := res.Runtime.Inspect(ctx, b.containerID)
if err != nil { if err != nil {
return nil, err return nil, err
} }

View file

@ -7,7 +7,7 @@ import (
"io" "io"
hepb "github.com/yaoapp/yao/tai/hostexec/pb" hepb "github.com/yaoapp/yao/tai/hostexec/pb"
"github.com/yaoapp/yao/tai/workspace" taiworkspace "github.com/yaoapp/yao/tai/workspace"
) )
// Host represents a Tai host machine execution environment. // Host represents a Tai host machine execution environment.
@ -44,12 +44,12 @@ func (h *Host) Exec(ctx context.Context, cmd []string, opts ...ExecOption) (*Exe
return nil, fmt.Errorf("sandbox: empty command") return nil, fmt.Errorf("sandbox: empty command")
} }
client, err := h.manager.getNode(h.nodeID) res, err := h.manager.getNode(h.nodeID)
if err != nil { if err != nil {
return nil, err return nil, err
} }
he := client.HostExec() he := res.HostExec
if he == nil { if he == nil {
return nil, fmt.Errorf("sandbox: host_exec not available on node %q", h.nodeID) return nil, fmt.Errorf("sandbox: host_exec not available on node %q", h.nodeID)
} }
@ -100,12 +100,12 @@ func (h *Host) Stream(ctx context.Context, cmd []string, opts ...ExecOption) (*E
return nil, fmt.Errorf("sandbox: empty command") return nil, fmt.Errorf("sandbox: empty command")
} }
client, err := h.manager.getNode(h.nodeID) res, err := h.manager.getNode(h.nodeID)
if err != nil { if err != nil {
return nil, err return nil, err
} }
he := client.HostExec() he := res.HostExec
if he == nil { if he == nil {
return nil, fmt.Errorf("sandbox: host_exec not available on node %q", h.nodeID) return nil, fmt.Errorf("sandbox: host_exec not available on node %q", h.nodeID)
} }
@ -189,21 +189,27 @@ func (h *Host) Stream(ctx context.Context, cmd []string, opts ...ExecOption) (*E
// VNC returns the VNC WebSocket URL for the Tai host machine. // VNC returns the VNC WebSocket URL for the Tai host machine.
// Uses the special __host__ identifier to route to localhost:5900 on the Tai server. // Uses the special __host__ identifier to route to localhost:5900 on the Tai server.
func (h *Host) VNC(ctx context.Context) (string, error) { func (h *Host) VNC(ctx context.Context) (string, error) {
client, err := h.manager.getNode(h.nodeID) res, err := h.manager.getNode(h.nodeID)
if err != nil { if err != nil {
return "", err return "", err
} }
return client.VNC().URL(ctx, "__host__") if res.VNC == nil {
return "", fmt.Errorf("sandbox: vnc not available on node %q", h.nodeID)
}
return res.VNC.URL(ctx, "__host__")
} }
// Proxy returns the HTTP URL for a service running on the Tai host machine. // Proxy returns the HTTP URL for a service running on the Tai host machine.
// Uses the special __host__ identifier to route to localhost:{port} on the Tai server. // Uses the special __host__ identifier to route to localhost:{port} on the Tai server.
func (h *Host) Proxy(ctx context.Context, port int, path string) (string, error) { func (h *Host) Proxy(ctx context.Context, port int, path string) (string, error) {
client, err := h.manager.getNode(h.nodeID) res, err := h.manager.getNode(h.nodeID)
if err != nil { if err != nil {
return "", err return "", err
} }
return client.Proxy().URL(ctx, "__host__", port, path) if res.Proxy == nil {
return "", fmt.Errorf("sandbox: proxy not available on node %q", h.nodeID)
}
return res.Proxy.URL(ctx, "__host__", port, path)
} }
// BindWorkplace binds a workspace to this host by ID. Subsequent calls to // BindWorkplace binds a workspace to this host by ID. Subsequent calls to
@ -213,15 +219,18 @@ func (h *Host) BindWorkplace(workspaceID string) {
} }
// Workplace returns the workspace FS bound to this host, or nil if unbound. // Workplace returns the workspace FS bound to this host, or nil if unbound.
func (h *Host) Workplace() workspace.FS { func (h *Host) Workplace() taiworkspace.FS {
if h.workplaceID == "" { if h.workplaceID == "" {
return nil return nil
} }
client, err := h.manager.getNode(h.nodeID) res, err := h.manager.getNode(h.nodeID)
if err != nil { if err != nil {
return nil return nil
} }
return client.Workspace(h.workplaceID) if res.Volume == nil {
return nil
}
return taiworkspace.New(res.Volume, h.workplaceID)
} }
// NodeID returns the node ID this Host belongs to. // NodeID returns the node ID this Host belongs to.

View file

@ -9,7 +9,6 @@ import (
"time" "time"
sandbox "github.com/yaoapp/yao/sandbox/v2" sandbox "github.com/yaoapp/yao/sandbox/v2"
"github.com/yaoapp/yao/tai"
) )
func setupHostManager(t *testing.T, tgt *hostExecTarget) *sandbox.Manager { func setupHostManager(t *testing.T, tgt *hostExecTarget) *sandbox.Manager {
@ -454,12 +453,12 @@ func findHostExecOnly(t *testing.T) *hostExecTarget {
for _, tgt := range hostExecTargets() { for _, tgt := range hostExecTargets() {
if tgt.IsWinNative { if tgt.IsWinNative {
addr := fmt.Sprintf("tai://%s", tgt.Addr) addr := fmt.Sprintf("tai://%s", tgt.Addr)
client, err := tai.New(addr) res, err := dialForTest(addr)
if err != nil { if err != nil {
continue continue
} }
hasNoSandbox := client.Sandbox() == nil hasNoSandbox := res.Runtime == nil
client.Close() res.Close()
if hasNoSandbox { if hasNoSandbox {
return &tgt return &tgt
} }

View file

@ -3,6 +3,7 @@ package jsapi_test
import ( import (
"fmt" "fmt"
"os" "os"
"strconv"
"strings" "strings"
"testing" "testing"
"time" "time"
@ -18,10 +19,9 @@ import (
) )
type testMode struct { type testMode struct {
Name string Name string
Addr string Addr string
TaiID string // filled by setupSandbox TaiID string
Options []tai.Option
} }
func testModes() []testMode { func testModes() []testMode {
@ -46,19 +46,71 @@ func setupSandbox(t *testing.T, m *testMode) {
reg := registry.Global() reg := registry.Global()
if reg == nil { if reg == nil {
registry.Init(nil) registry.Init(nil)
reg = registry.Global()
} }
client, err := tai.New(m.Addr, m.Options...) taiID, _ := registerForTest(t, m.Addr)
if err != nil { m.TaiID = taiID
t.Fatalf("tai.New: %v", err)
}
m.TaiID = client.TaiID()
sandbox.Init() sandbox.Init()
mgr := sandbox.M() mgr := sandbox.M()
t.Cleanup(func() { mgr.Close() }) t.Cleanup(func() { mgr.Close() })
} }
func registerForTest(t testing.TB, addr string, dialOps ...tai.DialOption) (string, *tai.ConnResources) {
t.Helper()
if registry.Global() == nil {
registry.Init(nil)
}
res, err := dialForTest(addr, dialOps...)
if err != nil {
t.Fatalf("dialForTest(%s): %v", addr, err)
}
taiID := taiIDFromAddr(addr)
reg := registry.Global()
reg.Register(&registry.TaiNode{TaiID: taiID, Mode: modeForAddr(addr)})
reg.SetResources(taiID, res)
t.Cleanup(func() { res.Close() })
return taiID, res
}
func dialForTest(addr string, dialOps ...tai.DialOption) (*tai.ConnResources, error) {
if addr == "local" || addr == "" {
return tai.DialLocal("", "", nil)
}
host, grpcPort := parseHostPort(addr)
ports := tai.Ports{GRPC: grpcPort}
return tai.DialRemote(host, ports, dialOps...)
}
func taiIDFromAddr(addr string) string {
if addr == "local" || addr == "" {
return "local"
}
addr = strings.TrimPrefix(addr, "tai://")
parts := strings.SplitN(addr, ":", 2)
return parts[0]
}
func modeForAddr(addr string) string {
if addr == "local" || addr == "" {
return "local"
}
return "direct"
}
func parseHostPort(addr string) (string, int) {
addr = strings.TrimPrefix(addr, "tai://")
parts := strings.SplitN(addr, ":", 2)
h := parts[0]
if len(parts) == 2 {
if p, err := strconv.Atoi(parts[1]); err == nil {
return h, p
}
}
return h, 19100
}
func runJS(t *testing.T, source string) interface{} { func runJS(t *testing.T, source string) interface{} {
t.Helper() t.Helper()
res, err := v8runtime.Call(v8runtime.CallOptions{ res, err := v8runtime.Call(v8runtime.CallOptions{

View file

@ -5,6 +5,7 @@ import (
"time" "time"
"github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/registry"
taitypes "github.com/yaoapp/yao/tai/types"
"rogchap.com/v8go" "rogchap.com/v8go"
) )
@ -63,17 +64,17 @@ func sbNodesByTeam(info *v8go.FunctionCallbackInfo) *v8go.Value {
return snapshotsToJSArray(v8ctx, snaps) return snapshotsToJSArray(v8ctx, snaps)
} }
// snapshotToJS converts a NodeSnapshot to a JS NodeInfo object. // snapshotToJS converts a NodeMeta to a JS NodeInfo object.
// Auth and YaoBase are excluded for security. // Auth and YaoBase are excluded for security.
func snapshotToJS(v8ctx *v8go.Context, snap *registry.NodeSnapshot) (*v8go.Value, error) { func snapshotToJS(v8ctx *v8go.Context, snap *taitypes.NodeMeta) (*v8go.Value, error) {
ports := make(map[string]interface{}, len(snap.Ports)) ports := map[string]interface{}{
for k, v := range snap.Ports { "grpc": snap.Ports.GRPC, "http": snap.Ports.HTTP,
ports[k] = v "vnc": snap.Ports.VNC, "docker": snap.Ports.Docker, "k8s": snap.Ports.K8s,
} }
caps := make(map[string]interface{}, len(snap.Capabilities)) caps := map[string]interface{}{
for k, v := range snap.Capabilities { "docker": snap.Capabilities.Docker, "k8s": snap.Capabilities.K8s,
caps[k] = v "host_exec": snap.Capabilities.HostExec,
} }
data, err := json.Marshal(map[string]interface{}{ data, err := json.Marshal(map[string]interface{}{
@ -103,17 +104,17 @@ func snapshotToJS(v8ctx *v8go.Context, snap *registry.NodeSnapshot) (*v8go.Value
return v8go.JSONParse(v8ctx, string(data)) return v8go.JSONParse(v8ctx, string(data))
} }
func snapshotsToJSArray(v8ctx *v8go.Context, snaps []registry.NodeSnapshot) *v8go.Value { func snapshotsToJSArray(v8ctx *v8go.Context, snaps []taitypes.NodeMeta) *v8go.Value {
items := make([]interface{}, 0, len(snaps)) items := make([]interface{}, 0, len(snaps))
for i := range snaps { for i := range snaps {
snap := &snaps[i] snap := &snaps[i]
ports := make(map[string]interface{}, len(snap.Ports)) ports := map[string]interface{}{
for k, v := range snap.Ports { "grpc": snap.Ports.GRPC, "http": snap.Ports.HTTP,
ports[k] = v "vnc": snap.Ports.VNC, "docker": snap.Ports.Docker, "k8s": snap.Ports.K8s,
} }
caps := make(map[string]interface{}, len(snap.Capabilities)) caps := map[string]interface{}{
for k, v := range snap.Capabilities { "docker": snap.Capabilities.Docker, "k8s": snap.Capabilities.K8s,
caps[k] = v "host_exec": snap.Capabilities.HostExec,
} }
items = append(items, map[string]interface{}{ items = append(items, map[string]interface{}{
"tai_id": snap.TaiID, "tai_id": snap.TaiID,

View file

@ -10,7 +10,8 @@ import (
"github.com/yaoapp/yao/config" "github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/tai" "github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/registry"
taisandbox "github.com/yaoapp/yao/tai/sandbox" tairuntime "github.com/yaoapp/yao/tai/runtime"
taitypes "github.com/yaoapp/yao/tai/types"
"github.com/yaoapp/yao/workspace" "github.com/yaoapp/yao/workspace"
) )
@ -38,11 +39,11 @@ func (m *Manager) Start(ctx context.Context) error {
m.ensureLocalNode(reg) m.ensureLocalNode(reg)
for _, snap := range reg.List() { for _, snap := range reg.List() {
client, err := m.getNode(snap.TaiID) res, err := m.getNode(snap.TaiID)
if err != nil { if err != nil {
continue continue
} }
m.recoverBoxes(ctx, snap.TaiID, client) m.recoverBoxes(ctx, snap.TaiID, res)
} }
loopCtx, cancel := context.WithCancel(ctx) loopCtx, cancel := context.WithCancel(ctx)
@ -61,7 +62,7 @@ func (m *Manager) ensureLocalNode(_ *registry.Registry) {
} }
// Nodes returns the list of registered Tai nodes from the registry. // Nodes returns the list of registered Tai nodes from the registry.
func (m *Manager) Nodes() []registry.NodeSnapshot { func (m *Manager) Nodes() []taitypes.NodeMeta {
reg := registry.Global() reg := registry.Global()
if reg == nil { if reg == nil {
return nil return nil
@ -89,26 +90,23 @@ func (m *Manager) Host(_ context.Context, nodeID string) (*Host, error) {
return nil, ErrNodeMissing return nil, ErrNodeMissing
} }
client, err := m.getNode(nodeID) res, err := m.getNode(nodeID)
if err != nil { if err != nil {
return nil, fmt.Errorf("sandbox: connect node %q: %w", nodeID, err) return nil, fmt.Errorf("sandbox: connect node %q: %w", nodeID, err)
} }
if client.HostExec() == nil { if res.HostExec == nil {
return nil, fmt.Errorf("sandbox: node %q has no host_exec capability", nodeID) return nil, fmt.Errorf("sandbox: node %q has no host_exec capability", nodeID)
} }
var sys SystemInfo sys := SystemInfo{
if snap, ok := tai.GetNodeSnapshot(nodeID); ok { OS: res.System.OS,
sys = SystemInfo{ Arch: res.System.Arch,
OS: snap.System.OS, Hostname: res.System.Hostname,
Arch: snap.System.Arch, NumCPU: res.System.NumCPU,
Hostname: snap.System.Hostname, TotalMem: res.System.TotalMem,
NumCPU: snap.System.NumCPU, Shell: res.System.Shell,
TotalMem: snap.System.TotalMem, TempDir: res.System.TempDir,
Shell: snap.System.Shell,
TempDir: snap.System.TempDir,
}
} }
return &Host{nodeID: nodeID, system: sys, manager: m}, nil return &Host{nodeID: nodeID, system: sys, manager: m}, nil
@ -165,24 +163,24 @@ func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Box, error)
id = fmt.Sprintf("sb-%d", time.Now().UnixNano()) id = fmt.Sprintf("sb-%d", time.Now().UnixNano())
} }
client, err := m.getNode(nodeID) res, err := m.getNode(nodeID)
if err != nil { if err != nil {
return nil, fmt.Errorf("sandbox: connect node %q: %w", nodeID, err) return nil, fmt.Errorf("sandbox: connect node %q: %w", nodeID, err)
} }
if client.Sandbox() == nil { if res.Runtime == nil {
return nil, fmt.Errorf("sandbox: node %q has no container runtime", nodeID) return nil, fmt.Errorf("sandbox: node %q has no container runtime", nodeID)
} }
taiOpts := m.buildTaiCreateOptions(opts, nodeID, id) taiOpts := m.buildTaiCreateOptions(opts, nodeID, id)
containerID, err := client.Sandbox().Create(ctx, taiOpts) containerID, err := res.Runtime.Create(ctx, taiOpts)
if err != nil { if err != nil {
return nil, fmt.Errorf("sandbox: create container: %w", err) return nil, fmt.Errorf("sandbox: create container: %w", err)
} }
if err := client.Sandbox().Start(ctx, containerID); err != nil { if err := res.Runtime.Start(ctx, containerID); err != nil {
client.Sandbox().Remove(ctx, containerID, true) res.Runtime.Remove(ctx, containerID, true)
return nil, fmt.Errorf("sandbox: start container: %w", err) return nil, fmt.Errorf("sandbox: start container: %w", err)
} }
@ -191,17 +189,14 @@ func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Box, error)
policy = Session policy = Session
} }
var sys SystemInfo sys := SystemInfo{
if snap, ok := tai.GetNodeSnapshot(nodeID); ok { OS: res.System.OS,
sys = SystemInfo{ Arch: res.System.Arch,
OS: snap.System.OS, Hostname: res.System.Hostname,
Arch: snap.System.Arch, NumCPU: res.System.NumCPU,
Hostname: snap.System.Hostname, TotalMem: res.System.TotalMem,
NumCPU: snap.System.NumCPU, Shell: res.System.Shell,
TotalMem: snap.System.TotalMem, TempDir: res.System.TempDir,
Shell: snap.System.Shell,
TempDir: snap.System.TempDir,
}
} }
box := &Box{ box := &Box{
@ -278,9 +273,9 @@ func (m *Manager) Remove(ctx context.Context, id string) error {
} }
b := v.(*Box) b := v.(*Box)
client, err := m.getNode(b.nodeID) res, err := m.getNode(b.nodeID)
if err == nil && client.Sandbox() != nil { if err == nil && res.Runtime != nil {
client.Sandbox().Remove(ctx, b.containerID, true) res.Runtime.Remove(ctx, b.containerID, true)
} }
m.boxes.Delete(id) m.boxes.Delete(id)
@ -303,8 +298,8 @@ func (m *Manager) Cleanup(ctx context.Context) error {
} }
case LongRunning: case LongRunning:
if timeout := b.idleTimeout(); timeout > 0 && idle > timeout { if timeout := b.idleTimeout(); timeout > 0 && idle > timeout {
if client, err := m.getNode(b.nodeID); err == nil && client.Sandbox() != nil { if res, err := m.getNode(b.nodeID); err == nil && res.Runtime != nil {
client.Sandbox().Stop(ctx, b.containerID, b.stopTimeout()) res.Runtime.Stop(ctx, b.containerID, b.stopTimeout())
} }
} }
if lifetime := b.maxLifetime(); lifetime > 0 && now.Sub(b.createdAt) > lifetime { if lifetime := b.maxLifetime(); lifetime > 0 && now.Sub(b.createdAt) > lifetime {
@ -339,15 +334,15 @@ func (m *Manager) cleanupLoop(ctx context.Context) {
} }
} }
func (m *Manager) getNode(name string) (*tai.Client, error) { func (m *Manager) getNode(name string) (*tai.ConnResources, error) {
client, ok := tai.GetClient(name) res, ok := tai.GetResources(name)
if !ok { if !ok {
return nil, ErrNodeNotFound return nil, ErrNodeNotFound
} }
return client, nil return res, nil
} }
func (m *Manager) buildTaiCreateOptions(opts CreateOptions, nodeID, sandboxID string) taisandbox.CreateOptions { func (m *Manager) buildTaiCreateOptions(opts CreateOptions, nodeID, sandboxID string) tairuntime.CreateOptions {
env := make(map[string]string) env := make(map[string]string)
reg := registry.Global() reg := registry.Global()
@ -385,9 +380,9 @@ func (m *Manager) buildTaiCreateOptions(opts CreateOptions, nodeID, sandboxID st
cmd := []string{"sh", "-c", "trap 'exit 0' TERM; while :; do sleep 86400 & wait $!; done"} cmd := []string{"sh", "-c", "trap 'exit 0' TERM; while :; do sleep 86400 & wait $!; done"}
var ports []taisandbox.PortMapping var ports []tairuntime.PortMapping
for _, p := range opts.Ports { for _, p := range opts.Ports {
ports = append(ports, taisandbox.PortMapping{ ports = append(ports, tairuntime.PortMapping{
ContainerPort: p.ContainerPort, ContainerPort: p.ContainerPort,
HostPort: p.HostPort, HostPort: p.HostPort,
HostIP: p.HostIP, HostIP: p.HostIP,
@ -413,7 +408,7 @@ func (m *Manager) buildTaiCreateOptions(opts CreateOptions, nodeID, sandboxID st
} }
} }
return taisandbox.CreateOptions{ return tairuntime.CreateOptions{
Name: sandboxID, Name: sandboxID,
Image: opts.Image, Image: opts.Image,
Cmd: cmd, Cmd: cmd,
@ -429,11 +424,11 @@ func (m *Manager) buildTaiCreateOptions(opts CreateOptions, nodeID, sandboxID st
} }
} }
func (m *Manager) recoverBoxes(ctx context.Context, nodeID string, client *tai.Client) { func (m *Manager) recoverBoxes(ctx context.Context, nodeID string, res *tai.ConnResources) {
if client.Sandbox() == nil { if res.Runtime == nil {
return return
} }
containers, err := client.Sandbox().List(ctx, taisandbox.ListOptions{ containers, err := res.Runtime.List(ctx, tairuntime.ListOptions{
All: true, All: true,
Labels: map[string]string{"managed-by": "yao-sandbox"}, Labels: map[string]string{"managed-by": "yao-sandbox"},
}) })
@ -473,37 +468,35 @@ func (m *Manager) recoverBoxes(ctx context.Context, nodeID string, client *tai.C
// ImageExists reports whether the given image ref exists on the target node. // ImageExists reports whether the given image ref exists on the target node.
func (m *Manager) ImageExists(ctx context.Context, nodeID, ref string) (bool, error) { func (m *Manager) ImageExists(ctx context.Context, nodeID, ref string) (bool, error) {
client, err := m.getNode(nodeID) res, err := m.getNode(nodeID)
if err != nil { if err != nil {
return false, err return false, err
} }
img := client.Image() if res.Image == nil {
if img == nil {
return true, nil return true, nil
} }
return img.Exists(ctx, ref) return res.Image.Exists(ctx, ref)
} }
// PullImage pulls an image to the target node, returning a channel of // PullImage pulls an image to the target node, returning a channel of
// real-time progress events. // real-time progress events.
func (m *Manager) PullImage(ctx context.Context, nodeID, ref string, opts ImagePullOptions) (<-chan taisandbox.PullProgress, error) { func (m *Manager) PullImage(ctx context.Context, nodeID, ref string, opts ImagePullOptions) (<-chan tairuntime.PullProgress, error) {
client, err := m.getNode(nodeID) res, err := m.getNode(nodeID)
if err != nil { if err != nil {
return nil, err return nil, err
} }
img := client.Image() if res.Image == nil {
if img == nil {
return nil, nil return nil, nil
} }
pullOpts := taisandbox.PullOptions{} pullOpts := tairuntime.PullOptions{}
if opts.Auth != nil { if opts.Auth != nil {
pullOpts.Auth = &taisandbox.RegistryAuth{ pullOpts.Auth = &tairuntime.RegistryAuth{
Username: opts.Auth.Username, Username: opts.Auth.Username,
Password: opts.Auth.Password, Password: opts.Auth.Password,
Server: opts.Auth.Server, Server: opts.Auth.Server,
} }
} }
return img.Pull(ctx, ref, pullOpts) return res.Image.Pull(ctx, ref, pullOpts)
} }
// EnsureImage checks whether the image exists on the node; if not, it // EnsureImage checks whether the image exists on the node; if not, it

View file

@ -0,0 +1,37 @@
//go:build containerized
package sandbox_test
import (
"fmt"
"os"
)
func init() {
extraNodeProviders = append(extraNodeProviders, containerizedNodes)
extraPurgeProviders = append(extraPurgeProviders, containerizedPurge)
}
func containerizedNodes() []nodeConfig {
host := os.Getenv("TAI_TEST_CONTAINERIZED_HOST")
if host == "" {
return nil
}
grpcPort := envPort("TAI_TEST_CONTAINERIZED_GRPC_PORT", 9200)
return []nodeConfig{{
Name: "containerized",
Addr: fmt.Sprintf("tai://%s:%d", host, grpcPort),
}}
}
func containerizedPurge() []purgeTarget {
host := os.Getenv("TAI_TEST_CONTAINERIZED_HOST")
if host == "" {
return nil
}
grpcPort := envPort("TAI_TEST_CONTAINERIZED_GRPC_PORT", 9200)
return []purgeTarget{{
name: "containerized",
addr: fmt.Sprintf("tai://%s:%d", host, grpcPort),
}}
}

View file

@ -0,0 +1,68 @@
//go:build k8s
package sandbox_test
import (
"fmt"
"os"
"github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/types"
)
func init() {
extraNodeProviders = append(extraNodeProviders, k8sNodes)
extraHostExecProviders = append(extraHostExecProviders, k8sHostExec)
extraPurgeProviders = append(extraPurgeProviders, k8sPurge)
}
func k8sNodes() []nodeConfig {
host := os.Getenv("TAI_TEST_K8S_HOST")
kubeconfig := os.Getenv("TAI_TEST_KUBECONFIG")
if host == "" || kubeconfig == "" {
return nil
}
grpcPort := envPort("TAI_TEST_K8S_GRPC_PORT", envPort("TAI_TEST_GRPC_PORT", 19100))
dialOps := []tai.DialOption{
tai.WithDialRuntime(types.K8s),
tai.WithDialKubeConfig(kubeconfig),
}
if ns := os.Getenv("TAI_TEST_K8S_NAMESPACE"); ns != "" {
dialOps = append(dialOps, tai.WithDialNamespace(ns))
}
return []nodeConfig{{
Name: "k8s",
Addr: fmt.Sprintf("tai://%s:%d", host, grpcPort),
DialOps: dialOps,
}}
}
func k8sHostExec() []hostExecTarget {
host := os.Getenv("TAI_TEST_K8S_HOST")
if host == "" {
return nil
}
grpcPort := envPort("TAI_TEST_K8S_GRPC_PORT", envPort("TAI_TEST_GRPC_PORT", 19100))
return []hostExecTarget{{Name: "k8s", Addr: fmt.Sprintf("%s:%d", host, grpcPort)}}
}
func k8sPurge() []purgeTarget {
host := os.Getenv("TAI_TEST_K8S_HOST")
kubeconfig := os.Getenv("TAI_TEST_KUBECONFIG")
if host == "" || kubeconfig == "" {
return nil
}
grpcPort := envPort("TAI_TEST_K8S_GRPC_PORT", envPort("TAI_TEST_GRPC_PORT", 19100))
dialOps := []tai.DialOption{
tai.WithDialRuntime(types.K8s),
tai.WithDialKubeConfig(kubeconfig),
}
if ns := os.Getenv("TAI_TEST_K8S_NAMESPACE"); ns != "" {
dialOps = append(dialOps, tai.WithDialNamespace(ns))
}
return []purgeTarget{{
name: "k8s",
addr: fmt.Sprintf("tai://%s:%d", host, grpcPort),
dialOps: dialOps,
}}
}

View file

@ -0,0 +1,39 @@
//go:build remote
package sandbox_test
import (
"os"
"strings"
)
func init() {
extraNodeProviders = append(extraNodeProviders, remoteNodes)
extraHostExecProviders = append(extraHostExecProviders, remoteHostExec)
extraPurgeProviders = append(extraPurgeProviders, remotePurge)
}
func remoteNodes() []nodeConfig {
addr := os.Getenv("SANDBOX_TEST_REMOTE_ADDR")
if addr == "" {
return nil
}
return []nodeConfig{{Name: "remote", Addr: addr}}
}
func remoteHostExec() []hostExecTarget {
addr := os.Getenv("SANDBOX_TEST_REMOTE_ADDR")
if addr == "" {
return nil
}
addr = strings.TrimPrefix(addr, "tai://")
return []hostExecTarget{{Name: "remote", Addr: addr}}
}
func remotePurge() []purgeTarget {
addr := os.Getenv("SANDBOX_TEST_REMOTE_ADDR")
if addr == "" {
return nil
}
return []purgeTarget{{name: "remote", addr: addr}}
}

View file

@ -14,7 +14,7 @@ import (
sandbox "github.com/yaoapp/yao/sandbox/v2" sandbox "github.com/yaoapp/yao/sandbox/v2"
"github.com/yaoapp/yao/tai" "github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/registry"
taisandbox "github.com/yaoapp/yao/tai/sandbox" tairuntime "github.com/yaoapp/yao/tai/runtime"
"github.com/yaoapp/yao/workspace" "github.com/yaoapp/yao/workspace"
) )
@ -25,62 +25,70 @@ var k8sSem = make(chan struct{}, 2)
// when many tests finish at once. // when many tests finish at once.
var k8sCleanupMu sync.Mutex var k8sCleanupMu sync.Mutex
// ---------------------------------------------------------------------------
// Build-tag extension points.
// Each tag file (testutils_remote_test.go, testutils_k8s_test.go, …) appends
// provider functions in its init(). This lets tags compose freely:
//
// go test ./sandbox/v2/... → local only
// go test -tags remote ./sandbox/v2/... → local + remote
// go test -tags "remote,k8s" ./sandbox/v2/... → local + remote + k8s
// go test -tags "remote,containerized,k8s,wintest" → all
//
// ---------------------------------------------------------------------------
var (
extraNodeProviders []func() []nodeConfig
extraHostExecProviders []func() []hostExecTarget
extraPurgeProviders []func() []purgeTarget
)
func TestMain(m *testing.M) { func TestMain(m *testing.M) {
purgeStaleContainers() purgeStaleContainers()
os.Exit(m.Run()) os.Exit(m.Run())
} }
// purgeStaleContainers removes leftover sb-* containers/pods from previous // ---------------------------------------------------------------------------
// test runs across all configured nodes (Docker + K8s). // Purge stale containers from previous runs
// ---------------------------------------------------------------------------
type purgeTarget struct {
name string
addr string
dialOps []tai.DialOption
}
func purgeStaleContainers() { func purgeStaleContainers() {
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel() defer cancel()
type target struct { type target struct {
name string name string
addr string addr string
opts []tai.Option dialOps []tai.DialOption
} }
var targets []target var targets []target
targets = append(targets, target{name: "local", addr: testLocalAddr()}) targets = append(targets, target{name: "local", addr: testLocalAddr()})
if addr := os.Getenv("SANDBOX_TEST_REMOTE_ADDR"); addr != "" { for _, fn := range extraPurgeProviders {
targets = append(targets, target{name: "remote", addr: addr}) for _, extra := range fn() {
} targets = append(targets, target{name: extra.name, addr: extra.addr, dialOps: extra.dialOps})
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 { for _, tgt := range targets {
client, err := tai.New(tgt.addr, tgt.opts...) res, err := dialForTest(tgt.addr, tgt.dialOps...)
if err != nil { if err != nil {
continue continue
} }
sb := client.Sandbox() sb := res.Runtime
if sb == nil { if sb == nil {
client.Close() res.Close()
continue continue
} }
containers, err := sb.List(ctx, taisandbox.ListOptions{All: true}) containers, err := sb.List(ctx, tairuntime.ListOptions{All: true})
if err != nil { if err != nil {
client.Close() res.Close()
continue continue
} }
for _, c := range containers { for _, c := range containers {
@ -94,57 +102,57 @@ func purgeStaleContainers() {
sb.Remove(ctx, id, true) sb.Remove(ctx, id, true)
log.Printf("[purge] %s: removed stale container %s", tgt.name, id) log.Printf("[purge] %s: removed stale container %s", tgt.name, id)
} }
client.Close() res.Close()
} }
} }
// ---------------------------------------------------------------------------
// Node / HostExec configuration
// ---------------------------------------------------------------------------
type nodeConfig struct { type nodeConfig struct {
Name string // human-readable label for t.Run (e.g. "remote", "k8s") Name string
Addr string Addr string
TaiID string // actual registry key, filled after tai.New TaiID string
Options []tai.Option DialOps []tai.DialOption
} }
// testNodes returns all available node configurations for multi-mode testing. type hostExecTarget struct {
Name string
Addr string
TaiID string
IsWinNative bool
}
// testNodes returns node configs. "local" is always present; other
// environments are injected by build-tag files via extraNodeProviders.
func testNodes() []nodeConfig { func testNodes() []nodeConfig {
nodes := []nodeConfig{ nodes := []nodeConfig{
{Name: "local", Addr: testLocalAddr()}, {Name: "local", Addr: testLocalAddr()},
} }
if addr := os.Getenv("SANDBOX_TEST_REMOTE_ADDR"); addr != "" { for _, fn := range extraNodeProviders {
nodes = append(nodes, nodeConfig{Name: "remote", Addr: addr}) nodes = append(nodes, fn()...)
}
if host := os.Getenv("TAI_TEST_CONTAINERIZED_HOST"); host != "" {
grpcPort := envPort("TAI_TEST_CONTAINERIZED_GRPC_PORT", 9200)
addr := fmt.Sprintf("tai://%s:%d", host, grpcPort)
nodes = append(nodes, nodeConfig{Name: "containerized", Addr: addr})
}
if host := os.Getenv("TAI_TEST_K8S_HOST"); host != "" {
kubeconfig := os.Getenv("TAI_TEST_KUBECONFIG")
if kubeconfig == "" {
return nodes
}
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,
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))
}
nodes = append(nodes, nodeConfig{Name: "k8s", Addr: addr, Options: opts})
} }
return nodes return nodes
} }
// hostExecTargets returns HostExec targets. Populated entirely by
// build-tag files via extraHostExecProviders.
func hostExecTargets() []hostExecTarget {
var targets []hostExecTarget
for _, fn := range extraHostExecProviders {
targets = append(targets, fn()...)
}
return targets
}
// ---------------------------------------------------------------------------
// Skip helpers
// ---------------------------------------------------------------------------
func skipIfNoDocker(t *testing.T) { func skipIfNoDocker(t *testing.T) {
t.Helper() t.Helper()
addr := testLocalAddr() if testLocalAddr() == "" {
if addr == "" {
t.Skip("SANDBOX_TEST_LOCAL_ADDR not set, skipping Docker tests") t.Skip("SANDBOX_TEST_LOCAL_ADDR not set, skipping Docker tests")
} }
} }
@ -156,32 +164,6 @@ func skipIfNoTai(t *testing.T) {
} }
} }
type hostExecTarget struct {
Name string
Addr string // host:port (without tai:// prefix)
TaiID string // filled after registration
IsWinNative bool
}
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) { func skipIfNoHostExec(t *testing.T) {
t.Helper() t.Helper()
if len(hostExecTargets()) == 0 { if len(hostExecTargets()) == 0 {
@ -189,6 +171,10 @@ func skipIfNoHostExec(t *testing.T) {
} }
} }
// ---------------------------------------------------------------------------
// Command helpers (Windows HostExec command translation)
// ---------------------------------------------------------------------------
func linuxCmd(tgt hostExecTarget, cmd string, args ...string) (string, []string) { func linuxCmd(tgt hostExecTarget, cmd string, args ...string) (string, []string) {
if tgt.IsWinNative { if tgt.IsWinNative {
switch cmd { switch cmd {
@ -214,6 +200,10 @@ func linuxCmd(tgt hostExecTarget, cmd string, args ...string) (string, []string)
return cmd, args return cmd, args
} }
// ---------------------------------------------------------------------------
// Environment helpers
// ---------------------------------------------------------------------------
func testLocalAddr() string { func testLocalAddr() string {
if addr := os.Getenv("SANDBOX_TEST_LOCAL_ADDR"); addr != "" { if addr := os.Getenv("SANDBOX_TEST_LOCAL_ADDR"); addr != "" {
return addr return addr
@ -237,41 +227,88 @@ func envPort(key string, fallback int) int {
return fallback return fallback
} }
// registerNode creates a tai.Client and registers it in the global registry. // ---------------------------------------------------------------------------
// It fills pc.TaiID with the actual registry key returned by tai.New. // Dial + Register helper (replaces old tai.New)
func registerNode(t *testing.T, pc *nodeConfig) { // ---------------------------------------------------------------------------
t.Helper()
reg := registry.Global() // dialForTest calls DialLocal or DialRemote based on the address.
if reg == nil { func dialForTest(addr string, dialOps ...tai.DialOption) (*tai.ConnResources, error) {
if addr == "local" || addr == "" {
return tai.DialLocal("", "", nil)
}
host, grpcPort := parseHostPort(addr)
ports := tai.Ports{GRPC: grpcPort}
return tai.DialRemote(host, ports, dialOps...)
}
// registerForTest dials and registers a node in the registry. Returns the
// taiID. On failure it calls t.Fatalf.
func registerForTest(t testing.TB, addr string, dialOps ...tai.DialOption) (string, *tai.ConnResources) {
t.Helper()
if registry.Global() == nil {
registry.Init(nil) registry.Init(nil)
} }
res, err := dialForTest(addr, dialOps...)
client, err := tai.New(pc.Addr, pc.Options...)
if err != nil { if err != nil {
t.Fatalf("tai.New(%s): %v", pc.Addr, err) t.Fatalf("dialForTest(%s): %v", addr, err)
} }
pc.TaiID = client.TaiID() taiID := taiIDFromAddr(addr)
t.Cleanup(func() { client.Close() }) reg := registry.Global()
reg.Register(&registry.TaiNode{TaiID: taiID, Mode: modeForAddr(addr)})
reg.SetResources(taiID, res)
return taiID, res
}
func taiIDFromAddr(addr string) string {
if addr == "local" || addr == "" {
return "local"
}
addr = strings.TrimPrefix(addr, "tai://")
host, _ := parseHostPort(addr)
return host
}
func modeForAddr(addr string) string {
if addr == "local" || addr == "" {
return "local"
}
return "direct"
}
func parseHostPort(addr string) (string, int) {
addr = strings.TrimPrefix(addr, "tai://")
parts := strings.SplitN(addr, ":", 2)
h := parts[0]
if len(parts) == 2 {
if p, err := strconv.Atoi(parts[1]); err == nil {
return h, p
}
}
return h, 19100
}
// ---------------------------------------------------------------------------
// Manager / Box setup helpers
// ---------------------------------------------------------------------------
func registerNode(t *testing.T, pc *nodeConfig) {
t.Helper()
taiID, res := registerForTest(t, pc.Addr, pc.DialOps...)
pc.TaiID = taiID
t.Cleanup(func() { res.Close() })
} }
func setupManager(t *testing.T, nodes ...nodeConfig) (*sandbox.Manager, []nodeConfig) { func setupManager(t *testing.T, nodes ...nodeConfig) (*sandbox.Manager, []nodeConfig) {
t.Helper() t.Helper()
if registry.Global() == nil {
reg := registry.Global()
if reg == nil {
registry.Init(nil) registry.Init(nil)
} }
_ = reg
out := make([]nodeConfig, len(nodes)) out := make([]nodeConfig, len(nodes))
copy(out, nodes) copy(out, nodes)
for i := range out { for i := range out {
client, err := tai.New(out[i].Addr, out[i].Options...) taiID, _ := registerForTest(t, out[i].Addr, out[i].DialOps...)
if err != nil { out[i].TaiID = taiID
t.Fatalf("tai.New(%s): %v", out[i].Addr, err)
}
out[i].TaiID = client.TaiID()
} }
sandbox.Init() sandbox.Init()
@ -287,8 +324,6 @@ func setupManagerForNode(t *testing.T, pc *nodeConfig) *sandbox.Manager {
return m return m
} }
// setupManagerWithWorkspace creates a sandbox Manager and returns
// the global workspace.Manager (which uses the registry for client lookups).
func setupManagerWithWorkspace(t *testing.T, pc *nodeConfig) (*sandbox.Manager, *workspace.Manager) { func setupManagerWithWorkspace(t *testing.T, pc *nodeConfig) (*sandbox.Manager, *workspace.Manager) {
t.Helper() t.Helper()
sbm := setupManagerForNode(t, pc) sbm := setupManagerForNode(t, pc)
@ -364,3 +399,6 @@ func createTestBox(t *testing.T, m *sandbox.Manager, pc nodeConfig, opts ...func
}) })
return box return box
} }
// Ensure imports are used.
var _ = fmt.Sprintf

View file

@ -0,0 +1,20 @@
//go:build wintest
package sandbox_test
import "os"
func init() {
extraHostExecProviders = append(extraHostExecProviders, winHostExec)
}
func winHostExec() []hostExecTarget {
var targets []hostExecTarget
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
}

View file

@ -11,22 +11,23 @@ import (
tai "github.com/yaoapp/yao/tai" tai "github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/registry"
"github.com/yaoapp/yao/tai/taiid" "github.com/yaoapp/yao/tai/taiid"
"github.com/yaoapp/yao/tai/types"
) )
// authenticateBearer validates a Bearer token and returns the caller's identity. // authenticateBearer validates a Bearer token and returns the caller's identity.
// Package-level var so tests can inject a mock without an OAuth service. // Package-level var so tests can inject a mock without an OAuth service.
var authenticateBearer = authenticateBearerDefault var authenticateBearer = authenticateBearerDefault
func authenticateBearerDefault(token string) (registry.AuthInfo, error) { func authenticateBearerDefault(token string) (types.AuthInfo, error) {
svc := oauth.OAuth svc := oauth.OAuth
if svc == nil { if svc == nil {
return registry.AuthInfo{}, fmt.Errorf("oauth service not initialized") return types.AuthInfo{}, fmt.Errorf("oauth service not initialized")
} }
result, err := svc.AuthenticateToken(oauth.AuthInput{AccessToken: token}) result, err := svc.AuthenticateToken(oauth.AuthInput{AccessToken: token})
if err != nil { if err != nil {
return registry.AuthInfo{}, err return types.AuthInfo{}, err
} }
info := registry.AuthInfo{} info := types.AuthInfo{}
if result.Info != nil { if result.Info != nil {
info.Subject = result.Info.Subject info.Subject = result.Info.Subject
info.UserID = result.Info.UserID info.UserID = result.Info.UserID
@ -86,15 +87,15 @@ func extractBearer(r *http.Request) string {
// registerRequest is the JSON body for POST /tai-nodes/register. // registerRequest is the JSON body for POST /tai-nodes/register.
type registerRequest struct { type registerRequest struct {
NodeID string `json:"node_id,omitempty"` NodeID string `json:"node_id,omitempty"`
ClientID string `json:"client_id,omitempty"` ClientID string `json:"client_id,omitempty"`
MachineID string `json:"machine_id"` MachineID string `json:"machine_id"`
DisplayName string `json:"display_name,omitempty"` DisplayName string `json:"display_name,omitempty"`
Version string `json:"version"` Version string `json:"version"`
Addr string `json:"addr"` Addr string `json:"addr"`
Ports map[string]int `json:"ports"` Ports map[string]int `json:"ports"`
Capabilities map[string]bool `json:"capabilities"` Capabilities map[string]bool `json:"capabilities"`
System registry.SystemInfo `json:"system"` System types.SystemInfo `json:"system"`
} }
// heartbeatRequest is the JSON body for POST /tai-nodes/heartbeat. // heartbeatRequest is the JSON body for POST /tai-nodes/heartbeat.
@ -162,8 +163,8 @@ func HandleRegister(c *gin.Context) {
System: req.System, System: req.System,
Mode: "direct", Mode: "direct",
Addr: addr, Addr: addr,
Ports: req.Ports, Ports: portsFromMap(req.Ports),
Capabilities: req.Capabilities, Capabilities: capsFromMap(req.Capabilities),
} }
reg.Register(node) reg.Register(node)
slog.Info("[register] node registered via API", slog.Info("[register] node registered via API",
@ -180,7 +181,7 @@ func HandleRegister(c *gin.Context) {
if strings.HasPrefix(addr, "tai://") { if strings.HasPrefix(addr, "tai://") {
slog.Info("[register] launching connectRegisteredNode goroutine", slog.Info("[register] launching connectRegisteredNode goroutine",
"tai_id", resolvedTaiID, "addr", addr) "tai_id", resolvedTaiID, "addr", addr)
go connectRegisteredNode(resolvedTaiID, addr, reg) go connectRegisteredNode(resolvedTaiID, addr, portsFromMap(req.Ports), reg)
} }
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
@ -278,45 +279,50 @@ func HandleUnregister(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"status": "unregistered"}) c.JSON(http.StatusOK, gin.H{"status": "unregistered"})
} }
// connectRegisteredNode dials the self-registered Tai node via gRPC, func portsFromMap(m map[string]int) types.Ports {
// creates a tai.Client, and binds it to the node's TaiID in the registry. return types.Ports{
// initRemote internally registers a redundant "host-port" entry; we remove GRPC: m["grpc"],
// it so that the registry contains only the canonical taiID. HTTP: m["http"],
func connectRegisteredNode(taiID, addr string, reg *registry.Registry) { VNC: m["vnc"],
Docker: m["docker"],
K8s: m["k8s"],
}
}
func capsFromMap(m map[string]bool) types.Capabilities {
return types.Capabilities{
Docker: m["docker"],
K8s: m["k8s"],
HostExec: m["host_exec"],
}
}
// connectRegisteredNode dials the Tai node via DialRemote and binds the
// returned ConnResources to the taiID in the registry. No double-registration.
func connectRegisteredNode(taiID, addr string, ports types.Ports, reg *registry.Registry) {
slog.Info("[connect] start", "tai_id", taiID, "addr", addr) slog.Info("[connect] start", "tai_id", taiID, "addr", addr)
client, err := tai.New(addr) host := extractHost(addr)
if err != nil { if host == "" {
slog.Warn("[connect] tai.New FAILED", slog.Warn("[connect] failed to extract host from addr", "addr", addr)
"tai_id", taiID, "addr", addr, "err", err)
allAfterFail := reg.List()
slog.Info("[connect] registry after tai.New failure", "total", len(allAfterFail))
for _, s := range allAfterFail {
slog.Info("[connect] node", "tai_id", s.TaiID, "mode", s.Mode, "addr", s.Addr)
}
return return
} }
autoID := client.TaiID() res, err := tai.DialRemote(host, ports)
slog.Info("[connect] tai.New OK", "tai_id", taiID, "autoID", autoID) if err != nil {
slog.Warn("[connect] DialRemote failed",
allAfterNew := reg.List() "tai_id", taiID, "addr", addr, "err", err)
slog.Info("[connect] registry after tai.New", "total", len(allAfterNew)) return
for _, s := range allAfterNew {
slog.Info("[connect] node", "tai_id", s.TaiID, "mode", s.Mode, "addr", s.Addr)
} }
if autoID != "" && autoID != taiID { reg.SetResources(taiID, res)
slog.Info("[connect] removing redundant autoID", "autoID", autoID)
reg.Unregister(autoID)
}
reg.SetClient(taiID, client)
allFinal := reg.List()
slog.Info("[connect] registry FINAL", "total", len(allFinal))
for _, s := range allFinal {
slog.Info("[connect] node", "tai_id", s.TaiID, "mode", s.Mode, "addr", s.Addr)
}
slog.Info("[connect] done", "tai_id", taiID) slog.Info("[connect] done", "tai_id", taiID)
} }
func extractHost(addr string) string {
addr = strings.TrimPrefix(addr, "tai://")
if idx := strings.LastIndex(addr, ":"); idx > 0 {
return addr[:idx]
}
return addr
}

View file

@ -9,6 +9,7 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/registry"
"github.com/yaoapp/yao/tai/types"
) )
func init() { func init() {
@ -20,8 +21,8 @@ func setupTest() func() {
registry.SetGlobalForTest(r) registry.SetGlobalForTest(r)
origAuth := authenticateBearer origAuth := authenticateBearer
authenticateBearer = func(token string) (registry.AuthInfo, error) { authenticateBearer = func(token string) (types.AuthInfo, error) {
return registry.AuthInfo{ return types.AuthInfo{
Subject: "sub-001", Subject: "sub-001",
UserID: "user-alice", UserID: "user-alice",
ClientID: "tai-abc123", ClientID: "tai-abc123",
@ -53,7 +54,7 @@ func TestHandleRegister_Success(t *testing.T) {
Addr: "192.168.1.100", Addr: "192.168.1.100",
Ports: map[string]int{"grpc": 19100, "http": 8099}, Ports: map[string]int{"grpc": 19100, "http": 8099},
Capabilities: map[string]bool{"docker": true, "host_exec": false}, Capabilities: map[string]bool{"docker": true, "host_exec": false},
System: registry.SystemInfo{ System: types.SystemInfo{
OS: "linux", Arch: "amd64", Hostname: "docker-host-01", NumCPU: 16, OS: "linux", Arch: "amd64", Hostname: "docker-host-01", NumCPU: 16,
}, },
} }
@ -114,7 +115,7 @@ func TestHandleRegister_ServerGeneratedTaiID(t *testing.T) {
Addr: "192.168.1.200", Addr: "192.168.1.200",
Ports: map[string]int{"grpc": 19100}, Ports: map[string]int{"grpc": 19100},
Capabilities: map[string]bool{"docker": true}, Capabilities: map[string]bool{"docker": true},
System: registry.SystemInfo{OS: "darwin", Arch: "arm64", Hostname: "mac-01", NumCPU: 12}, System: types.SystemInfo{OS: "darwin", Arch: "arm64", Hostname: "mac-01", NumCPU: 12},
} }
w := httptest.NewRecorder() w := httptest.NewRecorder()
@ -207,7 +208,7 @@ func TestHandleHeartbeat_Success(t *testing.T) {
reg.Register(&registry.TaiNode{ reg.Register(&registry.TaiNode{
TaiID: "tai-abc123", TaiID: "tai-abc123",
Mode: "direct", Mode: "direct",
Auth: registry.AuthInfo{ClientID: "tai-abc123"}, Auth: types.AuthInfo{ClientID: "tai-abc123"},
}) })
w := httptest.NewRecorder() w := httptest.NewRecorder()
@ -232,7 +233,7 @@ func TestHandleHeartbeat_WrongOwner(t *testing.T) {
reg.Register(&registry.TaiNode{ reg.Register(&registry.TaiNode{
TaiID: "tai-other", TaiID: "tai-other",
Mode: "direct", Mode: "direct",
Auth: registry.AuthInfo{ClientID: "different-client"}, Auth: types.AuthInfo{ClientID: "different-client"},
}) })
w := httptest.NewRecorder() w := httptest.NewRecorder()
@ -275,7 +276,7 @@ func TestHandleUnregister_Success(t *testing.T) {
reg.Register(&registry.TaiNode{ reg.Register(&registry.TaiNode{
TaiID: "tai-abc123", TaiID: "tai-abc123",
Mode: "direct", Mode: "direct",
Auth: registry.AuthInfo{ClientID: "tai-abc123"}, Auth: types.AuthInfo{ClientID: "tai-abc123"},
}) })
w := httptest.NewRecorder() w := httptest.NewRecorder()
@ -303,7 +304,7 @@ func TestHandleUnregister_WrongOwner(t *testing.T) {
reg.Register(&registry.TaiNode{ reg.Register(&registry.TaiNode{
TaiID: "tai-other", TaiID: "tai-other",
Mode: "direct", Mode: "direct",
Auth: registry.AuthInfo{ClientID: "different-client"}, Auth: types.AuthInfo{ClientID: "different-client"},
}) })
w := httptest.NewRecorder() w := httptest.NewRecorder()

62
tai/conn.go Normal file
View file

@ -0,0 +1,62 @@
package tai
import (
"errors"
"net"
hepb "github.com/yaoapp/yao/tai/hostexec/pb"
"github.com/yaoapp/yao/tai/proxy"
"github.com/yaoapp/yao/tai/runtime"
"github.com/yaoapp/yao/tai/types"
"github.com/yaoapp/yao/tai/vnc"
"github.com/yaoapp/yao/tai/volume"
"google.golang.org/grpc"
)
// ConnResources holds bare connection resources for a Tai node.
// Returned by Dial* functions. Caller (usually registry) is responsible
// for calling Close() when the node disconnects or resources are replaced.
type ConnResources struct {
GRPCConn *grpc.ClientConn
Runtime runtime.Runtime
Image runtime.Image
HostExec hepb.HostExecClient
Volume volume.Volume
Proxy proxy.Proxy
VNC vnc.VNC
Caps types.Capabilities
System types.SystemInfo
Ports types.Ports
Version string
DataDir string // host-side data dir (local mode only)
// Tunnel mode: local listeners that bridge to Tai via WS.
Listeners []net.Listener
}
// Close releases all held resources. Safe to call with nil fields.
func (r *ConnResources) Close() error {
if r == nil {
return nil
}
var errs []error
if r.Runtime != nil {
if err := r.Runtime.Close(); err != nil {
errs = append(errs, err)
}
}
if r.Volume != nil {
if err := r.Volume.Close(); err != nil {
errs = append(errs, err)
}
}
for _, ln := range r.Listeners {
ln.Close()
}
if r.GRPCConn != nil {
if err := r.GRPCConn.Close(); err != nil {
errs = append(errs, err)
}
}
return errors.Join(errs...)
}

387
tai/dial.go Normal file
View file

@ -0,0 +1,387 @@
package tai
import (
"context"
"fmt"
"net"
"net/http"
"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/runtime"
sipb "github.com/yaoapp/yao/tai/serverinfo/pb"
"github.com/yaoapp/yao/tai/types"
"github.com/yaoapp/yao/tai/vnc"
"github.com/yaoapp/yao/tai/volume"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/keepalive"
)
// DialRemote establishes connections to a remote Tai node via gRPC (direct mode).
// Does NOT interact with the registry. Caller must call ConnResources.Close().
func DialRemote(host string, ports types.Ports, opts ...DialOption) (*ConnResources, error) {
cfg := &dialConfig{ports: mergedPorts(ports)}
for _, o := range opts {
o.applyDial(cfg)
}
grpcAddr := fmt.Sprintf("%s:%d", host, cfg.ports.GRPC)
conn, err := dialGRPC(grpcAddr)
if err != nil {
return nil, fmt.Errorf("grpc dial %s: %w", grpcAddr, err)
}
return buildResources(conn, cfg, &remoteEnv{host: host, httpClient: cfg.httpClient})
}
// DialTunnel establishes connections to a Tai node through the WebSocket tunnel.
// Requires the node to already be registered in the registry (online).
// Does NOT call registry.SetResources. Caller must call ConnResources.Close().
func DialTunnel(taiID string, reg *registry.Registry, opts ...DialOption) (*ConnResources, error) {
node, ok := reg.Get(taiID)
if !ok || node.Status != "online" {
return nil, fmt.Errorf("tai node %s not online", taiID)
}
cfg := &dialConfig{
ports: types.Ports{
GRPC: intOr(node.Ports.GRPC, 19100),
HTTP: intOr(node.Ports.HTTP, 8099),
VNC: intOr(node.Ports.VNC, 16080),
Docker: intOr(node.Ports.Docker, 12375),
K8s: intOr(node.Ports.K8s, 16443),
},
}
for _, o := range opts {
o.applyDial(cfg)
}
grpcLn, err := reg.OpenLocalListener(taiID, cfg.ports.GRPC)
if err != nil {
return nil, fmt.Errorf("open grpc tunnel listener: %w", err)
}
conn, err := dialGRPC("passthrough:///" + grpcLn.Addr().String())
if err != nil {
grpcLn.Close()
return nil, fmt.Errorf("grpc dial tunnel %s: %w", grpcLn.Addr(), err)
}
env := &tunnelEnv{
taiID: taiID,
yaoBase: node.YaoBase,
reg: reg,
regCaps: node.Capabilities,
listeners: []net.Listener{grpcLn},
}
res, err := buildResources(conn, cfg, env)
if err != nil {
grpcLn.Close()
conn.Close()
return nil, err
}
res.Listeners = env.listeners
return res, nil
}
// DialLocal establishes connections to the local Docker daemon.
// Does NOT interact with the registry. Caller must call ConnResources.Close().
func DialLocal(addr string, dataDir string, vol volume.Volume) (*ConnResources, error) {
sb, err := runtime.NewLocal(addr)
if err != nil && vol == nil {
return nil, err
}
res := &ConnResources{DataDir: dataDir}
if sb != nil {
res.Runtime = sb
res.Image = runtime.NewDockerImage(runtime.DockerCli(sb))
res.Proxy = proxy.NewLocal(sb)
res.VNC = vnc.NewLocal(sb)
}
if vol != nil {
res.Volume = vol
} else {
if dataDir == "" {
dataDir = "/tmp/tai-volumes"
}
res.DataDir = dataDir
res.Volume = volume.NewLocal(dataDir)
}
return res, nil
}
// ---------------------------------------------------------------------------
// Shared build logic
// ---------------------------------------------------------------------------
// dialEnv abstracts the mode-specific differences (remote vs tunnel) that
// buildResources needs.
type dialEnv interface {
fallbackCaps() map[string]bool
mergeCaps(discovered map[string]bool) types.Capabilities
// listenAddr opens or formats a host:port address for the given port.
// Tunnel mode opens a local listener; remote mode formats host:port.
listenAddr(port int) (string, error)
newProxy(ports types.Ports) proxy.Proxy
newVNC(ports types.Ports) vnc.VNC
}
// buildResources constructs a ConnResources from an established gRPC
// connection. Shared by DialRemote and DialTunnel.
func buildResources(conn *grpc.ClientConn, cfg *dialConfig, env dialEnv) (*ConnResources, error) {
info, err := discoverInfo(conn, cfg)
if err != nil {
info = &discoveredInfo{Capabilities: env.fallbackCaps()}
}
caps := env.mergeCaps(info.Capabilities)
res := &ConnResources{
GRPCConn: conn,
HostExec: hepb.NewHostExecClient(conn),
Volume: volume.NewRemote(conn),
Caps: caps,
System: info.System,
Ports: cfg.ports,
Version: info.Version,
}
if cfg.runtime == types.K8s || (!caps.Docker && caps.K8s) {
if cfg.kubeConfig != "" {
k8sPort := cfg.ports.K8s
if k8sPort == 0 {
k8sPort = 16443
}
addr, err := env.listenAddr(k8sPort)
if err == nil {
sb, err := runtime.NewK8s(addr, runtime.K8sOption{
Namespace: cfg.namespace,
KubeConfig: cfg.kubeConfig,
})
if err == nil {
res.Runtime = sb
res.Image = runtime.NewK8sImage()
}
}
}
} else if caps.Docker {
dockerPort := cfg.ports.Docker
if dockerPort == 0 {
dockerPort = 12375
}
addr, err := env.listenAddr(dockerPort)
if err == nil {
sb, err := runtime.NewDocker("tcp://" + addr)
if err == nil {
res.Runtime = sb
res.Image = runtime.NewDockerImage(runtime.DockerCli(sb))
}
}
}
if res.Runtime != nil {
res.Proxy = env.newProxy(cfg.ports)
res.VNC = env.newVNC(cfg.ports)
}
return res, nil
}
// ---------------------------------------------------------------------------
// remoteEnv — direct TCP connections
// ---------------------------------------------------------------------------
type remoteEnv struct {
host string
httpClient *http.Client
}
func (e *remoteEnv) fallbackCaps() map[string]bool {
return map[string]bool{"docker": true}
}
func (e *remoteEnv) mergeCaps(discovered map[string]bool) types.Capabilities {
return types.Capabilities{
Docker: discovered["docker"],
K8s: discovered["k8s"],
HostExec: discovered["host_exec"],
}
}
func (e *remoteEnv) listenAddr(port int) (string, error) {
return fmt.Sprintf("%s:%d", e.host, port), nil
}
func (e *remoteEnv) newProxy(ports types.Ports) proxy.Proxy {
return proxy.NewRemote(e.host, ports.HTTP, e.httpClient)
}
func (e *remoteEnv) newVNC(ports types.Ports) vnc.VNC {
return vnc.NewRemote(e.host, ports.VNC, e.httpClient)
}
// ---------------------------------------------------------------------------
// tunnelEnv — connections via WebSocket tunnel
// ---------------------------------------------------------------------------
type tunnelEnv struct {
taiID string
yaoBase string
reg *registry.Registry
regCaps types.Capabilities
listeners []net.Listener
}
func (e *tunnelEnv) fallbackCaps() map[string]bool {
return make(map[string]bool)
}
func (e *tunnelEnv) mergeCaps(discovered map[string]bool) types.Capabilities {
return types.Capabilities{
Docker: discovered["docker"] || e.regCaps.Docker,
K8s: discovered["k8s"] || e.regCaps.K8s,
HostExec: discovered["host_exec"] || e.regCaps.HostExec,
}
}
func (e *tunnelEnv) listenAddr(port int) (string, error) {
ln, err := e.reg.OpenLocalListener(e.taiID, port)
if err != nil {
return "", err
}
e.listeners = append(e.listeners, ln)
return ln.Addr().String(), nil
}
func (e *tunnelEnv) newProxy(_ types.Ports) proxy.Proxy {
return proxy.NewTunnel(e.taiID, e.yaoBase)
}
func (e *tunnelEnv) newVNC(_ types.Ports) vnc.VNC {
return vnc.NewTunnel(e.taiID, e.yaoBase)
}
// ---------------------------------------------------------------------------
// Dial options
// ---------------------------------------------------------------------------
// DialOption configures a Dial* call.
type DialOption interface {
applyDial(*dialConfig)
}
type dialOptionFunc func(*dialConfig)
func (f dialOptionFunc) applyDial(c *dialConfig) { f(c) }
// WithDialRuntime selects the container runtime for the dial call.
func WithDialRuntime(rt types.Runtime) DialOption {
return dialOptionFunc(func(c *dialConfig) { c.runtime = rt })
}
// WithDialKubeConfig sets the kubeconfig for K8s runtime.
func WithDialKubeConfig(path string) DialOption {
return dialOptionFunc(func(c *dialConfig) { c.kubeConfig = path })
}
// WithDialNamespace sets the K8s namespace.
func WithDialNamespace(ns string) DialOption {
return dialOptionFunc(func(c *dialConfig) { c.namespace = ns })
}
// WithDialHTTPClient sets a custom HTTP client for proxy/VNC.
func WithDialHTTPClient(hc *http.Client) DialOption {
return dialOptionFunc(func(c *dialConfig) { c.httpClient = hc })
}
type dialConfig struct {
runtime types.Runtime
ports types.Ports
kubeConfig string
namespace string
httpClient *http.Client
userPorts types.Ports
}
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
func dialGRPC(target string) (*grpc.ClientConn, error) {
return grpc.NewClient(target,
grpc.WithTransportCredentials(insecure.NewCredentials()),
grpc.WithKeepaliveParams(keepalive.ClientParameters{
Time: 20 * time.Second,
Timeout: 5 * time.Second,
PermitWithoutStream: true,
}),
)
}
// ---------------------------------------------------------------------------
// ServerInfo discovery (shared by DialRemote / DialTunnel)
// ---------------------------------------------------------------------------
type discoveredInfo struct {
Capabilities map[string]bool
System types.SystemInfo
Version string
}
func discoverInfo(conn *grpc.ClientConn, cfg *dialConfig) (*discoveredInfo, 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 nil, err
}
up := cfg.userPorts
if p := int(resp.Ports["http"]); p > 0 && up.HTTP == 0 {
cfg.ports.HTTP = p
}
if p := int(resp.Ports["docker"]); p > 0 && up.Docker == 0 {
cfg.ports.Docker = p
}
if p := int(resp.Ports["vnc"]); p > 0 && up.VNC == 0 {
cfg.ports.VNC = p
}
if p := int(resp.Ports["k8s"]); p > 0 && up.K8s == 0 {
cfg.ports.K8s = p
}
caps := resp.Capabilities
if caps == nil {
caps = make(map[string]bool)
}
var sys types.SystemInfo
if s := resp.System; s != nil {
sys = types.SystemInfo{
OS: s.Os,
Arch: s.Arch,
Hostname: s.Hostname,
NumCPU: int(s.NumCpu),
TotalMem: s.TotalMem,
Shell: s.Shell,
TempDir: s.TempDir,
}
}
return &discoveredInfo{
Capabilities: caps,
System: sys,
Version: resp.Version,
}, nil
}

View file

@ -6,7 +6,7 @@ import (
"net/http" "net/http"
"strings" "strings"
"github.com/yaoapp/yao/tai/sandbox" "github.com/yaoapp/yao/tai/runtime"
) )
// Proxy resolves HTTP service URLs for containers. // Proxy resolves HTTP service URLs for containers.
@ -98,11 +98,11 @@ func (t *tunnelProxy) Healthz(_ context.Context) error {
// --- Local implementation --- // --- Local implementation ---
type localProxy struct { type localProxy struct {
sb sandbox.Sandbox sb runtime.Runtime
} }
// NewLocal creates a Proxy that resolves host ports via sandbox.Inspect. // NewLocal creates a Proxy that resolves host ports via runtime.Inspect.
func NewLocal(sb sandbox.Sandbox) Proxy { func NewLocal(sb runtime.Runtime) Proxy {
return &localProxy{sb: sb} return &localProxy{sb: sb}
} }

View file

@ -10,7 +10,7 @@ import (
"time" "time"
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
"github.com/yaoapp/yao/tai/sandbox" "github.com/yaoapp/yao/tai/runtime"
) )
func TestRemoteURL(t *testing.T) { func TestRemoteURL(t *testing.T) {
@ -72,10 +72,10 @@ func TestRemoteHealthzFail(t *testing.T) {
func TestLocalURL(t *testing.T) { func TestLocalURL(t *testing.T) {
mock := &mockSandbox{ mock := &mockSandbox{
inspectFn: func(ctx context.Context, id string) (*sandbox.ContainerInfo, error) { inspectFn: func(ctx context.Context, id string) (*runtime.ContainerInfo, error) {
return &sandbox.ContainerInfo{ return &runtime.ContainerInfo{
ID: id, ID: id,
Ports: []sandbox.PortMapping{ Ports: []runtime.PortMapping{
{ContainerPort: 3000, HostPort: 32768, HostIP: "127.0.0.1", Protocol: "tcp"}, {ContainerPort: 3000, HostPort: 32768, HostIP: "127.0.0.1", Protocol: "tcp"},
{ContainerPort: 8080, HostPort: 32769, HostIP: "127.0.0.1", Protocol: "tcp"}, {ContainerPort: 8080, HostPort: 32769, HostIP: "127.0.0.1", Protocol: "tcp"},
}, },
@ -98,8 +98,8 @@ func TestLocalURL(t *testing.T) {
func TestLocalURLPortNotFound(t *testing.T) { func TestLocalURLPortNotFound(t *testing.T) {
mock := &mockSandbox{ mock := &mockSandbox{
inspectFn: func(ctx context.Context, id string) (*sandbox.ContainerInfo, error) { inspectFn: func(ctx context.Context, id string) (*runtime.ContainerInfo, error) {
return &sandbox.ContainerInfo{ID: id}, nil return &runtime.ContainerInfo{ID: id}, nil
}, },
} }
@ -112,7 +112,7 @@ func TestLocalURLPortNotFound(t *testing.T) {
func TestLocalURLInspectError(t *testing.T) { func TestLocalURLInspectError(t *testing.T) {
mock := &mockSandbox{ mock := &mockSandbox{
inspectFn: func(ctx context.Context, id string) (*sandbox.ContainerInfo, error) { inspectFn: func(ctx context.Context, id string) (*runtime.ContainerInfo, error) {
return nil, fmt.Errorf("not found") return nil, fmt.Errorf("not found")
}, },
} }
@ -260,12 +260,12 @@ func TestConnectSSE_Non200(t *testing.T) {
} }
} }
// mockSandbox implements sandbox.Sandbox for testing. // mockSandbox implements runtime.Sandbox for testing.
type mockSandbox struct { type mockSandbox struct {
inspectFn func(ctx context.Context, id string) (*sandbox.ContainerInfo, error) inspectFn func(ctx context.Context, id string) (*runtime.ContainerInfo, error)
} }
func (m *mockSandbox) Create(ctx context.Context, opts sandbox.CreateOptions) (string, error) { func (m *mockSandbox) Create(ctx context.Context, opts runtime.CreateOptions) (string, error) {
return "", nil return "", nil
} }
func (m *mockSandbox) Start(ctx context.Context, id string) error { return nil } func (m *mockSandbox) Start(ctx context.Context, id string) error { return nil }
@ -273,19 +273,19 @@ func (m *mockSandbox) Stop(ctx context.Context, id string, timeout time.Duration
return nil return nil
} }
func (m *mockSandbox) Remove(ctx context.Context, id string, force bool) error { return nil } func (m *mockSandbox) Remove(ctx context.Context, id string, force bool) error { return nil }
func (m *mockSandbox) Exec(ctx context.Context, id string, cmd []string, opts sandbox.ExecOptions) (*sandbox.ExecResult, error) { func (m *mockSandbox) Exec(ctx context.Context, id string, cmd []string, opts runtime.ExecOptions) (*runtime.ExecResult, error) {
return nil, nil return nil, nil
} }
func (m *mockSandbox) ExecStream(ctx context.Context, id string, cmd []string, opts sandbox.ExecOptions) (*sandbox.StreamHandle, error) { func (m *mockSandbox) ExecStream(ctx context.Context, id string, cmd []string, opts runtime.ExecOptions) (*runtime.StreamHandle, error) {
return nil, nil return nil, nil
} }
func (m *mockSandbox) Inspect(ctx context.Context, id string) (*sandbox.ContainerInfo, error) { func (m *mockSandbox) Inspect(ctx context.Context, id string) (*runtime.ContainerInfo, error) {
if m.inspectFn != nil { if m.inspectFn != nil {
return m.inspectFn(ctx, id) return m.inspectFn(ctx, id)
} }
return &sandbox.ContainerInfo{ID: id}, nil return &runtime.ContainerInfo{ID: id}, nil
} }
func (m *mockSandbox) List(ctx context.Context, opts sandbox.ListOptions) ([]sandbox.ContainerInfo, error) { func (m *mockSandbox) List(ctx context.Context, opts runtime.ListOptions) ([]runtime.ContainerInfo, error) {
return nil, nil return nil, nil
} }
func (m *mockSandbox) Close() error { return nil } func (m *mockSandbox) Close() error { return nil }

View file

@ -12,109 +12,48 @@ import (
"sync" "sync"
"time" "time"
"github.com/gorilla/websocket" "github.com/yaoapp/yao/tai/types"
) )
// 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"`
Shell string `json:"shell,omitempty"`
TempDir string `json:"temp_dir,omitempty"`
}
// TaiNode represents a registered Tai instance (direct or tunnel). // TaiNode represents a registered Tai instance (direct or tunnel).
// Internal use only; external callers receive NodeSnapshot via Get()/List(). // Internal use only; external callers receive types.NodeMeta via Get()/List().
type TaiNode struct { type TaiNode struct {
TaiID string TaiID string
MachineID string MachineID string
Version string Version string
Auth AuthInfo Auth types.AuthInfo
System SystemInfo System types.SystemInfo
Mode string // "direct" | "tunnel" Mode string // "direct" | "tunnel"
Addr string // direct mode: "tai-host"; tunnel mode: empty Addr string // direct mode: "tai-host"; tunnel mode: empty
YaoBase string // Yao server base URL reported by Tai (tunnel mode) YaoBase string // Yao server base URL reported by Tai (tunnel mode)
Ports map[string]int // {"grpc":19100, "http":8099, "vnc":16080, "docker":12375} Ports types.Ports
Capabilities map[string]bool Capabilities types.Capabilities
ControlConn *websocket.Conn registerStream any // taipb.TaiTunnel_RegisterServer (stored as any to avoid import cycle)
connMu sync.Mutex // protects ControlConn writes
Status string // "online" | "offline" | "connecting" Status string // "online" | "offline" | "connecting"
ConnectedAt time.Time ConnectedAt time.Time
LastPing time.Time LastPing time.Time
DisplayName string // optional human-readable name for UI DisplayName string // optional human-readable name for UI
client any // *tai.Client; stored as any to avoid import cycle resources any // *tai.ConnResources; stored as any to avoid import cycle
localListeners map[int]*tunnelListener localListeners map[int]*tunnelListener
} }
// NodeSnapshot is a read-only copy of TaiNode fields safe to use outside locks. func (n *TaiNode) meta() types.NodeMeta {
type NodeSnapshot struct { return types.NodeMeta{
TaiID string
MachineID string
Version string
Auth AuthInfo
System SystemInfo
Mode string
Addr string
YaoBase string
Ports map[string]int
Capabilities map[string]bool
Status string
ConnectedAt time.Time
LastPing time.Time
DisplayName string
client any
}
func (n *TaiNode) snapshot() NodeSnapshot {
ports := make(map[string]int, len(n.Ports))
for k, v := range n.Ports {
ports[k] = v
}
caps := make(map[string]bool, len(n.Capabilities))
for k, v := range n.Capabilities {
caps[k] = v
}
return NodeSnapshot{
TaiID: n.TaiID, MachineID: n.MachineID, Version: n.Version, TaiID: n.TaiID, MachineID: n.MachineID, Version: n.Version,
Auth: n.Auth, System: n.System, Auth: n.Auth, System: n.System,
Mode: n.Mode, Addr: n.Addr, YaoBase: n.YaoBase, Mode: n.Mode, Addr: n.Addr, YaoBase: n.YaoBase,
Ports: ports, Capabilities: caps, Ports: n.Ports, Capabilities: n.Capabilities,
Status: n.Status, ConnectedAt: n.ConnectedAt, LastPing: n.LastPing, Status: n.Status, ConnectedAt: n.ConnectedAt, LastPing: n.LastPing,
DisplayName: n.DisplayName, DisplayName: n.DisplayName,
client: n.client,
} }
} }
// Client returns the associated *tai.Client (as any to avoid import cycle).
// Callers should type-assert: snap.Client().(*tai.Client).
func (s *NodeSnapshot) Client() any { return s.client }
// AuthInfo holds Yao user authorization extracted from OAuth token.
type AuthInfo struct {
Subject string
UserID string
ClientID string
Scope string
TeamID string
TenantID string
}
// pendingChannel represents a channel awaiting Tai's data WS connection.
type pendingChannel struct {
taiID string
result chan net.Conn
timer *time.Timer
}
// tunnelListener wraps a TCP listener that bridges each accepted connection // tunnelListener wraps a TCP listener that bridges each accepted connection
// through the WS tunnel to a specific Tai port. // through the tunnel to a specific Tai port.
type tunnelListener struct { type tunnelListener struct {
listener net.Listener listener net.Listener
taiID string taiID string
@ -127,12 +66,17 @@ var (
once sync.Once once sync.Once
) )
// BridgeFunc bridges a local TCP connection to a target port on a tunnel node.
// Set via SetBridgeFunc once the gRPC tunnel handler is ready.
type BridgeFunc func(taiID string, targetPort int, localConn net.Conn)
// Registry manages all Tai nodes (direct and tunnel). // Registry manages all Tai nodes (direct and tunnel).
type Registry struct { type Registry struct {
mu sync.RWMutex mu sync.RWMutex
nodes map[string]*TaiNode nodes map[string]*TaiNode
pending map[string]*pendingChannel logger *slog.Logger
logger *slog.Logger bridgeFn BridgeFunc
bridgeMu sync.RWMutex
} }
// Init initializes the global registry singleton. // Init initializes the global registry singleton.
@ -142,9 +86,8 @@ func Init(logger *slog.Logger) {
logger = slog.Default() logger = slog.Default()
} }
global = &Registry{ global = &Registry{
nodes: make(map[string]*TaiNode), nodes: make(map[string]*TaiNode),
pending: make(map[string]*pendingChannel), logger: logger,
logger: logger,
} }
}) })
} }
@ -188,7 +131,8 @@ func (r *Registry) Register(node *TaiNode) {
"tai_id", node.TaiID, "mode", node.Mode, "version", node.Version) "tai_id", node.TaiID, "mode", node.Mode, "version", node.Version)
} }
// Unregister removes a Tai node and closes its local listeners and control connection. // Unregister removes a Tai node, closes its local listeners,
// and any held ConnResources.
func (r *Registry) Unregister(taiID string) { func (r *Registry) Unregister(taiID string) {
r.mu.Lock() r.mu.Lock()
node, ok := r.nodes[taiID] node, ok := r.nodes[taiID]
@ -197,63 +141,43 @@ func (r *Registry) Unregister(taiID string) {
tl.cancel() tl.cancel()
tl.listener.Close() tl.listener.Close()
} }
node.connMu.Lock()
if node.ControlConn != nil {
node.ControlConn.Close()
node.ControlConn = nil
}
node.connMu.Unlock()
delete(r.nodes, taiID) delete(r.nodes, taiID)
} }
r.mu.Unlock() r.mu.Unlock()
if ok { if ok {
if node.resources != nil {
if closer, ok := node.resources.(ResourceCloser); ok {
closer.Close()
}
}
r.logger.Info("tai node unregistered", "tai_id", taiID) r.logger.Info("tai node unregistered", "tai_id", taiID)
} }
} }
// Get returns a snapshot of a Tai node by ID. Returns nil, false if not found. // Get returns the metadata of a Tai node by ID. Returns nil, false if not found.
func (r *Registry) Get(taiID string) (*NodeSnapshot, bool) { func (r *Registry) Get(taiID string) (*types.NodeMeta, bool) {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
n, ok := r.nodes[taiID] n, ok := r.nodes[taiID]
if !ok { if !ok {
return nil, false return nil, false
} }
snap := n.snapshot() m := n.meta()
return &snap, true return &m, true
} }
// List returns snapshots of all registered Tai nodes. // List returns metadata of all registered Tai nodes.
func (r *Registry) List() []NodeSnapshot { func (r *Registry) List() []types.NodeMeta {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
result := make([]NodeSnapshot, 0, len(r.nodes)) result := make([]types.NodeMeta, 0, len(r.nodes))
for _, n := range r.nodes { for _, n := range r.nodes {
result = append(result, n.snapshot()) result = append(result, n.meta())
} }
return result return result
} }
// WriteControlJSON sends a JSON message on the node's control channel
// with proper serialization. Returns error if node not found or not tunnel.
func (r *Registry) WriteControlJSON(taiID string, v interface{}) error {
r.mu.RLock()
node := r.nodes[taiID]
r.mu.RUnlock()
if node == nil {
return fmt.Errorf("tai node %s not found", taiID)
}
node.connMu.Lock()
defer node.connMu.Unlock()
if node.ControlConn == nil {
return fmt.Errorf("tai node %s has no active control channel", taiID)
}
return node.ControlConn.WriteJSON(v)
}
// UpdatePing records a heartbeat timestamp. // UpdatePing records a heartbeat timestamp.
func (r *Registry) UpdatePing(taiID string) { func (r *Registry) UpdatePing(taiID string) {
r.mu.Lock() r.mu.Lock()
@ -263,20 +187,77 @@ func (r *Registry) UpdatePing(taiID string) {
} }
} }
// SetClient associates a *tai.Client with a registered node. // ResourceCloser is implemented by *tai.ConnResources to allow the registry
// Called by tai.New() after successful initialization. // to close resources without importing the tai package (avoids import cycle).
func (r *Registry) SetClient(taiID string, c any) { type ResourceCloser interface {
Close() error
}
// SetResources binds connection resources to a registered node.
// If the node already has resources, the old ones are closed asynchronously.
// The node status is set to "online".
func (r *Registry) SetResources(taiID string, res any) {
r.mu.Lock()
defer r.mu.Unlock()
n, ok := r.nodes[taiID]
if !ok {
return
}
if n.resources != nil {
if closer, ok := n.resources.(ResourceCloser); ok {
go closer.Close()
}
}
n.resources = res
n.Status = "online"
}
// GetResources returns the *tai.ConnResources for a node (as any).
// Callers should type-assert to *tai.ConnResources.
func (r *Registry) GetResources(taiID string) (any, bool) {
r.mu.RLock()
defer r.mu.RUnlock()
n, ok := r.nodes[taiID]
if !ok || n.resources == nil {
return nil, false
}
return n.resources, true
}
// SetBridgeFunc sets the function used by OpenLocalListener to bridge
// TCP connections through the gRPC tunnel (Forward stream).
func (r *Registry) SetBridgeFunc(fn BridgeFunc) {
r.bridgeMu.Lock()
defer r.bridgeMu.Unlock()
r.bridgeFn = fn
}
// SetRegisterStream stores the gRPC Register stream for a tunnel node.
func (r *Registry) SetRegisterStream(taiID string, stream any) {
r.mu.Lock() r.mu.Lock()
defer r.mu.Unlock() defer r.mu.Unlock()
if n, ok := r.nodes[taiID]; ok { if n, ok := r.nodes[taiID]; ok {
n.client = c n.registerStream = stream
} }
} }
// GetRegisterStream returns the gRPC Register stream for a tunnel node.
func (r *Registry) GetRegisterStream(taiID string) any {
r.mu.RLock()
defer r.mu.RUnlock()
if n, ok := r.nodes[taiID]; ok {
return n.registerStream
}
return nil
}
// GenerateChannelID creates a random channel ID for Forward stream matching.
func GenerateChannelID() (string, error) {
return generateChannelID()
}
// FindTaiIDByAuthClient returns the TaiID of the first node whose // FindTaiIDByAuthClient returns the TaiID of the first node whose
// Auth.ClientID matches the given OAuth client ID. Returns "" if not found. // Auth.ClientID matches the given OAuth client ID. Returns "" if not found.
// This is needed because Tai's data channel authenticates with its OAuth
// ClientID, which may differ from the server-assigned TaiID.
func (r *Registry) FindTaiIDByAuthClient(clientID string) string { func (r *Registry) FindTaiIDByAuthClient(clientID string) string {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
@ -288,28 +269,28 @@ func (r *Registry) FindTaiIDByAuthClient(clientID string) string {
return "" return ""
} }
// ListByTeam returns snapshots of all nodes belonging to the given team. // ListByTeam returns metadata of all nodes belonging to the given team.
func (r *Registry) ListByTeam(teamID string) []NodeSnapshot { func (r *Registry) ListByTeam(teamID string) []types.NodeMeta {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
var result []NodeSnapshot var result []types.NodeMeta
for _, n := range r.nodes { for _, n := range r.nodes {
if n.Auth.TeamID == teamID { if n.Auth.TeamID == teamID {
result = append(result, n.snapshot()) result = append(result, n.meta())
} }
} }
return result return result
} }
// ListByUser returns snapshots of all nodes registered by the given user // ListByUser returns metadata of all nodes registered by the given user
// that are NOT associated with any team. // that are NOT associated with any team.
func (r *Registry) ListByUser(userID string) []NodeSnapshot { func (r *Registry) ListByUser(userID string) []types.NodeMeta {
r.mu.RLock() r.mu.RLock()
defer r.mu.RUnlock() defer r.mu.RUnlock()
var result []NodeSnapshot var result []types.NodeMeta
for _, n := range r.nodes { for _, n := range r.nodes {
if n.Auth.TeamID == "" && n.Auth.UserID == userID { if n.Auth.TeamID == "" && n.Auth.UserID == userID {
result = append(result, n.snapshot()) result = append(result, n.meta())
} }
} }
return result return result
@ -362,85 +343,6 @@ func (r *Registry) checkHealth(timeout, cleanupAfter time.Duration) {
} }
} }
// 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.
func (r *Registry) RequestChannel(taiID string, targetPort int) (string, chan net.Conn, error) {
r.mu.RLock()
node := r.nodes[taiID]
r.mu.RUnlock()
if node == nil {
return "", nil, fmt.Errorf("tai node %s not found", taiID)
}
if node.Mode != "tunnel" {
return "", nil, fmt.Errorf("tai node %s is not a tunnel node", taiID)
}
node.connMu.Lock()
hasConn := node.ControlConn != nil
node.connMu.Unlock()
if !hasConn {
return "", nil, fmt.Errorf("tai node %s has no active control channel", taiID)
}
channelID, err := generateChannelID()
if err != nil {
return "", nil, fmt.Errorf("generate channel_id: %w", err)
}
resultCh := make(chan net.Conn, 1)
timer := time.AfterFunc(30*time.Second, func() {
r.mu.Lock()
if pc, ok := r.pending[channelID]; ok {
close(pc.result)
delete(r.pending, channelID)
}
r.mu.Unlock()
})
r.mu.Lock()
r.pending[channelID] = &pendingChannel{taiID: taiID, result: resultCh, timer: timer}
r.mu.Unlock()
msg := map[string]interface{}{
"type": "open",
"channel_id": channelID,
"target_port": targetPort,
}
if err := r.WriteControlJSON(taiID, msg); err != nil {
r.mu.Lock()
delete(r.pending, channelID)
r.mu.Unlock()
timer.Stop()
return "", nil, fmt.Errorf("send open command: %w", err)
}
return channelID, resultCh, nil
}
// AcceptDataChannel resolves a pending channel when Tai connects its data WS.
// The taiID must match the node that requested the channel via RequestChannel.
func (r *Registry) AcceptDataChannel(channelID, taiID string, conn net.Conn) error {
r.mu.Lock()
pc, ok := r.pending[channelID]
if ok {
delete(r.pending, channelID)
}
r.mu.Unlock()
if !ok {
return fmt.Errorf("no pending channel for %s", channelID)
}
if pc.taiID != taiID {
pc.timer.Stop()
close(pc.result)
return fmt.Errorf("channel %s: tai_id mismatch (expected %s, got %s)", channelID, pc.taiID, taiID)
}
pc.timer.Stop()
pc.result <- conn
return nil
}
// OpenLocalListener creates a localhost TCP listener that tunnels every // OpenLocalListener creates a localhost TCP listener that tunnels every
// accepted connection to the specified port on the given Tai node. // accepted connection to the specified port on the given Tai node.
// Returns the listener address (e.g. "127.0.0.1:54321"). // Returns the listener address (e.g. "127.0.0.1:54321").
@ -486,37 +388,17 @@ func (r *Registry) OpenLocalListener(taiID string, targetPort int) (net.Listener
} }
func (r *Registry) bridgeTunnelConn(taiID string, targetPort int, localConn net.Conn) { func (r *Registry) bridgeTunnelConn(taiID string, targetPort int, localConn net.Conn) {
channelID, resultCh, err := r.RequestChannel(taiID, targetPort) r.bridgeMu.RLock()
if err != nil { fn := r.bridgeFn
localConn.Close() r.bridgeMu.RUnlock()
r.logger.Error("request channel failed", "tai_id", taiID, "port", targetPort, "err", err)
if fn != nil {
fn(taiID, targetPort, localConn)
return return
} }
remoteConn, ok := <-resultCh localConn.Close()
if !ok || remoteConn == nil { r.logger.Error("no bridge function configured", "tai_id", taiID, "port", targetPort)
localConn.Close()
r.logger.Error("data channel timeout", "tai_id", taiID, "channel_id", channelID)
return
}
bridgeTCP(localConn, remoteConn)
}
// bridgeTCP copies bytes bidirectionally between two net.Conn, closing both when done.
func bridgeTCP(a, b net.Conn) {
var wg sync.WaitGroup
wg.Add(2)
cp := func(dst, src net.Conn) {
defer wg.Done()
io.Copy(dst, src)
dst.Close()
}
go cp(a, b)
go cp(b, a)
wg.Wait()
} }
func generateChannelID() (string, error) { func generateChannelID() (string, error) {

View file

@ -2,23 +2,18 @@ package registry
import ( import (
"log/slog" "log/slog"
"net"
"net/http"
"net/http/httptest"
"strings"
"sync" "sync"
"testing" "testing"
"time" "time"
"github.com/gorilla/websocket" "github.com/yaoapp/yao/tai/types"
) )
// newTestRegistry creates a standalone registry for testing (bypasses global singleton). // newTestRegistry creates a standalone registry for testing (bypasses global singleton).
func newTestRegistry() *Registry { func newTestRegistry() *Registry {
return &Registry{ return &Registry{
nodes: make(map[string]*TaiNode), nodes: make(map[string]*TaiNode),
pending: make(map[string]*pendingChannel), logger: slog.Default(),
logger: slog.Default(),
} }
} }
@ -29,7 +24,7 @@ func TestRegister_SetsFieldsAndOnline(t *testing.T) {
MachineID: "m-abc", MachineID: "m-abc",
Version: "1.0.0", Version: "1.0.0",
Mode: "tunnel", Mode: "tunnel",
Ports: map[string]int{"grpc": 19100}, Ports: types.Ports{GRPC: 19100},
} }
r.Register(node) r.Register(node)
@ -117,14 +112,14 @@ func TestSnapshot_DeepCopy(t *testing.T) {
r := newTestRegistry() r := newTestRegistry()
r.Register(&TaiNode{ r.Register(&TaiNode{
TaiID: "tai-001", TaiID: "tai-001",
Ports: map[string]int{"grpc": 19100, "http": 8099}, Ports: types.Ports{GRPC: 19100, HTTP: 8099},
}) })
snap, _ := r.Get("tai-001") snap, _ := r.Get("tai-001")
snap.Ports["grpc"] = 0 snap.Ports.GRPC = 0
snap2, _ := r.Get("tai-001") snap2, _ := r.Get("tai-001")
if snap2.Ports["grpc"] != 19100 { if snap2.Ports.GRPC != 19100 {
t.Error("snapshot modification leaked into registry node") t.Error("snapshot modification leaked into registry node")
} }
} }
@ -146,98 +141,6 @@ func TestUpdatePing_NonexistentNode(t *testing.T) {
r.UpdatePing("ghost") r.UpdatePing("ghost")
} }
func TestWriteControlJSON_NoNode(t *testing.T) {
r := newTestRegistry()
err := r.WriteControlJSON("missing", map[string]string{"type": "test"})
if err == nil {
t.Fatal("expected error for missing node")
}
}
func TestWriteControlJSON_NilConn(t *testing.T) {
r := newTestRegistry()
r.Register(&TaiNode{TaiID: "tai-001"})
err := r.WriteControlJSON("tai-001", map[string]string{"type": "test"})
if err == nil {
t.Fatal("expected error for nil ControlConn")
}
}
func TestRequestChannel_NotFound(t *testing.T) {
r := newTestRegistry()
_, _, err := r.RequestChannel("ghost", 19100)
if err == nil {
t.Fatal("expected error for missing node")
}
}
func TestRequestChannel_DirectMode(t *testing.T) {
r := newTestRegistry()
r.Register(&TaiNode{TaiID: "tai-001", Mode: "direct"})
_, _, err := r.RequestChannel("tai-001", 19100)
if err == nil {
t.Fatal("expected error for direct-mode node")
}
}
func TestAcceptDataChannel_NotPending(t *testing.T) {
r := newTestRegistry()
pipe1, pipe2 := net.Pipe()
defer pipe1.Close()
defer pipe2.Close()
err := r.AcceptDataChannel("unknown-channel", "tai-001", pipe1)
if err == nil {
t.Fatal("expected error for non-pending channel")
}
}
func TestAcceptDataChannel_TaiIDMismatch(t *testing.T) {
r := newTestRegistry()
resultCh := make(chan net.Conn, 1)
timer := time.AfterFunc(5*time.Second, func() {})
r.mu.Lock()
r.pending["ch-001"] = &pendingChannel{taiID: "tai-owner", result: resultCh, timer: timer}
r.mu.Unlock()
pipe1, pipe2 := net.Pipe()
defer pipe1.Close()
defer pipe2.Close()
err := r.AcceptDataChannel("ch-001", "tai-intruder", pipe1)
if err == nil {
t.Fatal("expected error for tai_id mismatch")
}
}
func TestAcceptDataChannel_Success(t *testing.T) {
r := newTestRegistry()
resultCh := make(chan net.Conn, 1)
timer := time.AfterFunc(5*time.Second, func() {})
r.mu.Lock()
r.pending["ch-002"] = &pendingChannel{taiID: "tai-001", result: resultCh, timer: timer}
r.mu.Unlock()
pipe1, pipe2 := net.Pipe()
defer pipe2.Close()
if err := r.AcceptDataChannel("ch-002", "tai-001", pipe1); err != nil {
t.Fatalf("AcceptDataChannel: %v", err)
}
select {
case conn := <-resultCh:
if conn == nil {
t.Fatal("expected non-nil conn")
}
conn.Close()
case <-time.After(time.Second):
t.Fatal("timeout waiting for conn on resultCh")
}
}
func TestGenerateChannelID_Unique(t *testing.T) { func TestGenerateChannelID_Unique(t *testing.T) {
seen := make(map[string]bool) seen := make(map[string]bool)
for i := 0; i < 100; i++ { for i := 0; i < 100; i++ {
@ -255,26 +158,6 @@ func TestGenerateChannelID_Unique(t *testing.T) {
} }
} }
func TestBridgeTCP(t *testing.T) {
a1, a2 := net.Pipe()
b1, b2 := net.Pipe()
go bridgeTCP(a2, b1)
msg := []byte("hello tunnel")
go func() {
a1.Write(msg)
a1.Close()
}()
buf := make([]byte, 64)
n, _ := b2.Read(buf)
if string(buf[:n]) != "hello tunnel" {
t.Errorf("got %q, want %q", buf[:n], "hello tunnel")
}
b2.Close()
}
func TestConcurrentRegisterGet(t *testing.T) { func TestConcurrentRegisterGet(t *testing.T) {
r := newTestRegistry() r := newTestRegistry()
var wg sync.WaitGroup var wg sync.WaitGroup
@ -298,164 +181,6 @@ func TestConcurrentRegisterGet(t *testing.T) {
wg.Wait() wg.Wait()
} }
func TestWriteControlJSON_Success(t *testing.T) {
done := make(chan map[string]string, 1)
srv := newWSServer(func(conn *websocket.Conn) {
var msg map[string]string
conn.ReadJSON(&msg)
done <- msg
conn.Close()
})
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err != nil {
t.Fatalf("dial: %v", err)
}
r := newTestRegistry()
r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn})
payload := map[string]string{"type": "test", "data": "hello"}
if err := r.WriteControlJSON("tai-001", payload); err != nil {
t.Fatalf("WriteControlJSON: %v", err)
}
select {
case got := <-done:
if got["type"] != "test" {
t.Errorf("type = %q, want test", got["type"])
}
if got["data"] != "hello" {
t.Errorf("data = %q, want hello", got["data"])
}
case <-time.After(2 * time.Second):
t.Fatal("timeout waiting for server to receive message")
}
}
func TestRequestChannel_Success(t *testing.T) {
openCh := make(chan map[string]interface{}, 1)
srv := newWSServer(func(conn *websocket.Conn) {
var msg map[string]interface{}
conn.ReadJSON(&msg)
openCh <- msg
time.Sleep(time.Second)
conn.Close()
})
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err != nil {
t.Fatalf("dial: %v", err)
}
r := newTestRegistry()
r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn})
channelID, resultCh, err := r.RequestChannel("tai-001", 19100)
if err != nil {
t.Fatalf("RequestChannel: %v", err)
}
if channelID == "" {
t.Fatal("channelID should not be empty")
}
if len(channelID) != 64 {
t.Errorf("channelID len = %d, want 64", len(channelID))
}
if resultCh == nil {
t.Fatal("resultCh should not be nil")
}
select {
case cmd := <-openCh:
if cmd["type"] != "open" {
t.Errorf("cmd type = %v, want open", cmd["type"])
}
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) != 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")
}
}
func TestRequestChannel_NoControlConn(t *testing.T) {
r := newTestRegistry()
r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel"})
_, _, err := r.RequestChannel("tai-001", 19100)
if err == nil {
t.Fatal("expected error for nil ControlConn")
}
}
func TestOpenLocalListener_Success(t *testing.T) {
r := newTestRegistry()
controlCh := make(chan map[string]interface{}, 1)
srv := newWSServer(func(conn *websocket.Conn) {
for {
var msg map[string]interface{}
if err := conn.ReadJSON(&msg); err != nil {
return
}
controlCh <- msg
}
})
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
wsConn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err != nil {
t.Fatalf("dial: %v", err)
}
r.Register(&TaiNode{TaiID: "tai-001", Mode: "tunnel", ControlConn: wsConn})
ln, err := r.OpenLocalListener("tai-001", 19100)
if err != nil {
t.Fatalf("OpenLocalListener: %v", err)
}
defer ln.Close()
addr := ln.Addr().String()
if addr == "" {
t.Fatal("listener address should not be empty")
}
if !strings.HasPrefix(addr, "127.0.0.1:") {
t.Errorf("addr = %q, want 127.0.0.1:*", addr)
}
conn, err := net.DialTimeout("tcp", addr, time.Second)
if err != nil {
t.Fatalf("connect to local listener: %v", err)
}
defer conn.Close()
select {
case cmd := <-controlCh:
if cmd["type"] != "open" {
t.Errorf("open cmd type = %v, want open", cmd["type"])
}
if _, ok := cmd["channel_id"].(string); !ok {
t.Error("open cmd missing channel_id")
}
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")
}
}
func TestOpenLocalListener_NodeNotFound(t *testing.T) { func TestOpenLocalListener_NodeNotFound(t *testing.T) {
r := newTestRegistry() r := newTestRegistry()
_, err := r.OpenLocalListener("ghost", 19100) _, err := r.OpenLocalListener("ghost", 19100)
@ -464,22 +189,11 @@ func TestOpenLocalListener_NodeNotFound(t *testing.T) {
} }
} }
func newWSServer(handler func(*websocket.Conn)) *httptest.Server {
up := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
conn, err := up.Upgrade(w, r, nil)
if err != nil {
return
}
handler(conn)
}))
}
func TestRegister_SystemInfo(t *testing.T) { func TestRegister_SystemInfo(t *testing.T) {
r := newTestRegistry() r := newTestRegistry()
r.Register(&TaiNode{ r.Register(&TaiNode{
TaiID: "tai-001", TaiID: "tai-001",
System: SystemInfo{ System: types.SystemInfo{
OS: "linux", OS: "linux",
Arch: "amd64", Arch: "amd64",
Hostname: "docker-host-01", Hostname: "docker-host-01",
@ -507,9 +221,9 @@ func TestRegister_SystemInfo(t *testing.T) {
func TestListByTeam(t *testing.T) { func TestListByTeam(t *testing.T) {
r := newTestRegistry() r := newTestRegistry()
r.Register(&TaiNode{TaiID: "tai-a", Auth: AuthInfo{TeamID: "team-dev"}}) r.Register(&TaiNode{TaiID: "tai-a", Auth: types.AuthInfo{TeamID: "team-dev"}})
r.Register(&TaiNode{TaiID: "tai-b", Auth: AuthInfo{TeamID: "team-dev"}}) r.Register(&TaiNode{TaiID: "tai-b", Auth: types.AuthInfo{TeamID: "team-dev"}})
r.Register(&TaiNode{TaiID: "tai-c", Auth: AuthInfo{TeamID: "team-ops"}}) r.Register(&TaiNode{TaiID: "tai-c", Auth: types.AuthInfo{TeamID: "team-ops"}})
devNodes := r.ListByTeam("team-dev") devNodes := r.ListByTeam("team-dev")
if len(devNodes) != 2 { if len(devNodes) != 2 {
@ -604,11 +318,11 @@ func TestStartHealthCheck_PingKeepsAlive(t *testing.T) {
} }
} }
func TestNodeSnapshot_AuthInfo(t *testing.T) { func TestNodeMeta_AuthInfo(t *testing.T) {
r := newTestRegistry() r := newTestRegistry()
r.Register(&TaiNode{ r.Register(&TaiNode{
TaiID: "tai-001", TaiID: "tai-001",
Auth: AuthInfo{ Auth: types.AuthInfo{
Subject: "user123", Subject: "user123",
ClientID: "tai-001", ClientID: "tai-001",
Scope: "tai:tunnel", Scope: "tai:tunnel",

View file

@ -2,17 +2,14 @@ package registry
import ( import (
"log/slog" "log/slog"
"net"
"time"
) )
// NewForTest creates a standalone Registry for use in tests. // NewForTest creates a standalone Registry for use in tests.
// Not intended for production use. // Not intended for production use.
func NewForTest() *Registry { func NewForTest() *Registry {
return &Registry{ return &Registry{
nodes: make(map[string]*TaiNode), nodes: make(map[string]*TaiNode),
pending: make(map[string]*pendingChannel), logger: slog.Default(),
logger: slog.Default(),
} }
} }
@ -21,11 +18,3 @@ func NewForTest() *Registry {
func SetGlobalForTest(r *Registry) { func SetGlobalForTest(r *Registry) {
global = r global = r
} }
// SetPendingForTest injects a pending channel entry for testing.
// Not intended for production use.
func (r *Registry) SetPendingForTest(channelID, taiID string, result chan net.Conn, timer *time.Timer) {
r.mu.Lock()
defer r.mu.Unlock()
r.pending[channelID] = &pendingChannel{taiID: taiID, result: result, timer: timer}
}

View file

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

View file

@ -1,4 +1,4 @@
package sandbox package runtime
import ( import (
"context" "context"
@ -12,9 +12,9 @@ type dockerSandbox struct {
core dockerCore core dockerCore
} }
// NewDocker creates a Sandbox backed by Docker SDK through Tai's Docker API proxy. // NewDocker creates a Runtime backed by Docker SDK through Tai's Docker API proxy.
// addr should be "tcp://tai-host:12375". // addr should be "tcp://tai-host:12375".
func NewDocker(addr string) (Sandbox, error) { func NewDocker(addr string) (Runtime, error) {
cli, err := client.NewClientWithOpts( cli, err := client.NewClientWithOpts(
client.WithHost(addr), client.WithHost(addr),
client.WithAPIVersionNegotiation(), client.WithAPIVersionNegotiation(),

View file

@ -1,4 +1,4 @@
package sandbox package runtime
import ( import (
"bytes" "bytes"
@ -15,7 +15,7 @@ import (
"github.com/docker/go-connections/nat" "github.com/docker/go-connections/nat"
) )
// dockerCore contains Docker SDK operations shared by both Local and Docker (via Tai) sandboxes. // dockerCore contains Docker SDK operations shared by both Local and Docker (via Tai) runtimes.
type dockerCore struct { type dockerCore struct {
cli *client.Client cli *client.Client
} }

View file

@ -1,4 +1,4 @@
package sandbox package runtime
import ( import (
"context" "context"

View file

@ -1,4 +1,4 @@
package sandbox package runtime
import ( import (
"context" "context"
@ -14,7 +14,7 @@ import (
) )
// dockerImage implements Image using the Docker SDK. // dockerImage implements Image using the Docker SDK.
// Shared by both local and dockerSandbox (via Tai proxy) modes. // Shared by both local and docker (via Tai proxy) runtime modes.
type dockerImage struct { type dockerImage struct {
cli *client.Client cli *client.Client
} }

View file

@ -1,4 +1,4 @@
package sandbox package runtime
import "context" import "context"

View file

@ -1,4 +1,4 @@
package sandbox package runtime
import ( import (
"bytes" "bytes"
@ -20,7 +20,7 @@ import (
"k8s.io/client-go/tools/remotecommand" "k8s.io/client-go/tools/remotecommand"
) )
// K8sOption configures a K8s sandbox. // K8sOption configures a K8s runtime.
type K8sOption struct { type K8sOption struct {
Namespace string // default "default" Namespace string // default "default"
KubeConfig string // path to kubeconfig file KubeConfig string // path to kubeconfig file
@ -33,10 +33,10 @@ type k8sSandbox struct {
labels map[string]string labels map[string]string
} }
// NewK8s creates a Sandbox backed by Kubernetes via Tai's TCP proxy. // NewK8s creates a Runtime backed by Kubernetes via Tai's TCP proxy.
// addr should be "host:port" pointing to Tai's K8s proxy endpoint. // addr should be "host:port" pointing to Tai's K8s proxy endpoint.
// kubeConfigPath must be an absolute path or will be resolved relative to the caller's working directory. // kubeConfigPath must be an absolute path or will be resolved relative to the caller's working directory.
func NewK8s(addr string, opts ...K8sOption) (Sandbox, error) { func NewK8s(addr string, opts ...K8sOption) (Runtime, error) {
ns := "default" ns := "default"
var kubeConfigPath string var kubeConfigPath string
if len(opts) > 0 { if len(opts) > 0 {
@ -56,7 +56,7 @@ func NewK8s(addr string, opts ...K8sOption) (Sandbox, error) {
} }
if kubeConfigPath == "" { if kubeConfigPath == "" {
return nil, fmt.Errorf("kubeconfig path is required for K8s sandbox") return nil, fmt.Errorf("kubeconfig path is required for K8s runtime")
} }
cfg, err := clientcmd.BuildConfigFromFlags("", kubeConfigPath) cfg, err := clientcmd.BuildConfigFromFlags("", kubeConfigPath)

View file

@ -1,4 +1,4 @@
package sandbox package runtime
import ( import (
"context" "context"
@ -12,9 +12,9 @@ type local struct {
core dockerCore core dockerCore
} }
// NewLocal creates a Sandbox backed by a direct Docker daemon connection. // NewLocal creates a Runtime backed by a direct Docker daemon connection.
// addr can be "unix:///var/run/docker.sock", "tcp://host:port", or "" for platform default. // addr can be "unix:///var/run/docker.sock", "tcp://host:port", or "" for platform default.
func NewLocal(addr string) (Sandbox, error) { func NewLocal(addr string) (Runtime, error) {
opts := []client.Opt{client.WithAPIVersionNegotiation()} opts := []client.Opt{client.WithAPIVersionNegotiation()}
if addr != "" { if addr != "" {
opts = append(opts, client.WithHost(addr)) opts = append(opts, client.WithHost(addr))

View file

@ -1,4 +1,4 @@
package sandbox package runtime
import ( import (
"context" "context"
@ -53,7 +53,7 @@ func TestHelpers(t *testing.T) {
}) })
} }
func TestLocalSandbox(t *testing.T) { func TestLocalRuntime(t *testing.T) {
sb, err := NewLocal("") sb, err := NewLocal("")
if err != nil { if err != nil {
t.Skipf("Docker not available: %v", err) t.Skipf("Docker not available: %v", err)
@ -258,7 +258,7 @@ func TestLocalCreateWithEnvAndWorkDir(t *testing.T) {
} }
} }
func TestDockerSandboxViaTai(t *testing.T) { func TestDockerRuntimeViaTai(t *testing.T) {
addr := taiTestDocker() addr := taiTestDocker()
sb, err := NewDocker(addr) sb, err := NewDocker(addr)
if err != nil { if err != nil {
@ -352,7 +352,7 @@ func TestPortStr(t *testing.T) {
} }
} }
func TestK8sSandbox(t *testing.T) { func TestK8sRuntime(t *testing.T) {
host := taiTestK8sHost() host := taiTestK8sHost()
port := taiTestK8sPort() port := taiTestK8sPort()
kubeconfig := taiTestKubeConfig() kubeconfig := taiTestKubeConfig()
@ -496,7 +496,7 @@ func TestK8sBuildResourcesPartial(t *testing.T) {
} }
} }
func TestK8sSandboxStopAndRemove(t *testing.T) { func TestK8sRuntimeStopAndRemove(t *testing.T) {
host := taiTestK8sHost() host := taiTestK8sHost()
port := taiTestK8sPort() port := taiTestK8sPort()
kubeconfig := taiTestKubeConfig() kubeconfig := taiTestKubeConfig()

View file

@ -1,4 +1,4 @@
package sandbox package runtime
import ( import (
"context" "context"
@ -6,9 +6,9 @@ import (
"time" "time"
) )
// Sandbox manages container lifecycle. // Runtime manages container lifecycle.
// Local connects directly to a Docker daemon; Docker/Containerd/K8s connect via Tai proxy. // Local connects directly to a Docker daemon; Docker/Containerd/K8s connect via Tai proxy.
type Sandbox interface { type Runtime interface {
Create(ctx context.Context, opts CreateOptions) (string, error) Create(ctx context.Context, opts CreateOptions) (string, error)
Start(ctx context.Context, id string) error Start(ctx context.Context, id string) error
Stop(ctx context.Context, id string, timeout time.Duration) error Stop(ctx context.Context, id string, timeout time.Duration) error

View file

@ -1,38 +1,16 @@
package tai package tai
import ( import (
"context"
"fmt"
"net"
"net/http"
"net/url"
"strconv"
"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/registry"
"github.com/yaoapp/yao/tai/sandbox" "github.com/yaoapp/yao/tai/types"
sipb "github.com/yaoapp/yao/tai/serverinfo/pb"
"github.com/yaoapp/yao/tai/vnc"
"github.com/yaoapp/yao/tai/volume" "github.com/yaoapp/yao/tai/volume"
"github.com/yaoapp/yao/tai/workspace"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
) )
// Runtime selects which container runtime to use via Tai. // Type aliases kept at package level for convenience.
type Runtime int type Runtime = types.Runtime
type Ports = types.Ports
const ( // Option configures RegisterLocal.
Docker Runtime = iota
K8s
)
func (r Runtime) apply(c *config) { c.runtime = r }
// Option configures a Client.
type Option interface { type Option interface {
apply(*config) apply(*config)
} }
@ -41,60 +19,19 @@ type optionFunc func(*config)
func (f optionFunc) apply(c *config) { f(c) } func (f optionFunc) apply(c *config) { f(c) }
// Ports configures service ports for Tai server.
type Ports struct {
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.
// Ports set here take precedence over server-reported values from ServerInfo.
func WithPorts(p Ports) Option {
return optionFunc(func(c *config) {
c.ports = p
c.userPorts = p
})
}
// WithHTTPClient sets a custom HTTP client for proxy and VNC health checks.
func WithHTTPClient(hc *http.Client) Option {
return optionFunc(func(c *config) { c.httpClient = hc })
}
// WithDataDir sets the workspace root directory for Local mode. // WithDataDir sets the workspace root directory for Local mode.
func WithDataDir(dir string) Option { func WithDataDir(dir string) Option {
return optionFunc(func(c *config) { c.dataDir = dir }) return optionFunc(func(c *config) { c.dataDir = dir })
} }
// WithKubeConfig sets the kubeconfig file path for K8s runtime. // WithVolume injects a custom Volume implementation (useful for testing).
// Supports both absolute and relative paths (relative paths are resolved to absolute).
func WithKubeConfig(path string) Option {
return optionFunc(func(c *config) { c.kubeConfig = path })
}
// WithNamespace sets the namespace for K8s runtime. Default is "default".
func WithNamespace(ns string) Option {
return optionFunc(func(c *config) { c.namespace = ns })
}
// WithVolume injects a custom Volume implementation.
// Useful for testing workspace operations without Docker.
func WithVolume(vol volume.Volume) Option { func WithVolume(vol volume.Volume) Option {
return optionFunc(func(c *config) { c.volume = vol }) return optionFunc(func(c *config) { c.volume = vol })
} }
type config struct { type config struct {
runtime Runtime dataDir string
ports Ports volume volume.Volume
userPorts Ports // tracks explicitly set ports (zero = not set by user)
httpClient *http.Client
dataDir string
kubeConfig string
namespace string
volume volume.Volume // override volume (for testing without Docker)
} }
func defaultPorts() Ports { func defaultPorts() Ports {
@ -125,504 +62,15 @@ func mergedPorts(p Ports) Ports {
return d return d
} }
// Client provides unified access to all Tai SDK sub-packages. func intOr(v, fallback int) int {
type Client struct { if v > 0 {
scheme string // "tai", "docker", or "tunnel" return v
host string
addr string
taiID string // registry key — set by initLocal/initRemote/initTunnel
ports Ports
dataDir string // host-side data directory for local volume
vol volume.Volume
sb sandbox.Sandbox
img sandbox.Image
prx proxy.Proxy
vc vnc.VNC
he hepb.HostExecClient
grpcConn *grpc.ClientConn
// tunnel mode: local listeners that bridge to Tai via WS
tunnelListeners []net.Listener
}
// New creates a Client based on the address protocol:
//
// "local" → Local mode, platform default Docker socket
// "docker://addr" → Local mode, specified Docker daemon
// "tai://host" → Remote mode via Tai Server
//
// Empty string is not allowed — use "local" for default local Docker.
func New(addr string, opts ...Option) (*Client, error) {
cfg := &config{ports: defaultPorts()}
for _, o := range opts {
o.apply(cfg)
}
cfg.ports = mergedPorts(cfg.ports)
scheme, host, dockerAddr, grpcPort, err := parseAddr(addr)
if err != nil {
return nil, err
}
if grpcPort > 0 {
cfg.ports.GRPC = grpcPort
}
c := &Client{
scheme: scheme,
host: host,
addr: dockerAddr,
ports: cfg.ports,
}
switch scheme {
case "docker":
return c.initLocal(cfg)
case "tai":
return c.initRemote(cfg)
case "tunnel":
return c.initTunnel(cfg)
default:
return nil, fmt.Errorf("unsupported scheme: %s", scheme)
}
}
func (c *Client) initLocal(cfg *config) (*Client, error) {
sb, err := sandbox.NewLocal(c.addr)
if err != nil && cfg.volume == nil {
return nil, err
}
if sb != nil {
c.sb = sb
c.img = sandbox.NewDockerImage(sandbox.DockerCli(sb))
c.prx = proxy.NewLocal(sb)
c.vc = vnc.NewLocal(sb)
}
if cfg.volume != nil {
c.vol = cfg.volume
c.dataDir = cfg.dataDir
} else {
dataDir := cfg.dataDir
if dataDir == "" {
dataDir = "/tmp/tai-volumes"
}
c.dataDir = dataDir
c.vol = volume.NewLocal(dataDir)
}
if reg := registry.Global(); reg != nil {
id := c.host
if id == "" {
id = c.addr
}
if id == "" {
id = "local"
}
c.taiID = id
reg.Register(&registry.TaiNode{
TaiID: id,
Mode: "local",
Addr: c.addr,
})
reg.SetClient(id, c)
}
return c, nil
}
func (c *Client) initRemote(cfg *config) (*Client, error) {
grpcAddr := fmt.Sprintf("%s:%d", c.host, c.ports.GRPC)
conn, err := grpc.NewClient(grpcAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
return nil, fmt.Errorf("grpc dial %s: %w", grpcAddr, err)
}
c.grpcConn = conn
c.he = hepb.NewHostExecClient(conn)
info, err := c.discoverServerInfo(conn, cfg)
if err != nil {
info = &discoveredInfo{Capabilities: map[string]bool{"docker": true}}
}
hasDocker := info.Capabilities["docker"]
hasK8s := info.Capabilities["k8s"]
hasHostExec := info.Capabilities["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)
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 = 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 {
c.sb = sb
c.img = sandbox.NewK8sImage()
}
} else if hasDocker {
dockerPort := c.ports.Docker
if dockerPort == 0 {
dockerPort = 12375
}
sbAddr := fmt.Sprintf("tcp://%s:%d", c.host, dockerPort)
sb, err := sandbox.NewDocker(sbAddr)
if err == nil {
c.sb = sb
c.img = sandbox.NewDockerImage(sandbox.DockerCli(sb))
}
}
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 {
id := fmt.Sprintf("%s-%d", c.host, c.ports.GRPC)
c.taiID = id
reg.Register(&registry.TaiNode{
TaiID: id,
Mode: "direct",
Version: info.Version,
System: info.System,
Capabilities: info.Capabilities,
Addr: fmt.Sprintf("tai://%s:%d", c.host, c.ports.GRPC),
Ports: map[string]int{
"grpc": c.ports.GRPC,
"http": c.ports.HTTP,
"vnc": c.ports.VNC,
"docker": c.ports.Docker,
"k8s": c.ports.K8s,
},
})
reg.SetClient(id, c)
}
return c, nil
}
func (c *Client) initTunnel(cfg *config) (*Client, error) {
reg := registry.Global()
if reg == nil {
return nil, fmt.Errorf("tai registry not initialized")
}
taiID := c.host // for tunnel:// scheme, host stores the taiID
c.taiID = taiID
node, ok := reg.Get(taiID)
if !ok || node.Status != "online" {
return nil, fmt.Errorf("tai node %s not online", taiID)
}
c.ports = Ports{
GRPC: nodePort(node.Ports, "grpc", 19100),
HTTP: nodePort(node.Ports, "http", 8099),
VNC: nodePort(node.Ports, "vnc", 16080),
Docker: nodePort(node.Ports, "docker", 12375),
K8s: nodePort(node.Ports, "k8s", 16443),
}
grpcLn, err := reg.OpenLocalListener(taiID, c.ports.GRPC)
if err != nil {
return nil, fmt.Errorf("open grpc tunnel listener: %w", err)
}
c.tunnelListeners = append(c.tunnelListeners, grpcLn)
grpcAddr := grpcLn.Addr().String()
conn, err := grpc.NewClient("passthrough:///"+grpcAddr,
grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
grpcLn.Close()
return nil, fmt.Errorf("grpc dial tunnel %s: %w", grpcAddr, err)
}
c.grpcConn = conn
c.he = hepb.NewHostExecClient(conn)
c.vol = volume.NewRemote(conn)
info, err := c.discoverServerInfo(conn, cfg)
if err != nil {
info = &discoveredInfo{Capabilities: map[string]bool{"docker": true}}
}
hasDocker := info.Capabilities["docker"]
hasK8s := info.Capabilities["k8s"]
hasHostExec := info.Capabilities["host_exec"]
if !hasDocker && !hasK8s && !hasHostExec {
c.closeTunnelListeners()
conn.Close()
return nil, fmt.Errorf("tai %s: no capabilities available via tunnel (docker/k8s/host_exec all false)", taiID)
}
if cfg.runtime == K8s || (!hasDocker && hasK8s) {
k8sLn, err := reg.OpenLocalListener(taiID, c.ports.K8s)
if err == nil {
c.tunnelListeners = append(c.tunnelListeners, k8sLn)
sbAddr := k8sLn.Addr().String()
sb, err := sandbox.NewK8s(sbAddr, sandbox.K8sOption{
Namespace: cfg.namespace,
KubeConfig: cfg.kubeConfig,
})
if err == nil {
c.sb = sb
c.img = sandbox.NewK8sImage()
}
}
} else 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)
}
reg.SetClient(taiID, c)
return c, nil
}
func (c *Client) closeTunnelListeners() {
for _, ln := range c.tunnelListeners {
ln.Close()
}
c.tunnelListeners = nil
}
func nodePort(ports map[string]int, key string, fallback int) int {
if p, ok := ports[key]; ok && p > 0 {
return p
} }
return fallback return fallback
} }
// Close releases all resources.
func (c *Client) Close() error {
var errs []error
if c.sb != nil {
if err := c.sb.Close(); err != nil {
errs = append(errs, err)
}
}
if c.vol != nil {
if err := c.vol.Close(); err != nil {
errs = append(errs, err)
}
}
if c.grpcConn != nil {
if err := c.grpcConn.Close(); err != nil {
errs = append(errs, err)
}
}
c.closeTunnelListeners()
if c.taiID != "" {
if reg := registry.Global(); reg != nil {
reg.Unregister(c.taiID)
}
}
if len(errs) > 0 {
return fmt.Errorf("close: %v", errs)
}
return nil
}
// Volume returns the Volume IO layer. Never nil.
func (c *Client) Volume() volume.Volume { return c.vol }
// DataDir returns the host-side data directory used by the local volume.
// Empty for remote (Tai gRPC) connections — the Tai server manages paths.
func (c *Client) DataDir() string { return c.dataDir }
// Host returns the raw host parsed from the address (IP or hostname).
func (c *Client) Host() string { return c.host }
// TaiID returns the registry key for this client.
func (c *Client) TaiID() string { return c.taiID }
// Workspace returns an fs.FS-compatible filesystem for the given session.
func (c *Client) Workspace(sessionID string) workspace.FS {
return workspace.New(c.vol, sessionID)
}
// 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.
// 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.
// Nil when the Tai server has no container runtime.
func (c *Client) Proxy() proxy.Proxy { return c.prx }
// 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" }
func parseAddr(addr string) (scheme, host, dockerAddr string, grpcPort int, err error) {
addr = strings.TrimSpace(addr)
if addr == "" {
return "", "", "", 0, fmt.Errorf("empty address: use \"local\" for default Docker daemon")
}
if addr == "local" {
return "docker", "", "", 0, nil
}
// Bare IP or host(:port) without scheme → normalise before url.Parse,
// which misparses bare addresses (treats them as path, not host).
if !strings.Contains(addr, "://") {
if isLocalHost(addr) {
return "docker", "", "", 0, nil
}
// host:port — split carefully (IPv6 like [::1]:19100 is already handled above)
h := addr
if idx := strings.LastIndex(addr, ":"); idx > 0 {
h = addr[:idx]
}
if isLocalHost(h) {
return "docker", "", "", 0, nil
}
addr = "tai://" + addr
}
u, parseErr := url.Parse(addr)
if parseErr != nil {
return "", "", "", 0, fmt.Errorf("parse addr %q: %w", addr, parseErr)
}
switch u.Scheme {
case "tai":
hostname := u.Hostname()
if hostname == "" {
return "", "", "", 0, fmt.Errorf("tai:// requires a host")
}
if portStr := u.Port(); portStr != "" {
if p, convErr := strconv.Atoi(portStr); convErr == nil && p > 0 {
grpcPort = p
}
}
return "tai", hostname, "", grpcPort, nil
case "tunnel":
taiID := u.Host
if taiID == "" {
return "", "", "", 0, fmt.Errorf("tunnel:// requires a tai ID")
}
return "tunnel", taiID, "", 0, nil
case "docker":
return "docker", "", addr, 0, nil
case "unix":
return "docker", "", addr, 0, nil
case "tcp":
return "docker", "", addr, 0, nil
case "npipe":
return "docker", "", addr, 0, nil
default:
return "", "", "", 0, fmt.Errorf("unsupported scheme %q in addr %q", u.Scheme, addr)
}
}
func isLocalHost(h string) bool {
return h == "127.0.0.1" || h == "localhost" || h == "::1"
}
type discoveredInfo struct {
Capabilities map[string]bool
System registry.SystemInfo
Version string
}
// discoverServerInfo calls ServerInfo.GetInfo on the remote Tai server, merges
// discovered ports into c.ports, and returns capabilities + system info.
// Ports explicitly set via WithPorts take precedence over server-reported values.
func (c *Client) discoverServerInfo(conn *grpc.ClientConn, cfg *config) (*discoveredInfo, 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 nil, err
}
up := cfg.userPorts
if p := int(resp.Ports["http"]); p > 0 && up.HTTP == 0 {
c.ports.HTTP = p
}
if p := int(resp.Ports["docker"]); p > 0 && up.Docker == 0 {
c.ports.Docker = p
}
if p := int(resp.Ports["vnc"]); p > 0 && up.VNC == 0 {
c.ports.VNC = p
}
if p := int(resp.Ports["k8s"]); p > 0 && up.K8s == 0 {
c.ports.K8s = p
}
caps := resp.Capabilities
if caps == nil {
caps = make(map[string]bool)
}
var sys registry.SystemInfo
if s := resp.System; s != nil {
sys = registry.SystemInfo{
OS: s.Os,
Arch: s.Arch,
Hostname: s.Hostname,
NumCPU: int(s.NumCpu),
TotalMem: s.TotalMem,
Shell: s.Shell,
TempDir: s.TempDir,
}
}
return &discoveredInfo{
Capabilities: caps,
System: sys,
Version: resp.Version,
}, nil
}
// RegisterLocal probes the local Docker environment and, if reachable, // RegisterLocal probes the local Docker environment and, if reachable,
// creates a Client and registers it as the "local" node in the registry. // registers it as the "local" node in the registry with ConnResources.
// Returns true if a local node was successfully registered. // Returns true if a local node was successfully registered.
// Silently returns false if Docker is not available — this is not an error. // Silently returns false if Docker is not available — this is not an error.
func RegisterLocal(opts ...Option) bool { func RegisterLocal(opts ...Option) bool {
@ -634,34 +82,40 @@ func RegisterLocal(opts ...Option) bool {
return true return true
} }
c, err := New("local", opts...) cfg := &config{}
for _, o := range opts {
o.apply(cfg)
}
res, err := DialLocal("", cfg.dataDir, cfg.volume)
if err != nil { if err != nil {
return false return false
} }
_ = c // registered by initLocal → reg.Register + reg.SetClient
reg.Register(&registry.TaiNode{
TaiID: "local",
Mode: "local",
})
reg.SetResources("local", res)
return true return true
} }
// GetClient returns a registered *Client by taiID from the global registry. // GetResources returns the ConnResources for a registered Tai node.
func GetClient(taiID string) (*Client, bool) { func GetResources(taiID string) (*ConnResources, bool) {
reg := registry.Global() reg := registry.Global()
if reg == nil { if reg == nil {
return nil, false return nil, false
} }
snap, ok := reg.Get(taiID) raw, ok := reg.GetResources(taiID)
if !ok { if !ok {
return nil, false return nil, false
} }
c, ok := snap.Client().(*Client) res, ok := raw.(*ConnResources)
if !ok || c == nil { return res, ok && res != nil
return nil, false
}
return c, true
} }
// GetNodeSnapshot returns the registry snapshot for a Tai node by ID. // GetNodeMeta returns the metadata for a registered Tai node by ID.
// Callers can inspect System, Capabilities, Mode and other registry-level fields. func GetNodeMeta(taiID string) (*types.NodeMeta, bool) {
func GetNodeSnapshot(taiID string) (*registry.NodeSnapshot, bool) {
reg := registry.Global() reg := registry.Global()
if reg == nil { if reg == nil {
return nil, false return nil, false

View file

@ -1,12 +1,13 @@
package tai package tai
import ( import (
"fmt"
"os" "os"
"strconv" "strconv"
"testing" "testing"
"github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/registry"
"github.com/yaoapp/yao/tai/types"
"github.com/yaoapp/yao/tai/volume"
) )
func taiTestHost() string { func taiTestHost() string {
@ -16,18 +17,6 @@ func taiTestHost() string {
return "127.0.0.1" 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 { func taiTestPorts() Ports {
return Ports{ return Ports{
Docker: envPort("TAI_TEST_DOCKER_PORT", 0), Docker: envPort("TAI_TEST_DOCKER_PORT", 0),
@ -45,62 +34,6 @@ func envPort(key string, fallback int) int {
return fallback return fallback
} }
func TestParseAddr(t *testing.T) {
tests := []struct {
addr string
wantScheme string
wantHost string
wantDocker string
wantGRPCPort int
wantErr bool
}{
{"", "", "", "", 0, true},
{"local", "docker", "", "", 0, false},
{"127.0.0.1", "docker", "", "", 0, false},
{"localhost", "docker", "", "", 0, false},
{"::1", "docker", "", "", 0, false},
{"docker:///var/run/docker.sock", "docker", "", "docker:///var/run/docker.sock", 0, false},
{"docker://192.168.1.50:2375", "docker", "", "docker://192.168.1.50:2375", 0, false},
{"unix:///var/run/docker.sock", "docker", "", "unix:///var/run/docker.sock", 0, false},
{"tcp://127.0.0.1:2375", "docker", "", "tcp://127.0.0.1:2375", 0, false},
{"npipe:////./pipe/docker_engine", "docker", "", "npipe:////./pipe/docker_engine", 0, false},
{"tai://192.168.1.100", "tai", "192.168.1.100", "", 0, false},
{"tai://10.0.0.5:9200", "tai", "10.0.0.5", "", 9200, false},
{"tai://", "", "", "", 0, true},
{"ftp://host", "", "", "", 0, true},
{" tai://host ", "tai", "host", "", 0, false},
// Bare non-local host → auto-prepend tai://
{"192.168.1.50", "tai", "192.168.1.50", "", 0, false},
{"192.168.1.50:9200", "tai", "192.168.1.50", "", 9200, false},
{"my-server", "tai", "my-server", "", 0, false},
{"my-server:9200", "tai", "my-server", "", 9200, false},
}
for _, tt := range tests {
t.Run(tt.addr, func(t *testing.T) {
scheme, host, dockerAddr, grpcPort, err := parseAddr(tt.addr)
if (err != nil) != tt.wantErr {
t.Fatalf("err = %v, wantErr = %v", err, tt.wantErr)
}
if err != nil {
return
}
if scheme != tt.wantScheme {
t.Errorf("scheme = %q, want %q", scheme, tt.wantScheme)
}
if host != tt.wantHost {
t.Errorf("host = %q, want %q", host, tt.wantHost)
}
if dockerAddr != tt.wantDocker {
t.Errorf("dockerAddr = %q, want %q", dockerAddr, tt.wantDocker)
}
if grpcPort != tt.wantGRPCPort {
t.Errorf("grpcPort = %d, want %d", grpcPort, tt.wantGRPCPort)
}
})
}
}
func TestMergedPorts(t *testing.T) { func TestMergedPorts(t *testing.T) {
p := mergedPorts(Ports{HTTP: 8888}) p := mergedPorts(Ports{HTTP: 8888})
if p.HTTP != 8888 { if p.HTTP != 8888 {
@ -127,97 +60,73 @@ func TestMergedPortsAll(t *testing.T) {
} }
} }
func TestOptions(t *testing.T) { func TestDialLocalSuccess(t *testing.T) {
cfg := &config{ports: defaultPorts()} res, err := DialLocal("", t.TempDir(), nil)
WithPorts(Ports{HTTP: 9999}).apply(cfg)
if cfg.ports.HTTP != 9999 {
t.Errorf("WithPorts: HTTP = %d", cfg.ports.HTTP)
}
if cfg.userPorts.HTTP != 9999 {
t.Errorf("WithPorts: userPorts.HTTP = %d", cfg.userPorts.HTTP)
}
WithDataDir("/data").apply(cfg)
if cfg.dataDir != "/data" {
t.Errorf("WithDataDir = %q", cfg.dataDir)
}
WithHTTPClient(nil).apply(cfg)
Docker.apply(cfg)
if cfg.runtime != Docker {
t.Error("Docker option failed")
}
K8s.apply(cfg)
if cfg.runtime != K8s {
t.Error("K8s option failed")
}
}
func TestNewEmptyAddr(t *testing.T) {
_, err := New("")
if err == nil {
t.Error("expected error for empty addr")
}
}
func TestNewLocal(t *testing.T) {
c, err := New("local")
if err != nil { if err != nil {
t.Skipf("Docker not available: %v", err) t.Skipf("Docker not available: %v", err)
} }
defer c.Close() defer res.Close()
if !c.IsLocal() { if res.Volume == nil {
t.Error("expected IsLocal = true")
}
if c.Volume() == nil {
t.Error("Volume should not be nil") t.Error("Volume should not be nil")
} }
if c.Sandbox() == nil { if res.Runtime == nil {
t.Error("Sandbox should not be nil") t.Error("Runtime should not be nil")
}
if c.Proxy() == nil {
t.Error("Proxy should not be nil")
}
if c.VNC() == nil {
t.Error("VNC should not be nil")
}
// Test Workspace accessor
ws := c.Workspace("test-session")
if ws == nil {
t.Error("Workspace should not be nil")
} }
} }
func TestNewLocalWithDataDir(t *testing.T) { func TestDialLocalWithVolume(t *testing.T) {
dir := t.TempDir() dir := t.TempDir()
c, err := New("local", WithDataDir(dir)) vol := volume.NewLocal(dir)
res, err := DialLocal("", dir, vol)
if err != nil { if err != nil {
t.Skipf("Docker not available: %v", err) t.Skipf("Docker not available: %v", err)
} }
defer c.Close() defer res.Close()
if !c.IsLocal() { if res.DataDir != dir {
t.Error("expected IsLocal = true") t.Errorf("DataDir = %q, want %q", res.DataDir, dir)
}
if res.Volume == nil {
t.Error("Volume should not be nil")
} }
} }
func TestNewLocalExplicitSocket(t *testing.T) { func TestDialLocalExplicitSocket(t *testing.T) {
c, err := New("unix:///var/run/docker.sock") res, err := DialLocal("unix:///var/run/docker.sock", t.TempDir(), nil)
if err != nil { if err != nil {
t.Skipf("Docker not available: %v", err) t.Skipf("Docker not available: %v", err)
} }
defer c.Close() defer res.Close()
if !c.IsLocal() { if res.Runtime == nil {
t.Error("expected IsLocal = true for unix socket") t.Error("Runtime should not be nil for explicit unix socket")
} }
} }
func TestNewRemoteK8s(t *testing.T) { func TestDialRemoteDocker(t *testing.T) {
host := taiTestHost()
grpcPort := envPort("TAI_TEST_GRPC_PORT", 19100)
ports := taiTestPorts()
ports.GRPC = grpcPort
res, err := DialRemote(host, ports)
if err != nil {
t.Skipf("Tai not available at %s:%d: %v", host, grpcPort, err)
}
defer res.Close()
t.Logf("remote docker: host=%s ports=%+v", host, res.Ports)
if res.Volume == nil {
t.Error("Volume should not be nil")
}
if res.Runtime == nil {
t.Error("Runtime should not be nil")
}
}
func TestDialRemoteK8s(t *testing.T) {
host := os.Getenv("TAI_TEST_K8S_HOST") host := os.Getenv("TAI_TEST_K8S_HOST")
kubeconfig := os.Getenv("TAI_TEST_KUBECONFIG") kubeconfig := os.Getenv("TAI_TEST_KUBECONFIG")
if host == "" || kubeconfig == "" { if host == "" || kubeconfig == "" {
@ -232,119 +141,31 @@ func TestNewRemoteK8s(t *testing.T) {
VNC: envPort("TAI_TEST_K8S_VNC_PORT", 16080), VNC: envPort("TAI_TEST_K8S_VNC_PORT", 16080),
} }
c, err := New(fmt.Sprintf("tai://%s:%d", host, grpcPort), K8s, res, err := DialRemote(host, ports,
WithPorts(ports), WithDialRuntime(types.K8s),
WithKubeConfig(kubeconfig), WithDialKubeConfig(kubeconfig),
WithNamespace("default"), WithDialNamespace("default"),
) )
if err != nil { if err != nil {
t.Skipf("Tai K8s not available: %v", err) t.Skipf("Tai K8s not available: %v", err)
} }
defer c.Close() defer res.Close()
if c.IsLocal() { if res.Runtime == nil {
t.Error("expected IsLocal = false") t.Error("Runtime should not be nil")
}
if c.Sandbox() == nil {
t.Error("Sandbox should not be nil")
} }
} }
func TestNewRemoteK8sMissingKubeConfig(t *testing.T) { func TestDialRemoteK8sMissingKubeConfig(t *testing.T) {
_, err := New("tai://127.0.0.1", K8s) host := taiTestHost()
grpcPort := envPort("TAI_TEST_GRPC_PORT", 19100)
_, err := DialRemote(host, Ports{GRPC: grpcPort}, WithDialRuntime(types.K8s))
if err == nil { if err == nil {
t.Error("expected error for missing kubeconfig") t.Skip("Tai happened to be reachable; test only valid when gRPC is up")
} }
} }
func TestWithKubeConfigAndNamespace(t *testing.T) {
cfg := &config{ports: defaultPorts()}
WithKubeConfig("/path/to/kubeconfig").apply(cfg)
if cfg.kubeConfig != "/path/to/kubeconfig" {
t.Errorf("WithKubeConfig = %q", cfg.kubeConfig)
}
WithNamespace("test-ns").apply(cfg)
if cfg.namespace != "test-ns" {
t.Errorf("WithNamespace = %q", cfg.namespace)
}
}
func TestNewInvalidScheme(t *testing.T) {
_, err := New("ftp://host")
if err == nil {
t.Error("expected error for ftp://")
}
}
func TestNewRemoteDocker(t *testing.T) {
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://")
}
if c.Volume() == nil {
t.Error("Volume should not be nil")
}
if c.Sandbox() == nil {
t.Error("Sandbox should not be nil")
}
if c.Proxy() == nil {
t.Error("Proxy should not be nil")
}
if c.VNC() == nil {
t.Error("VNC should not be nil")
}
ws := c.Workspace("test")
if ws == nil {
t.Error("Workspace should not be nil")
}
}
func TestDiscoverPorts(t *testing.T) {
addr := taiRemoteAddr()
c, err := New(addr)
if err != nil {
t.Skipf("Tai not available at %s: %v", addr, err)
}
defer c.Close()
t.Logf("client resolved: GRPC=%d HTTP=%d VNC=%d Docker=%d K8s=%d",
c.ports.GRPC, c.ports.HTTP, c.ports.VNC, c.ports.Docker, c.ports.K8s)
if c.ports.GRPC == 0 {
t.Error("GRPC port should be discovered (non-zero)")
}
if c.ports.HTTP == 0 {
t.Error("HTTP port should be discovered (non-zero)")
}
}
func TestDiscoverPortsWithUserOverride(t *testing.T) {
addr := taiRemoteAddr()
c, err := New(addr, WithPorts(Ports{HTTP: 9999}))
if err != nil {
t.Skipf("Tai not available at %s: %v", addr, err)
}
defer c.Close()
if c.ports.HTTP != 9999 {
t.Errorf("HTTP = %d, want 9999 (user override should take precedence)", c.ports.HTTP)
}
if c.ports.GRPC == 0 {
t.Error("GRPC port should still be discovered (non-zero)")
}
t.Logf("ports: GRPC=%d HTTP=%d(user) VNC=%d Docker=%d",
c.ports.GRPC, c.ports.HTTP, c.ports.VNC, c.ports.Docker)
}
func TestRegisterLocal(t *testing.T) { func TestRegisterLocal(t *testing.T) {
registry.Init(nil) registry.Init(nil)
reg := registry.Global() reg := registry.Global()
@ -355,39 +176,37 @@ func TestRegisterLocal(t *testing.T) {
t.Skip("Docker not available, skipping RegisterLocal test") t.Skip("Docker not available, skipping RegisterLocal test")
} }
snap, found := reg.Get("local") meta, found := reg.Get("local")
if !found { if !found {
t.Fatal("expected 'local' node in registry after RegisterLocal") t.Fatal("expected 'local' node in registry after RegisterLocal")
} }
if snap.Mode != "local" { if meta.Mode != "local" {
t.Errorf("mode = %q, want 'local'", snap.Mode) t.Errorf("mode = %q, want 'local'", meta.Mode)
} }
if snap.Status != "online" { if meta.Status != "online" {
t.Errorf("status = %q, want 'online'", snap.Status) t.Errorf("status = %q, want 'online'", meta.Status)
} }
c, got := GetClient("local") res, got := GetResources("local")
if !got { if !got {
t.Fatal("GetClient('local') returned false after RegisterLocal") t.Fatal("GetResources('local') returned false after RegisterLocal")
} }
if c.DataDir() != dir { if res.DataDir != dir {
t.Errorf("DataDir = %q, want %q", c.DataDir(), dir) t.Errorf("DataDir = %q, want %q", res.DataDir, dir)
} }
if c.Sandbox() == nil { if res.Runtime == nil {
t.Error("local client Sandbox should not be nil") t.Error("local resources Runtime should not be nil")
} }
// Idempotent: second call should return true without error
ok2 := RegisterLocal(WithDataDir(dir)) ok2 := RegisterLocal(WithDataDir(dir))
if !ok2 { if !ok2 {
t.Error("second RegisterLocal should return true (idempotent)") t.Error("second RegisterLocal should return true (idempotent)")
} }
c.Close() res.Close()
} }
func TestRegisterLocal_NoRegistry(t *testing.T) { func TestRegisterLocal_NoRegistry(t *testing.T) {
// RegisterLocal without a registry should return false, not panic
origReg := registry.Global() origReg := registry.Global()
defer func() { defer func() {
if origReg != nil { if origReg != nil {
@ -395,26 +214,26 @@ func TestRegisterLocal_NoRegistry(t *testing.T) {
} }
}() }()
// registry.Global() returns the singleton; we can't un-init it,
// but we can verify RegisterLocal returns true (registry exists from
// other tests) or false gracefully.
ok := RegisterLocal() ok := RegisterLocal()
// Just verify it doesn't panic; result depends on Docker availability
_ = ok _ = ok
} }
func TestRegisterLocal_NoDocker(t *testing.T) { func TestRegisterLocal_NoDocker(t *testing.T) {
registry.Init(nil) registry.Init(nil)
// Use an unreachable Docker socket to ensure failure
ok := RegisterLocal(WithDataDir(t.TempDir())) ok := RegisterLocal(WithDataDir(t.TempDir()))
if !ok { if !ok {
// Expected when Docker is not available — just ensure no panic
return return
} }
// If Docker happens to be available, that's also fine res, got := GetResources("local")
c, _ := GetClient("local") if got && res != nil {
if c != nil { res.Close()
c.Close() }
}
func TestConnResourcesCloseNil(t *testing.T) {
var r *ConnResources
if err := r.Close(); err != nil {
t.Errorf("Close on nil should return nil, got %v", err)
} }
} }

128
tai/tunnel/forward.go Normal file
View file

@ -0,0 +1,128 @@
package tunnel
import (
"bytes"
"io"
"net/http"
"strings"
"github.com/gin-gonic/gin"
"github.com/yaoapp/yao/tai/tunnel/taipb"
"github.com/yaoapp/yao/tai/types"
)
// HandleForward handles HTTP/VNC/any TCP-level forwarding through the gRPC tunnel.
// Route: ANY /tai/:taiID/proxy/*path and GET /tai/:taiID/vnc/*path
//
// It hijacks the browser's raw TCP connection, asks Tai to open a Forward stream
// to the resolved target port, rewrites the request path, and then performs
// bidirectional byte-level bridging. No protocol parsing beyond HTTP hijack.
func (h *TunnelHandler) HandleForward(c *gin.Context) {
logger := h.logger
reg := h.reg
taiID := c.Param("taiID")
node, ok := reg.Get(taiID)
if !ok || node.Status != "online" {
c.JSON(http.StatusBadGateway, gin.H{"error": "tai node not available"})
return
}
targetPort := resolveTargetPort(c, node)
if targetPort == 0 {
c.JSON(http.StatusBadRequest, gin.H{"error": "cannot resolve target port"})
return
}
hijacker, ok := c.Writer.(http.Hijacker)
if !ok {
c.JSON(http.StatusInternalServerError, gin.H{"error": "hijack not supported"})
return
}
browserConn, bufrw, err := hijacker.Hijack()
if err != nil {
logger.Error("hijack failed", "err", err)
return
}
defer browserConn.Close()
fwd, err := h.RequestForward(taiID, targetPort)
if err != nil {
logger.Error("request forward failed",
"tai_id", taiID, "port", targetPort, "err", err)
browserConn.Write([]byte("HTTP/1.1 502 Bad Gateway\r\n\r\n"))
return
}
rewrittenReq := rewriteRequest(c.Request, taiID)
var reqBuf bytes.Buffer
rewrittenReq.Write(&reqBuf)
if bufrw.Reader.Buffered() > 0 {
buffered, _ := bufrw.Peek(bufrw.Reader.Buffered())
reqBuf.Write(buffered)
}
if err := fwd.Send(&taipb.ForwardData{Data: reqBuf.Bytes()}); err != nil {
logger.Error("send initial request", "err", err)
return
}
streamConn := newForwardConn(fwd)
bridgeTCP(
&netConnAdapter{ReadWriteCloser: browserConn},
streamConn,
)
}
// HandleForwardLazy is a gin.HandlerFunc that resolves the global TunnelHandler
// at call time (not registration time), so routes can be registered before the
// gRPC server starts.
func HandleForwardLazy(c *gin.Context) {
h := GlobalHandler()
if h == nil {
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "tunnel handler not initialized"})
return
}
h.HandleForward(c)
}
// resolveTargetPort determines the Tai-side port from the route pattern.
func resolveTargetPort(c *gin.Context, node *types.NodeMeta) int {
path := c.Request.URL.Path
if strings.Contains(path, "/vnc/") {
if node.Ports.VNC != 0 {
return node.Ports.VNC
}
return 16080
}
if strings.Contains(path, "/proxy/") {
if node.Ports.HTTP != 0 {
return node.Ports.HTTP
}
return 8099
}
return 0
}
// rewriteRequest clones the request and strips everything up to and including
// /tai/:taiID from the path, handling any baseURL prefix (e.g. /v1/tai/abc/proxy/x → /proxy/x).
func rewriteRequest(orig *http.Request, taiID string) *http.Request {
r := orig.Clone(orig.Context())
marker := "/tai/" + taiID
if idx := strings.Index(r.URL.Path, marker); idx >= 0 {
r.URL.Path = r.URL.Path[idx+len(marker):]
if r.URL.Path == "" {
r.URL.Path = "/"
}
}
r.RequestURI = r.URL.RequestURI()
return r
}
// netConnAdapter wraps an io.ReadWriteCloser as needed by bridgeTCP.
type netConnAdapter struct {
io.ReadWriteCloser
}

286
tai/tunnel/forward_test.go Normal file
View file

@ -0,0 +1,286 @@
package tunnel
import (
"net/http"
"net/http/httptest"
"net/url"
"testing"
"github.com/gin-gonic/gin"
"github.com/yaoapp/yao/tai/registry"
"github.com/yaoapp/yao/tai/types"
)
func init() {
gin.SetMode(gin.TestMode)
}
func TestResolveTargetPort_VNC(t *testing.T) {
tests := []struct {
name string
path string
vncPort int
wantPort int
}{
{"default_vnc", "/tai/abc/vnc/websockify", 0, 16080},
{"custom_vnc", "/tai/abc/vnc/websockify", 5900, 5900},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = &http.Request{URL: &url.URL{Path: tt.path}}
node := &types.NodeMeta{Ports: types.Ports{VNC: tt.vncPort}}
got := resolveTargetPort(c, node)
if got != tt.wantPort {
t.Errorf("resolveTargetPort = %d, want %d", got, tt.wantPort)
}
})
}
}
func TestResolveTargetPort_Proxy(t *testing.T) {
tests := []struct {
name string
path string
httpPort int
wantPort int
}{
{"default_proxy", "/tai/abc/proxy/api/v1/foo", 0, 8099},
{"custom_proxy", "/tai/abc/proxy/api/v1/foo", 9090, 9090},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = &http.Request{URL: &url.URL{Path: tt.path}}
node := &types.NodeMeta{Ports: types.Ports{HTTP: tt.httpPort}}
got := resolveTargetPort(c, node)
if got != tt.wantPort {
t.Errorf("resolveTargetPort = %d, want %d", got, tt.wantPort)
}
})
}
}
func TestResolveTargetPort_Unknown(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = &http.Request{URL: &url.URL{Path: "/tai/abc/unknown/something"}}
node := &types.NodeMeta{}
got := resolveTargetPort(c, node)
if got != 0 {
t.Errorf("resolveTargetPort = %d, want 0", got)
}
}
func TestRewriteRequest(t *testing.T) {
tests := []struct {
name string
origPath string
taiID string
wantPath string
wantURI string
}{
{
"proxy_path",
"/tai/abc123/proxy/api/v1/data",
"abc123",
"/proxy/api/v1/data",
"/proxy/api/v1/data",
},
{
"vnc_path",
"/tai/node-1/vnc/websockify",
"node-1",
"/vnc/websockify",
"/vnc/websockify",
},
{
"with_query",
"/tai/node-1/proxy/api?foo=bar",
"node-1",
"/proxy/api",
"/proxy/api?foo=bar",
},
{
"exact_prefix",
"/tai/node-1",
"node-1",
"/",
"/",
},
{
"with_base_url",
"/v1/tai/node-1/proxy/api/v1/data",
"node-1",
"/proxy/api/v1/data",
"/proxy/api/v1/data",
},
{
"with_base_url_vnc",
"/v1/tai/abc123/vnc/__host__/ws",
"abc123",
"/vnc/__host__/ws",
"/vnc/__host__/ws",
},
{
"no_match",
"/other/path",
"node-1",
"/other/path",
"/other/path",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
u, _ := url.Parse("http://localhost" + tt.origPath)
orig := &http.Request{
Method: "GET",
URL: u,
RequestURI: u.RequestURI(),
Host: "localhost",
Header: http.Header{},
}
got := rewriteRequest(orig, tt.taiID)
if got.URL.Path != tt.wantPath {
t.Errorf("path = %q, want %q", got.URL.Path, tt.wantPath)
}
if got.RequestURI != tt.wantURI {
t.Errorf("requestURI = %q, want %q", got.RequestURI, tt.wantURI)
}
if got == orig {
t.Error("rewriteRequest should return a clone, not the original")
}
})
}
}
func TestRewriteRequest_PreservesHeaders(t *testing.T) {
u, _ := url.Parse("http://localhost/tai/node-1/vnc/websockify")
orig := &http.Request{
Method: "GET",
URL: u,
RequestURI: u.RequestURI(),
Host: "localhost",
Header: http.Header{
"Connection": {"Upgrade"},
"Upgrade": {"websocket"},
},
}
got := rewriteRequest(orig, "node-1")
if got.Header.Get("Connection") != "Upgrade" {
t.Error("expected Connection header preserved")
}
if got.Header.Get("Upgrade") != "websocket" {
t.Error("expected Upgrade header preserved")
}
}
func TestHandleForwardLazy_NilHandler(t *testing.T) {
old := globalHandler
globalHandler = nil
defer func() { globalHandler = old }()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("GET", "/tai/abc/proxy/test", nil)
HandleForwardLazy(c)
if w.Code != http.StatusServiceUnavailable {
t.Errorf("expected 503, got %d", w.Code)
}
}
func TestHandleForward_NodeNotFound(t *testing.T) {
reg := registry.NewForTest()
h := NewTunnelHandler(reg)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("GET", "/tai/nonexistent/proxy/api", nil)
c.Params = gin.Params{{Key: "taiID", Value: "nonexistent"}}
h.HandleForward(c)
if w.Code != http.StatusBadGateway {
t.Errorf("expected 502, got %d", w.Code)
}
}
func TestHandleForward_NodeOffline(t *testing.T) {
reg := registry.NewForTest()
h := NewTunnelHandler(reg)
reg.Register(&registry.TaiNode{
TaiID: "offline-node",
Mode: "tunnel",
Ports: types.Ports{HTTP: 8099},
})
// Manually set status to offline via a Get() — the node is online by default
// after Register, but we need an offline one. We'll use Unregister + re-register
// pattern. Actually, let's just test with a node that doesn't exist:
// the NodeNotFound test above covers that case. Instead, test zero port.
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("GET", "/tai/offline-node/unknown/foo", nil)
c.Params = gin.Params{{Key: "taiID", Value: "offline-node"}}
h.HandleForward(c)
if w.Code != http.StatusBadRequest {
t.Errorf("expected 400 for unresolvable port, got %d", w.Code)
}
}
func TestHandleForwardLazy_WithHandler(t *testing.T) {
reg := registry.NewForTest()
old := globalHandler
globalHandler = NewTunnelHandler(reg)
defer func() { globalHandler = old }()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Request = httptest.NewRequest("GET", "/tai/missing/proxy/api", nil)
c.Params = gin.Params{{Key: "taiID", Value: "missing"}}
HandleForwardLazy(c)
if w.Code != http.StatusBadGateway {
t.Errorf("expected 502, got %d", w.Code)
}
}
func TestHandleForward_ViaRealHTTP(t *testing.T) {
reg := registry.NewForTest()
h := NewTunnelHandler(reg)
reg.Register(&registry.TaiNode{
TaiID: "http-node",
Mode: "tunnel",
Ports: types.Ports{HTTP: 8099},
})
router := gin.New()
router.Any("/tai/:taiID/proxy/*path", func(c *gin.Context) { h.HandleForward(c) })
srv := httptest.NewServer(router)
defer srv.Close()
resp, err := http.Get(srv.URL + "/tai/http-node/proxy/api")
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
// RequestForward will fail (no register stream) → hijacked conn gets "502"
// or the response will be a 502 written before hijack.
// Since hijack happens, the actual HTTP status may not be set normally.
// We just verify no panic and the request completes.
if resp.StatusCode == 200 {
t.Error("expected non-200 response for failed forward")
}
}

314
tai/tunnel/grpc_handler.go Normal file
View file

@ -0,0 +1,314 @@
package tunnel
import (
"fmt"
"io"
"log/slog"
"net"
"sync"
"time"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/peer"
"github.com/yaoapp/yao/grpc/auth"
tai "github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/registry"
"github.com/yaoapp/yao/tai/taiid"
"github.com/yaoapp/yao/tai/tunnel/taipb"
"github.com/yaoapp/yao/tai/types"
)
var globalHandler *TunnelHandler
// GlobalHandler returns the global TunnelHandler instance set by NewTunnelHandler.
func GlobalHandler() *TunnelHandler { return globalHandler }
// TunnelHandler implements the TaiTunnel gRPC service.
type TunnelHandler struct {
taipb.UnimplementedTaiTunnelServer
reg *registry.Registry
pending sync.Map // channel_id → chan taipb.TaiTunnel_ForwardServer
logger *slog.Logger
}
// NewTunnelHandler creates a TunnelHandler backed by the given registry.
// It also registers a bridge function so that OpenLocalListener uses
// gRPC Forward streams instead of WS data channels.
func NewTunnelHandler(reg *registry.Registry) *TunnelHandler {
h := &TunnelHandler{
reg: reg,
logger: slog.Default(),
}
reg.SetBridgeFunc(h.bridgeConn)
globalHandler = h
return h
}
// Register implements the control-plane stream (Tai → Yao).
func (h *TunnelHandler) Register(stream taipb.TaiTunnel_RegisterServer) error {
msg, err := stream.Recv()
if err != nil {
return fmt.Errorf("recv register: %w", err)
}
if msg.Type != "register" {
return fmt.Errorf("expected register, got %q", msg.Type)
}
if msg.NodeId == "" || msg.MachineId == "" {
return fmt.Errorf("register: node_id and machine_id required")
}
resolvedTaiID, err := taiid.Generate(msg.MachineId, msg.NodeId)
if err != nil {
return fmt.Errorf("taiid: %w", err)
}
authInfo := authInfoFromStream(stream)
remoteIP := ""
if p, ok := peer.FromContext(stream.Context()); ok {
if host, _, err := net.SplitHostPort(p.Addr.String()); err == nil {
remoteIP = host
}
}
node := &registry.TaiNode{
TaiID: resolvedTaiID,
MachineID: msg.MachineId,
Version: msg.Version,
DisplayName: msg.DisplayName,
Auth: authInfo,
System: systemFromProto(msg.System),
Mode: "tunnel",
Addr: "tunnel://" + remoteIP,
Ports: portsFromProto(msg.Ports),
Capabilities: capsFromProto(msg.Caps),
}
h.reg.Register(node)
h.reg.SetRegisterStream(resolvedTaiID, stream)
defer func() {
h.reg.Unregister(resolvedTaiID)
h.logger.Info("tai gRPC tunnel disconnected", "tai_id", resolvedTaiID)
}()
if err := stream.Send(&taipb.TunnelControl{
Type: "registered",
TaiId: resolvedTaiID,
}); err != nil {
return fmt.Errorf("send registered: %w", err)
}
h.logger.Info("tai gRPC tunnel connected", "tai_id", resolvedTaiID, "version", msg.Version)
go h.connectTunnelNode(resolvedTaiID)
for {
ctrl, err := stream.Recv()
if err != nil {
if err == io.EOF {
return nil
}
return err
}
switch ctrl.Type {
case "ping":
h.reg.UpdatePing(resolvedTaiID)
if err := stream.Send(&taipb.TunnelControl{Type: "pong"}); err != nil {
return err
}
}
}
}
// Forward implements the data-plane stream (Tai → Yao).
func (h *TunnelHandler) Forward(stream taipb.TaiTunnel_ForwardServer) error {
md, ok := metadata.FromIncomingContext(stream.Context())
if !ok {
return fmt.Errorf("missing metadata")
}
vals := md.Get("channel_id")
if len(vals) == 0 || vals[0] == "" {
return fmt.Errorf("missing channel_id in metadata")
}
channelID := vals[0]
if ch, ok := h.pending.LoadAndDelete(channelID); ok {
ch.(chan taipb.TaiTunnel_ForwardServer) <- stream
} else {
return fmt.Errorf("no pending channel for %s", channelID)
}
<-stream.Context().Done()
return nil
}
// RequestForward sends an "open" command to Tai via the Register stream and
// waits for Tai to call back with a Forward stream. Returns the Forward stream.
func (h *TunnelHandler) RequestForward(taiID string, targetPort int) (taipb.TaiTunnel_ForwardServer, error) {
stream := h.reg.GetRegisterStream(taiID)
if stream == nil {
return nil, fmt.Errorf("tai %s: no active register stream", taiID)
}
channelID, err := registry.GenerateChannelID()
if err != nil {
return nil, fmt.Errorf("generate channel_id: %w", err)
}
waitCh := make(chan taipb.TaiTunnel_ForwardServer, 1)
h.pending.Store(channelID, waitCh)
defer h.pending.Delete(channelID)
regStream, ok := stream.(taipb.TaiTunnel_RegisterServer)
if !ok {
return nil, fmt.Errorf("tai %s: register stream type mismatch", taiID)
}
if err := regStream.Send(&taipb.TunnelControl{
Type: "open",
ChannelId: channelID,
TargetPort: int32(targetPort),
}); err != nil {
return nil, fmt.Errorf("send open: %w", err)
}
select {
case fwd := <-waitCh:
return fwd, nil
case <-time.After(10 * time.Second):
return nil, fmt.Errorf("tai %s: forward timeout (10s)", taiID)
case <-regStream.Context().Done():
return nil, fmt.Errorf("tai %s: register stream closed while waiting for forward", taiID)
}
}
// connectTunnelNode establishes gRPC resources to the Tai node through the tunnel.
func (h *TunnelHandler) connectTunnelNode(taiID string) {
res, err := tai.DialTunnel(taiID, h.reg)
if err != nil {
h.logger.Warn("failed to connect tunnel node",
"tai_id", taiID, "err", err)
return
}
h.reg.SetResources(taiID, res)
h.logger.Info("tunnel node resources connected", "tai_id", taiID)
}
// bridgeConn bridges a local TCP connection to a Tai port via gRPC Forward stream.
// Called by registry.OpenLocalListener for each accepted TCP connection.
func (h *TunnelHandler) bridgeConn(taiID string, targetPort int, localConn net.Conn) {
fwd, err := h.RequestForward(taiID, targetPort)
if err != nil {
localConn.Close()
h.logger.Error("request forward failed",
"tai_id", taiID, "port", targetPort, "err", err)
return
}
streamConn := newForwardConn(fwd)
bridgeTCP(localConn, streamConn)
}
// forwardConn wraps a Forward stream as a net.Conn-like reader/writer.
type forwardConn struct {
stream taipb.TaiTunnel_ForwardServer
buf []byte
}
func newForwardConn(stream taipb.TaiTunnel_ForwardServer) *forwardConn {
return &forwardConn{stream: stream}
}
func (c *forwardConn) Read(p []byte) (int, error) {
if len(c.buf) > 0 {
n := copy(p, c.buf)
c.buf = c.buf[n:]
return n, nil
}
msg, err := c.stream.Recv()
if err != nil {
return 0, err
}
n := copy(p, msg.Data)
if n < len(msg.Data) {
c.buf = msg.Data[n:]
}
return n, nil
}
func (c *forwardConn) Write(p []byte) (int, error) {
if err := c.stream.Send(&taipb.ForwardData{Data: p}); err != nil {
return 0, err
}
return len(p), nil
}
func (c *forwardConn) Close() error {
return nil
}
// bridgeTCP copies bytes bidirectionally, closing both sides when done.
func bridgeTCP(a, b io.ReadWriteCloser) {
var wg sync.WaitGroup
wg.Add(2)
cp := func(dst io.WriteCloser, src io.ReadCloser) {
defer wg.Done()
io.Copy(dst, src)
dst.Close()
}
go cp(a, b)
go cp(b, a)
wg.Wait()
}
// ── helpers ──────────────────────────────────────────────────────────────────
func authInfoFromStream(stream taipb.TaiTunnel_RegisterServer) types.AuthInfo {
info := auth.GetAuthorizedInfo(stream.Context())
if info == nil {
return types.AuthInfo{}
}
return types.AuthInfo{
Subject: info.Subject,
UserID: info.UserID,
ClientID: info.ClientID,
Scope: info.Scope,
TeamID: info.TeamID,
TenantID: info.TenantID,
}
}
func portsFromProto(p *taipb.Ports) types.Ports {
if p == nil {
return types.Ports{}
}
return types.Ports{
GRPC: int(p.Grpc),
HTTP: int(p.Http),
VNC: int(p.Vnc),
Docker: int(p.Docker),
K8s: int(p.K8S),
}
}
func capsFromProto(c *taipb.Capabilities) types.Capabilities {
if c == nil {
return types.Capabilities{}
}
return types.Capabilities{
Docker: c.Docker,
K8s: c.K8S,
HostExec: c.HostExec,
}
}
func systemFromProto(s *taipb.SystemInfo) types.SystemInfo {
if s == nil {
return types.SystemInfo{}
}
return types.SystemInfo{
OS: s.Os,
Arch: s.Arch,
Hostname: s.Hostname,
Shell: s.Shell,
}
}

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,56 @@
syntax = "proto3";
package tai.tunnel;
option go_package = "github.com/yaoapp/yao/tai/tunnel/taipb";
service TaiTunnel {
// Control plane: Tai Yao, register + keepalive + receive commands.
rpc Register(stream TunnelControl) returns (stream TunnelControl);
// Data plane: Tai Yao, raw TCP forwarding.
rpc Forward(stream ForwardData) returns (stream ForwardData);
}
message TunnelControl {
string type = 1; // "register" / "registered" / "open" / "ping" / "pong"
// Carried on "register" (Tai Yao)
string node_id = 2;
string machine_id = 3;
string display_name = 4;
string version = 5;
Ports ports = 6;
Capabilities caps = 7;
SystemInfo system = 8;
// Carried on "open" (Yao Tai)
string channel_id = 10;
int32 target_port = 11;
// Carried on "registered" (Yao Tai)
string tai_id = 20;
}
message ForwardData {
bytes data = 1;
}
message Ports {
int32 grpc = 1;
int32 http = 2;
int32 vnc = 3;
int32 docker = 4;
int32 k8s = 5;
}
message Capabilities {
bool docker = 1;
bool k8s = 2;
bool host_exec = 3;
}
message SystemInfo {
string os = 1;
string arch = 2;
string hostname = 3;
string shell = 4;
}

View file

@ -1,172 +0,0 @@
package tunnel
import (
"bufio"
"io"
"log/slog"
"net"
"net/http"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
"github.com/yaoapp/yao/tai/registry"
)
// HandleProxy handles HTTP reverse proxy requests for a tunnel-connected Tai:
// ANY /tai/:taiID/proxy/*path
// Opens a data channel to Tai's HTTP port, forwards the HTTP request,
// and streams the response back.
func HandleProxy(c *gin.Context) {
logger := slog.Default()
reg := registry.Global()
if reg == nil {
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "registry not initialized"})
return
}
taiID := c.Param("taiID")
node, ok := reg.Get(taiID)
if !ok || node.Status != "online" {
c.JSON(http.StatusBadGateway, gin.H{"error": "tai node not available"})
return
}
httpPort := node.Ports["http"]
if httpPort == 0 {
httpPort = 8099
}
channelID, resultCh, err := reg.RequestChannel(taiID, httpPort)
if err != nil {
logger.Error("request channel failed", "tai_id", taiID, "err", err)
c.JSON(http.StatusBadGateway, gin.H{"error": "tunnel channel failed"})
return
}
remoteConn, ok := <-resultCh
if !ok || remoteConn == nil {
logger.Error("data channel timeout", "tai_id", taiID, "channel_id", channelID)
c.JSON(http.StatusGatewayTimeout, gin.H{"error": "data channel timeout"})
return
}
defer remoteConn.Close()
path := c.Param("path")
outReq, err := http.NewRequestWithContext(c.Request.Context(), c.Request.Method, "http://tai-tunnel"+path, c.Request.Body)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": "build request failed"})
return
}
outReq.Header = c.Request.Header.Clone()
outReq.Host = c.Request.Host
if err := outReq.Write(remoteConn); err != nil {
logger.Error("write request to tunnel", "err", err)
c.JSON(http.StatusBadGateway, gin.H{"error": "write to tunnel failed"})
return
}
resp, err := http.ReadResponse(bufio.NewReader(remoteConn), outReq)
if err != nil {
logger.Error("read response from tunnel", "err", err)
c.JSON(http.StatusBadGateway, gin.H{"error": "read from tunnel failed"})
return
}
defer resp.Body.Close()
for k, vv := range resp.Header {
for _, v := range vv {
c.Writer.Header().Add(k, v)
}
}
c.Writer.WriteHeader(resp.StatusCode)
io.Copy(c.Writer, resp.Body)
}
// HandleVNC handles VNC WebSocket proxying for a tunnel-connected Tai:
// GET /tai/:taiID/vnc/*path
// Upgrades the client connection to WebSocket, opens a data channel to
// Tai's VNC port, and bridges the two WebSocket connections.
func HandleVNC(c *gin.Context) {
logger := slog.Default()
reg := registry.Global()
if reg == nil {
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "registry not initialized"})
return
}
taiID := c.Param("taiID")
node, ok := reg.Get(taiID)
if !ok || node.Status != "online" {
c.JSON(http.StatusBadGateway, gin.H{"error": "tai node not available"})
return
}
vncPort := node.Ports["vnc"]
if vncPort == 0 {
vncPort = 16080
}
channelID, resultCh, err := reg.RequestChannel(taiID, vncPort)
if err != nil {
logger.Error("request vnc channel failed", "tai_id", taiID, "err", err)
c.JSON(http.StatusBadGateway, gin.H{"error": "tunnel channel failed"})
return
}
clientConn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
if err != nil {
logger.Error("ws upgrade client failed", "err", err)
return
}
taiConn, ok := <-resultCh
if !ok || taiConn == nil {
logger.Error("vnc data channel timeout", "tai_id", taiID, "channel_id", channelID)
clientConn.Close()
return
}
bridgeWSToConn(clientConn, taiConn)
}
// bridgeWSToConn bridges a client WebSocket to a net.Conn (tunnel data channel).
func bridgeWSToConn(clientWS *websocket.Conn, taiConn net.Conn) {
done := make(chan struct{}, 2)
// client WS -> tai conn
go func() {
defer func() { done <- struct{}{} }()
for {
_, data, err := clientWS.ReadMessage()
if err != nil {
return
}
if _, err := taiConn.Write(data); err != nil {
return
}
}
}()
// tai conn -> client WS
go func() {
defer func() { done <- struct{}{} }()
buf := make([]byte, 32*1024)
for {
n, err := taiConn.Read(buf)
if n > 0 {
if wErr := clientWS.WriteMessage(websocket.BinaryMessage, buf[:n]); wErr != nil {
return
}
}
if err != nil {
return
}
}
}()
<-done
clientWS.Close()
taiConn.Close()
<-done
}

View file

@ -2,204 +2,14 @@ package tunnel
import ( import (
"fmt" "fmt"
"io"
"log/slog" "log/slog"
"net"
"net/http" "net/http"
"strings" "strings"
"sync"
"time"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
oauth "github.com/yaoapp/yao/openapi/oauth" oauth "github.com/yaoapp/yao/openapi/oauth"
tai "github.com/yaoapp/yao/tai" "github.com/yaoapp/yao/tai/types"
"github.com/yaoapp/yao/tai/registry"
"github.com/yaoapp/yao/tai/taiid"
) )
var upgrader = websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool { return true },
}
// HandleControl handles the Tai control channel WebSocket: GET /ws/tai.
// Authenticates via Bearer token, reads register + ping messages,
// and maintains the Tai node in the global registry.
func HandleControl(c *gin.Context) {
logger := slog.Default()
reg := registry.Global()
if reg == nil {
c.JSON(http.StatusServiceUnavailable, 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 := authenticateBearerFunc(bearer)
if err != nil {
logger.Warn("tunnel auth failed", "err", err)
c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"})
return
}
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
if err != nil {
logger.Error("ws upgrade failed", "err", err)
return
}
// Read the register message
var regMsg registerMessage
if err := conn.ReadJSON(&regMsg); err != nil {
logger.Error("read register message", "err", err)
conn.Close()
return
}
if regMsg.Type != "register" {
logger.Error("expected register message", "got", regMsg.Type)
conn.Close()
return
}
if regMsg.NodeID == "" || regMsg.MachineID == "" {
logger.Error("register message missing node_id or machine_id")
conn.Close()
return
}
resolvedTaiID, err := taiid.Generate(regMsg.MachineID, regMsg.NodeID)
if err != nil {
logger.Error("taiid generation failed", "err", err)
conn.Close()
return
}
addr := ""
if host, _, err := net.SplitHostPort(c.Request.RemoteAddr); err == nil {
addr = "tunnel://" + host
}
node := &registry.TaiNode{
TaiID: resolvedTaiID,
MachineID: regMsg.MachineID,
Version: regMsg.Version,
DisplayName: regMsg.DisplayName,
Auth: authInfo,
System: regMsg.System,
Mode: "tunnel",
Addr: addr,
YaoBase: regMsg.Server,
Ports: regMsg.Ports,
Capabilities: regMsg.Capabilities,
ControlConn: conn,
}
reg.Register(node)
defer func() {
reg.Unregister(resolvedTaiID)
logger.Info("tai tunnel disconnected", "tai_id", resolvedTaiID)
}()
if err := reg.WriteControlJSON(resolvedTaiID, map[string]string{"type": "registered", "tai_id": resolvedTaiID}); err != nil {
logger.Error("write registered response", "err", err)
return
}
logger.Info("tai tunnel connected", "tai_id", resolvedTaiID, "version", regMsg.Version)
go connectTunnelNode(resolvedTaiID, reg, logger)
for {
var msg controlMsg
if err := conn.ReadJSON(&msg); err != nil {
if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
logger.Debug("control channel read error", "err", err)
}
return
}
switch msg.Type {
case "ping":
reg.UpdatePing(resolvedTaiID)
if err := reg.WriteControlJSON(resolvedTaiID, map[string]string{"type": "pong"}); err != nil {
logger.Debug("pong write failed", "err", err)
return
}
default:
logger.Debug("unknown control message", "type", msg.Type)
}
}
}
// HandleData handles a Tai data channel WebSocket: GET /ws/tai/data/:channel_id.
// Authenticates via Bearer token, verifies the caller matches the pending
// channel's owner, then wraps the WS as a net.Conn for bidirectional bridging.
func HandleData(c *gin.Context) {
logger := slog.Default()
reg := registry.Global()
if reg == nil {
c.JSON(http.StatusServiceUnavailable, 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 := authenticateBearerFunc(bearer)
if err != nil {
logger.Warn("data channel auth failed", "err", err)
c.JSON(http.StatusUnauthorized, gin.H{"error": "authentication failed"})
return
}
channelID := c.Param("channel_id")
if channelID == "" {
c.JSON(http.StatusBadRequest, gin.H{"error": "missing channel_id"})
return
}
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
if err != nil {
logger.Error("ws data upgrade failed", "err", err)
return
}
resolvedTaiID := reg.FindTaiIDByAuthClient(authInfo.ClientID)
if resolvedTaiID == "" {
resolvedTaiID = authInfo.ClientID
}
wsConn := newWSConn(conn)
if err := reg.AcceptDataChannel(channelID, resolvedTaiID, wsConn); err != nil {
logger.Debug("accept data channel failed", "channel_id", channelID, "err", err,
"auth_client_id", authInfo.ClientID, "resolved_tai_id", resolvedTaiID)
conn.Close()
return
}
}
// registerMessage is the JSON structure for Tai's register message.
type registerMessage struct {
Type string `json:"type"`
NodeID string `json:"node_id,omitempty"`
ClientID string `json:"client_id,omitempty"`
MachineID string `json:"machine_id"`
DisplayName string `json:"display_name,omitempty"`
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.
type controlMsg struct {
Type string `json:"type"`
}
func extractBearer(r *http.Request) string { func extractBearer(r *http.Request) string {
auth := r.Header.Get("Authorization") auth := r.Header.Get("Authorization")
if len(auth) > 7 && strings.EqualFold(auth[:7], "bearer ") { if len(auth) > 7 && strings.EqualFold(auth[:7], "bearer ") {
@ -210,20 +20,20 @@ func extractBearer(r *http.Request) string {
var authenticateBearerFunc = authenticateBearerDefault var authenticateBearerFunc = authenticateBearerDefault
func authenticateBearerDefault(token string) (registry.AuthInfo, error) { func authenticateBearerDefault(token string) (types.AuthInfo, error) {
svc := oauth.OAuth svc := oauth.OAuth
if svc == nil { if svc == nil {
return registry.AuthInfo{}, fmt.Errorf("oauth service not initialized") return types.AuthInfo{}, fmt.Errorf("oauth service not initialized")
} }
result, err := svc.AuthenticateToken(oauth.AuthInput{ result, err := svc.AuthenticateToken(oauth.AuthInput{
AccessToken: token, AccessToken: token,
}) })
if err != nil { if err != nil {
return registry.AuthInfo{}, err return types.AuthInfo{}, err
} }
info := registry.AuthInfo{} info := types.AuthInfo{}
if result.Info != nil { if result.Info != nil {
info.Subject = result.Info.Subject info.Subject = result.Info.Subject
info.UserID = result.Info.UserID info.UserID = result.Info.UserID
@ -266,72 +76,20 @@ func authenticateBearerDefault(token string) (registry.AuthInfo, error) {
return info, nil return info, nil
} }
// wsConn wraps a gorilla/websocket.Conn to implement net.Conn for raw byte bridging. func portsFromMap(m map[string]int) types.Ports {
type wsConn struct { return types.Ports{
ws *websocket.Conn GRPC: m["grpc"],
reader io.Reader HTTP: m["http"],
mu sync.Mutex VNC: m["vnc"],
} Docker: m["docker"],
K8s: m["k8s"],
func newWSConn(ws *websocket.Conn) *wsConn {
return &wsConn{ws: ws}
}
func (c *wsConn) Read(p []byte) (int, error) {
for {
if c.reader != nil {
n, err := c.reader.Read(p)
if n > 0 {
return n, nil
}
c.reader = nil
if err != nil && err != io.EOF {
return 0, err
}
}
_, reader, err := c.ws.NextReader()
if err != nil {
return 0, err
}
c.reader = reader
} }
} }
func (c *wsConn) Write(p []byte) (int, error) { func capsFromMap(m map[string]bool) types.Capabilities {
c.mu.Lock() return types.Capabilities{
defer c.mu.Unlock() Docker: m["docker"],
err := c.ws.WriteMessage(websocket.BinaryMessage, p) K8s: m["k8s"],
if err != nil { HostExec: m["host_exec"],
return 0, err
} }
return len(p), nil
}
func (c *wsConn) Close() error {
return c.ws.Close()
}
func (c *wsConn) LocalAddr() net.Addr { return c.ws.LocalAddr() }
func (c *wsConn) RemoteAddr() net.Addr { return c.ws.RemoteAddr() }
func (c *wsConn) SetDeadline(t time.Time) error {
if err := c.ws.SetReadDeadline(t); err != nil {
return err
}
return c.ws.SetWriteDeadline(t)
}
func (c *wsConn) SetReadDeadline(t time.Time) error { return c.ws.SetReadDeadline(t) }
func (c *wsConn) SetWriteDeadline(t time.Time) error { return c.ws.SetWriteDeadline(t) }
// connectTunnelNode creates a tai.Client through the tunnel and binds it to the taiID.
func connectTunnelNode(taiID string, reg *registry.Registry, logger *slog.Logger) {
client, err := tai.New("tunnel://" + taiID)
if err != nil {
logger.Warn("failed to connect tunnel node",
"tai_id", taiID, "err", err)
return
}
_ = client // initTunnel already calls reg.SetClient(taiID, c)
logger.Info("tai client created for tunnel node", "tai_id", taiID)
} }

View file

@ -1,41 +1,13 @@
package tunnel package tunnel
import ( import (
"fmt"
"io"
"net"
"net/http" "net/http"
"net/http/httptest"
"strings"
"sync"
"testing" "testing"
"time"
"github.com/gin-gonic/gin" "github.com/yaoapp/yao/tai/tunnel/taipb"
"github.com/gorilla/websocket" "github.com/yaoapp/yao/tai/types"
"github.com/yaoapp/yao/tai/registry"
) )
func init() {
gin.SetMode(gin.TestMode)
}
func setupTestRegistry() *registry.Registry {
r := registry.NewForTest()
registry.SetGlobalForTest(r)
return r
}
func mockAuth(info registry.AuthInfo, authErr error) func() {
old := authenticateBearerFunc
authenticateBearerFunc = func(token string) (registry.AuthInfo, error) {
return info, authErr
}
return func() { authenticateBearerFunc = old }
}
// --- extractBearer ---
func TestExtractBearer(t *testing.T) { func TestExtractBearer(t *testing.T) {
tests := []struct { tests := []struct {
name string name string
@ -62,555 +34,90 @@ func TestExtractBearer(t *testing.T) {
} }
} }
// --- wsConn --- func TestPortsFromMap(t *testing.T) {
m := map[string]int{"grpc": 19100, "http": 8099, "vnc": 16080, "docker": 12375, "k8s": 16443}
func TestWSConn_EchoRoundTrip(t *testing.T) { p := portsFromMap(m)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if p.GRPC != 19100 || p.HTTP != 8099 || p.VNC != 16080 || p.Docker != 12375 || p.K8s != 16443 {
conn, err := upgrader.Upgrade(w, r, nil) t.Errorf("portsFromMap got %+v", p)
if err != nil {
return
}
defer conn.Close()
wc := newWSConn(conn)
buf := make([]byte, 256)
n, err := wc.Read(buf)
if err != nil {
return
}
wc.Write(buf[:n])
}))
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err != nil {
t.Fatalf("dial: %v", err)
}
defer conn.Close()
if resp.StatusCode != http.StatusSwitchingProtocols {
t.Errorf("handshake status = %d, want 101", resp.StatusCode)
}
msg := []byte("hello tunnel")
if err := conn.WriteMessage(websocket.BinaryMessage, msg); err != nil {
t.Fatalf("write: %v", err)
}
mt, reply, err := conn.ReadMessage()
if err != nil {
t.Fatalf("read: %v", err)
}
if mt != websocket.BinaryMessage {
t.Errorf("type = %d, want BinaryMessage(%d)", mt, websocket.BinaryMessage)
}
if string(reply) != "hello tunnel" {
t.Errorf("reply = %q, want %q", reply, "hello tunnel")
} }
} }
func TestWSConn_MultipleMessages(t *testing.T) { func TestPortsFromMap_Empty(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { p := portsFromMap(nil)
conn, err := upgrader.Upgrade(w, r, nil) if p.GRPC != 0 || p.HTTP != 0 {
if err != nil { t.Errorf("portsFromMap(nil) got %+v", p)
return
}
defer conn.Close()
wc := newWSConn(conn)
for i := 0; i < 3; i++ {
buf := make([]byte, 256)
n, err := wc.Read(buf)
if err != nil {
return
}
wc.Write(buf[:n])
}
}))
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err != nil {
t.Fatalf("dial: %v", err)
}
defer conn.Close()
for i, msg := range []string{"one", "two", "three"} {
conn.WriteMessage(websocket.BinaryMessage, []byte(msg))
_, reply, err := conn.ReadMessage()
if err != nil {
t.Fatalf("round %d read: %v", i, err)
}
if string(reply) != msg {
t.Errorf("round %d: got %q, want %q", i, reply, msg)
}
} }
} }
func TestWSConn_ImplementsNetConn(t *testing.T) { func TestCapsFromMap(t *testing.T) {
var _ net.Conn = (*wsConn)(nil) m := map[string]bool{"docker": true, "k8s": false, "host_exec": true}
} c := capsFromMap(m)
if !c.Docker || c.K8s || !c.HostExec {
func TestWSConn_LocalRemoteAddr(t *testing.T) { t.Errorf("capsFromMap got %+v", c)
addrCh := make(chan [2]net.Addr, 1)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
return
}
wc := newWSConn(conn)
addrCh <- [2]net.Addr{wc.LocalAddr(), wc.RemoteAddr()}
wc.Close()
}))
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err != nil {
t.Fatalf("dial: %v", err)
}
defer conn.Close()
select {
case addrs := <-addrCh:
if addrs[0] == nil {
t.Error("LocalAddr should not be nil")
}
if addrs[1] == nil {
t.Error("RemoteAddr should not be nil")
}
case <-time.After(2 * time.Second):
t.Fatal("timeout waiting for addresses")
} }
} }
// --- HandleControl --- func TestCapsFromMap_Empty(t *testing.T) {
c := capsFromMap(nil)
func newGinRouter() *gin.Engine { if c.Docker || c.K8s || c.HostExec {
r := gin.New() t.Errorf("capsFromMap(nil) got %+v", c)
r.GET("/ws/tai", HandleControl) }
r.GET("/ws/tai/data/:channel_id", HandleData)
return r
} }
func TestHandleControl_NoRegistry(t *testing.T) { func TestPortsFromProto(t *testing.T) {
registry.SetGlobalForTest(nil) pp := &taipb.Ports{Grpc: 19100, Http: 8099, Vnc: 16080, Docker: 12375, K8S: 16443}
defer setupTestRegistry() p := portsFromProto(pp)
if p.GRPC != 19100 || p.HTTP != 8099 || p.VNC != 16080 || p.Docker != 12375 || p.K8s != 16443 {
t.Errorf("portsFromProto got %+v", p)
}
}
restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil) func TestPortsFromProto_Nil(t *testing.T) {
defer restore() p := portsFromProto(nil)
if p != (types.Ports{}) {
t.Errorf("portsFromProto(nil) = %+v", p)
}
}
srv := httptest.NewServer(newGinRouter()) func TestCapsFromProto(t *testing.T) {
defer srv.Close() cp := &taipb.Capabilities{Docker: true, K8S: false, HostExec: true}
c := capsFromProto(cp)
if !c.Docker || c.K8s || !c.HostExec {
t.Errorf("capsFromProto got %+v", c)
}
}
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai" func TestCapsFromProto_Nil(t *testing.T) {
_, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{ c := capsFromProto(nil)
"Authorization": []string{"Bearer test-token"}, if c != (types.Capabilities{}) {
}) t.Errorf("capsFromProto(nil) = %+v", c)
}
}
func TestSystemFromProto(t *testing.T) {
sp := &taipb.SystemInfo{Os: "linux", Arch: "amd64", Hostname: "host1", Shell: "bash"}
s := systemFromProto(sp)
if s.OS != "linux" || s.Arch != "amd64" || s.Hostname != "host1" || s.Shell != "bash" {
t.Errorf("systemFromProto got %+v", s)
}
}
func TestSystemFromProto_Nil(t *testing.T) {
s := systemFromProto(nil)
if s != (types.SystemInfo{}) {
t.Errorf("systemFromProto(nil) = %+v", s)
}
}
func TestAuthenticateBearerDefault_NoOAuth(t *testing.T) {
_, err := authenticateBearerDefault("some-token")
if err == nil { if err == nil {
t.Fatal("expected dial to fail when registry is nil") t.Fatal("expected error when oauth service is nil")
}
if resp != nil && resp.StatusCode != http.StatusServiceUnavailable {
t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusServiceUnavailable)
} }
} }
func TestHandleControl_NoAuth(t *testing.T) { func TestAuthenticateBearerFunc_IsDefault(t *testing.T) {
setupTestRegistry() if authenticateBearerFunc == nil {
t.Fatal("authenticateBearerFunc should be set")
srv := httptest.NewServer(newGinRouter())
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai"
_, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err == nil {
t.Fatal("expected dial to fail without auth")
}
if resp != nil && resp.StatusCode != http.StatusUnauthorized {
t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusUnauthorized)
}
}
func TestHandleControl_AuthFailed(t *testing.T) {
setupTestRegistry()
restore := mockAuth(registry.AuthInfo{}, fmt.Errorf("bad token"))
defer restore()
srv := httptest.NewServer(newGinRouter())
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai"
_, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
"Authorization": []string{"Bearer bad-token"},
})
if err == nil {
t.Fatal("expected dial to fail with bad auth")
}
if resp != nil && resp.StatusCode != http.StatusUnauthorized {
t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusUnauthorized)
}
}
func TestHandleControl_RegisterAndPing(t *testing.T) {
reg := setupTestRegistry()
restore := mockAuth(registry.AuthInfo{
ClientID: "tai-001",
Subject: "user-test",
Scope: "tai:tunnel",
}, nil)
defer restore()
srv := httptest.NewServer(newGinRouter())
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai"
conn, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
"Authorization": []string{"Bearer valid-token"},
})
if err != nil {
t.Fatalf("dial: %v", err)
}
defer conn.Close()
if resp.StatusCode != http.StatusSwitchingProtocols {
t.Errorf("handshake = %d, want 101", resp.StatusCode)
}
regMsg := registerMessage{
Type: "register",
NodeID: "9100",
MachineID: "m-test",
Version: "2.0",
Ports: map[string]int{"grpc": 9100},
}
if err := conn.WriteJSON(regMsg); err != nil {
t.Fatalf("write register: %v", err)
}
var registered map[string]string
if err := conn.ReadJSON(&registered); err != nil {
t.Fatalf("read registered: %v", err)
}
if registered["type"] != "registered" {
t.Errorf("response type = %q, want registered", registered["type"])
}
gotTaiID := registered["tai_id"]
if gotTaiID == "" || len(gotTaiID) < 5 || gotTaiID[:4] != "tai-" {
t.Errorf("response tai_id = %q, want server-generated tai-xxx", gotTaiID)
}
snap, ok := reg.Get(gotTaiID)
if !ok {
t.Fatal("node not found in registry after register")
}
if snap.Status != "online" {
t.Errorf("Status = %q, want online", snap.Status)
}
if snap.MachineID != "m-test" {
t.Errorf("MachineID = %q, want m-test", snap.MachineID)
}
if snap.Version != "2.0" {
t.Errorf("Version = %q, want 2.0", snap.Version)
}
if snap.Mode != "tunnel" {
t.Errorf("Mode = %q, want tunnel", snap.Mode)
}
if snap.Auth.ClientID != "tai-001" {
t.Errorf("Auth.ClientID = %q, want tai-001", snap.Auth.ClientID)
}
if snap.Auth.Subject != "user-test" {
t.Errorf("Auth.Subject = %q, want user-test", snap.Auth.Subject)
}
if snap.Ports["grpc"] != 9100 {
t.Errorf("Ports[grpc] = %d, want 9100", snap.Ports["grpc"])
}
time.Sleep(10 * time.Millisecond)
if err := conn.WriteJSON(map[string]string{"type": "ping"}); err != nil {
t.Fatalf("write ping: %v", err)
}
// Read messages until we get the pong; connectTunnelNode may inject
// "open" messages (with numeric fields) before our pong arrives.
var gotPong bool
for i := 0; i < 10; i++ {
var msg map[string]interface{}
if err := conn.ReadJSON(&msg); err != nil {
t.Fatalf("read message: %v", err)
}
if msg["type"] == "pong" {
gotPong = true
break
}
}
if !gotPong {
t.Error("did not receive pong after ping")
}
snap2, _ := reg.Get(gotTaiID)
if !snap2.LastPing.After(snap.LastPing) {
t.Error("LastPing should be updated after ping")
}
conn.WriteMessage(websocket.CloseMessage,
websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
time.Sleep(100 * time.Millisecond)
if _, ok := reg.Get("tai-001"); ok {
t.Error("node should be unregistered after connection close")
}
}
func TestHandleControl_BadRegisterType(t *testing.T) {
setupTestRegistry()
restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil)
defer restore()
srv := httptest.NewServer(newGinRouter())
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai"
conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
"Authorization": []string{"Bearer valid-token"},
})
if err != nil {
t.Fatalf("dial: %v", err)
}
defer conn.Close()
conn.WriteJSON(map[string]string{"type": "not-register"})
_, _, readErr := conn.ReadMessage()
if readErr == nil {
t.Error("expected connection to close for bad register type")
}
}
func TestHandleControl_MissingTaiID(t *testing.T) {
setupTestRegistry()
restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil)
defer restore()
srv := httptest.NewServer(newGinRouter())
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai"
conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
"Authorization": []string{"Bearer valid-token"},
})
if err != nil {
t.Fatalf("dial: %v", err)
}
defer conn.Close()
conn.WriteJSON(map[string]string{"type": "register"})
_, _, readErr := conn.ReadMessage()
if readErr == nil {
t.Error("expected connection to close for missing tai_id")
}
}
// --- HandleData ---
func TestHandleData_NoAuth(t *testing.T) {
setupTestRegistry()
srv := httptest.NewServer(newGinRouter())
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/ch-001"
_, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
if err == nil {
t.Fatal("expected dial to fail without auth")
}
if resp != nil && resp.StatusCode != http.StatusUnauthorized {
t.Errorf("status = %d, want %d", resp.StatusCode, http.StatusUnauthorized)
}
}
func TestHandleData_AcceptSuccess(t *testing.T) {
reg := setupTestRegistry()
restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil)
defer restore()
resultCh := make(chan net.Conn, 1)
timer := time.AfterFunc(5*time.Second, func() {})
reg.SetPendingForTest("ch-test-123", "tai-001", resultCh, timer)
srv := httptest.NewServer(newGinRouter())
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/ch-test-123"
conn, resp, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
"Authorization": []string{"Bearer valid-token"},
})
if err != nil {
t.Fatalf("dial: %v", err)
}
defer conn.Close()
if resp.StatusCode != http.StatusSwitchingProtocols {
t.Errorf("status = %d, want 101", resp.StatusCode)
}
select {
case c := <-resultCh:
if c == nil {
t.Fatal("expected non-nil conn from resultCh")
}
c.Close()
case <-time.After(2 * time.Second):
t.Fatal("timeout waiting for conn on resultCh")
}
}
func TestHandleData_ChannelNotPending(t *testing.T) {
setupTestRegistry()
restore := mockAuth(registry.AuthInfo{ClientID: "tai-001"}, nil)
defer restore()
srv := httptest.NewServer(newGinRouter())
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/nonexistent"
conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
"Authorization": []string{"Bearer valid-token"},
})
if err != nil {
return
}
defer conn.Close()
_, _, readErr := conn.ReadMessage()
if readErr == nil {
t.Error("expected connection to close for non-pending channel")
}
}
func TestHandleData_TaiIDMismatch(t *testing.T) {
reg := setupTestRegistry()
restore := mockAuth(registry.AuthInfo{ClientID: "tai-intruder"}, nil)
defer restore()
resultCh := make(chan net.Conn, 1)
timer := time.AfterFunc(5*time.Second, func() {})
reg.SetPendingForTest("ch-mismatch", "tai-owner", resultCh, timer)
srv := httptest.NewServer(newGinRouter())
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/ch-mismatch"
conn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
"Authorization": []string{"Bearer valid-token"},
})
if err != nil {
return
}
defer conn.Close()
_, _, readErr := conn.ReadMessage()
if readErr == nil {
t.Error("expected connection to close for tai_id mismatch")
}
}
// --- Full open-channel flow ---
func TestHandleControl_OpenChannelAndBridge(t *testing.T) {
reg := setupTestRegistry()
restore := mockAuth(registry.AuthInfo{
ClientID: "tai-001",
Subject: "user-test",
}, nil)
defer restore()
srv := httptest.NewServer(newGinRouter())
defer srv.Close()
wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai"
ctrlConn, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
"Authorization": []string{"Bearer valid-token"},
})
if err != nil {
t.Fatalf("dial control: %v", err)
}
defer ctrlConn.Close()
ctrlConn.WriteJSON(registerMessage{
Type: "register",
NodeID: "9100",
MachineID: "m-test",
Ports: map[string]int{"grpc": 9100},
})
var registered map[string]string
if err := ctrlConn.ReadJSON(&registered); err != nil {
t.Fatalf("read registered: %v", err)
}
if registered["type"] != "registered" {
t.Fatalf("expected registered, got %v", registered)
}
taiID := registered["tai_id"]
var wg sync.WaitGroup
wg.Add(1)
var requestErr error
var channelConn net.Conn
go func() {
defer wg.Done()
_, resultCh, err := reg.RequestChannel(taiID, 9100)
if err != nil {
requestErr = err
return
}
channelConn = <-resultCh
}()
time.Sleep(50 * time.Millisecond)
var openCmd map[string]interface{}
if err := ctrlConn.ReadJSON(&openCmd); err != nil {
t.Fatalf("read open cmd: %v", err)
}
if openCmd["type"] != "open" {
t.Errorf("open type = %v, want open", openCmd["type"])
}
channelID, ok := openCmd["channel_id"].(string)
if !ok || channelID == "" {
t.Fatalf("missing channel_id: %v", openCmd)
}
if tp, ok := openCmd["target_port"].(float64); !ok || int(tp) != 9100 {
t.Errorf("target_port = %v, want 9100", openCmd["target_port"])
}
dataURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws/tai/data/" + channelID
dataConn, _, err := websocket.DefaultDialer.Dial(dataURL, http.Header{
"Authorization": []string{"Bearer valid-token"},
})
if err != nil {
t.Fatalf("dial data: %v", err)
}
defer dataConn.Close()
wg.Wait()
if requestErr != nil {
t.Fatalf("RequestChannel: %v", requestErr)
}
if channelConn == nil {
t.Fatal("expected non-nil conn from RequestChannel")
}
defer channelConn.Close()
payload := []byte("grpc-payload-test")
dataConn.WriteMessage(websocket.BinaryMessage, payload)
buf := make([]byte, 256)
n, err := channelConn.Read(buf)
if err != nil && err != io.EOF {
t.Fatalf("read bridged: %v", err)
}
if string(buf[:n]) != "grpc-payload-test" {
t.Errorf("bridged data = %q, want %q", buf[:n], "grpc-payload-test")
} }
} }

View file

@ -0,0 +1,500 @@
// Code generated by protoc-gen-go. DO NOT EDIT.
// versions:
// protoc-gen-go v1.36.11
// protoc v4.25.0
// source: tunnel.proto
package taipb
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 TunnelControl struct {
state protoimpl.MessageState `protogen:"open.v1"`
Type string `protobuf:"bytes,1,opt,name=type,proto3" json:"type,omitempty"` // "register" / "registered" / "open" / "ping" / "pong"
// Carried on "register" (Tai → Yao)
NodeId string `protobuf:"bytes,2,opt,name=node_id,json=nodeId,proto3" json:"node_id,omitempty"`
MachineId string `protobuf:"bytes,3,opt,name=machine_id,json=machineId,proto3" json:"machine_id,omitempty"`
DisplayName string `protobuf:"bytes,4,opt,name=display_name,json=displayName,proto3" json:"display_name,omitempty"`
Version string `protobuf:"bytes,5,opt,name=version,proto3" json:"version,omitempty"`
Ports *Ports `protobuf:"bytes,6,opt,name=ports,proto3" json:"ports,omitempty"`
Caps *Capabilities `protobuf:"bytes,7,opt,name=caps,proto3" json:"caps,omitempty"`
System *SystemInfo `protobuf:"bytes,8,opt,name=system,proto3" json:"system,omitempty"`
// Carried on "open" (Yao → Tai)
ChannelId string `protobuf:"bytes,10,opt,name=channel_id,json=channelId,proto3" json:"channel_id,omitempty"`
TargetPort int32 `protobuf:"varint,11,opt,name=target_port,json=targetPort,proto3" json:"target_port,omitempty"`
// Carried on "registered" (Yao → Tai)
TaiId string `protobuf:"bytes,20,opt,name=tai_id,json=taiId,proto3" json:"tai_id,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *TunnelControl) Reset() {
*x = TunnelControl{}
mi := &file_tunnel_proto_msgTypes[0]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *TunnelControl) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*TunnelControl) ProtoMessage() {}
func (x *TunnelControl) ProtoReflect() protoreflect.Message {
mi := &file_tunnel_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 TunnelControl.ProtoReflect.Descriptor instead.
func (*TunnelControl) Descriptor() ([]byte, []int) {
return file_tunnel_proto_rawDescGZIP(), []int{0}
}
func (x *TunnelControl) GetType() string {
if x != nil {
return x.Type
}
return ""
}
func (x *TunnelControl) GetNodeId() string {
if x != nil {
return x.NodeId
}
return ""
}
func (x *TunnelControl) GetMachineId() string {
if x != nil {
return x.MachineId
}
return ""
}
func (x *TunnelControl) GetDisplayName() string {
if x != nil {
return x.DisplayName
}
return ""
}
func (x *TunnelControl) GetVersion() string {
if x != nil {
return x.Version
}
return ""
}
func (x *TunnelControl) GetPorts() *Ports {
if x != nil {
return x.Ports
}
return nil
}
func (x *TunnelControl) GetCaps() *Capabilities {
if x != nil {
return x.Caps
}
return nil
}
func (x *TunnelControl) GetSystem() *SystemInfo {
if x != nil {
return x.System
}
return nil
}
func (x *TunnelControl) GetChannelId() string {
if x != nil {
return x.ChannelId
}
return ""
}
func (x *TunnelControl) GetTargetPort() int32 {
if x != nil {
return x.TargetPort
}
return 0
}
func (x *TunnelControl) GetTaiId() string {
if x != nil {
return x.TaiId
}
return ""
}
type ForwardData struct {
state protoimpl.MessageState `protogen:"open.v1"`
Data []byte `protobuf:"bytes,1,opt,name=data,proto3" json:"data,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *ForwardData) Reset() {
*x = ForwardData{}
mi := &file_tunnel_proto_msgTypes[1]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *ForwardData) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*ForwardData) ProtoMessage() {}
func (x *ForwardData) ProtoReflect() protoreflect.Message {
mi := &file_tunnel_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 ForwardData.ProtoReflect.Descriptor instead.
func (*ForwardData) Descriptor() ([]byte, []int) {
return file_tunnel_proto_rawDescGZIP(), []int{1}
}
func (x *ForwardData) GetData() []byte {
if x != nil {
return x.Data
}
return nil
}
type Ports struct {
state protoimpl.MessageState `protogen:"open.v1"`
Grpc int32 `protobuf:"varint,1,opt,name=grpc,proto3" json:"grpc,omitempty"`
Http int32 `protobuf:"varint,2,opt,name=http,proto3" json:"http,omitempty"`
Vnc int32 `protobuf:"varint,3,opt,name=vnc,proto3" json:"vnc,omitempty"`
Docker int32 `protobuf:"varint,4,opt,name=docker,proto3" json:"docker,omitempty"`
K8S int32 `protobuf:"varint,5,opt,name=k8s,proto3" json:"k8s,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *Ports) Reset() {
*x = Ports{}
mi := &file_tunnel_proto_msgTypes[2]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *Ports) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*Ports) ProtoMessage() {}
func (x *Ports) ProtoReflect() protoreflect.Message {
mi := &file_tunnel_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 Ports.ProtoReflect.Descriptor instead.
func (*Ports) Descriptor() ([]byte, []int) {
return file_tunnel_proto_rawDescGZIP(), []int{2}
}
func (x *Ports) GetGrpc() int32 {
if x != nil {
return x.Grpc
}
return 0
}
func (x *Ports) GetHttp() int32 {
if x != nil {
return x.Http
}
return 0
}
func (x *Ports) GetVnc() int32 {
if x != nil {
return x.Vnc
}
return 0
}
func (x *Ports) GetDocker() int32 {
if x != nil {
return x.Docker
}
return 0
}
func (x *Ports) GetK8S() int32 {
if x != nil {
return x.K8S
}
return 0
}
type Capabilities struct {
state protoimpl.MessageState `protogen:"open.v1"`
Docker bool `protobuf:"varint,1,opt,name=docker,proto3" json:"docker,omitempty"`
K8S bool `protobuf:"varint,2,opt,name=k8s,proto3" json:"k8s,omitempty"`
HostExec bool `protobuf:"varint,3,opt,name=host_exec,json=hostExec,proto3" json:"host_exec,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *Capabilities) Reset() {
*x = Capabilities{}
mi := &file_tunnel_proto_msgTypes[3]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *Capabilities) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*Capabilities) ProtoMessage() {}
func (x *Capabilities) ProtoReflect() protoreflect.Message {
mi := &file_tunnel_proto_msgTypes[3]
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 Capabilities.ProtoReflect.Descriptor instead.
func (*Capabilities) Descriptor() ([]byte, []int) {
return file_tunnel_proto_rawDescGZIP(), []int{3}
}
func (x *Capabilities) GetDocker() bool {
if x != nil {
return x.Docker
}
return false
}
func (x *Capabilities) GetK8S() bool {
if x != nil {
return x.K8S
}
return false
}
func (x *Capabilities) GetHostExec() bool {
if x != nil {
return x.HostExec
}
return false
}
type SystemInfo struct {
state protoimpl.MessageState `protogen:"open.v1"`
Os string `protobuf:"bytes,1,opt,name=os,proto3" json:"os,omitempty"`
Arch string `protobuf:"bytes,2,opt,name=arch,proto3" json:"arch,omitempty"`
Hostname string `protobuf:"bytes,3,opt,name=hostname,proto3" json:"hostname,omitempty"`
Shell string `protobuf:"bytes,4,opt,name=shell,proto3" json:"shell,omitempty"`
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *SystemInfo) Reset() {
*x = SystemInfo{}
mi := &file_tunnel_proto_msgTypes[4]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *SystemInfo) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*SystemInfo) ProtoMessage() {}
func (x *SystemInfo) ProtoReflect() protoreflect.Message {
mi := &file_tunnel_proto_msgTypes[4]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use SystemInfo.ProtoReflect.Descriptor instead.
func (*SystemInfo) Descriptor() ([]byte, []int) {
return file_tunnel_proto_rawDescGZIP(), []int{4}
}
func (x *SystemInfo) GetOs() string {
if x != nil {
return x.Os
}
return ""
}
func (x *SystemInfo) GetArch() string {
if x != nil {
return x.Arch
}
return ""
}
func (x *SystemInfo) GetHostname() string {
if x != nil {
return x.Hostname
}
return ""
}
func (x *SystemInfo) GetShell() string {
if x != nil {
return x.Shell
}
return ""
}
var File_tunnel_proto protoreflect.FileDescriptor
const file_tunnel_proto_rawDesc = "" +
"\n" +
"\ftunnel.proto\x12\n" +
"tai.tunnel\"\xf6\x02\n" +
"\rTunnelControl\x12\x12\n" +
"\x04type\x18\x01 \x01(\tR\x04type\x12\x17\n" +
"\anode_id\x18\x02 \x01(\tR\x06nodeId\x12\x1d\n" +
"\n" +
"machine_id\x18\x03 \x01(\tR\tmachineId\x12!\n" +
"\fdisplay_name\x18\x04 \x01(\tR\vdisplayName\x12\x18\n" +
"\aversion\x18\x05 \x01(\tR\aversion\x12'\n" +
"\x05ports\x18\x06 \x01(\v2\x11.tai.tunnel.PortsR\x05ports\x12,\n" +
"\x04caps\x18\a \x01(\v2\x18.tai.tunnel.CapabilitiesR\x04caps\x12.\n" +
"\x06system\x18\b \x01(\v2\x16.tai.tunnel.SystemInfoR\x06system\x12\x1d\n" +
"\n" +
"channel_id\x18\n" +
" \x01(\tR\tchannelId\x12\x1f\n" +
"\vtarget_port\x18\v \x01(\x05R\n" +
"targetPort\x12\x15\n" +
"\x06tai_id\x18\x14 \x01(\tR\x05taiId\"!\n" +
"\vForwardData\x12\x12\n" +
"\x04data\x18\x01 \x01(\fR\x04data\"k\n" +
"\x05Ports\x12\x12\n" +
"\x04grpc\x18\x01 \x01(\x05R\x04grpc\x12\x12\n" +
"\x04http\x18\x02 \x01(\x05R\x04http\x12\x10\n" +
"\x03vnc\x18\x03 \x01(\x05R\x03vnc\x12\x16\n" +
"\x06docker\x18\x04 \x01(\x05R\x06docker\x12\x10\n" +
"\x03k8s\x18\x05 \x01(\x05R\x03k8s\"U\n" +
"\fCapabilities\x12\x16\n" +
"\x06docker\x18\x01 \x01(\bR\x06docker\x12\x10\n" +
"\x03k8s\x18\x02 \x01(\bR\x03k8s\x12\x1b\n" +
"\thost_exec\x18\x03 \x01(\bR\bhostExec\"b\n" +
"\n" +
"SystemInfo\x12\x0e\n" +
"\x02os\x18\x01 \x01(\tR\x02os\x12\x12\n" +
"\x04arch\x18\x02 \x01(\tR\x04arch\x12\x1a\n" +
"\bhostname\x18\x03 \x01(\tR\bhostname\x12\x14\n" +
"\x05shell\x18\x04 \x01(\tR\x05shell2\x92\x01\n" +
"\tTaiTunnel\x12D\n" +
"\bRegister\x12\x19.tai.tunnel.TunnelControl\x1a\x19.tai.tunnel.TunnelControl(\x010\x01\x12?\n" +
"\aForward\x12\x17.tai.tunnel.ForwardData\x1a\x17.tai.tunnel.ForwardData(\x010\x01B(Z&github.com/yaoapp/yao/tai/tunnel/taipbb\x06proto3"
var (
file_tunnel_proto_rawDescOnce sync.Once
file_tunnel_proto_rawDescData []byte
)
func file_tunnel_proto_rawDescGZIP() []byte {
file_tunnel_proto_rawDescOnce.Do(func() {
file_tunnel_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_tunnel_proto_rawDesc), len(file_tunnel_proto_rawDesc)))
})
return file_tunnel_proto_rawDescData
}
var file_tunnel_proto_msgTypes = make([]protoimpl.MessageInfo, 5)
var file_tunnel_proto_goTypes = []any{
(*TunnelControl)(nil), // 0: tai.tunnel.TunnelControl
(*ForwardData)(nil), // 1: tai.tunnel.ForwardData
(*Ports)(nil), // 2: tai.tunnel.Ports
(*Capabilities)(nil), // 3: tai.tunnel.Capabilities
(*SystemInfo)(nil), // 4: tai.tunnel.SystemInfo
}
var file_tunnel_proto_depIdxs = []int32{
2, // 0: tai.tunnel.TunnelControl.ports:type_name -> tai.tunnel.Ports
3, // 1: tai.tunnel.TunnelControl.caps:type_name -> tai.tunnel.Capabilities
4, // 2: tai.tunnel.TunnelControl.system:type_name -> tai.tunnel.SystemInfo
0, // 3: tai.tunnel.TaiTunnel.Register:input_type -> tai.tunnel.TunnelControl
1, // 4: tai.tunnel.TaiTunnel.Forward:input_type -> tai.tunnel.ForwardData
0, // 5: tai.tunnel.TaiTunnel.Register:output_type -> tai.tunnel.TunnelControl
1, // 6: tai.tunnel.TaiTunnel.Forward:output_type -> tai.tunnel.ForwardData
5, // [5:7] is the sub-list for method output_type
3, // [3:5] is the sub-list for method input_type
3, // [3:3] is the sub-list for extension type_name
3, // [3:3] is the sub-list for extension extendee
0, // [0:3] is the sub-list for field type_name
}
func init() { file_tunnel_proto_init() }
func file_tunnel_proto_init() {
if File_tunnel_proto != nil {
return
}
type x struct{}
out := protoimpl.TypeBuilder{
File: protoimpl.DescBuilder{
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
RawDescriptor: unsafe.Slice(unsafe.StringData(file_tunnel_proto_rawDesc), len(file_tunnel_proto_rawDesc)),
NumEnums: 0,
NumMessages: 5,
NumExtensions: 0,
NumServices: 1,
},
GoTypes: file_tunnel_proto_goTypes,
DependencyIndexes: file_tunnel_proto_depIdxs,
MessageInfos: file_tunnel_proto_msgTypes,
}.Build()
File_tunnel_proto = out.File
file_tunnel_proto_goTypes = nil
file_tunnel_proto_depIdxs = nil
}

View file

@ -0,0 +1,151 @@
// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
// versions:
// - protoc-gen-go-grpc v1.6.1
// - protoc v4.25.0
// source: tunnel.proto
package taipb
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 (
TaiTunnel_Register_FullMethodName = "/tai.tunnel.TaiTunnel/Register"
TaiTunnel_Forward_FullMethodName = "/tai.tunnel.TaiTunnel/Forward"
)
// TaiTunnelClient is the client API for TaiTunnel service.
//
// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream.
type TaiTunnelClient interface {
// Control plane: Tai → Yao, register + keepalive + receive commands.
Register(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[TunnelControl, TunnelControl], error)
// Data plane: Tai → Yao, raw TCP forwarding.
Forward(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[ForwardData, ForwardData], error)
}
type taiTunnelClient struct {
cc grpc.ClientConnInterface
}
func NewTaiTunnelClient(cc grpc.ClientConnInterface) TaiTunnelClient {
return &taiTunnelClient{cc}
}
func (c *taiTunnelClient) Register(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[TunnelControl, TunnelControl], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &TaiTunnel_ServiceDesc.Streams[0], TaiTunnel_Register_FullMethodName, cOpts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[TunnelControl, TunnelControl]{ClientStream: stream}
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 TaiTunnel_RegisterClient = grpc.BidiStreamingClient[TunnelControl, TunnelControl]
func (c *taiTunnelClient) Forward(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[ForwardData, ForwardData], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &TaiTunnel_ServiceDesc.Streams[1], TaiTunnel_Forward_FullMethodName, cOpts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[ForwardData, ForwardData]{ClientStream: stream}
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 TaiTunnel_ForwardClient = grpc.BidiStreamingClient[ForwardData, ForwardData]
// TaiTunnelServer is the server API for TaiTunnel service.
// All implementations must embed UnimplementedTaiTunnelServer
// for forward compatibility.
type TaiTunnelServer interface {
// Control plane: Tai → Yao, register + keepalive + receive commands.
Register(grpc.BidiStreamingServer[TunnelControl, TunnelControl]) error
// Data plane: Tai → Yao, raw TCP forwarding.
Forward(grpc.BidiStreamingServer[ForwardData, ForwardData]) error
mustEmbedUnimplementedTaiTunnelServer()
}
// UnimplementedTaiTunnelServer 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 UnimplementedTaiTunnelServer struct{}
func (UnimplementedTaiTunnelServer) Register(grpc.BidiStreamingServer[TunnelControl, TunnelControl]) error {
return status.Error(codes.Unimplemented, "method Register not implemented")
}
func (UnimplementedTaiTunnelServer) Forward(grpc.BidiStreamingServer[ForwardData, ForwardData]) error {
return status.Error(codes.Unimplemented, "method Forward not implemented")
}
func (UnimplementedTaiTunnelServer) mustEmbedUnimplementedTaiTunnelServer() {}
func (UnimplementedTaiTunnelServer) testEmbeddedByValue() {}
// UnsafeTaiTunnelServer may be embedded to opt out of forward compatibility for this service.
// Use of this interface is not recommended, as added methods to TaiTunnelServer will
// result in compilation errors.
type UnsafeTaiTunnelServer interface {
mustEmbedUnimplementedTaiTunnelServer()
}
func RegisterTaiTunnelServer(s grpc.ServiceRegistrar, srv TaiTunnelServer) {
// If the following call panics, it indicates UnimplementedTaiTunnelServer 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(&TaiTunnel_ServiceDesc, srv)
}
func _TaiTunnel_Register_Handler(srv interface{}, stream grpc.ServerStream) error {
return srv.(TaiTunnelServer).Register(&grpc.GenericServerStream[TunnelControl, TunnelControl]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type TaiTunnel_RegisterServer = grpc.BidiStreamingServer[TunnelControl, TunnelControl]
func _TaiTunnel_Forward_Handler(srv interface{}, stream grpc.ServerStream) error {
return srv.(TaiTunnelServer).Forward(&grpc.GenericServerStream[ForwardData, ForwardData]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type TaiTunnel_ForwardServer = grpc.BidiStreamingServer[ForwardData, ForwardData]
// TaiTunnel_ServiceDesc is the grpc.ServiceDesc for TaiTunnel service.
// It's only intended for direct use with grpc.RegisterService,
// and not to be introspected or modified (even as a copy)
var TaiTunnel_ServiceDesc = grpc.ServiceDesc{
ServiceName: "tai.tunnel.TaiTunnel",
HandlerType: (*TaiTunnelServer)(nil),
Methods: []grpc.MethodDesc{},
Streams: []grpc.StreamDesc{
{
StreamName: "Register",
Handler: _TaiTunnel_Register_Handler,
ServerStreams: true,
ClientStreams: true,
},
{
StreamName: "Forward",
Handler: _TaiTunnel_Forward_Handler,
ServerStreams: true,
ClientStreams: true,
},
},
Metadata: "tunnel.proto",
}

67
tai/types/types.go Normal file
View file

@ -0,0 +1,67 @@
package types
import "time"
// Runtime selects which container runtime to use via Tai.
type Runtime int
const (
Docker Runtime = iota
K8s
)
// Ports configures service ports for Tai server.
type Ports struct {
GRPC int `json:"grpc"`
HTTP int `json:"http"`
VNC int `json:"vnc"`
Docker int `json:"docker"`
K8s int `json:"k8s"`
}
// Capabilities describes what features a Tai node supports.
type Capabilities struct {
Docker bool `json:"docker"`
K8s bool `json:"k8s"`
HostExec bool `json:"host_exec"`
}
// 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"`
Shell string `json:"shell,omitempty"`
TempDir string `json:"temp_dir,omitempty"`
}
// AuthInfo holds Yao user authorization extracted from OAuth token.
type AuthInfo struct {
Subject string
UserID string
ClientID string
Scope string
TeamID string
TenantID string
}
// NodeMeta is the read-only metadata snapshot of a registered Tai node.
// Carries no runtime resource references.
type NodeMeta struct {
TaiID string
MachineID string
Version string
Auth AuthInfo
System SystemInfo
Mode string // "direct" | "tunnel" | "local"
Addr string
YaoBase string
Ports Ports
Capabilities Capabilities
Status string // "online" | "offline" | "connecting"
ConnectedAt time.Time
LastPing time.Time
DisplayName string
}

View file

@ -6,7 +6,7 @@ import (
"net/http" "net/http"
"strings" "strings"
"github.com/yaoapp/yao/tai/sandbox" "github.com/yaoapp/yao/tai/runtime"
) )
const defaultVNCContainerPort = 6080 const defaultVNCContainerPort = 6080
@ -78,11 +78,11 @@ func (t *tunnelVNC) Ping(_ context.Context, _ string) error {
// --- Local implementation --- // --- Local implementation ---
type localVNC struct { type localVNC struct {
sb sandbox.Sandbox sb runtime.Runtime
} }
// NewLocal creates a VNC that resolves host VNC ports via sandbox.Inspect. // NewLocal creates a VNC that resolves host VNC ports via runtime.Inspect.
func NewLocal(sb sandbox.Sandbox) VNC { func NewLocal(sb runtime.Runtime) VNC {
return &localVNC{sb: sb} return &localVNC{sb: sb}
} }

View file

@ -8,7 +8,7 @@ import (
"testing" "testing"
"time" "time"
"github.com/yaoapp/yao/tai/sandbox" "github.com/yaoapp/yao/tai/runtime"
) )
func TestRemoteURL(t *testing.T) { func TestRemoteURL(t *testing.T) {
@ -63,10 +63,10 @@ func TestRemotePingError(t *testing.T) {
func TestLocalURL(t *testing.T) { func TestLocalURL(t *testing.T) {
mock := &mockSandbox{ mock := &mockSandbox{
inspectFn: func(ctx context.Context, id string) (*sandbox.ContainerInfo, error) { inspectFn: func(ctx context.Context, id string) (*runtime.ContainerInfo, error) {
return &sandbox.ContainerInfo{ return &runtime.ContainerInfo{
ID: id, ID: id,
Ports: []sandbox.PortMapping{ Ports: []runtime.PortMapping{
{ContainerPort: 6080, HostPort: 49152, HostIP: "127.0.0.1", Protocol: "tcp"}, {ContainerPort: 6080, HostPort: 49152, HostIP: "127.0.0.1", Protocol: "tcp"},
}, },
}, nil }, nil
@ -86,10 +86,10 @@ func TestLocalURL(t *testing.T) {
func TestLocalURLEmptyHostIP(t *testing.T) { func TestLocalURLEmptyHostIP(t *testing.T) {
mock := &mockSandbox{ mock := &mockSandbox{
inspectFn: func(ctx context.Context, id string) (*sandbox.ContainerInfo, error) { inspectFn: func(ctx context.Context, id string) (*runtime.ContainerInfo, error) {
return &sandbox.ContainerInfo{ return &runtime.ContainerInfo{
ID: id, ID: id,
Ports: []sandbox.PortMapping{ Ports: []runtime.PortMapping{
{ContainerPort: 6080, HostPort: 49152, HostIP: "", Protocol: "tcp"}, {ContainerPort: 6080, HostPort: 49152, HostIP: "", Protocol: "tcp"},
}, },
}, nil }, nil
@ -109,8 +109,8 @@ func TestLocalURLEmptyHostIP(t *testing.T) {
func TestLocalURLPortNotFound(t *testing.T) { func TestLocalURLPortNotFound(t *testing.T) {
mock := &mockSandbox{ mock := &mockSandbox{
inspectFn: func(ctx context.Context, id string) (*sandbox.ContainerInfo, error) { inspectFn: func(ctx context.Context, id string) (*runtime.ContainerInfo, error) {
return &sandbox.ContainerInfo{ID: id}, nil return &runtime.ContainerInfo{ID: id}, nil
}, },
} }
@ -123,7 +123,7 @@ func TestLocalURLPortNotFound(t *testing.T) {
func TestLocalURLInspectError(t *testing.T) { func TestLocalURLInspectError(t *testing.T) {
mock := &mockSandbox{ mock := &mockSandbox{
inspectFn: func(ctx context.Context, id string) (*sandbox.ContainerInfo, error) { inspectFn: func(ctx context.Context, id string) (*runtime.ContainerInfo, error) {
return nil, fmt.Errorf("not found") return nil, fmt.Errorf("not found")
}, },
} }
@ -157,10 +157,10 @@ func TestLocalPingSuccess(t *testing.T) {
} }
mock := &mockSandbox{ mock := &mockSandbox{
inspectFn: func(ctx context.Context, id string) (*sandbox.ContainerInfo, error) { inspectFn: func(ctx context.Context, id string) (*runtime.ContainerInfo, error) {
return &sandbox.ContainerInfo{ return &runtime.ContainerInfo{
ID: id, ID: id,
Ports: []sandbox.PortMapping{ Ports: []runtime.PortMapping{
{ContainerPort: 6080, HostPort: port, HostIP: "127.0.0.1", Protocol: "tcp"}, {ContainerPort: 6080, HostPort: port, HostIP: "127.0.0.1", Protocol: "tcp"},
}, },
}, nil }, nil
@ -175,7 +175,7 @@ func TestLocalPingSuccess(t *testing.T) {
func TestLocalPingError(t *testing.T) { func TestLocalPingError(t *testing.T) {
mock := &mockSandbox{ mock := &mockSandbox{
inspectFn: func(ctx context.Context, id string) (*sandbox.ContainerInfo, error) { inspectFn: func(ctx context.Context, id string) (*runtime.ContainerInfo, error) {
return nil, fmt.Errorf("not found") return nil, fmt.Errorf("not found")
}, },
} }
@ -186,12 +186,12 @@ func TestLocalPingError(t *testing.T) {
} }
} }
// mockSandbox implements sandbox.Sandbox for testing. // mockSandbox implements runtime.Sandbox for testing.
type mockSandbox struct { type mockSandbox struct {
inspectFn func(ctx context.Context, id string) (*sandbox.ContainerInfo, error) inspectFn func(ctx context.Context, id string) (*runtime.ContainerInfo, error)
} }
func (m *mockSandbox) Create(ctx context.Context, opts sandbox.CreateOptions) (string, error) { func (m *mockSandbox) Create(ctx context.Context, opts runtime.CreateOptions) (string, error) {
return "", nil return "", nil
} }
func (m *mockSandbox) Start(ctx context.Context, id string) error { return nil } func (m *mockSandbox) Start(ctx context.Context, id string) error { return nil }
@ -199,19 +199,19 @@ func (m *mockSandbox) Stop(ctx context.Context, id string, timeout time.Duration
return nil return nil
} }
func (m *mockSandbox) Remove(ctx context.Context, id string, force bool) error { return nil } func (m *mockSandbox) Remove(ctx context.Context, id string, force bool) error { return nil }
func (m *mockSandbox) Exec(ctx context.Context, id string, cmd []string, opts sandbox.ExecOptions) (*sandbox.ExecResult, error) { func (m *mockSandbox) Exec(ctx context.Context, id string, cmd []string, opts runtime.ExecOptions) (*runtime.ExecResult, error) {
return nil, nil return nil, nil
} }
func (m *mockSandbox) ExecStream(ctx context.Context, id string, cmd []string, opts sandbox.ExecOptions) (*sandbox.StreamHandle, error) { func (m *mockSandbox) ExecStream(ctx context.Context, id string, cmd []string, opts runtime.ExecOptions) (*runtime.StreamHandle, error) {
return nil, nil return nil, nil
} }
func (m *mockSandbox) Inspect(ctx context.Context, id string) (*sandbox.ContainerInfo, error) { func (m *mockSandbox) Inspect(ctx context.Context, id string) (*runtime.ContainerInfo, error) {
if m.inspectFn != nil { if m.inspectFn != nil {
return m.inspectFn(ctx, id) return m.inspectFn(ctx, id)
} }
return &sandbox.ContainerInfo{ID: id}, nil return &runtime.ContainerInfo{ID: id}, nil
} }
func (m *mockSandbox) List(ctx context.Context, opts sandbox.ListOptions) ([]sandbox.ContainerInfo, error) { func (m *mockSandbox) List(ctx context.Context, opts runtime.ListOptions) ([]runtime.ContainerInfo, error) {
return nil, nil return nil, nil
} }
func (m *mockSandbox) Close() error { return nil } func (m *mockSandbox) Close() error { return nil }

View file

@ -2,6 +2,7 @@ package jsapi_test
import ( import (
"os" "os"
"strconv"
"strings" "strings"
"testing" "testing"
"time" "time"
@ -34,19 +35,48 @@ func setupForMode(t *testing.T, m testMode) {
test.Prepare(t, config.Conf) test.Prepare(t, config.Conf)
registry.Init(nil) registry.Init(nil)
var client *tai.Client
var err error
if m.Addr == "local" { if m.Addr == "local" {
dataDir := t.TempDir() dataDir := t.TempDir()
vol := volume.NewLocal(dataDir) vol := volume.NewLocal(dataDir)
client, err = tai.New("local", tai.WithVolume(vol), tai.WithDataDir(dataDir)) res, err := tai.DialLocal("", dataDir, vol)
if err != nil {
t.Fatalf("DialLocal: %v", err)
}
reg := registry.Global()
reg.Register(&registry.TaiNode{TaiID: "local", Mode: "local"})
reg.SetResources("local", res)
t.Cleanup(func() { res.Close() })
} else { } else {
client, err = tai.New(m.Addr) host, grpcPort := parseHostPort(m.Addr)
ports := tai.Ports{GRPC: grpcPort}
res, err := tai.DialRemote(host, ports)
if err != nil {
t.Fatalf("DialRemote(%s): %v", m.Addr, err)
}
taiID := taiIDFromAddr(m.Addr)
reg := registry.Global()
reg.Register(&registry.TaiNode{TaiID: taiID, Mode: "direct"})
reg.SetResources(taiID, res)
t.Cleanup(func() { res.Close() })
} }
if err != nil { }
t.Fatalf("tai.New(%s): %v", m.Addr, err)
func taiIDFromAddr(addr string) string {
addr = strings.TrimPrefix(addr, "tai://")
parts := strings.SplitN(addr, ":", 2)
return parts[0]
}
func parseHostPort(addr string) (string, int) {
addr = strings.TrimPrefix(addr, "tai://")
parts := strings.SplitN(addr, ":", 2)
h := parts[0]
if len(parts) == 2 {
if p, err := strconv.Atoi(parts[1]); err == nil {
return h, p
}
} }
t.Cleanup(func() { client.Close() }) return h, 19100
} }
func setupGlobal(t *testing.T) { func setupGlobal(t *testing.T) {

View file

@ -8,6 +8,7 @@ import (
"github.com/yaoapp/yao/tai" "github.com/yaoapp/yao/tai"
"github.com/yaoapp/yao/tai/registry" "github.com/yaoapp/yao/tai/registry"
taitypes "github.com/yaoapp/yao/tai/types"
"github.com/yaoapp/yao/tai/volume" "github.com/yaoapp/yao/tai/volume"
taiworkspace "github.com/yaoapp/yao/tai/workspace" taiworkspace "github.com/yaoapp/yao/tai/workspace"
) )
@ -20,7 +21,6 @@ func M() *Manager {
} }
// Manager owns workspace CRUD, file I/O, and node management. // Manager owns workspace CRUD, file I/O, and node management.
// All node/client lookups go through tai.GetClient → registry.
type Manager struct{} type Manager struct{}
// NewManager creates a workspace manager. // NewManager creates a workspace manager.
@ -34,7 +34,7 @@ func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Workspace, e
return nil, ErrNodeMissing return nil, ErrNodeMissing
} }
client, ok := tai.GetClient(opts.Node) res, ok := tai.GetResources(opts.Node)
if !ok { if !ok {
return nil, ErrNodeOffline return nil, ErrNodeOffline
} }
@ -55,8 +55,7 @@ func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Workspace, e
UpdatedAt: now, UpdatedAt: now,
} }
vol := client.Volume() vol := res.Volume
if err := vol.MkdirAll(ctx, id, "."); err != nil { if err := vol.MkdirAll(ctx, id, "."); err != nil {
return nil, fmt.Errorf("workspace: create directory: %w", err) return nil, fmt.Errorf("workspace: create directory: %w", err)
} }
@ -73,14 +72,13 @@ func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Workspace, e
} }
// Get returns a workspace by ID. // Get returns a workspace by ID.
// Scans all registered nodes.
func (m *Manager) Get(ctx context.Context, id string) (*Workspace, error) { func (m *Manager) Get(ctx context.Context, id string) (*Workspace, error) {
for _, snap := range listNodes() { for _, snap := range listNodes() {
client, ok := tai.GetClient(snap.TaiID) res, ok := tai.GetResources(snap.TaiID)
if !ok { if !ok {
continue continue
} }
ws, err := readMeta(ctx, client, id) ws, err := readMeta(ctx, res.Volume, id)
if err != nil { if err != nil {
continue continue
} }
@ -99,11 +97,11 @@ func (m *Manager) List(ctx context.Context, opts ListOptions) ([]*Workspace, err
if opts.Node != "" && snap.TaiID != opts.Node { if opts.Node != "" && snap.TaiID != opts.Node {
continue continue
} }
client, ok := tai.GetClient(snap.TaiID) res, ok := tai.GetResources(snap.TaiID)
if !ok { if !ok {
continue continue
} }
entries, err := client.Volume().ListDir(ctx, "", ".") entries, err := res.Volume.ListDir(ctx, "", ".")
if err != nil { if err != nil {
continue continue
} }
@ -111,7 +109,7 @@ func (m *Manager) List(ctx context.Context, opts ListOptions) ([]*Workspace, err
if !e.IsDir { if !e.IsDir {
continue continue
} }
ws, err := readMeta(ctx, client, e.Path) ws, err := readMeta(ctx, res.Volume, e.Path)
if err != nil { if err != nil {
continue continue
} }
@ -128,9 +126,8 @@ func (m *Manager) List(ctx context.Context, opts ListOptions) ([]*Workspace, err
} }
// Update modifies workspace metadata (Name, Labels). // Update modifies workspace metadata (Name, Labels).
// Node and Owner are immutable after creation.
func (m *Manager) Update(ctx context.Context, id string, opts UpdateOptions) (*Workspace, error) { func (m *Manager) Update(ctx context.Context, id string, opts UpdateOptions) (*Workspace, error) {
ws, client, err := m.resolve(ctx, id) ws, vol, err := m.resolve(ctx, id)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -147,7 +144,7 @@ func (m *Manager) Update(ctx context.Context, id string, opts UpdateOptions) (*W
if err != nil { if err != nil {
return nil, err return nil, err
} }
if err := client.Volume().WriteFile(ctx, id, metadataFile, data, 0644); err != nil { if err := vol.WriteFile(ctx, id, metadataFile, data, 0644); err != nil {
return nil, fmt.Errorf("workspace: write metadata: %w", err) return nil, fmt.Errorf("workspace: write metadata: %w", err)
} }
return ws, nil return ws, nil
@ -155,12 +152,10 @@ func (m *Manager) Update(ctx context.Context, id string, opts UpdateOptions) (*W
// Delete removes workspace storage from the node. // Delete removes workspace storage from the node.
func (m *Manager) Delete(ctx context.Context, id string, force bool) error { func (m *Manager) Delete(ctx context.Context, id string, force bool) error {
_, client, err := m.resolve(ctx, id) _, vol, err := m.resolve(ctx, id)
if err != nil { if err != nil {
return err return err
} }
vol := client.Volume()
if err := vol.Remove(ctx, id, ".", true); err != nil { if err := vol.Remove(ctx, id, ".", true); err != nil {
return fmt.Errorf("workspace: remove: %w", err) return fmt.Errorf("workspace: remove: %w", err)
} }
@ -182,39 +177,39 @@ func (m *Manager) Nodes() []NodeInfo {
// FS returns an fs.FS-compatible filesystem for the given workspace. // FS returns an fs.FS-compatible filesystem for the given workspace.
func (m *Manager) FS(ctx context.Context, id string) (taiworkspace.FS, error) { func (m *Manager) FS(ctx context.Context, id string) (taiworkspace.FS, error) {
_, client, err := m.resolve(ctx, id) _, vol, err := m.resolve(ctx, id)
if err != nil { if err != nil {
return nil, err return nil, err
} }
return client.Workspace(id), nil return taiworkspace.New(vol, id), nil
} }
// ReadFile reads a file from the workspace. // ReadFile reads a file from the workspace.
func (m *Manager) ReadFile(ctx context.Context, id string, path string) ([]byte, error) { func (m *Manager) ReadFile(ctx context.Context, id string, path string) ([]byte, error) {
_, client, err := m.resolve(ctx, id) _, vol, err := m.resolve(ctx, id)
if err != nil { if err != nil {
return nil, err return nil, err
} }
data, _, err := client.Volume().ReadFile(ctx, id, path) data, _, err := vol.ReadFile(ctx, id, path)
return data, err return data, err
} }
// WriteFile writes a file to the workspace. // WriteFile writes a file to the workspace.
func (m *Manager) WriteFile(ctx context.Context, id string, path string, data []byte, perm os.FileMode) error { func (m *Manager) WriteFile(ctx context.Context, id string, path string, data []byte, perm os.FileMode) error {
_, client, err := m.resolve(ctx, id) _, vol, err := m.resolve(ctx, id)
if err != nil { if err != nil {
return err return err
} }
return client.Volume().WriteFile(ctx, id, path, data, perm) return vol.WriteFile(ctx, id, path, data, perm)
} }
// ListDir lists entries in a workspace directory. // ListDir lists entries in a workspace directory.
func (m *Manager) ListDir(ctx context.Context, id string, path string) ([]DirEntry, error) { func (m *Manager) ListDir(ctx context.Context, id string, path string) ([]DirEntry, error) {
_, client, err := m.resolve(ctx, id) _, vol, err := m.resolve(ctx, id)
if err != nil { if err != nil {
return nil, err return nil, err
} }
entries, err := client.Volume().ListDir(ctx, id, path) entries, err := vol.ListDir(ctx, id, path)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -231,42 +226,41 @@ func (m *Manager) ListDir(ctx context.Context, id string, path string) ([]DirEnt
// Remove deletes a file or directory from the workspace. // Remove deletes a file or directory from the workspace.
func (m *Manager) Remove(ctx context.Context, id string, path string) error { func (m *Manager) Remove(ctx context.Context, id string, path string) error {
_, client, err := m.resolve(ctx, id) _, vol, err := m.resolve(ctx, id)
if err != nil { if err != nil {
return err return err
} }
return client.Volume().Remove(ctx, id, path, true) return vol.Remove(ctx, id, path, true)
} }
// Rename renames a file or directory within the workspace. // Rename renames a file or directory within the workspace.
func (m *Manager) Rename(ctx context.Context, id string, oldPath, newPath string) error { func (m *Manager) Rename(ctx context.Context, id string, oldPath, newPath string) error {
_, client, err := m.resolve(ctx, id) _, vol, err := m.resolve(ctx, id)
if err != nil { if err != nil {
return err return err
} }
return client.Volume().Rename(ctx, id, oldPath, newPath) return vol.Rename(ctx, id, oldPath, newPath)
} }
// MkdirAll creates a directory (and parents) in the workspace. // MkdirAll creates a directory (and parents) in the workspace.
func (m *Manager) MkdirAll(ctx context.Context, id string, path string) error { func (m *Manager) MkdirAll(ctx context.Context, id string, path string) error {
_, client, err := m.resolve(ctx, id) _, vol, err := m.resolve(ctx, id)
if err != nil { if err != nil {
return err return err
} }
return client.Volume().MkdirAll(ctx, id, path) return vol.MkdirAll(ctx, id, path)
} }
// Volume returns the Volume interface for the node hosting the given workspace. // Volume returns the Volume interface for the node hosting the given workspace.
func (m *Manager) Volume(ctx context.Context, id string) (volume.Volume, string, error) { func (m *Manager) Volume(ctx context.Context, id string) (volume.Volume, string, error) {
_, client, err := m.resolve(ctx, id) _, vol, err := m.resolve(ctx, id)
if err != nil { if err != nil {
return nil, "", err return nil, "", err
} }
return client.Volume(), id, nil return vol, id, nil
} }
// NodeForWorkspace returns the node name for a given workspace ID. // NodeForWorkspace returns the node name for a given workspace ID.
// Used by sandbox.Manager to route container creation to the correct pool.
func (m *Manager) NodeForWorkspace(ctx context.Context, id string) (string, error) { func (m *Manager) NodeForWorkspace(ctx context.Context, id string) (string, error) {
ws, _, err := m.resolve(ctx, id) ws, _, err := m.resolve(ctx, id)
if err != nil { if err != nil {
@ -275,47 +269,55 @@ func (m *Manager) NodeForWorkspace(ctx context.Context, id string) (string, erro
return ws.Node, nil return ws.Node, nil
} }
// MountPath returns the host-side directory path for a workspace, // MountPath returns the host-side directory path for a workspace.
// suitable for use as a Docker bind mount source.
func (m *Manager) MountPath(ctx context.Context, id string) (string, error) { func (m *Manager) MountPath(ctx context.Context, id string) (string, error) {
_, client, err := m.resolve(ctx, id) _, vol, err := m.resolve(ctx, id)
if err != nil { if err != nil {
return "", err return "", err
} }
dataDir := client.DataDir() _ = vol
if dataDir == "" { for _, snap := range listNodes() {
return "", nil res, ok := tai.GetResources(snap.TaiID)
if !ok {
continue
}
if res.Volume == vol {
if res.DataDir == "" {
return "", nil
}
return res.DataDir + "/" + id, nil
}
} }
return dataDir + "/" + id, nil return "", nil
} }
// --- internal --- // --- internal ---
// resolve finds the workspace and its tai.Client by scanning all registered nodes. // resolve finds the workspace and its Volume by scanning all registered nodes.
func (m *Manager) resolve(ctx context.Context, id string) (*Workspace, *tai.Client, error) { func (m *Manager) resolve(ctx context.Context, id string) (*Workspace, volume.Volume, error) {
for _, snap := range listNodes() { for _, snap := range listNodes() {
client, ok := tai.GetClient(snap.TaiID) res, ok := tai.GetResources(snap.TaiID)
if !ok { if !ok {
continue continue
} }
ws, err := readMeta(ctx, client, id) ws, err := readMeta(ctx, res.Volume, id)
if err != nil { if err != nil {
continue continue
} }
return ws, client, nil return ws, res.Volume, nil
} }
return nil, nil, ErrNotFound return nil, nil, ErrNotFound
} }
func readMeta(ctx context.Context, client *tai.Client, id string) (*Workspace, error) { func readMeta(ctx context.Context, vol volume.Volume, id string) (*Workspace, error) {
data, _, err := client.Volume().ReadFile(ctx, id, metadataFile) data, _, err := vol.ReadFile(ctx, id, metadataFile)
if err != nil { if err != nil {
return nil, err return nil, err
} }
return unmarshalMeta(data) return unmarshalMeta(data)
} }
func listNodes() []registry.NodeSnapshot { func listNodes() []taitypes.NodeMeta {
reg := registry.Global() reg := registry.Global()
if reg == nil { if reg == nil {
return nil return nil

View file

@ -4,6 +4,7 @@ import (
"context" "context"
"net/url" "net/url"
"os" "os"
"strconv"
"strings" "strings"
"testing" "testing"
"time" "time"
@ -60,32 +61,52 @@ func ensureRegistry(tb testing.TB) {
func setupManagerForPool(tb testing.TB, pc poolConfig) *workspace.Manager { func setupManagerForPool(tb testing.TB, pc poolConfig) *workspace.Manager {
tb.Helper() tb.Helper()
ensureRegistry(tb) ensureRegistry(tb)
registerClient(tb, pc) registerForTest(tb, pc)
return workspace.NewManager() return workspace.NewManager()
} }
func registerClient(tb testing.TB, pc poolConfig) *tai.Client { func registerForTest(tb testing.TB, pc poolConfig) {
tb.Helper() tb.Helper()
if pc.Addr == "local" { if pc.Addr == "local" {
return localClient(tb, tb.TempDir()) registerLocalForTest(tb, tb.TempDir())
return
} }
client, err := tai.New(pc.Addr) host, grpcPort := parseHostPort(pc.Addr)
ports := tai.Ports{GRPC: grpcPort}
res, err := tai.DialRemote(host, ports)
if err != nil { if err != nil {
tb.Fatalf("tai.New(%s): %v", pc.Addr, err) tb.Fatalf("DialRemote(%s): %v", pc.Addr, err)
} }
tb.Cleanup(func() { client.Close() }) taiID := taiIDFromAddr(pc.Addr)
return client reg := registry.Global()
reg.Register(&registry.TaiNode{TaiID: taiID, Mode: "direct"})
reg.SetResources(taiID, res)
tb.Cleanup(func() { res.Close() })
} }
func localClient(tb testing.TB, dataDir string) *tai.Client { func registerLocalForTest(tb testing.TB, dataDir string) {
tb.Helper() tb.Helper()
vol := volume.NewLocal(dataDir) vol := volume.NewLocal(dataDir)
client, err := tai.New("local", tai.WithVolume(vol), tai.WithDataDir(dataDir)) res, err := tai.DialLocal("", dataDir, vol)
if err != nil { if err != nil {
tb.Fatalf("tai.New local: %v", err) tb.Fatalf("DialLocal: %v", err)
} }
tb.Cleanup(func() { client.Close() }) reg := registry.Global()
return client reg.Register(&registry.TaiNode{TaiID: "local", Mode: "local"})
reg.SetResources("local", res)
tb.Cleanup(func() { res.Close() })
}
func parseHostPort(addr string) (string, int) {
addr = strings.TrimPrefix(addr, "tai://")
parts := strings.SplitN(addr, ":", 2)
h := parts[0]
if len(parts) == 2 {
if p, err := strconv.Atoi(parts[1]); err == nil {
return h, p
}
}
return h, 19100
} }
func setupManagerMultiNode(t *testing.T) (*workspace.Manager, string, string) { func setupManagerMultiNode(t *testing.T) (*workspace.Manager, string, string) {
@ -94,19 +115,26 @@ func setupManagerMultiNode(t *testing.T) (*workspace.Manager, string, string) {
dir1 := t.TempDir() dir1 := t.TempDir()
vol1 := volume.NewLocal(dir1) vol1 := volume.NewLocal(dir1)
_, err := tai.New("docker://node-a", tai.WithVolume(vol1), tai.WithDataDir(dir1)) res1, err := tai.DialLocal("", dir1, vol1)
if err != nil { if err != nil {
t.Fatalf("tai.New node-a: %v", err) t.Fatalf("DialLocal node-a: %v", err)
} }
reg := registry.Global()
reg.Register(&registry.TaiNode{TaiID: "node-a", Mode: "local"})
reg.SetResources("node-a", res1)
t.Cleanup(func() { res1.Close() })
dir2 := t.TempDir() dir2 := t.TempDir()
vol2 := volume.NewLocal(dir2) vol2 := volume.NewLocal(dir2)
_, err = tai.New("docker://node-b", tai.WithVolume(vol2), tai.WithDataDir(dir2)) res2, err := tai.DialLocal("", dir2, vol2)
if err != nil { if err != nil {
t.Fatalf("tai.New node-b: %v", err) t.Fatalf("DialLocal node-b: %v", err)
} }
reg.Register(&registry.TaiNode{TaiID: "node-b", Mode: "local"})
reg.SetResources("node-b", res2)
t.Cleanup(func() { res2.Close() })
return workspace.NewManager(), "docker://node-a", "docker://node-b" return workspace.NewManager(), "node-a", "node-b"
} }
func createWorkspace(tb testing.TB, m *workspace.Manager, node string, opts ...func(*workspace.CreateOptions)) *workspace.Workspace { func createWorkspace(tb testing.TB, m *workspace.Manager, node string, opts ...func(*workspace.CreateOptions)) *workspace.Workspace {