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

Implement gRPC support in the Yao SDK

Browse source 
- Add gRPC server configuration to the application, allowing for gRPC communication.
- Introduce new Makefile targets for gRPC unit testing and proto code generation.
- Update CI workflows to include gRPC tests with SQLite as the transport layer.
- Refactor the sandbox design to support multi-node capabilities and improve isolation.
- Enhance the service layer to facilitate internal request forwarding for gRPC APIs.

This commit lays the groundwork for integrating gRPC into the Yao SDK, improving performance and scalability.
This commit is contained in:
Max 2026-03-04 13:17:48 +08:00
parent 7772cc588f
commit 6e68efaba3
38 changed files with 7556 additions and 1348 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1722,3 +1722,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!'
});

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

@ -1262,3 +1262,110 @@ jobs:
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)

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 数据库配置

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 |

274
grpc/IMPL.md Normal file
View file

@ -0,0 +1,274 @@
# 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. | ⏳ Pending |
| Tai `server/server.go` | Remove `YaoUpstream` from `Config`. Gateway init no longer needs an address. | ⏳ Pending |
Connection cache: `sync.Map[string, *grpc.ClientConn]` — lazy dial on first request per upstream, reuse thereafter. No eviction needed (upstream count ≈ 1 in practice). `GracefulStop` 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/grpc.go` | `Dial(YAO_GRPC_ADDR)`, method wrappers mirroring server | ⏳ Pending |
| `tai/grpc/auth.go` | Read `YAO_TOKEN` / `YAO_REFRESH_TOKEN` / `YAO_SANDBOX_ID` / `YAO_GRPC_UPSTREAM` from env. Attach as metadata on every call: Bearer token, `x-refresh-token`, `x-sandbox-id`, `x-grpc-upstream` (if set, for Tai relay). Read `SendHeader` for rotated tokens, update in memory. | ⏳ Pending |
| `tai/grpc/cmd/main.go` | Stdio MCP server: JSON-RPC → gRPC. Replaces `yao-bridge`. `yao-grpc version` prints version/commit/build time (via `-ldflags`), for container debugging. | ⏳ Pending |
| `tai/grpc/grpc_test.go` | Tests | ⏳ Pending |
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 — backend (`yao login`) ⏳
Depends on: Phase 1. Independent — can parallel with Phase 2-5.
| Task | Detail | Status |
|------|--------|--------|
| `oauth/device.go` | Implement `DeviceAuthorization()` — generate `device_code` + `user_code`, store with expiry | ⏳ Pending |
| `oauth/token.go` | Device code store/get/consume helpers | ⏳ Pending |
| `oauth/core.go` | Add `GrantTypeDeviceCode` case → `handleDeviceCodeGrant()` (poll returns `authorization_pending` / token) | ⏳ Pending |
| `cmd/yao/login.go` | `yao login --server <url>` → device flow → poll token endpoint → save `~/.yao/credentials` | ⏳ Pending |
| `cmd/yao/logout.go` | Revoke + delete credentials | ⏳ Pending |
| `cmd/yao/run.go` | Credentials exist → gRPC; otherwise local. Non-silent mode prints `⟶ user@host (gRPC)` header before execution (same line position as existing `Run: process.name`). Silent mode (`-s`) keeps pure output — no connection info, for shell scripting. | ⏳ Pending |
Deliverable: `yao login` + `yao run` via gRPC (backend complete, auth page in Phase 7).
### Phase 7: Device Flow — CUI auth page (frontend) ⏳
Depends on: Phase 6 (backend endpoints ready). This is a **frontend-only** task in the 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` and clicks Authorize. Uses `AuthLayout` + `AuthInput` + `AuthButton` from existing `pages/auth/components/`. | ⏳ Pending |
| `pages/auth/device/index.less` | Styles, follow `pages/auth/entry/index.less` pattern | ⏳ Pending |
**Implementation details:**
- Framework: React + UmiJS Max + Ant Design + MobX (same as all auth pages)
- Layout: Wrap with `AuthLayout` (logo + theme switch), same as `/auth/entry`
- Components reuse: `AuthInput` for `user_code` input, `AuthButton` for submit, from `pages/auth/components/`
- Page export: `export default observer(DeviceAuth)` (same pattern as `pages/auth/entry/index.tsx`)
- API: `window.$app.openapi` → call backend `POST /oauth/device/authorize` with `{ user_code }`, bearer token from current session
- Auth: User must be logged in (redirect to `/auth/entry` if not). After authorizing, show success message and close/redirect
- i18n: Use `useIntl()` hook for text, support `zh-CN` / `en-US`
- Flow: User opens URL from CLI prompt → logs in if needed → enters user_code → clicks Authorize → backend binds device_code to user → CLI poll gets token
Deliverable: `/auth/device` page in CUI. User can authorize CLI device login from browser.
## V2 Phases
### Phase 8: `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 9: Base streaming handlers ⏳
Depends on: Phase 8.
| 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 (Tai) Phase 6
(handlers) (LLM/Agent) │ (device backend)
▼ │
Phase 5 ▼
(yao-grpc) Phase 7
(CUI auth page)
--- V2 ---
Phase 8 (gou/stream)
│
▼
Phase 9 (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)
}

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

@ -0,0 +1,220 @@
package testutils
import (
"context"
"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 = "127.0.0.1"
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]
}
// 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)
}

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
}

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 {