From daa4da763bcd601a6e410bafe1979ac918140a3d Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 4 Mar 2026 14:58:03 +0800 Subject: [PATCH] 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. --- grpc/IMPL.md | 143 ++++++++---- tai/grpc/auth.go | 147 ++++++++++++ tai/grpc/cmd/main.go | 263 ++++++++++++++++++++++ tai/grpc/grpc.go | 240 ++++++++++++++++++++ tai/grpc/grpc_test.go | 175 +++++++++++++++ tai/grpc/integration_test.go | 420 +++++++++++++++++++++++++++++++++++ 6 files changed, 1345 insertions(+), 43 deletions(-) create mode 100644 tai/grpc/auth.go create mode 100644 tai/grpc/cmd/main.go create mode 100644 tai/grpc/grpc.go create mode 100644 tai/grpc/grpc_test.go create mode 100644 tai/grpc/integration_test.go diff --git a/grpc/IMPL.md b/grpc/IMPL.md index fefc76f0..adad4644 100644 --- 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 ` → 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 ` — 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 ` 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 | `` | gRPC via specified credentials, no TUI, pure output | +| — | No | `` | 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) ``` diff --git a/tai/grpc/auth.go b/tai/grpc/auth.go new file mode 100644 index 00000000..557005e9 --- /dev/null +++ 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 +} diff --git a/tai/grpc/cmd/main.go b/tai/grpc/cmd/main.go new file mode 100644 index 00000000..e7ab70b9 --- /dev/null +++ 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 ") + 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 \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 +} diff --git a/tai/grpc/grpc.go b/tai/grpc/grpc.go new file mode 100644 index 00000000..4f538bcd --- /dev/null +++ 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 +} diff --git a/tai/grpc/grpc_test.go b/tai/grpc/grpc_test.go new file mode 100644 index 00000000..47057843 --- /dev/null +++ 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()) +} diff --git a/tai/grpc/integration_test.go b/tai/grpc/integration_test.go new file mode 100644 index 00000000..02a36a51 --- /dev/null +++ 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]) + } +}