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

Merge pull request #1486 from trheyi/main

Browse source 
Full-Stack gRPC Integration & OAuth Device Authorization
This commit is contained in:
Max 2026-03-04 17:32:30 +08:00 • committed by GitHub
commit 861243fbe9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
70 changed files with 11477 additions and 1664 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1537,6 +1537,16 @@ jobs:
# =============================================================================
TaiTest:
runs-on: ubuntu-latest
services:
mongodb:
image: mongo:6.0
ports:
- 27017:27017
env:
MONGO_INITDB_ROOT_USERNAME: root
MONGO_INITDB_ROOT_PASSWORD: 123456
MONGO_INITDB_DATABASE: test
strategy:
matrix:
go: ["1.25"]
@ -1647,11 +1657,28 @@ jobs:
with:
ref: ${{ env.HEAD }}
- name: Setup Apple Private Key
run: |
mkdir -p ../app/openapi/certs/apple
echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8
- name: Setup Go ${{ matrix.go }}
uses: actions/setup-go@v5
with:
go-version: ${{ matrix.go }}
- name: Start Redis
run: docker run --name redis --publish 6379:6379 --detach redis:6
- name: Setup Go Tools
run: make tools
- name: Setup ENV (SQLite)
run: |
mkdir -p ${{ github.WORKSPACE }}/../app/db
echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV
echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV
- name: Pull Tai & Test Images
run: |
docker pull yaoapp/tai:latest
@ -1677,14 +1704,42 @@ jobs:
-p 8080:8080 -p 9100:9100 -p 2375:2375 -p 6080:6080 -p 6443:6443 \
-e TAI_K8S_UPSTREAM="tcp://${K3D_IP}:6443" \
yaoapp/tai:latest
TAI_HTTP_READY=false
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1; then
echo "Tai is ready"
echo "Tai HTTP is ready"
TAI_HTTP_READY=true
break
fi
echo "Waiting for Tai... ($i)"
echo "Waiting for Tai HTTP... ($i)"
sleep 1
done
if [ "$TAI_HTTP_READY" != "true" ]; then
echo "::error::Tai HTTP failed to become ready within 30s"
echo "--- Tai container logs ---"
docker logs tai 2>&1 || true
echo "--- Tai container status ---"
docker inspect tai --format='{{.State.Status}} exit={{.State.ExitCode}}' || true
exit 1
fi
TAI_GRPC_READY=false
for i in $(seq 1 15); do
if nc -z 127.0.0.1 9100 2>/dev/null; then
echo "Tai gRPC is ready"
TAI_GRPC_READY=true
break
fi
echo "Waiting for Tai gRPC... ($i)"
sleep 1
done
if [ "$TAI_GRPC_READY" != "true" ]; then
echo "::error::Tai gRPC failed to become ready within 15s"
echo "--- Tai container logs ---"
docker logs tai 2>&1 || true
exit 1
fi
- name: Generate kubeconfig for Tai K8s proxy
run: |
@ -1702,6 +1757,7 @@ jobs:
TAI_TEST_K8S_HOST: "127.0.0.1"
TAI_TEST_K8S_PORT: "6443"
TAI_TEST_KUBECONFIG: "${{ runner.temp }}/kubeconfig-tai.yml"
TAI_TEST_HOST_IP: "172.17.0.1"
run: make unit-test-tai
- name: Codecov Report
@ -1722,3 +1778,172 @@ jobs:
issue_number: issue_number,
body: '✅ Tai SDK Tests passed!'
});
# =============================================================================
# gRPC Tests - Run once with SQLite (transport layer, no DB matrix needed)
# =============================================================================
GRPCTest:
runs-on: ubuntu-latest
services:
mongodb:
image: mongo:6.0
ports:
- 27017:27017
env:
MONGO_INITDB_ROOT_USERNAME: root
MONGO_INITDB_ROOT_PASSWORD: 123456
MONGO_INITDB_DATABASE: test
strategy:
matrix:
go: ["1.25"]
if: >
${{ github.event.workflow_run.event == 'pull_request' &&
github.event.workflow_run.conclusion == 'success' }}
steps:
- name: "Download artifact"
uses: actions/github-script@v7
with:
script: |
var artifacts = await github.rest.actions.listWorkflowRunArtifacts({
owner: context.repo.owner,
repo: context.repo.repo,
run_id: ${{github.event.workflow_run.id }},
});
var matchArtifact = artifacts.data.artifacts.filter((artifact) => {
return artifact.name == "pr"
})[0];
var download = await github.rest.actions.downloadArtifact({
owner: context.repo.owner,
repo: context.repo.repo,
artifact_id: matchArtifact.id,
archive_format: 'zip',
});
var fs = require('fs');
fs.writeFileSync('${{github.workspace}}/pr.zip', Buffer.from(download.data));
- name: "Read NR & SHA"
run: |
unzip pr.zip
cat NR
cat SHA
echo HEAD=$(cat SHA) >> $GITHUB_ENV
echo NR=$(cat NR) >> $GITHUB_ENV
- name: "Comment on PR"
uses: actions/github-script@v7
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const { NR } = process.env
var issue_number = NR;
await github.rest.issues.createComment({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: issue_number,
body: '🤖 gRPC Tests running with SQLite...'
});
- name: Checkout Kun
uses: actions/checkout@v4
with:
repository: yaoapp/kun
path: kun
- name: Checkout Xun
uses: actions/checkout@v4
with:
repository: yaoapp/xun
path: xun
- name: Checkout Gou
uses: actions/checkout@v4
with:
repository: yaoapp/gou
path: gou
- name: Checkout V8Go
uses: actions/checkout@v4
with:
repository: yaoapp/v8go
path: v8go
- name: Unzip libv8
run: |
files=$(find ./v8go -name "libv8*.zip")
for file in $files; do
dir=$(dirname "$file")
echo "Extracting $file to directory $dir"
unzip -o -d $dir $file
rm -rf $dir/__MACOSX
done
- name: Checkout Demo App
uses: actions/checkout@v4
with:
repository: yaoapp/yao-dev-app
path: app
- name: Checkout Extension
uses: actions/checkout@v4
with:
repository: yaoapp/yao-extensions-dev
path: extension
- name: Move Dependencies
run: |
mv kun ../
mv xun ../
mv gou ../
mv v8go ../
mv app ../
mv extension ../
- name: Checkout pull request HEAD commit
uses: actions/checkout@v4
with:
ref: ${{ env.HEAD }}
- name: Setup Apple Private Key
run: |
mkdir -p ../app/openapi/certs/apple
echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8
- name: Setup Go ${{ matrix.go }}
uses: actions/setup-go@v5
with:
go-version: ${{ matrix.go }}
- name: Start Redis
run: docker run --name redis --publish 6379:6379 --detach redis:6
- name: Setup Go Tools
run: make tools
- name: Setup ENV (SQLite)
run: |
mkdir -p ${{ github.WORKSPACE }}/../app/db
echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV
echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV
- name: Run gRPC Tests
run: make unit-test-grpc
- name: Codecov Report
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
- name: "Comment on PR - gRPC Tests Done"
uses: actions/github-script@v7
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const { NR } = process.env
var issue_number = NR;
await github.rest.issues.createComment({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: issue_number,
body: '✅ gRPC Tests passed!'
});

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

@ -1139,6 +1139,16 @@ jobs:
# =============================================================================
tai-test:
runs-on: ubuntu-latest
services:
mongodb:
image: mongo:6.0
ports:
- 27017:27017
env:
MONGO_INITDB_ROOT_USERNAME: root
MONGO_INITDB_ROOT_PASSWORD: 123456
MONGO_INITDB_DATABASE: test
strategy:
matrix:
go: ["1.25"]
@ -1201,11 +1211,28 @@ jobs:
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Apple Private Key
run: |
mkdir -p ../app/openapi/certs/apple
echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8
- name: Setup Go ${{ matrix.go }}
uses: actions/setup-go@v5
with:
go-version: ${{ matrix.go }}
- name: Start Redis
run: docker run --name redis --publish 6379:6379 --detach redis:6
- name: Setup Go Tools
run: make tools
- name: Setup ENV (SQLite)
run: |
mkdir -p ${{ github.WORKSPACE }}/../app/db
echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV
echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV
- name: Pull Tai & Test Images
run: |
docker pull yaoapp/tai:latest
@ -1231,14 +1258,42 @@ jobs:
-p 8080:8080 -p 9100:9100 -p 2375:2375 -p 6080:6080 -p 6443:6443 \
-e TAI_K8S_UPSTREAM="tcp://${K3D_IP}:6443" \
yaoapp/tai:latest
TAI_HTTP_READY=false
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1; then
echo "Tai is ready"
echo "Tai HTTP is ready"
TAI_HTTP_READY=true
break
fi
echo "Waiting for Tai... ($i)"
echo "Waiting for Tai HTTP... ($i)"
sleep 1
done
if [ "$TAI_HTTP_READY" != "true" ]; then
echo "::error::Tai HTTP failed to become ready within 30s"
echo "--- Tai container logs ---"
docker logs tai 2>&1 || true
echo "--- Tai container status ---"
docker inspect tai --format='{{.State.Status}} exit={{.State.ExitCode}}' || true
exit 1
fi
TAI_GRPC_READY=false
for i in $(seq 1 15); do
if nc -z 127.0.0.1 9100 2>/dev/null; then
echo "Tai gRPC is ready"
TAI_GRPC_READY=true
break
fi
echo "Waiting for Tai gRPC... ($i)"
sleep 1
done
if [ "$TAI_GRPC_READY" != "true" ]; then
echo "::error::Tai gRPC failed to become ready within 15s"
echo "--- Tai container logs ---"
docker logs tai 2>&1 || true
exit 1
fi
- name: Generate kubeconfig for Tai K8s proxy
run: |
@ -1256,9 +1311,117 @@ jobs:
TAI_TEST_K8S_HOST: "127.0.0.1"
TAI_TEST_K8S_PORT: "6443"
TAI_TEST_KUBECONFIG: "${{ runner.temp }}/kubeconfig-tai.yml"
TAI_TEST_HOST_IP: "172.17.0.1"
run: make unit-test-tai
- name: Codecov Report
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
# =============================================================================
# gRPC Tests - Run once with SQLite (transport layer, no DB matrix needed)
# =============================================================================
grpc-test:
runs-on: ubuntu-latest
services:
mongodb:
image: mongo:6.0
ports:
- 27017:27017
env:
MONGO_INITDB_ROOT_USERNAME: root
MONGO_INITDB_ROOT_PASSWORD: 123456
MONGO_INITDB_DATABASE: test
strategy:
matrix:
go: ["1.25"]
steps:
- name: Checkout Kun
uses: actions/checkout@v4
with:
repository: ${{ env.REPO_KUN }}
path: kun
- name: Checkout Xun
uses: actions/checkout@v4
with:
repository: ${{ env.REPO_XUN }}
path: xun
- name: Checkout Gou
uses: actions/checkout@v4
with:
repository: ${{ env.REPO_GOU }}
path: gou
- name: Checkout V8Go
uses: actions/checkout@v4
with:
repository: yaoapp/v8go
path: v8go
- name: Unzip libv8
run: |
files=$(find ./v8go -name "libv8*.zip")
for file in $files; do
dir=$(dirname "$file")
echo "Extracting $file to directory $dir"
unzip -o -d $dir $file
rm -rf $dir/__MACOSX
done
- name: Checkout Demo App
uses: actions/checkout@v4
with:
repository: yaoapp/yao-dev-app
path: app
- name: Checkout Extension
uses: actions/checkout@v4
with:
repository: yaoapp/yao-extensions-dev
path: extension
- name: Move Dependencies
run: |
mv kun ../
mv xun ../
mv gou ../
mv v8go ../
mv app ../
mv extension ../
- name: Checkout Code
uses: actions/checkout@v4
- name: Setup Apple Private Key
run: |
mkdir -p ../app/openapi/certs/apple
echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8
- name: Setup Go ${{ matrix.go }}
uses: actions/setup-go@v5
with:
go-version: ${{ matrix.go }}
- name: Start Redis
run: docker run --name redis --publish 6379:6379 --detach redis:6
- name: Setup Go Tools
run: make tools
- name: Setup ENV (SQLite)
run: |
mkdir -p ${{ github.WORKSPACE }}/../app/db
echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV
echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV
- name: Run gRPC Tests
run: make unit-test-grpc
- name: Codecov Report
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}

43
Makefile
View file

@ -11,8 +11,8 @@ OS := $(shell uname)
# ROOT_DIR := $(shell dirname $(realpath $(firstword $(MAKEFILE_LIST))))
TESTFOLDER := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|registry' | awk '!/\/tests\// || /openapi\/tests/')
# Core tests (exclude AI-related: agent, aigc, openai, KB, sandbox, registry, and integrations which require external services)
TESTFOLDER_CORE := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|agent|kb|sandbox|integrations|registry|tai' | awk '!/\/tests\// || /openapi\/tests/')
# Core tests (exclude AI-related: agent, aigc, openai, KB, sandbox, registry, grpc, and integrations which require external services)
TESTFOLDER_CORE := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|agent|kb|sandbox|integrations|registry|tai|grpc' | awk '!/\/tests\// || /openapi\/tests/')
# Agent tests (agent, aigc) - exclude agent/search/handlers/web (requires external API keys) and robot packages (tested in robot job)
TESTFOLDER_AGENT := $(shell $(GO) list ./agent/... ./aigc/... | grep -vE 'agent/search/handlers/web|agent/robot/')
# KB tests (kb)
@ -23,6 +23,8 @@ TESTFOLDER_ROBOT := $(shell $(GO) list ./agent/robot/... | grep -vE 'agent/robot
TESTFOLDER_SANDBOX := $(shell $(GO) list ./sandbox/...)
# Tai SDK tests (requires Tai container with Docker socket)
TESTFOLDER_TAI := $(shell $(GO) list ./tai/...)
# gRPC tests
TESTFOLDER_GRPC := $(shell $(GO) list ./grpc/...)
TESTTAGS ?= ""
# TESTWIDGETS := $(shell $(GO) list ./widgets/...)
@ -285,6 +287,43 @@ unit-test-tai:
@echo "All Tai SDK tests passed"
@echo "============================================="
# Proto codegen
.PHONY: proto
proto:
protoc --go_out=. --go_opt=paths=source_relative \
--go-grpc_out=. --go-grpc_opt=paths=source_relative \
grpc/pb/yao.proto
# gRPC Unit Test
.PHONY: unit-test-grpc
unit-test-grpc:
echo "mode: count" > coverage.out
for d in $(TESTFOLDER_GRPC); do \
$(GO) test -tags $(TESTTAGS) -v -timeout=10m \
-covermode=count -coverprofile=profile.out \
-coverpkg=$$(echo $$d | sed "s/\/test$$//g") \
-skip='TestMemoryLeak|TestIsolateDisposal|TestLeak_|TestScenario_' \
$$d > tmp.out; \
cat tmp.out; \
if grep -q "^--- FAIL" tmp.out; then \
rm tmp.out; \
exit 1; \
elif grep -q "build failed" tmp.out; then \
rm tmp.out; \
exit 1; \
elif grep -q "setup failed" tmp.out; then \
rm tmp.out; \
exit 1; \
elif grep -q "runtime error" tmp.out; then \
rm tmp.out; \
exit 1; \
fi; \
if [ -f profile.out ]; then \
cat profile.out | grep -v "mode:" >> coverage.out; \
rm profile.out; \
fi; \
done
# Benchmark Test
.PHONY: benchmark
benchmark:

153
agent/context/grpc.go Normal file
View file

@ -0,0 +1,153 @@
package context
import (
"context"
"encoding/json"
"fmt"
"net/http"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/store"
"github.com/yaoapp/yao/openapi/oauth/types"
)
// GRPCAgentInput holds the raw inputs from a gRPC AgentStream request.
type GRPCAgentInput struct {
AssistantID string
Messages []byte
Options []byte
AuthInfo *types.AuthorizedInfo
Cache store.Store
Writer http.ResponseWriter
}
// GetGRPCAgentRequest parses a gRPC agent request and creates a Context + Options,
// mirroring openapi.go GetCompletionRequest.
//
// Flow: validate → parse messages → parse options → build Context → build Options → register interrupt
func GetGRPCAgentRequest(parent context.Context, input GRPCAgentInput) ([]Message, *Context, *Options, error) {
if input.AssistantID == "" {
return nil, nil, nil, fmt.Errorf("assistant_id is required")
}
messages, err := parseGRPCMessages(input.Messages)
if err != nil {
return nil, nil, nil, err
}
var rawOpts map[string]interface{}
if len(input.Options) > 0 {
if err := json.Unmarshal(input.Options, &rawOpts); err != nil {
return nil, nil, nil, fmt.Errorf("invalid options JSON: %w", err)
}
}
chatID := getChatIDFromOpts(rawOpts)
ctx := New(parent, input.AuthInfo, chatID)
ctx.Cache = input.Cache
ctx.Writer = input.Writer
ctx.AssistantID = input.AssistantID
ctx.Locale = getStringOpt(rawOpts, "locale")
ctx.Theme = getStringOpt(rawOpts, "theme")
ctx.Referer = getRefererOpt(rawOpts)
ctx.Accept = getAcceptOpt(rawOpts)
ctx.Route = getStringOpt(rawOpts, "route")
ctx.Metadata = getMapOpt(rawOpts, "metadata")
ctx.Client = Client{Type: "grpc"}
opts := &Options{
Context: parent,
Skip: getSkipOpt(rawOpts),
Mode: getStringOpt(rawOpts, "mode"),
}
if connectorID := getStringOpt(rawOpts, "connector"); connectorID != "" {
if _, err := connector.Select(connectorID); err == nil {
opts.Connector = connectorID
}
}
ctx.Interrupt = NewInterruptController()
if err := Register(ctx); err != nil {
return nil, nil, nil, fmt.Errorf("failed to register context: %w", err)
}
ctx.Interrupt.Start(ctx.ID)
return messages, ctx, opts, nil
}
func parseGRPCMessages(raw []byte) ([]Message, error) {
if len(raw) == 0 {
return nil, fmt.Errorf("messages are required")
}
var messages []Message
if err := json.Unmarshal(raw, &messages); err != nil {
return nil, fmt.Errorf("invalid messages JSON: %w", err)
}
if len(messages) == 0 {
return nil, fmt.Errorf("messages must not be empty")
}
return messages, nil
}
func getChatIDFromOpts(opts map[string]interface{}) string {
if opts != nil {
if v, ok := opts["chat_id"].(string); ok && v != "" {
return v
}
}
return GenChatID()
}
func getStringOpt(opts map[string]interface{}, key string) string {
if opts == nil {
return ""
}
v, _ := opts[key].(string)
return v
}
func getRefererOpt(opts map[string]interface{}) string {
r := getStringOpt(opts, "referer")
if r != "" {
return validateReferer(r)
}
return RefererAPI
}
func getAcceptOpt(opts map[string]interface{}) Accept {
a := getStringOpt(opts, "accept")
if a != "" {
return validateAccept(a)
}
return AcceptStandard
}
func getMapOpt(opts map[string]interface{}, key string) map[string]interface{} {
if opts == nil {
return nil
}
v, _ := opts[key].(map[string]interface{})
return v
}
func getSkipOpt(opts map[string]interface{}) *Skip {
if opts == nil {
return nil
}
raw, ok := opts["skip"]
if !ok {
return nil
}
data, err := json.Marshal(raw)
if err != nil {
return nil
}
var skip Skip
if err := json.Unmarshal(data, &skip); err != nil {
return nil
}
return &skip
}

6
agent/llm/jsapi.go
View file

@ -201,6 +201,12 @@ func parseToolCalls(toolCalls []interface{}) []agentContext.ToolCall {
}
// buildCompletionOptions creates CompletionOptions from JS opts map
// BuildCompletionOptions builds CompletionOptions from a connector and raw opts map.
// Exported for reuse by gRPC handlers.
func BuildCompletionOptions(conn connector.Connector, opts map[string]interface{}) *agentContext.CompletionOptions {
return buildCompletionOptions(conn, opts)
}
func buildCompletionOptions(conn connector.Connector, opts map[string]interface{}) *agentContext.CompletionOptions {
// Get capabilities from connector
capabilities := GetCapabilitiesFromConn(conn)

118
cmd/credential.go Normal file
View file

@ -0,0 +1,118 @@
package cmd
import (
"encoding/base64"
"encoding/json"
"fmt"
"os"
"path/filepath"
"time"
)
// Credential represents the stored OAuth credential for gRPC mode.
type Credential struct {
Server string `json:"server"`
GRPCAddr string `json:"grpc_addr,omitempty"`
AccessToken string `json:"access_token"`
RefreshToken string `json:"refresh_token,omitempty"`
Scope string `json:"scope,omitempty"`
User string `json:"user,omitempty"`
ExpiresAt string `json:"expires_at,omitempty"`
}
// Expired returns true if the credential has an expires_at in the past.
func (c *Credential) Expired() bool {
if c.ExpiresAt == "" {
return false
}
t, err := time.Parse(time.RFC3339, c.ExpiresAt)
if err != nil {
return false
}
return time.Now().After(t)
}
func credentialPath() (string, error) {
home, err := os.UserHomeDir()
if err != nil {
return "", fmt.Errorf("cannot determine home directory: %w", err)
}
return filepath.Join(home, ".yao", "credentials"), nil
}
// LoadCredential reads and decodes ~/.yao/credentials. Returns nil if the file
// does not exist.
func LoadCredential() (*Credential, error) {
path, err := credentialPath()
if err != nil {
return nil, err
}
raw, err := os.ReadFile(path)
if os.IsNotExist(err) {
return nil, nil
}
if err != nil {
return nil, fmt.Errorf("read credentials: %w", err)
}
decoded, err := base64.StdEncoding.DecodeString(string(raw))
if err != nil {
return nil, fmt.Errorf("decode credentials: %w", err)
}
var cred Credential
if err := json.Unmarshal(decoded, &cred); err != nil {
return nil, fmt.Errorf("unmarshal credentials: %w", err)
}
return &cred, nil
}
// LoadCredentialFrom reads and decodes a credential file from a custom path.
func LoadCredentialFrom(path string) (*Credential, error) {
raw, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("read credentials from %s: %w", path, err)
}
decoded, err := base64.StdEncoding.DecodeString(string(raw))
if err != nil {
return nil, fmt.Errorf("decode credentials: %w", err)
}
var cred Credential
if err := json.Unmarshal(decoded, &cred); err != nil {
return nil, fmt.Errorf("unmarshal credentials: %w", err)
}
return &cred, nil
}
// SaveCredential encodes and writes the credential to ~/.yao/credentials.
func SaveCredential(cred *Credential) error {
path, err := credentialPath()
if err != nil {
return err
}
dir := filepath.Dir(path)
if err := os.MkdirAll(dir, 0700); err != nil {
return fmt.Errorf("create directory %s: %w", dir, err)
}
data, err := json.Marshal(cred)
if err != nil {
return fmt.Errorf("marshal credentials: %w", err)
}
encoded := base64.StdEncoding.EncodeToString(data)
if err := os.WriteFile(path, []byte(encoded), 0600); err != nil {
return fmt.Errorf("write credentials: %w", err)
}
return nil
}
// RemoveCredential deletes ~/.yao/credentials.
func RemoveCredential() error {
path, err := credentialPath()
if err != nil {
return err
}
if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("remove credentials: %w", err)
}
return nil
}

342
cmd/login.go Normal file
View file

