Update gRPC implementation in Tai SDK

- Mark Phase 4 and Phase 5 as complete, indicating successful removal of fixed upstream connections and the introduction of dynamic routing based on request metadata. - Implement `TokenManager` for handling authentication tokens and update gRPC client methods to support new features. - Enhance tests for dynamic routing and token management, achieving significant coverage improvements. - Begin preparations for Phase 6, outlining the structure for OAuth Device Flow and related tasks. This commit finalizes key gRPC features and sets the stage for upcoming authentication enhancements.
2026-03-04 14:58:03 +08:00 · 2026-03-04 14:58:03 +08:00 · daa4da763b
commit daa4da763b
parent 6e68efaba3
6 changed files with 1345 additions and 43 deletions
--- a/grpc/IMPL.md
+++ b/grpc/IMPL.md
@ -161,7 +161,7 @@ Depends on: Phase 1. No code dependency on Phase 2 — can parallel.

 Deliverable: LLM (unary + stream) and Agent streaming via gRPC.

-### Phase 4: Tai gateway change (Tai repo) ⏳
+### Phase 4: Tai gateway change (Tai repo) ✅

 Depends on: Phase 1 (need proto definitions for testing). yao-grpc depends on this.

@ -169,79 +169,131 @@ Tai gateway currently dials a fixed `YaoUpstream` at startup. New behavior: yao-

 | 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 |
+| Tai `gateway/gateway.go` | Remove fixed `upstream *grpc.ClientConn`. On each request, read `x-grpc-upstream` from metadata → lookup/create conn from `sync.Map` cache (key = address string) → forward. Typical deployment has 1 upstream, cache stays tiny. | ✅ Done |
+| Tai `server/server.go` | Remove `YaoUpstream` from `Config`. Gateway init no longer needs an address. | ✅ Done |
+| Tai `main.go` | Remove `--yao` flag, `TAI_YAO_UPSTREAM` env var, YAML `yao` field, and required check. | ✅ Done |
+| Tai `gateway/gateway_test.go` | Updated tests: dynamic routing, missing metadata → InvalidArgument, metadata forwarding (x-grpc-upstream stripped), upstream error propagation, multiple upstreams, connection cache. Coverage: 88.8%. | ✅ Done |

-Connection cache: `sync.Map[string, *grpc.ClientConn]` — lazy dial on first request per upstream, reuse thereafter. No eviction needed (upstream count ≈ 1 in practice). `GracefulStop` closes all cached connections.
+Connection cache: `sync.Map[string, *grpc.ClientConn]` — lazy dial on first request per upstream, reuse thereafter. No eviction needed (upstream count ≈ 1 in practice). `Close` closes all cached connections.

 Deliverable: Tai starts without Yao address. Forwards based on request metadata.

-### Phase 5: yao-grpc container client ⏳
+### 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 |
+| `tai/grpc/auth.go` | `TokenManager`: read `YAO_TOKEN` / `YAO_REFRESH_TOKEN` / `YAO_SANDBOX_ID` / `YAO_GRPC_UPSTREAM` from env. `YAO_GRPC_TAI=enable` triggers Tai relay mode (requires `YAO_GRPC_UPSTREAM`). Attach as gRPC metadata on every call via unary + stream interceptors. Auto-refresh from response headers. | ✅ Done |
+| `tai/grpc/grpc.go` | `Client`: `Dial(addr, TokenManager)`, `NewFromEnv()`. Method wrappers for all RPCs: Run, Shell, API, MCP (list/call/resources/read), ChatCompletions, ChatCompletionsStream, AgentStream, Healthz. | ✅ Done |
+| `tai/grpc/cmd/main.go` | Stdio MCP server: JSON-RPC → gRPC. `yao-grpc version` prints version/commit/build time (via `-ldflags`). `yao-grpc serve` reads stdin JSON-RPC, dispatches to gRPC client. | ✅ Done |
+| `tai/grpc/grpc_test.go` + `integration_test.go` | Black-box tests (package `grpc_test`). Unit: TokenManager metadata attachment, env parsing, refresh handling. Integration: real Yao gRPC server, all method wrappers, token refresh, auth rejection. Coverage: 83.9%. | ✅ Done |

 Container token issuance uses existing `oauth.MakeAccessToken` / `oauth.MakeRefreshToken` — called by sandbox Manager at container creation, injected as env vars. Revoke on Remove. No new auth code needed on the issuance side.

 Deliverable: `go build -o yao-grpc ./tai/grpc/cmd`.

-### Phase 6: Device Flow — backend (`yao login`) ⏳
+### Phase 6: Device Flow + CLI auth ⏳

-Depends on: Phase 1. Independent — can parallel with Phase 2-5.
+Depends on: Phase 1. Three sub-phases with sequential dependency: 6.1 → 6.2 → 6.3.
+
+#### Phase 6.1: OAuth Device Flow backend ⏳
+
+Backend endpoints for RFC 8628 Device Authorization Grant. Scaffolding already in place (`types.DeviceAuthorizationResponse`, `GrantTypeDeviceCode`, error codes, route registration).

 | Task | Detail | Status |
 |------|--------|--------|
-| `oauth/device.go` | Implement `DeviceAuthorization()` — generate `device_code` + `user_code`, store with expiry | ⏳ Pending |
-| `oauth/token.go` | Device code store/get/consume helpers | ⏳ Pending |
-| `oauth/core.go` | Add `GrantTypeDeviceCode` case → `handleDeviceCodeGrant()` (poll returns `authorization_pending` / token) | ⏳ Pending |
-| `cmd/yao/login.go` | `yao login --server <url>` → device flow → poll token endpoint → save `~/.yao/credentials` | ⏳ Pending |
-| `cmd/yao/logout.go` | Revoke + delete credentials | ⏳ Pending |
-| `cmd/yao/run.go` | Credentials exist → gRPC; otherwise local. Non-silent mode prints `⟶ user@host (gRPC)` header before execution (same line position as existing `Run: process.name`). Silent mode (`-s`) keeps pure output — no connection info, for shell scripting. | ⏳ Pending |
+| `oauth/token.go` | `deviceCodeKey`, `storeDeviceCode`, `getDeviceCodeData`, `consumeDeviceCode` — device_code storage/retrieval/consumption helpers using existing store infrastructure | ⏳ Pending |
+| `oauth/device.go` | Implement `DeviceAuthorization()` — generate `device_code` + `user_code` (crypto/rand), store with `DeviceCodeLifetime` expiry, return `DeviceAuthorizationResponse` | ⏳ Pending |
+| `oauth/core.go` | Add `case types.GrantTypeDeviceCode` → `handleDeviceCodeGrant()` — poll returns `authorization_pending` / `slow_down` / token | ⏳ Pending |
+| `openapi/oauth.go` | Replace hardcoded `oauthDeviceAuthorization` handler → call `openapi.OAuth.DeviceAuthorization()`. Add user authorization callback endpoint (`POST /oauth/device/authorize` — binds device_code to authenticated user). Fix discovery path (`/oauth/device` vs `/oauth/device_authorization`) | ⏳ Pending |

