From 6e68efaba3bd6fe757fa0a1618c7478704ce79ed Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 4 Mar 2026 13:17:48 +0800 Subject: [PATCH] Implement gRPC support in the Yao SDK - 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. --- .github/workflows/pr-test.yml | 169 ++++ .github/workflows/unit-test.yml | 107 ++ Makefile | 43 +- agent/context/grpc.go | 153 +++ agent/llm/jsapi.go | 6 + cmd/start.go | 15 + config/types.go | 56 +- grpc/DESIGN.md | 448 ++++++++ grpc/IMPL.md | 274 +++++ grpc/TEST.md | 426 ++++++++ grpc/agent/agent.go | 93 ++ grpc/agent/agent_test.go | 206 ++++ grpc/api/api.go | 69 ++ grpc/api/api_test.go | 135 +++ grpc/auth/endpoint.go | 66 ++ grpc/auth/endpoint_test.go | 136 +++ grpc/auth/guard.go | 169 ++++ grpc/auth/guard_test.go | 169 ++++ grpc/auth/scope.go | 14 + grpc/grpc.go | 173 ++++ grpc/health/health.go | 15 + grpc/health/health_test.go | 22 + grpc/llm/llm.go | 165 +++ grpc/llm/llm_test.go | 265 +++++ grpc/mcp/mcp.go | 102 ++ grpc/mcp/mcp_test.go | 264 +++++ grpc/pb/yao.pb.go | 1316 ++++++++++++++++++++++++ grpc/pb/yao.proto | 149 +++ grpc/pb/yao_grpc.pb.go | 606 +++++++++++ grpc/run/run.go | 79 ++ grpc/run/run_test.go | 155 +++ grpc/shell/shell.go | 91 ++ grpc/shell/shell_test.go | 166 +++ grpc/tests/testutils/testutils.go | 220 ++++ openapi/oauth/authenticate.go | 207 ++++ sandbox/DESIGN.md | 1570 +++++------------------------ sandbox/SPEC.md | 579 +++++++++++ service/service.go | 6 + 38 files changed, 7556 insertions(+), 1348 deletions(-) create mode 100644 agent/context/grpc.go create mode 100644 grpc/DESIGN.md create mode 100644 grpc/IMPL.md create mode 100644 grpc/TEST.md create mode 100644 grpc/agent/agent.go create mode 100644 grpc/agent/agent_test.go create mode 100644 grpc/api/api.go create mode 100644 grpc/api/api_test.go create mode 100644 grpc/auth/endpoint.go create mode 100644 grpc/auth/endpoint_test.go create mode 100644 grpc/auth/guard.go create mode 100644 grpc/auth/guard_test.go create mode 100644 grpc/auth/scope.go create mode 100644 grpc/grpc.go create mode 100644 grpc/health/health.go create mode 100644 grpc/health/health_test.go create mode 100644 grpc/llm/llm.go create mode 100644 grpc/llm/llm_test.go create mode 100644 grpc/mcp/mcp.go create mode 100644 grpc/mcp/mcp_test.go create mode 100644 grpc/pb/yao.pb.go create mode 100644 grpc/pb/yao.proto create mode 100644 grpc/pb/yao_grpc.pb.go create mode 100644 grpc/run/run.go create mode 100644 grpc/run/run_test.go create mode 100644 grpc/shell/shell.go create mode 100644 grpc/shell/shell_test.go create mode 100644 grpc/tests/testutils/testutils.go create mode 100644 openapi/oauth/authenticate.go create mode 100644 sandbox/SPEC.md diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index c9e07eda..c3c8ab5f 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -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!' + }); diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index 7d3696e2..1c5b36a4 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -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 }} diff --git a/Makefile b/Makefile index 6c704606..68790267 100644 --- a/Makefile +++ b/Makefile @@ -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: diff --git a/agent/context/grpc.go b/agent/context/grpc.go new file mode 100644 index 00000000..f4f2e8f1 --- /dev/null +++ b/agent/context/grpc.go @@ -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 +} diff --git a/agent/llm/jsapi.go b/agent/llm/jsapi.go index 4acbcddd..2439dd7f 100644 --- a/agent/llm/jsapi.go +++ b/agent/llm/jsapi.go @@ -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) diff --git a/cmd/start.go b/cmd/start.go index d663505b..18d2e367 100644 --- a/cmd/start.go +++ b/cmd/start.go @@ -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 { diff --git a/config/types.go b/config/types.go index 16e81bb5..8618fabc 100644 --- a/config/types.go +++ b/config/types.go @@ -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 (/plugins /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 (/plugins /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 ๆ•ฐๆฎๅบ“้…็ฝฎ diff --git a/grpc/DESIGN.md b/grpc/DESIGN.md new file mode 100644 index 00000000..d5dbc976 --- /dev/null +++ b/grpc/DESIGN.md @@ -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 | diff --git a/grpc/IMPL.md b/grpc/IMPL.md new file mode 100644 index 00000000..fefc76f0 --- /dev/null +++ b/grpc/IMPL.md @@ -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 ` โ†’ 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) +``` diff --git a/grpc/TEST.md b/grpc/TEST.md new file mode 100644 index 00000000..3f947753 --- /dev/null +++ b/grpc/TEST.md @@ -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/ +``` diff --git a/grpc/agent/agent.go b/grpc/agent/agent.go new file mode 100644 index 00000000..bd29cf3f --- /dev/null +++ b/grpc/agent/agent.go @@ -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() {} diff --git a/grpc/agent/agent_test.go b/grpc/agent/agent_test.go new file mode 100644 index 00000000..d06edeea --- /dev/null +++ b/grpc/agent/agent_test.go @@ -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) +} diff --git a/grpc/api/api.go b/grpc/api/api.go new file mode 100644 index 00000000..fcf5d91f --- /dev/null +++ b/grpc/api/api.go @@ -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 +} diff --git a/grpc/api/api_test.go b/grpc/api/api_test.go new file mode 100644 index 00000000..cd03df17 --- /dev/null +++ b/grpc/api/api_test.go @@ -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) + } +} diff --git a/grpc/auth/endpoint.go b/grpc/auth/endpoint.go new file mode 100644 index 00000000..67fcaf6e --- /dev/null +++ b/grpc/auth/endpoint.go @@ -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" + } +} diff --git a/grpc/auth/endpoint_test.go b/grpc/auth/endpoint_test.go new file mode 100644 index 00000000..edbe3c1e --- /dev/null +++ b/grpc/auth/endpoint_test.go @@ -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) +} diff --git a/grpc/auth/guard.go b/grpc/auth/guard.go new file mode 100644 index 00000000..c2abccf3 --- /dev/null +++ b/grpc/auth/guard.go @@ -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 +} diff --git a/grpc/auth/guard_test.go b/grpc/auth/guard_test.go new file mode 100644 index 00000000..093d03d6 --- /dev/null +++ b/grpc/auth/guard_test.go @@ -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()) +} diff --git a/grpc/auth/scope.go b/grpc/auth/scope.go new file mode 100644 index 00000000..4957c482 --- /dev/null +++ b/grpc/auth/scope.go @@ -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/"}}, + ) +} diff --git a/grpc/grpc.go b/grpc/grpc.go new file mode 100644 index 00000000..22755275 --- /dev/null +++ b/grpc/grpc.go @@ -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 +} diff --git a/grpc/health/health.go b/grpc/health/health.go new file mode 100644 index 00000000..fadf1ce6 --- /dev/null +++ b/grpc/health/health.go @@ -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 +} diff --git a/grpc/health/health_test.go b/grpc/health/health_test.go new file mode 100644 index 00000000..740e83d9 --- /dev/null +++ b/grpc/health/health_test.go @@ -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) +} diff --git a/grpc/llm/llm.go b/grpc/llm/llm.go new file mode 100644 index 00000000..93461b59 --- /dev/null +++ b/grpc/llm/llm.go @@ -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 +} diff --git a/grpc/llm/llm_test.go b/grpc/llm/llm_test.go new file mode 100644 index 00000000..0106adc6 --- /dev/null +++ b/grpc/llm/llm_test.go @@ -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) +} diff --git a/grpc/mcp/mcp.go b/grpc/mcp/mcp.go new file mode 100644 index 00000000..62ac1a08 --- /dev/null +++ b/grpc/mcp/mcp.go @@ -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 +} diff --git a/grpc/mcp/mcp_test.go b/grpc/mcp/mcp_test.go new file mode 100644 index 00000000..2b6a8e9a --- /dev/null +++ b/grpc/mcp/mcp_test.go @@ -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()) +} diff --git a/grpc/pb/yao.pb.go b/grpc/pb/yao.pb.go new file mode 100644 index 00000000..cea122a1 --- /dev/null +++ b/grpc/pb/yao.pb.go @@ -0,0 +1,1316 @@ +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.36.11 +// protoc v4.25.0 +// source: yao.proto + +package pb + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + reflect "reflect" + sync "sync" + unsafe "unsafe" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +type RunRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Process string `protobuf:"bytes,1,opt,name=process,proto3" json:"process,omitempty"` + Args []byte `protobuf:"bytes,2,opt,name=args,proto3" json:"args,omitempty"` // JSON-encoded argument array + Timeout int32 `protobuf:"varint,3,opt,name=timeout,proto3" json:"timeout,omitempty"` // seconds, 0 = server default + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RunRequest) Reset() { + *x = RunRequest{} + mi := &file_yao_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RunRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RunRequest) ProtoMessage() {} + +func (x *RunRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[0] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RunRequest.ProtoReflect.Descriptor instead. +func (*RunRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{0} +} + +func (x *RunRequest) GetProcess() string { + if x != nil { + return x.Process + } + return "" +} + +func (x *RunRequest) GetArgs() []byte { + if x != nil { + return x.Args + } + return nil +} + +func (x *RunRequest) GetTimeout() int32 { + if x != nil { + return x.Timeout + } + return 0 +} + +type RunResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Data []byte `protobuf:"bytes,1,opt,name=data,proto3" json:"data,omitempty"` // JSON-encoded result + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RunResponse) Reset() { + *x = RunResponse{} + mi := &file_yao_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RunResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RunResponse) ProtoMessage() {} + +func (x *RunResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[1] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RunResponse.ProtoReflect.Descriptor instead. +func (*RunResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{1} +} + +func (x *RunResponse) GetData() []byte { + if x != nil { + return x.Data + } + return nil +} + +type Chunk struct { + state protoimpl.MessageState `protogen:"open.v1"` + Data []byte `protobuf:"bytes,1,opt,name=data,proto3" json:"data,omitempty"` + Done bool `protobuf:"varint,2,opt,name=done,proto3" json:"done,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *Chunk) Reset() { + *x = Chunk{} + mi := &file_yao_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *Chunk) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*Chunk) ProtoMessage() {} + +func (x *Chunk) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[2] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use Chunk.ProtoReflect.Descriptor instead. +func (*Chunk) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{2} +} + +func (x *Chunk) GetData() []byte { + if x != nil { + return x.Data + } + return nil +} + +func (x *Chunk) GetDone() bool { + if x != nil { + return x.Done + } + return false +} + +type ShellRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Command string `protobuf:"bytes,1,opt,name=command,proto3" json:"command,omitempty"` + Args []string `protobuf:"bytes,2,rep,name=args,proto3" json:"args,omitempty"` + Env map[string]string `protobuf:"bytes,3,rep,name=env,proto3" json:"env,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + Timeout int32 `protobuf:"varint,4,opt,name=timeout,proto3" json:"timeout,omitempty"` // seconds, 0 = default 30s + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ShellRequest) Reset() { + *x = ShellRequest{} + mi := &file_yao_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ShellRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ShellRequest) ProtoMessage() {} + +func (x *ShellRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[3] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ShellRequest.ProtoReflect.Descriptor instead. +func (*ShellRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{3} +} + +func (x *ShellRequest) GetCommand() string { + if x != nil { + return x.Command + } + return "" +} + +func (x *ShellRequest) GetArgs() []string { + if x != nil { + return x.Args + } + return nil +} + +func (x *ShellRequest) GetEnv() map[string]string { + if x != nil { + return x.Env + } + return nil +} + +func (x *ShellRequest) GetTimeout() int32 { + if x != nil { + return x.Timeout + } + return 0 +} + +type ShellResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Stdout []byte `protobuf:"bytes,1,opt,name=stdout,proto3" json:"stdout,omitempty"` + Stderr []byte `protobuf:"bytes,2,opt,name=stderr,proto3" json:"stderr,omitempty"` + ExitCode int32 `protobuf:"varint,3,opt,name=exit_code,json=exitCode,proto3" json:"exit_code,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ShellResponse) Reset() { + *x = ShellResponse{} + mi := &file_yao_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ShellResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ShellResponse) ProtoMessage() {} + +func (x *ShellResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[4] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ShellResponse.ProtoReflect.Descriptor instead. +func (*ShellResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{4} +} + +func (x *ShellResponse) GetStdout() []byte { + if x != nil { + return x.Stdout + } + return nil +} + +func (x *ShellResponse) GetStderr() []byte { + if x != nil { + return x.Stderr + } + return nil +} + +func (x *ShellResponse) GetExitCode() int32 { + if x != nil { + return x.ExitCode + } + return 0 +} + +type APIRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Method string `protobuf:"bytes,1,opt,name=method,proto3" json:"method,omitempty"` // HTTP method + Path string `protobuf:"bytes,2,opt,name=path,proto3" json:"path,omitempty"` // openapi path + Headers map[string]string `protobuf:"bytes,3,rep,name=headers,proto3" json:"headers,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + Body []byte `protobuf:"bytes,4,opt,name=body,proto3" json:"body,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *APIRequest) Reset() { + *x = APIRequest{} + mi := &file_yao_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *APIRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*APIRequest) ProtoMessage() {} + +func (x *APIRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[5] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use APIRequest.ProtoReflect.Descriptor instead. +func (*APIRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{5} +} + +func (x *APIRequest) GetMethod() string { + if x != nil { + return x.Method + } + return "" +} + +func (x *APIRequest) GetPath() string { + if x != nil { + return x.Path + } + return "" +} + +func (x *APIRequest) GetHeaders() map[string]string { + if x != nil { + return x.Headers + } + return nil +} + +func (x *APIRequest) GetBody() []byte { + if x != nil { + return x.Body + } + return nil +} + +type APIResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Status int32 `protobuf:"varint,1,opt,name=status,proto3" json:"status,omitempty"` // HTTP status code + Headers map[string]string `protobuf:"bytes,2,rep,name=headers,proto3" json:"headers,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + Body []byte `protobuf:"bytes,3,opt,name=body,proto3" json:"body,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *APIResponse) Reset() { + *x = APIResponse{} + mi := &file_yao_proto_msgTypes[6] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *APIResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*APIResponse) ProtoMessage() {} + +func (x *APIResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[6] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use APIResponse.ProtoReflect.Descriptor instead. +func (*APIResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{6} +} + +func (x *APIResponse) GetStatus() int32 { + if x != nil { + return x.Status + } + return 0 +} + +func (x *APIResponse) GetHeaders() map[string]string { + if x != nil { + return x.Headers + } + return nil +} + +func (x *APIResponse) GetBody() []byte { + if x != nil { + return x.Body + } + return nil +} + +type MCPListRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPListRequest) Reset() { + *x = MCPListRequest{} + mi := &file_yao_proto_msgTypes[7] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPListRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPListRequest) ProtoMessage() {} + +func (x *MCPListRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[7] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPListRequest.ProtoReflect.Descriptor instead. +func (*MCPListRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{7} +} + +func (x *MCPListRequest) GetSessionId() string { + if x != nil { + return x.SessionId + } + return "" +} + +type MCPListResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Tools []byte `protobuf:"bytes,1,opt,name=tools,proto3" json:"tools,omitempty"` // JSON array of tool definitions + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPListResponse) Reset() { + *x = MCPListResponse{} + mi := &file_yao_proto_msgTypes[8] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPListResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPListResponse) ProtoMessage() {} + +func (x *MCPListResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[8] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPListResponse.ProtoReflect.Descriptor instead. +func (*MCPListResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{8} +} + +func (x *MCPListResponse) GetTools() []byte { + if x != nil { + return x.Tools + } + return nil +} + +type MCPCallRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` + Tool string `protobuf:"bytes,2,opt,name=tool,proto3" json:"tool,omitempty"` + Arguments []byte `protobuf:"bytes,3,opt,name=arguments,proto3" json:"arguments,omitempty"` // JSON-encoded arguments + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPCallRequest) Reset() { + *x = MCPCallRequest{} + mi := &file_yao_proto_msgTypes[9] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPCallRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPCallRequest) ProtoMessage() {} + +func (x *MCPCallRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[9] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPCallRequest.ProtoReflect.Descriptor instead. +func (*MCPCallRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{9} +} + +func (x *MCPCallRequest) GetSessionId() string { + if x != nil { + return x.SessionId + } + return "" +} + +func (x *MCPCallRequest) GetTool() string { + if x != nil { + return x.Tool + } + return "" +} + +func (x *MCPCallRequest) GetArguments() []byte { + if x != nil { + return x.Arguments + } + return nil +} + +type MCPCallResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Result []byte `protobuf:"bytes,1,opt,name=result,proto3" json:"result,omitempty"` // JSON-encoded result + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPCallResponse) Reset() { + *x = MCPCallResponse{} + mi := &file_yao_proto_msgTypes[10] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPCallResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPCallResponse) ProtoMessage() {} + +func (x *MCPCallResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[10] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPCallResponse.ProtoReflect.Descriptor instead. +func (*MCPCallResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{10} +} + +func (x *MCPCallResponse) GetResult() []byte { + if x != nil { + return x.Result + } + return nil +} + +type MCPResourcesResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Resources []byte `protobuf:"bytes,1,opt,name=resources,proto3" json:"resources,omitempty"` // JSON array of resource definitions + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPResourcesResponse) Reset() { + *x = MCPResourcesResponse{} + mi := &file_yao_proto_msgTypes[11] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPResourcesResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPResourcesResponse) ProtoMessage() {} + +func (x *MCPResourcesResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[11] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPResourcesResponse.ProtoReflect.Descriptor instead. +func (*MCPResourcesResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{11} +} + +func (x *MCPResourcesResponse) GetResources() []byte { + if x != nil { + return x.Resources + } + return nil +} + +type MCPResourceRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` + Uri string `protobuf:"bytes,2,opt,name=uri,proto3" json:"uri,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPResourceRequest) Reset() { + *x = MCPResourceRequest{} + mi := &file_yao_proto_msgTypes[12] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPResourceRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPResourceRequest) ProtoMessage() {} + +func (x *MCPResourceRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[12] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPResourceRequest.ProtoReflect.Descriptor instead. +func (*MCPResourceRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{12} +} + +func (x *MCPResourceRequest) GetSessionId() string { + if x != nil { + return x.SessionId + } + return "" +} + +func (x *MCPResourceRequest) GetUri() string { + if x != nil { + return x.Uri + } + return "" +} + +type MCPResourceResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Contents []byte `protobuf:"bytes,1,opt,name=contents,proto3" json:"contents,omitempty"` // JSON-encoded resource contents + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPResourceResponse) Reset() { + *x = MCPResourceResponse{} + mi := &file_yao_proto_msgTypes[13] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPResourceResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPResourceResponse) ProtoMessage() {} + +func (x *MCPResourceResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[13] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPResourceResponse.ProtoReflect.Descriptor instead. +func (*MCPResourceResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{13} +} + +func (x *MCPResourceResponse) GetContents() []byte { + if x != nil { + return x.Contents + } + return nil +} + +type ChatRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Connector string `protobuf:"bytes,1,opt,name=connector,proto3" json:"connector,omitempty"` // connector ID + Messages []byte `protobuf:"bytes,2,opt,name=messages,proto3" json:"messages,omitempty"` // JSON-encoded message array + Options []byte `protobuf:"bytes,3,opt,name=options,proto3" json:"options,omitempty"` // JSON-encoded options + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ChatRequest) Reset() { + *x = ChatRequest{} + mi := &file_yao_proto_msgTypes[14] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ChatRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ChatRequest) ProtoMessage() {} + +func (x *ChatRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[14] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ChatRequest.ProtoReflect.Descriptor instead. +func (*ChatRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{14} +} + +func (x *ChatRequest) GetConnector() string { + if x != nil { + return x.Connector + } + return "" +} + +func (x *ChatRequest) GetMessages() []byte { + if x != nil { + return x.Messages + } + return nil +} + +func (x *ChatRequest) GetOptions() []byte { + if x != nil { + return x.Options + } + return nil +} + +type ChatResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Data []byte `protobuf:"bytes,1,opt,name=data,proto3" json:"data,omitempty"` // JSON-encoded completion result + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ChatResponse) Reset() { + *x = ChatResponse{} + mi := &file_yao_proto_msgTypes[15] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ChatResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ChatResponse) ProtoMessage() {} + +func (x *ChatResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[15] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ChatResponse.ProtoReflect.Descriptor instead. +func (*ChatResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{15} +} + +func (x *ChatResponse) GetData() []byte { + if x != nil { + return x.Data + } + return nil +} + +type ChatChunk struct { + state protoimpl.MessageState `protogen:"open.v1"` + Data []byte `protobuf:"bytes,1,opt,name=data,proto3" json:"data,omitempty"` // JSON-encoded chunk + Done bool `protobuf:"varint,2,opt,name=done,proto3" json:"done,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ChatChunk) Reset() { + *x = ChatChunk{} + mi := &file_yao_proto_msgTypes[16] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ChatChunk) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ChatChunk) ProtoMessage() {} + +func (x *ChatChunk) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[16] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ChatChunk.ProtoReflect.Descriptor instead. +func (*ChatChunk) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{16} +} + +func (x *ChatChunk) GetData() []byte { + if x != nil { + return x.Data + } + return nil +} + +func (x *ChatChunk) GetDone() bool { + if x != nil { + return x.Done + } + return false +} + +type AgentRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + AssistantId string `protobuf:"bytes,1,opt,name=assistant_id,json=assistantId,proto3" json:"assistant_id,omitempty"` + Messages []byte `protobuf:"bytes,2,opt,name=messages,proto3" json:"messages,omitempty"` // JSON-encoded message array + Options []byte `protobuf:"bytes,3,opt,name=options,proto3" json:"options,omitempty"` // JSON-encoded options + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *AgentRequest) Reset() { + *x = AgentRequest{} + mi := &file_yao_proto_msgTypes[17] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *AgentRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*AgentRequest) ProtoMessage() {} + +func (x *AgentRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[17] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use AgentRequest.ProtoReflect.Descriptor instead. +func (*AgentRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{17} +} + +func (x *AgentRequest) GetAssistantId() string { + if x != nil { + return x.AssistantId + } + return "" +} + +func (x *AgentRequest) GetMessages() []byte { + if x != nil { + return x.Messages + } + return nil +} + +func (x *AgentRequest) GetOptions() []byte { + if x != nil { + return x.Options + } + return nil +} + +// Each chunk carries JSON-serialized agent/output/message.Message. +type AgentChunk struct { + state protoimpl.MessageState `protogen:"open.v1"` + Data []byte `protobuf:"bytes,1,opt,name=data,proto3" json:"data,omitempty"` + Done bool `protobuf:"varint,2,opt,name=done,proto3" json:"done,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *AgentChunk) Reset() { + *x = AgentChunk{} + mi := &file_yao_proto_msgTypes[18] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *AgentChunk) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*AgentChunk) ProtoMessage() {} + +func (x *AgentChunk) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[18] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use AgentChunk.ProtoReflect.Descriptor instead. +func (*AgentChunk) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{18} +} + +func (x *AgentChunk) GetData() []byte { + if x != nil { + return x.Data + } + return nil +} + +func (x *AgentChunk) GetDone() bool { + if x != nil { + return x.Done + } + return false +} + +type Empty struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *Empty) Reset() { + *x = Empty{} + mi := &file_yao_proto_msgTypes[19] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *Empty) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*Empty) ProtoMessage() {} + +func (x *Empty) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[19] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use Empty.ProtoReflect.Descriptor instead. +func (*Empty) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{19} +} + +type HealthzResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Status string `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *HealthzResponse) Reset() { + *x = HealthzResponse{} + mi := &file_yao_proto_msgTypes[20] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *HealthzResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HealthzResponse) ProtoMessage() {} + +func (x *HealthzResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[20] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use HealthzResponse.ProtoReflect.Descriptor instead. +func (*HealthzResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{20} +} + +func (x *HealthzResponse) GetStatus() string { + if x != nil { + return x.Status + } + return "" +} + +var File_yao_proto protoreflect.FileDescriptor + +const file_yao_proto_rawDesc = "" + + "\n" + + "\tyao.proto\x12\x03yao\"T\n" + + "\n" + + "RunRequest\x12\x18\n" + + "\aprocess\x18\x01 \x01(\tR\aprocess\x12\x12\n" + + "\x04args\x18\x02 \x01(\fR\x04args\x12\x18\n" + + "\atimeout\x18\x03 \x01(\x05R\atimeout\"!\n" + + "\vRunResponse\x12\x12\n" + + "\x04data\x18\x01 \x01(\fR\x04data\"/\n" + + "\x05Chunk\x12\x12\n" + + "\x04data\x18\x01 \x01(\fR\x04data\x12\x12\n" + + "\x04done\x18\x02 \x01(\bR\x04done\"\xbc\x01\n" + + "\fShellRequest\x12\x18\n" + + "\acommand\x18\x01 \x01(\tR\acommand\x12\x12\n" + + "\x04args\x18\x02 \x03(\tR\x04args\x12,\n" + + "\x03env\x18\x03 \x03(\v2\x1a.yao.ShellRequest.EnvEntryR\x03env\x12\x18\n" + + "\atimeout\x18\x04 \x01(\x05R\atimeout\x1a6\n" + + "\bEnvEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\\\n" + + "\rShellResponse\x12\x16\n" + + "\x06stdout\x18\x01 \x01(\fR\x06stdout\x12\x16\n" + + "\x06stderr\x18\x02 \x01(\fR\x06stderr\x12\x1b\n" + + "\texit_code\x18\x03 \x01(\x05R\bexitCode\"\xc0\x01\n" + + "\n" + + "APIRequest\x12\x16\n" + + "\x06method\x18\x01 \x01(\tR\x06method\x12\x12\n" + + "\x04path\x18\x02 \x01(\tR\x04path\x126\n" + + "\aheaders\x18\x03 \x03(\v2\x1c.yao.APIRequest.HeadersEntryR\aheaders\x12\x12\n" + + "\x04body\x18\x04 \x01(\fR\x04body\x1a:\n" + + "\fHeadersEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\xae\x01\n" + + "\vAPIResponse\x12\x16\n" + + "\x06status\x18\x01 \x01(\x05R\x06status\x127\n" + + "\aheaders\x18\x02 \x03(\v2\x1d.yao.APIResponse.HeadersEntryR\aheaders\x12\x12\n" + + "\x04body\x18\x03 \x01(\fR\x04body\x1a:\n" + + "\fHeadersEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"/\n" + + "\x0eMCPListRequest\x12\x1d\n" + + "\n" + + "session_id\x18\x01 \x01(\tR\tsessionId\"'\n" + + "\x0fMCPListResponse\x12\x14\n" + + "\x05tools\x18\x01 \x01(\fR\x05tools\"a\n" + + "\x0eMCPCallRequest\x12\x1d\n" + + "\n" + + "session_id\x18\x01 \x01(\tR\tsessionId\x12\x12\n" + + "\x04tool\x18\x02 \x01(\tR\x04tool\x12\x1c\n" + + "\targuments\x18\x03 \x01(\fR\targuments\")\n" + + "\x0fMCPCallResponse\x12\x16\n" + + "\x06result\x18\x01 \x01(\fR\x06result\"4\n" + + "\x14MCPResourcesResponse\x12\x1c\n" + + "\tresources\x18\x01 \x01(\fR\tresources\"E\n" + + "\x12MCPResourceRequest\x12\x1d\n" + + "\n" + + "session_id\x18\x01 \x01(\tR\tsessionId\x12\x10\n" + + "\x03uri\x18\x02 \x01(\tR\x03uri\"1\n" + + "\x13MCPResourceResponse\x12\x1a\n" + + "\bcontents\x18\x01 \x01(\fR\bcontents\"a\n" + + "\vChatRequest\x12\x1c\n" + + "\tconnector\x18\x01 \x01(\tR\tconnector\x12\x1a\n" + + "\bmessages\x18\x02 \x01(\fR\bmessages\x12\x18\n" + + "\aoptions\x18\x03 \x01(\fR\aoptions\"\"\n" + + "\fChatResponse\x12\x12\n" + + "\x04data\x18\x01 \x01(\fR\x04data\"3\n" + + "\tChatChunk\x12\x12\n" + + "\x04data\x18\x01 \x01(\fR\x04data\x12\x12\n" + + "\x04done\x18\x02 \x01(\bR\x04done\"g\n" + + "\fAgentRequest\x12!\n" + + "\fassistant_id\x18\x01 \x01(\tR\vassistantId\x12\x1a\n" + + "\bmessages\x18\x02 \x01(\fR\bmessages\x12\x18\n" + + "\aoptions\x18\x03 \x01(\fR\aoptions\"4\n" + + "\n" + + "AgentChunk\x12\x12\n" + + "\x04data\x18\x01 \x01(\fR\x04data\x12\x12\n" + + "\x04done\x18\x02 \x01(\bR\x04done\"\a\n" + + "\x05Empty\")\n" + + "\x0fHealthzResponse\x12\x16\n" + + "\x06status\x18\x01 \x01(\tR\x06status2\xb8\x05\n" + + "\x03Yao\x12(\n" + + "\x03Run\x12\x0f.yao.RunRequest\x1a\x10.yao.RunResponse\x12'\n" + + "\x06Stream\x12\x0f.yao.RunRequest\x1a\n" + + ".yao.Chunk0\x01\x12.\n" + + "\x05Shell\x12\x11.yao.ShellRequest\x1a\x12.yao.ShellResponse\x12.\n" + + "\vShellStream\x12\x11.yao.ShellRequest\x1a\n" + + ".yao.Chunk0\x01\x12(\n" + + "\x03API\x12\x0f.yao.APIRequest\x1a\x10.yao.APIResponse\x129\n" + + "\fMCPListTools\x12\x13.yao.MCPListRequest\x1a\x14.yao.MCPListResponse\x128\n" + + "\vMCPCallTool\x12\x13.yao.MCPCallRequest\x1a\x14.yao.MCPCallResponse\x12B\n" + + "\x10MCPListResources\x12\x13.yao.MCPListRequest\x1a\x19.yao.MCPResourcesResponse\x12D\n" + + "\x0fMCPReadResource\x12\x17.yao.MCPResourceRequest\x1a\x18.yao.MCPResourceResponse\x126\n" + + "\x0fChatCompletions\x12\x10.yao.ChatRequest\x1a\x11.yao.ChatResponse\x12;\n" + + "\x15ChatCompletionsStream\x12\x10.yao.ChatRequest\x1a\x0e.yao.ChatChunk0\x01\x123\n" + + "\vAgentStream\x12\x11.yao.AgentRequest\x1a\x0f.yao.AgentChunk0\x01\x12+\n" + + "\aHealthz\x12\n" + + ".yao.Empty\x1a\x14.yao.HealthzResponseB\x1fZ\x1dgithub.com/yaoapp/yao/grpc/pbb\x06proto3" + +var ( + file_yao_proto_rawDescOnce sync.Once + file_yao_proto_rawDescData []byte +) + +func file_yao_proto_rawDescGZIP() []byte { + file_yao_proto_rawDescOnce.Do(func() { + file_yao_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_yao_proto_rawDesc), len(file_yao_proto_rawDesc))) + }) + return file_yao_proto_rawDescData +} + +var file_yao_proto_msgTypes = make([]protoimpl.MessageInfo, 24) +var file_yao_proto_goTypes = []any{ + (*RunRequest)(nil), // 0: yao.RunRequest + (*RunResponse)(nil), // 1: yao.RunResponse + (*Chunk)(nil), // 2: yao.Chunk + (*ShellRequest)(nil), // 3: yao.ShellRequest + (*ShellResponse)(nil), // 4: yao.ShellResponse + (*APIRequest)(nil), // 5: yao.APIRequest + (*APIResponse)(nil), // 6: yao.APIResponse + (*MCPListRequest)(nil), // 7: yao.MCPListRequest + (*MCPListResponse)(nil), // 8: yao.MCPListResponse + (*MCPCallRequest)(nil), // 9: yao.MCPCallRequest + (*MCPCallResponse)(nil), // 10: yao.MCPCallResponse + (*MCPResourcesResponse)(nil), // 11: yao.MCPResourcesResponse + (*MCPResourceRequest)(nil), // 12: yao.MCPResourceRequest + (*MCPResourceResponse)(nil), // 13: yao.MCPResourceResponse + (*ChatRequest)(nil), // 14: yao.ChatRequest + (*ChatResponse)(nil), // 15: yao.ChatResponse + (*ChatChunk)(nil), // 16: yao.ChatChunk + (*AgentRequest)(nil), // 17: yao.AgentRequest + (*AgentChunk)(nil), // 18: yao.AgentChunk + (*Empty)(nil), // 19: yao.Empty + (*HealthzResponse)(nil), // 20: yao.HealthzResponse + nil, // 21: yao.ShellRequest.EnvEntry + nil, // 22: yao.APIRequest.HeadersEntry + nil, // 23: yao.APIResponse.HeadersEntry +} +var file_yao_proto_depIdxs = []int32{ + 21, // 0: yao.ShellRequest.env:type_name -> yao.ShellRequest.EnvEntry + 22, // 1: yao.APIRequest.headers:type_name -> yao.APIRequest.HeadersEntry + 23, // 2: yao.APIResponse.headers:type_name -> yao.APIResponse.HeadersEntry + 0, // 3: yao.Yao.Run:input_type -> yao.RunRequest + 0, // 4: yao.Yao.Stream:input_type -> yao.RunRequest + 3, // 5: yao.Yao.Shell:input_type -> yao.ShellRequest + 3, // 6: yao.Yao.ShellStream:input_type -> yao.ShellRequest + 5, // 7: yao.Yao.API:input_type -> yao.APIRequest + 7, // 8: yao.Yao.MCPListTools:input_type -> yao.MCPListRequest + 9, // 9: yao.Yao.MCPCallTool:input_type -> yao.MCPCallRequest + 7, // 10: yao.Yao.MCPListResources:input_type -> yao.MCPListRequest + 12, // 11: yao.Yao.MCPReadResource:input_type -> yao.MCPResourceRequest + 14, // 12: yao.Yao.ChatCompletions:input_type -> yao.ChatRequest + 14, // 13: yao.Yao.ChatCompletionsStream:input_type -> yao.ChatRequest + 17, // 14: yao.Yao.AgentStream:input_type -> yao.AgentRequest + 19, // 15: yao.Yao.Healthz:input_type -> yao.Empty + 1, // 16: yao.Yao.Run:output_type -> yao.RunResponse + 2, // 17: yao.Yao.Stream:output_type -> yao.Chunk + 4, // 18: yao.Yao.Shell:output_type -> yao.ShellResponse + 2, // 19: yao.Yao.ShellStream:output_type -> yao.Chunk + 6, // 20: yao.Yao.API:output_type -> yao.APIResponse + 8, // 21: yao.Yao.MCPListTools:output_type -> yao.MCPListResponse + 10, // 22: yao.Yao.MCPCallTool:output_type -> yao.MCPCallResponse + 11, // 23: yao.Yao.MCPListResources:output_type -> yao.MCPResourcesResponse + 13, // 24: yao.Yao.MCPReadResource:output_type -> yao.MCPResourceResponse + 15, // 25: yao.Yao.ChatCompletions:output_type -> yao.ChatResponse + 16, // 26: yao.Yao.ChatCompletionsStream:output_type -> yao.ChatChunk + 18, // 27: yao.Yao.AgentStream:output_type -> yao.AgentChunk + 20, // 28: yao.Yao.Healthz:output_type -> yao.HealthzResponse + 16, // [16:29] is the sub-list for method output_type + 3, // [3:16] is the sub-list for method input_type + 3, // [3:3] is the sub-list for extension type_name + 3, // [3:3] is the sub-list for extension extendee + 0, // [0:3] is the sub-list for field type_name +} + +func init() { file_yao_proto_init() } +func file_yao_proto_init() { + if File_yao_proto != nil { + return + } + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + GoPackagePath: reflect.TypeOf(x{}).PkgPath(), + RawDescriptor: unsafe.Slice(unsafe.StringData(file_yao_proto_rawDesc), len(file_yao_proto_rawDesc)), + NumEnums: 0, + NumMessages: 24, + NumExtensions: 0, + NumServices: 1, + }, + GoTypes: file_yao_proto_goTypes, + DependencyIndexes: file_yao_proto_depIdxs, + MessageInfos: file_yao_proto_msgTypes, + }.Build() + File_yao_proto = out.File + file_yao_proto_goTypes = nil + file_yao_proto_depIdxs = nil +} diff --git a/grpc/pb/yao.proto b/grpc/pb/yao.proto new file mode 100644 index 00000000..f120151f --- /dev/null +++ b/grpc/pb/yao.proto @@ -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 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 headers = 3; + bytes body = 4; +} + +message APIResponse { + int32 status = 1; // HTTP status code + map 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; +} diff --git a/grpc/pb/yao_grpc.pb.go b/grpc/pb/yao_grpc.pb.go new file mode 100644 index 00000000..db5b07d6 --- /dev/null +++ b/grpc/pb/yao_grpc.pb.go @@ -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", +} diff --git a/grpc/run/run.go b/grpc/run/run.go new file mode 100644 index 00000000..fc5b0820 --- /dev/null +++ b/grpc/run/run.go @@ -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, + }) +} diff --git a/grpc/run/run_test.go b/grpc/run/run_test.go new file mode 100644 index 00000000..fa64bfd2 --- /dev/null +++ b/grpc/run/run_test.go @@ -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) +} diff --git a/grpc/shell/shell.go b/grpc/shell/shell.go new file mode 100644 index 00000000..fe51e0bb --- /dev/null +++ b/grpc/shell/shell.go @@ -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 +} diff --git a/grpc/shell/shell_test.go b/grpc/shell/shell_test.go new file mode 100644 index 00000000..c87c2aff --- /dev/null +++ b/grpc/shell/shell_test.go @@ -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) +} diff --git a/grpc/tests/testutils/testutils.go b/grpc/tests/testutils/testutils.go new file mode 100644 index 00000000..d7abab80 --- /dev/null +++ b/grpc/tests/testutils/testutils.go @@ -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) +} diff --git a/openapi/oauth/authenticate.go b/openapi/oauth/authenticate.go new file mode 100644 index 00000000..288ff2ad --- /dev/null +++ b/openapi/oauth/authenticate.go @@ -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 +} diff --git a/sandbox/DESIGN.md b/sandbox/DESIGN.md index 500a9f45..19fc29fd 100644 --- a/sandbox/DESIGN.md +++ b/sandbox/DESIGN.md @@ -1,1379 +1,305 @@ -# Sandbox Design +# Sandbox Refactoring Design -## 1. Overview +## Background -Sandbox provides **persistent Docker containers** as isolated execution environments for external CLI agents like Claude Code. +The current `sandbox.Manager` was built as a quick prototype for the Claude coding agent. It directly depends on the local Docker client, uses bind mounts for file IO, and Unix sockets for IPC. This limits it to single-node, local-only operation. -### Why Docker? +This document outlines the refactoring plan to make sandbox a production-grade, multi-node capable system built on top of the Tai SDK (`yao/tai`). -- **Persistence**: Claude installs dependencies (npm, pip, apt), which must persist across sessions -- **Cross-platform**: Works on Linux, macOS, and Windows -- **Strong isolation**: Process, filesystem, and network isolation -- **Mature ecosystem**: Well-documented, easy to maintain - -### Architecture +## Architecture ``` -โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” -โ”‚ Yao Server โ”‚ -โ”‚ โ”‚ -โ”‚ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ -โ”‚ โ”‚ Sandbox Manager โ”‚ โ”‚ -โ”‚ โ”‚ โ”‚ โ”‚ -โ”‚ โ”‚ containers: map[containerName]*Container โ”‚ โ”‚ -โ”‚ โ”‚ โ”‚ โ”‚ -โ”‚ โ”‚ - GetOrCreate(userID, chatID) โ†’ get or create container โ”‚ โ”‚ -โ”‚ โ”‚ - Exec(containerName, cmd) โ†’ execute command in container โ”‚ โ”‚ -โ”‚ โ”‚ - Stop(containerName) โ†’ stop container (preserve data) โ”‚ โ”‚ -โ”‚ โ”‚ - Remove(containerName) โ†’ delete container โ”‚ โ”‚ -โ”‚ โ”‚ โ”‚ โ”‚ -โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ -โ”‚ โ”‚ โ”‚ -โ”‚ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ผโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ -โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ -โ”‚ โ–ผ โ–ผ โ–ผ โ”‚ -โ”‚ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ -โ”‚ โ”‚ Container-A โ”‚ โ”‚ Container-B โ”‚ โ”‚ Container-C โ”‚ โ”‚ -โ”‚ โ”‚ (user1-chat1)โ”‚ โ”‚ (user1-chat2)โ”‚ โ”‚ (user2-chat1)โ”‚ โ”‚ -โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ -โ”‚ โ”‚ - Claude CLI โ”‚ โ”‚ - Claude CLI โ”‚ โ”‚ - Claude CLI โ”‚ โ”‚ -โ”‚ โ”‚ - Node.js โ”‚ โ”‚ - Python โ”‚ โ”‚ - Go โ”‚ โ”‚ -โ”‚ โ”‚ - User code โ”‚ โ”‚ - User code โ”‚ โ”‚ - User code โ”‚ โ”‚ -โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ -โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ -โ”‚ โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€ โ”‚ -โ”‚ Unix Socket IPC โ”‚ -โ”‚ (one socket per container) โ”‚ -โ”‚ โ”‚ -โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ +tai.Client Single connection to a Tai endpoint (or local Docker) + โ”‚ sandbox / volume / proxy / vnc low-level APIs + โ”‚ +sandbox.Manager Business layer + Lifecycle, user isolation, TTL, cleanup, IPC ``` ---- +`sandbox.Manager` takes a single `tai.Client` at construction time. Scaling is handled externally by the container runtime โ€” K8s scheduler for pod placement and node scaling, Docker for single-host. The SDK does not manage multiple endpoints or do any scheduling. -## 2. Container Lifecycle +### Why Tai stays in the K8s path +K8s handles pod scheduling and container lifecycle, but it does **not** provide: + +| Capability | K8s native? | What you'd need without Tai | +|-----------|-------------|---------------------------| +| File sync to/from container | No | PVC + init container or sidecar | +| HTTP preview proxy | No | Ingress + Service per sandbox | +| VNC access | No | VNC sidecar + Service + Ingress | +| gRPC IPC relay (container โ†’ Yao) | No | Pod must reach Yao directly (network policy, Service) | + +Tai bundles all four behind a single endpoint. Bypassing Tai to "direct-connect" K8s only covers pod CRUD and exec โ€” you'd still need to solve file IO, preview, VNC, and IPC separately, which means either deploying Tai anyway or assembling equivalent infrastructure from K8s primitives. + +The SDK's `NewK8s()` already supports direct kube-apiserver connection (pass empty `addr`), but this is only useful for bare compute scenarios with no file sync or web preview requirements. + +### sandbox.Manager (yao/sandbox) + +High-level business layer on top of `tai.Client`. Manages container lifecycle, user/session isolation, file operations, and IPC. + +**Responsibilities:** +- Create / get / start / stop / remove sandboxes +- Lifecycle policies: one-shot, session-bound, long-running, persistent +- Per-user and global container limits +- Idle timeout and cleanup +- File operations (via `tai.Client.Volume()` for remote, bind mount for local) +- IPC relay to Yao gRPC server + +### Yao gRPC Server (yao/grpc) + +General-purpose gRPC service exposed by the Yao process. Not limited to sandbox IPC โ€” it exposes Yao's process execution capability to any gRPC client. + +**Clients:** +- Container-internal MCP tools (via Tai Gateway relay) +- `yao run --remote` CLI +- Other Yao instances (future node-to-node) + +**IPC path (replacing Unix socket):** ``` -Create โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ–บ Running โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ–บ Stopped โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ–บ Removed -(docker create) (docker start) (docker stop) (docker rm) - โ”‚ โ”‚ โ”‚ โ”‚ - โ”‚ โ”‚ โ”‚ โ”‚ - โ–ผ โ–ผ โ–ผ โ–ผ -First request Execute tasks Idle timeout Cleanup policy - (persistent) (data preserved) (manual/scheduled) +Container process โ†’ yao-bridge (tai/bridge/) โ†’ Tai Gateway (:9100 gRPC) โ†’ Yao gRPC Server (:9099) + โ”‚ + process.Run(...) ``` -### Container States +Tai does **not** know Yao gRPC address at startup. The upstream is passed per-container via `CreateRequest.GRPCUpstream` โ€” Tai records the mapping and routes relay traffic by source container. This keeps Tai stateless and allows one Tai to serve multiple Yao instances. -| State | Description | -| --------- | ---------------------------------------------- | -| `created` | Container created, not started | -| `running` | Container running, can execute commands | -| `stopped` | Container stopped, data preserved, can restart | -| `removed` | Container deleted | +## Authentication -### Naming Convention +The gRPC server reuses the existing `openapi/oauth` service โ€” no new auth system needed. -``` -yao-sandbox-{userID}-{chatID} +### What already exists -Example: yao-sandbox-u123-c456 -``` +| Capability | Module | Reuse | Needs changes | +|------------|--------|-------|---------------| +| JWT sign (RS256) | `oauth.MakeAccessToken()` | Issue tokens for gRPC clients | None โ€” supports custom scope/subject/extraClaims | +| JWT verify | `oauth.VerifyToken(token string)` | Validate Bearer token in interceptor | None โ€” pure string input, no Gin dependency | +| Signing certs | `oauth.SigningCertificates` | Same keypair for HTTP and gRPC | None | +| Identity | `TokenClaims` (Subject/ClientID/Scope) | gRPC request context | None | +| Scope/ACL | `acl.Scope.Check(*AccessRequest)` | Method-level access control | None โ€” only needs `(Method, Path, Scopes)`, no Gin dependency | +| Scope registration | `acl.Register(...)` | gRPC scopes via same pattern | None โ€” add `grpc:*` scope definitions in `init()` | +| Client auth | `ClientProvider` | `client_credentials` grant for CLI/containers | None | +| Token revocation | `oauth.Revoke(ctx, token, hint)` | Container token cleanup | None | +| Device Flow scaffolding | `types.DeviceAuthorizationResponse`, `GrantTypeDeviceCode`, error codes | CLI `yao login` | Implement `DeviceAuthorization()` (currently stub) | ---- +**Key insight**: `authorized.SetInfo/GetInfo` are Gin-bound, but gRPC does NOT need them. The gRPC interceptor builds `AccessRequest` directly from JWT claims and calls `ScopeManager.Check` โ€” bypasses the full `Enforce` chain (client/team/member), which is HTTP multi-tenant only. -## 3. IPC Communication +**Impact on existing code: zero.** All gRPC auth is purely additive (~80 lines interceptor + scope registration). Device Flow adds ~190 lines new code + ~10 lines to existing `Token()` switch. -### Problem - -Claude CLI runs inside the sandbox but needs to call Yao's MCP Tools (Yao Processes) which run outside. - -### Solution: Unix Socket + MCP JSON-RPC - -``` -โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” -โ”‚ Docker Container โ”‚ -โ”‚ โ”‚ -โ”‚ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ -โ”‚ โ”‚ Claude CLI โ”‚ โ”‚ -โ”‚ โ”‚ โ”‚ โ”‚ โ”‚ -โ”‚ โ”‚ โ”‚ .mcp.json: "yao" โ†’ stdio โ”‚ โ”‚ -โ”‚ โ”‚ โ–ผ โ”‚ โ”‚ -โ”‚ โ”‚ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ โ”‚ -โ”‚ โ”‚ โ”‚ yao-bridge (lightweight binary) โ”‚ โ”‚ โ”‚ -โ”‚ โ”‚ โ”‚ stdin/stdout โ†” /tmp/yao.sock โ”‚ โ”‚ โ”‚ -โ”‚ โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ โ”‚ -โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ -โ”‚ โ”‚ โ”‚ -โ”‚ /tmp/yao.sock โ”‚ -โ”‚ โ”‚ โ”‚ -โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ผโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ - โ”‚ - Unix Socket (bind mount) - โ”‚ -โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ผโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” -โ”‚ โ”‚ โ”‚ -โ”‚ {YAO_DATA_ROOT}/sandbox/ipc/{sessionID}.sock โ”‚ -โ”‚ โ”‚ โ”‚ -โ”‚ โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ–ผโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ” โ”‚ -โ”‚ โ”‚ IPC Server (goroutine) โ”‚ โ”‚ -โ”‚ โ”‚ โ”‚ โ”‚ -โ”‚ โ”‚ MCP JSON-RPC Methods: โ”‚ โ”‚ -โ”‚ โ”‚ - initialize โ†’ handshake โ”‚ โ”‚ -โ”‚ โ”‚ - tools/list โ†’ return authorized Yao MCP tools โ”‚ โ”‚ -โ”‚ โ”‚ - tools/call โ†’ execute process.New(name, args...) โ”‚ โ”‚ -โ”‚ โ”‚ - resources/list โ†’ list Yao resources โ”‚ โ”‚ -โ”‚ โ”‚ - resources/read โ†’ read Yao resource โ”‚ โ”‚ -โ”‚ โ”‚ โ”‚ โ”‚ -โ”‚ โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ โ”‚ -โ”‚ โ”‚ -โ”‚ Yao Server โ”‚ -โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜ -``` - -### Protocol - -- **Format**: MCP standard JSON-RPC 2.0 over NDJSON (newline-delimited JSON) -- **Transport**: Unix Socket -- **Bridge**: `yao-bridge` binary converts stdio โ†” socket - ---- - -## 4. Core Interfaces - -### 4.1 Sandbox Manager +### gRPC interceptor ```go -// sandbox/manager.go -package sandbox - -type Manager struct { - mu sync.Mutex // Protects creation - containers sync.Map // containerName โ†’ *Container - running int32 // Running container count - ipcManager *ipc.Manager - dockerClient *docker.Client - config *Config -} - -var ErrTooManyContainers = errors.New("sandbox: too many running containers, please try again later") - -type Config struct { - Image string // Docker image, default: yao/sandbox:latest - WorkspaceRoot string // Host workspace root directory - IPCDir string // IPC socket directory - MaxContainers int // Maximum concurrent containers - IdleTimeout time.Duration // Idle timeout before stopping container - MaxMemory string // Memory limit, e.g., "2g" - MaxCPU float64 // CPU limit, e.g., 1.0 -} - -type Container struct { - ID string - Name string // yao-sandbox-{userID}-{chatID} - UserID string - ChatID string - Status string // created, running, stopped - CreatedAt time.Time - LastUsedAt time.Time - IPCSession *ipc.Session +func authInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { + md, _ := metadata.FromIncomingContext(ctx) + token := extractBearer(md) + claims, err := oauth.OAuth.VerifyToken(token) + if err != nil { + return nil, status.Errorf(codes.Unauthenticated, "invalid token") + } + ctx = withClaims(ctx, claims) + return handler(ctx, req) } ``` -### 4.2 Manager Methods +`oauth.OAuth` is a global singleton initialized at Yao startup. The gRPC server simply references it โ€” same signing keys, same token format, same user/client model. -```go -// GetOrCreate returns existing container or creates new one -// Returns ErrTooManyContainers if limit exceeded -func (m *Manager) GetOrCreate(ctx context.Context, userID, chatID string) (*Container, error) +### Token flow by client type -// Stream executes command and returns stdout reader -func (m *Manager) Stream(ctx context.Context, containerName string, cmd []string, opts *ExecOptions) (io.ReadCloser, error) +| Client | How it gets a token | +|--------|-------------------| +| Container MCP tool | `oauth.MakeAccessToken(clientID, "grpc:mcp grpc:run", userID, 900)` injected as `YAO_TOKEN` env var. `yao-bridge` auto-refreshes via `YAO_REFRESH_TOKEN`. Manager revokes refresh token on container Remove. | +| `yao run` CLI | `yao login` โ†’ OAuth Device Authorization Grant โ†’ token saved to `~/.yao/credentials`. Logged in = gRPC, not logged in = local. | +| Yao-to-Yao | Pre-shared service token or `client_credentials` | -// Exec executes command and waits for completion -func (m *Manager) Exec(ctx context.Context, containerName string, cmd []string, opts *ExecOptions) (*ExecResult, error) +## Network Security -// Stop stops container but preserves data -func (m *Manager) Stop(ctx context.Context, containerName string) error +### Yao โ†” Tai communication -// Start starts a stopped container -func (m *Manager) Start(ctx context.Context, containerName string) error +Tai exposes Docker Engine API (:2375), K8s API (:6443), gRPC Volume (:9100), HTTP proxy (:8080), and VNC (:6080). These are raw protocol proxies with **no built-in auth** โ€” security is handled at the network layer. -// Remove deletes container and its data -func (m *Manager) Remove(ctx context.Context, containerName string) error +| Deployment | Strategy | +|-----------|----------| +| Same host (local) | Bind to `127.0.0.1` or Unix socket, no exposure | +| Same VPC / LAN | Firewall rules / security groups, private subnet only | +| Cross-network | VPN / WireGuard tunnel, or mTLS termination at Tai | -// List returns all containers for a user -func (m *Manager) List(ctx context.Context, userID string) ([]*Container, error) +### gRPC Server listen policy -// Cleanup stops idle containers -func (m *Manager) Cleanup(ctx context.Context) error +The Yao gRPC server (:9099) supports configurable listen address: + +| Scenario | Listen | Why | +|----------|--------|-----| +| Local dev | `127.0.0.1:9099` | Only local containers reach it | +| Production (same host) | `127.0.0.1:9099` | Tai on same machine forwards via loopback | +| Production (multi-node) | `0.0.0.0:9099` + IP allowlist | Remote Tai nodes need access | + +### IP allowlist (gRPC server) + +For multi-node deployment where gRPC must listen on `0.0.0.0`, the server should support an IP/CIDR allowlist: + +``` +YAO_GRPC_LISTEN=0.0.0.0:9099 +YAO_GRPC_ALLOW=10.0.0.0/8,172.16.0.0/12,192.168.0.0/16 ``` -### 4.3 Filesystem Methods +Enforcement is a simple gRPC interceptor that runs **before** the auth interceptor: ```go -// WriteFile writes content to a file in container -func (m *Manager) WriteFile(ctx context.Context, containerName, path string, content []byte) error - -// ReadFile reads content from a file in container -func (m *Manager) ReadFile(ctx context.Context, containerName, path string) ([]byte, error) - -// ListDir lists directory contents in container -func (m *Manager) ListDir(ctx context.Context, containerName, path string) ([]FileInfo, error) - -// Stat returns file info -func (m *Manager) Stat(ctx context.Context, containerName, path string) (*FileInfo, error) - -// MkDir creates directory in container -func (m *Manager) MkDir(ctx context.Context, containerName, path string) error - -// Remove removes file or directory in container -func (m *Manager) RemoveFile(ctx context.Context, containerName, path string) error - -// CopyToContainer copies file/directory from host to container -func (m *Manager) CopyToContainer(ctx context.Context, containerName, hostPath, containerPath string) error - -// CopyFromContainer copies file/directory from container to host -func (m *Manager) CopyFromContainer(ctx context.Context, containerName, containerPath, hostPath string) error - -// FileInfo represents file metadata -type FileInfo struct { - Name string - Path string - Size int64 - Mode os.FileMode - ModTime time.Time - IsDir bool -} -``` - -### 4.4 ExecOptions - -```go -type ExecOptions struct { - WorkDir string // Working directory inside container - Env map[string]string // Environment variables - Stdin io.Reader // Standard input - Timeout time.Duration // Execution timeout -} - -type ExecResult struct { - ExitCode int - Stdout string - Stderr string -} -``` - ---- - -## 5. IPC System - -### 5.1 IPC Session - -```go -// ipc/session.go -package ipc - -type Session struct { - ID string // Usually equals chatID - SocketPath string // {IPCDir}/{id}.sock - Listener net.Listener - Conn net.Conn - Context *AgentContext - MCPTools map[string]*MCPTool - cancel context.CancelFunc -} - -type AgentContext struct { - UserID string - ChatID string - Locale string -} - -type MCPTool struct { - Name string - Description string - Process string // Yao process name - InputSchema json.RawMessage // JSON Schema -} -``` - -### 5.2 IPC Manager - -```go -// ipc/manager.go -type Manager struct { - sessions sync.Map // sessionID โ†’ *Session - sockDir string // {YAO_DATA_ROOT}/sandbox/ipc/ -} - -// Create creates new IPC session -func (m *Manager) Create(ctx context.Context, sessionID string, agentCtx *AgentContext, mcpTools map[string]*MCPTool) (*Session, error) - -// Close closes IPC session and cleans up -func (m *Manager) Close(sessionID string) error - -// Get returns existing session -func (m *Manager) Get(sessionID string) (*Session, bool) -``` - -### 5.3 JSON-RPC Message Handling - -```go -// JSON-RPC request structure -type JSONRPCRequest struct { - JSONRPC string `json:"jsonrpc"` - ID interface{} `json:"id,omitempty"` - Method string `json:"method"` - Params json.RawMessage `json:"params,omitempty"` -} - -// JSON-RPC response structure -type JSONRPCResponse struct { - JSONRPC string `json:"jsonrpc"` - ID interface{} `json:"id,omitempty"` - Result interface{} `json:"result,omitempty"` - Error *JSONRPCError `json:"error,omitempty"` -} - -type JSONRPCError struct { - Code int `json:"code"` - Message string `json:"message"` - Data interface{} `json:"data,omitempty"` -} -``` - -### 5.4 Session Message Loop - -```go -func (s *Session) serve(ctx context.Context) { - defer s.cleanup() - - for { - select { - case <-ctx.Done(): - return - default: +func ipAllowInterceptor(allowedCIDRs []*net.IPNet) grpc.UnaryServerInterceptor { + return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { + peer, _ := peer.FromContext(ctx) + if !isAllowed(peer.Addr, allowedCIDRs) { + return nil, status.Errorf(codes.PermissionDenied, "ip not allowed") } - - conn, err := s.Listener.Accept() - if err != nil { - continue - } - s.Conn = conn - s.handleConnection(ctx, conn) - } -} - -func (s *Session) handleConnection(ctx context.Context, conn net.Conn) { - defer conn.Close() - - scanner := bufio.NewScanner(conn) - for scanner.Scan() { - select { - case <-ctx.Done(): - return - default: - } - - line := scanner.Text() - response := s.handleMessage(line) - if response != "" { - conn.Write([]byte(response + "\n")) - } - } -} - -func (s *Session) handleMessage(line string) string { - var req JSONRPCRequest - if err := json.Unmarshal([]byte(line), &req); err != nil { - return s.errorResponse(nil, -32700, "Parse error") - } - - switch req.Method { - case "initialize": - return s.handleInitialize(req) - case "initialized": - return "" // notification, no response - case "tools/list": - return s.handleListTools(req) - case "tools/call": - return s.handleCallTool(req) - case "resources/list": - return s.handleListResources(req) - case "resources/read": - return s.handleReadResource(req) - default: - return s.errorResponse(req.ID, -32601, "Method not found") + return handler(ctx, req) } } ``` -### 5.5 Tool Call Handler +Defense in depth: IP allowlist is the first gate, OAuth token is the second. Both must pass. + +### Tai side security + +Tai itself does not need auth โ€” it trusts its network boundary. Recommended: +- Docker: Tai container runs on a private network, ports not exposed to public +- K8s: Tai runs as a DaemonSet or Deployment, service only accessible within cluster +- If Tai must be exposed, put it behind a reverse proxy (nginx/envoy) with mTLS or VPN + +### Tai high availability + +Tai is a single endpoint, but all its services except VNC WebSocket are stateless. Avoiding single-point-of-failure is a deployment concern, not an SDK concern. + +| Deployment | HA strategy | +|-----------|-------------| +| Docker single-host | Docker restart policy (`--restart=always`), Tai failure = transient | +| K8s Deployment | `replicas: N` + K8s Service load balancing, liveness probe on `/healthz` | +| K8s DaemonSet | One Tai per node, pod talks to local Tai via node-local Service | + +VNC uses WebSocket long connections โ€” if Tai restarts, active VNC sessions drop and the client reconnects. Stateless services (K8s proxy, Docker proxy, Volume gRPC, HTTP proxy) recover transparently behind a Service. + +The SDK `tai.Client` connects to a single address. In K8s this address is a Service VIP โ€” Tai replicas behind it are invisible to the SDK. + +## Container Lifecycle + +Lifecycle is managed by `sandbox.Manager`, not by tai.Client. + +| Policy | TTL | Behavior | +|--------|-----|----------| +| One-shot | 0 | Destroyed immediately after execution | +| Session | Minutes | Alive while user is active, cleaned up on idle timeout | +| Long-running | Hours/Days | User workspace, recoverable, cleaned up after extended idle | +| Persistent | None | User-managed, never auto-cleaned | + +## File Operations + +| Mode | tai.Client | File IO | +|------|-----------|---------| +| Local | `tai.New("")` | Bind mount, direct host filesystem | +| Remote | `tai.New("tai://host")` | `tai.Client.Volume()` via gRPC | + +Local mode preserves bind mount for performance. Remote mode uses `tai/volume` (gRPC + lz4 compression). `sandbox.Manager` routes based on `client.IsLocal()`. + +## Agent Layer Adaptation + +The agent layer (`agent/assistant`, `agent/sandbox`, `agent/context`) currently hardcodes local-only assumptions. It needs to be adapted to work with the new `sandbox.Manager` backed by `tai.Client`. + +### Current coupling + +``` +agent/assistant/sandbox.go + โ”‚ + โ”œโ”€ GetSandboxManager() Global singleton, local Docker only + โ”œโ”€ initSandbox() Creates executor, calls manager.GetOrCreate() + โ”œโ”€ BuildMCPConfigForSandbox() Hardcodes /tmp/yao.sock for yao-bridge + โ””โ”€ loadMCPToolsForIPC() Loads MCP tools, injects into IPC session + +agent/sandbox/claude/executor.go + โ”‚ + โ”œโ”€ manager.GetOrCreate() Direct Docker container creation + โ”œโ”€ manager.Stream() Docker exec + attach + โ””โ”€ manager.Remove() Docker container removal + +agent/context/jsapi_sandbox.go + โ”‚ + โ”œโ”€ ReadFile() Host filesystem via bind mount path translation + โ”œโ”€ WriteFile() Docker CopyToContainer + โ””โ”€ Exec() Docker exec +``` + +### What changes + +| Component | Before | After | +|-----------|--------|-------| +| `GetSandboxManager()` | Global singleton, `docker.NewClientWithOpts(FromEnv)` | Initialized with a `tai.Client` from Yao config | +| Container creation | `dockerClient.ContainerCreate()` | `tai.Client.Sandbox().Create()` | +| Container exec | `dockerClient.ContainerExecCreate/Start/Attach` | `tai.Client.Sandbox().Exec()` | +| File read | Host path via bind mount (`containerPathToHost`) | Local: bind mount (same). Remote: `tai.Client.Volume().Read()` | +| File write | `dockerClient.CopyToContainer` | Local: bind mount. Remote: `tai.Client.Volume().Write()` | +| IPC | Unix socket bind mount + yao-bridge | Local: Unix socket (same). Remote: Tai gRPC relay โ†’ Yao gRPC server | +| MCP config | `{args: ["/tmp/yao.sock"]}` hardcoded | Local: socket path from config. Remote: gRPC endpoint injected as env var | +| VNC | `vncproxy.NewProxy(nil)` local assumption | `tai.Client.VNC().URL()` | +| Cleanup | `dockerClient.ContainerRemove` | `tai.Client.Sandbox().Remove()` | + +### IPC migration detail + +**Local mode** (same host): Unix socket preserved โ€” zero overhead, no change needed. + +**Remote mode** (via Tai): +``` +Container process โ†’ yao-bridge (tai/bridge/) โ†’ Tai relay (:9100 gRPC) โ†’ Yao gRPC Server +``` + +`yao-bridge` source lives in `yao/tai/bridge/` โ€” it's a Tai SDK client (consumes Tai relay), shares gRPC deps with `tai/`, and is version-locked with the Tai protocol. Built via `go build ./tai/bridge/cmd/yao-bridge`. + +Bridge mode determined by env var: + +``` +YAO_IPC_MODE=socket YAO_IPC_ADDR=/tmp/yao.sock # local +YAO_IPC_MODE=grpc YAO_IPC_ADDR=tai-host:9100 # remote +``` + +In gRPC mode, bridge also reads `YAO_TOKEN` / `YAO_REFRESH_TOKEN` and handles automatic token refresh (see grpc/DESIGN.md Container token section). + +Tai relay upstream is NOT configured at Tai startup. Manager passes `GRPCUpstream` per-container in `CreateRequest` โ€” Tai records the mapping and routes by source container. One Tai can serve containers from different Yao instances. + +`BuildMCPConfigForSandbox()` sets the env vars based on `client.IsLocal()`. + +### SandboxExecutor interface + +The `agent/context/jsapi_sandbox.go` `SandboxExecutor` interface stays the same โ€” it's already abstract. Implementation behind it changes: ```go -func (s *Session) handleCallTool(req JSONRPCRequest) string { - var params struct { - Name string `json:"name"` - Arguments map[string]interface{} `json:"arguments"` - } - json.Unmarshal(req.Params, ¶ms) - - // Check authorization - tool, ok := s.MCPTools[params.Name] - if !ok { - return s.errorResponse(req.ID, -32602, "Tool not found or not authorized") - } - - // Execute Yao Process - proc := process.New(tool.Process, params.Arguments) - proc.WithContext(s.Context) - - if err := proc.Execute(); err != nil { - return s.toolErrorResponse(req.ID, params.Name, err) - } - defer proc.Release() - - result := proc.Value() - return s.toolSuccessResponse(req.ID, result) +type SandboxExecutor interface { + ReadFile(path string) ([]byte, error) + WriteFile(path string, data []byte) error + ListDir(path string) ([]FileInfo, error) + Exec(cmd string, args ...string) (string, error) + GetWorkDir() string + GetSandboxID() string + GetVNCUrl() (string, error) } ``` ---- +Hooks (`ctx.sandbox.ReadFile()`, etc.) work unchanged. The executor routes to bind mount or `tai.Client.Volume()` internally. -## 6. Docker Container Management +### Agent lifecycle policy -### 6.1 NewManager Constructor +Currently: sandbox created on chat start, removed on chat end (`defer sandboxCleanup`). -```go -func NewManager(config *Config) (*Manager, error) { - // Initialize Docker client - cli, err := docker.NewClientWithOpts(docker.FromEnv, docker.WithAPIVersionNegotiation()) - if err != nil { - return nil, fmt.Errorf("failed to create Docker client: %w", err) - } - - // Ping Docker to verify connection - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - if _, err := cli.Ping(ctx); err != nil { - return nil, fmt.Errorf("Docker not available: %w", err) - } - - // Ensure directories exist - os.MkdirAll(config.WorkspaceRoot, 0755) - os.MkdirAll(config.IPCDir, 0755) - - m := &Manager{ - dockerClient: cli, - config: config, - ipcManager: ipc.NewManager(config.IPCDir), - } - - // Start cleanup loop - go m.startCleanupLoop(context.Background()) - - return m, nil -} -``` - -### 6.2 GetOrCreate with Limit Check - -```go -func (m *Manager) GetOrCreate(ctx context.Context, userID, chatID string) (*Container, error) { - containerName := fmt.Sprintf("yao-sandbox-%s-%s", userID, chatID) - - // Check if container already exists (fast path) - if c, ok := m.containers.Load(containerName); ok { - container := c.(*Container) - container.LastUsedAt = time.Now() - return container, nil - } - - // Use mutex for creation to avoid race condition - m.mu.Lock() - defer m.mu.Unlock() - - // Double-check after acquiring lock - if c, ok := m.containers.Load(containerName); ok { - container := c.(*Container) - container.LastUsedAt = time.Now() - return container, nil - } - - // Check running container limit - if m.running >= int32(m.config.MaxContainers) { - return nil, ErrTooManyContainers - } - - // Create new container - container, err := m.createContainer(ctx, userID, chatID) - if err != nil { - return nil, err - } - - // Store and increment counter - m.containers.Store(containerName, container) - m.running++ - - return container, nil -} -``` - -### 6.3 Create Container (internal) - -```go -func (m *Manager) createContainer(ctx context.Context, userID, chatID string) (*Container, error) { - containerName := fmt.Sprintf("yao-sandbox-%s-%s", userID, chatID) - - // Ensure image exists, auto-pull if not - if err := m.ensureImage(ctx, m.config.Image); err != nil { - return nil, err - } - - // Workspace directory - workspaceHost := filepath.Join(m.config.WorkspaceRoot, userID, chatID) - os.MkdirAll(workspaceHost, 0755) - - // IPC socket path - sessionID := chatID - ipcSocketHost := filepath.Join(m.config.IPCDir, sessionID+".sock") - - // Create container - resp, err := m.dockerClient.ContainerCreate(ctx, - &container.Config{ - Image: m.config.Image, - Cmd: []string{"sleep", "infinity"}, // Keep running - WorkingDir: "/workspace", - Env: []string{ - "YAO_IPC_SOCKET=/tmp/yao.sock", - }, - }, - &container.HostConfig{ - Binds: []string{ - workspaceHost + ":/workspace", - ipcSocketHost + ":/tmp/yao.sock", - }, - Resources: container.Resources{ - Memory: parseMemory(m.config.MaxMemory), - NanoCPUs: int64(m.config.MaxCPU * 1e9), - }, - SecurityOpt: []string{"no-new-privileges"}, - CapDrop: []string{"ALL"}, - }, - nil, nil, containerName, - ) - - if err != nil { - return nil, err - } - - return &Container{ - ID: resp.ID, - Name: containerName, - UserID: userID, - ChatID: chatID, - Status: "created", - CreatedAt: time.Now(), - }, nil -} - -// ensureImage ensures the image exists locally, pulls if not -func (m *Manager) ensureImage(ctx context.Context, imageName string) error { - // Check if image exists locally - _, _, err := m.dockerClient.ImageInspectWithRaw(ctx, imageName) - if err == nil { - return nil // Image exists - } - - // Image not found, pull it - reader, err := m.dockerClient.ImagePull(ctx, imageName, image.PullOptions{}) - if err != nil { - return fmt.Errorf("failed to pull image %s: %w", imageName, err) - } - defer reader.Close() - - // Wait for pull to complete - io.Copy(io.Discard, reader) - return nil -} -``` - -### 6.4 Ensure Running - -```go -func (m *Manager) ensureRunning(ctx context.Context, containerName string) error { - c, ok := m.containers.Load(containerName) - if !ok { - return fmt.Errorf("container not found: %s", containerName) - } - cont := c.(*Container) - - if cont.Status == "running" { - return nil - } - - // Start the container - if err := m.dockerClient.ContainerStart(ctx, cont.ID, container.StartOptions{}); err != nil { - return err - } - - m.mu.Lock() - cont.Status = "running" - cont.LastUsedAt = time.Now() - m.mu.Unlock() - - return nil -} -``` - -### 6.5 Execute Command (Streaming) - -```go -func (m *Manager) Stream(ctx context.Context, containerName string, cmd []string, opts *ExecOptions) (io.ReadCloser, error) { - // Ensure container is running - if err := m.ensureRunning(ctx, containerName); err != nil { - return nil, err - } - - // Get container - c, _ := m.containers.Load(containerName) - cont := c.(*Container) - - // Create exec instance - execConfig := container.ExecOptions{ - Cmd: cmd, - WorkingDir: opts.WorkDir, - Env: mapToSlice(opts.Env), - AttachStdout: true, - AttachStderr: true, - } - - execResp, err := m.dockerClient.ContainerExecCreate(ctx, cont.ID, execConfig) - if err != nil { - return nil, err - } - - // Attach to exec - attachResp, err := m.dockerClient.ContainerExecAttach(ctx, execResp.ID, container.ExecAttachOptions{}) - if err != nil { - return nil, err - } - - return attachResp.Reader, nil -} - -// Exec executes command and waits for completion -func (m *Manager) Exec(ctx context.Context, containerName string, cmd []string, opts *ExecOptions) (*ExecResult, error) { - if opts == nil { - opts = &ExecOptions{} - } - - reader, err := m.Stream(ctx, containerName, cmd, opts) - if err != nil { - return nil, err - } - defer reader.Close() - - // Read all output - output, err := io.ReadAll(reader) - if err != nil { - return nil, err - } - - // TODO: Parse stdout/stderr from Docker multiplexed stream - // TODO: Get exit code from ContainerExecInspect - - return &ExecResult{ - ExitCode: 0, - Stdout: string(output), - Stderr: "", - }, nil -} -``` - -### 6.6 Filesystem Operations - -```go -// WriteFile writes content to a file in container using docker cp -func (m *Manager) WriteFile(ctx context.Context, containerName, path string, content []byte) error { - c, ok := m.containers.Load(containerName) - if !ok { - return fmt.Errorf("container not found: %s", containerName) - } - cont := c.(*Container) - // Create a tar archive with the file - var buf bytes.Buffer - tw := tar.NewWriter(&buf) - - hdr := &tar.Header{ - Name: filepath.Base(path), - Mode: 0644, - Size: int64(len(content)), - } - tw.WriteHeader(hdr) - tw.Write(content) - tw.Close() - - // Copy to container - return m.dockerClient.CopyToContainer(ctx, cont.ID, filepath.Dir(path), &buf, container.CopyToContainerOptions{}) -} - -// ReadFile reads content from a file in container -func (m *Manager) ReadFile(ctx context.Context, containerName, path string) ([]byte, error) { - c, ok := m.containers.Load(containerName) - if !ok { - return nil, fmt.Errorf("container not found: %s", containerName) - } - cont := c.(*Container) - - reader, _, err := m.dockerClient.CopyFromContainer(ctx, cont.ID, path) - if err != nil { - return nil, err - } - defer reader.Close() - - // Extract from tar - tr := tar.NewReader(reader) - _, err = tr.Next() - if err != nil { - return nil, err - } - - return io.ReadAll(tr) -} - -// ListDir lists directory contents -func (m *Manager) ListDir(ctx context.Context, containerName, path string) ([]FileInfo, error) { - result, err := m.Exec(ctx, containerName, []string{"ls", "-la", "--time-style=+%s", path}, nil) - if err != nil { - return nil, err - } - - return parseLS(result.Stdout), nil -} - -// Stat returns file info -func (m *Manager) Stat(ctx context.Context, containerName, path string) (*FileInfo, error) { - result, err := m.Exec(ctx, containerName, []string{"stat", "--format=%n|%s|%f|%Y|%F", path}, nil) - if err != nil { - return nil, err - } - return parseStat(result.Stdout), nil -} - -// MkDir creates directory in container -func (m *Manager) MkDir(ctx context.Context, containerName, path string) error { - _, err := m.Exec(ctx, containerName, []string{"mkdir", "-p", path}, nil) - return err -} - -// RemoveFile removes file or directory in container -func (m *Manager) RemoveFile(ctx context.Context, containerName, path string) error { - _, err := m.Exec(ctx, containerName, []string{"rm", "-rf", path}, nil) - return err -} - -// CopyToContainer copies from host to container -func (m *Manager) CopyToContainer(ctx context.Context, containerName, hostPath, containerPath string) error { - c, ok := m.containers.Load(containerName) - if !ok { - return fmt.Errorf("container not found: %s", containerName) - } - cont := c.(*Container) - - // Create tar archive from host path - archive, err := createTarFromPath(hostPath) - if err != nil { - return err - } - defer archive.Close() - - return m.dockerClient.CopyToContainer(ctx, cont.ID, containerPath, archive, container.CopyToContainerOptions{}) -} - -// CopyFromContainer copies from container to host -func (m *Manager) CopyFromContainer(ctx context.Context, containerName, containerPath, hostPath string) error { - c, ok := m.containers.Load(containerName) - if !ok { - return fmt.Errorf("container not found: %s", containerName) - } - cont := c.(*Container) - - reader, _, err := m.dockerClient.CopyFromContainer(ctx, cont.ID, containerPath) - if err != nil { - return err - } - defer reader.Close() - - return extractTarToPath(reader, hostPath) -} -``` - -### 6.7 Cleanup Strategy - -```go -func (m *Manager) startCleanupLoop(ctx context.Context) { - ticker := time.NewTicker(5 * time.Minute) - defer ticker.Stop() - - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - m.Cleanup(ctx) - } - } -} - -func (m *Manager) Cleanup(ctx context.Context) error { - now := time.Now() - - m.containers.Range(func(key, value interface{}) bool { - containerName := key.(string) - c := value.(*Container) - - // Stop idle containers - if c.Status == "running" && now.Sub(c.LastUsedAt) > m.config.IdleTimeout { - m.Stop(ctx, containerName) - } - - return true - }) - - return nil -} - -func (m *Manager) Start(ctx context.Context, containerName string) error { - return m.ensureRunning(ctx, containerName) -} - -func (m *Manager) Stop(ctx context.Context, containerName string) error { - c, ok := m.containers.Load(containerName) - if !ok { - return nil - } - cont := c.(*Container) - - if err := m.dockerClient.ContainerStop(ctx, cont.ID, container.StopOptions{}); err != nil { - return err - } - - // Update status, decrement running count - m.mu.Lock() - if cont.Status == "running" { - cont.Status = "stopped" - m.running-- - } - m.mu.Unlock() - - return nil -} - -func (m *Manager) Remove(ctx context.Context, containerName string) error { - // Stop first if running - m.Stop(ctx, containerName) - - c, ok := m.containers.Load(containerName) - if !ok { - return nil - } - cont := c.(*Container) - - if err := m.dockerClient.ContainerRemove(ctx, cont.ID, container.RemoveOptions{}); err != nil { - return err - } - - // Remove from map - m.containers.Delete(containerName) - - return nil -} - -func (m *Manager) List(ctx context.Context, userID string) ([]*Container, error) { - var result []*Container - prefix := fmt.Sprintf("yao-sandbox-%s-", userID) - - m.containers.Range(func(key, value interface{}) bool { - containerName := key.(string) - if strings.HasPrefix(containerName, prefix) { - result = append(result, value.(*Container)) - } - return true - }) - - return result, nil -} -``` - -### 6.8 Helper Functions - -```go -// mapToSlice converts map to []string for env vars -func mapToSlice(m map[string]string) []string { - if m == nil { - return nil - } - result := make([]string, 0, len(m)) - for k, v := range m { - result = append(result, k+"="+v) - } - return result -} - -// parseMemory converts string like "2g" to bytes -func parseMemory(s string) int64 { - // Implementation: parse "2g" โ†’ 2*1024*1024*1024 - // Use Docker's units package or manual parsing - return 0 // placeholder -} - -// parseLS parses ls -la output to []FileInfo -func parseLS(output string) []FileInfo { - // Implementation: parse ls output lines - return nil // placeholder -} - -// parseStat parses stat output to *FileInfo -func parseStat(output string) *FileInfo { - // Implementation: parse stat --format output - return nil // placeholder -} - -// createTarFromPath creates a tar archive from a host path -func createTarFromPath(hostPath string) (io.ReadCloser, error) { - // Implementation: walk directory, create tar entries - return nil, nil // placeholder -} - -// extractTarToPath extracts a tar archive to a host path -func extractTarToPath(reader io.Reader, hostPath string) error { - // Implementation: read tar entries, write to disk - return nil // placeholder -} -``` - ---- - -## 7. yao-bridge - -Lightweight binary inside container that bridges stdio to Unix socket. - -```go -// cmd/yao-bridge/main.go -package main - -import ( - "io" - "net" - "os" -) - -func main() { - if len(os.Args) < 2 { - os.Exit(1) - } - - sockPath := os.Args[1] - - // Connect to Unix socket - conn, err := net.Dial("unix", sockPath) - if err != nil { - os.Exit(1) - } - defer conn.Close() - - // stdin โ†’ socket - go func() { - io.Copy(conn, os.Stdin) - conn.(*net.UnixConn).CloseWrite() - }() - - // socket โ†’ stdout - io.Copy(os.Stdout, conn) -} -``` - -Build as static binary and include in Docker image. - ---- - -## 8. Docker Image - -### 8.1 Image Naming Convention - -``` -yao/sandbox-{tool}:{variant} - -Examples: - yao/sandbox-claude:latest # Claude CLI + Node.js + Python (default) - yao/sandbox-claude:full # + Go - yao/sandbox-cursor:latest # Cursor CLI + Node.js + Python (future) -``` - -### 8.2 Source Directory Structure - -``` -sandbox/ -โ”œโ”€โ”€ docker/ -โ”‚ โ”œโ”€โ”€ base/ -โ”‚ โ”‚ โ””โ”€โ”€ Dockerfile.base # Common base image -โ”‚ โ”œโ”€โ”€ claude/ -โ”‚ โ”‚ โ”œโ”€โ”€ Dockerfile # Default: Claude + Node + Python -โ”‚ โ”‚ โ””โ”€โ”€ Dockerfile.full # + Go -โ”‚ โ”œโ”€โ”€ cursor/ # Future -โ”‚ โ”‚ โ””โ”€โ”€ Dockerfile -โ”‚ โ”œโ”€โ”€ build.sh -โ”‚ โ””โ”€โ”€ scripts/ -โ”‚ โ””โ”€โ”€ entrypoint.sh -โ”œโ”€โ”€ bridge/ -โ”‚ โ””โ”€โ”€ main.go # yao-bridge source -โ”œโ”€โ”€ ipc/ -โ”‚ โ”œโ”€โ”€ manager.go -โ”‚ โ””โ”€โ”€ session.go -โ”œโ”€โ”€ manager.go -โ”œโ”€โ”€ config.go -โ””โ”€โ”€ types.go -``` - -### 8.3 Base Image - -```dockerfile -# sandbox/docker/base/Dockerfile.base -FROM ubuntu:22.04 - -# Base tools -RUN apt-get update && apt-get install -y \ - curl \ - git \ - ca-certificates \ - && rm -rf /var/lib/apt/lists/* - -# yao-bridge (common to all tools) -COPY yao-bridge /usr/local/bin/yao-bridge -RUN chmod +x /usr/local/bin/yao-bridge - -# Working directory -WORKDIR /workspace - -# Non-root user -RUN useradd -m -s /bin/bash sandbox -USER sandbox - -CMD ["sleep", "infinity"] -``` - -### 8.4 Claude Tool Images - -```dockerfile -# sandbox/docker/claude/Dockerfile -# Default image: Claude CLI + Node.js + Python -FROM yao/sandbox-base:latest - -USER root - -# Claude CLI -RUN curl -fsSL https://claude.ai/install.sh | sh - -# Node.js 20 -RUN curl -fsSL https://deb.nodesource.com/setup_20.x | bash - \ - && apt-get install -y nodejs \ - && rm -rf /var/lib/apt/lists/* - -# Python 3.11 -RUN apt-get update && apt-get install -y \ - python3.11 \ - python3-pip \ - && rm -rf /var/lib/apt/lists/* - -USER sandbox -``` - -```dockerfile -# sandbox/docker/claude/Dockerfile.full -# Full image: + Go -FROM yao/sandbox-claude:latest - -USER root - -# Go 1.23 -RUN curl -fsSL https://go.dev/dl/go1.23.linux-amd64.tar.gz | tar -C /usr/local -xzf - \ - && ln -s /usr/local/go/bin/go /usr/local/bin/go - -USER sandbox -``` - -### 8.5 Build Script - -```bash -#!/bin/bash -# sandbox/docker/build.sh - -set -e - -TOOL=${1:-claude} - -# Build yao-bridge -cd ../bridge -CGO_ENABLED=0 go build -o ../docker/yao-bridge . -cd ../docker - -# Build base image -docker build -t yao/sandbox-base:latest -f base/Dockerfile.base . - -# Build tool-specific images -case $TOOL in - claude) - docker build -t yao/sandbox-claude:latest -f claude/Dockerfile . - docker build -t yao/sandbox-claude:full -f claude/Dockerfile.full . - ;; - cursor) - docker build -t yao/sandbox-cursor:latest -f cursor/Dockerfile . - ;; - all) - $0 claude - $0 cursor - ;; -esac - -echo "Images built for tool: $TOOL" -``` - -### 8.6 Image Variants - -| Image | Tool | Size | Pre-installed | -| --------------------------- | ------ | ------ | ----------------------------------- | -| `yao/sandbox-base:latest` | - | ~200MB | git, curl, yao-bridge | -| `yao/sandbox-claude:latest` | Claude | ~700MB | Claude CLI, Node.js 20, Python 3.11 | -| `yao/sandbox-claude:full` | Claude | ~1.3GB | + Go 1.23 | -| `yao/sandbox-cursor:latest` | Cursor | ~700MB | Cursor CLI, Node.js 20, Python 3.11 | - -Default: `yao/sandbox-claude:latest` (includes Node + Python) - ---- - -## 9. Configuration - -### 9.1 Environment Variables - -| Env Variable | Default | Description | -| -------------------------- | ----------------------------------- | ------------------------- | -| `YAO_SANDBOX_IMAGE` | `yao/sandbox-claude:latest` | Default Docker image | -| `YAO_SANDBOX_WORKSPACE` | `{YAO_DATA_ROOT}/sandbox/workspace` | Workspace root directory | -| `YAO_SANDBOX_IPC` | `{YAO_DATA_ROOT}/sandbox/ipc` | IPC socket directory | -| `YAO_SANDBOX_MAX` | `100` | Max concurrent containers | -| `YAO_SANDBOX_IDLE_TIMEOUT` | `30m` | Idle timeout | -| `YAO_SANDBOX_MEMORY` | `2g` | Default memory limit | -| `YAO_SANDBOX_CPU` | `1.0` | Default CPU limit | - -### 9.2 Go Config Struct - -```go -// sandbox/config.go -type Config struct { - Image string `json:"image,omitempty" env:"YAO_SANDBOX_IMAGE" envDefault:"yao/sandbox-claude:latest"` - WorkspaceRoot string `json:"workspace_root,omitempty" env:"YAO_SANDBOX_WORKSPACE"` - IPCDir string `json:"ipc_dir,omitempty" env:"YAO_SANDBOX_IPC"` - MaxContainers int `json:"max_containers,omitempty" env:"YAO_SANDBOX_MAX" envDefault:"100"` - IdleTimeout time.Duration `json:"idle_timeout,omitempty" env:"YAO_SANDBOX_IDLE_TIMEOUT" envDefault:"30m"` - MaxMemory string `json:"max_memory,omitempty" env:"YAO_SANDBOX_MEMORY" envDefault:"2g"` - MaxCPU float64 `json:"max_cpu,omitempty" env:"YAO_SANDBOX_CPU" envDefault:"1.0"` -} - -// Init sets defaults based on Yao config -func (c *Config) Init(dataRoot string) { - if c.WorkspaceRoot == "" { - c.WorkspaceRoot = filepath.Join(dataRoot, "sandbox", "workspace") - } - if c.IPCDir == "" { - c.IPCDir = filepath.Join(dataRoot, "sandbox", "ipc") - } -} -``` - -### 9.3 app.yao (optional override) +New: lifecycle policy set per-assistant config: ```yaml sandbox: - image: "yao/sandbox-claude:full" - max_memory: "4g" + lifecycle: session # one-shot | session | long-running | persistent + idle_timeout: 30m + image: yaoapp/workspace:latest ``` -### 9.4 Assistant-level Configuration (package.yao) +`initSandbox()` passes the policy to `sandbox.Manager`, which enforces TTL and cleanup. `sandboxCleanup()` only disconnects the executor โ€” the Manager decides whether to actually remove the container based on policy. -```yaml -name: "My Coder" -type: claude +## Migration Path -sandbox: - image: "yao/sandbox-claude:full" # Override image - max_memory: "4g" # Override memory limit -``` - -### 9.5 Image Resolution - -``` -1. If package.yao sandbox.image is set โ†’ use it -2. Else if type is set โ†’ use yao/sandbox-{type}:latest -3. Else โ†’ use YAO_SANDBOX_IMAGE (or app.yao sandbox.image) -``` - ---- - -## 10. Data Persistence - -### Directory Structure - -``` -{YAO_DATA_ROOT}/sandbox/ -โ”œโ”€โ”€ workspace/ -โ”‚ โ””โ”€โ”€ {userID}/ -โ”‚ โ”œโ”€โ”€ {chatID-1}/ # Mounted as /workspace in container -โ”‚ โ”‚ โ”œโ”€โ”€ .mcp.json # MCP configuration -โ”‚ โ”‚ โ”œโ”€โ”€ .claude/ # Claude configuration -โ”‚ โ”‚ โ”‚ โ””โ”€โ”€ skills/ # Skills symlink -โ”‚ โ”‚ โ”œโ”€โ”€ project/ # User project code -โ”‚ โ”‚ โ””โ”€โ”€ node_modules/ # Installed dependencies -โ”‚ โ”‚ -โ”‚ โ””โ”€โ”€ {chatID-2}/ -โ”‚ โ””โ”€โ”€ ... -โ”‚ -โ””โ”€โ”€ ipc/ - โ”œโ”€โ”€ {chatID-1}.sock # IPC socket - โ””โ”€โ”€ {chatID-2}.sock -``` - -### What Persists - -| Item | Location | Persists | -| ------------------ | --------------------- | ------------------------------ | -| User code | `/workspace/` | โœ… Yes (host mount) | -| Installed packages | Container filesystem | โœ… Yes (container persists) | -| Claude config | `/workspace/.claude/` | โœ… Yes | -| IPC socket | `/tmp/yao.sock` | โŒ No (recreated each session) | - ---- - -## 11. Stability Guarantees - -| Concern | Solution | -| ----------------------- | ----------------------------------------------------------------- | -| **Container isolation** | One container per user+chat | -| **IPC isolation** | One socket per session | -| **Resource limits** | Docker memory/CPU limits | -| **Idle cleanup** | Auto-stop after timeout (preserve data) | -| **Data persistence** | Workspace directory mount, container preserves installed packages | -| **Connection handling** | Goroutine detects EOF, auto-cleanup | -| **Concurrency safety** | sync.Map + dedicated goroutines | - ---- - -## 12. Claude Executor Integration - -### Execution Flow - -```go -func (e *ClaudeExecutor) Stream(ctx *context.Context, messages []context.Message, opts ...*context.Options) (*context.Response, error) { - - // 1. Get or create sandbox container - container, err := e.SandboxManager.GetOrCreate(ctx, ctx.User.ID, ctx.ChatID) - if err != nil { - return nil, fmt.Errorf("failed to get sandbox: %w", err) - } - - // 2. Create IPC session - mcpTools := e.getMCPTools() - ipcSession, err := e.IPCManager.Create(ctx, ctx.ChatID, &AgentContext{ - UserID: ctx.User.ID, - ChatID: ctx.ChatID, - Locale: ctx.Locale, - }, mcpTools) - if err != nil { - return nil, fmt.Errorf("failed to create IPC session: %w", err) - } - defer e.IPCManager.Close(ctx.ChatID) - - // 3. Generate .mcp.json - if err := e.writeMCPConfig(ctx, container); err != nil { - return nil, err - } - - // 4. Setup skills - if err := e.setupSkills(container); err != nil { - return nil, err - } - - // 5. Build Claude CLI arguments - args := e.buildArgs(ctx, messages) - - // 6. Execute Claude CLI in container - stdout, err := e.SandboxManager.Stream(ctx, container.Name, - append([]string{"claude"}, args...), - &ExecOptions{ - WorkDir: "/workspace", - Env: e.buildEnvMap(ctx), - Timeout: e.getTimeout(), - }, - ) - if err != nil { - return nil, err - } - defer stdout.Close() - - // 7. Parse stream-json output - return e.parseClaudeOutput(ctx, stdout) -} -``` - -### MCP Configuration Generation - -```go -func (e *ClaudeExecutor) writeMCPConfig(ctx context.Context, container *Container) error { - config := map[string]interface{}{ - "mcpServers": map[string]interface{}{ - "yao": map[string]interface{}{ - "command": "yao-bridge", - "args": []string{"/tmp/yao.sock"}, - }, - }, - } - - // Add other MCP servers (external stdio/sse) - for _, server := range e.Assistant.MCP.Servers { - if server.Transport != "process" { - config["mcpServers"].(map[string]interface{})[server.Name] = server.ToClaudeConfig() - } - } - - data, _ := json.MarshalIndent(config, "", " ") - return e.SandboxManager.WriteFile(ctx, container.Name, "/workspace/.mcp.json", data) -} -``` - ---- - -## 13. Security Considerations - -| Layer | Measures | -| -------------- | ------------------------------------------------------ | -| **Filesystem** | Only workspace mounted, host filesystem not accessible | -| **Network** | Can be restricted with `--network none` if needed | -| **Privileges** | `--cap-drop ALL`, `no-new-privileges` | -| **Resources** | Memory and CPU limits | -| **User** | Non-root user inside container | -| **IPC** | Per-session socket, authorized tools only | - ---- - -## 14. Summary - -| Aspect | Description | -| -------------------------- | ----------------------------------------- | -| **Core approach** | Persistent Docker containers | -| **Communication** | Unix Socket + MCP JSON-RPC | -| **Container granularity** | One container per user+chat | -| **Data persistence** | Workspace mount + container filesystem | -| **Dependency persistence** | npm/pip packages persist in container | -| **Security isolation** | Full isolation between users and sessions | -| **Resource control** | Memory, CPU, idle timeout | -| **Estimated code** | ~1200 lines | +1. **Phase 1:** Yao gRPC server โ€” expose process execution, replace Unix socket IPC +2. **Phase 2:** `sandbox.Manager` refactoring โ€” replace Docker client with `tai.Client`, unified file ops, new lifecycle model +3. **Phase 3:** Agent layer adaptation โ€” executor uses new Manager, IPC mode switch, lifecycle policy +4. **Phase 4:** `yao run --remote` โ€” CLI calls remote Yao via gRPC +5. **Phase 5:** Workspace persistence โ€” browser preview, service exposure, delivery diff --git a/sandbox/SPEC.md b/sandbox/SPEC.md new file mode 100644 index 00000000..c7613eda --- /dev/null +++ b/sandbox/SPEC.md @@ -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 | diff --git a/service/service.go b/service/service.go index 3af013c1..59f21df2 100644 --- a/service/service.go +++ b/service/service.go @@ -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 {