@ -0,0 +1,342 @@
package cmd
import (
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"os"
"strings"
"time"
"github.com/fatih/color"
"github.com/spf13/cobra"
"github.com/yaoapp/yao/engine"
)
var loginServer string
var loginCmd = &cobra.Command{
Use: "login",
Short: L("Login to remote Yao server"),
Long: L("Login to remote Yao server using device authorization flow"),
Run: func(cmd *cobra.Command, args []string) {
if loginServer == "" {
color.Red(L("Missing --server flag\n"))
fmt.Println(" yao login --server https://yaoagents.com")
os.Exit(1)
}
serverURL := strings.TrimRight(loginServer, "/")
// 1. Discover OAuth endpoints via well-known metadata
endpoints, err := discoverEndpoints(serverURL)
if err != nil {
color.Red(" %s %s\n", L("Server discovery failed:"), err)
os.Exit(1)
}
// 2. Compute deterministic client_id from machine fingerprint
machine, err := engine.GetMachineID()
if err != nil {
color.Red("Failed to compute machine ID: %s\n", err)
os.Exit(1)
}
clientID := machine.ID
// 3. Register the client (idempotent for same client_id)
if endpoints.RegistrationEndpoint != "" {
if err := registerClient(endpoints.RegistrationEndpoint, clientID); err != nil {
color.Red("Client registration failed: %s\n", err)
os.Exit(1)
}
}
// 4. Start device authorization
deviceResp, err := requestDeviceAuthorization(endpoints.DeviceAuthorizationEndpoint, clientID)
if err != nil {
color.Red("Device authorization failed: %s\n", err)
os.Exit(1)
}
// 5. Display the code to the user
dashboard := endpoints.Dashboard
if dashboard == "" {
dashboard = "/admin"
}
verifyURI := strings.TrimRight(serverURL, "/") + dashboard + "/auth/device"
verifyURIComplete := verifyURI + "?user_code=" + deviceResp.UserCode
fmt.Println()
color.White(" %s %s\n",
L("Open:"),
color.CyanString(verifyURIComplete))
fmt.Println()
color.White(" %s %s\n",
L("Or visit:"),
color.CyanString(verifyURI))
color.White(" %s %s\n",
L("Enter code:"),
color.YellowString(deviceResp.UserCode))
fmt.Println()
// 6. Poll for token
interval := deviceResp.Interval
if interval < 5 {
interval = 5
}
color.White(" %s", L("Waiting for authorization..."))
tokenResp, err := pollForToken(endpoints.TokenEndpoint, clientID, deviceResp.DeviceCode, interval, deviceResp.ExpiresIn)
if err != nil {
fmt.Println()
color.Red("\n %s %s\n", L("Login failed:"), err)
os.Exit(1)
}
// 6. Save credential
expiresAt := ""
if tokenResp.ExpiresIn > 0 {
expiresAt = time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second).UTC().Format(time.RFC3339)
}
cred := &Credential{
Server: serverURL,
GRPCAddr: endpoints.GRPCAddr,
AccessToken: tokenResp.AccessToken,
RefreshToken: tokenResp.RefreshToken,
Scope: tokenResp.Scope,
User: parseJWTSubject(tokenResp.AccessToken),
ExpiresAt: expiresAt,
}
if err := SaveCredential(cred); err != nil {
color.Red("\n Failed to save credentials: %s\n", err)
os.Exit(1)
}
fmt.Print("\033[2J\033[H")
color.Green(" ✓ %s\n", L("Login successful"))
color.White(" %s %s\n", L("Server:"), serverURL)
if cred.GRPCAddr != "" {
color.White(" %s %s\n", L("gRPC:"), cred.GRPCAddr)
}
if cred.User != "" {
color.White(" %s %s\n", L("User:"), cred.User)
}
if cred.ExpiresAt != "" {
color.White(" %s %s\n", L("Expires:"), cred.ExpiresAt)
}
fmt.Println()
},
}
func init() {
loginCmd.PersistentFlags().StringVar(&loginServer, "server", "", L("Remote Yao server URL"))
}
// --- types ---
type oauthEndpoints struct {
RegistrationEndpoint string `json:"registration_endpoint"`
DeviceAuthorizationEndpoint string `json:"device_authorization_endpoint"`
TokenEndpoint string `json:"token_endpoint"`
RevocationEndpoint string `json:"revocation_endpoint"`
Dashboard string `json:"-"`
GRPCAddr string `json:"-"`
}
type deviceAuthResponse struct {
DeviceCode string `json:"device_code"`
UserCode string `json:"user_code"`
VerificationURI string `json:"verification_uri"`
VerificationURIComplete string `json:"verification_uri_complete"`
ExpiresIn int `json:"expires_in"`
Interval int `json:"interval"`
}
type tokenResponse struct {
AccessToken string `json:"access_token"`
TokenType string `json:"token_type"`
ExpiresIn int `json:"expires_in"`
RefreshToken string `json:"refresh_token"`
Scope string `json:"scope"`
}
type oauthError struct {
Error string `json:"error"`
ErrorDescription string `json:"error_description"`
}
// --- HTTP helpers ---
// discoverEndpoints fetches OAuth endpoint URLs from /.well-known/yao,
// using the openapi base prefix to construct correct API paths.
func discoverEndpoints(serverURL string) (*oauthEndpoints, error) {
return discoverFromYaoMetadata(serverURL)
}
type yaoMetadataResponse struct {
OpenAPI string `json:"openapi"`
Dashboard string `json:"dashboard"`
GRPC string `json:"grpc"`
}
func discoverFromYaoMetadata(serverURL string) (*oauthEndpoints, error) {
resp, err := http.Get(serverURL + "/.well-known/yao")
if err != nil {
return nil, fmt.Errorf("network error: %w", err)
}
defer resp.Body.Close()
body, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("/.well-known/yao returned %d", resp.StatusCode)
}
var meta yaoMetadataResponse
if err := json.Unmarshal(body, &meta); err != nil {
return nil, fmt.Errorf("invalid /.well-known/yao response: %w", err)
}
base := strings.TrimRight(serverURL, "/") + meta.OpenAPI
return &oauthEndpoints{
RegistrationEndpoint: base + "/oauth/register",
DeviceAuthorizationEndpoint: base + "/oauth/device_authorization",
TokenEndpoint: base + "/oauth/token",
RevocationEndpoint: base + "/oauth/revoke",
Dashboard: meta.Dashboard,
GRPCAddr: meta.GRPC,
}, nil
}
func registerClient(endpoint, clientID string) error {
body := fmt.Sprintf(
`{"client_id":"%s","client_name":"yao-cli","grant_types":["urn:ietf:params:oauth:grant-type:device_code"],"token_endpoint_auth_method":"none","redirect_uris":["http://localhost"]}`,
clientID,
)
resp, err := http.Post(endpoint, "application/json", strings.NewReader(body))
if err != nil {
return fmt.Errorf("network error: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusCreated {
return nil
}
respBody, _ := io.ReadAll(resp.Body)
var oerr oauthError
if json.Unmarshal(respBody, &oerr) == nil && oerr.Error == "invalid_client_metadata" {
return nil // client already registered, idempotent
}
return fmt.Errorf("registration returned %d: %s", resp.StatusCode, string(respBody))
}
func requestDeviceAuthorization(endpoint, clientID string) (*deviceAuthResponse, error) {
data := url.Values{
"client_id": {clientID},
"scope": {"grpc:run grpc:stream grpc:shell grpc:mcp grpc:llm grpc:agent"},
}
resp, err := http.PostForm(endpoint, data)
if err != nil {
return nil, fmt.Errorf("network error: %w", err)
}
defer resp.Body.Close()
respBody, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK {
var oerr oauthError
json.Unmarshal(respBody, &oerr)
if oerr.ErrorDescription != "" {
return nil, fmt.Errorf("%s", oerr.ErrorDescription)
}
return nil, fmt.Errorf("server returned %d: %s", resp.StatusCode, string(respBody))
}
var result deviceAuthResponse
if err := json.Unmarshal(respBody, &result); err != nil {
return nil, fmt.Errorf("invalid response: %w", err)
}
return &result, nil
}
func pollForToken(endpoint, clientID, deviceCode string, interval, expiresIn int) (*tokenResponse, error) {
deadline := time.Now().Add(time.Duration(expiresIn) * time.Second)
ticker := time.NewTicker(time.Duration(interval) * time.Second)
defer ticker.Stop()
for range ticker.C {
if time.Now().After(deadline) {
return nil, fmt.Errorf("device code expired")
}
data := url.Values{
"grant_type": {"urn:ietf:params:oauth:grant-type:device_code"},
"client_id": {clientID},
"device_code": {deviceCode},
}
resp, err := http.PostForm(endpoint, data)
if err != nil {
continue
}
respBody, _ := io.ReadAll(resp.Body)
resp.Body.Close()
if resp.StatusCode == http.StatusOK {
var tok tokenResponse
if err := json.Unmarshal(respBody, &tok); err != nil {
return nil, fmt.Errorf("invalid token response: %w", err)
}
return &tok, nil
}
var oerr oauthError
json.Unmarshal(respBody, &oerr)
switch oerr.Error {
case "authorization_pending":
fmt.Print(".")
continue
case "slow_down":
interval += 5
ticker.Reset(time.Duration(interval) * time.Second)
continue
case "expired_token":
return nil, fmt.Errorf("device code expired")
case "access_denied":
return nil, fmt.Errorf("authorization denied by user")
default:
desc := oerr.ErrorDescription
if desc == "" {
desc = oerr.Error
}
return nil, fmt.Errorf("%s", desc)
}
}
return nil, fmt.Errorf("device code expired")
}
// parseJWTSubject extracts the "sub" claim from a JWT access token
// without verifying the signature (display-only).
func parseJWTSubject(token string) string {
parts := strings.Split(token, ".")
if len(parts) != 3 {
return ""
}
payload, err := base64.RawURLEncoding.DecodeString(parts[1])
if err != nil {
return ""
}
var claims struct {
Sub string `json:"sub"`
}
if json.Unmarshal(payload, &claims) != nil {
return ""
}
return claims.Sub
}

80
cmd/logout.go Normal file
View file

@ -0,0 +1,80 @@
package cmd
import (
"net/http"
"net/url"
"os"
"strings"
"github.com/fatih/color"
"github.com/spf13/cobra"
)
var logoutCmd = &cobra.Command{
Use: "logout",
Short: L("Logout from remote Yao server"),
Long: L("Revoke token and remove stored credentials"),
Run: func(cmd *cobra.Command, args []string) {
cred, err := LoadCredential()
if err != nil {
color.Red(" %s %s\n", L("Failed to read credentials:"), err)
os.Exit(1)
}
if cred == nil {
color.Yellow(" %s\n", L("Not logged in"))
return
}
// Best-effort token revocation via discovery
if cred.AccessToken != "" && cred.Server != "" {
if ep, err := discoverEndpoints(cred.Server); err == nil && ep.RevocationEndpoint != "" {
revokeToken(ep.RevocationEndpoint, cred.AccessToken)
}
}
if err := RemoveCredential(); err != nil {
color.Red(" %s %s\n", L("Failed to remove credentials:"), err)
os.Exit(1)
}
color.Green(" ✓ %s\n", L("Logged out"))
if cred.Server != "" {
color.White(" %s %s\n", L("Server:"), cred.Server)
}
},
}
func revokeToken(endpoint, token string) {
data := url.Values{"token": {token}}
req, err := http.NewRequest("POST", endpoint, strings.NewReader(data.Encode()))
if err != nil {
return
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
http.DefaultClient.Do(req)
}
func init() {
// Add i18n entries
langs["Login to remote Yao server"] = "登录远程 Yao 服务器"
langs["Login to remote Yao server using device authorization flow"] = "使用设备授权流程登录远程 Yao 服务器"
langs["Remote Yao server URL"] = "远程 Yao 服务器地址"
langs["Logout from remote Yao server"] = "登出远程 Yao 服务器"
langs["Revoke token and remove stored credentials"] = "撤销令牌并移除存储的凭证"
langs["Missing --server flag"] = "缺少 --server 参数"
langs["Open:"] = "打开:"
langs["Or visit:"] = "或访问:"
langs["Enter code:"] = "输入设备码:"
langs["Waiting for authorization..."] = "等待授权..."
langs["Login failed:"] = "登录失败:"
langs["Login successful"] = "登录成功"
langs["Server:"] = "服务器:"
langs["Scope:"] = "授权范围:"
langs["Failed to read credentials:"] = "读取凭证失败:"
langs["Not logged in"] = "未登录"
langs["Failed to remove credentials:"] = "移除凭证失败:"
langs["Logged out"] = "已登出"
langs["Path to credentials file"] = "凭证文件路径"
langs["Failed to load credentials:"] = "加载凭证失败:"
langs["Server discovery failed:"] = "服务发现失败:"
}

2
cmd/root.go
View file

@ -189,6 +189,8 @@ func init() {
inspectCmd,
startCmd,
runCmd,
loginCmd,
logoutCmd,
// getCmd,
// dumpCmd,
// restoreCmd,

392
cmd/run.go
View file

@ -18,173 +18,271 @@ import (
"github.com/yaoapp/yao/engine"
ischedule "github.com/yaoapp/yao/schedule"
"github.com/yaoapp/yao/share"
taigrpc "github.com/yaoapp/yao/tai/grpc"
itask "github.com/yaoapp/yao/task"
)
var runSilent = false
var runAuthPath string
var runCmd = &cobra.Command{
Use: "run",
Short: L("Execute process"),
Long: L("Execute process"),
Run: func(cmd *cobra.Command, args []string) {
defer share.SessionStop()
defer plugin.KillAll()
defer func() {
err := exception.Catch(recover())
if err != nil {
if !runSilent {
color.Red(L("Fatal: %s\n"), err.Error())
return
}
fmt.Printf("%s\n", err.Error())
}
}()
// Resolve credential: --auth flag > ~/.yao/credentials > nil (local mode)
cred := resolveCredential()
// Auto-detect app root if not specified
if appPath == "" {
cwd, err := os.Getwd()
if err == nil {
if root, err := findAppRootFromPath(cwd); err == nil {
appPath = root
}
}
}
Boot()
// Set Runtime Mode
config.Conf.Runtime.Mode = "standard"
cfg := config.Conf
cfg.Session.IsCLI = true
if len(args) < 1 {
if !runSilent {
color.Red(L("Not enough arguments\n"))
color.White(share.BUILDNAME + " help\n")
return
}
fmt.Print(L("Not enough arguments\n"))
if cred != nil {
runGRPC(cred, args)
return
}
loadWarnings, err := engine.Load(cfg, engine.LoadOption{Action: "run"})
if err != nil {
if !runSilent {
color.Red(L("Engine: %s\n"), err.Error())
return
}
fmt.Printf("%s\n", err.Error())
return
}
name := args[0]
if !runSilent {
color.Green(L("Run: %s\n"), name)
}
pargs := []interface{}{}
for i, arg := range args {
if i == 0 {
continue
}
// Parse the arguments
if strings.HasPrefix(arg, "::") {
arg := strings.TrimPrefix(arg, "::")
var v interface{}
err := jsoniter.Unmarshal([]byte(arg), &v)
if err != nil {
color.Red(L("Arguments: %s\n"), err.Error())
return
}
pargs = append(pargs, v)
if !runSilent {
color.White("args[%d]: %s\n", i-1, arg)
}
} else if strings.HasPrefix(arg, "\\::") {
arg := "::" + strings.TrimPrefix(arg, "\\::")
pargs = append(pargs, arg)
if !runSilent {
color.White("args[%d]: %s\n", i-1, arg)
}
} else {
pargs = append(pargs, arg)
if !runSilent {
color.White("args[%d]: %s\n", i-1, arg)
}
}
}
// Start Tasks
itask.Start()
defer itask.Stop()
// Start Schedules
ischedule.Start()
defer ischedule.Stop()
process := process.NewWithContext(context.Background(), name, pargs...)
res, err := process.Exec()
if err != nil {
if !runSilent {
color.Red(L("Process: %s\n"), fmt.Sprintf("%s", strings.TrimPrefix(err.Error(), "Exception|404:")))
return
}
fmt.Printf("%s\n", err.Error())
return
}
if !runSilent {
if len(loadWarnings) > 0 {
fmt.Println(color.YellowString("---------------------------------"))
fmt.Println(color.YellowString(L("Warnings")))
fmt.Println(color.YellowString("---------------------------------"))
for _, warning := range loadWarnings {
fmt.Println(color.YellowString("[%s] %s", warning.Widget, warning.Error))
}
fmt.Printf("\n")
}
color.White("--------------------------------------\n")
color.White(L("%s Response\n"), name)
color.White("--------------------------------------\n")
helper.Dump(res)
color.White("--------------------------------------\n")
color.Green(L("✨DONE✨\n"))
return
}
// Silent mode output
switch res.(type) {
case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64, bool:
fmt.Printf("%v\n", res)
return
case string, []byte:
fmt.Printf("%s\n", res)
return
default:
txt, err := jsoniter.Marshal(res)
if err != nil {
fmt.Printf("%s\n", err.Error())
}
fmt.Printf("%s\n", txt)
}
runLocal(args)
},
}
func init() {
runCmd.PersistentFlags().BoolVarP(&runSilent, "silent", "s", false, L("Silent mode"))
runCmd.PersistentFlags().StringVar(&runAuthPath, "auth", "", L("Path to credentials file"))
}
// resolveCredential loads credential from --auth flag or default path.
func resolveCredential() *Credential {
if runAuthPath != "" {
cred, err := LoadCredentialFrom(runAuthPath)
if err != nil {
color.Red(" %s %s\n", L("Failed to load credentials:"), err)
os.Exit(1)
}
return cred
}
cred, _ := LoadCredential()
return cred
}
// runGRPC executes a process via the remote gRPC server.
func runGRPC(cred *Credential, args []string) {
if len(args) < 1 {
if !runSilent {
color.Red(L("Not enough arguments\n"))
color.White(share.BUILDNAME + " help\n")
} else {
fmt.Print(L("Not enough arguments\n"))
}
os.Exit(1)
}
if cred.GRPCAddr == "" {
color.Red(" %s\n", L("No gRPC address in credentials. Please re-login."))
os.Exit(1)
}
name := args[0]
if !runSilent {
color.Green(L("Run: %s gRPC: %s\n"), name, cred.GRPCAddr)
}
pargs := parseRunArgs(args)
argsJSON, err := jsoniter.Marshal(pargs)
if err != nil {
color.Red(" %s %s\n", L("Arguments:"), err.Error())
os.Exit(1)
}
tm := taigrpc.NewTokenManager(cred.AccessToken, cred.RefreshToken, "", "")
client, err := taigrpc.Dial(cred.GRPCAddr, tm)
if err != nil {
color.Red(" %s %s\n", L("gRPC connect failed:"), err.Error())
os.Exit(1)
}
defer client.Close()
data, err := client.Run(context.Background(), name, argsJSON, 0)
if err != nil {
if !runSilent {
color.Red(" %s %s\n", L("Process:"), err.Error())
} else {
fmt.Printf("%s\n", err.Error())
}
os.Exit(1)
}
if !runSilent {
color.White("--------------------------------------\n")
color.White(L("%s Response\n"), name)
color.White("--------------------------------------\n")
var res interface{}
if jsoniter.Unmarshal(data, &res) == nil {
helper.Dump(res)
} else {
fmt.Printf("%s\n", data)
}
color.White("--------------------------------------\n")
fmt.Printf("\033[32m✨DONE✨\033[0m \033[90mgRPC: %s\033[0m\n", cred.GRPCAddr)
} else {
fmt.Printf("%s\n", data)
}
}
// runLocal executes a process locally (existing behavior).
func runLocal(args []string) {
defer share.SessionStop()
defer plugin.KillAll()
defer func() {
err := exception.Catch(recover())
if err != nil {
if !runSilent {
color.Red(L("Fatal: %s\n"), err.Error())
return
}
fmt.Printf("%s\n", err.Error())
}
}()
// Auto-detect app root if not specified
if appPath == "" {
cwd, err := os.Getwd()
if err == nil {
if root, err := findAppRootFromPath(cwd); err == nil {
appPath = root
}
}
}
Boot()
// Set Runtime Mode
config.Conf.Runtime.Mode = "standard"
cfg := config.Conf
cfg.Session.IsCLI = true
if len(args) < 1 {
if !runSilent {
color.Red(L("Not enough arguments\n"))
color.White(share.BUILDNAME + " help\n")
return
}
fmt.Print(L("Not enough arguments\n"))
return
}
loadWarnings, err := engine.Load(cfg, engine.LoadOption{Action: "run"})
if err != nil {
if !runSilent {
color.Red(L("Engine: %s\n"), err.Error())
return
}
fmt.Printf("%s\n", err.Error())
return
}
name := args[0]
if !runSilent {
color.Green(L("Run: %s\n"), name)
}
pargs := parseRunArgs(args)
// Start Tasks
itask.Start()
defer itask.Stop()
// Start Schedules
ischedule.Start()
defer ischedule.Stop()
p := process.NewWithContext(context.Background(), name, pargs...)
res, err := p.Exec()
if err != nil {
if !runSilent {
color.Red(L("Process: %s\n"), fmt.Sprintf("%s", strings.TrimPrefix(err.Error(), "Exception|404:")))
return
}
fmt.Printf("%s\n", err.Error())
return
}
if !runSilent {
if len(loadWarnings) > 0 {
fmt.Println(color.YellowString("---------------------------------"))
fmt.Println(color.YellowString(L("Warnings")))
fmt.Println(color.YellowString("---------------------------------"))
for _, warning := range loadWarnings {
fmt.Println(color.YellowString("[%s] %s", warning.Widget, warning.Error))
}
fmt.Printf("\n")
}
color.White("--------------------------------------\n")
color.White(L("%s Response\n"), name)
color.White("--------------------------------------\n")
helper.Dump(res)
color.White("--------------------------------------\n")
color.Green(L("✨DONE✨\n"))
return
}
// Silent mode output
switch res.(type) {
case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64, bool:
fmt.Printf("%v\n", res)
return
case string, []byte:
fmt.Printf("%s\n", res)
return
default:
txt, err := jsoniter.Marshal(res)
if err != nil {
fmt.Printf("%s\n", err.Error())
}
fmt.Printf("%s\n", txt)
}
}
// parseRunArgs parses the CLI arguments into process arguments, handling :: prefixed JSON.
func parseRunArgs(args []string) []interface{} {
pargs := []interface{}{}
for i, arg := range args {
if i == 0 {
continue
}
if strings.HasPrefix(arg, "::") {
raw := strings.TrimPrefix(arg, "::")
var v interface{}
err := jsoniter.Unmarshal([]byte(raw), &v)
if err != nil {
color.Red(L("Arguments: %s\n"), err.Error())
return pargs
}
pargs = append(pargs, v)
if !runSilent {
color.White("args[%d]: %s\n", i-1, raw)
}
} else if strings.HasPrefix(arg, "\\::") {
cleaned := "::" + strings.TrimPrefix(arg, "\\::")
pargs = append(pargs, cleaned)
if !runSilent {
color.White("args[%d]: %s\n", i-1, cleaned)
}
} else {
pargs = append(pargs, arg)
if !runSilent {
color.White("args[%d]: %s\n", i-1, arg)
}
}
}
return pargs
}
// findAppRootFromPath finds the Yao application root directory by looking for app.yao

15
cmd/start.go
View file

@ -28,6 +28,8 @@ import (
agentcontext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/engine"
yaogrpc "github.com/yaoapp/yao/grpc"
_ "github.com/yaoapp/yao/grpc/auth"
"github.com/yaoapp/yao/openapi"
ischedule "github.com/yaoapp/yao/schedule"
"github.com/yaoapp/yao/service"
@ -190,6 +192,12 @@ var startCmd = &cobra.Command{
}
}
// Print gRPC listen addresses
grpcAddrs := yaogrpc.Addr()
for _, addr := range grpcAddrs {
fmt.Println(color.WhiteString(L("Listening")), color.GreenString(" %s (gRPC)", addr))
}
fmt.Println(color.WhiteString("\n---------------------------------"))
fmt.Println(color.WhiteString(L("Access Points")))
fmt.Println(color.WhiteString("---------------------------------"))
@ -235,6 +243,13 @@ var startCmd = &cobra.Command{
os.Exit(1)
}
// Start gRPC Server (after HTTP, LIFO shutdown: gRPC stops before HTTP)
if grpcErr := yaogrpc.StartServer(config.Conf); grpcErr != nil {
fmt.Println(color.RedString(L("gRPC: %s"), grpcErr.Error()))
os.Exit(1)
}
defer yaogrpc.Stop()
// Start watching
watchDone := make(chan uint8, 1)
if mode == "development" && !startDisableWatching {

56
config/types.go
View file

@ -2,30 +2,38 @@ package config
// Config 象传应用引擎配置
type Config struct {
Mode string `json:"mode,omitempty" env:"YAO_ENV" envDefault:"production"` // The start mode production/development
AppSource string `json:"app,omitempty" env:"YAO_APP_SOURCE"` // The Application Source Root Path default same as Root
Root string `json:"root,omitempty" env:"YAO_ROOT" envDefault:"."` // The Application Root Path
Lang string `json:"lang,omitempty" env:"YAO_LANG" envDefault:"en-us"` // Default language setting
TimeZone string `json:"timezone,omitempty" env:"YAO_TIMEZONE"` // Default TimeZone
DataRoot string `json:"data_root,omitempty" env:"YAO_DATA_ROOT" envDefault:""` // The data root path
ExtensionRoot string `json:"extension_root,omitempty" env:"YAO_EXTENSION_ROOT" envDefault:""` // Plugin, Wasm root PATH, Default is <YAO_ROOT> (<YAO_ROOT>/plugins <YAO_ROOT>/wasms)
Host string `json:"host,omitempty" env:"YAO_HOST" envDefault:"0.0.0.0"` // The server host
Port int `json:"port,omitempty" env:"YAO_PORT" envDefault:"5099"` // The server port
Cert string `json:"cert,omitempty" env:"YAO_CERT"` // The HTTPS certificate path
Key string `json:"key,omitempty" env:"YAO_KEY"` // The HTTPS certificate key path
Log string `json:"log,omitempty" env:"YAO_LOG"` // The log file path
LogMode string `json:"log_mode,omitempty" env:"YAO_LOG_MODE" envDefault:"TEXT"` // The log mode TEXT|JSON
LogMaxSize int `json:"log_max_size,omitempty" env:"YAO_LOG_MAX_SIZE" envDefault:"100"` // The max log size in MB, the default is 100
LogMaxAage int `json:"log_max_age,omitempty" env:"YAO_LOG_MAX_AGE" envDefault:"7"` // The max log age in day, the default is 7
LogMaxBackups int `json:"log_max_backups" env:"YAO_LOG_MAX_BACKUPS" envDefault:"3"` // The max log backups, the default is 3
LogLocalTime bool `json:"log_local_time" env:"YAO_LOG_LOCAL_TIME" envDefault:"true"`
JWTSecret string `json:"jwt_secret,omitempty" env:"YAO_JWT_SECRET"` // The JWT Secret
DB Database `json:"db,omitempty"` // The database config
AllowFrom []string `json:"allowfrom,omitempty" envSeparator:"|" env:"YAO_ALLOW_FROM"` // Domain list the separator is |
Session Session `json:"session,omitempty"` // Session Config
Runtime Runtime `json:"runtime,omitempty"` // Runtime config
Trace Trace `json:"trace,omitempty"` // Trace config
Registry string `json:"registry,omitempty" env:"YAO_REGISTRY" envDefault:"https://registry.yaoagents.com"` // The package registry server URL
Mode string `json:"mode,omitempty" env:"YAO_ENV" envDefault:"production"` // The start mode production/development
AppSource string `json:"app,omitempty" env:"YAO_APP_SOURCE"` // The Application Source Root Path default same as Root
Root string `json:"root,omitempty" env:"YAO_ROOT" envDefault:"."` // The Application Root Path
Lang string `json:"lang,omitempty" env:"YAO_LANG" envDefault:"en-us"` // Default language setting
TimeZone string `json:"timezone,omitempty" env:"YAO_TIMEZONE"` // Default TimeZone
DataRoot string `json:"data_root,omitempty" env:"YAO_DATA_ROOT" envDefault:""` // The data root path
ExtensionRoot string `json:"extension_root,omitempty" env:"YAO_EXTENSION_ROOT" envDefault:""` // Plugin, Wasm root PATH, Default is <YAO_ROOT> (<YAO_ROOT>/plugins <YAO_ROOT>/wasms)
Host string `json:"host,omitempty" env:"YAO_HOST" envDefault:"0.0.0.0"` // The server host
Port int `json:"port,omitempty" env:"YAO_PORT" envDefault:"5099"` // The server port
Cert string `json:"cert,omitempty" env:"YAO_CERT"` // The HTTPS certificate path
Key string `json:"key,omitempty" env:"YAO_KEY"` // The HTTPS certificate key path
Log string `json:"log,omitempty" env:"YAO_LOG"` // The log file path
LogMode string `json:"log_mode,omitempty" env:"YAO_LOG_MODE" envDefault:"TEXT"` // The log mode TEXT|JSON
LogMaxSize int `json:"log_max_size,omitempty" env:"YAO_LOG_MAX_SIZE" envDefault:"100"` // The max log size in MB, the default is 100
LogMaxAage int `json:"log_max_age,omitempty" env:"YAO_LOG_MAX_AGE" envDefault:"7"` // The max log age in day, the default is 7
LogMaxBackups int `json:"log_max_backups" env:"YAO_LOG_MAX_BACKUPS" envDefault:"3"` // The max log backups, the default is 3
LogLocalTime bool `json:"log_local_time" env:"YAO_LOG_LOCAL_TIME" envDefault:"true"`
JWTSecret string `json:"jwt_secret,omitempty" env:"YAO_JWT_SECRET"` // The JWT Secret
DB Database `json:"db,omitempty"` // The database config
AllowFrom []string `json:"allowfrom,omitempty" envSeparator:"|" env:"YAO_ALLOW_FROM"` // Domain list the separator is |
Session Session `json:"session,omitempty"` // Session Config
Runtime Runtime `json:"runtime,omitempty"` // Runtime config
Trace Trace `json:"trace,omitempty"` // Trace config
Registry string `json:"registry,omitempty" env:"YAO_REGISTRY" envDefault:"https://registry.yaoagents.com"` // The package registry server URL
GRPC GRPCConfig `json:"grpc,omitempty"`
}
// GRPCConfig gRPC server configuration
type GRPCConfig struct {
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
Port int `json:"port,omitempty" env:"YAO_GRPC_PORT" envDefault:"9099"` // Listen port shared by all addresses
}
// Database 数据库配置

88
engine/machine.go Normal file
View file

@ -0,0 +1,88 @@
package engine
import (
"crypto/sha256"
"fmt"
"net"
"os"
"runtime"
"strings"
"sync"
"github.com/yaoapp/gou/process"
)
// MachineInfo contains deterministic machine identification.
type MachineInfo struct {
ID string `json:"id"` // "yao-cli-{hash32}" deterministic client ID
Hostname string `json:"hostname"` // OS hostname
Platform string `json:"platform"` // runtime.GOOS: "darwin", "linux", "windows"
}
var (
cachedMachineInfo *MachineInfo
machineOnce sync.Once
machineErr error
)
func init() {
process.Register("utils.app.MachineID", processMachineID)
}
// GetMachineID returns a deterministic machine fingerprint.
// The result is cached after the first call.
func GetMachineID() (*MachineInfo, error) {
machineOnce.Do(func() {
cachedMachineInfo, machineErr = computeMachineID()
})
return cachedMachineInfo, machineErr
}
func computeMachineID() (*MachineInfo, error) {
hostname, _ := os.Hostname()
raw, err := platformMachineID()
if err != nil || strings.TrimSpace(raw) == "" {
raw = fallbackMachineID(hostname)
}
hash := sha256.Sum256([]byte(raw))
id := fmt.Sprintf("yao-cli-%x", hash[:16]) // 32 hex chars
return &MachineInfo{
ID: id,
Hostname: hostname,
Platform: runtime.GOOS,
}, nil
}
func fallbackMachineID(hostname string) string {
mac := firstHardwareAddr()
return hostname + ":" + mac
}
func firstHardwareAddr() string {
ifaces, err := net.Interfaces()
if err != nil {
return "unknown"
}
for _, iface := range ifaces {
if iface.Flags&net.FlagLoopback != 0 || len(iface.HardwareAddr) == 0 {
continue
}
return iface.HardwareAddr.String()
}
return "unknown"
}
func processMachineID(p *process.Process) interface{} {
info, err := GetMachineID()
if err != nil {
return map[string]interface{}{"error": err.Error()}
}
return map[string]interface{}{
"id": info.ID,
"hostname": info.Hostname,
"platform": info.Platform,
}
}

24
engine/machine_darwin.go Normal file
View file

@ -0,0 +1,24 @@
//go:build darwin
package engine
import (
"os/exec"
"strings"
)
func platformMachineID() (string, error) {
out, err := exec.Command("ioreg", "-rd1", "-c", "IOPlatformExpertDevice").Output()
if err != nil {
return "", err
}
for _, line := range strings.Split(string(out), "\n") {
if strings.Contains(line, "IOPlatformUUID") {
parts := strings.SplitN(line, `"`, 4)
if len(parts) >= 4 {
return strings.TrimSpace(parts[3]), nil
}
}
}
return "", nil
}

21
engine/machine_linux.go Normal file
View file

@ -0,0 +1,21 @@
//go:build linux
package engine
import (
"os"
"strings"
)
func platformMachineID() (string, error) {
for _, path := range []string{"/etc/machine-id", "/var/lib/dbus/machine-id"} {
data, err := os.ReadFile(path)
if err == nil {
id := strings.TrimSpace(string(data))
if id != "" {
return id, nil
}
}
}
return "", nil
}

57
engine/machine_test.go Normal file
View file

@ -0,0 +1,57 @@
package engine
import (
"strings"
"testing"
)
func TestGetMachineID_Deterministic(t *testing.T) {
info1, err := GetMachineID()
if err != nil {
t.Fatalf("GetMachineID() returned error: %v", err)
}
info2, err := GetMachineID()
if err != nil {
t.Fatalf("GetMachineID() second call returned error: %v", err)
}
if info1.ID != info2.ID {
t.Errorf("GetMachineID() not deterministic: %q != %q", info1.ID, info2.ID)
}
}
func TestGetMachineID_Format(t *testing.T) {
info, err := GetMachineID()
if err != nil {
t.Fatalf("GetMachineID() returned error: %v", err)
}
if !strings.HasPrefix(info.ID, "yao-cli-") {
t.Errorf("ID should have prefix 'yao-cli-', got %q", info.ID)
}
// "yao-cli-" (8) + 32 hex chars = 40
if len(info.ID) != 40 {
t.Errorf("ID should be 40 chars, got %d: %q", len(info.ID), info.ID)
}
if info.Hostname == "" {
t.Error("Hostname should not be empty")
}
if info.Platform == "" {
t.Error("Platform should not be empty")
}
}
func TestGetMachineID_NonEmpty(t *testing.T) {
info, err := GetMachineID()
if err != nil {
t.Fatalf("GetMachineID() returned error: %v", err)
}
if info.ID == "" {
t.Error("ID should not be empty")
}
}

21
engine/machine_windows.go Normal file
View file

@ -0,0 +1,21 @@
//go:build windows
package engine
import (
"golang.org/x/sys/windows/registry"
)
func platformMachineID() (string, error) {
k, err := registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\Microsoft\Cryptography`, registry.READ|registry.WOW64_64KEY)
if err != nil {
return "", err
}
defer k.Close()
val, _, err := k.GetStringValue("MachineGuid")
if err != nil {
return "", err
}
return val, nil
}

2
go.mod
View file

@ -51,6 +51,7 @@ require (
go.mongodb.org/mongo-driver v1.17.3
golang.org/x/crypto v0.48.0
golang.org/x/net v0.50.0
golang.org/x/sys v0.41.0
golang.org/x/text v0.34.0
google.golang.org/grpc v1.78.0
google.golang.org/protobuf v1.36.11
@ -235,7 +236,6 @@ require (
golang.org/x/mod v0.33.0 // indirect
golang.org/x/oauth2 v0.32.0 // indirect
golang.org/x/sync v0.19.0 // indirect
golang.org/x/sys v0.41.0 // indirect
golang.org/x/term v0.40.0 // indirect
golang.org/x/time v0.14.0 // indirect
golang.org/x/tools v0.42.0 // indirect

448
grpc/DESIGN.md Normal file
View file

@ -0,0 +1,448 @@
# Yao gRPC Server
General-purpose gRPC gateway for the Yao process. Shares OAuth + ACL scope system with openapi — one token, two protocols.
## Services
| Layer | Method | Purpose | Scope |
|-------|--------|---------|-------|
| **Base** | `Run` | Execute Yao process, return result | `grpc:run` |
| | `Stream` | Execute Yao process, stream output | `grpc:stream` |
| | `Shell` | Execute system command, wait for result | `grpc:shell` |
| | `ShellStream` | Execute system command, stream stdout/stderr | `grpc:shell` |
| **API** | `API` | Proxy to openapi, any endpoint | openapi's own scopes |
| **MCP** | `MCPListTools` | List MCP tools for a session | `grpc:mcp` |
| | `MCPCallTool` | Call MCP tool → process.Exec() | `grpc:mcp` |
| | `MCPListResources` | List MCP resources | `grpc:mcp` |
| | `MCPReadResource` | Read MCP resource | `grpc:mcp` |
| **LLM** | `ChatCompletions` | Send messages to LLM, get response | `grpc:llm` |
| | `ChatCompletionsStream` | Stream LLM response (SSE → gRPC stream) | `grpc:llm` |
| **Agent** | `AgentStream` | Call agent, stream response | `grpc:agent` |
## Clients
- Container MCP tools (via Tai gRPC relay)
- `yao run` CLI (after `yao login`)
- Yao-to-Yao (cross-node process execution)
## Auth
Same as openapi. gRPC auth interceptor reuses the same `guard.Authenticate` logic — including automatic token refresh when access token is expired but refresh token is valid.
```
metadata (Bearer + x-refresh-token)
→ VerifyToken
→ expired? → TryRefresh (same as guard.go) → new tokens in response metadata
→ extract scopes → acl.Scope.Check(method, path, scopes)
```
### Infrastructure reuse assessment
Existing openapi/oauth infrastructure can be reused for gRPC with **zero modifications**:
| Component | Reusable as-is | Notes |
|-----------|---------------|-------|
| `VerifyToken(token string)` | Yes | Pure string input, no Gin dependency |
| `MakeAccessToken(clientID, scope, subject, expiresIn, extraClaims...)` | Yes | Supports custom scope/subject for container tokens |
| `MakeRefreshToken(...)` | Yes | Same as above |
| `Revoke(ctx, token, tokenTypeHint)` | Yes | For container token cleanup on Remove |
| `ScopeManager.Check(req *AccessRequest)` | Yes | Only needs `(Method, Path, Scopes)` — no Gin dependency |
| `acl.Register(...)` | Yes | gRPC scopes registered via same pattern |
The `authorized.SetInfo` / `authorized.GetInfo` are Gin-bound but **not needed** — gRPC interceptor builds `AccessRequest` directly from JWT claims. Full `Enforce` chain (client/team/member) is HTTP multi-tenant only; gRPC uses `VerifyToken → ScopeManager.Check` which is sufficient.
New code required: ~80 lines (interceptor + scope registration). Existing code changes: **zero**.
### CLI auth: `yao login` / `yao logout`
OAuth 2.0 Device Authorization Grant. No `--remote` flag needed — logged in = gRPC, not logged in = local.
```
$ yao login --server https://yao.example.com
请访问: https://yao.example.com/device
输入代码: ABCD-1234
等待授权... ✓ (token saved to ~/.yao/credentials)
$ yao run models.user.Find '{"id":1}' ← auto gRPC
$ yao logout
```
Requires two new openapi endpoints:
- `POST /oauth/device/authorize` — issue device_code + user_code
- `POST /oauth/device/token` — poll for access_token
Token scope: based on user's role, e.g. `grpc:run grpc:stream grpc:shell grpc:llm grpc:agent grpc:mcp`.
**Implementation cost**: ~190 lines new code, ~10 lines changes to existing code.
Scaffolding already in place — `types.DeviceAuthorizationResponse`, `GrantTypeDeviceCode`, error codes (`ErrorAuthorizationPending`, `ErrorSlowDown`), `DeviceCodeLifetime` config, `DeviceAuthorization()` method signature, and HTTP route are all pre-defined. Core work:
1. Implement `DeviceAuthorization()` in `device.go` (currently returns `nil, nil`)
2. Add device_code store/get/consume helpers in `token.go`
3. Add `GrantTypeDeviceCode` case to `Token()` switch in `core.go` (1 case branch)
4. Implement `handleDeviceCodeGrant()` in `core.go`
5. Add user authorization callback handler
6. Fix discovery endpoint path inconsistency (`/oauth/device` vs `/oauth/device_authorization`)
Risk: **very low** — all additions are in isolated code paths, no changes to existing `authorization_code` / `client_credentials` / `refresh_token` flows.
### Container token
Container images and `yao-grpc` (`yao/tai/grpc/`) are ours — it handles token refresh automatically.
```
Manager creates container
├─ oauth.MakeAccessToken(subject=userID, scope="grpc:mcp grpc:run")
├─ oauth.MakeRefreshToken(...)
└─ tai.Client.Sandbox().Create(CreateRequest{
Env: {
YAO_TOKEN, YAO_REFRESH_TOKEN, YAO_SANDBOX_ID,
YAO_GRPC_ADDR, // where to connect
YAO_GRPC_UPSTREAM, // remote only: where Tai should forward to
},
})
Local: YAO_GRPC_ADDR=127.0.0.1:9099 (direct to Yao, no upstream needed)
Remote: YAO_GRPC_ADDR=tai-host:9100 YAO_GRPC_UPSTREAM=yao-host:9099
yao-grpc (tai/grpc/, container-internal)
├─ reads YAO_GRPC_ADDR + YAO_TOKEN + YAO_REFRESH_TOKEN + YAO_SANDBOX_ID from env
├─ if YAO_GRPC_UPSTREAM set: attaches x-grpc-upstream metadata (tells Tai where to forward)
├─ every call: Bearer token + x-refresh-token + x-sandbox-id in gRPC metadata
├─ server auth interceptor reuses guard.Authenticate logic:
│ token valid → pass through
│ token expired + refresh token present → auto rotate (same as HTTP guard)
│ new tokens returned via response metadata (x-access-token, x-refresh-token)
├─ yao-grpc reads response metadata, updates tokens in memory
└─ transparent to caller, no separate refresh RPC needed
```
- access_token: short TTL (15m)
- refresh_token: no expiry (valid until container removed)
- Manager revokes refresh_token on container Remove
- Tai does NOT know Yao address at startup — yao-grpc carries target in request metadata
### Virtual endpoint mapping
| gRPC | Virtual endpoint |
|------|-----------------|
| Run("models.user.Find") | `POST /grpc/run/models.user.Find` |
| Stream("flows.report") | `POST /grpc/stream/flows.report` |
| Shell | `POST /grpc/shell` |
| ShellStream | `POST /grpc/shell` (same) |
| API(POST, /kb/collections) | `POST /kb/collections` (real openapi path) |
| MCPListTools | `GET /grpc/mcp/tools` |
| MCPCallTool("search") | `POST /grpc/mcp/call/search` |
| MCPListResources | `GET /grpc/mcp/resources` |
| MCPReadResource("uri") | `GET /grpc/mcp/resources/read` |
| ChatCompletions | `POST /grpc/llm/completions` |
| ChatCompletionsStream | `POST /grpc/llm/completions` (same) |
| AgentStream("robot-id") | `POST /grpc/agent/robot-id` |
API method uses the **actual openapi path** — no virtual mapping needed, scope check is identical to HTTP.
### Scope registration
```go
func init() {
acl.Register(
&acl.ScopeDefinition{Name: "grpc:run", Endpoints: []string{"POST /grpc/run/*"}},
&acl.ScopeDefinition{Name: "grpc:stream", Endpoints: []string{"POST /grpc/stream/*"}},
&acl.ScopeDefinition{Name: "grpc:shell", Endpoints: []string{"POST /grpc/shell"}},
&acl.ScopeDefinition{Name: "grpc:mcp", Endpoints: []string{"GET /grpc/mcp/tools", "POST /grpc/mcp/call/*", "GET /grpc/mcp/resources", "GET /grpc/mcp/resources/read"}},
&acl.ScopeDefinition{Name: "grpc:llm", Endpoints: []string{"POST /grpc/llm/completions"}},
&acl.ScopeDefinition{Name: "grpc:agent", Endpoints: []string{"POST /grpc/agent/*"}},
)
}
```
## Network
### Server listen config
| Env | Default | Purpose |
|-----|---------|---------|
| `YAO_GRPC_HOST` | `127.0.0.1` | Comma-separated bind addresses. |
| `YAO_GRPC_PORT` | `9099` | Listen port (shared by all addresses). |
| `YAO_GRPC` | _(unset)_ | Set `off` to explicitly disable gRPC server. |
gRPC server **defaults to enabled** (`127.0.0.1:9099`) — sandbox container callbacks depend on it.
`YAO_GRPC_HOST` accepts one or more addresses separated by `,`. Each address gets its own `net.Listener`; all listeners feed into the same `grpc.Server` (gRPC supports multiple `Serve` calls on one server).
| Scenario | Config | Effect |
|----------|--------|--------|
| Local dev / default | _(nothing to set)_ | `127.0.0.1:9099` — loopback, sandbox works out of box |
| LAN multi-NIC | `YAO_GRPC_HOST=192.168.10.1,10.0.0.1` | Binds each internal IP |
| Open | `YAO_GRPC_HOST=0.0.0.0` | All interfaces |
| Disabled | `YAO_GRPC=off` | gRPC server not started (pure API gateway, no sandbox) |
When multiple addresses are given, the server creates one goroutine per listener. Shutdown (`grpc.GracefulStop`) drains all listeners.
Config lives in `config.Config.GRPC` (type `GRPCConfig`), same pattern as `Host`/`Port` for HTTP.
### Startup
gRPC server starts **after** HTTP server in `cmd/start.go`, as a parallel goroutine:
```
engine.Load → itask.Start → ischedule.Start → service.Start (HTTP) → grpc.StartServer (gRPC)
```
gRPC server starts by default. Set `YAO_GRPC=off` to explicitly disable (no-op startup). Any other value or unset means enabled.
Shutdown: `defer grpc.Stop()` in `cmd/start.go`, called before HTTP stop for graceful drain.
### Access control
Local: containers and CLI connect via loopback. Remote: only Tai relay connects (address known from `YAO_TAI_ADDR`). All callers carry OAuth tokens — no IP allowlist needed.
Interceptor chain: auth → ACL → handler.
Public methods (skip auth): `Healthz`. Auth interceptor checks method name and passes through.
## IPC Path (replacing Unix socket)
All modes use gRPC — no Unix socket fallback. One code path, local and remote.
```
Local: Container → yao-grpc → Yao gRPC 127.0.0.1:9099
Remote: Container → yao-grpc → Tai :9100 relay → Yao gRPC :9099
```
`yao-grpc` reads `YAO_GRPC_ADDR` from env and connects. Local containers point directly at the Yao gRPC server on loopback; remote containers point at the Tai relay. No mode switch, no branching.
### Tai relay routing
Tai does **not** know the Yao gRPC address at startup. yao-grpc tells Tai where to forward on every request via metadata:
```
Manager.Create(sandbox)
├─ oauth.MakeAccessToken(...)
├─ oauth.MakeRefreshToken(...)
└─ tai.Client.Sandbox().Create(CreateRequest{
Env: {
YAO_TOKEN, YAO_REFRESH_TOKEN,
YAO_GRPC_ADDR: "tai-host:9100",
YAO_GRPC_UPSTREAM: "yao-host:9099",
},
})
```
yao-grpc reads `YAO_GRPC_UPSTREAM` from env and attaches it as `x-grpc-upstream` metadata on every request to Tai. Tai gateway reads this metadata and forwards to the specified address. No per-container state in Tai, no lookup table — pure transparent proxy. One Tai can serve containers from different Yao instances because each request carries its own target.
For local mode, no Tai relay — Manager injects `YAO_GRPC_ADDR=127.0.0.1:9099` directly (no `YAO_GRPC_UPSTREAM` needed).
### yao-grpc (container client)
`yao-grpc` is the in-container gRPC client binary. Replaces the old `yao-bridge`. Lives in `yao/tai/grpc/`:
```
yao/tai/grpc/
├── grpc.go // gRPC client: connect, forward MCP/process calls
├── auth.go // token management: read env, auto-refresh
├── grpc_test.go
└── cmd/
└── main.go
```
Rationale for placing in `yao/tai`:
- Consumes Tai relay — same layer as `tai/proxy`, `tai/volume`
- Shares gRPC deps already in `yao/tai`
- Version-locked with Tai SDK and server protocol
- Built in same CI: `go build -o yao-grpc ./tai/grpc/cmd`
Pure client — no signing keys, no `oauth` package dependency. Reads `YAO_TOKEN` + `YAO_REFRESH_TOKEN` + `YAO_SANDBOX_ID` from env, attaches all three as gRPC metadata on every call. Token refresh is transparent — server auto-rotates expired tokens (same logic as HTTP guard) and returns new tokens via response metadata.
## Proto
```protobuf
service Yao {
// Base
rpc Run(RunRequest) returns (RunResponse);
rpc Stream(RunRequest) returns (stream Chunk);
rpc Shell(ShellRequest) returns (ShellResponse);
rpc ShellStream(ShellRequest) returns (stream Chunk);
// API gateway
rpc API(APIRequest) returns (APIResponse);
// MCP
rpc MCPListTools(MCPListRequest) returns (MCPListResponse);
rpc MCPCallTool(MCPCallRequest) returns (MCPCallResponse);
rpc MCPListResources(MCPListRequest) returns (MCPResourcesResponse);
rpc MCPReadResource(MCPResourceRequest) returns (MCPResourceResponse);
// AI - LLM
rpc ChatCompletions(ChatRequest) returns (ChatResponse);
rpc ChatCompletionsStream(ChatRequest) returns (stream ChatChunk);
// AI - Agent
rpc AgentStream(AgentRequest) returns (stream AgentChunk);
// Health
rpc Healthz(Empty) returns (HealthzResponse);
}
```
### LLM layer
`ChatCompletions` and `ChatCompletionsStream` call the existing `llm.ChatCompletions` process (`agent/llm/process.go`). It auto-detects connector type (openai/anthropic/etc.), selects the appropriate provider, and returns OpenAI-compatible format.
```
gRPC ChatCompletions(connector, messages, opts)
→ process.Exec("llm.ChatCompletions", connector, messages, opts)
→ agent/llm.New(conn, opts) → provider.Stream/Post → response
gRPC ChatCompletionsStream(connector, messages, opts)
→ same path, with streaming callback → gRPC stream chunks
```
The caller specifies a connector ID. The `llm.ChatCompletions` process resolves it via `connector.Select()`, creates the LLM instance, and executes. Streaming version passes a callback that forwards chunks to the gRPC stream.
### Agent layer
`AgentStream` wraps `agent/robots/:id/completions` — resolves robot → host assistant → runs agent pipeline → streams output. Only stream method — agent output is inherently streamed; non-stream callers simply consume all chunks. Internally calls `assistant.Stream()` with `ctx.Writer` set to nil (or noop) when the caller doesn't need incremental output.
```
gRPC AgentStream(agent_id, messages) → resolve robot → assistant.Stream() → stream chunks
```
This enables container-internal agents to call other agents without HTTP, and remote `yao` instances to orchestrate agent pipelines cross-node.
`AgentChunk` carries `agent/output/message.Message` — the same DSL used by HTTP SSE streaming. Each chunk is one JSON-serialized `Message`:
```protobuf
message AgentChunk {
bytes data = 1; // JSON-encoded agent/output/message.Message
bool done = 2;
}
```
The `Message` structure uses `Type` + `Props` to express all content types (text, thinking, tool_call, error, action, event, image, audio, video). Streaming control fields (`chunk_id`, `message_id`, `block_id`, `thread_id`) and delta fields (`delta`, `delta_path`, `delta_action`) are preserved as-is over gRPC — the client merges chunks using the same logic as CUI's SSE consumer.
### Shell execution context
`Shell` and `ShellStream` execute commands in the **Yao host process**, not inside a sandbox container. This is by design — the scope `grpc:shell` is a privileged capability, not granted to container tokens by default. Container-internal commands run via `tai.Client.Sandbox().Exec()`, which is a different path (not exposed as a gRPC method).
See [pb/yao.proto](./pb/yao.proto) for full message definitions.
## Process & Stream (gou foundation)
gRPC `Run` and `Stream` map to two parallel systems in `gou`:
```
gou/process/ — execute once, return result → gRPC Run
gou/stream/ — execute once, push chunks → gRPC Stream
```
### gou/process (existing, unchanged)
```go
type Handler func(process *Process) interface{}
process.Register("scripts", handler)
p := process.New("scripts.foo.bar", args...)
p.Execute()
result := p.Value()
```
### gou/stream (new package, parallel to process)
```go
type Handler func(ctx context.Context, process *Process, send func([]byte) error) error
stream.Register("scripts", handler)
s := stream.New("scripts.foo.bar", args...)
s.Execute(ctx, func(chunk []byte) error { ... })
```
`stream.Process` mirrors `process.Process` fields (Name, Group, Method, ID, Args, Global, Sid, Authorized) but `ctx` is a first-class parameter, not buried in a struct field.
`send` returns error when the receiver disconnects — handler should stop.
### Fallback
If a stream handler is not registered for a name but a process handler exists, `stream.Execute` falls back to: run the process handler once, JSON-marshal the result, call `send` once.
### Registration
```go
// gou/process — existing
process.Register("models", modelsHandler)
process.Register("scripts", scriptsHandler)
// gou/stream — new, same namespace
stream.Register("scripts", scriptsStreamHandler)
stream.Register("llm", llmStreamHandler)
```
Same naming convention. A process name can have both a process handler and a stream handler.
### gRPC mapping
```go
func (s *yaoServer) Run(ctx context.Context, req *pb.RunRequest) (*pb.RunResponse, error) {
p := process.NewWithContext(ctx, req.Process, args...)
if err := p.Execute(); err != nil { return nil, err }
data, _ := json.Marshal(p.Value())
return &pb.RunResponse{Result: data}, nil
}
func (s *yaoServer) Stream(req *pb.RunRequest, grpcStream pb.Yao_StreamServer) error {
st := stream.New(req.Process, args...)
return st.Execute(grpcStream.Context(), func(chunk []byte) error {
return grpcStream.Send(&pb.Chunk{Data: chunk})
})
}
```
### V8 integration
Both are exposed as top-level globals in JavaScript, parallel:
```go
// gou/runtime/v8/isolate.go MakeTemplate
template.Set("Process", processModule.ExportFunction(iso)) // existing
template.Set("Stream", streamModule.ExportFunction(iso)) // new
```
**JS calling Go stream** (JS is consumer):
```javascript
Stream("llm.chat.completions", function(chunk) {
log.Info(chunk)
return 1 // 1=continue, 0=stop
}, { model: "gpt-4", messages: [...] })
```
**JS script as stream handler** (JS is producer):
```javascript
// scripts/report.js — registered via stream.Register("scripts", ...)
function generate(args, send) {
send("part 1")
send("part 2")
}
```
V8 runtime registers both:
```go
func init() {
process.Register("scripts", processScripts) // existing
stream.Register("scripts", processScriptsStream) // new
}
```
`processScriptsStream` calls `script.ExecStream(ctx, p, send)` which injects `send` into the V8 global before executing the script method.
### Impact on existing code
| Component | Changes |
|-----------|---------|
| `gou/process/` | None |
| `gou/stream/` | New package (~150 lines) |
| `gou/runtime/v8/process.go` | +1 line: `stream.Register(...)` |
| `gou/runtime/v8/script.go` | +`ExecStream` method |
| `gou/runtime/v8/isolate.go` | +1 line: `template.Set("Stream", ...)` |
| `gou/runtime/v8/functions/` | +`stream/` module for JS→Go stream consumption |

340
grpc/IMPL.md Normal file
View file

@ -0,0 +1,340 @@
# Yao gRPC Server — Implementation Plan
Design: [DESIGN.md](./DESIGN.md)
## Scope
**V1**: Auth + unary RPCs + LLM/Agent streaming + container client.
**V2**: Base streaming (`Stream`, `ShellStream`) + `gou/stream` package + V8 integration.
## Package Structure
```
grpc/
├── grpc.go // StartServer, config, server lifecycle
├── pb/
│ ├── yao.proto
│ ├── yao.pb.go // generated
│ └── yao_grpc.pb.go // generated
├── auth/
│ ├── guard.go // unary + stream interceptor (calls oauth.VerifyToken, ScopeManager.Check)
│ ├── endpoint.go // gRPC method → virtual HTTP endpoint mapping
│ └── scope.go // init() acl.Register for grpc:* scopes
├── run/
│ └── run.go // Run handler
├── shell/
│ └── shell.go // Shell, ShellStream (V2) handlers
├── api/
│ └── api.go // API proxy handler
├── mcp/
│ └── mcp.go // MCPListTools, MCPCallTool, MCPListResources, MCPReadResource
├── llm/
│ └── llm.go // ChatCompletions, ChatCompletionsStream
├── agent/
│ └── agent.go // AgentStream
└── health/
└── health.go // Healthz
```
Container client:
```
tai/grpc/
├── grpc.go // gRPC client, Dial, method wrappers
├── auth.go // read env tokens, attach metadata, handle refresh
├── grpc_test.go
└── cmd/
└── main.go // yao-grpc binary entry
```
## V1 Phases
### Phase 0: Proto + codegen ✅
No dependency.
| Task | Detail | Status |
|------|--------|--------|
| `grpc/pb/yao.proto` | All 14 RPCs + all message types. V2 methods (`Stream`, `ShellStream`) included in proto, handler left `Unimplemented`. | ✅ Done |
| codegen | `protoc` → `pb/*.pb.go` + `pb/*_grpc.pb.go` | ✅ Done |
### Phase 1: Auth + server skeleton ✅
Depends on: Phase 0. Auth is ~80 lines new code calling existing `openapi/oauth` functions.
| Task | Detail | Status |
|------|--------|--------|
| `grpc/auth/scope.go` | `init()` — `acl.Register` 6 gRPC scope definitions | ✅ Done |
| `grpc/auth/endpoint.go` | Map gRPC method + request params → virtual HTTP endpoint for ACL (e.g. `Run("models.user.Find")` → `POST /grpc/run/models.user.Find`) | ✅ Done |
| `grpc/auth/guard.go` | Extract Bearer from metadata → `oauth.AuthenticateToken` (pure, no gin) → ACL scope check. Skip `Healthz`. New tokens via `SendHeader`. | ✅ Done |
| `openapi/oauth/authenticate.go` | `AuthenticateToken(AuthInput) → AuthResult` — gin-free auth core. `refreshTokenDirect`, `buildAuthInfo`. Shares `refreshGates` with `TryRefreshToken`. | ✅ Done |
| `grpc/grpc.go` | `StartServer(cfg)` — `grpc.NewServer` with interceptor, register service, listen. See **Server config & startup** below. | ✅ Done |
| `grpc/health/health.go` | `Healthz` → `{status: "ok"}` | ✅ Done |
| `config/types.go` | Add `GRPC` field to `Config` struct — see config below | ✅ Done |
| `cmd/start.go` | After `service.Start(config.Conf)` (HTTP ready), call `grpc.StartServer(config.Conf)` in goroutine. Print gRPC listen address in Access Points block. `defer grpc.Stop()` in shutdown path. | ✅ Done |
**Server config & startup:**
Config struct addition (`config/types.go`):
```go
type Config struct {
// ... existing fields ...
GRPC GRPCConfig `json:"grpc,omitempty"`
}
type GRPCConfig struct {
Enabled string `json:"enabled,omitempty" env:"YAO_GRPC"`
Host string `json:"host,omitempty" env:"YAO_GRPC_HOST" envDefault:"127.0.0.1"`
Port int `json:"port,omitempty" env:"YAO_GRPC_PORT" envDefault:"9099"`
}
```
- **Default** — `127.0.0.1:9099`, enabled. Sandbox callbacks work out of box.
- `YAO_GRPC_HOST=192.168.10.1,10.0.0.1` — comma-separated, binds each IP for multi-NIC LAN
- `YAO_GRPC_HOST=0.0.0.0` — all interfaces
- `YAO_GRPC=off` — explicitly disable gRPC server
`grpc.StartServer` implementation:
```go
func StartServer(cfg config.Config) error {
if strings.ToLower(cfg.GRPC.Enabled) == "off" {
log.Info("gRPC server disabled (YAO_GRPC=off)")
return nil
}
hosts := strings.Split(cfg.GRPC.Host, ",")
for _, host := range hosts {
addr := net.JoinHostPort(strings.TrimSpace(host), strconv.Itoa(cfg.GRPC.Port))
lis, err := net.Listen("tcp", addr)
// ... error handling ...
go server.Serve(lis) // one goroutine per listener, same grpc.Server
}
return nil
}
```
Startup sequence in `cmd/start.go`:
```
engine.Load(cfg)
itask.Start()
ischedule.Start()
service.Start(cfg) // HTTP server
grpc.StartServer(cfg) // gRPC server (after HTTP, parallel goroutine)
// ... event loop ...
defer grpc.Stop() // GracefulStop drains all listeners (no-op if not started)
```
`cmd/start.go` prints each gRPC listen address:
```
Listening 0.0.0.0:5099 (HTTP)
Listening 192.168.10.1:9099 (gRPC)
Listening 10.0.0.1:9099 (gRPC)
```
Deliverable: Server starts, Healthz works, unauthenticated calls rejected, token refresh via metadata works.
### Phase 2: Base + API + MCP handlers ✅
Depends on: Phase 1.
| Task | Detail | Status |
|------|--------|--------|
| `grpc/run/run.go` | `Run` — `process.New(req.Process, args...).Exec()`. Injects `AuthorizedInfo` via `p.WithSID()` + `p.WithAuthorized()`. | ✅ Done |
| `grpc/shell/shell.go` | `Shell` — `exec.CommandContext` in host process. **Security**: refuse execution if Yao process is running as root (`os.Getuid() == 0` → `PermissionDenied`). Timeout: use request `timeout` field, default 30s, capped by server max. | ✅ Done |
| `grpc/api/api.go` | `API` — build `http.Request`, call openapi internally | ✅ Done |
| `grpc/mcp/mcp.go` | `MCPListTools`, `MCPCallTool`, `MCPListResources`, `MCPReadResource` | ✅ Done |
Deliverable: Base + API + MCP methods work with valid tokens.
### Phase 3: LLM + Agent handlers ✅
Depends on: Phase 1. No code dependency on Phase 2 — can parallel.
| Task | Detail | Status |
|------|--------|--------|
| `grpc/llm/llm.go` | `ChatCompletions` / `ChatCompletionsStream` — direct call to `agent/llm` (`connector.Select` → `llm.New` → `Stream`). Uses `agent/llm.BuildCompletionOptions`. Constructs `agent/context.Context` with `AuthorizedInfo`. | ✅ Done |
| `grpc/agent/agent.go` | `AgentStream` — `assistant.Get` → `ast.Stream` with `grpcStreamWriter` adapter bridging `http.ResponseWriter` to gRPC `ServerStreamingServer`. Constructs `agent/context.Context` with `AuthorizedInfo`. | ✅ Done |
Deliverable: LLM (unary + stream) and Agent streaming via gRPC.
### Phase 4: Tai gateway change (Tai repo) ✅
Depends on: Phase 1 (need proto definitions for testing). yao-grpc depends on this.
Tai gateway currently dials a fixed `YaoUpstream` at startup. New behavior: yao-grpc tells Tai where to forward via request metadata (`x-grpc-upstream`). Tai reads the target address and proxies to it — removes `YaoUpstream` startup config.
| Task | Detail | Status |
|------|--------|--------|
| Tai `gateway/gateway.go` | Remove fixed `upstream *grpc.ClientConn`. On each request, read `x-grpc-upstream` from metadata → lookup/create conn from `sync.Map` cache (key = address string) → forward. Typical deployment has 1 upstream, cache stays tiny. | ✅ Done |
| Tai `server/server.go` | Remove `YaoUpstream` from `Config`. Gateway init no longer needs an address. | ✅ Done |
| Tai `main.go` | Remove `--yao` flag, `TAI_YAO_UPSTREAM` env var, YAML `yao` field, and required check. | ✅ Done |
| Tai `gateway/gateway_test.go` | Updated tests: dynamic routing, missing metadata → InvalidArgument, metadata forwarding (x-grpc-upstream stripped), upstream error propagation, multiple upstreams, connection cache. Coverage: 88.8%. | ✅ Done |
Connection cache: `sync.Map[string, *grpc.ClientConn]` — lazy dial on first request per upstream, reuse thereafter. No eviction needed (upstream count ≈ 1 in practice). `Close` closes all cached connections.
Deliverable: Tai starts without Yao address. Forwards based on request metadata.
### Phase 5: yao-grpc container client ✅
Depends on: Phase 1 (server + auth), Phase 4 (Tai gateway accepts `x-grpc-upstream`).
| Task | Detail | Status |
|------|--------|--------|
| `tai/grpc/auth.go` | `TokenManager`: read `YAO_TOKEN` / `YAO_REFRESH_TOKEN` / `YAO_SANDBOX_ID` / `YAO_GRPC_UPSTREAM` from env. `YAO_GRPC_TAI=enable` triggers Tai relay mode (requires `YAO_GRPC_UPSTREAM`). Attach as gRPC metadata on every call via unary + stream interceptors. Auto-refresh from response headers. | ✅ Done |
| `tai/grpc/grpc.go` | `Client`: `Dial(addr, TokenManager)`, `NewFromEnv()`. Method wrappers for all RPCs: Run, Shell, API, MCP (list/call/resources/read), ChatCompletions, ChatCompletionsStream, AgentStream, Healthz. | ✅ Done |
| `tai/grpc/cmd/main.go` | Stdio MCP server: JSON-RPC → gRPC. `yao-grpc version` prints version/commit/build time (via `-ldflags`). `yao-grpc serve` reads stdin JSON-RPC, dispatches to gRPC client. | ✅ Done |
| `tai/grpc/grpc_test.go` + `integration_test.go` | Black-box tests (package `grpc_test`). Unit: TokenManager metadata attachment, env parsing, refresh handling. Integration: real Yao gRPC server, all method wrappers, token refresh, auth rejection. Coverage: 83.9%. | ✅ Done |
Container token issuance uses existing `oauth.MakeAccessToken` / `oauth.MakeRefreshToken` — called by sandbox Manager at container creation, injected as env vars. Revoke on Remove. No new auth code needed on the issuance side.
Deliverable: `go build -o yao-grpc ./tai/grpc/cmd`.
### Phase 6: Device Flow + CLI auth ✅
Depends on: Phase 1. Three sub-phases with sequential dependency: 6.1 → 6.2 → 6.3.
#### Phase 6.1: OAuth Device Flow backend ✅
Backend endpoints for RFC 8628 Device Authorization Grant. Scaffolding already in place (`types.DeviceAuthorizationResponse`, `GrantTypeDeviceCode`, error codes, route registration).
| Task | Detail | Status |
|------|--------|--------|
| `engine/machine.go` + platform files | `GetMachineID()` Go API + `utils.app.MachineID` process — cross-platform (macOS/Linux/Windows) deterministic machine fingerprint | ✅ Done |
| `oauth/token.go` | `deviceCodeKey`, `userCodeKey`, `storeDeviceCode`, `getDeviceCodeData`, `authorizeDeviceCode`, `consumeDeviceCode` — device_code + user_code storage/retrieval/consumption helpers | ✅ Done |
| `oauth/device.go` | Implement `DeviceAuthorization()` + `AuthorizeDevice()` + `generateUserCode()` — generate codes (gonanoid, XXXX-XXXX format), validate client + grant type, store, return `DeviceAuthorizationResponse` | ✅ Done |
| `oauth/core.go` | Add `case types.GrantTypeDeviceCode` → `handleDeviceCodeGrant()` — poll returns `authorization_pending` / `expired_token` / token | ✅ Done |
| `openapi/oauth.go` | Replace stub `oauthDeviceAuthorization` handler → call `DeviceAuthorization()`. Add `POST /oauth/device/authorize` → `oauthDeviceAuthorize` (bearer token + user_code → authorize device). | ✅ Done |
| `oauth/discovery.go` | Fix path: `/oauth/device` → `/oauth/device_authorization` | ✅ Done |
| `oauth/oauth.go` | Config defaults: `DeviceCodeLength=8`, `UserCodeLength=8`, `DeviceCodeInterval=5s`, `DeviceFlowEnabled=true`, `DynamicClientRegistrationEnabled=true` | ✅ Done |
| `openapi/tests/oauth/device_test.go` | Full test suite: device auth success/error, token polling (pending/invalid), end-to-end flow | ✅ Done |
Deliverable: Device flow endpoints functional — `POST /oauth/device_authorization` issues codes, `POST /oauth/token` with `grant_type=device_code` polls status. `POST /oauth/device/authorize` allows authenticated user to authorize device.
#### Phase 6.2: CUI auth/device page (frontend) ✅
Depends on: Phase 6.1 (backend endpoints). Frontend-only task in **CUI repo**.
Route: `/auth/device` (Umi convention-based routing → `pages/auth/device/index.tsx`)
| Task | Detail | Status |
|------|--------|--------|
| `pages/auth/device/index.tsx` | Device authorization page. User enters `user_code`, clicks Authorize. Uses `AuthLayout` + `AuthInput` + `AuthButton` from existing `pages/auth/components/`. Three states: input, success, error. i18n (zh/en), light/dark, system CSS variables only. | ✅ Done |
| `pages/auth/device/index.less` | Styles, follow `pages/auth/entry/mfa/index.less` pattern. Full responsive + dark theme. | ✅ Done |
| `openapi/user/auth.ts` | `AuthorizeDevice(userCode)` method — `POST /oauth/device/authorize` | ✅ Done |
| `layouts/index.tsx` | Register `['auth_device', '/auth/device']` in `STANDALONE_PAGES` | ✅ Done |
Implementation:
- Framework: React + UmiJS Max + Ant Design + MobX (same as all auth pages)
- Layout: `AuthLayout` (logo + theme switch), same as `/auth/entry`
- Components: reuse `AuthInput` for `user_code` input, `AuthButton` for submit
- Page export: `export default observer(DeviceAuth)`
- API: `window.$app.openapi` → `POST /oauth/device/authorize` with `{ user_code }`, bearer token from session
- Auth: must be logged in (redirect to `/auth/entry` if not). After authorizing, show success and close/redirect
- i18n: `useIntl()`, `zh-CN` / `en-US`
Deliverable: `/auth/device` page. User authorizes CLI device login from browser.
#### Phase 6.3: CLI commands + TUI status bar ✅
Depends on: Phase 6.1 (backend) + Phase 6.2 (CUI page for end-to-end `yao login`).
**Credentials file** (`~/.yao/credentials`): base64-encoded JSON.
```json
{
"server": "https://yao.example.com",
"access_token": "eyJ...",
"refresh_token": "eyJ...",
"scope": "grpc:run grpc:stream grpc:shell grpc:llm grpc:agent grpc:mcp",
"user": "admin@example.com",
"expires_at": "2026-03-05T10:00:00Z"
}
```
Stored as: `base64(json) → ~/.yao/credentials`. Prevents casual `cat` exposure.
| Task | Detail | Status |
|------|--------|--------|
| `cmd/credential.go` | `Credential` struct, `LoadCredential`, `LoadCredentialFrom`, `SaveCredential`, `RemoveCredential` — base64-encoded JSON read/write to `~/.yao/credentials` | ✅ Done |
| `cmd/login.go` | `yao login --server <url>` — compute machine ID → `POST /oauth/register` (dynamic client) → `POST /oauth/device_authorization` → color-print device code + verification URL → poll `POST /oauth/token` with interval + slow_down handling → save to `~/.yao/credentials` | ✅ Done |
| `cmd/logout.go` | `yao logout` — read credentials, best-effort `POST /oauth/revoke`, delete `~/.yao/credentials` | ✅ Done |
| `cmd/run.go` | Detect credentials → gRPC mode vs local mode. `--auth <path>` flag loads alternate credentials file. `-s` (silent) mode: no TUI. gRPC mode renders TUI status bar then calls remote (gRPC call wiring pending Phase 4/5 integration). Local mode unchanged. | ✅ Done |
| `cmd/tui_status.go` | lipgloss `RenderStatusBar(cred)` — one-line persistent bar: `user (gRPC) │ scope: run,stream,...`. Rounded border, colored connection info. Hidden in silent mode. | ✅ Done |
| `cmd/root.go` | Register `loginCmd`, `logoutCmd` in root command | ✅ Done |
| i18n | All new strings have zh-CN translations via `langs` map | ✅ Done |
**`yao run` behavior matrix:**
| Credentials | `-s` flag | `--auth` flag | Behavior |
|-------------|-----------|---------------|----------|
| None | — | — | Local execution (current behavior) |
| `~/.yao/credentials` | No | — | gRPC + TUI status bar |
| `~/.yao/credentials` | Yes | — | gRPC, no TUI, pure output |
| — | Yes | `<path>` | gRPC via specified credentials, no TUI, pure output |
| — | No | `<path>` | gRPC via specified credentials + TUI status bar |
**TUI status bar** (bubbletea, `cmd/tui_status.go`):
```
┌─ admin@yao.example.com (gRPC) │ scope: run,stream,shell,llm,agent,mcp ─┐
```
- Top-line, persistent during execution
- lipgloss styled (dim border, colored connection info)
- Process output renders below, unaffected
- Hidden in silent mode (`-s`)
Deliverable: `yao login` + `yao logout` + `yao run` via gRPC with TUI status bar.
## V2 Phases
### Phase 7: `gou/stream` package ⏳
| Task | Detail | Status |
|------|--------|--------|
| `gou/stream/` | ~150 lines. `Handler`, `Process`, `Register`, `New`, `Execute`. Fallback to process. | ⏳ Pending |
| V8 | `stream.Register("scripts", ...)`, `ExecStream`, `template.Set("Stream", ...)`, JS `Stream()` global | ⏳ Pending |
### Phase 8: Base streaming handlers ⏳
Depends on: Phase 7.
| Task | Detail | Status |
|------|--------|--------|
| `grpc/run/run.go` | Add `Stream` handler — `stream.New(req.Process).Execute(ctx, send)` | ⏳ Pending |
| `grpc/shell/shell.go` | Add `ShellStream` handler — piped stdout → gRPC stream | ⏳ Pending |
## Dependency Graph
```
Phase 0 (proto) ✅
│
▼
Phase 1 (auth + server) ✅
│
├───────────┬───────────┬──────────────────────┐
▼ ▼ ▼ ▼
Phase 2 ✅ Phase 3 ✅ Phase 4 ✅ Phase 6 ✅ (device flow + CLI)
(handlers) (LLM/Agent) (Tai gateway) │
│ ┌───────┴───────┐
▼ ▼ ▼
Phase 5 ✅ 6.1 ✅ 6.2 ✅
(yao-grpc) (OAuth backend) (CUI page)
│ │
└───────┬───────┘
▼
6.3 ✅
(CMD + TUI)
--- V2 ---
Phase 7 (gou/stream)
│
▼
Phase 8 (Stream, ShellStream)
```

426
grpc/TEST.md Normal file
View file

@ -0,0 +1,426 @@
# Yao gRPC Server — Test Specification
Design: [DESIGN.md](./DESIGN.md) | Implementation: [IMPL.md](./IMPL.md)
## Principles
- **Black-box testing**: all `*_test.go` files use `package xxx_test` — tests only access exported API via gRPC client
- **Tests follow implementation**: `*_test.go` lives next to the code it tests (`grpc/auth/guard_test.go` beside `grpc/auth/guard.go`)
- **Real server**: every test starts a real gRPC server on a random TCP port, exercises the full interceptor → handler chain
- **Coverage > 80%**: per sub-package and overall
## Prerequisites
```bash
source $YAO_SOURCE_ROOT/env.local.sh
```
Required environment variables (same as existing Yao tests):
| Variable | Purpose |
|----------|---------|
| `YAO_TEST_APPLICATION` | Path to `yao-dev-app` |
| `YAO_DB_DRIVER` / `YAO_DB_PRIMARY` | Database connection |
| `YAO_JWT_SECRET` / `YAO_DB_AESKEY` | Crypto keys |
| `OPENAI_TEST_KEY` | LLM streaming tests |
| `ANTHROPIC_API_KEY` | LLM streaming tests (Anthropic) |
## Directory Structure
```
grpc/
├── grpc.go
├── tests/
│ └── testutils/
│ └── testutils.go # shared test utilities
├── auth/
│ ├── guard.go
│ ├── guard_test.go # package auth_test
│ ├── endpoint.go
│ ├── endpoint_test.go # package auth_test
│ └── scope.go
├── run/
│ ├── run.go
│ └── run_test.go # package run_test
├── shell/
│ ├── shell.go
│ └── shell_test.go # package shell_test
├── api/
│ ├── api.go
│ └── api_test.go # package api_test
├── mcp/
│ ├── mcp.go
│ └── mcp_test.go # package mcp_test
├── llm/
│ ├── llm.go
│ └── llm_test.go # package llm_test
├── agent/
│ ├── agent.go
│ └── agent_test.go # package agent_test
└── health/
├── health.go
└── health_test.go # package health_test
```
Tests live beside the code they verify. `grpc/tests/testutils/` is shared infrastructure only.
## testutils API
`grpc/tests/testutils/testutils.go` provides the test harness used by all sub-packages.
```go
package testutils
// Prepare initializes the full Yao runtime (DB, V8, models, scripts, etc.)
// then starts a real gRPC server on :0 (random port).
// Returns a connected grpc.ClientConn ready to create service clients.
//
// Internally calls:
// test.Prepare(t, config.Conf) — Yao runtime
// grpc.StartServer(cfg{Port:0}) — gRPC server
// grpc.Dial("127.0.0.1:port") — client connection
func Prepare(t *testing.T) *grpc.ClientConn
// Clean gracefully stops the gRPC server and tears down the Yao runtime.
// Always use with defer:
// conn := testutils.Prepare(t)
// defer testutils.Clean()
func Clean()
// Addr returns the gRPC server address "127.0.0.1:xxxxx".
func Addr() string
// ObtainAccessToken mints a token with the given scopes.
// Calls oauth.MakeAccessToken directly — no HTTP round-trip.
func ObtainAccessToken(t *testing.T, scopes ...string) string
// ObtainAccessTokenForUser mints a token for a specific user ID.
func ObtainAccessTokenForUser(t *testing.T, userID string, scopes ...string) string
// WithToken returns ctx with Bearer token in gRPC metadata.
func WithToken(ctx context.Context, token string) context.Context
// WithRefreshToken returns ctx with both Bearer and x-refresh-token metadata.
func WithRefreshToken(ctx context.Context, token, refreshToken string) context.Context
// WithSandboxMetadata returns ctx with x-sandbox-id and x-grpc-upstream metadata.
func WithSandboxMetadata(ctx context.Context, sandboxID, upstream string) context.Context
// NewClient creates a pb.YaoServiceClient from a connection.
func NewClient(conn *grpc.ClientConn) pb.YaoServiceClient
```
## How to Write a Test
### Standard pattern
Every test file follows this structure:
```go
// grpc/run/run_test.go
package run_test
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/yaoapp/yao/grpc/pb"
"github.com/yaoapp/yao/grpc/tests/testutils"
)
func TestRun_ProcessExec(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.Run(ctx, &pb.RunRequest{
Process: "utils.app.Ping",
})
assert.NoError(t, err)
assert.NotNil(t, resp.Data)
}
func TestRun_InvalidProcess(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.Run(ctx, &pb.RunRequest{Process: "nonexistent.process"})
assert.Error(t, err)
}
```
### Auth tests
Auth tests verify the interceptor chain through the gRPC client:
```go
// grpc/auth/guard_test.go
package auth_test
func TestAuth_NoToken_Rejected(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
_, err := client.Run(context.Background(), &pb.RunRequest{Process: "utils.app.Ping"})
st, _ := status.FromError(err)
assert.Equal(t, codes.Unauthenticated, st.Code())
}
func TestAuth_WrongScope_Denied(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.Run(ctx, &pb.RunRequest{Process: "utils.app.Ping"})
st, _ := status.FromError(err)
assert.Equal(t, codes.PermissionDenied, st.Code())
}
func TestAuth_TokenRefresh(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
// Mint an expired token + valid refresh token,
// send request with x-refresh-token metadata,
// verify response header contains x-new-access-token.
}
func TestHealthz_Public(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
resp, err := client.Healthz(context.Background(), &pb.Empty{})
assert.NoError(t, err)
assert.Equal(t, "ok", resp.Status)
}
```
### Streaming tests
```go
// grpc/llm/llm_test.go
package llm_test
func TestChatCompletionsStream(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:llm")
ctx := testutils.WithToken(context.Background(), token)
stream, err := client.ChatCompletionsStream(ctx, &pb.ChatRequest{
// ... model, messages, etc.
})
assert.NoError(t, err)
var chunks int
for {
chunk, err := stream.Recv()
if err == io.EOF {
break
}
assert.NoError(t, err)
chunks++
assert.NotEmpty(t, chunk.Data)
}
assert.Greater(t, chunks, 0)
}
```
```go
// grpc/agent/agent_test.go
package agent_test
func TestAgentStream(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:agent")
ctx := testutils.WithToken(context.Background(), token)
stream, err := client.AgentStream(ctx, &pb.AgentRequest{
RobotID: "test-robot",
// ...
})
assert.NoError(t, err)
var chunks int
for {
chunk, err := stream.Recv()
if err == io.EOF {
break
}
assert.NoError(t, err)
chunks++
// Each chunk carries JSON-serialized agent/output/message.Message
}
assert.Greater(t, chunks, 0)
}
func TestAgentStream_InvalidRobot(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:agent")
ctx := testutils.WithToken(context.Background(), token)
stream, err := client.AgentStream(ctx, &pb.AgentRequest{
RobotID: "nonexistent-robot",
})
// Either err on open or first Recv returns error
if err == nil {
_, err = stream.Recv()
}
assert.Error(t, err)
}
```
## Required Test Cases
Each sub-package must cover at minimum:
| Sub-package | Required cases |
|-------------|----------------|
| `auth` | valid token / no token (Unauthenticated) / expired token + refresh / wrong scope (PermissionDenied) / Healthz skips auth |
| `health` | Healthz returns ok without token |
| `run` | valid process / nonexistent process / bad arguments |
| `shell` | valid command / command not found / timeout |
| `api` | valid proxy / 404 endpoint |
| `mcp` | MCPListTools / MCPCallTool / MCPListResources / MCPReadResource |
| `llm` | ChatCompletions (unary) / ChatCompletionsStream (multiple chunks) / invalid model |
| `agent` | AgentStream (receives message chunks) / nonexistent robot ID |
## Makefile
Add to [Makefile](../Makefile):
```makefile
TESTFOLDER_GRPC := $(shell $(GO) list ./grpc/...)
.PHONY: unit-test-grpc
unit-test-grpc:
echo "mode: count" > coverage.out
for d in $(TESTFOLDER_GRPC); do \
$(GO) test -tags $(TESTTAGS) -v -timeout=10m \
-covermode=count -coverprofile=profile.out \
-coverpkg=$$(echo $$d | sed "s/\/test$$//g") \
-skip='TestMemoryLeak|TestIsolateDisposal|TestLeak_|TestScenario_' \
$$d > tmp.out; \
cat tmp.out; \
if grep -q "^--- FAIL" tmp.out; then \
rm tmp.out; \
exit 1; \
elif grep -q "build failed" tmp.out; then \
rm tmp.out; \
exit 1; \
elif grep -q "setup failed" tmp.out; then \
rm tmp.out; \
exit 1; \
elif grep -q "runtime error" tmp.out; then \
rm tmp.out; \
exit 1; \
fi; \
if [ -f profile.out ]; then \
cat profile.out | grep -v "mode:" >> coverage.out; \
rm profile.out; \
fi; \
done
```
Also add `|grpc` to the `TESTFOLDER_CORE` exclude pattern so core-test does not duplicate gRPC tests.
## CI Integration
Add `grpc-test` job to `unit-test.yml` and `pr-test.yml`:
```yaml
grpc-test:
runs-on: ubuntu-latest
services:
mongodb:
image: mongo:6.0
ports:
- 27017:27017
env:
MONGO_INITDB_ROOT_USERNAME: root
MONGO_INITDB_ROOT_PASSWORD: "123456"
MONGO_INITDB_DATABASE: test
strategy:
matrix:
go: ["1.25"]
steps:
# ... standard checkout + setup (same as core-test) ...
- name: Setup ENV (SQLite)
run: |
mkdir -p ${{ github.WORKSPACE }}/../app/db
echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV
echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV
- name: Run gRPC Tests
run: make unit-test-grpc
- name: Codecov Report
uses: codecov/codecov-action@v4
with:
token: ${{ secrets.CODECOV_TOKEN }}
```
Key decisions:
- SQLite only — gRPC is a transport layer, no need for MySQL matrix
- No Qdrant/Neo4j/MCP-everything services needed
- LLM/Agent streaming uses real `OPENAI_TEST_KEY` + `ANTHROPIC_API_KEY` (same secrets as agent-test job)
## Coverage
- Target: >80% per sub-package, >80% overall
- `grpc.go` (server lifecycle) covered indirectly via testutils.Prepare/Clean
- Coverage collected via `-coverprofile`, reported to Codecov
## Phase Test Schedule
Tests are written alongside implementation, not after:
| Phase | Test files | Repo |
|-------|------------|------|
| Phase 1 (auth + server) | `auth/guard_test.go`, `health/health_test.go` | yao |
| Phase 2 (handlers) | `run/run_test.go`, `shell/shell_test.go`, `api/api_test.go`, `mcp/mcp_test.go` | yao |
| Phase 3 (LLM + Agent) | `llm/llm_test.go`, `agent/agent_test.go` | yao |
| Phase 4 (Tai gateway) | Tai repo tests — gateway forwards `x-grpc-upstream`, conn cache reuse, missing metadata rejected | tai |
| Phase 5 (yao-grpc client) | `tai/grpc/grpc_test.go` — dial, method wrappers, token refresh via response metadata, `x-grpc-upstream` attachment | yao |
| Phase 6 (Device Flow) | `openapi/oauth/*_test.go` — DeviceAuthorization, device_code grant, poll pending/approved/expired | yao |
Each Phase PR must include tests for all new code. Coverage must meet threshold before merge.
## Running Tests
```bash
# All gRPC tests
make unit-test-grpc
# Single sub-package
go test -v ./grpc/auth/
# Single test
go test -v -run TestAuth_NoToken_Rejected ./grpc/auth/
```

93
grpc/agent/agent.go Normal file
View file

@ -0,0 +1,93 @@
package agent
import (
"net/http"
"strings"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/yaoapp/yao/agent"
"github.com/yaoapp/yao/agent/assistant"
agentContext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/grpc/auth"
"github.com/yaoapp/yao/grpc/pb"
)
// Handler implements the AgentStream gRPC method.
type Handler struct{}
// AgentStream resolves an assistant by ID and streams agent output as AgentChunk messages.
// Mirrors openapi/chat/completions.go GinCreateCompletions flow via context.GetGRPCAgentRequest.
func (h *Handler) AgentStream(req *pb.AgentRequest, stream grpc.ServerStreamingServer[pb.AgentChunk]) error {
ctx := stream.Context()
if req.AssistantId == "" {
return status.Error(codes.InvalidArgument, "assistant_id is required")
}
agentDSL := agent.GetAgent()
if agentDSL == nil {
return status.Error(codes.Internal, "agent DSL not initialized")
}
cache, err := agentDSL.GetCacheStore()
if err != nil {
return status.Errorf(codes.Internal, "failed to get cache store: %v", err)
}
messages, agentCtx, opts, err := agentContext.GetGRPCAgentRequest(ctx, agentContext.GRPCAgentInput{
AssistantID: req.AssistantId,
Messages: req.Messages,
Options: req.Options,
AuthInfo: auth.GetAuthorizedInfo(ctx),
Cache: cache,
Writer: &grpcStreamWriter{stream: stream, header: make(http.Header)},
})
if err != nil {
return toGRPCError(err)
}
defer agentCtx.Release()
ast, err := assistant.Get(agentCtx.AssistantID)
if err != nil {
return status.Errorf(codes.NotFound, "assistant not found: %v", err)
}
_, err = ast.Stream(agentCtx, messages, opts)
if err != nil {
return status.Errorf(codes.Internal, "agent stream failed: %v", err)
}
return stream.Send(&pb.AgentChunk{Done: true})
}
func toGRPCError(err error) error {
msg := err.Error()
if strings.Contains(msg, "is required") ||
strings.Contains(msg, "must not be empty") ||
strings.Contains(msg, "invalid") {
return status.Error(codes.InvalidArgument, msg)
}
return status.Error(codes.Internal, msg)
}
// grpcStreamWriter bridges agent/context.Writer (http.ResponseWriter) to gRPC stream.
type grpcStreamWriter struct {
stream grpc.ServerStreamingServer[pb.AgentChunk]
header http.Header
code int
}
func (w *grpcStreamWriter) Header() http.Header { return w.header }
func (w *grpcStreamWriter) WriteHeader(statusCode int) { w.code = statusCode }
func (w *grpcStreamWriter) Write(data []byte) (int, error) {
if err := w.stream.Send(&pb.AgentChunk{Data: data}); err != nil {
return 0, err
}
return len(data), nil
}
// Flush implements http.Flusher for streaming compatibility.
func (w *grpcStreamWriter) Flush() {}

206
grpc/agent/agent_test.go Normal file
View file

@ -0,0 +1,206 @@
package agent_test
import (
"context"
"encoding/json"
"io"
"os"
"testing"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/yaoapp/yao/grpc/pb"
"github.com/yaoapp/yao/grpc/tests/testutils"
)
func TestAgentStream_InvalidAssistant(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:agent")
ctx := testutils.WithToken(context.Background(), token)
msgs, _ := json.Marshal([]map[string]interface{}{
{"role": "user", "content": "hello"},
})
stream, err := client.AgentStream(ctx, &pb.AgentRequest{
AssistantId: "nonexistent-assistant-id",
Messages: msgs,
})
if err != nil {
st, _ := status.FromError(err)
assert.Equal(t, codes.NotFound, st.Code())
return
}
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.NotFound, st.Code())
}
func TestAgentStream_EmptyAssistantID(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:agent")
ctx := testutils.WithToken(context.Background(), token)
stream, err := client.AgentStream(ctx, &pb.AgentRequest{
AssistantId: "",
})
if err != nil {
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
return
}
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestAgentStream_EmptyMessages(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:agent")
ctx := testutils.WithToken(context.Background(), token)
msgs, _ := json.Marshal([]map[string]interface{}{})
stream, err := client.AgentStream(ctx, &pb.AgentRequest{
AssistantId: "some-assistant",
Messages: msgs,
})
if err != nil {
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
return
}
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestAgentStream_NilMessages(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:agent")
ctx := testutils.WithToken(context.Background(), token)
stream, err := client.AgentStream(ctx, &pb.AgentRequest{
AssistantId: "some-assistant",
Messages: nil,
})
if err != nil {
st, _ := status.FromError(err)
assert.NotEqual(t, codes.OK, st.Code())
return
}
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.NotEqual(t, codes.OK, st.Code())
}
func TestAgentStream_BadMessagesJSON(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:agent")
ctx := testutils.WithToken(context.Background(), token)
stream, err := client.AgentStream(ctx, &pb.AgentRequest{
AssistantId: "some-assistant",
Messages: []byte("{bad-json"),
})
if err != nil {
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
return
}
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestAgentStream_BadOptionsJSON(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:agent")
ctx := testutils.WithToken(context.Background(), token)
msgs, _ := json.Marshal([]map[string]interface{}{
{"role": "user", "content": "hello"},
})
stream, err := client.AgentStream(ctx, &pb.AgentRequest{
AssistantId: "some-assistant",
Messages: msgs,
Options: []byte("{bad-options"),
})
if err != nil {
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
return
}
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestAgentStream_RealAgent(t *testing.T) {
if os.Getenv("OPENAI_TEST_KEY") == "" {
t.Skip("OPENAI_TEST_KEY not set, skipping real agent test")
}
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:agent")
ctx := testutils.WithToken(context.Background(), token)
msgs, _ := json.Marshal([]map[string]interface{}{
{"role": "user", "content": "Say hello in one word."},
})
stream, err := client.AgentStream(ctx, &pb.AgentRequest{
AssistantId: "tests.nested.demo",
Messages: msgs,
})
if !assert.NoError(t, err) {
return
}
var chunks int
for {
chunk, err := stream.Recv()
if err == io.EOF {
break
}
if !assert.NoError(t, err) {
break
}
chunks++
if chunk.Done {
break
}
assert.NotEmpty(t, chunk.Data)
}
assert.Greater(t, chunks, 0)
}

69
grpc/api/api.go Normal file
View file

@ -0,0 +1,69 @@
package api
import (
"bytes"
"context"
"net/http"
"net/http/httptest"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
"github.com/yaoapp/yao/grpc/pb"
"github.com/yaoapp/yao/service"
)
// Handler implements the API gRPC method (internal HTTP proxy).
type Handler struct{}
// API proxies a gRPC request to the internal openapi HTTP router.
func (h *Handler) API(ctx context.Context, req *pb.APIRequest) (*pb.APIResponse, error) {
router := service.Router
if router == nil {
return nil, status.Error(codes.Unavailable, "HTTP router not initialized")
}
if req.Method == "" {
return nil, status.Error(codes.InvalidArgument, "method is required")
}
if req.Path == "" {
return nil, status.Error(codes.InvalidArgument, "path is required")
}
httpReq, err := http.NewRequestWithContext(ctx, req.Method, req.Path, bytes.NewReader(req.Body))
if err != nil {
return nil, status.Errorf(codes.Internal, "failed to build HTTP request: %v", err)
}
for k, v := range req.Headers {
httpReq.Header.Set(k, v)
}
// Forward Bearer token from gRPC metadata to HTTP Authorization header
// when the caller didn't explicitly set it.
if httpReq.Header.Get("Authorization") == "" {
if md, ok := metadata.FromIncomingContext(ctx); ok {
if vals := md.Get("authorization"); len(vals) > 0 {
httpReq.Header.Set("Authorization", vals[0])
}
}
}
w := httptest.NewRecorder()
router.ServeHTTP(w, httpReq)
result := w.Result()
defer result.Body.Close()
respHeaders := make(map[string]string, len(result.Header))
for k := range result.Header {
respHeaders[k] = result.Header.Get(k)
}
return &pb.APIResponse{
Status: int32(result.StatusCode),
Headers: respHeaders,
Body: w.Body.Bytes(),
}, nil
}

135
grpc/api/api_test.go Normal file
View file

@ -0,0 +1,135 @@
package api_test
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/yaoapp/yao/grpc/pb"
"github.com/yaoapp/yao/grpc/tests/testutils"
)
func TestAPI_Proxy(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
// The API method's ACL check uses the actual openapi path, so we grant all gRPC scopes.
// The openapi guard inside the HTTP router handles further auth via the forwarded Authorization header.
token := testutils.ObtainAccessToken(t,
"grpc:run", "grpc:stream", "grpc:shell", "grpc:mcp", "grpc:llm", "grpc:agent",
)
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.API(ctx, &pb.APIRequest{
Method: "GET",
Path: "/api/__yao/app/setting",
})
// The proxy itself should succeed (no gRPC error), even if the HTTP response
// is a non-200 status (e.g. 401 from openapi's own guard).
assert.NoError(t, err)
if assert.NotNil(t, resp) {
assert.Greater(t, resp.Status, int32(0))
assert.NotNil(t, resp.Body)
}
}
func TestAPI_NotFoundEndpoint(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t,
"grpc:run", "grpc:stream", "grpc:shell", "grpc:mcp", "grpc:llm", "grpc:agent",
)
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.API(ctx, &pb.APIRequest{
Method: "GET",
Path: "/api/this/does/not/exist",
})
assert.NoError(t, err)
if assert.NotNil(t, resp) {
assert.Equal(t, int32(404), resp.Status)
}
}
func TestAPI_MissingMethod(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.API(ctx, &pb.APIRequest{
Method: "",
Path: "/api/test",
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestAPI_MissingPath(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.API(ctx, &pb.APIRequest{
Method: "GET",
Path: "",
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestAPI_WithHeaders(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t,
"grpc:run", "grpc:stream", "grpc:shell", "grpc:mcp", "grpc:llm", "grpc:agent",
)
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.API(ctx, &pb.APIRequest{
Method: "GET",
Path: "/api/__yao/app/setting",
Headers: map[string]string{"X-Custom-Header": "test-value"},
})
assert.NoError(t, err)
assert.NotNil(t, resp)
assert.Greater(t, resp.Status, int32(0))
}
func TestAPI_PostWithBody(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t,
"grpc:run", "grpc:stream", "grpc:shell", "grpc:mcp", "grpc:llm", "grpc:agent",
)
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.API(ctx, &pb.APIRequest{
Method: "POST",
Path: "/api/this/does/not/exist",
Body: []byte(`{"key":"value"}`),
})
assert.NoError(t, err)
if assert.NotNil(t, resp) {
assert.Equal(t, int32(404), resp.Status)
}
}

66
grpc/auth/endpoint.go Normal file
View file

@ -0,0 +1,66 @@
package auth
import (
"fmt"
"strings"
"github.com/yaoapp/yao/grpc/pb"
)
// VirtualEndpoint maps a gRPC full method + request to a virtual HTTP endpoint for ACL.
// Returns the HTTP method and path used for scope-based access control.
func VirtualEndpoint(fullMethod string, req interface{}) (method string, path string) {
switch fullMethod {
case "/yao.Yao/Run":
if r, ok := req.(*pb.RunRequest); ok && r.Process != "" {
return "POST", "/grpc/run/" + r.Process
}
return "POST", "/grpc/run/"
case "/yao.Yao/Stream":
if r, ok := req.(*pb.RunRequest); ok && r.Process != "" {
return "POST", "/grpc/stream/" + r.Process
}
return "POST", "/grpc/stream/"
case "/yao.Yao/Shell", "/yao.Yao/ShellStream":
return "POST", "/grpc/shell"
case "/yao.Yao/API":
if r, ok := req.(*pb.APIRequest); ok && r.Path != "" {
m := strings.ToUpper(r.Method)
if m == "" {
m = "POST"
}
return m, r.Path
}
return "POST", "/"
case "/yao.Yao/MCPListTools":
return "GET", "/grpc/mcp/tools"
case "/yao.Yao/MCPCallTool":
if r, ok := req.(*pb.MCPCallRequest); ok && r.Tool != "" {
return "POST", "/grpc/mcp/call/" + r.Tool
}
return "POST", "/grpc/mcp/call/"
case "/yao.Yao/MCPListResources":
return "GET", "/grpc/mcp/resources"
case "/yao.Yao/MCPReadResource":
return "GET", "/grpc/mcp/resources/read"
case "/yao.Yao/ChatCompletions", "/yao.Yao/ChatCompletionsStream":
return "POST", "/grpc/llm/completions"
case "/yao.Yao/AgentStream":
if r, ok := req.(*pb.AgentRequest); ok && r.AssistantId != "" {
return "POST", fmt.Sprintf("/grpc/agent/%s", r.AssistantId)
}
return "POST", "/grpc/agent/"
default:
return "POST", "/grpc/unknown"
}
}

136
grpc/auth/endpoint_test.go Normal file
View file

@ -0,0 +1,136 @@
package auth_test
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/grpc/auth"
"github.com/yaoapp/yao/grpc/pb"
)
func TestVirtualEndpoint_Run(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/Run", &pb.RunRequest{Process: "models.user.Find"})
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/run/models.user.Find", path)
}
func TestVirtualEndpoint_Stream(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/Stream", &pb.RunRequest{Process: "flows.report"})
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/stream/flows.report", path)
}
func TestVirtualEndpoint_Shell(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/Shell", &pb.ShellRequest{Command: "ls"})
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/shell", path)
}
func TestVirtualEndpoint_ShellStream(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/ShellStream", &pb.ShellRequest{Command: "ls"})
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/shell", path)
}
func TestVirtualEndpoint_API(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/API", &pb.APIRequest{Method: "GET", Path: "/kb/collections"})
assert.Equal(t, "GET", method)
assert.Equal(t, "/kb/collections", path)
}
func TestVirtualEndpoint_MCPListTools(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/MCPListTools", &pb.MCPListRequest{SessionId: "abc"})
assert.Equal(t, "GET", method)
assert.Equal(t, "/grpc/mcp/tools", path)
}
func TestVirtualEndpoint_MCPCallTool(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/MCPCallTool", &pb.MCPCallRequest{Tool: "search"})
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/mcp/call/search", path)
}
func TestVirtualEndpoint_MCPListResources(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/MCPListResources", &pb.MCPListRequest{})
assert.Equal(t, "GET", method)
assert.Equal(t, "/grpc/mcp/resources", path)
}
func TestVirtualEndpoint_MCPReadResource(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/MCPReadResource", &pb.MCPResourceRequest{Uri: "file://test"})
assert.Equal(t, "GET", method)
assert.Equal(t, "/grpc/mcp/resources/read", path)
}
func TestVirtualEndpoint_ChatCompletions(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/ChatCompletions", &pb.ChatRequest{Connector: "openai"})
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/llm/completions", path)
}
func TestVirtualEndpoint_ChatCompletionsStream(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/ChatCompletionsStream", &pb.ChatRequest{})
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/llm/completions", path)
}
func TestVirtualEndpoint_AgentStream(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/AgentStream", &pb.AgentRequest{AssistantId: "my-robot"})
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/agent/my-robot", path)
}
func TestVirtualEndpoint_Unknown(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/NonExistent", nil)
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/unknown", path)
}
func TestVirtualEndpoint_RunNilReq(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/Run", nil)
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/run/", path)
}
func TestVirtualEndpoint_RunEmptyProcess(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/Run", &pb.RunRequest{Process: ""})
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/run/", path)
}
func TestVirtualEndpoint_StreamNilReq(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/Stream", nil)
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/stream/", path)
}
func TestVirtualEndpoint_APINilReq(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/API", nil)
assert.Equal(t, "POST", method)
assert.Equal(t, "/", path)
}
func TestVirtualEndpoint_APIEmptyMethod(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/API", &pb.APIRequest{Method: "", Path: "/test"})
assert.Equal(t, "POST", method)
assert.Equal(t, "/test", path)
}
func TestVirtualEndpoint_MCPCallToolNilReq(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/MCPCallTool", nil)
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/mcp/call/", path)
}
func TestVirtualEndpoint_AgentStreamNilReq(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/AgentStream", nil)
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/agent/", path)
}
func TestVirtualEndpoint_AgentStreamEmptyID(t *testing.T) {
method, path := auth.VirtualEndpoint("/yao.Yao/AgentStream", &pb.AgentRequest{AssistantId: ""})
assert.Equal(t, "POST", method)
assert.Equal(t, "/grpc/agent/", path)
}

169
grpc/auth/guard.go Normal file
View file

@ -0,0 +1,169 @@
package auth
import (
"context"
"strings"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
"github.com/yaoapp/yao/openapi/oauth"
"github.com/yaoapp/yao/openapi/oauth/acl"
"github.com/yaoapp/yao/openapi/oauth/types"
)
const (
healthzMethod = "/yao.Yao/Healthz"
apiMethod = "/yao.Yao/API"
metaAuthorization = "authorization"
metaRefreshToken = "x-refresh-token"
metaAccessToken = "x-access-token"
metaSandboxID = "x-sandbox-id"
metaSessionID = "x-session-id"
)
type authCtxKey struct{}
// WithAuthorizedInfo stores AuthorizedInfo in context for downstream handlers.
func WithAuthorizedInfo(ctx context.Context, info *types.AuthorizedInfo) context.Context {
return context.WithValue(ctx, authCtxKey{}, info)
}
// GetAuthorizedInfo retrieves AuthorizedInfo from context (set by the interceptor).
func GetAuthorizedInfo(ctx context.Context) *types.AuthorizedInfo {
info, _ := ctx.Value(authCtxKey{}).(*types.AuthorizedInfo)
return info
}
// UnaryInterceptor is the gRPC unary server interceptor for authentication and authorization.
func UnaryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
if info.FullMethod == healthzMethod {
return handler(ctx, req)
}
ctx, err := authenticate(ctx, info.FullMethod, req)
if err != nil {
return nil, err
}
return handler(ctx, req)
}
// StreamInterceptor is the gRPC stream server interceptor for authentication and authorization.
// For streaming RPCs, the request object is not available at intercept time,
// so ACL scope check uses the method-level virtual path (without request-specific IDs).
func StreamInterceptor(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
if info.FullMethod == healthzMethod {
return handler(srv, ss)
}
ctx, err := authenticate(ss.Context(), info.FullMethod, nil)
if err != nil {
return err
}
return handler(srv, &wrappedStream{ServerStream: ss, ctx: ctx})
}
// authenticate calls oauth.Service.AuthenticateToken directly — no gin/HTTP shim.
func authenticate(ctx context.Context, fullMethod string, req interface{}) (context.Context, error) {
md, ok := metadata.FromIncomingContext(ctx)
if !ok {
return ctx, status.Error(codes.Unauthenticated, "missing metadata")
}
svc := oauth.OAuth
if svc == nil {
return ctx, status.Error(codes.Internal, "oauth service not initialized")
}
bearer := extractBearer(md)
if bearer == "" {
return ctx, status.Error(codes.Unauthenticated, "missing authorization token")
}
result, err := svc.AuthenticateToken(oauth.AuthInput{
AccessToken: bearer,
RefreshToken: extractMeta(md, metaRefreshToken),
SessionID: extractMeta(md, metaSessionID),
})
if err != nil {
return ctx, status.Error(codes.Unauthenticated, err.Error())
}
ctx = WithAuthorizedInfo(ctx, result.Info)
if result.NewAccessToken != "" {
_ = grpc.SendHeader(ctx, metadata.Pairs(
metaAccessToken, result.NewAccessToken,
metaRefreshToken, result.NewRefreshToken,
))
}
// ACL scope check — skip for API proxy (the openapi router does its own auth).
if fullMethod != apiMethod {
httpMethod, httpPath := VirtualEndpoint(fullMethod, req)
scopes := strings.Fields(result.Info.Scope)
enforcer := getACLEnforcer()
if enforcer != nil && enforcer.Scope != nil {
decision := enforcer.Scope.Check(&acl.AccessRequest{
Method: httpMethod,
Path: httpPath,
Scopes: scopes,
})
if !decision.Allowed {
return ctx, status.Errorf(codes.PermissionDenied, "insufficient scope: %s", decision.Reason)
}
}
}
return ctx, nil
}
// getACLEnforcer returns the ACL enforcer if available and enabled.
func getACLEnforcer() *acl.ACL {
if acl.Global == nil {
return nil
}
enforcer, ok := acl.Global.(*acl.ACL)
if !ok || enforcer == nil {
return nil
}
if !enforcer.Config.Enabled {
return nil
}
return enforcer
}
func extractBearer(md metadata.MD) string {
vals := md.Get(metaAuthorization)
if len(vals) == 0 {
return ""
}
parts := strings.SplitN(vals[0], " ", 2)
if len(parts) == 2 && strings.EqualFold(parts[0], "bearer") {
return parts[1]
}
return vals[0]
}
func extractMeta(md metadata.MD, key string) string {
vals := md.Get(key)
if len(vals) == 0 {
return ""
}
return vals[0]
}
// wrappedStream wraps grpc.ServerStream with a custom context.
type wrappedStream struct {
grpc.ServerStream
ctx context.Context
}
func (w *wrappedStream) Context() context.Context {
return w.ctx
}

169
grpc/auth/guard_test.go Normal file
View file

@ -0,0 +1,169 @@
package auth_test
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/metadata"
"google.golang.org/grpc/status"
"github.com/yaoapp/yao/grpc/pb"
"github.com/yaoapp/yao/grpc/tests/testutils"
)
func TestAuth_NoToken_Rejected(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
_, err := client.Run(context.Background(), &pb.RunRequest{Process: "utils.app.Ping"})
assert.Error(t, err)
st, ok := status.FromError(err)
assert.True(t, ok)
assert.Equal(t, codes.Unauthenticated, st.Code())
}
func TestAuth_ValidToken(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.Run(ctx, &pb.RunRequest{Process: "utils.app.Ping"})
// Run returns Unimplemented (handler stub), not an auth error
st, ok := status.FromError(err)
assert.True(t, ok)
assert.NotEqual(t, codes.Unauthenticated, st.Code())
assert.NotEqual(t, codes.PermissionDenied, st.Code())
}
func TestAuth_WrongScope_Denied(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.Run(ctx, &pb.RunRequest{Process: "utils.app.Ping"})
assert.Error(t, err)
st, ok := status.FromError(err)
assert.True(t, ok)
assert.Equal(t, codes.PermissionDenied, st.Code())
}
func TestAuth_TokenRefresh(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
expiredToken := testutils.ObtainExpiredAccessToken(t, "grpc:run")
refreshToken := testutils.ObtainRefreshToken(t, "grpc:run")
ctx := testutils.WithRefreshToken(context.Background(), expiredToken, refreshToken)
// The call should succeed (auth interceptor refreshes the token)
_, err := client.Run(ctx, &pb.RunRequest{Process: "utils.app.Ping"})
st, ok := status.FromError(err)
assert.True(t, ok)
// Should not be an auth error — either Unimplemented (handler stub) or OK
assert.NotEqual(t, codes.Unauthenticated, st.Code())
assert.NotEqual(t, codes.PermissionDenied, st.Code())
}
func TestHealthz_Public(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
resp, err := client.Healthz(context.Background(), &pb.Empty{})
assert.NoError(t, err)
assert.NotNil(t, resp)
assert.Equal(t, "ok", resp.Status)
}
func TestAuth_InvalidBearerFormat(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
ctx := metadata.AppendToOutgoingContext(context.Background(), "authorization", "not-a-valid-token-at-all")
_, err := client.Run(ctx, &pb.RunRequest{Process: "utils.app.Ping"})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.Unauthenticated, st.Code())
}
func TestAuth_StreamInterceptor_NoToken(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
stream, err := client.ChatCompletionsStream(context.Background(), &pb.ChatRequest{
Connector: "openai",
})
if err != nil {
st, _ := status.FromError(err)
assert.Equal(t, codes.Unauthenticated, st.Code())
return
}
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.Unauthenticated, st.Code())
}
func TestAuth_StreamInterceptor_ValidToken(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:agent")
ctx := testutils.WithToken(context.Background(), token)
stream, err := client.AgentStream(ctx, &pb.AgentRequest{
AssistantId: "nonexistent",
Messages: []byte(`[{"role":"user","content":"hi"}]`),
})
if err != nil {
st, _ := status.FromError(err)
assert.NotEqual(t, codes.Unauthenticated, st.Code())
assert.NotEqual(t, codes.PermissionDenied, st.Code())
return
}
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.NotEqual(t, codes.Unauthenticated, st.Code())
assert.NotEqual(t, codes.PermissionDenied, st.Code())
}
func TestAuth_StreamInterceptor_WrongScope(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
stream, err := client.AgentStream(ctx, &pb.AgentRequest{
AssistantId: "test",
Messages: []byte(`[{"role":"user","content":"hi"}]`),
})
if err != nil {
st, _ := status.FromError(err)
assert.Equal(t, codes.PermissionDenied, st.Code())
return
}
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.PermissionDenied, st.Code())
}

14
grpc/auth/scope.go Normal file
View file

@ -0,0 +1,14 @@
package auth
import "github.com/yaoapp/yao/openapi/oauth/acl"
func init() {
acl.Register(
&acl.ScopeDefinition{Name: "grpc:run", Endpoints: []string{"POST /grpc/run/*", "POST /grpc/run/"}},
&acl.ScopeDefinition{Name: "grpc:stream", Endpoints: []string{"POST /grpc/stream/*", "POST /grpc/stream/"}},
&acl.ScopeDefinition{Name: "grpc:shell", Endpoints: []string{"POST /grpc/shell"}},
&acl.ScopeDefinition{Name: "grpc:mcp", Endpoints: []string{"GET /grpc/mcp/tools", "POST /grpc/mcp/call/*", "POST /grpc/mcp/call/", "GET /grpc/mcp/resources", "GET /grpc/mcp/resources/read"}},
&acl.ScopeDefinition{Name: "grpc:llm", Endpoints: []string{"POST /grpc/llm/completions"}},
&acl.ScopeDefinition{Name: "grpc:agent", Endpoints: []string{"POST /grpc/agent/*", "POST /grpc/agent/"}},
)
}

173
grpc/grpc.go Normal file
View file

@ -0,0 +1,173 @@
package grpc
import (
"context"
"net"
"strconv"
"strings"
"sync"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/config"
agenthandler "github.com/yaoapp/yao/grpc/agent"
apihandler "github.com/yaoapp/yao/grpc/api"
"github.com/yaoapp/yao/grpc/auth"
"github.com/yaoapp/yao/grpc/health"
llmhandler "github.com/yaoapp/yao/grpc/llm"
mcphandler "github.com/yaoapp/yao/grpc/mcp"
"github.com/yaoapp/yao/grpc/pb"
runhandler "github.com/yaoapp/yao/grpc/run"
shellhandler "github.com/yaoapp/yao/grpc/shell"
)
var (
mu sync.Mutex
server *grpc.Server
listeners []net.Listener
addrs []string
)
type yaoServer struct {
pb.UnimplementedYaoServer
health health.Handler
run runhandler.Handler
shell shellhandler.Handler
api apihandler.Handler
mcp mcphandler.Handler
llm llmhandler.Handler
agent agenthandler.Handler
}
// ── Health ───────────────────────────────────────────────────────────────────
func (s *yaoServer) Healthz(ctx context.Context, req *pb.Empty) (*pb.HealthzResponse, error) {
return s.health.Healthz(ctx, req)
}
// ── Base ─────────────────────────────────────────────────────────────────────
func (s *yaoServer) Run(ctx context.Context, req *pb.RunRequest) (*pb.RunResponse, error) {
return s.run.Run(ctx, req)
}
func (s *yaoServer) Shell(ctx context.Context, req *pb.ShellRequest) (*pb.ShellResponse, error) {
return s.shell.Shell(ctx, req)
}
// V2 stubs — Stream and ShellStream depend on gou/stream package.
func (s *yaoServer) Stream(req *pb.RunRequest, stream grpc.ServerStreamingServer[pb.Chunk]) error {
return status.Error(codes.Unimplemented, "Stream not implemented (V2)")
}
func (s *yaoServer) ShellStream(req *pb.ShellRequest, stream grpc.ServerStreamingServer[pb.Chunk]) error {
return status.Error(codes.Unimplemented, "ShellStream not implemented (V2)")
}
// ── API ──────────────────────────────────────────────────────────────────────
func (s *yaoServer) API(ctx context.Context, req *pb.APIRequest) (*pb.APIResponse, error) {
return s.api.API(ctx, req)
}
// ── MCP ──────────────────────────────────────────────────────────────────────
func (s *yaoServer) MCPListTools(ctx context.Context, req *pb.MCPListRequest) (*pb.MCPListResponse, error) {
return s.mcp.MCPListTools(ctx, req)
}
func (s *yaoServer) MCPCallTool(ctx context.Context, req *pb.MCPCallRequest) (*pb.MCPCallResponse, error) {
return s.mcp.MCPCallTool(ctx, req)
}
func (s *yaoServer) MCPListResources(ctx context.Context, req *pb.MCPListRequest) (*pb.MCPResourcesResponse, error) {
return s.mcp.MCPListResources(ctx, req)
}
func (s *yaoServer) MCPReadResource(ctx context.Context, req *pb.MCPResourceRequest) (*pb.MCPResourceResponse, error) {
return s.mcp.MCPReadResource(ctx, req)
}
// ── LLM ──────────────────────────────────────────────────────────────────────
func (s *yaoServer) ChatCompletions(ctx context.Context, req *pb.ChatRequest) (*pb.ChatResponse, error) {
return s.llm.ChatCompletions(ctx, req)
}
func (s *yaoServer) ChatCompletionsStream(req *pb.ChatRequest, stream grpc.ServerStreamingServer[pb.ChatChunk]) error {
return s.llm.ChatCompletionsStream(req, stream)
}
// ── Agent ────────────────────────────────────────────────────────────────────
func (s *yaoServer) AgentStream(req *pb.AgentRequest, stream grpc.ServerStreamingServer[pb.AgentChunk]) error {
return s.agent.AgentStream(req, stream)
}
// ── Server lifecycle ─────────────────────────────────────────────────────────
// StartServer initializes and starts the gRPC server based on config.
// It supports multiple bind addresses and returns immediately (listeners run in goroutines).
func StartServer(cfg config.Config) error {
if strings.ToLower(cfg.GRPC.Enabled) == "off" {
log.Info("gRPC server disabled (YAO_GRPC=off)")
return nil
}
mu.Lock()
defer mu.Unlock()
server = grpc.NewServer(
grpc.ChainUnaryInterceptor(auth.UnaryInterceptor),
grpc.ChainStreamInterceptor(auth.StreamInterceptor),
)
pb.RegisterYaoServer(server, &yaoServer{})
hosts := strings.Split(cfg.GRPC.Host, ",")
port := strconv.Itoa(cfg.GRPC.Port)
for _, h := range hosts {
addr := net.JoinHostPort(strings.TrimSpace(h), port)
lis, err := net.Listen("tcp", addr)
if err != nil {
Stop()
return err
}
listeners = append(listeners, lis)
addrs = append(addrs, lis.Addr().String())
log.Info("gRPC server listening on %s", lis.Addr().String())
go func(l net.Listener) {
if err := server.Serve(l); err != nil {
log.Error("gRPC server error on %s: %s", l.Addr().String(), err.Error())
}
}(lis)
}
return nil
}
// Stop gracefully stops the gRPC server. Safe to call if server was never started.
func Stop() {
mu.Lock()
defer mu.Unlock()
if server != nil {
server.GracefulStop()
server = nil
}
listeners = nil
addrs = nil
}
// Addr returns all addresses the gRPC server is listening on.
func Addr() []string {
mu.Lock()
defer mu.Unlock()
result := make([]string, len(addrs))
copy(result, addrs)
return result
}

15
grpc/health/health.go Normal file
View file

@ -0,0 +1,15 @@
package health
import (
"context"
"github.com/yaoapp/yao/grpc/pb"
)
// Handler implements the Healthz RPC.
type Handler struct{}
// Healthz returns server health status. This method is public (no auth required).
func (h *Handler) Healthz(ctx context.Context, req *pb.Empty) (*pb.HealthzResponse, error) {
return &pb.HealthzResponse{Status: "ok"}, nil
}

22
grpc/health/health_test.go Normal file
View file

@ -0,0 +1,22 @@
package health_test
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/grpc/pb"
"github.com/yaoapp/yao/grpc/tests/testutils"
)
func TestHealthz_ReturnsOk(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
resp, err := client.Healthz(context.Background(), &pb.Empty{})
assert.NoError(t, err)
assert.NotNil(t, resp)
assert.Equal(t, "ok", resp.Status)
}

165
grpc/llm/llm.go Normal file
View file

@ -0,0 +1,165 @@
package llm
import (
"context"
"encoding/json"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/yaoapp/gou/connector"
agentContext "github.com/yaoapp/yao/agent/context"
agentLLM "github.com/yaoapp/yao/agent/llm"
"github.com/yaoapp/yao/agent/output/message"
"github.com/yaoapp/yao/grpc/auth"
"github.com/yaoapp/yao/grpc/pb"
)
// Handler implements the LLM gRPC methods.
type Handler struct{}
// ChatCompletions sends messages to an LLM connector and returns the full response (unary).
func (h *Handler) ChatCompletions(ctx context.Context, req *pb.ChatRequest) (*pb.ChatResponse, error) {
if req.Connector == "" {
return nil, status.Error(codes.InvalidArgument, "connector is required")
}
llmInstance, completionOpts, ctxMessages, agentCtx, err := prepareLLMCall(ctx, req)
if err != nil {
return nil, err
}
defer agentCtx.Release()
noopHandler := func(chunkType message.StreamChunkType, data []byte) int { return 0 }
response, err := llmInstance.Stream(agentCtx, ctxMessages, completionOpts, noopHandler)
if err != nil {
return nil, status.Errorf(codes.Internal, "LLM call failed: %v", err)
}
data, err := json.Marshal(toOpenAIFormat(response))
if err != nil {
return nil, status.Errorf(codes.Internal, "failed to marshal LLM response: %v", err)
}
return &pb.ChatResponse{Data: data}, nil
}
// ChatCompletionsStream sends messages to an LLM connector and streams response chunks.
func (h *Handler) ChatCompletionsStream(req *pb.ChatRequest, stream grpc.ServerStreamingServer[pb.ChatChunk]) error {
ctx := stream.Context()
if req.Connector == "" {
return status.Error(codes.InvalidArgument, "connector is required")
}
llmInstance, completionOpts, ctxMessages, agentCtx, err := prepareLLMCall(ctx, req)
if err != nil {
return err
}
defer agentCtx.Release()
streamHandler := func(chunkType message.StreamChunkType, data []byte) int {
if ctx.Err() != nil {
return 1
}
if chunkType == message.ChunkText || chunkType == message.ChunkThinking {
if sendErr := stream.Send(&pb.ChatChunk{Data: data}); sendErr != nil {
return 1
}
}
return 0
}
_, err = llmInstance.Stream(agentCtx, ctxMessages, completionOpts, streamHandler)
if err != nil {
return status.Errorf(codes.Internal, "LLM stream failed: %v", err)
}
return stream.Send(&pb.ChatChunk{Done: true})
}
// prepareLLMCall builds the LLM instance, messages, and agent context from the gRPC request.
// Mirrors agent/llm/process.go ProcessChatCompletions logic without the process wrapper.
func prepareLLMCall(ctx context.Context, req *pb.ChatRequest) (agentLLM.LLM, *agentContext.CompletionOptions, []agentContext.Message, *agentContext.Context, error) {
ctxMessages, err := parseMessagesToContext(req.Messages)
if err != nil {
return nil, nil, nil, nil, err
}
var opts map[string]interface{}
if len(req.Options) > 0 {
if err := json.Unmarshal(req.Options, &opts); err != nil {
return nil, nil, nil, nil, status.Errorf(codes.InvalidArgument, "invalid options JSON: %v", err)
}
}
conn, err := connector.Select(req.Connector)
if err != nil {
return nil, nil, nil, nil, status.Errorf(codes.NotFound, "connector %s not found: %v", req.Connector, err)
}
completionOpts := agentLLM.BuildCompletionOptions(conn, opts)
llmInstance, err := agentLLM.New(conn, completionOpts)
if err != nil {
return nil, nil, nil, nil, status.Errorf(codes.Internal, "failed to create LLM: %v", err)
}
authInfo := auth.GetAuthorizedInfo(ctx)
chatID := agentContext.GenChatID()
agentCtx := agentContext.New(ctx, authInfo, chatID)
return llmInstance, completionOpts, ctxMessages, agentCtx, nil
}
// parseMessagesToContext converts raw JSON message bytes to []agentContext.Message via JSON round-trip.
func parseMessagesToContext(raw []byte) ([]agentContext.Message, error) {
if len(raw) == 0 {
return nil, status.Error(codes.InvalidArgument, "messages are required")
}
var messages []agentContext.Message
if err := json.Unmarshal(raw, &messages); err != nil {
return nil, status.Errorf(codes.InvalidArgument, "invalid messages JSON: %v", err)
}
if len(messages) == 0 {
return nil, status.Error(codes.InvalidArgument, "messages must not be empty")
}
return messages, nil
}
// toOpenAIFormat converts CompletionResponse to OpenAI chat.completions format.
func toOpenAIFormat(resp *agentContext.CompletionResponse) map[string]interface{} {
if resp == nil {
return map[string]interface{}{"choices": []interface{}{}}
}
msgMap := map[string]interface{}{
"role": resp.Role,
"content": resp.Content,
}
if len(resp.ToolCalls) > 0 {
msgMap["tool_calls"] = resp.ToolCalls
}
choice := map[string]interface{}{
"index": 0,
"message": msgMap,
"finish_reason": "stop",
}
result := map[string]interface{}{
"id": resp.ID,
"object": "chat.completion",
"created": resp.Created,
"model": resp.Model,
"choices": []interface{}{choice},
}
if resp.Usage != nil {
result["usage"] = resp.Usage
}
return result
}

265
grpc/llm/llm_test.go Normal file
View file

@ -0,0 +1,265 @@
package llm_test
import (
"context"
"encoding/json"
"io"
"os"
"testing"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/yaoapp/yao/grpc/pb"
"github.com/yaoapp/yao/grpc/tests/testutils"
)
func TestChatCompletions_InvalidConnector(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:llm")
ctx := testutils.WithToken(context.Background(), token)
msgs, _ := json.Marshal([]map[string]interface{}{
{"role": "user", "content": "hello"},
})
_, err := client.ChatCompletions(ctx, &pb.ChatRequest{
Connector: "nonexistent-connector",
Messages: msgs,
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.NotFound, st.Code())
}
func TestChatCompletions_EmptyConnector(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:llm")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.ChatCompletions(ctx, &pb.ChatRequest{
Connector: "",
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestChatCompletionsStream_EmptyConnector(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:llm")
ctx := testutils.WithToken(context.Background(), token)
stream, err := client.ChatCompletionsStream(ctx, &pb.ChatRequest{
Connector: "",
})
if err != nil {
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
return
}
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestChatCompletions_BadMessagesJSON(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:llm")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.ChatCompletions(ctx, &pb.ChatRequest{
Connector: "openai",
Messages: []byte("{bad-json"),
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestChatCompletions_EmptyMessages(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:llm")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.ChatCompletions(ctx, &pb.ChatRequest{
Connector: "openai",
Messages: nil,
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestChatCompletions_EmptyMessageArray(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:llm")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.ChatCompletions(ctx, &pb.ChatRequest{
Connector: "openai",
Messages: []byte("[]"),
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestChatCompletions_BadOptionsJSON(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:llm")
ctx := testutils.WithToken(context.Background(), token)
msgs, _ := json.Marshal([]map[string]interface{}{
{"role": "user", "content": "hello"},
})
_, err := client.ChatCompletions(ctx, &pb.ChatRequest{
Connector: "openai",
Messages: msgs,
Options: []byte("{bad-options"),
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestChatCompletionsStream_InvalidConnector(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:llm")
ctx := testutils.WithToken(context.Background(), token)
msgs, _ := json.Marshal([]map[string]interface{}{
{"role": "user", "content": "hello"},
})
stream, err := client.ChatCompletionsStream(ctx, &pb.ChatRequest{
Connector: "nonexistent-connector",
Messages: msgs,
})
if err != nil {
st, _ := status.FromError(err)
assert.Equal(t, codes.NotFound, st.Code())
return
}
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.NotFound, st.Code())
}
func TestChatCompletionsStream_BadMessages(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:llm")
ctx := testutils.WithToken(context.Background(), token)
stream, err := client.ChatCompletionsStream(ctx, &pb.ChatRequest{
Connector: "openai",
Messages: []byte("{bad-json"),
})
if err != nil {
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
return
}
_, err = stream.Recv()
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
// TestChatCompletions_RealLLM tests against a real LLM if OPENAI_TEST_KEY is set.
func TestChatCompletions_RealLLM(t *testing.T) {
if os.Getenv("OPENAI_TEST_KEY") == "" {
t.Skip("OPENAI_TEST_KEY not set, skipping real LLM test")
}
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:llm")
ctx := testutils.WithToken(context.Background(), token)
msgs, _ := json.Marshal([]map[string]interface{}{
{"role": "user", "content": "Say hello in one word."},
})
resp, err := client.ChatCompletions(ctx, &pb.ChatRequest{
Connector: "gpt-4o-mini",
Messages: msgs,
})
assert.NoError(t, err)
if assert.NotNil(t, resp) {
assert.NotEmpty(t, resp.Data)
}
}
// TestChatCompletionsStream_RealLLM tests streaming against a real LLM if OPENAI_TEST_KEY is set.
func TestChatCompletionsStream_RealLLM(t *testing.T) {
if os.Getenv("OPENAI_TEST_KEY") == "" {
t.Skip("OPENAI_TEST_KEY not set, skipping real LLM stream test")
}
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:llm")
ctx := testutils.WithToken(context.Background(), token)
msgs, _ := json.Marshal([]map[string]interface{}{
{"role": "user", "content": "Count from 1 to 3."},
})
stream, err := client.ChatCompletionsStream(ctx, &pb.ChatRequest{
Connector: "gpt-4o-mini",
Messages: msgs,
})
assert.NoError(t, err)
var chunks int
for {
chunk, err := stream.Recv()
if err == io.EOF {
break
}
if !assert.NoError(t, err) {
break
}
chunks++
if chunk.Done {
break
}
assert.NotEmpty(t, chunk.Data)
}
assert.Greater(t, chunks, 0)
}

102
grpc/mcp/mcp.go Normal file
View file

@ -0,0 +1,102 @@
package mcp
import (
"context"
"encoding/json"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
goumcp "github.com/yaoapp/gou/mcp"
"github.com/yaoapp/yao/grpc/pb"
)
// Handler implements the MCP gRPC methods.
type Handler struct{}
// MCPListTools lists all available MCP tools for a given session.
func (h *Handler) MCPListTools(ctx context.Context, req *pb.MCPListRequest) (*pb.MCPListResponse, error) {
client, err := goumcp.Select(req.SessionId)
if err != nil {
return nil, status.Errorf(codes.NotFound, "MCP client not found: %v", err)
}
resp, err := client.ListTools(ctx, "")
if err != nil {
return nil, status.Errorf(codes.Internal, "ListTools failed: %v", err)
}
data, err := json.Marshal(resp.Tools)
if err != nil {
return nil, status.Errorf(codes.Internal, "failed to marshal tools: %v", err)
}
return &pb.MCPListResponse{Tools: data}, nil
}
// MCPCallTool calls an MCP tool by name with the provided arguments.
func (h *Handler) MCPCallTool(ctx context.Context, req *pb.MCPCallRequest) (*pb.MCPCallResponse, error) {
client, err := goumcp.Select(req.SessionId)
if err != nil {
return nil, status.Errorf(codes.NotFound, "MCP client not found: %v", err)
}
var args interface{}
if len(req.Arguments) > 0 {
if err := json.Unmarshal(req.Arguments, &args); err != nil {
return nil, status.Errorf(codes.InvalidArgument, "invalid arguments JSON: %v", err)
}
}
resp, err := client.CallTool(ctx, req.Tool, args)
if err != nil {
return nil, status.Errorf(codes.Internal, "CallTool failed: %v", err)
}
data, err := json.Marshal(resp)
if err != nil {
return nil, status.Errorf(codes.Internal, "failed to marshal result: %v", err)
}
return &pb.MCPCallResponse{Result: data}, nil
}
// MCPListResources lists all available MCP resources for a given session.
func (h *Handler) MCPListResources(ctx context.Context, req *pb.MCPListRequest) (*pb.MCPResourcesResponse, error) {
client, err := goumcp.Select(req.SessionId)
if err != nil {
return nil, status.Errorf(codes.NotFound, "MCP client not found: %v", err)
}
resp, err := client.ListResources(ctx, "")
if err != nil {
return nil, status.Errorf(codes.Internal, "ListResources failed: %v", err)
}
data, err := json.Marshal(resp.Resources)
if err != nil {
return nil, status.Errorf(codes.Internal, "failed to marshal resources: %v", err)
}
return &pb.MCPResourcesResponse{Resources: data}, nil
}
// MCPReadResource reads a specific MCP resource by URI.
func (h *Handler) MCPReadResource(ctx context.Context, req *pb.MCPResourceRequest) (*pb.MCPResourceResponse, error) {
client, err := goumcp.Select(req.SessionId)
if err != nil {
return nil, status.Errorf(codes.NotFound, "MCP client not found: %v", err)
}
resp, err := client.ReadResource(ctx, req.Uri)
if err != nil {
return nil, status.Errorf(codes.Internal, "ReadResource failed: %v", err)
}
data, err := json.Marshal(resp.Contents)
if err != nil {
return nil, status.Errorf(codes.Internal, "failed to marshal contents: %v", err)
}
return &pb.MCPResourceResponse{Contents: data}, nil
}

264
grpc/mcp/mcp_test.go Normal file
View file

@ -0,0 +1,264 @@
package mcp_test
import (
"context"
"encoding/json"
"testing"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/yaoapp/yao/grpc/pb"
"github.com/yaoapp/yao/grpc/tests/testutils"
)
const echoSession = "echo"
// --- MCPListTools ---
func TestMCPListTools_Success(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.MCPListTools(ctx, &pb.MCPListRequest{
SessionId: echoSession,
})
assert.NoError(t, err)
if assert.NotNil(t, resp) {
assert.NotEmpty(t, resp.Tools)
var tools []map[string]interface{}
err := json.Unmarshal(resp.Tools, &tools)
assert.NoError(t, err)
assert.GreaterOrEqual(t, len(tools), 3, "echo MCP defines ping, status, echo")
names := make(map[string]bool)
for _, tool := range tools {
if n, ok := tool["name"].(string); ok {
names[n] = true
}
}
assert.True(t, names["ping"], "should contain ping tool")
assert.True(t, names["status"], "should contain status tool")
assert.True(t, names["echo"], "should contain echo tool")
}
}
func TestMCPListTools_InvalidSession(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.MCPListTools(ctx, &pb.MCPListRequest{
SessionId: "nonexistent-session",
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.NotFound, st.Code())
}
// --- MCPCallTool ---
func TestMCPCallTool_Ping(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
args, _ := json.Marshal(map[string]interface{}{"count": 2, "message": "ping"})
resp, err := client.MCPCallTool(ctx, &pb.MCPCallRequest{
SessionId: echoSession,
Tool: "ping",
Arguments: args,
})
assert.NoError(t, err)
if assert.NotNil(t, resp) {
assert.NotEmpty(t, resp.Result)
var result map[string]interface{}
err := json.Unmarshal(resp.Result, &result)
assert.NoError(t, err)
}
}
func TestMCPCallTool_Echo(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
args, _ := json.Marshal(map[string]interface{}{"message": "hello", "uppercase": true})
resp, err := client.MCPCallTool(ctx, &pb.MCPCallRequest{
SessionId: echoSession,
Tool: "echo",
Arguments: args,
})
assert.NoError(t, err)
if assert.NotNil(t, resp) {
assert.NotEmpty(t, resp.Result)
}
}
func TestMCPCallTool_NilArgs(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.MCPCallTool(ctx, &pb.MCPCallRequest{
SessionId: echoSession,
Tool: "ping",
Arguments: nil,
})
assert.NoError(t, err)
if assert.NotNil(t, resp) {
assert.NotEmpty(t, resp.Result)
}
}
func TestMCPCallTool_InvalidSession(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.MCPCallTool(ctx, &pb.MCPCallRequest{
SessionId: "nonexistent-session",
Tool: "some-tool",
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.NotFound, st.Code())
}
func TestMCPCallTool_BadArgs(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.MCPCallTool(ctx, &pb.MCPCallRequest{
SessionId: echoSession,
Tool: "ping",
Arguments: []byte("{not-json"),
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
// --- MCPListResources ---
func TestMCPListResources_Success(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.MCPListResources(ctx, &pb.MCPListRequest{
SessionId: echoSession,
})
assert.NoError(t, err)
if assert.NotNil(t, resp) {
assert.NotEmpty(t, resp.Resources)
var resources []map[string]interface{}
err := json.Unmarshal(resp.Resources, &resources)
assert.NoError(t, err)
assert.GreaterOrEqual(t, len(resources), 2, "echo MCP defines info and health resources")
}
}
func TestMCPListResources_InvalidSession(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.MCPListResources(ctx, &pb.MCPListRequest{
SessionId: "nonexistent-session",
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.NotFound, st.Code())
}
// --- MCPReadResource ---
func TestMCPReadResource_Success(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.MCPReadResource(ctx, &pb.MCPResourceRequest{
SessionId: echoSession,
Uri: "echo://info",
})
assert.NoError(t, err)
if assert.NotNil(t, resp) {
assert.NotEmpty(t, resp.Contents)
var contents []map[string]interface{}
err := json.Unmarshal(resp.Contents, &contents)
assert.NoError(t, err)
assert.Greater(t, len(contents), 0)
}
}
func TestMCPReadResource_InvalidSession(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.MCPReadResource(ctx, &pb.MCPResourceRequest{
SessionId: "nonexistent-session",
Uri: "echo://info",
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.NotFound, st.Code())
}
func TestMCPReadResource_NotFoundURI(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:mcp")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.MCPReadResource(ctx, &pb.MCPResourceRequest{
SessionId: echoSession,
Uri: "echo://nonexistent",
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.Internal, st.Code())
}

1316
grpc/pb/yao.pb.go Normal file
View file

File diff suppressed because it is too large Load diff

149
grpc/pb/yao.proto Normal file
View file

@ -0,0 +1,149 @@
syntax = "proto3";
package yao;
option go_package = "github.com/yaoapp/yao/grpc/pb";
// Yao gRPC gateway. Shares OAuth + ACL scope system with openapi.
service Yao {
// Base
rpc Run(RunRequest) returns (RunResponse);
rpc Stream(RunRequest) returns (stream Chunk);
rpc Shell(ShellRequest) returns (ShellResponse);
rpc ShellStream(ShellRequest) returns (stream Chunk);
// API gateway
rpc API(APIRequest) returns (APIResponse);
// MCP
rpc MCPListTools(MCPListRequest) returns (MCPListResponse);
rpc MCPCallTool(MCPCallRequest) returns (MCPCallResponse);
rpc MCPListResources(MCPListRequest) returns (MCPResourcesResponse);
rpc MCPReadResource(MCPResourceRequest) returns (MCPResourceResponse);
// AI - LLM
rpc ChatCompletions(ChatRequest) returns (ChatResponse);
rpc ChatCompletionsStream(ChatRequest) returns (stream ChatChunk);
// AI - Agent
rpc AgentStream(AgentRequest) returns (stream AgentChunk);
// Health
rpc Healthz(Empty) returns (HealthzResponse);
}
// ── Base ─────────────────────────────────────────────────────────────────────
message RunRequest {
string process = 1;
bytes args = 2; // JSON-encoded argument array
int32 timeout = 3; // seconds, 0 = server default
}
message RunResponse {
bytes data = 1; // JSON-encoded result
}
message Chunk {
bytes data = 1;
bool done = 2;
}
message ShellRequest {
string command = 1;
repeated string args = 2;
map<string,string> env = 3;
int32 timeout = 4; // seconds, 0 = default 30s
}
message ShellResponse {
bytes stdout = 1;
bytes stderr = 2;
int32 exit_code = 3;
}
// ── API gateway ──────────────────────────────────────────────────────────────
message APIRequest {
string method = 1; // HTTP method
string path = 2; // openapi path
map<string,string> headers = 3;
bytes body = 4;
}
message APIResponse {
int32 status = 1; // HTTP status code
map<string,string> headers = 2;
bytes body = 3;
}
// ── MCP ──────────────────────────────────────────────────────────────────────
message MCPListRequest {
string session_id = 1;
}
message MCPListResponse {
bytes tools = 1; // JSON array of tool definitions
}
message MCPCallRequest {
string session_id = 1;
string tool = 2;
bytes arguments = 3; // JSON-encoded arguments
}
message MCPCallResponse {
bytes result = 1; // JSON-encoded result
}
message MCPResourcesResponse {
bytes resources = 1; // JSON array of resource definitions
}
message MCPResourceRequest {
string session_id = 1;
string uri = 2;
}
message MCPResourceResponse {
bytes contents = 1; // JSON-encoded resource contents
}
// ── LLM ──────────────────────────────────────────────────────────────────────
message ChatRequest {
string connector = 1; // connector ID
bytes messages = 2; // JSON-encoded message array
bytes options = 3; // JSON-encoded options
}
message ChatResponse {
bytes data = 1; // JSON-encoded completion result
}
message ChatChunk {
bytes data = 1; // JSON-encoded chunk
bool done = 2;
}
// ── Agent ────────────────────────────────────────────────────────────────────
message AgentRequest {
string assistant_id = 1;
bytes messages = 2; // JSON-encoded message array
bytes options = 3; // JSON-encoded options
}
// Each chunk carries JSON-serialized agent/output/message.Message.
message AgentChunk {
bytes data = 1;
bool done = 2;
}
// ── Health ───────────────────────────────────────────────────────────────────
message Empty {}
message HealthzResponse {
string status = 1;
}

606
grpc/pb/yao_grpc.pb.go Normal file
View file

@ -0,0 +1,606 @@
// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
// versions:
// - protoc-gen-go-grpc v1.6.1
// - protoc v4.25.0
// source: yao.proto
package pb
import (
context "context"
grpc "google.golang.org/grpc"
codes "google.golang.org/grpc/codes"
status "google.golang.org/grpc/status"
)
// This is a compile-time assertion to ensure that this generated file
// is compatible with the grpc package it is being compiled against.
// Requires gRPC-Go v1.64.0 or later.
const _ = grpc.SupportPackageIsVersion9
const (
Yao_Run_FullMethodName = "/yao.Yao/Run"
Yao_Stream_FullMethodName = "/yao.Yao/Stream"
Yao_Shell_FullMethodName = "/yao.Yao/Shell"
Yao_ShellStream_FullMethodName = "/yao.Yao/ShellStream"
Yao_API_FullMethodName = "/yao.Yao/API"
Yao_MCPListTools_FullMethodName = "/yao.Yao/MCPListTools"
Yao_MCPCallTool_FullMethodName = "/yao.Yao/MCPCallTool"
Yao_MCPListResources_FullMethodName = "/yao.Yao/MCPListResources"
Yao_MCPReadResource_FullMethodName = "/yao.Yao/MCPReadResource"
Yao_ChatCompletions_FullMethodName = "/yao.Yao/ChatCompletions"
Yao_ChatCompletionsStream_FullMethodName = "/yao.Yao/ChatCompletionsStream"
Yao_AgentStream_FullMethodName = "/yao.Yao/AgentStream"
Yao_Healthz_FullMethodName = "/yao.Yao/Healthz"
)
// YaoClient is the client API for Yao 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.
//
// Yao gRPC gateway. Shares OAuth + ACL scope system with openapi.
type YaoClient interface {
// Base
Run(ctx context.Context, in *RunRequest, opts ...grpc.CallOption) (*RunResponse, error)
Stream(ctx context.Context, in *RunRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[Chunk], error)
Shell(ctx context.Context, in *ShellRequest, opts ...grpc.CallOption) (*ShellResponse, error)
ShellStream(ctx context.Context, in *ShellRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[Chunk], error)
// API gateway
API(ctx context.Context, in *APIRequest, opts ...grpc.CallOption) (*APIResponse, error)
// MCP
MCPListTools(ctx context.Context, in *MCPListRequest, opts ...grpc.CallOption) (*MCPListResponse, error)
MCPCallTool(ctx context.Context, in *MCPCallRequest, opts ...grpc.CallOption) (*MCPCallResponse, error)
MCPListResources(ctx context.Context, in *MCPListRequest, opts ...grpc.CallOption) (*MCPResourcesResponse, error)
MCPReadResource(ctx context.Context, in *MCPResourceRequest, opts ...grpc.CallOption) (*MCPResourceResponse, error)
// AI - LLM
ChatCompletions(ctx context.Context, in *ChatRequest, opts ...grpc.CallOption) (*ChatResponse, error)
ChatCompletionsStream(ctx context.Context, in *ChatRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ChatChunk], error)
// AI - Agent
AgentStream(ctx context.Context, in *AgentRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[AgentChunk], error)
// Health
Healthz(ctx context.Context, in *Empty, opts ...grpc.CallOption) (*HealthzResponse, error)
}
type yaoClient struct {
cc grpc.ClientConnInterface
}
func NewYaoClient(cc grpc.ClientConnInterface) YaoClient {
return &yaoClient{cc}
}
func (c *yaoClient) Run(ctx context.Context, in *RunRequest, opts ...grpc.CallOption) (*RunResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(RunResponse)
err := c.cc.Invoke(ctx, Yao_Run_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
func (c *yaoClient) Stream(ctx context.Context, in *RunRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[Chunk], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &Yao_ServiceDesc.Streams[0], Yao_Stream_FullMethodName, cOpts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[RunRequest, Chunk]{ClientStream: stream}
if err := x.ClientStream.SendMsg(in); err != nil {
return nil, err
}
if err := x.ClientStream.CloseSend(); err != nil {
return nil, err
}
return x, nil
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type Yao_StreamClient = grpc.ServerStreamingClient[Chunk]
func (c *yaoClient) Shell(ctx context.Context, in *ShellRequest, opts ...grpc.CallOption) (*ShellResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(ShellResponse)
err := c.cc.Invoke(ctx, Yao_Shell_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
func (c *yaoClient) ShellStream(ctx context.Context, in *ShellRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[Chunk], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &Yao_ServiceDesc.Streams[1], Yao_ShellStream_FullMethodName, cOpts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[ShellRequest, Chunk]{ClientStream: stream}
if err := x.ClientStream.SendMsg(in); err != nil {
return nil, err
}
if err := x.ClientStream.CloseSend(); err != nil {
return nil, err
}
return x, nil
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type Yao_ShellStreamClient = grpc.ServerStreamingClient[Chunk]
func (c *yaoClient) API(ctx context.Context, in *APIRequest, opts ...grpc.CallOption) (*APIResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(APIResponse)
err := c.cc.Invoke(ctx, Yao_API_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
func (c *yaoClient) MCPListTools(ctx context.Context, in *MCPListRequest, opts ...grpc.CallOption) (*MCPListResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(MCPListResponse)
err := c.cc.Invoke(ctx, Yao_MCPListTools_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
func (c *yaoClient) MCPCallTool(ctx context.Context, in *MCPCallRequest, opts ...grpc.CallOption) (*MCPCallResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(MCPCallResponse)
err := c.cc.Invoke(ctx, Yao_MCPCallTool_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
func (c *yaoClient) MCPListResources(ctx context.Context, in *MCPListRequest, opts ...grpc.CallOption) (*MCPResourcesResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(MCPResourcesResponse)
err := c.cc.Invoke(ctx, Yao_MCPListResources_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
func (c *yaoClient) MCPReadResource(ctx context.Context, in *MCPResourceRequest, opts ...grpc.CallOption) (*MCPResourceResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(MCPResourceResponse)
err := c.cc.Invoke(ctx, Yao_MCPReadResource_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
func (c *yaoClient) ChatCompletions(ctx context.Context, in *ChatRequest, opts ...grpc.CallOption) (*ChatResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(ChatResponse)
err := c.cc.Invoke(ctx, Yao_ChatCompletions_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
func (c *yaoClient) ChatCompletionsStream(ctx context.Context, in *ChatRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ChatChunk], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &Yao_ServiceDesc.Streams[2], Yao_ChatCompletionsStream_FullMethodName, cOpts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[ChatRequest, ChatChunk]{ClientStream: stream}
if err := x.ClientStream.SendMsg(in); err != nil {
return nil, err
}
if err := x.ClientStream.CloseSend(); err != nil {
return nil, err
}
return x, nil
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type Yao_ChatCompletionsStreamClient = grpc.ServerStreamingClient[ChatChunk]
func (c *yaoClient) AgentStream(ctx context.Context, in *AgentRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[AgentChunk], error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
stream, err := c.cc.NewStream(ctx, &Yao_ServiceDesc.Streams[3], Yao_AgentStream_FullMethodName, cOpts...)
if err != nil {
return nil, err
}
x := &grpc.GenericClientStream[AgentRequest, AgentChunk]{ClientStream: stream}
if err := x.ClientStream.SendMsg(in); err != nil {
return nil, err
}
if err := x.ClientStream.CloseSend(); err != nil {
return nil, err
}
return x, nil
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type Yao_AgentStreamClient = grpc.ServerStreamingClient[AgentChunk]
func (c *yaoClient) Healthz(ctx context.Context, in *Empty, opts ...grpc.CallOption) (*HealthzResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(HealthzResponse)
err := c.cc.Invoke(ctx, Yao_Healthz_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
// YaoServer is the server API for Yao service.
// All implementations must embed UnimplementedYaoServer
// for forward compatibility.
//
// Yao gRPC gateway. Shares OAuth + ACL scope system with openapi.
type YaoServer interface {
// Base
Run(context.Context, *RunRequest) (*RunResponse, error)
Stream(*RunRequest, grpc.ServerStreamingServer[Chunk]) error
Shell(context.Context, *ShellRequest) (*ShellResponse, error)
ShellStream(*ShellRequest, grpc.ServerStreamingServer[Chunk]) error
// API gateway
API(context.Context, *APIRequest) (*APIResponse, error)
// MCP
MCPListTools(context.Context, *MCPListRequest) (*MCPListResponse, error)
MCPCallTool(context.Context, *MCPCallRequest) (*MCPCallResponse, error)
MCPListResources(context.Context, *MCPListRequest) (*MCPResourcesResponse, error)
MCPReadResource(context.Context, *MCPResourceRequest) (*MCPResourceResponse, error)
// AI - LLM
ChatCompletions(context.Context, *ChatRequest) (*ChatResponse, error)
ChatCompletionsStream(*ChatRequest, grpc.ServerStreamingServer[ChatChunk]) error
// AI - Agent
AgentStream(*AgentRequest, grpc.ServerStreamingServer[AgentChunk]) error
// Health
Healthz(context.Context, *Empty) (*HealthzResponse, error)
mustEmbedUnimplementedYaoServer()
}
// UnimplementedYaoServer 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 UnimplementedYaoServer struct{}
func (UnimplementedYaoServer) Run(context.Context, *RunRequest) (*RunResponse, error) {
return nil, status.Error(codes.Unimplemented, "method Run not implemented")
}
func (UnimplementedYaoServer) Stream(*RunRequest, grpc.ServerStreamingServer[Chunk]) error {
return status.Error(codes.Unimplemented, "method Stream not implemented")
}
func (UnimplementedYaoServer) Shell(context.Context, *ShellRequest) (*ShellResponse, error) {
return nil, status.Error(codes.Unimplemented, "method Shell not implemented")
}
func (UnimplementedYaoServer) ShellStream(*ShellRequest, grpc.ServerStreamingServer[Chunk]) error {
return status.Error(codes.Unimplemented, "method ShellStream not implemented")
}
func (UnimplementedYaoServer) API(context.Context, *APIRequest) (*APIResponse, error) {
return nil, status.Error(codes.Unimplemented, "method API not implemented")
}
func (UnimplementedYaoServer) MCPListTools(context.Context, *MCPListRequest) (*MCPListResponse, error) {
return nil, status.Error(codes.Unimplemented, "method MCPListTools not implemented")
}
func (UnimplementedYaoServer) MCPCallTool(context.Context, *MCPCallRequest) (*MCPCallResponse, error) {
return nil, status.Error(codes.Unimplemented, "method MCPCallTool not implemented")
}
func (UnimplementedYaoServer) MCPListResources(context.Context, *MCPListRequest) (*MCPResourcesResponse, error) {
return nil, status.Error(codes.Unimplemented, "method MCPListResources not implemented")
}
func (UnimplementedYaoServer) MCPReadResource(context.Context, *MCPResourceRequest) (*MCPResourceResponse, error) {
return nil, status.Error(codes.Unimplemented, "method MCPReadResource not implemented")
}
func (UnimplementedYaoServer) ChatCompletions(context.Context, *ChatRequest) (*ChatResponse, error) {
return nil, status.Error(codes.Unimplemented, "method ChatCompletions not implemented")
}
func (UnimplementedYaoServer) ChatCompletionsStream(*ChatRequest, grpc.ServerStreamingServer[ChatChunk]) error {
return status.Error(codes.Unimplemented, "method ChatCompletionsStream not implemented")
}
func (UnimplementedYaoServer) AgentStream(*AgentRequest, grpc.ServerStreamingServer[AgentChunk]) error {
return status.Error(codes.Unimplemented, "method AgentStream not implemented")
}
func (UnimplementedYaoServer) Healthz(context.Context, *Empty) (*HealthzResponse, error) {
return nil, status.Error(codes.Unimplemented, "method Healthz not implemented")
}
func (UnimplementedYaoServer) mustEmbedUnimplementedYaoServer() {}
func (UnimplementedYaoServer) testEmbeddedByValue() {}
// UnsafeYaoServer may be embedded to opt out of forward compatibility for this service.
// Use of this interface is not recommended, as added methods to YaoServer will
// result in compilation errors.
type UnsafeYaoServer interface {
mustEmbedUnimplementedYaoServer()
}
func RegisterYaoServer(s grpc.ServiceRegistrar, srv YaoServer) {
// If the following call panics, it indicates UnimplementedYaoServer 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(&Yao_ServiceDesc, srv)
}
func _Yao_Run_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(RunRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(YaoServer).Run(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: Yao_Run_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(YaoServer).Run(ctx, req.(*RunRequest))
}
return interceptor(ctx, in, info, handler)
}
func _Yao_Stream_Handler(srv interface{}, stream grpc.ServerStream) error {
m := new(RunRequest)
if err := stream.RecvMsg(m); err != nil {
return err
}
return srv.(YaoServer).Stream(m, &grpc.GenericServerStream[RunRequest, Chunk]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type Yao_StreamServer = grpc.ServerStreamingServer[Chunk]
func _Yao_Shell_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(ShellRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(YaoServer).Shell(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: Yao_Shell_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(YaoServer).Shell(ctx, req.(*ShellRequest))
}
return interceptor(ctx, in, info, handler)
}
func _Yao_ShellStream_Handler(srv interface{}, stream grpc.ServerStream) error {
m := new(ShellRequest)
if err := stream.RecvMsg(m); err != nil {
return err
}
return srv.(YaoServer).ShellStream(m, &grpc.GenericServerStream[ShellRequest, Chunk]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type Yao_ShellStreamServer = grpc.ServerStreamingServer[Chunk]
func _Yao_API_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(APIRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(YaoServer).API(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: Yao_API_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(YaoServer).API(ctx, req.(*APIRequest))
}
return interceptor(ctx, in, info, handler)
}
func _Yao_MCPListTools_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(MCPListRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(YaoServer).MCPListTools(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: Yao_MCPListTools_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(YaoServer).MCPListTools(ctx, req.(*MCPListRequest))
}
return interceptor(ctx, in, info, handler)
}
func _Yao_MCPCallTool_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(MCPCallRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(YaoServer).MCPCallTool(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: Yao_MCPCallTool_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(YaoServer).MCPCallTool(ctx, req.(*MCPCallRequest))
}
return interceptor(ctx, in, info, handler)
}
func _Yao_MCPListResources_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(MCPListRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(YaoServer).MCPListResources(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: Yao_MCPListResources_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(YaoServer).MCPListResources(ctx, req.(*MCPListRequest))
}
return interceptor(ctx, in, info, handler)
}
func _Yao_MCPReadResource_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(MCPResourceRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(YaoServer).MCPReadResource(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: Yao_MCPReadResource_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(YaoServer).MCPReadResource(ctx, req.(*MCPResourceRequest))
}
return interceptor(ctx, in, info, handler)
}
func _Yao_ChatCompletions_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(ChatRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(YaoServer).ChatCompletions(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: Yao_ChatCompletions_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(YaoServer).ChatCompletions(ctx, req.(*ChatRequest))
}
return interceptor(ctx, in, info, handler)
}
func _Yao_ChatCompletionsStream_Handler(srv interface{}, stream grpc.ServerStream) error {
m := new(ChatRequest)
if err := stream.RecvMsg(m); err != nil {
return err
}
return srv.(YaoServer).ChatCompletionsStream(m, &grpc.GenericServerStream[ChatRequest, ChatChunk]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type Yao_ChatCompletionsStreamServer = grpc.ServerStreamingServer[ChatChunk]
func _Yao_AgentStream_Handler(srv interface{}, stream grpc.ServerStream) error {
m := new(AgentRequest)
if err := stream.RecvMsg(m); err != nil {
return err
}
return srv.(YaoServer).AgentStream(m, &grpc.GenericServerStream[AgentRequest, AgentChunk]{ServerStream: stream})
}
// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
type Yao_AgentStreamServer = grpc.ServerStreamingServer[AgentChunk]
func _Yao_Healthz_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(Empty)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(YaoServer).Healthz(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: Yao_Healthz_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(YaoServer).Healthz(ctx, req.(*Empty))
}
return interceptor(ctx, in, info, handler)
}
// Yao_ServiceDesc is the grpc.ServiceDesc for Yao service.
// It's only intended for direct use with grpc.RegisterService,
// and not to be introspected or modified (even as a copy)
var Yao_ServiceDesc = grpc.ServiceDesc{
ServiceName: "yao.Yao",
HandlerType: (*YaoServer)(nil),
Methods: []grpc.MethodDesc{
{
MethodName: "Run",
Handler: _Yao_Run_Handler,
},
{
MethodName: "Shell",
Handler: _Yao_Shell_Handler,
},
{
MethodName: "API",
Handler: _Yao_API_Handler,
},
{
MethodName: "MCPListTools",
Handler: _Yao_MCPListTools_Handler,
},
{
MethodName: "MCPCallTool",
Handler: _Yao_MCPCallTool_Handler,
},
{
MethodName: "MCPListResources",
Handler: _Yao_MCPListResources_Handler,
},
{
MethodName: "MCPReadResource",
Handler: _Yao_MCPReadResource_Handler,
},
{
MethodName: "ChatCompletions",
Handler: _Yao_ChatCompletions_Handler,
},
{
MethodName: "Healthz",
Handler: _Yao_Healthz_Handler,
},
},
Streams: []grpc.StreamDesc{
{
StreamName: "Stream",
Handler: _Yao_Stream_Handler,
ServerStreams: true,
},
{
StreamName: "ShellStream",
Handler: _Yao_ShellStream_Handler,
ServerStreams: true,
},
{
StreamName: "ChatCompletionsStream",
Handler: _Yao_ChatCompletionsStream_Handler,
ServerStreams: true,
},
{
StreamName: "AgentStream",
Handler: _Yao_AgentStream_Handler,
ServerStreams: true,
},
},
Metadata: "yao.proto",
}

79
grpc/run/run.go Normal file
View file

@ -0,0 +1,79 @@
package run
import (
"context"
"encoding/json"
"time"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/yao/grpc/auth"
"github.com/yaoapp/yao/grpc/pb"
)
// Handler implements the Run gRPC method.
type Handler struct{}
// Run executes a Yao process by name and returns the JSON-encoded result.
func (h *Handler) Run(ctx context.Context, req *pb.RunRequest) (*pb.RunResponse, error) {
if req.Process == "" {
return nil, status.Error(codes.InvalidArgument, "process name is required")
}
if req.Timeout > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, time.Duration(req.Timeout)*time.Second)
defer cancel()
}
var args []interface{}
if len(req.Args) > 0 {
if err := json.Unmarshal(req.Args, &args); err != nil {
return nil, status.Errorf(codes.InvalidArgument, "invalid args JSON: %v", err)
}
}
p, err := process.Of(req.Process, args...)
if err != nil {
return nil, status.Errorf(codes.Internal, "process error: %v", err)
}
p.WithContext(ctx)
injectAuth(p, ctx)
if err := p.Execute(); err != nil {
if ctx.Err() == context.DeadlineExceeded {
return nil, status.Error(codes.DeadlineExceeded, "process execution timed out")
}
return nil, status.Errorf(codes.Internal, "process execution failed: %v", err)
}
defer p.Release()
val := p.Value()
data, err := json.Marshal(val)
if err != nil {
return nil, status.Errorf(codes.Internal, "failed to marshal result: %v", err)
}
return &pb.RunResponse{Data: data}, nil
}
// injectAuth propagates AuthorizedInfo from the gRPC context into the Process.
func injectAuth(p *process.Process, ctx context.Context) {
authInfo := auth.GetAuthorizedInfo(ctx)
if authInfo == nil {
return
}
p.WithSID(authInfo.SessionID)
p.WithAuthorized(&process.AuthorizedInfo{
Subject: authInfo.Subject,
ClientID: authInfo.ClientID,
Scope: authInfo.Scope,
SessionID: authInfo.SessionID,
UserID: authInfo.UserID,
TeamID: authInfo.TeamID,
TenantID: authInfo.TenantID,
})
}

155
grpc/run/run_test.go Normal file
View file

@ -0,0 +1,155 @@
package run_test
import (
"context"
"encoding/json"
"testing"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/yaoapp/yao/grpc/pb"
"github.com/yaoapp/yao/grpc/tests/testutils"
)
func TestRun_ProcessExec(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.Run(ctx, &pb.RunRequest{
Process: "utils.app.Ping",
})
assert.NoError(t, err)
assert.NotNil(t, resp)
assert.NotEmpty(t, resp.Data)
}
func TestRun_WithArgs(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
args, _ := json.Marshal([]interface{}{"hello", " world"})
resp, err := client.Run(ctx, &pb.RunRequest{
Process: "utils.str.Concat",
Args: args,
})
assert.NoError(t, err)
if assert.NotNil(t, resp) {
assert.NotEmpty(t, resp.Data)
var result string
err = json.Unmarshal(resp.Data, &result)
assert.NoError(t, err)
assert.Equal(t, "hello world", result)
}
}
func TestRun_InvalidProcess(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.Run(ctx, &pb.RunRequest{Process: "nonexistent.process.here"})
assert.Error(t, err)
}
func TestRun_EmptyProcessName(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.Run(ctx, &pb.RunRequest{Process: ""})
assert.Error(t, err)
}
func TestRun_BadArgsJSON(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.Run(ctx, &pb.RunRequest{
Process: "utils.app.Ping",
Args: []byte("{not-json"),
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestRun_WithTimeout(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.Run(ctx, &pb.RunRequest{
Process: "utils.app.Ping",
Timeout: 30,
})
assert.NoError(t, err)
assert.NotNil(t, resp)
}
func TestRun_EmptyProcessName_StatusCode(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.Run(ctx, &pb.RunRequest{Process: ""})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestRun_InvalidProcess_StatusCode(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.Run(ctx, &pb.RunRequest{Process: "nonexistent.process.here"})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.Internal, st.Code())
}
func TestRun_NilArgs(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:run")
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.Run(ctx, &pb.RunRequest{
Process: "utils.app.Ping",
Args: nil,
})
assert.NoError(t, err)
assert.NotNil(t, resp)
}

91
grpc/shell/shell.go Normal file
View file

@ -0,0 +1,91 @@
package shell
import (
"bytes"
"context"
"errors"
"os"
"os/exec"
"syscall"
"time"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/yaoapp/yao/grpc/pb"
)
const (
defaultTimeout = 30 * time.Second
maxTimeout = 300 * time.Second
)
// Handler implements the Shell gRPC method.
type Handler struct{}
// Shell executes a system command in the host process and returns stdout/stderr/exit code.
func (h *Handler) Shell(ctx context.Context, req *pb.ShellRequest) (*pb.ShellResponse, error) {
if os.Getuid() == 0 {
return nil, status.Error(codes.PermissionDenied, "shell execution refused when running as root")
}
if req.Command == "" {
return nil, status.Error(codes.InvalidArgument, "command is required")
}
timeout := defaultTimeout
if req.Timeout > 0 {
timeout = time.Duration(req.Timeout) * time.Second
if timeout > maxTimeout {
timeout = maxTimeout
}
}
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
cmd := exec.CommandContext(ctx, req.Command, req.Args...)
if len(req.Env) > 0 {
env := os.Environ()
for k, v := range req.Env {
env = append(env, k+"="+v)
}
cmd.Env = env
}
var stdout, stderr bytes.Buffer
cmd.Stdout = &stdout
cmd.Stderr = &stderr
err := cmd.Run()
resp := &pb.ShellResponse{
Stdout: stdout.Bytes(),
Stderr: stderr.Bytes(),
ExitCode: 0,
}
if err != nil {
if ctx.Err() == context.DeadlineExceeded {
return nil, status.Error(codes.DeadlineExceeded, "command timed out")
}
var exitErr *exec.ExitError
if errors.As(err, &exitErr) {
if ws, ok := exitErr.Sys().(syscall.WaitStatus); ok {
resp.ExitCode = int32(ws.ExitStatus())
} else {
resp.ExitCode = int32(exitErr.ExitCode())
}
return resp, nil
}
if errors.Is(err, exec.ErrNotFound) {
return nil, status.Errorf(codes.NotFound, "command not found: %s", req.Command)
}
return nil, status.Errorf(codes.Internal, "command execution failed: %v", err)
}
return resp, nil
}

166
grpc/shell/shell_test.go Normal file
View file

@ -0,0 +1,166 @@
package shell_test
import (
"context"
"runtime"
"testing"
"github.com/stretchr/testify/assert"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"github.com/yaoapp/yao/grpc/pb"
"github.com/yaoapp/yao/grpc/tests/testutils"
)
func TestShell_Echo(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:shell")
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.Shell(ctx, &pb.ShellRequest{
Command: "echo",
Args: []string{"hello"},
})
assert.NoError(t, err)
assert.NotNil(t, resp)
assert.Contains(t, string(resp.Stdout), "hello")
assert.Equal(t, int32(0), resp.ExitCode)
}
func TestShell_CommandNotFound(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:shell")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.Shell(ctx, &pb.ShellRequest{
Command: "this_command_does_not_exist_xyz",
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.NotFound, st.Code())
}
func TestShell_Timeout(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("sleep command not available on Windows")
}
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:shell")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.Shell(ctx, &pb.ShellRequest{
Command: "sleep",
Args: []string{"10"},
Timeout: 1,
})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.DeadlineExceeded, st.Code())
}
func TestShell_EmptyCommand(t *testing.T) {
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:shell")
ctx := testutils.WithToken(context.Background(), token)
_, err := client.Shell(ctx, &pb.ShellRequest{Command: ""})
assert.Error(t, err)
st, _ := status.FromError(err)
assert.Equal(t, codes.InvalidArgument, st.Code())
}
func TestShell_NonZeroExit(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("false command not available on Windows")
}
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:shell")
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.Shell(ctx, &pb.ShellRequest{
Command: "false",
})
assert.NoError(t, err)
assert.NotEqual(t, int32(0), resp.ExitCode)
}
func TestShell_WithEnv(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("printenv not available on Windows")
}
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:shell")
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.Shell(ctx, &pb.ShellRequest{
Command: "printenv",
Args: []string{"TEST_GRPC_VAR"},
Env: map[string]string{"TEST_GRPC_VAR": "grpc_value"},
})
assert.NoError(t, err)
assert.Contains(t, string(resp.Stdout), "grpc_value")
}
func TestShell_MaxTimeoutCapped(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("echo not available on Windows")
}
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:shell")
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.Shell(ctx, &pb.ShellRequest{
Command: "echo",
Args: []string{"ok"},
Timeout: 9999,
})
assert.NoError(t, err)
assert.Contains(t, string(resp.Stdout), "ok")
}
func TestShell_Stderr(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("bash not available on Windows")
}
conn := testutils.Prepare(t)
defer testutils.Clean()
client := testutils.NewClient(conn)
token := testutils.ObtainAccessToken(t, "grpc:shell")
ctx := testutils.WithToken(context.Background(), token)
resp, err := client.Shell(ctx, &pb.ShellRequest{
Command: "bash",
Args: []string{"-c", "echo error_msg >&2"},
})
assert.NoError(t, err)
assert.Contains(t, string(resp.Stderr), "error_msg")
assert.Equal(t, int32(0), resp.ExitCode)
}

242
grpc/tests/testutils/testutils.go Normal file
View file

@ -0,0 +1,242 @@
package testutils
import (
"context"
"net"
"os"
"strings"
"testing"
"github.com/gin-gonic/gin"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/metadata"
gouapi "github.com/yaoapp/gou/api"
"github.com/yaoapp/gou/model"
"github.com/yaoapp/gou/query"
"github.com/yaoapp/gou/query/gou"
"github.com/yaoapp/xun/capsule"
yaoagent "github.com/yaoapp/yao/agent"
"github.com/yaoapp/yao/agent/caller"
agentllm "github.com/yaoapp/yao/agent/llm"
"github.com/yaoapp/yao/config"
yaogrpc "github.com/yaoapp/yao/grpc"
_ "github.com/yaoapp/yao/grpc/auth"
"github.com/yaoapp/yao/grpc/pb"
"github.com/yaoapp/yao/kb"
"github.com/yaoapp/yao/openapi"
"github.com/yaoapp/yao/openapi/oauth"
"github.com/yaoapp/yao/service"
"github.com/yaoapp/yao/test"
_ "github.com/yaoapp/gou/encoding"
_ "github.com/yaoapp/gou/text"
_ "github.com/yaoapp/yao/agent/assistant"
)
// Prepare initializes the Yao runtime (DB, V8, models, stores, scripts),
// loads the OpenAPI server (which bootstraps oauth.OAuth and acl.Global),
// sets up the HTTP router for API proxy tests,
// then starts a real gRPC server on a random port.
// Returns a connected grpc.ClientConn ready to create service clients.
func Prepare(t *testing.T) *grpc.ClientConn {
t.Helper()
cfg := config.Conf
cfg.GRPC.Port = 0
cfg.GRPC.Host = "0.0.0.0"
cfg.GRPC.Enabled = ""
test.Prepare(t, config.Conf)
if openapi.Server == nil {
if _, err := openapi.Load(config.Conf); err != nil {
t.Fatalf("failed to load OpenAPI server: %v", err)
}
}
// Load KB (required for agent KB features).
if _, err := kb.Load(config.Conf); err != nil {
t.Logf("warning: failed to load KB: %v", err)
}
// Load agent DSL (required for AgentStream handler).
if yaoagent.GetAgent() == nil {
if err := yaoagent.Load(config.Conf); err != nil {
t.Logf("warning: failed to load agent DSL: %v", err)
}
}
// Register JSAPI factories (idempotent, needed because Go init order is not guaranteed).
caller.SetJSAPIFactory()
agentllm.SetJSAPIFactory()
// Register default query engine (required for DB search).
if _, has := query.Engines["default"]; !has && capsule.Global != nil {
query.Register("default", &gou.Query{
Query: capsule.Query(),
GetTableName: func(s string) string {
if mod, has := model.Models[s]; has {
return mod.MetaData.Table.Name
}
return s
},
AESKey: config.Conf.DB.AESKey,
})
}
// Set up the HTTP router so grpc/api can proxy requests internally.
if service.Router == nil {
router := gin.New()
if openapi.Server != nil {
gouapi.SetRoutes(router, openapi.Server.Config.BaseURL)
gouapi.BuildRouteTable()
openapi.Server.Attach(router)
}
service.Router = router
}
if err := yaogrpc.StartServer(cfg); err != nil {
t.Fatalf("failed to start gRPC server: %v", err)
}
addrs := yaogrpc.Addr()
if len(addrs) == 0 {
t.Fatal("gRPC server has no listen address")
}
conn, err := grpc.NewClient(addrs[0], grpc.WithTransportCredentials(insecure.NewCredentials()))
if err != nil {
t.Fatalf("failed to dial gRPC server: %v", err)
}
return conn
}
// Clean stops the gRPC server and tears down the Yao runtime.
func Clean() {
yaogrpc.Stop()
service.Router = nil
openapi.Server = nil
test.Clean()
}
// Addr returns the gRPC server listen address.
func Addr() string {
addrs := yaogrpc.Addr()
if len(addrs) == 0 {
return ""
}
return addrs[0]
}
// RelayAddr returns the gRPC address reachable from a Docker container.
// When TAI_TEST_HOST_IP is set (e.g. to the docker bridge gateway),
// it replaces the host portion so that the Tai container can reach the
// Yao gRPC server running on the CI host.
func RelayAddr() string {
addr := Addr()
if addr == "" {
return ""
}
hostIP := os.Getenv("TAI_TEST_HOST_IP")
if hostIP == "" {
return addr
}
_, port, err := net.SplitHostPort(addr)
if err != nil {
return addr
}
return hostIP + ":" + port
}
// ObtainAccessToken mints a token with the given scopes via oauth.MakeAccessToken.
func ObtainAccessToken(t *testing.T, scopes ...string) string {
t.Helper()
svc := oauth.OAuth
if svc == nil {
t.Fatal("oauth service not initialized")
}
scope := strings.Join(scopes, " ")
token, err := svc.MakeAccessToken("grpc-test", scope, "test-user", 3600)
if err != nil {
t.Fatalf("failed to make access token: %v", err)
}
return token
}
// ObtainAccessTokenForUser mints a token for a specific user ID.
func ObtainAccessTokenForUser(t *testing.T, userID string, scopes ...string) string {
t.Helper()
svc := oauth.OAuth
if svc == nil {
t.Fatal("oauth service not initialized")
}
scope := strings.Join(scopes, " ")
token, err := svc.MakeAccessToken("grpc-test", scope, userID, 3600)
if err != nil {
t.Fatalf("failed to make access token: %v", err)
}
return token
}
// ObtainExpiredAccessToken mints an already-expired token (TTL=1s already elapsed).
func ObtainExpiredAccessToken(t *testing.T, scopes ...string) string {
t.Helper()
svc := oauth.OAuth
if svc == nil {
t.Fatal("oauth service not initialized")
}
scope := strings.Join(scopes, " ")
token, err := svc.MakeAccessToken("grpc-test", scope, "test-user", -1)
if err != nil {
t.Fatalf("failed to make expired access token: %v", err)
}
return token
}
// ObtainRefreshToken mints a refresh token.
func ObtainRefreshToken(t *testing.T, scopes ...string) string {
t.Helper()
svc := oauth.OAuth
if svc == nil {
t.Fatal("oauth service not initialized")
}
scope := strings.Join(scopes, " ")
token, err := svc.MakeRefreshToken("grpc-test", scope, "test-user", 0)
if err != nil {
t.Fatalf("failed to make refresh token: %v", err)
}
return token
}
// WithToken attaches a Bearer token to the context via gRPC metadata.
func WithToken(ctx context.Context, token string) context.Context {
return metadata.AppendToOutgoingContext(ctx, "authorization", "Bearer "+token)
}
// WithRefreshToken attaches both Bearer and x-refresh-token to the context.
func WithRefreshToken(ctx context.Context, token, refreshToken string) context.Context {
return metadata.AppendToOutgoingContext(ctx,
"authorization", "Bearer "+token,
"x-refresh-token", refreshToken,
)
}
// WithSandboxMetadata attaches x-sandbox-id and x-grpc-upstream metadata.
func WithSandboxMetadata(ctx context.Context, sandboxID, upstream string) context.Context {
return metadata.AppendToOutgoingContext(ctx,
"x-sandbox-id", sandboxID,
"x-grpc-upstream", upstream,
)
}
// NewClient creates a pb.YaoClient from a connection.
func NewClient(conn *grpc.ClientConn) pb.YaoClient {
return pb.NewYaoClient(conn)
}

1
openapi/config.go
View file

@ -388,6 +388,7 @@ func (config *Config) OAuthConfig(appConfig config.Config) (*oauth.Config, error
Cache: cacheStore,
Store: dataStore,
IssuerURL: config.OAuth.IssuerURL,
BaseURL: config.BaseURL,
Signing: signingConfig, // Use the converted signing config
Token: config.OAuth.Token,
Security: config.OAuth.Security,

113
openapi/oauth.go
View file

@ -58,6 +58,7 @@ func (openapi *OpenAPI) attachOAuth(base *gin.RouterGroup) {
// Device Authorization Flow - RFC 8628
oauth.POST("/device_authorization", openapi.oauthDeviceAuthorization)
oauth.POST("/device/authorize", openapi.oauthDeviceAuthorize)
// Pushed Authorization Request - RFC 9126
oauth.POST("/par", openapi.oauthPushedAuthorizationRequest)
@ -523,22 +524,101 @@ func (openapi *OpenAPI) oauthDeleteClient(c *gin.Context) {
// oauthDeviceAuthorization handles device authorization - RFC 8628
func (openapi *OpenAPI) oauthDeviceAuthorization(c *gin.Context) {
clientID := c.PostForm("client_id")
if clientID == "" {
response.RespondWithError(c, response.StatusBadRequest, response.ErrInvalidRequest)
response.RespondWithSecureError(c, response.StatusBadRequest, response.ErrInvalidRequest)
return
}
// TODO: Implement device authorization logic
deviceResponse := &response.DeviceAuthorizationResponse{
DeviceCode: "generated-device-code",
UserCode: "USER-CODE",
VerificationURI: "https://example.com/device",
ExpiresIn: 900, // 15 minutes
Interval: 5, // 5 seconds
scope := c.PostForm("scope")
oauthService := openapi.OAuth
res, err := oauthService.DeviceAuthorization(c, clientID, scope)
if err != nil {
if oauthErr, ok := err.(*response.ErrorResponse); ok {
response.RespondWithSecureError(c, response.StatusBadRequest, oauthErr)
} else {
response.RespondWithSecureError(c, response.StatusBadRequest, response.ErrInvalidRequest)
}
return
}
response.RespondWithSuccess(c, response.StatusOK, deviceResponse)
response.RespondWithSecureSuccess(c, response.StatusOK, res)
}
// oauthDeviceAuthorize allows an authenticated user to authorize a pending device code.
func (openapi *OpenAPI) oauthDeviceAuthorize(c *gin.Context) {
tokenStr := extractBearerToken(c)
if tokenStr == "" {
response.RespondWithSecureError(c, response.StatusUnauthorized, &response.ErrorResponse{
Code: types.ErrorInvalidGrant,
ErrorDescription: "Bearer token required",
})
return
}
svc, ok := openapi.OAuth.(*oauth.Service)
if !ok {
response.RespondWithSecureError(c, response.StatusInternalServerError, &response.ErrorResponse{
Code: types.ErrorServerError,
ErrorDescription: "OAuth service unavailable",
})
return
}
tokenClaims, err := svc.VerifyToken(tokenStr)
if err != nil || tokenClaims == nil {
response.RespondWithSecureError(c, response.StatusUnauthorized, &response.ErrorResponse{
Code: types.ErrorInvalidGrant,
ErrorDescription: "Invalid or expired token",
})
return
}
if tokenClaims.Subject == "" {
response.RespondWithSecureError(c, response.StatusUnauthorized, &response.ErrorResponse{
Code: types.ErrorInvalidGrant,
ErrorDescription: "Token has no subject",
})
return
}
extraClaims := tokenClaims.Extra
if extraClaims == nil {
extraClaims = make(map[string]interface{})
}
if tokenClaims.TeamID != "" {
extraClaims["team_id"] = tokenClaims.TeamID
}
if tokenClaims.TenantID != "" {
extraClaims["tenant_id"] = tokenClaims.TenantID
}
userCode := c.PostForm("user_code")
if userCode == "" {
userCode = c.Query("user_code")
}
if userCode == "" {
var body struct {
UserCode string `json:"user_code"`
}
if c.ShouldBindJSON(&body) == nil {
userCode = body.UserCode
}
}
if userCode == "" {
response.RespondWithSecureError(c, response.StatusBadRequest, response.ErrInvalidRequest)
return
}
if err := svc.AuthorizeDevice(c, userCode, tokenClaims.Subject, extraClaims); err != nil {
if oauthErr, ok := err.(*response.ErrorResponse); ok {
response.RespondWithSecureError(c, response.StatusBadRequest, oauthErr)
} else {
response.RespondWithSecureError(c, response.StatusBadRequest, response.ErrInvalidGrant)
}
return
}
response.RespondWithSecureSuccess(c, response.StatusOK, map[string]string{"status": "authorized"})
}
// oauthPushedAuthorizationRequest handles PAR - RFC 9126
@ -640,3 +720,16 @@ func (openapi *OpenAPI) getParam(c *gin.Context, key string) string {
// Then try to get from POST form data (POST request)
return c.PostForm(key)
}
// extractBearerToken reads the access token from Authorization header or cookie,
// matching the same logic as guard.getAccessToken.
func extractBearerToken(c *gin.Context) string {
if auth := c.GetHeader("Authorization"); strings.HasPrefix(auth, "Bearer ") {
return strings.TrimPrefix(auth, "Bearer ")
}
cookieName := response.GetCookieName("access_token")
if cookie, err := c.Cookie(cookieName); err == nil && cookie != "" {
return strings.TrimPrefix(cookie, "Bearer ")
}
return ""
}

207
openapi/oauth/authenticate.go Normal file
View file

@ -0,0 +1,207 @@
package oauth
import (
"errors"
"fmt"
"strings"
"time"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/openapi/oauth/types"
)
// AuthInput contains the raw tokens extracted from the transport layer
// (HTTP headers/cookies or gRPC metadata). No framework dependency.
type AuthInput struct {
AccessToken string
RefreshToken string
SessionID string
}
// AuthResult holds the outcome of a successful authentication.
type AuthResult struct {
Claims *types.TokenClaims
Info *types.AuthorizedInfo
NewAccessToken string // non-empty when token refresh occurred
NewRefreshToken string // non-empty when token refresh occurred
}
// AuthenticateToken performs token verification and optional refresh
// without any gin/HTTP dependency. The caller is responsible for
// extracting tokens from the transport and delivering refreshed tokens
// back to the client.
func (s *Service) AuthenticateToken(input AuthInput) (*AuthResult, error) {
token := input.AccessToken
// API Key resolution (same as getAccessToken in guard.go)
if s.isAPIKey(token) {
token = s.getAccessTokenFromAPIKey(token)
}
token = strings.TrimPrefix(token, "Bearer ")
if token == "" {
return nil, fmt.Errorf("%s", types.ErrTokenMissing.Error())
}
var newAccessToken, newRefreshToken string
claims, err := s.VerifyToken(token)
if err != nil {
expiredClaims, expErr := s.VerifyTokenAllowExpired(token)
if expErr != nil || expiredClaims == nil {
return nil, fmt.Errorf("%s", types.ErrInvalidToken.Error())
}
if !expiredClaims.ExpiresAt.IsZero() && expiredClaims.ExpiresAt.Before(time.Now()) {
newClaims, access, refresh, refreshErr := s.refreshTokenDirect(input.RefreshToken, expiredClaims)
if refreshErr != nil {
if errors.Is(refreshErr, errRefreshInProgress) || errors.Is(refreshErr, errRefreshAlreadyDone) {
claims = expiredClaims
} else {
log.Error("[OAuth] Token refresh failed: %v", refreshErr)
return nil, fmt.Errorf("%s", types.ErrInvalidRefreshToken.Error())
}
} else {
claims = newClaims
newAccessToken = access
newRefreshToken = refresh
}
} else {
return nil, fmt.Errorf("%s", types.ErrInvalidToken.Error())
}
}
info := s.buildAuthInfo(claims, input.SessionID)
return &AuthResult{
Claims: claims,
Info: info,
NewAccessToken: newAccessToken,
NewRefreshToken: newRefreshToken,
}, nil
}
// refreshTokenDirect performs token rotation without any gin/HTTP dependency.
// It shares the same refreshGates concurrency control as TryRefreshToken.
// Returns (newClaims, newAccessToken, newRefreshToken, error).
func (s *Service) refreshTokenDirect(refreshToken string, expiredClaims *types.TokenClaims) (*types.TokenClaims, string, string, error) {
if refreshToken == "" {
return nil, "", "", fmt.Errorf("refresh token missing")
}
gate := &refreshGate{done: make(chan struct{})}
if actual, loaded := refreshGates.LoadOrStore(refreshToken, gate); loaded {
existing := actual.(*refreshGate)
select {
case <-existing.done:
return nil, "", "", errRefreshAlreadyDone
default:
return nil, "", "", errRefreshInProgress
}
}
defer func() {
close(gate.done)
time.AfterFunc(30*time.Second, func() {
refreshGates.CompareAndDelete(refreshToken, gate)
})
}()
refreshClaims, err := s.VerifyRefreshToken(refreshToken)
if err != nil {
return nil, "", "", fmt.Errorf("invalid or expired refresh token: %w", err)
}
var accessTTL time.Duration
if expiredClaims != nil && !expiredClaims.IssuedAt.IsZero() && !expiredClaims.ExpiresAt.IsZero() {
accessTTL = expiredClaims.ExpiresAt.Sub(expiredClaims.IssuedAt)
}
if accessTTL <= 0 {
accessTTL = s.config.Token.AccessTokenLifetime
}
if accessTTL <= 0 {
accessTTL = time.Hour
}
sourceClaims := expiredClaims
if sourceClaims == nil {
sourceClaims = refreshClaims
}
extraClaims := sourceClaims.Extra
if extraClaims == nil {
extraClaims = make(map[string]interface{})
}
if sourceClaims.TeamID != "" {
extraClaims["team_id"] = sourceClaims.TeamID
}
if sourceClaims.TenantID != "" {
extraClaims["tenant_id"] = sourceClaims.TenantID
}
s.revokeRefreshToken(refreshToken)
var refreshRemainingSeconds int
if !refreshClaims.ExpiresAt.IsZero() {
refreshRemainingSeconds = int(time.Until(refreshClaims.ExpiresAt).Seconds())
if refreshRemainingSeconds <= 0 {
return nil, "", "", fmt.Errorf("refresh token already expired after revocation")
}
} else {
refreshTTL := s.config.Token.RefreshTokenLifetime
if refreshTTL == 0 {
refreshTTL = 24 * time.Hour
}
refreshRemainingSeconds = int(refreshTTL.Seconds())
}
newRefreshToken, err := s.MakeRefreshToken(
sourceClaims.ClientID,
sourceClaims.Scope,
sourceClaims.Subject,
refreshRemainingSeconds,
extraClaims,
)
if err != nil {
return nil, "", "", fmt.Errorf("failed to issue new refresh token: %w", err)
}
newTokenStr, err := s.MakeAccessToken(
sourceClaims.ClientID,
sourceClaims.Scope,
sourceClaims.Subject,
int(accessTTL.Seconds()),
extraClaims,
)
if err != nil {
return nil, "", "", fmt.Errorf("failed to issue access token: %w", err)
}
newClaims, err := s.VerifyToken(newTokenStr)
if err != nil {
return nil, "", "", fmt.Errorf("failed to verify refreshed token: %w", err)
}
log.Info("[OAuth] Token rotated for subject %s (access + refresh)", sourceClaims.Subject)
return newClaims, newTokenStr, newRefreshToken, nil
}
// buildAuthInfo constructs AuthorizedInfo directly from token claims,
// equivalent to the SetInfo+GetInfo round-trip through gin.Context.
func (s *Service) buildAuthInfo(claims *types.TokenClaims, sessionID string) *types.AuthorizedInfo {
info := &types.AuthorizedInfo{
Subject: claims.Subject,
ClientID: claims.ClientID,
Scope: claims.Scope,
SessionID: sessionID,
TeamID: claims.TeamID,
TenantID: claims.TenantID,
}
userID, err := s.UserID(claims.ClientID, claims.Subject)
if err == nil && userID != "" {
info.UserID = userID
}
return info
}

98
openapi/oauth/core.go
View file

@ -6,6 +6,7 @@ import (
"time"
"github.com/yaoapp/yao/openapi/oauth/types"
"go.mongodb.org/mongo-driver/bson/primitive"
)
// AuthorizationServer returns the authorization server endpoint URL
@ -131,6 +132,8 @@ func (s *Service) Token(ctx context.Context, grantType string, code string, clie
return s.handleClientCredentialsGrant(ctx, client)
case types.GrantTypeRefreshToken:
return s.handleRefreshTokenGrant(ctx, client, code) // code is refresh token in this case
case types.GrantTypeDeviceCode:
return s.handleDeviceCodeGrant(ctx, client, code) // code is device_code in this case
default:
return nil, &types.ErrorResponse{
Code: types.ErrorUnsupportedGrantType,
@ -615,3 +618,98 @@ func (s *Service) validatePKCE(ctx context.Context, client *types.ClientInfo, co
return nil
}
// handleDeviceCodeGrant handles the device_code grant type (RFC 8628 Section 3.4).
func (s *Service) handleDeviceCodeGrant(ctx context.Context, client *types.ClientInfo, deviceCode string) (*types.Token, error) {
if !s.config.Features.DeviceFlowEnabled {
return nil, &types.ErrorResponse{
Code: types.ErrorUnsupportedGrantType,
ErrorDescription: "Device flow is not enabled",
}
}
codeData, err := s.getDeviceCodeData(deviceCode)
if err != nil {
return nil, err
}
storedClientID, _ := codeData["client_id"].(string)
if storedClientID != client.ClientID {
return nil, &types.ErrorResponse{
Code: types.ErrorInvalidGrant,
ErrorDescription: "Device code was issued to a different client",
}
}
expiresAt, _ := codeData["expires_at"].(int64)
if expiresAt == 0 {
if f, ok := codeData["expires_at"].(float64); ok {
expiresAt = int64(f)
}
}
if expiresAt > 0 && time.Now().Unix() > expiresAt {
s.consumeDeviceCode(deviceCode)
return nil, &types.ErrorResponse{
Code: types.ErrorExpiredToken,
ErrorDescription: "Device code has expired",
}
}
status, _ := codeData["status"].(string)
switch status {
case "pending":
return nil, &types.ErrorResponse{
Code: types.ErrorAuthorizationPending,
ErrorDescription: "The authorization request is still pending",
}
case "authorized":
scope, _ := codeData["scope"].(string)
subject, _ := codeData["subject"].(string)
s.consumeDeviceCode(deviceCode)
var extraClaims map[string]interface{}
if ec, ok := codeData["extra_claims"]; ok {
switch v := ec.(type) {
case map[string]interface{}:
extraClaims = v
case primitive.M:
extraClaims = map[string]interface{}(v)
}
}
expiresIn := int(s.config.Token.AccessTokenLifetime.Seconds())
accessToken, err := s.generateAccessTokenWithScope(client.ClientID, scope, subject, expiresIn, extraClaims)
if err != nil {
return nil, &types.ErrorResponse{
Code: types.ErrorServerError,
ErrorDescription: "Failed to generate access token",
}
}
token := &types.Token{
AccessToken: accessToken,
TokenType: "Bearer",
ExpiresIn: expiresIn,
}
if types.Contains(client.GrantTypes, types.GrantTypeRefreshToken) {
refreshToken, err := s.generateRefreshToken(client.ClientID, scope, subject, 0, extraClaims)
if err != nil {
return nil, &types.ErrorResponse{
Code: types.ErrorServerError,
ErrorDescription: "Failed to generate refresh token",
}
}
token.RefreshToken = refreshToken
}
return token, nil
default:
return nil, &types.ErrorResponse{
Code: types.ErrorInvalidGrant,
ErrorDescription: "Invalid device code status",
}
}
}

116
openapi/oauth/device.go
View file

@ -2,13 +2,121 @@ package oauth
import (
"context"
"fmt"
"strings"
gonanoid "github.com/matoous/go-nanoid/v2"
"github.com/yaoapp/yao/openapi/oauth/types"
)
// DeviceAuthorization initiates the device authorization flow
// This is used for devices with limited input capabilities
const userCodeAlphabet = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
// DeviceAuthorization initiates the device authorization flow (RFC 8628).
func (s *Service) DeviceAuthorization(ctx context.Context, clientID string, scope string) (*types.DeviceAuthorizationResponse, error) {
// TODO: Implement device authorization flow
return nil, nil
if !s.config.Features.DeviceFlowEnabled {
return nil, &types.ErrorResponse{
Code: types.ErrorUnsupportedGrantType,
ErrorDescription: "Device flow is not enabled",
}
}
client, err := s.clientProvider.GetClientByID(ctx, clientID)
if err != nil || client == nil {
return nil, &types.ErrorResponse{
Code: types.ErrorInvalidClient,
ErrorDescription: "Invalid client",
}
}
if !clientSupportsGrantType(client, types.GrantTypeDeviceCode) {
return nil, &types.ErrorResponse{
Code: types.ErrorUnauthorizedClient,
ErrorDescription: "Client does not support device code grant",
}
}
deviceCode, err := s.generateToken("dc", clientID)
if err != nil {
return nil, &types.ErrorResponse{
Code: types.ErrorServerError,
ErrorDescription: "Failed to generate device code",
}
}
userCode, err := s.generateUserCode()
if err != nil {
return nil, &types.ErrorResponse{
Code: types.ErrorServerError,
ErrorDescription: "Failed to generate user code",
}
}
if err := s.storeDeviceCode(deviceCode, userCode, clientID, scope); err != nil {
return nil, &types.ErrorResponse{
Code: types.ErrorServerError,
ErrorDescription: "Failed to store device code",
}
}
verificationURI := fmt.Sprintf("%s/auth/device", s.config.IssuerURL)
verificationURIComplete := fmt.Sprintf("%s?user_code=%s", verificationURI, userCode)
return &types.DeviceAuthorizationResponse{
DeviceCode: deviceCode,
UserCode: userCode,
VerificationURI: verificationURI,
VerificationURIComplete: verificationURIComplete,
ExpiresIn: int(s.config.Token.DeviceCodeLifetime.Seconds()),
Interval: int(s.config.Token.DeviceCodeInterval.Seconds()),
}, nil
}
// AuthorizeDevice allows an authenticated user to authorize a device code via user_code.
func (s *Service) AuthorizeDevice(ctx context.Context, userCode string, subject string, extraClaims ...map[string]interface{}) error {
if !s.config.Features.DeviceFlowEnabled {
return &types.ErrorResponse{
Code: types.ErrorUnsupportedGrantType,
ErrorDescription: "Device flow is not enabled",
}
}
normalized := strings.ToUpper(strings.ReplaceAll(userCode, "-", ""))
formatted := normalized
if len(normalized) == 8 {
formatted = normalized[:4] + "-" + normalized[4:]
}
var claims map[string]interface{}
if len(extraClaims) > 0 {
claims = extraClaims[0]
}
return s.authorizeDeviceCode(formatted, subject, claims)
}
// generateUserCode generates a user-friendly code formatted as XXXX-XXXX.
func (s *Service) generateUserCode() (string, error) {
length := s.config.Token.UserCodeLength
if length <= 0 {
length = 8
}
raw, err := gonanoid.Generate(userCodeAlphabet, length)
if err != nil {
return "", err
}
if len(raw) == 8 {
return raw[:4] + "-" + raw[4:], nil
}
return raw, nil
}
func clientSupportsGrantType(client *types.ClientInfo, grantType string) bool {
if client == nil || len(client.GrantTypes) == 0 {
return false
}
for _, gt := range client.GrantTypes {
if gt == grantType {
return true
}
}
return false
}

5
openapi/oauth/discovery.go
View file

@ -6,6 +6,7 @@ import (
"encoding/base64"
"fmt"
"math/big"
"strings"
"github.com/yaoapp/yao/openapi/oauth/types"
)
@ -53,7 +54,7 @@ func (s *Service) JWKS(ctx context.Context) (*types.JWKSResponse, error) {
// Endpoints returns a map of all available OAuth endpoints
// This provides endpoint discovery for clients
func (s *Service) Endpoints(ctx context.Context) (map[string]string, error) {
baseURL := s.config.IssuerURL
baseURL := strings.TrimRight(s.config.IssuerURL, "/") + s.config.BaseURL
endpoints := map[string]string{
"authorization_endpoint": fmt.Sprintf("%s/oauth/authorize", baseURL),
@ -63,7 +64,7 @@ func (s *Service) Endpoints(ctx context.Context) (map[string]string, error) {
"registration_endpoint": fmt.Sprintf("%s/oauth/register", baseURL),
"introspection_endpoint": fmt.Sprintf("%s/oauth/introspect", baseURL),
"revocation_endpoint": fmt.Sprintf("%s/oauth/revoke", baseURL),
"device_authorization_endpoint": fmt.Sprintf("%s/oauth/device", baseURL),
"device_authorization_endpoint": fmt.Sprintf("%s/oauth/device_authorization", baseURL),
"pushed_authorization_request_endpoint": fmt.Sprintf("%s/oauth/par", baseURL),
}

12
openapi/oauth/oauth.go
View file

@ -57,6 +57,7 @@ type Config struct {
// OAuth server metadata
IssuerURL string `json:"issuer_url"` // JWT token issuer URL
BaseURL string `json:"base_url"` // API route prefix (e.g. "/v1")
}
// FeatureFlags represents feature toggle configuration
@ -212,6 +213,15 @@ func setConfigDefaults(config *Config) error {
if config.Token.DeviceCodeLifetime == 0 {
config.Token.DeviceCodeLifetime = 15 * time.Minute
}
if config.Token.DeviceCodeLength == 0 {
config.Token.DeviceCodeLength = 8
}
if config.Token.UserCodeLength == 0 {
config.Token.UserCodeLength = 8
}
if config.Token.DeviceCodeInterval == 0 {
config.Token.DeviceCodeInterval = 5 * time.Second
}
if config.Token.AccessTokenFormat == "" {
config.Token.AccessTokenFormat = "jwt"
}
@ -257,6 +267,8 @@ func setConfigDefaults(config *Config) error {
config.Features.OAuth21Enabled = true
config.Features.PKCEEnforced = true
config.Features.RefreshTokenRotationEnabled = true
config.Features.DeviceFlowEnabled = true
config.Features.DynamicClientRegistrationEnabled = true
return nil
}

121
openapi/oauth/token.go
View file

@ -539,6 +539,127 @@ func (s *Service) consumeAuthorizationCode(code string) error {
return nil
}
// deviceCodeKey generates a key for device code storage
func (s *Service) deviceCodeKey(code string) string {
return fmt.Sprintf("%soauth:device_code:%s", s.prefix, code)
}
// userCodeKey generates a key for user code storage (reverse mapping)
func (s *Service) userCodeKey(code string) string {
return fmt.Sprintf("%soauth:user_code:%s", s.prefix, code)
}
// storeDeviceCode stores device code data and user_code -> device_code reverse mapping
func (s *Service) storeDeviceCode(deviceCode, userCode, clientID, scope string) error {
ttl := s.config.Token.DeviceCodeLifetime
codeData := map[string]interface{}{
"client_id": clientID,
"user_code": userCode,
"scope": scope,
"status": "pending",
"issued_at": time.Now().Unix(),
"expires_at": time.Now().Add(ttl).Unix(),
}
if err := s.store.Set(s.deviceCodeKey(deviceCode), codeData, ttl); err != nil {
return err
}
reverseData := map[string]interface{}{
"device_code": deviceCode,
}
return s.store.Set(s.userCodeKey(userCode), reverseData, ttl)
}
// getDeviceCodeData retrieves device code data from store
func (s *Service) getDeviceCodeData(deviceCode string) (map[string]interface{}, error) {
data, exists := s.store.Get(s.deviceCodeKey(deviceCode))
if !exists {
return nil, &types.ErrorResponse{
Code: types.ErrorExpiredToken,
ErrorDescription: "Device code not found or expired",
}
}
codeInfo, ok := data.(map[string]interface{})
if !ok {
if m, ok := data.(primitive.M); ok {
codeInfo = map[string]interface{}(m)
} else {
return nil, &types.ErrorResponse{
Code: types.ErrorServerError,
ErrorDescription: "Invalid device code data format",
}
}
}
return codeInfo, nil
}
// authorizeDeviceCode marks a device code as authorized via user_code lookup
func (s *Service) authorizeDeviceCode(userCode, subject string, extraClaims map[string]interface{}) error {
reverseData, exists := s.store.Get(s.userCodeKey(userCode))
if !exists {
return &types.ErrorResponse{
Code: types.ErrorInvalidGrant,
ErrorDescription: "Invalid or expired user code",
}
}
var deviceCode string
switch v := reverseData.(type) {
case map[string]interface{}:
deviceCode, _ = v["device_code"].(string)
case primitive.M:
deviceCode, _ = v["device_code"].(string)
}
if deviceCode == "" {
return &types.ErrorResponse{
Code: types.ErrorServerError,
ErrorDescription: "Invalid user code mapping",
}
}
codeData, err := s.getDeviceCodeData(deviceCode)
if err != nil {
return err
}
codeData["status"] = "authorized"
codeData["subject"] = subject
if extraClaims != nil {
codeData["extra_claims"] = extraClaims
}
// Re-store with remaining TTL
expiresAt, _ := codeData["expires_at"].(int64)
if expiresAt == 0 {
if f, ok := codeData["expires_at"].(float64); ok {
expiresAt = int64(f)
}
}
remaining := time.Until(time.Unix(expiresAt, 0))
if remaining <= 0 {
return &types.ErrorResponse{
Code: types.ErrorExpiredToken,
ErrorDescription: "Device code has expired",
}
}
return s.store.Set(s.deviceCodeKey(deviceCode), codeData, remaining)
}
// consumeDeviceCode deletes both device_code and user_code entries
func (s *Service) consumeDeviceCode(deviceCode string) error {
codeData, _ := s.getDeviceCodeData(deviceCode)
if codeData != nil {
if uc, ok := codeData["user_code"].(string); ok && uc != "" {
s.store.Del(s.userCodeKey(uc))
}
}
s.store.Del(s.deviceCodeKey(deviceCode))
return nil
}
// storeRefreshToken stores refresh token with metadata
func (s *Service) storeRefreshToken(refreshToken, clientID string) error {
tokenData := map[string]interface{}{

292
openapi/tests/oauth/device_test.go Normal file
View file

@ -0,0 +1,292 @@
package openapi_test
import (
"bytes"
"encoding/json"
"io"
"net/http"
"net/url"
"regexp"
"testing"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/yao/openapi"
"github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/openapi/tests/testutils"
)
// registerDeviceClient registers a device-flow-capable public client via HTTP POST to /oauth/register.
// Returns the client ID.
func registerDeviceClient(t *testing.T, serverURL, baseURL string) string {
t.Helper()
endpoint := serverURL + baseURL + "/oauth/register"
req := types.DynamicClientRegistrationRequest{
ClientName: "device-test-client",
RedirectURIs: []string{"http://localhost/device-callback"},
GrantTypes: []string{types.GrantTypeDeviceCode, types.GrantTypeRefreshToken},
TokenEndpointAuthMethod: types.TokenEndpointAuthNone,
}
jsonData, err := json.Marshal(req)
assert.NoError(t, err)
resp, err := http.Post(endpoint, "application/json", bytes.NewBuffer(jsonData))
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, http.StatusCreated, resp.StatusCode, "device client registration should succeed")
var regResp types.DynamicClientRegistrationResponse
err = json.NewDecoder(resp.Body).Decode(&regResp)
assert.NoError(t, err)
assert.NotEmpty(t, regResp.ClientID)
return regResp.ClientID
}
// registerConfidentialClient registers a confidential client with client_credentials grant.
// Returns clientID and clientSecret.
func registerConfidentialClient(t *testing.T, serverURL, baseURL string) (string, string) {
t.Helper()
endpoint := serverURL + baseURL + "/oauth/register"
req := types.DynamicClientRegistrationRequest{
ClientName: "confidential-token-client",
RedirectURIs: []string{"http://localhost/callback"},
GrantTypes: []string{types.GrantTypeClientCredentials},
TokenEndpointAuthMethod: types.TokenEndpointAuthBasic,
}
jsonData, err := json.Marshal(req)
assert.NoError(t, err)
resp, err := http.Post(endpoint, "application/json", bytes.NewBuffer(jsonData))
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, http.StatusCreated, resp.StatusCode, "confidential client registration should succeed")
var regResp types.DynamicClientRegistrationResponse
err = json.NewDecoder(resp.Body).Decode(&regResp)
assert.NoError(t, err)
assert.NotEmpty(t, regResp.ClientID)
assert.NotEmpty(t, regResp.ClientSecret)
return regResp.ClientID, regResp.ClientSecret
}
func TestDeviceAuthorization_Success(t *testing.T) {
serverURL := testutils.Prepare(t)
defer testutils.Clean()
baseURL := openapi.Server.Config.BaseURL
clientID := registerDeviceClient(t, serverURL, baseURL)
endpoint := serverURL + baseURL + "/oauth/device_authorization"
form := url.Values{}
form.Set("client_id", clientID)
resp, err := http.PostForm(endpoint, form)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
bodyBytes, err := io.ReadAll(resp.Body)
assert.NoError(t, err)
var devResp types.DeviceAuthorizationResponse
err = json.Unmarshal(bodyBytes, &devResp)
assert.NoError(t, err)
assert.NotEmpty(t, devResp.DeviceCode)
assert.NotEmpty(t, devResp.UserCode)
// user_code format XXXX-XXXX (9 chars including hyphen)
assert.Len(t, devResp.UserCode, 9)
assert.Regexp(t, regexp.MustCompile(`^[A-Z0-9]{4}-[A-Z0-9]{4}$`), devResp.UserCode)
assert.NotEmpty(t, devResp.VerificationURI)
assert.Greater(t, devResp.ExpiresIn, 0)
assert.Greater(t, devResp.Interval, 0)
}
func TestDeviceAuthorization_MissingClientID(t *testing.T) {
serverURL := testutils.Prepare(t)
defer testutils.Clean()
baseURL := openapi.Server.Config.BaseURL
endpoint := serverURL + baseURL + "/oauth/device_authorization"
form := url.Values{}
// no client_id
resp, err := http.PostForm(endpoint, form)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
}
func TestDeviceAuthorization_InvalidClient(t *testing.T) {
serverURL := testutils.Prepare(t)
defer testutils.Clean()
baseURL := openapi.Server.Config.BaseURL
endpoint := serverURL + baseURL + "/oauth/device_authorization"
form := url.Values{}
form.Set("client_id", "nonexistent")
resp, err := http.PostForm(endpoint, form)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
}
func TestDeviceToken_AuthorizationPending(t *testing.T) {
serverURL := testutils.Prepare(t)
defer testutils.Clean()
baseURL := openapi.Server.Config.BaseURL
clientID := registerDeviceClient(t, serverURL, baseURL)
// Get device code
devAuthEndpoint := serverURL + baseURL + "/oauth/device_authorization"
form := url.Values{}
form.Set("client_id", clientID)
resp, err := http.PostForm(devAuthEndpoint, form)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var devResp types.DeviceAuthorizationResponse
err = json.NewDecoder(resp.Body).Decode(&devResp)
assert.NoError(t, err)
assert.NotEmpty(t, devResp.DeviceCode)
// Poll token endpoint before user authorizes - should get authorization_pending
tokenEndpoint := serverURL + baseURL + "/oauth/token"
tokenForm := url.Values{}
tokenForm.Set("grant_type", types.GrantTypeDeviceCode)
tokenForm.Set("device_code", devResp.DeviceCode)
tokenForm.Set("client_id", clientID)
tokenResp, err := http.PostForm(tokenEndpoint, tokenForm)
assert.NoError(t, err)
defer tokenResp.Body.Close()
// RFC 8628: authorization_pending returns 400 with error
assert.Equal(t, http.StatusBadRequest, tokenResp.StatusCode)
bodyBytes, err := io.ReadAll(tokenResp.Body)
assert.NoError(t, err)
var errResp types.ErrorResponse
err = json.Unmarshal(bodyBytes, &errResp)
assert.NoError(t, err)
assert.Equal(t, types.ErrorAuthorizationPending, errResp.Code)
}
func TestDeviceToken_InvalidDeviceCode(t *testing.T) {
serverURL := testutils.Prepare(t)
defer testutils.Clean()
baseURL := openapi.Server.Config.BaseURL
clientID := registerDeviceClient(t, serverURL, baseURL)
tokenEndpoint := serverURL + baseURL + "/oauth/token"
tokenForm := url.Values{}
tokenForm.Set("grant_type", types.GrantTypeDeviceCode)
tokenForm.Set("device_code", "bogus-invalid-device-code")
tokenForm.Set("client_id", clientID)
resp, err := http.PostForm(tokenEndpoint, tokenForm)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, http.StatusBadRequest, resp.StatusCode)
bodyBytes, err := io.ReadAll(resp.Body)
assert.NoError(t, err)
var errResp types.ErrorResponse
err = json.Unmarshal(bodyBytes, &errResp)
assert.NoError(t, err)
assert.NotEmpty(t, errResp.Code)
}
func TestDeviceFlow_EndToEnd(t *testing.T) {
serverURL := testutils.Prepare(t)
defer testutils.Clean()
baseURL := openapi.Server.Config.BaseURL
// a. Register device client
deviceClientID := registerDeviceClient(t, serverURL, baseURL)
// b. POST /oauth/device_authorization -> get device_code + user_code
devAuthEndpoint := serverURL + baseURL + "/oauth/device_authorization"
form := url.Values{}
form.Set("client_id", deviceClientID)
resp, err := http.PostForm(devAuthEndpoint, form)
assert.NoError(t, err)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode)
var devResp types.DeviceAuthorizationResponse
err = json.NewDecoder(resp.Body).Decode(&devResp)
assert.NoError(t, err)
assert.NotEmpty(t, devResp.DeviceCode)
assert.NotEmpty(t, devResp.UserCode)
// c. Get bearer token: register confidential client, get token via client_credentials.
// Device authorize requires a token with subject; client_credentials tokens have no subject.
// Use ObtainAccessTokenWithRootPermission to get a token with subject for device authorize.
confClientID, confClientSecret := registerConfidentialClient(t, serverURL, baseURL)
tokenInfo := testutils.ObtainAccessTokenWithRootPermission(t, serverURL, confClientID, confClientSecret, "http://localhost/callback", "openid profile")
bearerToken := tokenInfo.AccessToken
tokenEndpoint := serverURL + baseURL + "/oauth/token"
// d. POST /oauth/device/authorize with bearer + user_code -> assert 200
deviceAuthorizeEndpoint := serverURL + baseURL + "/oauth/device/authorize"
authForm := url.Values{}
authForm.Set("user_code", devResp.UserCode)
authReq, err := http.NewRequest("POST", deviceAuthorizeEndpoint, bytes.NewBufferString(authForm.Encode()))
assert.NoError(t, err)
authReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
authReq.Header.Set("Authorization", "Bearer "+bearerToken)
authResp, err := http.DefaultClient.Do(authReq)
assert.NoError(t, err)
defer authResp.Body.Close()
assert.Equal(t, http.StatusOK, authResp.StatusCode, "device authorize should succeed")
// e. POST /oauth/token with device_code -> assert access_token returned
dcForm := url.Values{}
dcForm.Set("grant_type", types.GrantTypeDeviceCode)
dcForm.Set("device_code", devResp.DeviceCode)
dcForm.Set("client_id", deviceClientID)
dcResp, err := http.PostForm(tokenEndpoint, dcForm)
assert.NoError(t, err)
defer dcResp.Body.Close()
assert.Equal(t, http.StatusOK, dcResp.StatusCode, "device token exchange should succeed")
var finalToken struct {
AccessToken string `json:"access_token"`
TokenType string `json:"token_type"`
}
err = json.NewDecoder(dcResp.Body).Decode(&finalToken)
assert.NoError(t, err)
assert.NotEmpty(t, finalToken.AccessToken)
assert.Equal(t, "Bearer", finalToken.TokenType)
}

46
openapi/well-known.go
View file

@ -1,10 +1,14 @@
package openapi
import (
"fmt"
"net"
"os"
"strconv"
"strings"
"github.com/gin-gonic/gin"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/share"
)
@ -45,6 +49,7 @@ type YaoMetadata struct {
// Dashboard configuration
Dashboard string `json:"dashboard,omitempty"` // Admin dashboard root path
GRPC string `json:"grpc,omitempty"` // gRPC server address (e.g., "127.0.0.1:9099")
Optional map[string]interface{} `json:"optional,omitempty"` // Optional settings
// Developer information
@ -67,6 +72,7 @@ func (openapi *OpenAPI) yaoMetadata(c *gin.Context) {
IssuerURL: openapi.Config.OAuth.IssuerURL,
ServerURL: resolveServerURL(openapi.Config.OAuth.IssuerURL),
Dashboard: "/" + dashboard,
GRPC: resolveGRPCAddr(c),
Optional: share.App.Optional,
}
@ -97,8 +103,46 @@ func resolveServerURL(issuerURL string) string {
return ""
}
// resolveGRPCAddr returns the gRPC server address for client discovery.
// Uses the request Host's IP with the configured gRPC port.
func resolveGRPCAddr(c *gin.Context) string {
cfg := config.Conf.GRPC
if strings.ToLower(cfg.Enabled) == "off" {
return ""
}
port := cfg.Port
if port == 0 {
port = 9099
}
host := cfg.Host
if host == "" || host == "0.0.0.0" {
reqHost := c.Request.Host
h, _, err := net.SplitHostPort(reqHost)
if err != nil {
h = reqHost
}
host = h
} else if strings.Contains(host, ",") {
host = strings.TrimSpace(strings.Split(host, ",")[0])
}
return fmt.Sprintf("%s:%s", host, strconv.Itoa(port))
}
// oauthServerMetadata returns authorization server metadata - RFC 8414
func (openapi *OpenAPI) oauthServerMetadata(c *gin.Context) {}
func (openapi *OpenAPI) oauthServerMetadata(c *gin.Context) {
if openapi.OAuth == nil {
c.JSON(503, gin.H{"error": "OAuth service not available"})
return
}
metadata, err := openapi.OAuth.GetServerMetadata(c.Request.Context())
if err != nil {
c.JSON(500, gin.H{"error": err.Error()})
return
}
c.JSON(200, metadata)
}
// oauthOpenIDConfiguration returns OpenID Connect configuration
func (openapi *OpenAPI) oauthOpenIDConfiguration(c *gin.Context) {}

1570
sandbox/DESIGN.md
View file

File diff suppressed because it is too large Load diff

579
sandbox/SPEC.md Normal file
View file

@ -0,0 +1,579 @@
# Sandbox Functional Specification
Detailed interfaces, types, and behavior for the sandbox refactoring.
Architecture and rationale: see [DESIGN.md](./DESIGN.md).
---
## 1. sandbox.Manager
Replaces the current Docker-only Manager. Backed by a single `tai.Client`.
### Config
```go
type Config struct {
Image string // container image, default "yaoapp/workspace:latest"
MaxContainers int // global limit, default 100
IdleTimeout time.Duration // default cleanup interval, default 30m
MaxMemory string // per-container, e.g. "2g"
MaxCPU float64 // per-container, e.g. 1.0
ContainerWorkDir string // mount target inside container, default "/workspace"
ContainerUser string // empty = image default
}
```
Environment variable overrides remain the same (`YAO_SANDBOX_IMAGE`, etc.). `WorkspaceRoot` and `IPCDir` are removed — local paths derived from `tai.Client.IsLocal()` at runtime; remote mode uses `tai.Client.Volume()`.
### Constructor
```go
func NewManager(client *tai.Client, cfg *Config) (*Manager, error)
```
- Validates `client` is non-nil and healthy (calls `client.Sandbox().List()` as connectivity check)
- Starts background cleanup goroutine
- Returns ready Manager
### Manager struct
```go
type Manager struct {
client *tai.Client
config *Config
sandboxes sync.Map // name → *Sandbox
running atomic.Int32
ipc *IPCRouter // local: Unix socket manager, remote: gRPC stub
cleanup *time.Ticker
done chan struct{}
}
```
### Public methods
```go
// Lifecycle
func (m *Manager) GetOrCreate(ctx context.Context, opts GetOrCreateOptions) (*Sandbox, error)
func (m *Manager) Get(ctx context.Context, name string) (*Sandbox, error)
func (m *Manager) Start(ctx context.Context, name string) error
func (m *Manager) Stop(ctx context.Context, name string, timeout time.Duration) error
func (m *Manager) Remove(ctx context.Context, name string) error
func (m *Manager) List(ctx context.Context, filter ListFilter) ([]*Sandbox, error)
// Execution
func (m *Manager) Exec(ctx context.Context, name string, cmd []string, opts ExecOptions) (*ExecResult, error)
func (m *Manager) Stream(ctx context.Context, name string, cmd []string, opts ExecOptions) (io.ReadCloser, error)
func (m *Manager) KillProcess(ctx context.Context, name string, pattern string) error
// File operations (routes local/remote internally)
func (m *Manager) ReadFile(ctx context.Context, name string, path string) ([]byte, error)
func (m *Manager) WriteFile(ctx context.Context, name string, path string, data []byte) error
func (m *Manager) ListDir(ctx context.Context, name string, path string) ([]FileInfo, error)
func (m *Manager) Stat(ctx context.Context, name string, path string) (*FileInfo, error)
func (m *Manager) MkDir(ctx context.Context, name string, path string) error
func (m *Manager) RemoveFile(ctx context.Context, name string, path string) error
func (m *Manager) CopyToContainer(ctx context.Context, name string, hostPath, containerPath string) error
func (m *Manager) CopyFromContainer(ctx context.Context, name string, containerPath, hostPath string) error
// Info
func (m *Manager) IsLocal() bool
func (m *Manager) Close() error
```
---
## 2. Types
### Sandbox
```go
type Sandbox struct {
Name string
UserID string
ChatID string
Image string
Status Status
Lifecycle Lifecycle
CreatedAt time.Time
LastUsedAt time.Time
IP string
}
```
### Status
```go
type Status string
const (
StatusCreated Status = "created"
StatusRunning Status = "running"
StatusStopped Status = "stopped"
)
```
### Lifecycle
```go
type Lifecycle string
const (
LifecycleOneShot Lifecycle = "one-shot" // destroyed after execution
LifecycleSession Lifecycle = "session" // alive while user active, idle timeout
LifecycleLongRunning Lifecycle = "long-running" // hours/days, recoverable
LifecyclePersistent Lifecycle = "persistent" // never auto-cleaned
)
```
### GetOrCreateOptions
```go
type GetOrCreateOptions struct {
UserID string
ChatID string
Image string // override Config.Image
Lifecycle Lifecycle // default: LifecycleSession
Env map[string]string // injected into container
Cmd []string // override entrypoint
Memory string // override Config.MaxMemory
CPU float64 // override Config.MaxCPU
}
```
### ExecOptions / ExecResult
```go
type ExecOptions struct {
WorkDir string
Env map[string]string
Stdin io.Reader
Timeout time.Duration
}
type ExecResult struct {
ExitCode int
Stdout string
Stderr string
}
```
### ListFilter
```go
type ListFilter struct {
UserID string // empty = all users
Status Status // empty = all statuses
Lifecycle Lifecycle // empty = all policies
}
```
### FileInfo
```go
type FileInfo struct {
Name string
Path string
Size int64
Mode os.FileMode
ModTime time.Time
IsDir bool
}
```
### Errors
```go
var (
ErrTooManyContainers = errors.New("sandbox: container limit reached")
ErrNotFound = errors.New("sandbox: not found")
ErrNotRunning = errors.New("sandbox: not running")
ErrAlreadyExists = errors.New("sandbox: already exists")
)
```
---
## 3. Lifecycle State Machine
```
GetOrCreate()
│
▼
┌─────────┐
┌────────│ Created │
│ └────┬────┘
│ Start() │
│ ▼
│ ┌─────────┐ idle timeout / Stop()
│ │ Running │──────────────────┐
│ └────┬────┘ │
│ │ ▼
│ │ ┌─────────┐
│ │ │ Stopped │
│ │ └────┬────┘
│ │ Start() │
│ │ ┌───────────────────┘
│ │ │
│ ▼ ▼
│ Remove() from any state
│ │
│ ▼
│ [Destroyed]
│
└── one-shot: auto Remove() after Exec/Stream returns
```
### Cleanup rules
| Lifecycle | Trigger | Action |
|-----------|---------|--------|
| one-shot | Exec/Stream completes | Manager.Remove() immediately |
| session | `IdleTimeout` since `LastUsedAt` | Manager.Stop() then Remove() |
| long-running | `IdleTimeout * 24` since `LastUsedAt` | Manager.Stop() (not removed, can restart) |
| persistent | Never | No automatic action |
Background goroutine runs every `Config.IdleTimeout / 2`, scans `sandboxes`, applies rules.
### Touch
Every `Exec`, `Stream`, `ReadFile`, `WriteFile`, `ListDir` call updates `LastUsedAt`.
---
## 4. File Operations Routing
```go
func (m *Manager) ReadFile(ctx context.Context, name string, path string) ([]byte, error) {
if m.client.IsLocal() {
hostPath := m.hostPath(name, path)
return os.ReadFile(hostPath)
}
sessionID := m.sessionID(name)
ws := m.client.Workspace(sessionID)
f, err := ws.Open(path)
// ... read and return
}
```
| Operation | Local | Remote |
|-----------|-------|--------|
| ReadFile | `os.ReadFile(hostPath)` | `client.Workspace(session).Open(path)` |
| WriteFile | `os.WriteFile(hostPath)` | `client.Volume().Write(session, path, data)` |
| ListDir | `os.ReadDir(hostPath)` | `client.Volume().ReadDir(session, path)` |
| Stat | `os.Stat(hostPath)` | `client.Volume().Stat(session, path)` |
| MkDir | `os.MkdirAll(hostPath)` | `client.Volume().MkDir(session, path)` |
| RemoveFile | `os.RemoveAll(hostPath)` | `client.Volume().Remove(session, path)` |
| CopyToContainer | bind mount (noop, already on host) | `client.Volume().Write()` streamed |
| CopyFromContainer | bind mount (direct read) | `client.Volume().Read()` streamed |
`hostPath` = `dataDir/{userID}/{chatID}/{containerRelativePath}`
`sessionID` = `{userID}/{chatID}` (maps to volume session on Tai)
---
## 5. IPC Router
Abstracts local Unix socket vs remote gRPC relay. Manager creates the right one based on `client.IsLocal()`.
### Interface
```go
type IPCRouter interface {
Create(sessionID string, tools []MCPTool) (IPCSession, error)
Get(sessionID string) (IPCSession, error)
Close(sessionID string) error
CloseAll() error
}
type IPCSession interface {
SetTools(tools []MCPTool)
SetContext(ctx *AgentContext)
SocketPath() string // local only, empty for remote
GRPCAddr() string // remote only, empty for local
Close() error
}
```
### Local implementation
Same as current `ipc.Manager` — creates Unix socket per session, bind-mounts into container, yao-bridge connects to it.
### Remote implementation
No socket. Container receives `YAO_IPC_MODE=grpc` and `YAO_IPC_ADDR=tai-host:9100`. `yao-bridge` (`yao/tai/bridge/`) connects to Tai's gRPC relay, which forwards to Yao gRPC Server. Tai relay upstream is per-container via `CreateRequest.GRPCUpstream`, not a Tai startup parameter.
Tool registration: remote IPCSession sends tool list to Yao gRPC Server via a registration RPC at session creation.
### Container env injection
```go
func (m *Manager) buildContainerEnv(session IPCSession, userEnv map[string]string) map[string]string {
env := maps.Clone(userEnv)
if m.client.IsLocal() {
env["YAO_IPC_MODE"] = "socket"
env["YAO_IPC_ADDR"] = session.SocketPath()
} else {
env["YAO_IPC_MODE"] = "grpc"
env["YAO_IPC_ADDR"] = session.GRPCAddr()
env["YAO_TOKEN"] = m.issueAccessToken(session)
env["YAO_REFRESH_TOKEN"] = m.issueRefreshToken(session)
}
return env
}
// CreateRequest also carries GRPCUpstream for Tai relay routing (per-container, not per-Tai)
```
---
## 6. Yao gRPC Server
### Proto definition
```protobuf
syntax = "proto3";
package yao.v1;
service Yao {
rpc Exec(ExecRequest) returns (ExecResponse);
rpc StreamExec(ExecRequest) returns (stream ExecChunk);
// MCP tool registration (called by remote IPC sessions)
rpc RegisterTools(RegisterToolsRequest) returns (RegisterToolsResponse);
// Health
rpc Healthz(HealthzRequest) returns (HealthzResponse);
}
message ExecRequest {
string process = 1; // e.g. "models.user.Find"
bytes args = 2; // JSON-encoded arguments
string session = 3; // sandbox session ID for context
}
message ExecResponse {
bytes result = 1; // JSON-encoded result
string error = 2;
}
message ExecChunk {
bytes data = 1;
bool done = 2;
}
message RegisterToolsRequest {
string session = 1;
repeated MCPToolDef tools = 2;
}
message MCPToolDef {
string name = 1;
string description = 2;
string process = 3; // Yao process to call
bytes input_schema = 4; // JSON Schema
}
message RegisterToolsResponse {}
message HealthzRequest {}
message HealthzResponse {
string status = 1;
}
```
### Server startup
```go
func StartGRPCServer(cfg GRPCConfig) (*grpc.Server, error)
type GRPCConfig struct {
Listen string // "127.0.0.1:9099" or "0.0.0.0:9099"
AllowCIDR []string // IP allowlist, empty = no restriction
}
```
Interceptor chain: `ipAllowInterceptor` → `authInterceptor` → handler.
### Exec handler
```go
func (s *yaoServer) Exec(ctx context.Context, req *pb.ExecRequest) (*pb.ExecResponse, error) {
claims := claimsFromContext(ctx)
// ACL check: does this token have permission to call this process?
p := process.New(req.Process)
var args []interface{}
json.Unmarshal(req.Args, &args)
result, err := p.Exec(args...)
if err != nil {
return &pb.ExecResponse{Error: err.Error()}, nil
}
data, _ := json.Marshal(result)
return &pb.ExecResponse{Result: data}, nil
}
```
---
## 7. Agent Layer
### Assistant sandbox config
```yaml
# assistants/coder.yao
sandbox:
enabled: true
lifecycle: session
idle_timeout: 30m
image: yaoapp/workspace:latest
command: claude
memory: "4g"
cpu: 2.0
```
### Parsed config type
```go
type AssistantSandboxConfig struct {
Enabled bool `json:"enabled"`
Lifecycle Lifecycle `json:"lifecycle"`
IdleTimeout time.Duration `json:"idle_timeout"`
Image string `json:"image"`
Command string `json:"command"`
Memory string `json:"memory"`
CPU float64 `json:"cpu"`
}
```
### Init flow (new)
```go
func (a *Assistant) initSandbox(ctx context.Context) (*agentsandbox.Executor, error) {
mgr := GetSandboxManager() // global, initialized with tai.Client at Yao startup
sb, err := mgr.GetOrCreate(ctx, sandbox.GetOrCreateOptions{
UserID: a.userID,
ChatID: a.chatID,
Image: a.config.Sandbox.Image,
Lifecycle: a.config.Sandbox.Lifecycle,
Memory: a.config.Sandbox.Memory,
CPU: a.config.Sandbox.CPU,
})
// ...
executor := agentsandbox.New(mgr, sb, a.config.Sandbox.Command)
return executor, nil
}
```
### Cleanup (new)
```go
func (a *Assistant) sandboxCleanup(executor *agentsandbox.Executor) {
executor.Disconnect()
// Manager handles actual removal based on lifecycle policy.
// one-shot: already removed after Exec.
// session: will be cleaned up by background goroutine after idle timeout.
// long-running/persistent: stays.
}
```
### GetSandboxManager (new)
```go
var (
managerOnce sync.Once
manager *sandbox.Manager
)
func GetSandboxManager() *sandbox.Manager {
managerOnce.Do(func() {
client := config.GetTaiClient() // initialized at Yao startup from env/config
mgr, err := sandbox.NewManager(client, loadSandboxConfig())
if err != nil {
log.Fatal("sandbox manager init: %v", err)
}
manager = mgr
})
return manager
}
```
### Executor factory
```go
// agent/sandbox/executor.go
func New(mgr *sandbox.Manager, sb *sandbox.Sandbox, command string) Executor {
switch command {
case "claude":
return claude.NewExecutor(mgr, sb)
default:
return generic.NewExecutor(mgr, sb)
}
}
```
### Executor interface (unchanged)
```go
type Executor interface {
Stream(ctx context.Context, opts StreamOptions) (io.ReadCloser, error)
Disconnect() error
// Delegated to Manager internally
ReadFile(ctx context.Context, path string) ([]byte, error)
WriteFile(ctx context.Context, path string, data []byte) error
ListDir(ctx context.Context, path string) ([]FileInfo, error)
Exec(ctx context.Context, cmd []string) (string, error)
GetWorkDir() string
GetSandboxID() string
GetVNCUrl() string
}
```
Each method delegates to `mgr.ReadFile(ctx, sb.Name, path)` etc. The executor is a thin wrapper that knows the sandbox name.
---
## 8. Naming Convention
| Entity | Pattern | Example |
|--------|---------|---------|
| Container/Pod name | `yao-sb-{userID}-{chatID}` | `yao-sb-u123-c456` |
| Volume session | `{userID}/{chatID}` | `u123/c456` |
| IPC session | `{chatID}` | `c456` |
| Host workspace (local) | `{dataDir}/{userID}/{chatID}/` | `/data/u123/c456/` |
Prefix shortened from `yao-sandbox-` to `yao-sb-` for K8s DNS name length limit (63 chars).
---
## 9. Environment Variables
### Yao process
| Variable | Purpose | Default |
|----------|---------|---------|
| `YAO_TAI_ADDR` | Tai endpoint, e.g. `tai://10.0.0.1` or empty for local Docker | `""` (local) |
| `YAO_TAI_RUNTIME` | `docker` or `k8s` | `docker` |
| `YAO_TAI_KUBECONFIG` | Path to kubeconfig (K8s only) | |
| `YAO_TAI_NAMESPACE` | K8s namespace | `default` |
| `YAO_GRPC_LISTEN` | gRPC server listen address | `127.0.0.1:9099` |
| `YAO_GRPC_ALLOW` | CIDR allowlist, comma-separated | (empty = no filter) |
| `YAO_SANDBOX_IMAGE` | Default container image | `yaoapp/workspace:latest` |
| `YAO_SANDBOX_MAX` | Max containers | `100` |
| `YAO_SANDBOX_IDLE_TIMEOUT` | Idle timeout duration | `30m` |
| `YAO_SANDBOX_MEMORY` | Memory limit | `2g` |
| `YAO_SANDBOX_CPU` | CPU limit | `1.0` |
### Container-internal
| Variable | Purpose | Set by |
|----------|---------|--------|
| `YAO_IPC_MODE` | `socket` or `grpc` | Manager at creation |
| `YAO_IPC_ADDR` | Socket path or gRPC host:port | Manager at creation |
| `YAO_TOKEN` | JWT access token for gRPC auth (remote only, short TTL 15m) | Manager at creation |
| `YAO_REFRESH_TOKEN` | JWT refresh token (remote only, no expiry, revoked on Remove) | Manager at creation |

6
service/service.go
View file

@ -11,6 +11,10 @@ import (
"github.com/yaoapp/yao/share"
)
// Router holds the active gin.Engine so the gRPC API proxy can forward
// requests internally without an HTTP round-trip.
var Router *gin.Engine
// Start the yao service
func Start(cfg config.Config) (*http.Server, error) {
@ -24,6 +28,7 @@ func Start(cfg config.Config) (*http.Server, error) {
}
router := gin.New()
Router = router
router.Use(Middlewares...)
var apiRoot string
@ -68,6 +73,7 @@ func Start(cfg config.Config) (*http.Server, error) {
// Restart the yao service
func Restart(srv *http.Server, cfg config.Config) error {
router := gin.New()
Router = router
router.Use(Middlewares...)
if openapi.Server != nil {

147
tai/DESIGN.md
View file

@ -1,147 +0,0 @@
# Tai Go SDK
Go client library for [Tai](https://github.com/yaoapp/tai) — the universal runtime bridge for Yao Sandbox.
## Overview
Provides a unified API for container lifecycle, filesystem operations, HTTP proxy, and VNC access.
Supports two modes via a single entry point:
- **Local** (`docker://` or `""`) — direct Docker daemon connection
- **Remote** (`tai://host`) — via Tai Server proxy (Docker, K8s)
All sub-packages follow the same pattern: **interface + Remote/Local implementations**.
## Package Layout
```
yao/tai/
├── tai.go # Client, New(), Option, Close()
├── volume/ # Volume IO + Sync
├── workspace/ # Go fs.FS wrapper over volume.Volume
├── sandbox/ # Container lifecycle (Create/Start/Stop/Exec/Remove)
│ ├── sandbox.go # Interface + shared types
│ ├── local.go # Direct Docker socket
│ ├── docker.go # Docker via Tai proxy
│ ├── docker_core.go # Shared Docker SDK logic
│ └── k8s.go # Kubernetes via Tai TCP proxy
├── proxy/ # HTTP reverse proxy URL resolution
└── vnc/ # VNC WebSocket URL resolution
```
## Quick Start
```go
import "github.com/yaoapp/yao/tai"
// Local — default Docker socket
c, _ := tai.New("")
// Local — explicit address
c, _ := tai.New("docker:///var/run/docker.sock")
c, _ := tai.New("docker://192.168.1.50:2375")
// Remote — via Tai Server (Docker runtime, default)
c, _ := tai.New("tai://192.168.1.100")
// Remote — via Tai Server (K8s runtime)
c, _ := tai.New("tai://10.0.0.5", tai.K8s,
tai.WithKubeConfig("/path/to/kubeconfig.yml"),
tai.WithNamespace("sandbox"),
)
defer c.Close()
// Container lifecycle
id, _ := c.Sandbox().Create(ctx, sandbox.CreateOptions{
Image: "node:20",
Cmd: []string{"sleep", "infinity"},
})
c.Sandbox().Start(ctx, id)
// Filesystem
ws := c.Workspace("session-1")
ws.WriteFile("app.js", []byte("console.log('hi')"), 0644)
data, _ := ws.ReadFile("app.js")
// HTTP proxy URL
url, _ := c.Proxy().URL(ctx, id, 3000, "/api/health")
// VNC URL
vncURL, _ := c.VNC().URL(ctx, id)
```
## Address Protocol
| Prefix | Mode | Description |
|--------|------|-------------|
| `""` | Local | Platform default Docker socket |
| `docker://...` | Local | Direct Docker daemon (socket or TCP) |
| `tai://host` | Remote | Via Tai Server, all services proxied |
## Sub-Package Interfaces
### volume.Volume
File IO and directory sync between Yao and the container workspace.
- `ReadFile`, `WriteFile`, `Stat`, `ListDir`, `Remove`, `Rename`, `MkdirAll`
- `SyncPush` (Yao -> Tai), `SyncPull` (Tai -> Yao)
- **Remote**: gRPC to Tai `:9100`
- **Local**: direct disk IO under `dataDir/{sessionID}/`
### workspace.FS
Go `fs.FS`-compatible interface wrapping `volume.Volume`, adding write operations.
### sandbox.Sandbox
Container lifecycle: `Create`, `Start`, `Stop`, `Remove`, `Exec`, `Inspect`, `List`.
- **Local**: direct Docker socket, handles VNC port mapping and capabilities
- **Docker**: via Tai `:2375` (Docker Engine API proxy)
- **K8s**: via Tai `:6443` (kube-apiserver TCP proxy, single-container Pod per sandbox)
### proxy.Proxy
HTTP service URL resolution: `URL(ctx, containerID, port, path)`.
- **Remote**: `http://tai-host:8080/{id}:{port}/{path}`
- **Local**: `http://127.0.0.1:{hostPort}/{path}` via `sandbox.Inspect`
### vnc.VNC
VNC WebSocket URL resolution: `URL(ctx, containerID)`.
- **Remote**: `ws://tai-host:6080/vnc/{id}/ws`
- **Local**: `ws://127.0.0.1:{vncHostPort}/ws` via `sandbox.Inspect`
## Options
```go
tai.Docker // Docker runtime (default, can omit)
tai.K8s // Kubernetes runtime
tai.WithPorts(Ports{}) // custom port mapping
tai.WithHTTPClient(hc) // custom HTTP client
tai.WithDataDir(dir) // workspace root (Local mode)
tai.WithKubeConfig(path) // kubeconfig file path (K8s runtime)
tai.WithNamespace(ns) // namespace for K8s (default "default")
```
## Default Ports
| Service | Default Port |
|---------|-------------|
| gRPC (Volume + Gateway) | 9100 |
| HTTP Proxy | 8080 |
| VNC Router | 6080 |
| Docker API Proxy | 2375 |
| K8s API Proxy | 6443 |
## Dependencies
- `github.com/yaoapp/tai/volume/pb` — gRPC proto types
- `google.golang.org/grpc`
- `github.com/pierrec/lz4/v4` — sync compression
- `github.com/docker/docker` — Docker SDK
- `k8s.io/client-go` + `k8s.io/api` + `k8s.io/apimachinery` — Kubernetes SDK

111
tai/docs/README.md Normal file
View file

@ -0,0 +1,111 @@
# Tai SDK
Go client library for the [Tai](https://github.com/YaoApp/tai) runtime bridge. Provides unified access to container sandboxes, volume IO, HTTP proxy, and VNC routing — transparently working in both **Local** (direct Docker) and **Remote** (via Tai server) modes.
## Package Layout
| Package | Import Path | Description |
|---------|-------------|-------------|
| `tai` | `github.com/yaoapp/yao/tai` | Top-level client, `New()`, options, `Close()` |
| `sandbox` | `github.com/yaoapp/yao/tai/sandbox` | Container lifecycle (Create/Start/Stop/Exec/Remove) |
| `volume` | `github.com/yaoapp/yao/tai/volume` | File IO and directory sync |
| `workspace` | `github.com/yaoapp/yao/tai/workspace` | `fs.FS`-compatible filesystem over Volume |
| `proxy` | `github.com/yaoapp/yao/tai/proxy` | HTTP reverse proxy URL resolution |
| `vnc` | `github.com/yaoapp/yao/tai/vnc` | VNC WebSocket URL resolution |
## Quick Start
### Local Mode (direct Docker)
```go
c, err := tai.New("")
// or: tai.New("unix:///var/run/docker.sock")
// or: tai.New("tcp://192.168.1.50:2375")
defer c.Close()
id, _ := c.Sandbox().Create(ctx, sandbox.CreateOptions{
Name: "my-sandbox",
Image: "alpine:latest",
Cmd: []string{"sleep", "300"},
})
c.Sandbox().Start(ctx, id)
```
### Remote Mode (via Tai server, Docker runtime)
```go
c, err := tai.New("tai://192.168.1.100")
defer c.Close()
result, _ := c.Sandbox().Exec(ctx, id, []string{"echo", "hello"}, sandbox.ExecOptions{})
fmt.Println(result.Stdout) // "hello\n"
```
### Remote Mode (via Tai server, K8s runtime)
```go
c, err := tai.New("tai://192.168.1.100", tai.K8s,
tai.WithKubeConfig("/path/to/kubeconfig.yml"),
tai.WithNamespace("default"),
tai.WithPorts(tai.Ports{K8s: 6443}),
)
defer c.Close()
```
## Address Protocols
| Address | Mode | Description |
|---------|------|-------------|
| `""` | Local | Platform default Docker socket |
| `unix:///var/run/docker.sock` | Local | Explicit Unix socket |
| `tcp://host:port` | Local | Explicit TCP Docker daemon |
| `npipe:////./pipe/docker_engine` | Local | Windows named pipe |
| `docker://host:port` | Local | Docker scheme |
| `tai://host` | Remote | Connect via Tai server |
## Options
| Option | Description | Default |
|--------|-------------|---------|
| `WithPorts(Ports{...})` | Override Tai service ports | gRPC=9100, HTTP=8080, VNC=6080 |
| `WithHTTPClient(*http.Client)` | Custom HTTP client for proxy/VNC | `http.DefaultClient` |
| `WithDataDir(path)` | Volume storage root (Local mode) | `/tmp/tai-volumes` |
| `WithKubeConfig(path)` | Kubeconfig file path (K8s mode, **required**) | - |
| `WithNamespace(ns)` | K8s namespace | `"default"` |
## Default Ports
| Service | Port | Description |
|---------|------|-------------|
| gRPC | 9100 | Volume IO + Gateway |
| HTTP | 8080 | HTTP reverse proxy |
| VNC | 6080 | VNC WebSocket router |
| Docker | 2375 | Docker API proxy |
| K8s | 6443 | Kubernetes API proxy |
## Client API
```go
c.Volume() // volume.Volume
c.Workspace(sessionID) // workspace.FS
c.Sandbox() // sandbox.Sandbox
c.Proxy() // proxy.Proxy
c.VNC() // vnc.VNC
c.IsLocal() // bool
c.Close() // error
```
## Runtime Constants
```go
tai.Docker // default — use Docker runtime via Tai
tai.K8s // use Kubernetes runtime via Tai
```
## Sub-Package Documentation
- [sandbox.md](sandbox.md) — Container lifecycle management
- [volume.md](volume.md) — File IO and sync
- [workspace.md](workspace.md) — fs.FS-compatible filesystem
- [proxy.md](proxy.md) — HTTP reverse proxy
- [vnc.md](vnc.md) — VNC WebSocket routing

86
tai/docs/proxy.md Normal file
View file

@ -0,0 +1,86 @@
# Package `proxy`
HTTP reverse proxy URL resolution. Resolves service URLs for containers so that HTTP services running inside sandboxes can be accessed from the host.
## Interface
```go
type Proxy interface {
URL(ctx context.Context, containerID string, port int, path string) (string, error)
Healthz(ctx context.Context) error
}
```
## Implementations
| Implementation | Constructor | Mode | URL Pattern |
|----------------|-------------|------|-------------|
| **Remote** | `NewRemote(host, port, hc)` | Via Tai HTTP proxy | `http://tai-host:8080/{containerID}:{port}/{path}` |
| **Local** | `NewLocal(sb)` | Direct host port lookup | `http://127.0.0.1:{hostPort}/{path}` |
## Constructors
### NewRemote
```go
func NewRemote(host string, port int, hc *http.Client) Proxy
```
Creates a Proxy that routes through Tai's HTTP reverse proxy. URLs are constructed by combining the Tai server address with the container ID and port.
- `host` — Tai server hostname/IP
- `port` — Tai HTTP proxy port (default 8080)
- `hc` — custom HTTP client, `nil` uses `http.DefaultClient`
### NewLocal
```go
func NewLocal(sb sandbox.Sandbox) Proxy
```
Creates a Proxy that resolves URLs by inspecting the container's port mappings via `sandbox.Inspect`. Looks up the host port bound to the requested container port.
Returns an error if the requested port is not mapped.
## Methods
### URL
```go
URL(ctx context.Context, containerID string, port int, path string) (string, error)
```
Resolves an HTTP URL to reach a service running on `port` inside the given container.
**Remote example:** container `abc123` port `3000` path `/api/health`
→ `http://tai-host:8080/abc123:3000/api/health`
**Local example:** container `abc123` port `3000` mapped to host port `32768`
→ `http://127.0.0.1:32768/api/health`
### Healthz
```go
Healthz(ctx context.Context) error
```
Checks the health of the proxy backend.
- **Remote**: sends `GET /healthz` to the Tai HTTP proxy server
- **Local**: always returns `nil` (no external dependency)
## Example
```go
c, _ := tai.New("tai://192.168.1.100")
defer c.Close()
// Get URL for a web service running on port 3000
url, _ := c.Proxy().URL(ctx, containerID, 3000, "/api/status")
resp, _ := http.Get(url)
// Health check
if err := c.Proxy().Healthz(ctx); err != nil {
log.Fatal("Tai HTTP proxy is down:", err)
}
```

182
tai/docs/sandbox.md Normal file
View file

@ -0,0 +1,182 @@
# Package `sandbox`
Container lifecycle management. Provides a unified `Sandbox` interface with three implementations:
| Implementation | Constructor | Backend | Mode |
|----------------|-------------|---------|------|
| **Local** | `NewLocal(addr)` | Direct Docker daemon | Local |
| **Docker** | `NewDocker(addr)` | Docker via Tai proxy | Remote |
| **K8s** | `NewK8s(addr, opts)` | Kubernetes via Tai proxy | Remote |
## Interface
```go
type Sandbox interface {
Create(ctx context.Context, opts CreateOptions) (id string, err error)
Start(ctx context.Context, id string) error
Stop(ctx context.Context, id string, timeout time.Duration) error
Remove(ctx context.Context, id string, force bool) error
Exec(ctx context.Context, id string, cmd []string, opts ExecOptions) (*ExecResult, error)
Inspect(ctx context.Context, id string) (*ContainerInfo, error)
List(ctx context.Context, opts ListOptions) ([]ContainerInfo, error)
Close() error
}
```
## Constructors
### NewLocal
```go
func NewLocal(addr string) (Sandbox, error)
```
Connects directly to a Docker daemon. `addr` can be:
- `""` — platform default (Unix socket on Linux/macOS, named pipe on Windows)
- `"unix:///var/run/docker.sock"` — explicit Unix socket
- `"tcp://host:port"` — explicit TCP
Pings the daemon on creation; returns an error if unreachable.
### NewDocker
```go
func NewDocker(addr string) (Sandbox, error)
```
Connects to Docker Engine API through Tai's Docker proxy. `addr` should be `"tcp://tai-host:2375"`.
### NewK8s
```go
func NewK8s(addr string, opts ...K8sOption) (Sandbox, error)
```
Connects to Kubernetes through Tai's TCP proxy. Each sandbox maps to a single-container Pod.
**Parameters:**
- `addr` — `"host:port"` pointing to Tai's K8s proxy endpoint
- `opts.KubeConfig` — path to kubeconfig file (**required**). Relative paths are resolved to absolute.
- `opts.Namespace` — Kubernetes namespace (default `"default"`)
The constructor overrides the kubeconfig's `server` field to point at `addr`, enables insecure TLS (since Tai does TCP passthrough), and verifies connectivity by querying the namespace.
All pods created by K8s sandbox are labeled with `managed-by: yao-tai-sdk`.
## Types
### CreateOptions
```go
type CreateOptions struct {
Name string // container/pod name
Image string // container image
Cmd []string // entrypoint command
Env map[string]string // environment variables
Binds []string // volume binds (Docker only)
WorkingDir string // working directory
Memory int64 // memory limit in bytes, 0 = no limit
CPUs float64 // CPU limit, 0 = no limit
VNC bool // enable VNC port mapping (Local only)
Ports []PortMapping // port mappings (Docker only)
}
```
### PortMapping
```go
type PortMapping struct {
ContainerPort int // port inside the container
HostPort int // port on the host, 0 = random
HostIP string // host bind address, default "127.0.0.1"
Protocol string // "tcp" (default) or "udp"
}
```
### ContainerInfo
```go
type ContainerInfo struct {
ID string // container/pod ID
Name string // container/pod name
Image string // image name
Status string // "created", "running", "exited", "removing" (Docker)
// "Pending", "Running", "Succeeded", "Failed" (K8s)
IP string // container/pod IP address
Ports []PortMapping // mapped ports (Docker only)
}
```
### ExecOptions
```go
type ExecOptions struct {
WorkDir string // override working directory
Env map[string]string // additional environment variables
}
```
### ExecResult
```go
type ExecResult struct {
ExitCode int
Stdout string
Stderr string
}
```
### ListOptions
```go
type ListOptions struct {
All bool // include stopped containers
Labels map[string]string // filter by labels
}
```
### K8sOption
```go
type K8sOption struct {
Namespace string // default "default"
KubeConfig string // path to kubeconfig file (required)
}
```
## Behavioral Differences
| Behavior | Docker (Local/Remote) | K8s |
|----------|----------------------|-----|
| `Create` returns | container ID (hash) | pod name |
| `Start` | starts a stopped container | polls until pod leaves Pending (up to 30s) |
| `Stop` | stops with timeout, container persists | deletes the pod with grace period |
| `Remove(force=true)` | force-removes | deletes with grace period 0 |
| `Exec` | Docker exec API | `kubectl exec` via SPDY |
| `Inspect.Ports` | populated from Docker | always empty |
| `List` | filters by `tai-sdk=true` label | filters by `managed-by=yao-tai-sdk` label |
| `Binds` | supported | not supported |
| `VNC` flag | auto port-maps 6080 on macOS/Windows | not applicable |
## Example
```go
sb, _ := sandbox.NewLocal("")
defer sb.Close()
id, _ := sb.Create(ctx, sandbox.CreateOptions{
Name: "worker",
Image: "alpine:latest",
Cmd: []string{"sleep", "300"},
Env: map[string]string{"FOO": "bar"},
Memory: 256 * 1024 * 1024, // 256 MB
})
sb.Start(ctx, id)
result, _ := sb.Exec(ctx, id, []string{"echo", "$FOO"}, sandbox.ExecOptions{})
fmt.Println(result.Stdout)
sb.Stop(ctx, id, 10*time.Second)
sb.Remove(ctx, id, false)
```

94
tai/docs/vnc.md Normal file
View file

@ -0,0 +1,94 @@
# Package `vnc`
VNC WebSocket URL resolution. Resolves WebSocket URLs for VNC sessions running inside containers, enabling remote desktop access to sandbox environments.
## Interface
```go
type VNC interface {
URL(ctx context.Context, containerID string) (string, error)
Ping(ctx context.Context, containerID string) error
}
```
## Implementations
| Implementation | Constructor | Mode | URL Pattern |
|----------------|-------------|------|-------------|
| **Remote** | `NewRemote(host, port, hc)` | Via Tai VNC router | `ws://tai-host:6080/vnc/{containerID}/ws` |
| **Local** | `NewLocal(sb)` | Direct host port lookup | `ws://127.0.0.1:{hostPort}/ws` |
## Constructors
### NewRemote
```go
func NewRemote(host string, port int, hc *http.Client) VNC
```
Creates a VNC that routes through Tai's VNC WebSocket router.
- `host` — Tai server hostname/IP
- `port` — Tai VNC router port (default 6080)
- `hc` — custom HTTP client for Ping, `nil` uses `http.DefaultClient`
### NewLocal
```go
func NewLocal(sb sandbox.Sandbox) VNC
```
Creates a VNC that resolves URLs by inspecting the container's port mappings. Looks for container port **6080** (the standard noVNC port) in the port mappings.
Returns an error if port 6080 is not mapped. On macOS and Windows (Docker Desktop), the Local sandbox automatically maps port 6080 when `CreateOptions.VNC` is `true`.
## Methods
### URL
```go
URL(ctx context.Context, containerID string) (string, error)
```
Returns a WebSocket URL for connecting to the container's VNC session.
**Remote:** `ws://tai-host:6080/vnc/abc123/ws`
**Local:** `ws://127.0.0.1:32769/ws`
### Ping
```go
Ping(ctx context.Context, containerID string) error
```
Checks if the VNC endpoint is reachable by making an HTTP GET request to the WebSocket URL. Useful for verifying that the VNC server inside the container is ready before connecting a client.
- **Remote**: sends GET to `http://tai-host:6080/vnc/{containerID}/ws`
- **Local**: resolves the host port via Inspect, then sends GET
## Example
```go
c, _ := tai.New("tai://192.168.1.100")
defer c.Close()
// Create a sandbox with VNC enabled
id, _ := c.Sandbox().Create(ctx, sandbox.CreateOptions{
Name: "desktop",
Image: "yaoapp/sandbox-claude:latest",
VNC: true,
})
c.Sandbox().Start(ctx, id)
// Wait for VNC to be ready
for i := 0; i < 10; i++ {
if err := c.VNC().Ping(ctx, id); err == nil {
break
}
time.Sleep(time.Second)
}
// Get the WebSocket URL for a noVNC client
url, _ := c.VNC().URL(ctx, id)
fmt.Println(url) // ws://192.168.1.100:6080/vnc/desktop/ws
```

120
tai/docs/volume.md Normal file
View file

@ -0,0 +1,120 @@
# Package `volume`
File IO and directory synchronization. Provides a `Volume` interface with two implementations:
| Implementation | Constructor | Backend | Mode |
|----------------|-------------|---------|------|
| **Local** | `NewLocal(root)` | Direct filesystem | Local |
| **Remote** | `NewRemote(conn)` | gRPC to Tai :9100 | Remote |
## Interface
```go
type Volume interface {
ReadFile(ctx context.Context, sessionID, path string) (data []byte, perm os.FileMode, err error)
WriteFile(ctx context.Context, sessionID, path string, data []byte, perm os.FileMode) error
Stat(ctx context.Context, sessionID, path string) (*FileInfo, error)
ListDir(ctx context.Context, sessionID, path string) ([]FileInfo, error)
Remove(ctx context.Context, sessionID, path string, recursive bool) error
Rename(ctx context.Context, sessionID, oldPath, newPath string) error
MkdirAll(ctx context.Context, sessionID, path string) error
SyncPush(ctx context.Context, sessionID, localDir string, opts ...SyncOption) (*SyncResult, error)
SyncPull(ctx context.Context, sessionID, localDir string, opts ...SyncOption) (*SyncResult, error)
Close() error
}
```
All paths are **relative** to the session's workspace root. The `sessionID` identifies the workspace partition — in Local mode this maps to `<root>/<sessionID>/`, in Remote mode the Tai server manages the path.
## Constructors
### NewLocal
```go
func NewLocal(root string) Volume
```
Creates a Volume backed by the local filesystem. Files are stored under `<root>/<sessionID>/`.
### NewRemote
```go
func NewRemote(conn *grpc.ClientConn) Volume
```
Creates a Volume backed by Tai's gRPC Volume service. The connection should target Tai's gRPC port (default 9100). Uses lz4 compression for `SyncPush`/`SyncPull` bulk transfers.
## Types
### FileInfo
```go
type FileInfo struct {
Path string
Size int64
Mtime time.Time
Mode fs.FileMode
IsDir bool
}
```
### SyncResult
```go
type SyncResult struct {
FilesSynced int
BytesTransferred int64
Duration time.Duration
}
```
## Sync Options
```go
volume.WithForceFull() // skip snapshot cache, diff against actual disk
volume.WithExcludes("*.log", ".DS_Store") // glob patterns to exclude
```
## File Operations
| Method | Description |
|--------|-------------|
| `ReadFile` | Read file contents and permissions |
| `WriteFile` | Write file with specified permissions (creates parent dirs) |
| `Stat` | Get file/directory metadata |
| `ListDir` | List directory contents (one level) |
| `Remove` | Delete file or directory (`recursive=true` for tree) |
| `Rename` | Move/rename a file or directory |
| `MkdirAll` | Create directory tree |
## Sync Operations
| Method | Direction | Description |
|--------|-----------|-------------|
| `SyncPush` | local → remote | Upload a local directory to the session workspace |
| `SyncPull` | remote → local | Download the session workspace to a local directory |
Both sync methods use snapshot-based diffing to transfer only changed files. Use `WithForceFull()` to bypass the cache and force a full transfer.
Remote sync uses **lz4 compression** on the wire, streaming files via gRPC bidirectional streaming.
## Example
```go
vol := volume.NewLocal("/data/volumes")
defer vol.Close()
// Write a file
vol.WriteFile(ctx, "session-1", "main.py", []byte("print('hi')"), 0644)
// Read it back
data, perm, _ := vol.ReadFile(ctx, "session-1", "main.py")
// Sync a local directory to the session
result, _ := vol.SyncPush(ctx, "session-1", "/tmp/project",
volume.WithExcludes("node_modules", ".git"),
)
fmt.Printf("synced %d files (%d bytes)\n", result.FilesSynced, result.BytesTransferred)
```

93
tai/docs/workspace.md Normal file
View file

@ -0,0 +1,93 @@
# Package `workspace`
Provides an `fs.FS`-compatible filesystem abstraction over `volume.Volume`. This allows session workspaces to be used with any Go standard library function that accepts `fs.FS`, such as `fs.WalkDir`, `template.ParseFS`, or `http.FS`.
## Interface
```go
type FS interface {
fs.FS // Open(name) (fs.File, error)
fs.StatFS // Stat(name) (fs.FileInfo, error)
fs.ReadFileFS // ReadFile(name) ([]byte, error)
fs.ReadDirFS // ReadDir(name) ([]fs.DirEntry, error)
io.Closer // Close() error
WriteFile(name string, data []byte, perm os.FileMode) error
Remove(name string) error
RemoveAll(name string) error
Rename(oldname, newname string) error
MkdirAll(name string, perm os.FileMode) error
}
```
## Constructor
```go
func New(vol volume.Volume, sessionID string) FS
```
Creates an FS backed by the given Volume for the specified session. The returned FS works transparently whether `vol` is Local or Remote.
Typically accessed through the top-level client:
```go
c, _ := tai.New("tai://host")
ws := c.Workspace("session-123")
```
## Read Operations (fs.FS compatible)
All read operations comply with the `fs.FS` contract. Paths must be valid according to `fs.ValidPath` — forward slashes, no leading slash, no `..` segments.
| Method | Standard Interface | Description |
|--------|--------------------|-------------|
| `Open(name)` | `fs.FS` | Opens a file or directory |
| `Stat(name)` | `fs.StatFS` | Returns file metadata |
| `ReadFile(name)` | `fs.ReadFileFS` | Reads entire file contents |
| `ReadDir(name)` | `fs.ReadDirFS` | Lists directory entries |
`Open` returns an in-memory `fs.File` for regular files (entire content loaded on open) and a directory handle for directories.
## Write Operations
| Method | Description |
|--------|-------------|
| `WriteFile(name, data, perm)` | Write file contents with permissions |
| `Remove(name)` | Delete a single file or empty directory |
| `RemoveAll(name)` | Delete a file or directory tree recursively |
| `Rename(old, new)` | Move/rename a file or directory |
| `MkdirAll(name, perm)` | Create directory tree (perm currently unused) |
## Example
```go
c, _ := tai.New("tai://192.168.1.100")
defer c.Close()
ws := c.Workspace("project-abc")
// Write files
ws.WriteFile("src/main.go", []byte("package main"), 0644)
ws.MkdirAll("src/utils", 0755)
// Read with standard fs.FS
data, _ := fs.ReadFile(ws, "src/main.go")
// Walk the tree
fs.WalkDir(ws, ".", func(path string, d fs.DirEntry, err error) error {
fmt.Println(path)
return nil
})
// Use with Go templates
tmpl, _ := template.ParseFS(ws, "templates/*.html")
// Clean up
ws.RemoveAll("src")
```
## Implementation Notes
- `Open` on a regular file reads the entire content into memory. For large files, prefer `ReadFile` or Volume's `ReadFile` directly.
- `Close()` is a no-op — the underlying Volume's lifecycle is managed by the `tai.Client`.
- Path validation follows `fs.ValidPath` rules. Invalid paths return `fs.ErrInvalid`.

147
tai/grpc/auth.go Normal file
View file

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

263
tai/grpc/cmd/main.go Normal file
View file

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

240
tai/grpc/grpc.go Normal file
View file

@ -0,0 +1,240 @@
package grpc
import (
"context"
"fmt"
"io"
"os"
"github.com/yaoapp/yao/grpc/pb"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
)
// Client wraps a gRPC connection to a Yao server (direct or via Tai relay).
// TokenManager handles auth metadata attachment and token refresh automatically.
type Client struct {
conn *grpc.ClientConn
svc pb.YaoClient
token *TokenManager
}
// NewFromEnv reads YAO_GRPC_ADDR (required) and token env vars, dials the
// gRPC server, and returns a connected Client.
func NewFromEnv() (*Client, error) {
addr := os.Getenv("YAO_GRPC_ADDR")
if addr == "" {
return nil, fmt.Errorf("YAO_GRPC_ADDR is required")
}
tm, err := NewTokenManagerFromEnv()
if err != nil {
return nil, err
}
return Dial(addr, tm)
}
// Dial connects to the gRPC server at addr with the given TokenManager.
func Dial(addr string, tm *TokenManager) (*Client, error) {
opts := []grpc.DialOption{
grpc.WithTransportCredentials(insecure.NewCredentials()),
}
if tm != nil {
opts = append(opts,
grpc.WithUnaryInterceptor(tm.UnaryInterceptor()),
grpc.WithStreamInterceptor(tm.StreamInterceptor()),
)
}
conn, err := grpc.NewClient(addr, opts...)
if err != nil {
return nil, fmt.Errorf("dial %s: %w", addr, err)
}
return &Client{
conn: conn,
svc: pb.NewYaoClient(conn),
token: tm,
}, nil
}
// Close releases the gRPC connection.
func (c *Client) Close() error {
if c.conn != nil {
return c.conn.Close()
}
return nil
}
// Conn returns the underlying gRPC connection.
func (c *Client) Conn() *grpc.ClientConn { return c.conn }
// TokenManager returns the client's token manager.
func (c *Client) TokenManager() *TokenManager { return c.token }
// --- Base ---
// Run executes a Yao process and returns the JSON-encoded result.
func (c *Client) Run(ctx context.Context, process string, args []byte, timeout int32) ([]byte, error) {
resp, err := c.svc.Run(ctx, &pb.RunRequest{
Process: process,
Args: args,
Timeout: timeout,
})
if err != nil {
return nil, err
}
return resp.Data, nil
}
// Shell executes a system command and returns stdout, stderr, exit code.
func (c *Client) Shell(ctx context.Context, command string, args []string, env map[string]string, timeout int32) (*pb.ShellResponse, error) {
return c.svc.Shell(ctx, &pb.ShellRequest{
Command: command,
Args: args,
Env: env,
Timeout: timeout,
})
}
// --- API ---
// API proxies an HTTP request through the gRPC gateway.
func (c *Client) API(ctx context.Context, method, path string, headers map[string]string, body []byte) (*pb.APIResponse, error) {
return c.svc.API(ctx, &pb.APIRequest{
Method: method,
Path: path,
Headers: headers,
Body: body,
})
}
// --- MCP ---
// MCPListTools lists available MCP tools for a session.
func (c *Client) MCPListTools(ctx context.Context, sessionID string) ([]byte, error) {
resp, err := c.svc.MCPListTools(ctx, &pb.MCPListRequest{SessionId: sessionID})
if err != nil {
return nil, err
}
return resp.Tools, nil
}
// MCPCallTool calls an MCP tool and returns the JSON result.
func (c *Client) MCPCallTool(ctx context.Context, sessionID, tool string, arguments []byte) ([]byte, error) {
resp, err := c.svc.MCPCallTool(ctx, &pb.MCPCallRequest{
SessionId: sessionID,
Tool: tool,
Arguments: arguments,
})
if err != nil {
return nil, err
}
return resp.Result, nil
}
// MCPListResources lists available MCP resources for a session.
func (c *Client) MCPListResources(ctx context.Context, sessionID string) ([]byte, error) {
resp, err := c.svc.MCPListResources(ctx, &pb.MCPListRequest{SessionId: sessionID})
if err != nil {
return nil, err
}
return resp.Resources, nil
}
// MCPReadResource reads an MCP resource by URI.
func (c *Client) MCPReadResource(ctx context.Context, sessionID, uri string) ([]byte, error) {
resp, err := c.svc.MCPReadResource(ctx, &pb.MCPResourceRequest{
SessionId: sessionID,
Uri: uri,
})
if err != nil {
return nil, err
}
return resp.Contents, nil
}
// --- LLM ---
// ChatCompletions sends a chat completion request and returns the result.
func (c *Client) ChatCompletions(ctx context.Context, connector string, messages, options []byte) ([]byte, error) {
resp, err := c.svc.ChatCompletions(ctx, &pb.ChatRequest{
Connector: connector,
Messages: messages,
Options: options,
})
if err != nil {
return nil, err
}
return resp.Data, nil
}
// ChatCompletionsStream sends a streaming chat completion request.
// The callback receives each chunk's data; return a non-nil error to stop.
func (c *Client) ChatCompletionsStream(ctx context.Context, connector string, messages, options []byte, cb func(data []byte, done bool) error) error {
stream, err := c.svc.ChatCompletionsStream(ctx, &pb.ChatRequest{
Connector: connector,
Messages: messages,
Options: options,
})
if err != nil {
return err
}
for {
chunk, err := stream.Recv()
if err == io.EOF {
return nil
}
if err != nil {
return err
}
if err := cb(chunk.Data, chunk.Done); err != nil {
return err
}
if chunk.Done {
return nil
}
}
}
// --- Agent ---
// AgentStream calls an agent with streaming response.
// The callback receives each chunk's data; return a non-nil error to stop.
func (c *Client) AgentStream(ctx context.Context, assistantID string, messages, options []byte, cb func(data []byte, done bool) error) error {
stream, err := c.svc.AgentStream(ctx, &pb.AgentRequest{
AssistantId: assistantID,
Messages: messages,
Options: options,
})
if err != nil {
return err
}
for {
chunk, err := stream.Recv()
if err == io.EOF {
return nil
}
if err != nil {
return err
}
if err := cb(chunk.Data, chunk.Done); err != nil {
return err
}
if chunk.Done {
return nil
}
}
}
// --- Health ---
// Healthz checks the server health.
func (c *Client) Healthz(ctx context.Context) (string, error) {
resp, err := c.svc.Healthz(ctx, &pb.Empty{})
if err != nil {
return "", err
}
return resp.Status, nil
}

175
tai/grpc/grpc_test.go Normal file
View file

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

420
tai/grpc/integration_test.go Normal file
View file

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