-Deliverable: `yao login` + `yao run` via gRPC (backend complete, auth page in Phase 7).
+Deliverable: Device flow endpoints functional — `POST /oauth/device_authorization` issues codes, `POST /oauth/token` with `grant_type=device_code` polls status.

-### Phase 7: Device Flow — CUI auth page (frontend) ⏳
+#### Phase 6.2: CUI auth/device page (frontend) ⏳

-Depends on: Phase 6 (backend endpoints ready). This is a **frontend-only** task in the CUI repo.
+Depends on: Phase 6.1 (backend endpoints). Frontend-only task in **CUI repo**.

 Route: `/auth/device` (Umi convention-based routing → `pages/auth/device/index.tsx`)

 | Task | Detail | Status |
 |------|--------|--------|
-| `pages/auth/device/index.tsx` | Device authorization page. User enters `user_code` and clicks Authorize. Uses `AuthLayout` + `AuthInput` + `AuthButton` from existing `pages/auth/components/`. | ⏳ Pending |
+| `pages/auth/device/index.tsx` | Device authorization page. User enters `user_code`, 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:**
+Implementation:

 - 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
+- Layout: `AuthLayout` (logo + theme switch), same as `/auth/entry`
+- Components: reuse `AuthInput` for `user_code` input, `AuthButton` for submit
+- Page export: `export default observer(DeviceAuth)`
+- API: `window.$app.openapi` → `POST /oauth/device/authorize` with `{ user_code }`, bearer token from session
+- Auth: must be logged in (redirect to `/auth/entry` if not). After authorizing, show success and close/redirect
+- i18n: `useIntl()`, `zh-CN` / `en-US`

-Deliverable: `/auth/device` page in CUI. User can authorize CLI device login from browser.
+Deliverable: `/auth/device` page. User authorizes CLI device login from browser.
+
+#### Phase 6.3: CLI commands + TUI status bar ⏳
+
+Depends on: Phase 6.1 (backend) + Phase 6.2 (CUI page for end-to-end `yao login`).
+
+**Credentials file** (`~/.yao/credentials`): base64-encoded JSON.
+
+```json
+{
+  "server": "https://yao.example.com",
+  "access_token": "eyJ...",
+  "refresh_token": "eyJ...",
+  "scope": "grpc:run grpc:stream grpc:shell grpc:llm grpc:agent grpc:mcp",
+  "user": "admin@example.com",
+  "expires_at": "2026-03-05T10:00:00Z"
+}
+```
+
+Stored as: `base64(json) → ~/.yao/credentials`. Prevents casual `cat` exposure.
+
+| Task | Detail | Status |
+|------|--------|--------|
+| `cmd/login.go` | `yao login --server <url>` — call device authorization endpoint, color-print device code + verification URL (no TUI), poll token endpoint with interval, on success base64-encode and save to `~/.yao/credentials` | ⏳ Pending |
+| `cmd/logout.go` | `yao logout` — read credentials, revoke token via server, delete `~/.yao/credentials` | ⏳ Pending |
+| `cmd/run.go` | Detect credentials → gRPC mode vs local mode. `--auth <path>` flag loads alternate credentials file (for bash scripting). `-s` (silent) mode: no TUI, pure output. gRPC mode with terminal: bubbletea TUI status bar. | ⏳ Pending |
+| `cmd/tui_status.go` | bubbletea `StatusBarModel` — top-line persistent bar showing `user@host (gRPC)` + scope summary. Does not interfere with process output below. Uses existing bubbletea + lipgloss deps. | ⏳ Pending |
+
+**`yao run` behavior matrix:**
+
+| Credentials | `-s` flag | `--auth` flag | Behavior |
+|-------------|-----------|---------------|----------|
+| None | — | — | Local execution (current behavior) |
+| `~/.yao/credentials` | No | — | gRPC + TUI status bar |
+| `~/.yao/credentials` | Yes | — | gRPC, no TUI, pure output |
+| — | Yes | `<path>` | gRPC via specified credentials, no TUI, pure output |
+| — | No | `<path>` | gRPC via specified credentials + TUI status bar |
+
+**TUI status bar** (bubbletea, `cmd/tui_status.go`):
+
+```
+┌─ admin@yao.example.com (gRPC) │ scope: run,stream,shell,llm,agent,mcp ─┐
+```
+
+- Top-line, persistent during execution
+- lipgloss styled (dim border, colored connection info)
+- Process output renders below, unaffected
+- Hidden in silent mode (`-s`)
+
+Deliverable: `yao login` + `yao logout` + `yao run` via gRPC with TUI status bar.

 ## V2 Phases

-### Phase 8: `gou/stream` package ⏳
+### Phase 7: `gou/stream` package ⏳

 | Task | Detail | Status |
 |------|--------|--------|
 | `gou/stream/` | ~150 lines. `Handler`, `Process`, `Register`, `New`, `Execute`. Fallback to process. | ⏳ Pending |
 | V8 | `stream.Register("scripts", ...)`, `ExecStream`, `template.Set("Stream", ...)`, JS `Stream()` global | ⏳ Pending |

-### Phase 9: Base streaming handlers ⏳
+### Phase 8: Base streaming handlers ⏳

-Depends on: Phase 8.
+Depends on: Phase 7.

 | Task | Detail | Status |
 |------|--------|--------|
@ -256,19 +308,24 @@ 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)
+    ├───────────┬───────────┬──────────────────────┐
+    ▼           ▼           ▼                      ▼
+Phase 2 ✅   Phase 3 ✅  Phase 4 ✅           Phase 6 (device flow + CLI)
+(handlers)  (LLM/Agent)  (Tai gateway)            │
+                            │              ┌───────┴───────┐
+                            ▼              ▼               ▼
+                         Phase 5 ✅     6.1 OAuth       6.2 CUI page
+                         (yao-grpc)    (backend)       (frontend)
+                                           │               │
+                                           └───────┬───────┘
+                                                   ▼
+                                              6.3 CMD + TUI
+                                         (login/logout/run)

 --- V2 ---

-Phase 8 (gou/stream)
+Phase 7 (gou/stream)
    │
    ▼
-Phase 9 (Stream, ShellStream)
+Phase 8 (Stream, ShellStream)
 ```
--- a/tai/grpc/auth.go
+++ b/tai/grpc/auth.go
@ -0,0 +1,147 @@
+package grpc
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"sync"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+// TokenManager reads auth credentials from environment variables and attaches
+// them as gRPC metadata on every call. It also handles automatic token refresh
+// by reading new tokens from response headers.
+type TokenManager struct {
+	mu           sync.RWMutex
+	accessToken  string
+	refreshToken string
+	sandboxID    string
+	upstream     string // only set when YAO_GRPC_TAI=enable
+	taiMode      bool
+}
+
+// NewTokenManagerFromEnv creates a TokenManager from environment variables.
+// Returns an error if required variables are missing.
+func NewTokenManagerFromEnv() (*TokenManager, error) {
+	tm := &TokenManager{
+		accessToken:  os.Getenv("YAO_TOKEN"),
+		refreshToken: os.Getenv("YAO_REFRESH_TOKEN"),
+		sandboxID:    os.Getenv("YAO_SANDBOX_ID"),
+	}
+
+	if os.Getenv("YAO_GRPC_TAI") == "enable" {
+		tm.taiMode = true
+		tm.upstream = os.Getenv("YAO_GRPC_UPSTREAM")
+		if tm.upstream == "" {
+			return nil, fmt.Errorf("YAO_GRPC_TAI=enable but YAO_GRPC_UPSTREAM is not set")
+		}
+	}
+
+	return tm, nil
+}
+
+// NewTokenManager creates a TokenManager with explicit values (for testing).
+func NewTokenManager(accessToken, refreshToken, sandboxID, upstream string) *TokenManager {
+	return &TokenManager{
+		accessToken:  accessToken,
+		refreshToken: refreshToken,
+		sandboxID:    sandboxID,
+		upstream:     upstream,
+		taiMode:      upstream != "",
+	}
+}
+
+// AttachMetadata returns a context with auth credentials in gRPC metadata.
+func (tm *TokenManager) AttachMetadata(ctx context.Context) context.Context {
+	tm.mu.RLock()
+	defer tm.mu.RUnlock()
+
+	pairs := []string{}
+	if tm.accessToken != "" {
+		pairs = append(pairs, "authorization", "Bearer "+tm.accessToken)
+	}
+	if tm.refreshToken != "" {
+		pairs = append(pairs, "x-refresh-token", tm.refreshToken)
+	}
+	if tm.sandboxID != "" {
+		pairs = append(pairs, "x-sandbox-id", tm.sandboxID)
+	}
+	if tm.taiMode && tm.upstream != "" {
+		pairs = append(pairs, "x-grpc-upstream", tm.upstream)
+	}
+
+	if len(pairs) == 0 {
+		return ctx
+	}
+	return metadata.AppendToOutgoingContext(ctx, pairs...)
+}
+
+// HandleResponseHeaders reads new tokens from response headers and updates
+// the in-memory credentials. Call after each gRPC response.
+func (tm *TokenManager) HandleResponseHeaders(header metadata.MD) {
+	if header == nil {
+		return
+	}
+
+	tm.mu.Lock()
+	defer tm.mu.Unlock()
+
+	if vals := header.Get("x-access-token"); len(vals) > 0 && vals[0] != "" {
+		tm.accessToken = vals[0]
+	}
+	if vals := header.Get("x-refresh-token"); len(vals) > 0 && vals[0] != "" {
+		tm.refreshToken = vals[0]
+	}
+}
+
+// UnaryInterceptor returns a gRPC unary client interceptor that attaches
+// auth metadata and handles token refresh from response headers.
+func (tm *TokenManager) UnaryInterceptor() grpc.UnaryClientInterceptor {
+	return func(ctx context.Context, method string, req, reply any,
+		cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+
+		ctx = tm.AttachMetadata(ctx)
+
+		var header metadata.MD
+		opts = append(opts, grpc.Header(&header))
+
+		err := invoker(ctx, method, req, reply, cc, opts...)
+		tm.HandleResponseHeaders(header)
+		return err
+	}
+}
+
+// StreamInterceptor returns a gRPC stream client interceptor that attaches
+// auth metadata. Token refresh from stream headers is handled by the caller
+// via stream.Header().
+func (tm *TokenManager) StreamInterceptor() grpc.StreamClientInterceptor {
+	return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn,
+		method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
+
+		ctx = tm.AttachMetadata(ctx)
+		stream, err := streamer(ctx, desc, cc, method, opts...)
+		if err != nil {
+			return nil, err
+		}
+
+		if header, hErr := stream.Header(); hErr == nil {
+			tm.HandleResponseHeaders(header)
+		}
+
+		return stream, nil
+	}
+}
+
+// AccessToken returns the current access token (for testing/debugging).
+func (tm *TokenManager) AccessToken() string {
+	tm.mu.RLock()
+	defer tm.mu.RUnlock()
+	return tm.accessToken
+}
+
+// IsTaiMode returns whether the client is configured for Tai relay mode.
+func (tm *TokenManager) IsTaiMode() bool {
+	return tm.taiMode
+}
--- a/tai/grpc/cmd/main.go
+++ b/tai/grpc/cmd/main.go
@ -0,0 +1,263 @@
+package main
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+	"os/signal"
+	"syscall"
+
+	yaogrpc "github.com/yaoapp/yao/tai/grpc"
+)
+
+// Build-time variables set via -ldflags.
+var (
+	Version   = "dev"
+	Commit    = "none"
+	BuildTime = "unknown"
+)
+
+func main() {
+	if len(os.Args) < 2 {
+		fmt.Fprintln(os.Stderr, "Usage: yao-grpc <version|serve>")
+		os.Exit(1)
+	}
+
+	switch os.Args[1] {
+	case "version":
+		fmt.Printf("yao-grpc %s (commit: %s, built: %s)\n", Version, Commit, BuildTime)
+	case "serve":
+		if err := serve(); err != nil {
+			fmt.Fprintf(os.Stderr, "error: %v\n", err)
+			os.Exit(1)
+		}
+	default:
+		fmt.Fprintf(os.Stderr, "Unknown command: %s\nUsage: yao-grpc <version|serve>\n", os.Args[1])
+		os.Exit(1)
+	}
+}
+
+// jsonrpcRequest is a minimal JSON-RPC 2.0 request.
+type jsonrpcRequest struct {
+	JSONRPC string          `json:"jsonrpc"`
+	ID      json.RawMessage `json:"id,omitempty"`
+	Method  string          `json:"method"`
+	Params  json.RawMessage `json:"params,omitempty"`
+}
+
+// jsonrpcResponse is a minimal JSON-RPC 2.0 response.
+type jsonrpcResponse struct {
+	JSONRPC string          `json:"jsonrpc"`
+	ID      json.RawMessage `json:"id,omitempty"`
+	Result  json.RawMessage `json:"result,omitempty"`
+	Error   *jsonrpcError   `json:"error,omitempty"`
+}
+
+type jsonrpcError struct {
+	Code    int    `json:"code"`
+	Message string `json:"message"`
+}
+
+func serve() error {
+	client, err := yaogrpc.NewFromEnv()
+	if err != nil {
+		return err
+	}
+	defer client.Close()
+
+	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer cancel()
+
+	scanner := bufio.NewScanner(os.Stdin)
+	scanner.Buffer(make([]byte, 0, 4*1024*1024), 4*1024*1024)
+	encoder := json.NewEncoder(os.Stdout)
+
+	for scanner.Scan() {
+		select {
+		case <-ctx.Done():
+			return nil
+		default:
+		}
+
+		line := scanner.Bytes()
+		if len(line) == 0 {
+			continue
+		}
+
+		var req jsonrpcRequest
+		if err := json.Unmarshal(line, &req); err != nil {
+			encoder.Encode(jsonrpcResponse{
+				JSONRPC: "2.0",
+				Error:   &jsonrpcError{Code: -32700, Message: "parse error"},
+			})
+			continue
+		}
+
+		resp := dispatch(ctx, client, &req)
+		encoder.Encode(resp)
+	}
+
+	if err := scanner.Err(); err != nil && err != io.EOF {
+		return fmt.Errorf("stdin read: %w", err)
+	}
+	return nil
+}
+
+func dispatch(ctx context.Context, client *yaogrpc.Client, req *jsonrpcRequest) jsonrpcResponse {
+	base := jsonrpcResponse{JSONRPC: "2.0", ID: req.ID}
+
+	switch req.Method {
+	case "run":
+		return handleRun(ctx, client, req, base)
+	case "shell":
+		return handleShell(ctx, client, req, base)
+	case "mcp/list_tools":
+		return handleMCPListTools(ctx, client, req, base)
+	case "mcp/call_tool":
+		return handleMCPCallTool(ctx, client, req, base)
+	case "mcp/list_resources":
+		return handleMCPListResources(ctx, client, req, base)
+	case "mcp/read_resource":
+		return handleMCPReadResource(ctx, client, req, base)
+	case "healthz":
+		return handleHealthz(ctx, client, base)
+	default:
+		base.Error = &jsonrpcError{Code: -32601, Message: "method not found: " + req.Method}
+		return base
+	}
+}
+
+// --- handlers ---
+
+type runParams struct {
+	Process string          `json:"process"`
+	Args    json.RawMessage `json:"args,omitempty"`
+	Timeout int32           `json:"timeout,omitempty"`
+}
+
+func handleRun(ctx context.Context, c *yaogrpc.Client, req *jsonrpcRequest, base jsonrpcResponse) jsonrpcResponse {
+	var p runParams
+	if err := json.Unmarshal(req.Params, &p); err != nil {
+		base.Error = &jsonrpcError{Code: -32602, Message: "invalid params: " + err.Error()}
+		return base
+	}
+	data, err := c.Run(ctx, p.Process, p.Args, p.Timeout)
+	if err != nil {
+		base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
+		return base
+	}
+	base.Result = data
+	return base
+}
+
+type shellParams struct {
+	Command string            `json:"command"`
+	Args    []string          `json:"args,omitempty"`
+	Env     map[string]string `json:"env,omitempty"`
+	Timeout int32             `json:"timeout,omitempty"`
+}
+
+func handleShell(ctx context.Context, c *yaogrpc.Client, req *jsonrpcRequest, base jsonrpcResponse) jsonrpcResponse {
+	var p shellParams
+	if err := json.Unmarshal(req.Params, &p); err != nil {
+		base.Error = &jsonrpcError{Code: -32602, Message: "invalid params: " + err.Error()}
+		return base
+	}
+	resp, err := c.Shell(ctx, p.Command, p.Args, p.Env, p.Timeout)
+	if err != nil {
+		base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
+		return base
+	}
+	data, _ := json.Marshal(resp)
+	base.Result = data
+	return base
+}
+
+type mcpSessionParams struct {
+	SessionID string `json:"session_id"`
+}
+
+func handleMCPListTools(ctx context.Context, c *yaogrpc.Client, req *jsonrpcRequest, base jsonrpcResponse) jsonrpcResponse {
+	var p mcpSessionParams
+	if err := json.Unmarshal(req.Params, &p); err != nil {
+		base.Error = &jsonrpcError{Code: -32602, Message: "invalid params: " + err.Error()}
+		return base
+	}
+	data, err := c.MCPListTools(ctx, p.SessionID)
+	if err != nil {
+		base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
+		return base
+	}
+	base.Result = data
+	return base
+}
+
+type mcpCallParams struct {
+	SessionID string          `json:"session_id"`
+	Tool      string          `json:"tool"`
+	Arguments json.RawMessage `json:"arguments,omitempty"`
+}
+
+func handleMCPCallTool(ctx context.Context, c *yaogrpc.Client, req *jsonrpcRequest, base jsonrpcResponse) jsonrpcResponse {
+	var p mcpCallParams
+	if err := json.Unmarshal(req.Params, &p); err != nil {
+		base.Error = &jsonrpcError{Code: -32602, Message: "invalid params: " + err.Error()}
+		return base
+	}
+	data, err := c.MCPCallTool(ctx, p.SessionID, p.Tool, p.Arguments)
+	if err != nil {
+		base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
+		return base
+	}
+	base.Result = data
+	return base
+}
+
+func handleMCPListResources(ctx context.Context, c *yaogrpc.Client, req *jsonrpcRequest, base jsonrpcResponse) jsonrpcResponse {
+	var p mcpSessionParams
+	if err := json.Unmarshal(req.Params, &p); err != nil {
+		base.Error = &jsonrpcError{Code: -32602, Message: "invalid params: " + err.Error()}
+		return base
+	}
+	data, err := c.MCPListResources(ctx, p.SessionID)
+	if err != nil {
+		base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
+		return base
+	}
+	base.Result = data
+	return base
+}
+
+type mcpReadParams struct {
+	SessionID string `json:"session_id"`
+	URI       string `json:"uri"`
+}
+
+func handleMCPReadResource(ctx context.Context, c *yaogrpc.Client, req *jsonrpcRequest, base jsonrpcResponse) jsonrpcResponse {
+	var p mcpReadParams
+	if err := json.Unmarshal(req.Params, &p); err != nil {
+		base.Error = &jsonrpcError{Code: -32602, Message: "invalid params: " + err.Error()}
+		return base
+	}
+	data, err := c.MCPReadResource(ctx, p.SessionID, p.URI)
+	if err != nil {
+		base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
+		return base
+	}
+	base.Result = data
+	return base
+}
+
+func handleHealthz(ctx context.Context, c *yaogrpc.Client, base jsonrpcResponse) jsonrpcResponse {
+	status, err := c.Healthz(ctx)
+	if err != nil {
+		base.Error = &jsonrpcError{Code: -32000, Message: err.Error()}
+		return base
+	}
+	data, _ := json.Marshal(map[string]string{"status": status})
+	base.Result = data
+	return base
+}
--- a/tai/grpc/grpc.go
+++ b/tai/grpc/grpc.go
@ -0,0 +1,240 @@
+package grpc
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/yaoapp/yao/grpc/pb"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+// Client wraps a gRPC connection to a Yao server (direct or via Tai relay).
+// TokenManager handles auth metadata attachment and token refresh automatically.
+type Client struct {
+	conn  *grpc.ClientConn
+	svc   pb.YaoClient
+	token *TokenManager
+}
+
+// NewFromEnv reads YAO_GRPC_ADDR (required) and token env vars, dials the
+// gRPC server, and returns a connected Client.
+func NewFromEnv() (*Client, error) {
+	addr := os.Getenv("YAO_GRPC_ADDR")
+	if addr == "" {
+		return nil, fmt.Errorf("YAO_GRPC_ADDR is required")
+	}
+
+	tm, err := NewTokenManagerFromEnv()
+	if err != nil {
+		return nil, err
+	}
+
+	return Dial(addr, tm)
+}
+
+// Dial connects to the gRPC server at addr with the given TokenManager.
+func Dial(addr string, tm *TokenManager) (*Client, error) {
+	opts := []grpc.DialOption{
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+	}
+	if tm != nil {
+		opts = append(opts,
+			grpc.WithUnaryInterceptor(tm.UnaryInterceptor()),
+			grpc.WithStreamInterceptor(tm.StreamInterceptor()),
+		)
+	}
+
+	conn, err := grpc.NewClient(addr, opts...)
+	if err != nil {
+		return nil, fmt.Errorf("dial %s: %w", addr, err)
+	}
+
+	return &Client{
+		conn:  conn,
+		svc:   pb.NewYaoClient(conn),
+		token: tm,
+	}, nil
+}
+
+// Close releases the gRPC connection.
+func (c *Client) Close() error {
+	if c.conn != nil {
+		return c.conn.Close()
+	}
+	return nil
+}
+
+// Conn returns the underlying gRPC connection.
+func (c *Client) Conn() *grpc.ClientConn { return c.conn }
+
+// TokenManager returns the client's token manager.
+func (c *Client) TokenManager() *TokenManager { return c.token }
+
+// --- Base ---
+
+// Run executes a Yao process and returns the JSON-encoded result.
+func (c *Client) Run(ctx context.Context, process string, args []byte, timeout int32) ([]byte, error) {
+	resp, err := c.svc.Run(ctx, &pb.RunRequest{
+		Process: process,
+		Args:    args,
+		Timeout: timeout,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return resp.Data, nil
+}
+
+// Shell executes a system command and returns stdout, stderr, exit code.
+func (c *Client) Shell(ctx context.Context, command string, args []string, env map[string]string, timeout int32) (*pb.ShellResponse, error) {
+	return c.svc.Shell(ctx, &pb.ShellRequest{
+		Command: command,
+		Args:    args,
+		Env:     env,
+		Timeout: timeout,
+	})
+}
+
+// --- API ---
+
+// API proxies an HTTP request through the gRPC gateway.
+func (c *Client) API(ctx context.Context, method, path string, headers map[string]string, body []byte) (*pb.APIResponse, error) {
+	return c.svc.API(ctx, &pb.APIRequest{
+		Method:  method,
+		Path:    path,
+		Headers: headers,
+		Body:    body,
+	})
+}
+
+// --- MCP ---
+
+// MCPListTools lists available MCP tools for a session.
+func (c *Client) MCPListTools(ctx context.Context, sessionID string) ([]byte, error) {
+	resp, err := c.svc.MCPListTools(ctx, &pb.MCPListRequest{SessionId: sessionID})
+	if err != nil {
+		return nil, err
+	}
+	return resp.Tools, nil
+}
+
+// MCPCallTool calls an MCP tool and returns the JSON result.
+func (c *Client) MCPCallTool(ctx context.Context, sessionID, tool string, arguments []byte) ([]byte, error) {
+	resp, err := c.svc.MCPCallTool(ctx, &pb.MCPCallRequest{
+		SessionId: sessionID,
+		Tool:      tool,
+		Arguments: arguments,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return resp.Result, nil
+}
+
+// MCPListResources lists available MCP resources for a session.
+func (c *Client) MCPListResources(ctx context.Context, sessionID string) ([]byte, error) {
+	resp, err := c.svc.MCPListResources(ctx, &pb.MCPListRequest{SessionId: sessionID})
+	if err != nil {
+		return nil, err
+	}
+	return resp.Resources, nil
+}
+
+// MCPReadResource reads an MCP resource by URI.
+func (c *Client) MCPReadResource(ctx context.Context, sessionID, uri string) ([]byte, error) {
+	resp, err := c.svc.MCPReadResource(ctx, &pb.MCPResourceRequest{
+		SessionId: sessionID,
+		Uri:       uri,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return resp.Contents, nil
+}
+
+// --- LLM ---
+
+// ChatCompletions sends a chat completion request and returns the result.
+func (c *Client) ChatCompletions(ctx context.Context, connector string, messages, options []byte) ([]byte, error) {
+	resp, err := c.svc.ChatCompletions(ctx, &pb.ChatRequest{
+		Connector: connector,
+		Messages:  messages,
+		Options:   options,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return resp.Data, nil
+}
+
+// ChatCompletionsStream sends a streaming chat completion request.
+// The callback receives each chunk's data; return a non-nil error to stop.
+func (c *Client) ChatCompletionsStream(ctx context.Context, connector string, messages, options []byte, cb func(data []byte, done bool) error) error {
+	stream, err := c.svc.ChatCompletionsStream(ctx, &pb.ChatRequest{
+		Connector: connector,
+		Messages:  messages,
+		Options:   options,
+	})
+	if err != nil {
+		return err
+	}
+	for {
+		chunk, err := stream.Recv()
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+		if err := cb(chunk.Data, chunk.Done); err != nil {
+			return err
+		}
+		if chunk.Done {
+			return nil
+		}
+	}
+}
+
+// --- Agent ---
+
+// AgentStream calls an agent with streaming response.
+// The callback receives each chunk's data; return a non-nil error to stop.
+func (c *Client) AgentStream(ctx context.Context, assistantID string, messages, options []byte, cb func(data []byte, done bool) error) error {
+	stream, err := c.svc.AgentStream(ctx, &pb.AgentRequest{
+		AssistantId: assistantID,
+		Messages:    messages,
+		Options:     options,
+	})
+	if err != nil {
+		return err
+	}
+	for {
+		chunk, err := stream.Recv()
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+		if err := cb(chunk.Data, chunk.Done); err != nil {
+			return err
+		}
+		if chunk.Done {
+			return nil
+		}
+	}
+}
+
+// --- Health ---
+
+// Healthz checks the server health.
+func (c *Client) Healthz(ctx context.Context) (string, error) {
+	resp, err := c.svc.Healthz(ctx, &pb.Empty{})
+	if err != nil {
+		return "", err
+	}
+	return resp.Status, nil
+}
--- a/tai/grpc/grpc_test.go
+++ b/tai/grpc/grpc_test.go
@ -0,0 +1,175 @@
+package grpc_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc/metadata"
+
+	yaogrpc "github.com/yaoapp/yao/tai/grpc"
+)
+
+// ── TokenManager unit tests ──────────────────────────────────────────────────
+
+func TestTokenManager_AttachMetadata_WithAllFields(t *testing.T) {
+	tm := yaogrpc.NewTokenManager("tok", "ref", "sb-1", "yao:9099")
+	ctx := tm.AttachMetadata(context.Background())
+
+	md, ok := metadata.FromOutgoingContext(ctx)
+	require.True(t, ok)
+
+	assert.Equal(t, []string{"Bearer tok"}, md.Get("authorization"))
+	assert.Equal(t, []string{"ref"}, md.Get("x-refresh-token"))
+	assert.Equal(t, []string{"sb-1"}, md.Get("x-sandbox-id"))
+	assert.Equal(t, []string{"yao:9099"}, md.Get("x-grpc-upstream"))
+}
+
+func TestTokenManager_AttachMetadata_DirectMode(t *testing.T) {
+	tm := yaogrpc.NewTokenManager("tok", "ref", "sb-1", "")
+	ctx := tm.AttachMetadata(context.Background())
+
+	md, ok := metadata.FromOutgoingContext(ctx)
+	require.True(t, ok)
+
+	assert.Equal(t, []string{"Bearer tok"}, md.Get("authorization"))
+	assert.Empty(t, md.Get("x-grpc-upstream"), "direct mode should not set x-grpc-upstream")
+}
+
+func TestTokenManager_AttachMetadata_EmptyTokens(t *testing.T) {
+	tm := yaogrpc.NewTokenManager("", "", "", "")
+	ctx := tm.AttachMetadata(context.Background())
+
+	_, ok := metadata.FromOutgoingContext(ctx)
+	assert.False(t, ok, "empty tokens should not produce metadata")
+}
+
+func TestTokenManager_HandleResponseHeaders(t *testing.T) {
+	tm := yaogrpc.NewTokenManager("old-tok", "old-ref", "", "")
+
+	tm.HandleResponseHeaders(metadata.New(map[string]string{
+		"x-access-token":  "new-tok",
+		"x-refresh-token": "new-ref",
+	}))
+
+	assert.Equal(t, "new-tok", tm.AccessToken())
+
+	ctx := tm.AttachMetadata(context.Background())
+	md, _ := metadata.FromOutgoingContext(ctx)
+	assert.Equal(t, []string{"Bearer new-tok"}, md.Get("authorization"))
+	assert.Equal(t, []string{"new-ref"}, md.Get("x-refresh-token"))
+}
+
+func TestTokenManager_HandleResponseHeaders_Nil(t *testing.T) {
+	tm := yaogrpc.NewTokenManager("tok", "", "", "")
+	tm.HandleResponseHeaders(nil)
+	assert.Equal(t, "tok", tm.AccessToken())
+}
+
+func TestTokenManager_HandleResponseHeaders_EmptyValues(t *testing.T) {
+	tm := yaogrpc.NewTokenManager("tok", "ref", "", "")
+	tm.HandleResponseHeaders(metadata.New(map[string]string{
+		"x-access-token": "",
+	}))
+	assert.Equal(t, "tok", tm.AccessToken(), "empty header should not overwrite")
+}
+
+func TestTokenManager_IsTaiMode(t *testing.T) {
+	tmDirect := yaogrpc.NewTokenManager("tok", "", "", "")
+	assert.False(t, tmDirect.IsTaiMode())
+
+	tmTai := yaogrpc.NewTokenManager("tok", "", "", "tai:9100")
+	assert.True(t, tmTai.IsTaiMode())
+}
+
+func TestTokenManager_NewFromEnv_MissingUpstream(t *testing.T) {
+	t.Setenv("YAO_GRPC_TAI", "enable")
+	t.Setenv("YAO_GRPC_UPSTREAM", "")
+	t.Setenv("YAO_TOKEN", "tok")
+
+	_, err := yaogrpc.NewTokenManagerFromEnv()
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "YAO_GRPC_UPSTREAM")
+}
+
+func TestTokenManager_NewFromEnv_TaiEnabled(t *testing.T) {
+	t.Setenv("YAO_GRPC_TAI", "enable")
+	t.Setenv("YAO_GRPC_UPSTREAM", "yao:9099")
+	t.Setenv("YAO_TOKEN", "my-token")
+	t.Setenv("YAO_REFRESH_TOKEN", "my-refresh")
+	t.Setenv("YAO_SANDBOX_ID", "sb-42")
+
+	tm, err := yaogrpc.NewTokenManagerFromEnv()
+	require.NoError(t, err)
+	assert.True(t, tm.IsTaiMode())
+	assert.Equal(t, "my-token", tm.AccessToken())
+}
+
+func TestTokenManager_NewFromEnv_DirectMode(t *testing.T) {
+	t.Setenv("YAO_GRPC_TAI", "")
+	t.Setenv("YAO_GRPC_UPSTREAM", "")
+	t.Setenv("YAO_TOKEN", "tok")
+
+	tm, err := yaogrpc.NewTokenManagerFromEnv()
+	require.NoError(t, err)
+	assert.False(t, tm.IsTaiMode())
+}
+
+// ── Dial tests ───────────────────────────────────────────────────────────────
+
+func TestNewFromEnv_MissingAddr(t *testing.T) {
+	t.Setenv("YAO_GRPC_ADDR", "")
+	_, err := yaogrpc.NewFromEnv()
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "YAO_GRPC_ADDR")
+}
+
+func TestNewFromEnv_Success(t *testing.T) {
+	t.Setenv("YAO_GRPC_ADDR", "127.0.0.1:9099")
+	t.Setenv("YAO_TOKEN", "test-token")
+	t.Setenv("YAO_REFRESH_TOKEN", "test-refresh")
+	t.Setenv("YAO_SANDBOX_ID", "sb-1")
+	t.Setenv("YAO_GRPC_TAI", "")
+
+	c, err := yaogrpc.NewFromEnv()
+	require.NoError(t, err)
+	defer c.Close()
+
+	assert.NotNil(t, c.Conn())
+	assert.Equal(t, "test-token", c.TokenManager().AccessToken())
+}
+
+func TestNewFromEnv_TaiMode_MissingUpstream(t *testing.T) {
+	t.Setenv("YAO_GRPC_ADDR", "tai:9100")
+	t.Setenv("YAO_GRPC_TAI", "enable")
+	t.Setenv("YAO_GRPC_UPSTREAM", "")
+
+	_, err := yaogrpc.NewFromEnv()
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "YAO_GRPC_UPSTREAM")
+}
+
+func TestDial_WithNilTokenManager(t *testing.T) {
+	c, err := yaogrpc.Dial("127.0.0.1:0", nil)
+	require.NoError(t, err)
+	defer c.Close()
+
+	assert.NotNil(t, c.Conn())
+	assert.Nil(t, c.TokenManager())
+}
+
+func TestDial_WithTokenManager(t *testing.T) {
+	tm := yaogrpc.NewTokenManager("tok", "", "", "")
+	c, err := yaogrpc.Dial("127.0.0.1:0", tm)
+	require.NoError(t, err)
+	defer c.Close()
+
+	assert.NotNil(t, c.TokenManager())
+	assert.False(t, c.TokenManager().IsTaiMode())
+}
+
+func TestClient_Close_Nil(t *testing.T) {
+	c := &yaogrpc.Client{}
+	assert.NoError(t, c.Close())
+}
--- a/tai/grpc/integration_test.go
+++ b/tai/grpc/integration_test.go
@ -0,0 +1,420 @@
+package grpc_test
+
+import (
+	"context"
+	"encoding/json"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/yaoapp/yao/grpc/tests/testutils"
+	yaogrpc "github.com/yaoapp/yao/tai/grpc"
+)
+
+// Integration tests that start a real Yao gRPC server and test the tai/grpc
+// client through the full interceptor -> handler chain.
+
+func setupClient(t *testing.T, scopes ...string) *yaogrpc.Client {
+	t.Helper()
+
+	conn := testutils.Prepare(t)
+	t.Cleanup(func() {
+		conn.Close()
+		testutils.Clean()
+	})
+
+	addr := testutils.Addr()
+	token := testutils.ObtainAccessToken(t, scopes...)
+	refreshToken := testutils.ObtainRefreshToken(t, scopes...)
+
+	tm := yaogrpc.NewTokenManager(token, refreshToken, "test-sandbox", "")
+	client, err := yaogrpc.Dial(addr, tm)
+	require.NoError(t, err)
+	t.Cleanup(func() { client.Close() })
+
+	return client
+}
+
+// ── Healthz ──────────────────────────────────────────────────────────────────
+
+func TestIntegration_Healthz(t *testing.T) {
+	conn := testutils.Prepare(t)
+	defer func() {
+		conn.Close()
+		testutils.Clean()
+	}()
+
+	addr := testutils.Addr()
+	tm := yaogrpc.NewTokenManager("", "", "", "")
+	client, err := yaogrpc.Dial(addr, tm)
+	require.NoError(t, err)
+	defer client.Close()
+
+	status, err := client.Healthz(context.Background())
+	assert.NoError(t, err)
+	assert.Equal(t, "ok", status)
+}
+
+// ── Run ──────────────────────────────────────────────────────────────────────
+
+func TestIntegration_Run_Ping(t *testing.T) {
+	client := setupClient(t, "grpc:run")
+
+	data, err := client.Run(context.Background(), "utils.app.Ping", nil, 0)
+	assert.NoError(t, err)
+	assert.NotNil(t, data)
+}
+
+func TestIntegration_Run_InvalidProcess(t *testing.T) {
+	client := setupClient(t, "grpc:run")
+
+	_, err := client.Run(context.Background(), "nonexistent.process", nil, 0)
+	assert.Error(t, err)
+}
+
+func TestIntegration_Run_WithArgs(t *testing.T) {
+	client := setupClient(t, "grpc:run")
+
+	args, _ := json.Marshal([]any{"hello", "world"})
+	data, err := client.Run(context.Background(), "utils.app.Ping", args, 5)
+	assert.NoError(t, err)
+	assert.NotNil(t, data)
+}
+
+// ── Shell ────────────────────────────────────────────────────────────────────
+
+func TestIntegration_Shell_Echo(t *testing.T) {
+	client := setupClient(t, "grpc:shell")
+
+	resp, err := client.Shell(context.Background(), "echo", []string{"hello"}, nil, 5)
+	assert.NoError(t, err)
+	assert.Equal(t, int32(0), resp.ExitCode)
+	assert.Contains(t, string(resp.Stdout), "hello")
+}
+
+func TestIntegration_Shell_NotFound(t *testing.T) {
+	client := setupClient(t, "grpc:shell")
+
+	_, err := client.Shell(context.Background(), "nonexistent-command-xyz", nil, nil, 5)
+	assert.Error(t, err)
+}
+
+// ── MCP ──────────────────────────────────────────────────────────────────────
+
+func TestIntegration_MCPListTools(t *testing.T) {
+	client := setupClient(t, "grpc:mcp")
+
+	data, err := client.MCPListTools(context.Background(), "echo")
+	assert.NoError(t, err)
+	assert.NotNil(t, data)
+
+	var tools []any
+	assert.NoError(t, json.Unmarshal(data, &tools))
+	assert.Greater(t, len(tools), 0)
+}
+
+func TestIntegration_MCPCallTool(t *testing.T) {
+	client := setupClient(t, "grpc:mcp")
+
+	args, _ := json.Marshal(map[string]string{"message": "hi"})
+	data, err := client.MCPCallTool(context.Background(), "echo", "ping", args)
+	assert.NoError(t, err)
+	assert.NotNil(t, data)
+}
+
+func TestIntegration_MCPListResources(t *testing.T) {
+	client := setupClient(t, "grpc:mcp")
+
+	data, err := client.MCPListResources(context.Background(), "echo")
+	assert.NoError(t, err)
+	assert.NotNil(t, data)
+}
+
+func TestIntegration_MCPReadResource(t *testing.T) {
+	client := setupClient(t, "grpc:mcp")
+
+	data, err := client.MCPReadResource(context.Background(), "echo", "echo://info")
+	assert.NoError(t, err)
+	assert.NotNil(t, data)
+}
+
+// ── API ──────────────────────────────────────────────────────────────────────
+
+func TestIntegration_API_Proxy(t *testing.T) {
+	client := setupClient(t, "grpc:run", "grpc:mcp")
+
+	resp, err := client.API(context.Background(), "GET", "/api/__yao/app/setting", nil, nil)
+	assert.NoError(t, err)
+	assert.NotNil(t, resp)
+	// API proxy returns the response; the actual status depends on the route.
+	// A valid openapi path returns 200; anything else returns 404.
+	t.Logf("API proxy status: %d", resp.Status)
+}
+
+// ── LLM ──────────────────────────────────────────────────────────────────────
+
+func TestIntegration_ChatCompletions_InvalidConnector(t *testing.T) {
+	client := setupClient(t, "grpc:llm")
+
+	messages, _ := json.Marshal([]map[string]string{
+		{"role": "user", "content": "test"},
+	})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	_, err := client.ChatCompletions(ctx, "nonexistent-connector", messages, nil)
+	assert.Error(t, err)
+}
+
+func TestIntegration_ChatCompletionsStream_InvalidConnector(t *testing.T) {
+	client := setupClient(t, "grpc:llm")
+
+	messages, _ := json.Marshal([]map[string]string{
+		{"role": "user", "content": "test"},
+	})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	err := client.ChatCompletionsStream(ctx, "nonexistent-connector", messages, nil,
+		func(data []byte, done bool) error { return nil })
+	assert.Error(t, err)
+}
+
+func TestIntegration_ChatCompletions_EmptyMessages(t *testing.T) {
+	client := setupClient(t, "grpc:llm")
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	_, err := client.ChatCompletions(ctx, "default", nil, nil)
+	assert.Error(t, err)
+}
+
+// ── Agent ────────────────────────────────────────────────────────────────────
+
+func TestIntegration_AgentStream_InvalidRobot(t *testing.T) {
+	client := setupClient(t, "grpc:agent")
+
+	messages, _ := json.Marshal([]map[string]string{
+		{"role": "user", "content": "hello"},
+	})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	err := client.AgentStream(ctx, "nonexistent-robot-xyz", messages, nil,
+		func(data []byte, done bool) error { return nil })
+	assert.Error(t, err)
+}
+
+// ── Unauthenticated ──────────────────────────────────────────────────────────
+
+func TestIntegration_Run_NoToken(t *testing.T) {
+	conn := testutils.Prepare(t)
+	defer func() {
+		conn.Close()
+		testutils.Clean()
+	}()
+
+	addr := testutils.Addr()
+	tm := yaogrpc.NewTokenManager("", "", "", "")
+	client, err := yaogrpc.Dial(addr, tm)
+	require.NoError(t, err)
+	defer client.Close()
+
+	_, err = client.Run(context.Background(), "utils.app.Ping", nil, 0)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "Unauthenticated")
+}
+
+// ── Token Refresh via interceptor ────────────────────────────────────────────
+
+func TestIntegration_TokenRefresh(t *testing.T) {
+	conn := testutils.Prepare(t)
+	defer func() {
+		conn.Close()
+		testutils.Clean()
+	}()
+
+	addr := testutils.Addr()
+	scopes := []string{"grpc:run"}
+
+	expiredToken := testutils.ObtainExpiredAccessToken(t, scopes...)
+	refreshToken := testutils.ObtainRefreshToken(t, scopes...)
+
+	tm := yaogrpc.NewTokenManager(expiredToken, refreshToken, "sb-test", "")
+	client, err := yaogrpc.Dial(addr, tm)
+	require.NoError(t, err)
+	defer client.Close()
+
+	data, err := client.Run(context.Background(), "utils.app.Ping", nil, 0)
+	assert.NoError(t, err)
+	assert.NotNil(t, data)
+
+	newToken := tm.AccessToken()
+	if newToken != expiredToken {
+		t.Logf("token was refreshed: old=%s... new=%s...", expiredToken[:20], newToken[:20])
+	}
+}
+
+// ══════════════════════════════════════════════════════════════════════════════
+// Relay mode tests — client → Tai (:9100) → x-grpc-upstream → Yao gRPC
+// Requires TAI_TEST_GRPC env var (e.g. 127.0.0.1:9100) and a running Tai server.
+// ══════════════════════════════════════════════════════════════════════════════
+
+func setupRelayClient(t *testing.T, scopes ...string) *yaogrpc.Client {
+	t.Helper()
+
+	taiAddr := os.Getenv("TAI_TEST_GRPC")
+	if taiAddr == "" {
+		t.Skip("TAI_TEST_GRPC not set, skipping relay mode test")
+	}
+
+	conn := testutils.Prepare(t)
+	t.Cleanup(func() {
+		conn.Close()
+		testutils.Clean()
+	})
+
+	yaoAddr := testutils.Addr()
+	token := testutils.ObtainAccessToken(t, scopes...)
+	refreshToken := testutils.ObtainRefreshToken(t, scopes...)
+
+	// upstream = Yao gRPC address; taiMode = true
+	tm := yaogrpc.NewTokenManager(token, refreshToken, "relay-sandbox", yaoAddr)
+	client, err := yaogrpc.Dial(taiAddr, tm)
+	require.NoError(t, err)
+	t.Cleanup(func() { client.Close() })
+
+	return client
+}
+
+func TestRelay_Healthz(t *testing.T) {
+	taiAddr := os.Getenv("TAI_TEST_GRPC")
+	if taiAddr == "" {
+		t.Skip("TAI_TEST_GRPC not set")
+	}
+
+	conn := testutils.Prepare(t)
+	defer func() {
+		conn.Close()
+		testutils.Clean()
+	}()
+
+	yaoAddr := testutils.Addr()
+	tm := yaogrpc.NewTokenManager("", "", "", yaoAddr)
+	client, err := yaogrpc.Dial(taiAddr, tm)
+	require.NoError(t, err)
+	defer client.Close()
+
+	status, err := client.Healthz(context.Background())
+	assert.NoError(t, err)
+	assert.Equal(t, "ok", status)
+}
+
+func TestRelay_Run_Ping(t *testing.T) {
+	client := setupRelayClient(t, "grpc:run")
+
+	data, err := client.Run(context.Background(), "utils.app.Ping", nil, 0)
+	assert.NoError(t, err)
+	assert.NotNil(t, data)
+	t.Logf("relay Run result: %s", string(data))
+}
+
+func TestRelay_Run_InvalidProcess(t *testing.T) {
+	client := setupRelayClient(t, "grpc:run")
+
+	_, err := client.Run(context.Background(), "nonexistent.process", nil, 0)
+	assert.Error(t, err)
+}
+
+func TestRelay_Shell_Echo(t *testing.T) {
+	client := setupRelayClient(t, "grpc:shell")
+
+	resp, err := client.Shell(context.Background(), "echo", []string{"relay-test"}, nil, 5)
+	assert.NoError(t, err)
+	assert.Equal(t, int32(0), resp.ExitCode)
+	assert.Contains(t, string(resp.Stdout), "relay-test")
+}
+
+func TestRelay_MCPListTools(t *testing.T) {
+	client := setupRelayClient(t, "grpc:mcp")
+
+	data, err := client.MCPListTools(context.Background(), "echo")
+	assert.NoError(t, err)
+	assert.NotNil(t, data)
+
+	var tools []any
+	assert.NoError(t, json.Unmarshal(data, &tools))
+	assert.Greater(t, len(tools), 0)
+}
+
+func TestRelay_MCPCallTool(t *testing.T) {
+	client := setupRelayClient(t, "grpc:mcp")
+
+	args, _ := json.Marshal(map[string]string{"message": "relay"})
+	data, err := client.MCPCallTool(context.Background(), "echo", "ping", args)
+	assert.NoError(t, err)
+	assert.NotNil(t, data)
+}
+
+func TestRelay_Run_NoToken(t *testing.T) {
+	taiAddr := os.Getenv("TAI_TEST_GRPC")
+	if taiAddr == "" {
+		t.Skip("TAI_TEST_GRPC not set")
+	}
+
+	conn := testutils.Prepare(t)
+	defer func() {
+		conn.Close()
+		testutils.Clean()
+	}()
+
+	yaoAddr := testutils.Addr()
+	tm := yaogrpc.NewTokenManager("", "", "", yaoAddr)
+	client, err := yaogrpc.Dial(taiAddr, tm)
+	require.NoError(t, err)
+	defer client.Close()
+
+	_, err = client.Run(context.Background(), "utils.app.Ping", nil, 0)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "Unauthenticated")
+}
+
+func TestRelay_TokenRefresh(t *testing.T) {
+	taiAddr := os.Getenv("TAI_TEST_GRPC")
+	if taiAddr == "" {
+		t.Skip("TAI_TEST_GRPC not set")
+	}
+
+	conn := testutils.Prepare(t)
+	defer func() {
+		conn.Close()
+		testutils.Clean()
+	}()
+
+	yaoAddr := testutils.Addr()
+	scopes := []string{"grpc:run"}
+
+	expiredToken := testutils.ObtainExpiredAccessToken(t, scopes...)
+	refreshToken := testutils.ObtainRefreshToken(t, scopes...)
+
+	tm := yaogrpc.NewTokenManager(expiredToken, refreshToken, "relay-sb", yaoAddr)
+	client, err := yaogrpc.Dial(taiAddr, tm)
+	require.NoError(t, err)
+	defer client.Close()
+
+	data, err := client.Run(context.Background(), "utils.app.Ping", nil, 0)
+	assert.NoError(t, err)
+	assert.NotNil(t, data)
+
+	newToken := tm.AccessToken()
+	if newToken != expiredToken {
+		t.Logf("relay token refreshed: old=%s... new=%s...", expiredToken[:20], newToken[:20])
+	}
+}