From 7772cc588f3c974f35c78f2d6ad8856590e6cb4e Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 3 Mar 2026 23:50:47 +0800 Subject: [PATCH 1/7] Remove DESIGN.md file from Tai Go SDK, eliminating outdated documentation on the SDK's architecture, usage, and package layout. --- tai/DESIGN.md | 147 ---------------------------------- tai/docs/README.md | 111 ++++++++++++++++++++++++++ tai/docs/proxy.md | 86 ++++++++++++++++++++ tai/docs/sandbox.md | 182 ++++++++++++++++++++++++++++++++++++++++++ tai/docs/vnc.md | 94 ++++++++++++++++++++++ tai/docs/volume.md | 120 ++++++++++++++++++++++++++++ tai/docs/workspace.md | 93 +++++++++++++++++++++ 7 files changed, 686 insertions(+), 147 deletions(-) delete mode 100644 tai/DESIGN.md create mode 100644 tai/docs/README.md create mode 100644 tai/docs/proxy.md create mode 100644 tai/docs/sandbox.md create mode 100644 tai/docs/vnc.md create mode 100644 tai/docs/volume.md create mode 100644 tai/docs/workspace.md diff --git a/tai/DESIGN.md b/tai/DESIGN.md deleted file mode 100644 index 1565fa32..00000000 --- a/tai/DESIGN.md +++ /dev/null @@ -1,147 +0,0 @@ -# Tai Go SDK - -Go client library for [Tai](https://github.com/yaoapp/tai) — the universal runtime bridge for Yao Sandbox. - -## Overview - -Provides a unified API for container lifecycle, filesystem operations, HTTP proxy, and VNC access. -Supports two modes via a single entry point: - -- **Local** (`docker://` or `""`) — direct Docker daemon connection -- **Remote** (`tai://host`) — via Tai Server proxy (Docker, K8s) - -All sub-packages follow the same pattern: **interface + Remote/Local implementations**. - -## Package Layout - -``` -yao/tai/ -├── tai.go # Client, New(), Option, Close() -├── volume/ # Volume IO + Sync -├── workspace/ # Go fs.FS wrapper over volume.Volume -├── sandbox/ # Container lifecycle (Create/Start/Stop/Exec/Remove) -│ ├── sandbox.go # Interface + shared types -│ ├── local.go # Direct Docker socket -│ ├── docker.go # Docker via Tai proxy -│ ├── docker_core.go # Shared Docker SDK logic -│ └── k8s.go # Kubernetes via Tai TCP proxy -├── proxy/ # HTTP reverse proxy URL resolution -└── vnc/ # VNC WebSocket URL resolution -``` - -## Quick Start - -```go -import "github.com/yaoapp/yao/tai" - -// Local — default Docker socket -c, _ := tai.New("") - -// Local — explicit address -c, _ := tai.New("docker:///var/run/docker.sock") -c, _ := tai.New("docker://192.168.1.50:2375") - -// Remote — via Tai Server (Docker runtime, default) -c, _ := tai.New("tai://192.168.1.100") - -// Remote — via Tai Server (K8s runtime) -c, _ := tai.New("tai://10.0.0.5", tai.K8s, - tai.WithKubeConfig("/path/to/kubeconfig.yml"), - tai.WithNamespace("sandbox"), -) - -defer c.Close() - -// Container lifecycle -id, _ := c.Sandbox().Create(ctx, sandbox.CreateOptions{ - Image: "node:20", - Cmd: []string{"sleep", "infinity"}, -}) -c.Sandbox().Start(ctx, id) - -// Filesystem -ws := c.Workspace("session-1") -ws.WriteFile("app.js", []byte("console.log('hi')"), 0644) -data, _ := ws.ReadFile("app.js") - -// HTTP proxy URL -url, _ := c.Proxy().URL(ctx, id, 3000, "/api/health") - -// VNC URL -vncURL, _ := c.VNC().URL(ctx, id) -``` - -## Address Protocol - -| Prefix | Mode | Description | -|--------|------|-------------| -| `""` | Local | Platform default Docker socket | -| `docker://...` | Local | Direct Docker daemon (socket or TCP) | -| `tai://host` | Remote | Via Tai Server, all services proxied | - -## Sub-Package Interfaces - -### volume.Volume - -File IO and directory sync between Yao and the container workspace. - -- `ReadFile`, `WriteFile`, `Stat`, `ListDir`, `Remove`, `Rename`, `MkdirAll` -- `SyncPush` (Yao -> Tai), `SyncPull` (Tai -> Yao) -- **Remote**: gRPC to Tai `:9100` -- **Local**: direct disk IO under `dataDir/{sessionID}/` - -### workspace.FS - -Go `fs.FS`-compatible interface wrapping `volume.Volume`, adding write operations. - -### sandbox.Sandbox - -Container lifecycle: `Create`, `Start`, `Stop`, `Remove`, `Exec`, `Inspect`, `List`. - -- **Local**: direct Docker socket, handles VNC port mapping and capabilities -- **Docker**: via Tai `:2375` (Docker Engine API proxy) -- **K8s**: via Tai `:6443` (kube-apiserver TCP proxy, single-container Pod per sandbox) - -### proxy.Proxy - -HTTP service URL resolution: `URL(ctx, containerID, port, path)`. - -- **Remote**: `http://tai-host:8080/{id}:{port}/{path}` -- **Local**: `http://127.0.0.1:{hostPort}/{path}` via `sandbox.Inspect` - -### vnc.VNC - -VNC WebSocket URL resolution: `URL(ctx, containerID)`. - -- **Remote**: `ws://tai-host:6080/vnc/{id}/ws` -- **Local**: `ws://127.0.0.1:{vncHostPort}/ws` via `sandbox.Inspect` - -## Options - -```go -tai.Docker // Docker runtime (default, can omit) -tai.K8s // Kubernetes runtime -tai.WithPorts(Ports{}) // custom port mapping -tai.WithHTTPClient(hc) // custom HTTP client -tai.WithDataDir(dir) // workspace root (Local mode) -tai.WithKubeConfig(path) // kubeconfig file path (K8s runtime) -tai.WithNamespace(ns) // namespace for K8s (default "default") -``` - -## Default Ports - -| Service | Default Port | -|---------|-------------| -| gRPC (Volume + Gateway) | 9100 | -| HTTP Proxy | 8080 | -| VNC Router | 6080 | -| Docker API Proxy | 2375 | -| K8s API Proxy | 6443 | - -## Dependencies - -- `github.com/yaoapp/tai/volume/pb` — gRPC proto types -- `google.golang.org/grpc` -- `github.com/pierrec/lz4/v4` — sync compression -- `github.com/docker/docker` — Docker SDK -- `k8s.io/client-go` + `k8s.io/api` + `k8s.io/apimachinery` — Kubernetes SDK diff --git a/tai/docs/README.md b/tai/docs/README.md new file mode 100644 index 00000000..47848611 --- /dev/null +++ b/tai/docs/README.md @@ -0,0 +1,111 @@ +# Tai SDK + +Go client library for the [Tai](https://github.com/YaoApp/tai) runtime bridge. Provides unified access to container sandboxes, volume IO, HTTP proxy, and VNC routing — transparently working in both **Local** (direct Docker) and **Remote** (via Tai server) modes. + +## Package Layout + +| Package | Import Path | Description | +|---------|-------------|-------------| +| `tai` | `github.com/yaoapp/yao/tai` | Top-level client, `New()`, options, `Close()` | +| `sandbox` | `github.com/yaoapp/yao/tai/sandbox` | Container lifecycle (Create/Start/Stop/Exec/Remove) | +| `volume` | `github.com/yaoapp/yao/tai/volume` | File IO and directory sync | +| `workspace` | `github.com/yaoapp/yao/tai/workspace` | `fs.FS`-compatible filesystem over Volume | +| `proxy` | `github.com/yaoapp/yao/tai/proxy` | HTTP reverse proxy URL resolution | +| `vnc` | `github.com/yaoapp/yao/tai/vnc` | VNC WebSocket URL resolution | + +## Quick Start + +### Local Mode (direct Docker) + +```go +c, err := tai.New("") +// or: tai.New("unix:///var/run/docker.sock") +// or: tai.New("tcp://192.168.1.50:2375") +defer c.Close() + +id, _ := c.Sandbox().Create(ctx, sandbox.CreateOptions{ + Name: "my-sandbox", + Image: "alpine:latest", + Cmd: []string{"sleep", "300"}, +}) +c.Sandbox().Start(ctx, id) +``` + +### Remote Mode (via Tai server, Docker runtime) + +```go +c, err := tai.New("tai://192.168.1.100") +defer c.Close() + +result, _ := c.Sandbox().Exec(ctx, id, []string{"echo", "hello"}, sandbox.ExecOptions{}) +fmt.Println(result.Stdout) // "hello\n" +``` + +### Remote Mode (via Tai server, K8s runtime) + +```go +c, err := tai.New("tai://192.168.1.100", tai.K8s, + tai.WithKubeConfig("/path/to/kubeconfig.yml"), + tai.WithNamespace("default"), + tai.WithPorts(tai.Ports{K8s: 6443}), +) +defer c.Close() +``` + +## Address Protocols + +| Address | Mode | Description | +|---------|------|-------------| +| `""` | Local | Platform default Docker socket | +| `unix:///var/run/docker.sock` | Local | Explicit Unix socket | +| `tcp://host:port` | Local | Explicit TCP Docker daemon | +| `npipe:////./pipe/docker_engine` | Local | Windows named pipe | +| `docker://host:port` | Local | Docker scheme | +| `tai://host` | Remote | Connect via Tai server | + +## Options + +| Option | Description | Default | +|--------|-------------|---------| +| `WithPorts(Ports{...})` | Override Tai service ports | gRPC=9100, HTTP=8080, VNC=6080 | +| `WithHTTPClient(*http.Client)` | Custom HTTP client for proxy/VNC | `http.DefaultClient` | +| `WithDataDir(path)` | Volume storage root (Local mode) | `/tmp/tai-volumes` | +| `WithKubeConfig(path)` | Kubeconfig file path (K8s mode, **required**) | - | +| `WithNamespace(ns)` | K8s namespace | `"default"` | + +## Default Ports + +| Service | Port | Description | +|---------|------|-------------| +| gRPC | 9100 | Volume IO + Gateway | +| HTTP | 8080 | HTTP reverse proxy | +| VNC | 6080 | VNC WebSocket router | +| Docker | 2375 | Docker API proxy | +| K8s | 6443 | Kubernetes API proxy | + +## Client API + +```go +c.Volume() // volume.Volume +c.Workspace(sessionID) // workspace.FS +c.Sandbox() // sandbox.Sandbox +c.Proxy() // proxy.Proxy +c.VNC() // vnc.VNC +c.IsLocal() // bool +c.Close() // error +``` + +## Runtime Constants + +```go +tai.Docker // default — use Docker runtime via Tai +tai.K8s // use Kubernetes runtime via Tai +``` + +## Sub-Package Documentation + +- [sandbox.md](sandbox.md) — Container lifecycle management +- [volume.md](volume.md) — File IO and sync +- [workspace.md](workspace.md) — fs.FS-compatible filesystem +- [proxy.md](proxy.md) — HTTP reverse proxy +- [vnc.md](vnc.md) — VNC WebSocket routing diff --git a/tai/docs/proxy.md b/tai/docs/proxy.md new file mode 100644 index 00000000..97411a5c --- /dev/null +++ b/tai/docs/proxy.md @@ -0,0 +1,86 @@ +# Package `proxy` + +HTTP reverse proxy URL resolution. Resolves service URLs for containers so that HTTP services running inside sandboxes can be accessed from the host. + +## Interface + +```go +type Proxy interface { + URL(ctx context.Context, containerID string, port int, path string) (string, error) + Healthz(ctx context.Context) error +} +``` + +## Implementations + +| Implementation | Constructor | Mode | URL Pattern | +|----------------|-------------|------|-------------| +| **Remote** | `NewRemote(host, port, hc)` | Via Tai HTTP proxy | `http://tai-host:8080/{containerID}:{port}/{path}` | +| **Local** | `NewLocal(sb)` | Direct host port lookup | `http://127.0.0.1:{hostPort}/{path}` | + +## Constructors + +### NewRemote + +```go +func NewRemote(host string, port int, hc *http.Client) Proxy +``` + +Creates a Proxy that routes through Tai's HTTP reverse proxy. URLs are constructed by combining the Tai server address with the container ID and port. + +- `host` — Tai server hostname/IP +- `port` — Tai HTTP proxy port (default 8080) +- `hc` — custom HTTP client, `nil` uses `http.DefaultClient` + +### NewLocal + +```go +func NewLocal(sb sandbox.Sandbox) Proxy +``` + +Creates a Proxy that resolves URLs by inspecting the container's port mappings via `sandbox.Inspect`. Looks up the host port bound to the requested container port. + +Returns an error if the requested port is not mapped. + +## Methods + +### URL + +```go +URL(ctx context.Context, containerID string, port int, path string) (string, error) +``` + +Resolves an HTTP URL to reach a service running on `port` inside the given container. + +**Remote example:** container `abc123` port `3000` path `/api/health` +→ `http://tai-host:8080/abc123:3000/api/health` + +**Local example:** container `abc123` port `3000` mapped to host port `32768` +→ `http://127.0.0.1:32768/api/health` + +### Healthz + +```go +Healthz(ctx context.Context) error +``` + +Checks the health of the proxy backend. + +- **Remote**: sends `GET /healthz` to the Tai HTTP proxy server +- **Local**: always returns `nil` (no external dependency) + +## Example + +```go +c, _ := tai.New("tai://192.168.1.100") +defer c.Close() + +// Get URL for a web service running on port 3000 +url, _ := c.Proxy().URL(ctx, containerID, 3000, "/api/status") +resp, _ := http.Get(url) + +// Health check +if err := c.Proxy().Healthz(ctx); err != nil { + log.Fatal("Tai HTTP proxy is down:", err) +} +``` diff --git a/tai/docs/sandbox.md b/tai/docs/sandbox.md new file mode 100644 index 00000000..fd82e0d0 --- /dev/null +++ b/tai/docs/sandbox.md @@ -0,0 +1,182 @@ +# Package `sandbox` + +Container lifecycle management. Provides a unified `Sandbox` interface with three implementations: + +| Implementation | Constructor | Backend | Mode | +|----------------|-------------|---------|------| +| **Local** | `NewLocal(addr)` | Direct Docker daemon | Local | +| **Docker** | `NewDocker(addr)` | Docker via Tai proxy | Remote | +| **K8s** | `NewK8s(addr, opts)` | Kubernetes via Tai proxy | Remote | + +## Interface + +```go +type Sandbox interface { + Create(ctx context.Context, opts CreateOptions) (id string, err error) + Start(ctx context.Context, id string) error + Stop(ctx context.Context, id string, timeout time.Duration) error + Remove(ctx context.Context, id string, force bool) error + Exec(ctx context.Context, id string, cmd []string, opts ExecOptions) (*ExecResult, error) + Inspect(ctx context.Context, id string) (*ContainerInfo, error) + List(ctx context.Context, opts ListOptions) ([]ContainerInfo, error) + Close() error +} +``` + +## Constructors + +### NewLocal + +```go +func NewLocal(addr string) (Sandbox, error) +``` + +Connects directly to a Docker daemon. `addr` can be: +- `""` — platform default (Unix socket on Linux/macOS, named pipe on Windows) +- `"unix:///var/run/docker.sock"` — explicit Unix socket +- `"tcp://host:port"` — explicit TCP + +Pings the daemon on creation; returns an error if unreachable. + +### NewDocker + +```go +func NewDocker(addr string) (Sandbox, error) +``` + +Connects to Docker Engine API through Tai's Docker proxy. `addr` should be `"tcp://tai-host:2375"`. + +### NewK8s + +```go +func NewK8s(addr string, opts ...K8sOption) (Sandbox, error) +``` + +Connects to Kubernetes through Tai's TCP proxy. Each sandbox maps to a single-container Pod. + +**Parameters:** +- `addr` — `"host:port"` pointing to Tai's K8s proxy endpoint +- `opts.KubeConfig` — path to kubeconfig file (**required**). Relative paths are resolved to absolute. +- `opts.Namespace` — Kubernetes namespace (default `"default"`) + +The constructor overrides the kubeconfig's `server` field to point at `addr`, enables insecure TLS (since Tai does TCP passthrough), and verifies connectivity by querying the namespace. + +All pods created by K8s sandbox are labeled with `managed-by: yao-tai-sdk`. + +## Types + +### CreateOptions + +```go +type CreateOptions struct { + Name string // container/pod name + Image string // container image + Cmd []string // entrypoint command + Env map[string]string // environment variables + Binds []string // volume binds (Docker only) + WorkingDir string // working directory + Memory int64 // memory limit in bytes, 0 = no limit + CPUs float64 // CPU limit, 0 = no limit + VNC bool // enable VNC port mapping (Local only) + Ports []PortMapping // port mappings (Docker only) +} +``` + +### PortMapping + +```go +type PortMapping struct { + ContainerPort int // port inside the container + HostPort int // port on the host, 0 = random + HostIP string // host bind address, default "127.0.0.1" + Protocol string // "tcp" (default) or "udp" +} +``` + +### ContainerInfo + +```go +type ContainerInfo struct { + ID string // container/pod ID + Name string // container/pod name + Image string // image name + Status string // "created", "running", "exited", "removing" (Docker) + // "Pending", "Running", "Succeeded", "Failed" (K8s) + IP string // container/pod IP address + Ports []PortMapping // mapped ports (Docker only) +} +``` + +### ExecOptions + +```go +type ExecOptions struct { + WorkDir string // override working directory + Env map[string]string // additional environment variables +} +``` + +### ExecResult + +```go +type ExecResult struct { + ExitCode int + Stdout string + Stderr string +} +``` + +### ListOptions + +```go +type ListOptions struct { + All bool // include stopped containers + Labels map[string]string // filter by labels +} +``` + +### K8sOption + +```go +type K8sOption struct { + Namespace string // default "default" + KubeConfig string // path to kubeconfig file (required) +} +``` + +## Behavioral Differences + +| Behavior | Docker (Local/Remote) | K8s | +|----------|----------------------|-----| +| `Create` returns | container ID (hash) | pod name | +| `Start` | starts a stopped container | polls until pod leaves Pending (up to 30s) | +| `Stop` | stops with timeout, container persists | deletes the pod with grace period | +| `Remove(force=true)` | force-removes | deletes with grace period 0 | +| `Exec` | Docker exec API | `kubectl exec` via SPDY | +| `Inspect.Ports` | populated from Docker | always empty | +| `List` | filters by `tai-sdk=true` label | filters by `managed-by=yao-tai-sdk` label | +| `Binds` | supported | not supported | +| `VNC` flag | auto port-maps 6080 on macOS/Windows | not applicable | + +## Example + +```go +sb, _ := sandbox.NewLocal("") +defer sb.Close() + +id, _ := sb.Create(ctx, sandbox.CreateOptions{ + Name: "worker", + Image: "alpine:latest", + Cmd: []string{"sleep", "300"}, + Env: map[string]string{"FOO": "bar"}, + Memory: 256 * 1024 * 1024, // 256 MB +}) + +sb.Start(ctx, id) + +result, _ := sb.Exec(ctx, id, []string{"echo", "$FOO"}, sandbox.ExecOptions{}) +fmt.Println(result.Stdout) + +sb.Stop(ctx, id, 10*time.Second) +sb.Remove(ctx, id, false) +``` diff --git a/tai/docs/vnc.md b/tai/docs/vnc.md new file mode 100644 index 00000000..e1d9e425 --- /dev/null +++ b/tai/docs/vnc.md @@ -0,0 +1,94 @@ +# Package `vnc` + +VNC WebSocket URL resolution. Resolves WebSocket URLs for VNC sessions running inside containers, enabling remote desktop access to sandbox environments. + +## Interface + +```go +type VNC interface { + URL(ctx context.Context, containerID string) (string, error) + Ping(ctx context.Context, containerID string) error +} +``` + +## Implementations + +| Implementation | Constructor | Mode | URL Pattern | +|----------------|-------------|------|-------------| +| **Remote** | `NewRemote(host, port, hc)` | Via Tai VNC router | `ws://tai-host:6080/vnc/{containerID}/ws` | +| **Local** | `NewLocal(sb)` | Direct host port lookup | `ws://127.0.0.1:{hostPort}/ws` | + +## Constructors + +### NewRemote + +```go +func NewRemote(host string, port int, hc *http.Client) VNC +``` + +Creates a VNC that routes through Tai's VNC WebSocket router. + +- `host` — Tai server hostname/IP +- `port` — Tai VNC router port (default 6080) +- `hc` — custom HTTP client for Ping, `nil` uses `http.DefaultClient` + +### NewLocal + +```go +func NewLocal(sb sandbox.Sandbox) VNC +``` + +Creates a VNC that resolves URLs by inspecting the container's port mappings. Looks for container port **6080** (the standard noVNC port) in the port mappings. + +Returns an error if port 6080 is not mapped. On macOS and Windows (Docker Desktop), the Local sandbox automatically maps port 6080 when `CreateOptions.VNC` is `true`. + +## Methods + +### URL + +```go +URL(ctx context.Context, containerID string) (string, error) +``` + +Returns a WebSocket URL for connecting to the container's VNC session. + +**Remote:** `ws://tai-host:6080/vnc/abc123/ws` +**Local:** `ws://127.0.0.1:32769/ws` + +### Ping + +```go +Ping(ctx context.Context, containerID string) error +``` + +Checks if the VNC endpoint is reachable by making an HTTP GET request to the WebSocket URL. Useful for verifying that the VNC server inside the container is ready before connecting a client. + +- **Remote**: sends GET to `http://tai-host:6080/vnc/{containerID}/ws` +- **Local**: resolves the host port via Inspect, then sends GET + +## Example + +```go +c, _ := tai.New("tai://192.168.1.100") +defer c.Close() + +// Create a sandbox with VNC enabled +id, _ := c.Sandbox().Create(ctx, sandbox.CreateOptions{ + Name: "desktop", + Image: "yaoapp/sandbox-claude:latest", + VNC: true, +}) +c.Sandbox().Start(ctx, id) + +// Wait for VNC to be ready +for i := 0; i < 10; i++ { + if err := c.VNC().Ping(ctx, id); err == nil { + break + } + time.Sleep(time.Second) +} + +// Get the WebSocket URL for a noVNC client +url, _ := c.VNC().URL(ctx, id) +fmt.Println(url) // ws://192.168.1.100:6080/vnc/desktop/ws +``` diff --git a/tai/docs/volume.md b/tai/docs/volume.md new file mode 100644 index 00000000..15b75e67 --- /dev/null +++ b/tai/docs/volume.md @@ -0,0 +1,120 @@ +# Package `volume` + +File IO and directory synchronization. Provides a `Volume` interface with two implementations: + +| Implementation | Constructor | Backend | Mode | +|----------------|-------------|---------|------| +| **Local** | `NewLocal(root)` | Direct filesystem | Local | +| **Remote** | `NewRemote(conn)` | gRPC to Tai :9100 | Remote | + +## Interface + +```go +type Volume interface { + ReadFile(ctx context.Context, sessionID, path string) (data []byte, perm os.FileMode, err error) + WriteFile(ctx context.Context, sessionID, path string, data []byte, perm os.FileMode) error + Stat(ctx context.Context, sessionID, path string) (*FileInfo, error) + ListDir(ctx context.Context, sessionID, path string) ([]FileInfo, error) + Remove(ctx context.Context, sessionID, path string, recursive bool) error + Rename(ctx context.Context, sessionID, oldPath, newPath string) error + MkdirAll(ctx context.Context, sessionID, path string) error + + SyncPush(ctx context.Context, sessionID, localDir string, opts ...SyncOption) (*SyncResult, error) + SyncPull(ctx context.Context, sessionID, localDir string, opts ...SyncOption) (*SyncResult, error) + + Close() error +} +``` + +All paths are **relative** to the session's workspace root. The `sessionID` identifies the workspace partition — in Local mode this maps to `//`, in Remote mode the Tai server manages the path. + +## Constructors + +### NewLocal + +```go +func NewLocal(root string) Volume +``` + +Creates a Volume backed by the local filesystem. Files are stored under `//`. + +### NewRemote + +```go +func NewRemote(conn *grpc.ClientConn) Volume +``` + +Creates a Volume backed by Tai's gRPC Volume service. The connection should target Tai's gRPC port (default 9100). Uses lz4 compression for `SyncPush`/`SyncPull` bulk transfers. + +## Types + +### FileInfo + +```go +type FileInfo struct { + Path string + Size int64 + Mtime time.Time + Mode fs.FileMode + IsDir bool +} +``` + +### SyncResult + +```go +type SyncResult struct { + FilesSynced int + BytesTransferred int64 + Duration time.Duration +} +``` + +## Sync Options + +```go +volume.WithForceFull() // skip snapshot cache, diff against actual disk +volume.WithExcludes("*.log", ".DS_Store") // glob patterns to exclude +``` + +## File Operations + +| Method | Description | +|--------|-------------| +| `ReadFile` | Read file contents and permissions | +| `WriteFile` | Write file with specified permissions (creates parent dirs) | +| `Stat` | Get file/directory metadata | +| `ListDir` | List directory contents (one level) | +| `Remove` | Delete file or directory (`recursive=true` for tree) | +| `Rename` | Move/rename a file or directory | +| `MkdirAll` | Create directory tree | + +## Sync Operations + +| Method | Direction | Description | +|--------|-----------|-------------| +| `SyncPush` | local → remote | Upload a local directory to the session workspace | +| `SyncPull` | remote → local | Download the session workspace to a local directory | + +Both sync methods use snapshot-based diffing to transfer only changed files. Use `WithForceFull()` to bypass the cache and force a full transfer. + +Remote sync uses **lz4 compression** on the wire, streaming files via gRPC bidirectional streaming. + +## Example + +```go +vol := volume.NewLocal("/data/volumes") +defer vol.Close() + +// Write a file +vol.WriteFile(ctx, "session-1", "main.py", []byte("print('hi')"), 0644) + +// Read it back +data, perm, _ := vol.ReadFile(ctx, "session-1", "main.py") + +// Sync a local directory to the session +result, _ := vol.SyncPush(ctx, "session-1", "/tmp/project", + volume.WithExcludes("node_modules", ".git"), +) +fmt.Printf("synced %d files (%d bytes)\n", result.FilesSynced, result.BytesTransferred) +``` diff --git a/tai/docs/workspace.md b/tai/docs/workspace.md new file mode 100644 index 00000000..61b84f10 --- /dev/null +++ b/tai/docs/workspace.md @@ -0,0 +1,93 @@ +# Package `workspace` + +Provides an `fs.FS`-compatible filesystem abstraction over `volume.Volume`. This allows session workspaces to be used with any Go standard library function that accepts `fs.FS`, such as `fs.WalkDir`, `template.ParseFS`, or `http.FS`. + +## Interface + +```go +type FS interface { + fs.FS // Open(name) (fs.File, error) + fs.StatFS // Stat(name) (fs.FileInfo, error) + fs.ReadFileFS // ReadFile(name) ([]byte, error) + fs.ReadDirFS // ReadDir(name) ([]fs.DirEntry, error) + io.Closer // Close() error + + WriteFile(name string, data []byte, perm os.FileMode) error + Remove(name string) error + RemoveAll(name string) error + Rename(oldname, newname string) error + MkdirAll(name string, perm os.FileMode) error +} +``` + +## Constructor + +```go +func New(vol volume.Volume, sessionID string) FS +``` + +Creates an FS backed by the given Volume for the specified session. The returned FS works transparently whether `vol` is Local or Remote. + +Typically accessed through the top-level client: + +```go +c, _ := tai.New("tai://host") +ws := c.Workspace("session-123") +``` + +## Read Operations (fs.FS compatible) + +All read operations comply with the `fs.FS` contract. Paths must be valid according to `fs.ValidPath` — forward slashes, no leading slash, no `..` segments. + +| Method | Standard Interface | Description | +|--------|--------------------|-------------| +| `Open(name)` | `fs.FS` | Opens a file or directory | +| `Stat(name)` | `fs.StatFS` | Returns file metadata | +| `ReadFile(name)` | `fs.ReadFileFS` | Reads entire file contents | +| `ReadDir(name)` | `fs.ReadDirFS` | Lists directory entries | + +`Open` returns an in-memory `fs.File` for regular files (entire content loaded on open) and a directory handle for directories. + +## Write Operations + +| Method | Description | +|--------|-------------| +| `WriteFile(name, data, perm)` | Write file contents with permissions | +| `Remove(name)` | Delete a single file or empty directory | +| `RemoveAll(name)` | Delete a file or directory tree recursively | +| `Rename(old, new)` | Move/rename a file or directory | +| `MkdirAll(name, perm)` | Create directory tree (perm currently unused) | + +## Example + +```go +c, _ := tai.New("tai://192.168.1.100") +defer c.Close() + +ws := c.Workspace("project-abc") + +// Write files +ws.WriteFile("src/main.go", []byte("package main"), 0644) +ws.MkdirAll("src/utils", 0755) + +// Read with standard fs.FS +data, _ := fs.ReadFile(ws, "src/main.go") + +// Walk the tree +fs.WalkDir(ws, ".", func(path string, d fs.DirEntry, err error) error { + fmt.Println(path) + return nil +}) + +// Use with Go templates +tmpl, _ := template.ParseFS(ws, "templates/*.html") + +// Clean up +ws.RemoveAll("src") +``` + +## Implementation Notes + +- `Open` on a regular file reads the entire content into memory. For large files, prefer `ReadFile` or Volume's `ReadFile` directly. +- `Close()` is a no-op — the underlying Volume's lifecycle is managed by the `tai.Client`. +- Path validation follows `fs.ValidPath` rules. Invalid paths return `fs.ErrInvalid`. From 6e68efaba3bd6fe757fa0a1618c7478704ce79ed Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 4 Mar 2026 13:17:48 +0800 Subject: [PATCH 2/7] Implement gRPC support in the Yao SDK - Add gRPC server configuration to the application, allowing for gRPC communication. - Introduce new Makefile targets for gRPC unit testing and proto code generation. - Update CI workflows to include gRPC tests with SQLite as the transport layer. - Refactor the sandbox design to support multi-node capabilities and improve isolation. - Enhance the service layer to facilitate internal request forwarding for gRPC APIs. This commit lays the groundwork for integrating gRPC into the Yao SDK, improving performance and scalability. --- .github/workflows/pr-test.yml | 169 ++++ .github/workflows/unit-test.yml | 107 ++ Makefile | 43 +- agent/context/grpc.go | 153 +++ agent/llm/jsapi.go | 6 + cmd/start.go | 15 + config/types.go | 56 +- grpc/DESIGN.md | 448 ++++++++ grpc/IMPL.md | 274 +++++ grpc/TEST.md | 426 ++++++++ grpc/agent/agent.go | 93 ++ grpc/agent/agent_test.go | 206 ++++ grpc/api/api.go | 69 ++ grpc/api/api_test.go | 135 +++ grpc/auth/endpoint.go | 66 ++ grpc/auth/endpoint_test.go | 136 +++ grpc/auth/guard.go | 169 ++++ grpc/auth/guard_test.go | 169 ++++ grpc/auth/scope.go | 14 + grpc/grpc.go | 173 ++++ grpc/health/health.go | 15 + grpc/health/health_test.go | 22 + grpc/llm/llm.go | 165 +++ grpc/llm/llm_test.go | 265 +++++ grpc/mcp/mcp.go | 102 ++ grpc/mcp/mcp_test.go | 264 +++++ grpc/pb/yao.pb.go | 1316 ++++++++++++++++++++++++ grpc/pb/yao.proto | 149 +++ grpc/pb/yao_grpc.pb.go | 606 +++++++++++ grpc/run/run.go | 79 ++ grpc/run/run_test.go | 155 +++ grpc/shell/shell.go | 91 ++ grpc/shell/shell_test.go | 166 +++ grpc/tests/testutils/testutils.go | 220 ++++ openapi/oauth/authenticate.go | 207 ++++ sandbox/DESIGN.md | 1570 +++++------------------------ sandbox/SPEC.md | 579 +++++++++++ service/service.go | 6 + 38 files changed, 7556 insertions(+), 1348 deletions(-) create mode 100644 agent/context/grpc.go create mode 100644 grpc/DESIGN.md create mode 100644 grpc/IMPL.md create mode 100644 grpc/TEST.md create mode 100644 grpc/agent/agent.go create mode 100644 grpc/agent/agent_test.go create mode 100644 grpc/api/api.go create mode 100644 grpc/api/api_test.go create mode 100644 grpc/auth/endpoint.go create mode 100644 grpc/auth/endpoint_test.go create mode 100644 grpc/auth/guard.go create mode 100644 grpc/auth/guard_test.go create mode 100644 grpc/auth/scope.go create mode 100644 grpc/grpc.go create mode 100644 grpc/health/health.go create mode 100644 grpc/health/health_test.go create mode 100644 grpc/llm/llm.go create mode 100644 grpc/llm/llm_test.go create mode 100644 grpc/mcp/mcp.go create mode 100644 grpc/mcp/mcp_test.go create mode 100644 grpc/pb/yao.pb.go create mode 100644 grpc/pb/yao.proto create mode 100644 grpc/pb/yao_grpc.pb.go create mode 100644 grpc/run/run.go create mode 100644 grpc/run/run_test.go create mode 100644 grpc/shell/shell.go create mode 100644 grpc/shell/shell_test.go create mode 100644 grpc/tests/testutils/testutils.go create mode 100644 openapi/oauth/authenticate.go create mode 100644 sandbox/SPEC.md diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index c9e07eda..c3c8ab5f 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -1722,3 +1722,172 @@ jobs: issue_number: issue_number, body: '✅ Tai SDK Tests passed!' }); + + # ============================================================================= + # gRPC Tests - Run once with SQLite (transport layer, no DB matrix needed) + # ============================================================================= + GRPCTest: + runs-on: ubuntu-latest + services: + mongodb: + image: mongo:6.0 + ports: + - 27017:27017 + env: + MONGO_INITDB_ROOT_USERNAME: root + MONGO_INITDB_ROOT_PASSWORD: 123456 + MONGO_INITDB_DATABASE: test + + strategy: + matrix: + go: ["1.25"] + if: > + ${{ github.event.workflow_run.event == 'pull_request' && + github.event.workflow_run.conclusion == 'success' }} + steps: + - name: "Download artifact" + uses: actions/github-script@v7 + with: + script: | + var artifacts = await github.rest.actions.listWorkflowRunArtifacts({ + owner: context.repo.owner, + repo: context.repo.repo, + run_id: ${{github.event.workflow_run.id }}, + }); + var matchArtifact = artifacts.data.artifacts.filter((artifact) => { + return artifact.name == "pr" + })[0]; + var download = await github.rest.actions.downloadArtifact({ + owner: context.repo.owner, + repo: context.repo.repo, + artifact_id: matchArtifact.id, + archive_format: 'zip', + }); + var fs = require('fs'); + fs.writeFileSync('${{github.workspace}}/pr.zip', Buffer.from(download.data)); + + - name: "Read NR & SHA" + run: | + unzip pr.zip + cat NR + cat SHA + echo HEAD=$(cat SHA) >> $GITHUB_ENV + echo NR=$(cat NR) >> $GITHUB_ENV + + - name: "Comment on PR" + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const { NR } = process.env + var issue_number = NR; + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: issue_number, + body: '🤖 gRPC Tests running with SQLite...' + }); + + - name: Checkout Kun + uses: actions/checkout@v4 + with: + repository: yaoapp/kun + path: kun + + - name: Checkout Xun + uses: actions/checkout@v4 + with: + repository: yaoapp/xun + path: xun + + - name: Checkout Gou + uses: actions/checkout@v4 + with: + repository: yaoapp/gou + path: gou + + - name: Checkout V8Go + uses: actions/checkout@v4 + with: + repository: yaoapp/v8go + path: v8go + + - name: Unzip libv8 + run: | + files=$(find ./v8go -name "libv8*.zip") + for file in $files; do + dir=$(dirname "$file") + echo "Extracting $file to directory $dir" + unzip -o -d $dir $file + rm -rf $dir/__MACOSX + done + + - name: Checkout Demo App + uses: actions/checkout@v4 + with: + repository: yaoapp/yao-dev-app + path: app + + - name: Checkout Extension + uses: actions/checkout@v4 + with: + repository: yaoapp/yao-extensions-dev + path: extension + + - name: Move Dependencies + run: | + mv kun ../ + mv xun ../ + mv gou ../ + mv v8go ../ + mv app ../ + mv extension ../ + + - name: Checkout pull request HEAD commit + uses: actions/checkout@v4 + with: + ref: ${{ env.HEAD }} + + - name: Setup Apple Private Key + run: | + mkdir -p ../app/openapi/certs/apple + echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8 + + - name: Setup Go ${{ matrix.go }} + uses: actions/setup-go@v5 + with: + go-version: ${{ matrix.go }} + + - name: Start Redis + run: docker run --name redis --publish 6379:6379 --detach redis:6 + + - name: Setup Go Tools + run: make tools + + - name: Setup ENV (SQLite) + run: | + mkdir -p ${{ github.WORKSPACE }}/../app/db + echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV + echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV + + - name: Run gRPC Tests + run: make unit-test-grpc + + - name: Codecov Report + uses: codecov/codecov-action@v4 + with: + token: ${{ secrets.CODECOV_TOKEN }} + + - name: "Comment on PR - gRPC Tests Done" + uses: actions/github-script@v7 + with: + github-token: ${{ secrets.GITHUB_TOKEN }} + script: | + const { NR } = process.env + var issue_number = NR; + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: issue_number, + body: '✅ gRPC Tests passed!' + }); diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index 7d3696e2..1c5b36a4 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -1262,3 +1262,110 @@ jobs: uses: codecov/codecov-action@v4 with: token: ${{ secrets.CODECOV_TOKEN }} + + # ============================================================================= + # gRPC Tests - Run once with SQLite (transport layer, no DB matrix needed) + # ============================================================================= + grpc-test: + runs-on: ubuntu-latest + services: + mongodb: + image: mongo:6.0 + ports: + - 27017:27017 + env: + MONGO_INITDB_ROOT_USERNAME: root + MONGO_INITDB_ROOT_PASSWORD: 123456 + MONGO_INITDB_DATABASE: test + + strategy: + matrix: + go: ["1.25"] + steps: + - name: Checkout Kun + uses: actions/checkout@v4 + with: + repository: ${{ env.REPO_KUN }} + path: kun + + - name: Checkout Xun + uses: actions/checkout@v4 + with: + repository: ${{ env.REPO_XUN }} + path: xun + + - name: Checkout Gou + uses: actions/checkout@v4 + with: + repository: ${{ env.REPO_GOU }} + path: gou + + - name: Checkout V8Go + uses: actions/checkout@v4 + with: + repository: yaoapp/v8go + path: v8go + + - name: Unzip libv8 + run: | + files=$(find ./v8go -name "libv8*.zip") + for file in $files; do + dir=$(dirname "$file") + echo "Extracting $file to directory $dir" + unzip -o -d $dir $file + rm -rf $dir/__MACOSX + done + + - name: Checkout Demo App + uses: actions/checkout@v4 + with: + repository: yaoapp/yao-dev-app + path: app + + - name: Checkout Extension + uses: actions/checkout@v4 + with: + repository: yaoapp/yao-extensions-dev + path: extension + + - name: Move Dependencies + run: | + mv kun ../ + mv xun ../ + mv gou ../ + mv v8go ../ + mv app ../ + mv extension ../ + + - name: Checkout Code + uses: actions/checkout@v4 + + - name: Setup Apple Private Key + run: | + mkdir -p ../app/openapi/certs/apple + echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8 + + - name: Setup Go ${{ matrix.go }} + uses: actions/setup-go@v5 + with: + go-version: ${{ matrix.go }} + + - name: Start Redis + run: docker run --name redis --publish 6379:6379 --detach redis:6 + + - name: Setup Go Tools + run: make tools + + - name: Setup ENV (SQLite) + run: | + mkdir -p ${{ github.WORKSPACE }}/../app/db + echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV + echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV + + - name: Run gRPC Tests + run: make unit-test-grpc + + - name: Codecov Report + uses: codecov/codecov-action@v4 + with: + token: ${{ secrets.CODECOV_TOKEN }} diff --git a/Makefile b/Makefile index 6c704606..68790267 100644 --- a/Makefile +++ b/Makefile @@ -11,8 +11,8 @@ OS := $(shell uname) # ROOT_DIR := $(shell dirname $(realpath $(firstword $(MAKEFILE_LIST)))) TESTFOLDER := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|registry' | awk '!/\/tests\// || /openapi\/tests/') -# Core tests (exclude AI-related: agent, aigc, openai, KB, sandbox, registry, and integrations which require external services) -TESTFOLDER_CORE := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|agent|kb|sandbox|integrations|registry|tai' | awk '!/\/tests\// || /openapi\/tests/') +# Core tests (exclude AI-related: agent, aigc, openai, KB, sandbox, registry, grpc, and integrations which require external services) +TESTFOLDER_CORE := $(shell $(GO) list ./... | grep -vE 'examples|openai|aigc|neo|twilio|share*|agent|kb|sandbox|integrations|registry|tai|grpc' | awk '!/\/tests\// || /openapi\/tests/') # Agent tests (agent, aigc) - exclude agent/search/handlers/web (requires external API keys) and robot packages (tested in robot job) TESTFOLDER_AGENT := $(shell $(GO) list ./agent/... ./aigc/... | grep -vE 'agent/search/handlers/web|agent/robot/') # KB tests (kb) @@ -23,6 +23,8 @@ TESTFOLDER_ROBOT := $(shell $(GO) list ./agent/robot/... | grep -vE 'agent/robot TESTFOLDER_SANDBOX := $(shell $(GO) list ./sandbox/...) # Tai SDK tests (requires Tai container with Docker socket) TESTFOLDER_TAI := $(shell $(GO) list ./tai/...) +# gRPC tests +TESTFOLDER_GRPC := $(shell $(GO) list ./grpc/...) TESTTAGS ?= "" # TESTWIDGETS := $(shell $(GO) list ./widgets/...) @@ -285,6 +287,43 @@ unit-test-tai: @echo "All Tai SDK tests passed" @echo "=============================================" +# Proto codegen +.PHONY: proto +proto: + protoc --go_out=. --go_opt=paths=source_relative \ + --go-grpc_out=. --go-grpc_opt=paths=source_relative \ + grpc/pb/yao.proto + +# gRPC Unit Test +.PHONY: unit-test-grpc +unit-test-grpc: + echo "mode: count" > coverage.out + for d in $(TESTFOLDER_GRPC); do \ + $(GO) test -tags $(TESTTAGS) -v -timeout=10m \ + -covermode=count -coverprofile=profile.out \ + -coverpkg=$$(echo $$d | sed "s/\/test$$//g") \ + -skip='TestMemoryLeak|TestIsolateDisposal|TestLeak_|TestScenario_' \ + $$d > tmp.out; \ + cat tmp.out; \ + if grep -q "^--- FAIL" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "build failed" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "setup failed" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "runtime error" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + fi; \ + if [ -f profile.out ]; then \ + cat profile.out | grep -v "mode:" >> coverage.out; \ + rm profile.out; \ + fi; \ + done + # Benchmark Test .PHONY: benchmark benchmark: diff --git a/agent/context/grpc.go b/agent/context/grpc.go new file mode 100644 index 00000000..f4f2e8f1 --- /dev/null +++ b/agent/context/grpc.go @@ -0,0 +1,153 @@ +package context + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + + "github.com/yaoapp/gou/connector" + "github.com/yaoapp/gou/store" + "github.com/yaoapp/yao/openapi/oauth/types" +) + +// GRPCAgentInput holds the raw inputs from a gRPC AgentStream request. +type GRPCAgentInput struct { + AssistantID string + Messages []byte + Options []byte + AuthInfo *types.AuthorizedInfo + Cache store.Store + Writer http.ResponseWriter +} + +// GetGRPCAgentRequest parses a gRPC agent request and creates a Context + Options, +// mirroring openapi.go GetCompletionRequest. +// +// Flow: validate → parse messages → parse options → build Context → build Options → register interrupt +func GetGRPCAgentRequest(parent context.Context, input GRPCAgentInput) ([]Message, *Context, *Options, error) { + if input.AssistantID == "" { + return nil, nil, nil, fmt.Errorf("assistant_id is required") + } + + messages, err := parseGRPCMessages(input.Messages) + if err != nil { + return nil, nil, nil, err + } + + var rawOpts map[string]interface{} + if len(input.Options) > 0 { + if err := json.Unmarshal(input.Options, &rawOpts); err != nil { + return nil, nil, nil, fmt.Errorf("invalid options JSON: %w", err) + } + } + + chatID := getChatIDFromOpts(rawOpts) + ctx := New(parent, input.AuthInfo, chatID) + + ctx.Cache = input.Cache + ctx.Writer = input.Writer + ctx.AssistantID = input.AssistantID + ctx.Locale = getStringOpt(rawOpts, "locale") + ctx.Theme = getStringOpt(rawOpts, "theme") + ctx.Referer = getRefererOpt(rawOpts) + ctx.Accept = getAcceptOpt(rawOpts) + ctx.Route = getStringOpt(rawOpts, "route") + ctx.Metadata = getMapOpt(rawOpts, "metadata") + ctx.Client = Client{Type: "grpc"} + + opts := &Options{ + Context: parent, + Skip: getSkipOpt(rawOpts), + Mode: getStringOpt(rawOpts, "mode"), + } + + if connectorID := getStringOpt(rawOpts, "connector"); connectorID != "" { + if _, err := connector.Select(connectorID); err == nil { + opts.Connector = connectorID + } + } + + ctx.Interrupt = NewInterruptController() + if err := Register(ctx); err != nil { + return nil, nil, nil, fmt.Errorf("failed to register context: %w", err) + } + ctx.Interrupt.Start(ctx.ID) + + return messages, ctx, opts, nil +} + +func parseGRPCMessages(raw []byte) ([]Message, error) { + if len(raw) == 0 { + return nil, fmt.Errorf("messages are required") + } + var messages []Message + if err := json.Unmarshal(raw, &messages); err != nil { + return nil, fmt.Errorf("invalid messages JSON: %w", err) + } + if len(messages) == 0 { + return nil, fmt.Errorf("messages must not be empty") + } + return messages, nil +} + +func getChatIDFromOpts(opts map[string]interface{}) string { + if opts != nil { + if v, ok := opts["chat_id"].(string); ok && v != "" { + return v + } + } + return GenChatID() +} + +func getStringOpt(opts map[string]interface{}, key string) string { + if opts == nil { + return "" + } + v, _ := opts[key].(string) + return v +} + +func getRefererOpt(opts map[string]interface{}) string { + r := getStringOpt(opts, "referer") + if r != "" { + return validateReferer(r) + } + return RefererAPI +} + +func getAcceptOpt(opts map[string]interface{}) Accept { + a := getStringOpt(opts, "accept") + if a != "" { + return validateAccept(a) + } + return AcceptStandard +} + +func getMapOpt(opts map[string]interface{}, key string) map[string]interface{} { + if opts == nil { + return nil + } + v, _ := opts[key].(map[string]interface{}) + return v +} + +func getSkipOpt(opts map[string]interface{}) *Skip { + if opts == nil { + return nil + } + raw, ok := opts["skip"] + if !ok { + return nil + } + + data, err := json.Marshal(raw) + if err != nil { + return nil + } + var skip Skip + if err := json.Unmarshal(data, &skip); err != nil { + return nil + } + return &skip +} diff --git a/agent/llm/jsapi.go b/agent/llm/jsapi.go index 4acbcddd..2439dd7f 100644 --- a/agent/llm/jsapi.go +++ b/agent/llm/jsapi.go @@ -201,6 +201,12 @@ func parseToolCalls(toolCalls []interface{}) []agentContext.ToolCall { } // buildCompletionOptions creates CompletionOptions from JS opts map +// BuildCompletionOptions builds CompletionOptions from a connector and raw opts map. +// Exported for reuse by gRPC handlers. +func BuildCompletionOptions(conn connector.Connector, opts map[string]interface{}) *agentContext.CompletionOptions { + return buildCompletionOptions(conn, opts) +} + func buildCompletionOptions(conn connector.Connector, opts map[string]interface{}) *agentContext.CompletionOptions { // Get capabilities from connector capabilities := GetCapabilitiesFromConn(conn) diff --git a/cmd/start.go b/cmd/start.go index d663505b..18d2e367 100644 --- a/cmd/start.go +++ b/cmd/start.go @@ -28,6 +28,8 @@ import ( agentcontext "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/engine" + yaogrpc "github.com/yaoapp/yao/grpc" + _ "github.com/yaoapp/yao/grpc/auth" "github.com/yaoapp/yao/openapi" ischedule "github.com/yaoapp/yao/schedule" "github.com/yaoapp/yao/service" @@ -190,6 +192,12 @@ var startCmd = &cobra.Command{ } } + // Print gRPC listen addresses + grpcAddrs := yaogrpc.Addr() + for _, addr := range grpcAddrs { + fmt.Println(color.WhiteString(L("Listening")), color.GreenString(" %s (gRPC)", addr)) + } + fmt.Println(color.WhiteString("\n---------------------------------")) fmt.Println(color.WhiteString(L("Access Points"))) fmt.Println(color.WhiteString("---------------------------------")) @@ -235,6 +243,13 @@ var startCmd = &cobra.Command{ os.Exit(1) } + // Start gRPC Server (after HTTP, LIFO shutdown: gRPC stops before HTTP) + if grpcErr := yaogrpc.StartServer(config.Conf); grpcErr != nil { + fmt.Println(color.RedString(L("gRPC: %s"), grpcErr.Error())) + os.Exit(1) + } + defer yaogrpc.Stop() + // Start watching watchDone := make(chan uint8, 1) if mode == "development" && !startDisableWatching { diff --git a/config/types.go b/config/types.go index 16e81bb5..8618fabc 100644 --- a/config/types.go +++ b/config/types.go @@ -2,30 +2,38 @@ package config // Config 象传应用引擎配置 type Config struct { - Mode string `json:"mode,omitempty" env:"YAO_ENV" envDefault:"production"` // The start mode production/development - AppSource string `json:"app,omitempty" env:"YAO_APP_SOURCE"` // The Application Source Root Path default same as Root - Root string `json:"root,omitempty" env:"YAO_ROOT" envDefault:"."` // The Application Root Path - Lang string `json:"lang,omitempty" env:"YAO_LANG" envDefault:"en-us"` // Default language setting - TimeZone string `json:"timezone,omitempty" env:"YAO_TIMEZONE"` // Default TimeZone - DataRoot string `json:"data_root,omitempty" env:"YAO_DATA_ROOT" envDefault:""` // The data root path - ExtensionRoot string `json:"extension_root,omitempty" env:"YAO_EXTENSION_ROOT" envDefault:""` // Plugin, Wasm root PATH, Default is (/plugins /wasms) - Host string `json:"host,omitempty" env:"YAO_HOST" envDefault:"0.0.0.0"` // The server host - Port int `json:"port,omitempty" env:"YAO_PORT" envDefault:"5099"` // The server port - Cert string `json:"cert,omitempty" env:"YAO_CERT"` // The HTTPS certificate path - Key string `json:"key,omitempty" env:"YAO_KEY"` // The HTTPS certificate key path - Log string `json:"log,omitempty" env:"YAO_LOG"` // The log file path - LogMode string `json:"log_mode,omitempty" env:"YAO_LOG_MODE" envDefault:"TEXT"` // The log mode TEXT|JSON - LogMaxSize int `json:"log_max_size,omitempty" env:"YAO_LOG_MAX_SIZE" envDefault:"100"` // The max log size in MB, the default is 100 - LogMaxAage int `json:"log_max_age,omitempty" env:"YAO_LOG_MAX_AGE" envDefault:"7"` // The max log age in day, the default is 7 - LogMaxBackups int `json:"log_max_backups" env:"YAO_LOG_MAX_BACKUPS" envDefault:"3"` // The max log backups, the default is 3 - LogLocalTime bool `json:"log_local_time" env:"YAO_LOG_LOCAL_TIME" envDefault:"true"` - JWTSecret string `json:"jwt_secret,omitempty" env:"YAO_JWT_SECRET"` // The JWT Secret - DB Database `json:"db,omitempty"` // The database config - AllowFrom []string `json:"allowfrom,omitempty" envSeparator:"|" env:"YAO_ALLOW_FROM"` // Domain list the separator is | - Session Session `json:"session,omitempty"` // Session Config - Runtime Runtime `json:"runtime,omitempty"` // Runtime config - Trace Trace `json:"trace,omitempty"` // Trace config - Registry string `json:"registry,omitempty" env:"YAO_REGISTRY" envDefault:"https://registry.yaoagents.com"` // The package registry server URL + Mode string `json:"mode,omitempty" env:"YAO_ENV" envDefault:"production"` // The start mode production/development + AppSource string `json:"app,omitempty" env:"YAO_APP_SOURCE"` // The Application Source Root Path default same as Root + Root string `json:"root,omitempty" env:"YAO_ROOT" envDefault:"."` // The Application Root Path + Lang string `json:"lang,omitempty" env:"YAO_LANG" envDefault:"en-us"` // Default language setting + TimeZone string `json:"timezone,omitempty" env:"YAO_TIMEZONE"` // Default TimeZone + DataRoot string `json:"data_root,omitempty" env:"YAO_DATA_ROOT" envDefault:""` // The data root path + ExtensionRoot string `json:"extension_root,omitempty" env:"YAO_EXTENSION_ROOT" envDefault:""` // Plugin, Wasm root PATH, Default is (/plugins /wasms) + Host string `json:"host,omitempty" env:"YAO_HOST" envDefault:"0.0.0.0"` // The server host + Port int `json:"port,omitempty" env:"YAO_PORT" envDefault:"5099"` // The server port + Cert string `json:"cert,omitempty" env:"YAO_CERT"` // The HTTPS certificate path + Key string `json:"key,omitempty" env:"YAO_KEY"` // The HTTPS certificate key path + Log string `json:"log,omitempty" env:"YAO_LOG"` // The log file path + LogMode string `json:"log_mode,omitempty" env:"YAO_LOG_MODE" envDefault:"TEXT"` // The log mode TEXT|JSON + LogMaxSize int `json:"log_max_size,omitempty" env:"YAO_LOG_MAX_SIZE" envDefault:"100"` // The max log size in MB, the default is 100 + LogMaxAage int `json:"log_max_age,omitempty" env:"YAO_LOG_MAX_AGE" envDefault:"7"` // The max log age in day, the default is 7 + LogMaxBackups int `json:"log_max_backups" env:"YAO_LOG_MAX_BACKUPS" envDefault:"3"` // The max log backups, the default is 3 + LogLocalTime bool `json:"log_local_time" env:"YAO_LOG_LOCAL_TIME" envDefault:"true"` + JWTSecret string `json:"jwt_secret,omitempty" env:"YAO_JWT_SECRET"` // The JWT Secret + DB Database `json:"db,omitempty"` // The database config + AllowFrom []string `json:"allowfrom,omitempty" envSeparator:"|" env:"YAO_ALLOW_FROM"` // Domain list the separator is | + Session Session `json:"session,omitempty"` // Session Config + Runtime Runtime `json:"runtime,omitempty"` // Runtime config + Trace Trace `json:"trace,omitempty"` // Trace config + Registry string `json:"registry,omitempty" env:"YAO_REGISTRY" envDefault:"https://registry.yaoagents.com"` // The package registry server URL + GRPC GRPCConfig `json:"grpc,omitempty"` +} + +// GRPCConfig gRPC server configuration +type GRPCConfig struct { + Enabled string `json:"enabled,omitempty" env:"YAO_GRPC"` // Set "off" to disable gRPC server + Host string `json:"host,omitempty" env:"YAO_GRPC_HOST" envDefault:"127.0.0.1"` // Comma-separated bind addresses + Port int `json:"port,omitempty" env:"YAO_GRPC_PORT" envDefault:"9099"` // Listen port shared by all addresses } // Database 数据库配置 diff --git a/grpc/DESIGN.md b/grpc/DESIGN.md new file mode 100644 index 00000000..d5dbc976 --- /dev/null +++ b/grpc/DESIGN.md @@ -0,0 +1,448 @@ +# Yao gRPC Server + +General-purpose gRPC gateway for the Yao process. Shares OAuth + ACL scope system with openapi — one token, two protocols. + +## Services + +| Layer | Method | Purpose | Scope | +|-------|--------|---------|-------| +| **Base** | `Run` | Execute Yao process, return result | `grpc:run` | +| | `Stream` | Execute Yao process, stream output | `grpc:stream` | +| | `Shell` | Execute system command, wait for result | `grpc:shell` | +| | `ShellStream` | Execute system command, stream stdout/stderr | `grpc:shell` | +| **API** | `API` | Proxy to openapi, any endpoint | openapi's own scopes | +| **MCP** | `MCPListTools` | List MCP tools for a session | `grpc:mcp` | +| | `MCPCallTool` | Call MCP tool → process.Exec() | `grpc:mcp` | +| | `MCPListResources` | List MCP resources | `grpc:mcp` | +| | `MCPReadResource` | Read MCP resource | `grpc:mcp` | +| **LLM** | `ChatCompletions` | Send messages to LLM, get response | `grpc:llm` | +| | `ChatCompletionsStream` | Stream LLM response (SSE → gRPC stream) | `grpc:llm` | +| **Agent** | `AgentStream` | Call agent, stream response | `grpc:agent` | + +## Clients + +- Container MCP tools (via Tai gRPC relay) +- `yao run` CLI (after `yao login`) +- Yao-to-Yao (cross-node process execution) + +## Auth + +Same as openapi. gRPC auth interceptor reuses the same `guard.Authenticate` logic — including automatic token refresh when access token is expired but refresh token is valid. + +``` +metadata (Bearer + x-refresh-token) + → VerifyToken + → expired? → TryRefresh (same as guard.go) → new tokens in response metadata + → extract scopes → acl.Scope.Check(method, path, scopes) +``` + +### Infrastructure reuse assessment + +Existing openapi/oauth infrastructure can be reused for gRPC with **zero modifications**: + +| Component | Reusable as-is | Notes | +|-----------|---------------|-------| +| `VerifyToken(token string)` | Yes | Pure string input, no Gin dependency | +| `MakeAccessToken(clientID, scope, subject, expiresIn, extraClaims...)` | Yes | Supports custom scope/subject for container tokens | +| `MakeRefreshToken(...)` | Yes | Same as above | +| `Revoke(ctx, token, tokenTypeHint)` | Yes | For container token cleanup on Remove | +| `ScopeManager.Check(req *AccessRequest)` | Yes | Only needs `(Method, Path, Scopes)` — no Gin dependency | +| `acl.Register(...)` | Yes | gRPC scopes registered via same pattern | + +The `authorized.SetInfo` / `authorized.GetInfo` are Gin-bound but **not needed** — gRPC interceptor builds `AccessRequest` directly from JWT claims. Full `Enforce` chain (client/team/member) is HTTP multi-tenant only; gRPC uses `VerifyToken → ScopeManager.Check` which is sufficient. + +New code required: ~80 lines (interceptor + scope registration). Existing code changes: **zero**. + +### CLI auth: `yao login` / `yao logout` + +OAuth 2.0 Device Authorization Grant. No `--remote` flag needed — logged in = gRPC, not logged in = local. + +``` +$ yao login --server https://yao.example.com +请访问: https://yao.example.com/device +输入代码: ABCD-1234 +等待授权... ✓ (token saved to ~/.yao/credentials) + +$ yao run models.user.Find '{"id":1}' ← auto gRPC +$ yao logout +``` + +Requires two new openapi endpoints: +- `POST /oauth/device/authorize` — issue device_code + user_code +- `POST /oauth/device/token` — poll for access_token + +Token scope: based on user's role, e.g. `grpc:run grpc:stream grpc:shell grpc:llm grpc:agent grpc:mcp`. + +**Implementation cost**: ~190 lines new code, ~10 lines changes to existing code. +Scaffolding already in place — `types.DeviceAuthorizationResponse`, `GrantTypeDeviceCode`, error codes (`ErrorAuthorizationPending`, `ErrorSlowDown`), `DeviceCodeLifetime` config, `DeviceAuthorization()` method signature, and HTTP route are all pre-defined. Core work: + +1. Implement `DeviceAuthorization()` in `device.go` (currently returns `nil, nil`) +2. Add device_code store/get/consume helpers in `token.go` +3. Add `GrantTypeDeviceCode` case to `Token()` switch in `core.go` (1 case branch) +4. Implement `handleDeviceCodeGrant()` in `core.go` +5. Add user authorization callback handler +6. Fix discovery endpoint path inconsistency (`/oauth/device` vs `/oauth/device_authorization`) + +Risk: **very low** — all additions are in isolated code paths, no changes to existing `authorization_code` / `client_credentials` / `refresh_token` flows. + +### Container token + +Container images and `yao-grpc` (`yao/tai/grpc/`) are ours — it handles token refresh automatically. + +``` +Manager creates container + ├─ oauth.MakeAccessToken(subject=userID, scope="grpc:mcp grpc:run") + ├─ oauth.MakeRefreshToken(...) + └─ tai.Client.Sandbox().Create(CreateRequest{ + Env: { + YAO_TOKEN, YAO_REFRESH_TOKEN, YAO_SANDBOX_ID, + YAO_GRPC_ADDR, // where to connect + YAO_GRPC_UPSTREAM, // remote only: where Tai should forward to + }, + }) + + Local: YAO_GRPC_ADDR=127.0.0.1:9099 (direct to Yao, no upstream needed) + Remote: YAO_GRPC_ADDR=tai-host:9100 YAO_GRPC_UPSTREAM=yao-host:9099 + +yao-grpc (tai/grpc/, container-internal) + ├─ reads YAO_GRPC_ADDR + YAO_TOKEN + YAO_REFRESH_TOKEN + YAO_SANDBOX_ID from env + ├─ if YAO_GRPC_UPSTREAM set: attaches x-grpc-upstream metadata (tells Tai where to forward) + ├─ every call: Bearer token + x-refresh-token + x-sandbox-id in gRPC metadata + ├─ server auth interceptor reuses guard.Authenticate logic: + │ token valid → pass through + │ token expired + refresh token present → auto rotate (same as HTTP guard) + │ new tokens returned via response metadata (x-access-token, x-refresh-token) + ├─ yao-grpc reads response metadata, updates tokens in memory + └─ transparent to caller, no separate refresh RPC needed +``` + +- access_token: short TTL (15m) +- refresh_token: no expiry (valid until container removed) +- Manager revokes refresh_token on container Remove +- Tai does NOT know Yao address at startup — yao-grpc carries target in request metadata + +### Virtual endpoint mapping + +| gRPC | Virtual endpoint | +|------|-----------------| +| Run("models.user.Find") | `POST /grpc/run/models.user.Find` | +| Stream("flows.report") | `POST /grpc/stream/flows.report` | +| Shell | `POST /grpc/shell` | +| ShellStream | `POST /grpc/shell` (same) | +| API(POST, /kb/collections) | `POST /kb/collections` (real openapi path) | +| MCPListTools | `GET /grpc/mcp/tools` | +| MCPCallTool("search") | `POST /grpc/mcp/call/search` | +| MCPListResources | `GET /grpc/mcp/resources` | +| MCPReadResource("uri") | `GET /grpc/mcp/resources/read` | +| ChatCompletions | `POST /grpc/llm/completions` | +| ChatCompletionsStream | `POST /grpc/llm/completions` (same) | +| AgentStream("robot-id") | `POST /grpc/agent/robot-id` | + +API method uses the **actual openapi path** — no virtual mapping needed, scope check is identical to HTTP. + +### Scope registration + +```go +func init() { + acl.Register( + &acl.ScopeDefinition{Name: "grpc:run", Endpoints: []string{"POST /grpc/run/*"}}, + &acl.ScopeDefinition{Name: "grpc:stream", Endpoints: []string{"POST /grpc/stream/*"}}, + &acl.ScopeDefinition{Name: "grpc:shell", Endpoints: []string{"POST /grpc/shell"}}, + &acl.ScopeDefinition{Name: "grpc:mcp", Endpoints: []string{"GET /grpc/mcp/tools", "POST /grpc/mcp/call/*", "GET /grpc/mcp/resources", "GET /grpc/mcp/resources/read"}}, + &acl.ScopeDefinition{Name: "grpc:llm", Endpoints: []string{"POST /grpc/llm/completions"}}, + &acl.ScopeDefinition{Name: "grpc:agent", Endpoints: []string{"POST /grpc/agent/*"}}, + ) +} +``` + +## Network + +### Server listen config + +| Env | Default | Purpose | +|-----|---------|---------| +| `YAO_GRPC_HOST` | `127.0.0.1` | Comma-separated bind addresses. | +| `YAO_GRPC_PORT` | `9099` | Listen port (shared by all addresses). | +| `YAO_GRPC` | _(unset)_ | Set `off` to explicitly disable gRPC server. | + +gRPC server **defaults to enabled** (`127.0.0.1:9099`) — sandbox container callbacks depend on it. + +`YAO_GRPC_HOST` accepts one or more addresses separated by `,`. Each address gets its own `net.Listener`; all listeners feed into the same `grpc.Server` (gRPC supports multiple `Serve` calls on one server). + +| Scenario | Config | Effect | +|----------|--------|--------| +| Local dev / default | _(nothing to set)_ | `127.0.0.1:9099` — loopback, sandbox works out of box | +| LAN multi-NIC | `YAO_GRPC_HOST=192.168.10.1,10.0.0.1` | Binds each internal IP | +| Open | `YAO_GRPC_HOST=0.0.0.0` | All interfaces | +| Disabled | `YAO_GRPC=off` | gRPC server not started (pure API gateway, no sandbox) | + +When multiple addresses are given, the server creates one goroutine per listener. Shutdown (`grpc.GracefulStop`) drains all listeners. + +Config lives in `config.Config.GRPC` (type `GRPCConfig`), same pattern as `Host`/`Port` for HTTP. + +### Startup + +gRPC server starts **after** HTTP server in `cmd/start.go`, as a parallel goroutine: + +``` +engine.Load → itask.Start → ischedule.Start → service.Start (HTTP) → grpc.StartServer (gRPC) +``` + +gRPC server starts by default. Set `YAO_GRPC=off` to explicitly disable (no-op startup). Any other value or unset means enabled. + +Shutdown: `defer grpc.Stop()` in `cmd/start.go`, called before HTTP stop for graceful drain. + +### Access control + +Local: containers and CLI connect via loopback. Remote: only Tai relay connects (address known from `YAO_TAI_ADDR`). All callers carry OAuth tokens — no IP allowlist needed. + +Interceptor chain: auth → ACL → handler. + +Public methods (skip auth): `Healthz`. Auth interceptor checks method name and passes through. + +## IPC Path (replacing Unix socket) + +All modes use gRPC — no Unix socket fallback. One code path, local and remote. + +``` +Local: Container → yao-grpc → Yao gRPC 127.0.0.1:9099 +Remote: Container → yao-grpc → Tai :9100 relay → Yao gRPC :9099 +``` + +`yao-grpc` reads `YAO_GRPC_ADDR` from env and connects. Local containers point directly at the Yao gRPC server on loopback; remote containers point at the Tai relay. No mode switch, no branching. + +### Tai relay routing + +Tai does **not** know the Yao gRPC address at startup. yao-grpc tells Tai where to forward on every request via metadata: + +``` +Manager.Create(sandbox) + ├─ oauth.MakeAccessToken(...) + ├─ oauth.MakeRefreshToken(...) + └─ tai.Client.Sandbox().Create(CreateRequest{ + Env: { + YAO_TOKEN, YAO_REFRESH_TOKEN, + YAO_GRPC_ADDR: "tai-host:9100", + YAO_GRPC_UPSTREAM: "yao-host:9099", + }, + }) +``` + +yao-grpc reads `YAO_GRPC_UPSTREAM` from env and attaches it as `x-grpc-upstream` metadata on every request to Tai. Tai gateway reads this metadata and forwards to the specified address. No per-container state in Tai, no lookup table — pure transparent proxy. One Tai can serve containers from different Yao instances because each request carries its own target. + +For local mode, no Tai relay — Manager injects `YAO_GRPC_ADDR=127.0.0.1:9099` directly (no `YAO_GRPC_UPSTREAM` needed). + +### yao-grpc (container client) + +`yao-grpc` is the in-container gRPC client binary. Replaces the old `yao-bridge`. Lives in `yao/tai/grpc/`: + +``` +yao/tai/grpc/ +├── grpc.go // gRPC client: connect, forward MCP/process calls +├── auth.go // token management: read env, auto-refresh +├── grpc_test.go +└── cmd/ + └── main.go +``` + +Rationale for placing in `yao/tai`: +- Consumes Tai relay — same layer as `tai/proxy`, `tai/volume` +- Shares gRPC deps already in `yao/tai` +- Version-locked with Tai SDK and server protocol +- Built in same CI: `go build -o yao-grpc ./tai/grpc/cmd` + +Pure client — no signing keys, no `oauth` package dependency. Reads `YAO_TOKEN` + `YAO_REFRESH_TOKEN` + `YAO_SANDBOX_ID` from env, attaches all three as gRPC metadata on every call. Token refresh is transparent — server auto-rotates expired tokens (same logic as HTTP guard) and returns new tokens via response metadata. + +## Proto + +```protobuf +service Yao { + // Base + rpc Run(RunRequest) returns (RunResponse); + rpc Stream(RunRequest) returns (stream Chunk); + rpc Shell(ShellRequest) returns (ShellResponse); + rpc ShellStream(ShellRequest) returns (stream Chunk); + + // API gateway + rpc API(APIRequest) returns (APIResponse); + + // MCP + rpc MCPListTools(MCPListRequest) returns (MCPListResponse); + rpc MCPCallTool(MCPCallRequest) returns (MCPCallResponse); + rpc MCPListResources(MCPListRequest) returns (MCPResourcesResponse); + rpc MCPReadResource(MCPResourceRequest) returns (MCPResourceResponse); + + // AI - LLM + rpc ChatCompletions(ChatRequest) returns (ChatResponse); + rpc ChatCompletionsStream(ChatRequest) returns (stream ChatChunk); + + // AI - Agent + rpc AgentStream(AgentRequest) returns (stream AgentChunk); + + // Health + rpc Healthz(Empty) returns (HealthzResponse); +} +``` + +### LLM layer + +`ChatCompletions` and `ChatCompletionsStream` call the existing `llm.ChatCompletions` process (`agent/llm/process.go`). It auto-detects connector type (openai/anthropic/etc.), selects the appropriate provider, and returns OpenAI-compatible format. + +``` +gRPC ChatCompletions(connector, messages, opts) + → process.Exec("llm.ChatCompletions", connector, messages, opts) + → agent/llm.New(conn, opts) → provider.Stream/Post → response + +gRPC ChatCompletionsStream(connector, messages, opts) + → same path, with streaming callback → gRPC stream chunks +``` + +The caller specifies a connector ID. The `llm.ChatCompletions` process resolves it via `connector.Select()`, creates the LLM instance, and executes. Streaming version passes a callback that forwards chunks to the gRPC stream. + +### Agent layer + +`AgentStream` wraps `agent/robots/:id/completions` — resolves robot → host assistant → runs agent pipeline → streams output. Only stream method — agent output is inherently streamed; non-stream callers simply consume all chunks. Internally calls `assistant.Stream()` with `ctx.Writer` set to nil (or noop) when the caller doesn't need incremental output. + +``` +gRPC AgentStream(agent_id, messages) → resolve robot → assistant.Stream() → stream chunks +``` + +This enables container-internal agents to call other agents without HTTP, and remote `yao` instances to orchestrate agent pipelines cross-node. + +`AgentChunk` carries `agent/output/message.Message` — the same DSL used by HTTP SSE streaming. Each chunk is one JSON-serialized `Message`: + +```protobuf +message AgentChunk { + bytes data = 1; // JSON-encoded agent/output/message.Message + bool done = 2; +} +``` + +The `Message` structure uses `Type` + `Props` to express all content types (text, thinking, tool_call, error, action, event, image, audio, video). Streaming control fields (`chunk_id`, `message_id`, `block_id`, `thread_id`) and delta fields (`delta`, `delta_path`, `delta_action`) are preserved as-is over gRPC — the client merges chunks using the same logic as CUI's SSE consumer. + +### Shell execution context + +`Shell` and `ShellStream` execute commands in the **Yao host process**, not inside a sandbox container. This is by design — the scope `grpc:shell` is a privileged capability, not granted to container tokens by default. Container-internal commands run via `tai.Client.Sandbox().Exec()`, which is a different path (not exposed as a gRPC method). + +See [pb/yao.proto](./pb/yao.proto) for full message definitions. + +## Process & Stream (gou foundation) + +gRPC `Run` and `Stream` map to two parallel systems in `gou`: + +``` +gou/process/ — execute once, return result → gRPC Run +gou/stream/ — execute once, push chunks → gRPC Stream +``` + +### gou/process (existing, unchanged) + +```go +type Handler func(process *Process) interface{} + +process.Register("scripts", handler) +p := process.New("scripts.foo.bar", args...) +p.Execute() +result := p.Value() +``` + +### gou/stream (new package, parallel to process) + +```go +type Handler func(ctx context.Context, process *Process, send func([]byte) error) error + +stream.Register("scripts", handler) +s := stream.New("scripts.foo.bar", args...) +s.Execute(ctx, func(chunk []byte) error { ... }) +``` + +`stream.Process` mirrors `process.Process` fields (Name, Group, Method, ID, Args, Global, Sid, Authorized) but `ctx` is a first-class parameter, not buried in a struct field. + +`send` returns error when the receiver disconnects — handler should stop. + +### Fallback + +If a stream handler is not registered for a name but a process handler exists, `stream.Execute` falls back to: run the process handler once, JSON-marshal the result, call `send` once. + +### Registration + +```go +// gou/process — existing +process.Register("models", modelsHandler) +process.Register("scripts", scriptsHandler) + +// gou/stream — new, same namespace +stream.Register("scripts", scriptsStreamHandler) +stream.Register("llm", llmStreamHandler) +``` + +Same naming convention. A process name can have both a process handler and a stream handler. + +### gRPC mapping + +```go +func (s *yaoServer) Run(ctx context.Context, req *pb.RunRequest) (*pb.RunResponse, error) { + p := process.NewWithContext(ctx, req.Process, args...) + if err := p.Execute(); err != nil { return nil, err } + data, _ := json.Marshal(p.Value()) + return &pb.RunResponse{Result: data}, nil +} + +func (s *yaoServer) Stream(req *pb.RunRequest, grpcStream pb.Yao_StreamServer) error { + st := stream.New(req.Process, args...) + return st.Execute(grpcStream.Context(), func(chunk []byte) error { + return grpcStream.Send(&pb.Chunk{Data: chunk}) + }) +} +``` + +### V8 integration + +Both are exposed as top-level globals in JavaScript, parallel: + +```go +// gou/runtime/v8/isolate.go MakeTemplate +template.Set("Process", processModule.ExportFunction(iso)) // existing +template.Set("Stream", streamModule.ExportFunction(iso)) // new +``` + +**JS calling Go stream** (JS is consumer): + +```javascript +Stream("llm.chat.completions", function(chunk) { + log.Info(chunk) + return 1 // 1=continue, 0=stop +}, { model: "gpt-4", messages: [...] }) +``` + +**JS script as stream handler** (JS is producer): + +```javascript +// scripts/report.js — registered via stream.Register("scripts", ...) +function generate(args, send) { + send("part 1") + send("part 2") +} +``` + +V8 runtime registers both: + +```go +func init() { + process.Register("scripts", processScripts) // existing + stream.Register("scripts", processScriptsStream) // new +} +``` + +`processScriptsStream` calls `script.ExecStream(ctx, p, send)` which injects `send` into the V8 global before executing the script method. + +### Impact on existing code + +| Component | Changes | +|-----------|---------| +| `gou/process/` | None | +| `gou/stream/` | New package (~150 lines) | +| `gou/runtime/v8/process.go` | +1 line: `stream.Register(...)` | +| `gou/runtime/v8/script.go` | +`ExecStream` method | +| `gou/runtime/v8/isolate.go` | +1 line: `template.Set("Stream", ...)` | +| `gou/runtime/v8/functions/` | +`stream/` module for JS→Go stream consumption | diff --git a/grpc/IMPL.md b/grpc/IMPL.md new file mode 100644 index 00000000..fefc76f0 --- /dev/null +++ b/grpc/IMPL.md @@ -0,0 +1,274 @@ +# Yao gRPC Server — Implementation Plan + +Design: [DESIGN.md](./DESIGN.md) + +## Scope + +**V1**: Auth + unary RPCs + LLM/Agent streaming + container client. + +**V2**: Base streaming (`Stream`, `ShellStream`) + `gou/stream` package + V8 integration. + +## Package Structure + +``` +grpc/ +├── grpc.go // StartServer, config, server lifecycle +├── pb/ +│ ├── yao.proto +│ ├── yao.pb.go // generated +│ └── yao_grpc.pb.go // generated +├── auth/ +│ ├── guard.go // unary + stream interceptor (calls oauth.VerifyToken, ScopeManager.Check) +│ ├── endpoint.go // gRPC method → virtual HTTP endpoint mapping +│ └── scope.go // init() acl.Register for grpc:* scopes +├── run/ +│ └── run.go // Run handler +├── shell/ +│ └── shell.go // Shell, ShellStream (V2) handlers +├── api/ +│ └── api.go // API proxy handler +├── mcp/ +│ └── mcp.go // MCPListTools, MCPCallTool, MCPListResources, MCPReadResource +├── llm/ +│ └── llm.go // ChatCompletions, ChatCompletionsStream +├── agent/ +│ └── agent.go // AgentStream +└── health/ + └── health.go // Healthz +``` + +Container client: + +``` +tai/grpc/ +├── grpc.go // gRPC client, Dial, method wrappers +├── auth.go // read env tokens, attach metadata, handle refresh +├── grpc_test.go +└── cmd/ + └── main.go // yao-grpc binary entry +``` + +## V1 Phases + +### Phase 0: Proto + codegen ✅ + +No dependency. + +| Task | Detail | Status | +|------|--------|--------| +| `grpc/pb/yao.proto` | All 14 RPCs + all message types. V2 methods (`Stream`, `ShellStream`) included in proto, handler left `Unimplemented`. | ✅ Done | +| codegen | `protoc` → `pb/*.pb.go` + `pb/*_grpc.pb.go` | ✅ Done | + +### Phase 1: Auth + server skeleton ✅ + +Depends on: Phase 0. Auth is ~80 lines new code calling existing `openapi/oauth` functions. + +| Task | Detail | Status | +|------|--------|--------| +| `grpc/auth/scope.go` | `init()` — `acl.Register` 6 gRPC scope definitions | ✅ Done | +| `grpc/auth/endpoint.go` | Map gRPC method + request params → virtual HTTP endpoint for ACL (e.g. `Run("models.user.Find")` → `POST /grpc/run/models.user.Find`) | ✅ Done | +| `grpc/auth/guard.go` | Extract Bearer from metadata → `oauth.AuthenticateToken` (pure, no gin) → ACL scope check. Skip `Healthz`. New tokens via `SendHeader`. | ✅ Done | +| `openapi/oauth/authenticate.go` | `AuthenticateToken(AuthInput) → AuthResult` — gin-free auth core. `refreshTokenDirect`, `buildAuthInfo`. Shares `refreshGates` with `TryRefreshToken`. | ✅ Done | +| `grpc/grpc.go` | `StartServer(cfg)` — `grpc.NewServer` with interceptor, register service, listen. See **Server config & startup** below. | ✅ Done | +| `grpc/health/health.go` | `Healthz` → `{status: "ok"}` | ✅ Done | +| `config/types.go` | Add `GRPC` field to `Config` struct — see config below | ✅ Done | +| `cmd/start.go` | After `service.Start(config.Conf)` (HTTP ready), call `grpc.StartServer(config.Conf)` in goroutine. Print gRPC listen address in Access Points block. `defer grpc.Stop()` in shutdown path. | ✅ Done | + +**Server config & startup:** + +Config struct addition (`config/types.go`): + +```go +type Config struct { + // ... existing fields ... + GRPC GRPCConfig `json:"grpc,omitempty"` +} + +type GRPCConfig struct { + Enabled string `json:"enabled,omitempty" env:"YAO_GRPC"` + Host string `json:"host,omitempty" env:"YAO_GRPC_HOST" envDefault:"127.0.0.1"` + Port int `json:"port,omitempty" env:"YAO_GRPC_PORT" envDefault:"9099"` +} +``` + +- **Default** — `127.0.0.1:9099`, enabled. Sandbox callbacks work out of box. +- `YAO_GRPC_HOST=192.168.10.1,10.0.0.1` — comma-separated, binds each IP for multi-NIC LAN +- `YAO_GRPC_HOST=0.0.0.0` — all interfaces +- `YAO_GRPC=off` — explicitly disable gRPC server + +`grpc.StartServer` implementation: + +```go +func StartServer(cfg config.Config) error { + if strings.ToLower(cfg.GRPC.Enabled) == "off" { + log.Info("gRPC server disabled (YAO_GRPC=off)") + return nil + } + hosts := strings.Split(cfg.GRPC.Host, ",") + for _, host := range hosts { + addr := net.JoinHostPort(strings.TrimSpace(host), strconv.Itoa(cfg.GRPC.Port)) + lis, err := net.Listen("tcp", addr) + // ... error handling ... + go server.Serve(lis) // one goroutine per listener, same grpc.Server + } + return nil +} +``` + +Startup sequence in `cmd/start.go`: + +``` +engine.Load(cfg) +itask.Start() +ischedule.Start() +service.Start(cfg) // HTTP server +grpc.StartServer(cfg) // gRPC server (after HTTP, parallel goroutine) +// ... event loop ... +defer grpc.Stop() // GracefulStop drains all listeners (no-op if not started) +``` + +`cmd/start.go` prints each gRPC listen address: + +``` +Listening 0.0.0.0:5099 (HTTP) +Listening 192.168.10.1:9099 (gRPC) +Listening 10.0.0.1:9099 (gRPC) +``` + +Deliverable: Server starts, Healthz works, unauthenticated calls rejected, token refresh via metadata works. + +### Phase 2: Base + API + MCP handlers ✅ + +Depends on: Phase 1. + +| Task | Detail | Status | +|------|--------|--------| +| `grpc/run/run.go` | `Run` — `process.New(req.Process, args...).Exec()`. Injects `AuthorizedInfo` via `p.WithSID()` + `p.WithAuthorized()`. | ✅ Done | +| `grpc/shell/shell.go` | `Shell` — `exec.CommandContext` in host process. **Security**: refuse execution if Yao process is running as root (`os.Getuid() == 0` → `PermissionDenied`). Timeout: use request `timeout` field, default 30s, capped by server max. | ✅ Done | +| `grpc/api/api.go` | `API` — build `http.Request`, call openapi internally | ✅ Done | +| `grpc/mcp/mcp.go` | `MCPListTools`, `MCPCallTool`, `MCPListResources`, `MCPReadResource` | ✅ Done | + +Deliverable: Base + API + MCP methods work with valid tokens. + +### Phase 3: LLM + Agent handlers ✅ + +Depends on: Phase 1. No code dependency on Phase 2 — can parallel. + +| Task | Detail | Status | +|------|--------|--------| +| `grpc/llm/llm.go` | `ChatCompletions` / `ChatCompletionsStream` — direct call to `agent/llm` (`connector.Select` → `llm.New` → `Stream`). Uses `agent/llm.BuildCompletionOptions`. Constructs `agent/context.Context` with `AuthorizedInfo`. | ✅ Done | +| `grpc/agent/agent.go` | `AgentStream` — `assistant.Get` → `ast.Stream` with `grpcStreamWriter` adapter bridging `http.ResponseWriter` to gRPC `ServerStreamingServer`. Constructs `agent/context.Context` with `AuthorizedInfo`. | ✅ Done | + +Deliverable: LLM (unary + stream) and Agent streaming via gRPC. + +### Phase 4: Tai gateway change (Tai repo) ⏳ + +Depends on: Phase 1 (need proto definitions for testing). yao-grpc depends on this. + +Tai gateway currently dials a fixed `YaoUpstream` at startup. New behavior: yao-grpc tells Tai where to forward via request metadata (`x-grpc-upstream`). Tai reads the target address and proxies to it — removes `YaoUpstream` startup config. + +| Task | Detail | Status | +|------|--------|--------| +| Tai `gateway/gateway.go` | Remove fixed `upstream *grpc.ClientConn`. On each request, read `x-grpc-upstream` from metadata → lookup/create conn from `sync.Map` cache (key = address string) → forward. Typical deployment has 1 upstream, cache stays tiny. | ⏳ Pending | +| Tai `server/server.go` | Remove `YaoUpstream` from `Config`. Gateway init no longer needs an address. | ⏳ Pending | + +Connection cache: `sync.Map[string, *grpc.ClientConn]` — lazy dial on first request per upstream, reuse thereafter. No eviction needed (upstream count ≈ 1 in practice). `GracefulStop` closes all cached connections. + +Deliverable: Tai starts without Yao address. Forwards based on request metadata. + +### Phase 5: yao-grpc container client ⏳ + +Depends on: Phase 1 (server + auth), Phase 4 (Tai gateway accepts `x-grpc-upstream`). + +| Task | Detail | Status | +|------|--------|--------| +| `tai/grpc/grpc.go` | `Dial(YAO_GRPC_ADDR)`, method wrappers mirroring server | ⏳ Pending | +| `tai/grpc/auth.go` | Read `YAO_TOKEN` / `YAO_REFRESH_TOKEN` / `YAO_SANDBOX_ID` / `YAO_GRPC_UPSTREAM` from env. Attach as metadata on every call: Bearer token, `x-refresh-token`, `x-sandbox-id`, `x-grpc-upstream` (if set, for Tai relay). Read `SendHeader` for rotated tokens, update in memory. | ⏳ Pending | +| `tai/grpc/cmd/main.go` | Stdio MCP server: JSON-RPC → gRPC. Replaces `yao-bridge`. `yao-grpc version` prints version/commit/build time (via `-ldflags`), for container debugging. | ⏳ Pending | +| `tai/grpc/grpc_test.go` | Tests | ⏳ Pending | + +Container token issuance uses existing `oauth.MakeAccessToken` / `oauth.MakeRefreshToken` — called by sandbox Manager at container creation, injected as env vars. Revoke on Remove. No new auth code needed on the issuance side. + +Deliverable: `go build -o yao-grpc ./tai/grpc/cmd`. + +### Phase 6: Device Flow — backend (`yao login`) ⏳ + +Depends on: Phase 1. Independent — can parallel with Phase 2-5. + +| Task | Detail | Status | +|------|--------|--------| +| `oauth/device.go` | Implement `DeviceAuthorization()` — generate `device_code` + `user_code`, store with expiry | ⏳ Pending | +| `oauth/token.go` | Device code store/get/consume helpers | ⏳ Pending | +| `oauth/core.go` | Add `GrantTypeDeviceCode` case → `handleDeviceCodeGrant()` (poll returns `authorization_pending` / token) | ⏳ Pending | +| `cmd/yao/login.go` | `yao login --server ` → device flow → poll token endpoint → save `~/.yao/credentials` | ⏳ Pending | +| `cmd/yao/logout.go` | Revoke + delete credentials | ⏳ Pending | +| `cmd/yao/run.go` | Credentials exist → gRPC; otherwise local. Non-silent mode prints `⟶ user@host (gRPC)` header before execution (same line position as existing `Run: process.name`). Silent mode (`-s`) keeps pure output — no connection info, for shell scripting. | ⏳ Pending | + +Deliverable: `yao login` + `yao run` via gRPC (backend complete, auth page in Phase 7). + +### Phase 7: Device Flow — CUI auth page (frontend) ⏳ + +Depends on: Phase 6 (backend endpoints ready). This is a **frontend-only** task in the CUI repo. + +Route: `/auth/device` (Umi convention-based routing → `pages/auth/device/index.tsx`) + +| Task | Detail | Status | +|------|--------|--------| +| `pages/auth/device/index.tsx` | Device authorization page. User enters `user_code` and clicks Authorize. Uses `AuthLayout` + `AuthInput` + `AuthButton` from existing `pages/auth/components/`. | ⏳ Pending | +| `pages/auth/device/index.less` | Styles, follow `pages/auth/entry/index.less` pattern | ⏳ Pending | + +**Implementation details:** + +- Framework: React + UmiJS Max + Ant Design + MobX (same as all auth pages) +- Layout: Wrap with `AuthLayout` (logo + theme switch), same as `/auth/entry` +- Components reuse: `AuthInput` for `user_code` input, `AuthButton` for submit, from `pages/auth/components/` +- Page export: `export default observer(DeviceAuth)` (same pattern as `pages/auth/entry/index.tsx`) +- API: `window.$app.openapi` → call backend `POST /oauth/device/authorize` with `{ user_code }`, bearer token from current session +- Auth: User must be logged in (redirect to `/auth/entry` if not). After authorizing, show success message and close/redirect +- i18n: Use `useIntl()` hook for text, support `zh-CN` / `en-US` +- Flow: User opens URL from CLI prompt → logs in if needed → enters user_code → clicks Authorize → backend binds device_code to user → CLI poll gets token + +Deliverable: `/auth/device` page in CUI. User can authorize CLI device login from browser. + +## V2 Phases + +### Phase 8: `gou/stream` package ⏳ + +| Task | Detail | Status | +|------|--------|--------| +| `gou/stream/` | ~150 lines. `Handler`, `Process`, `Register`, `New`, `Execute`. Fallback to process. | ⏳ Pending | +| V8 | `stream.Register("scripts", ...)`, `ExecStream`, `template.Set("Stream", ...)`, JS `Stream()` global | ⏳ Pending | + +### Phase 9: Base streaming handlers ⏳ + +Depends on: Phase 8. + +| Task | Detail | Status | +|------|--------|--------| +| `grpc/run/run.go` | Add `Stream` handler — `stream.New(req.Process).Execute(ctx, send)` | ⏳ Pending | +| `grpc/shell/shell.go` | Add `ShellStream` handler — piped stdout → gRPC stream | ⏳ Pending | + +## Dependency Graph + +``` +Phase 0 (proto) ✅ + │ + ▼ +Phase 1 (auth + server) ✅ + │ + ├───────────┬───────────┬──────────────┐ + ▼ ▼ ▼ ▼ +Phase 2 ✅ Phase 3 ✅ Phase 4 (Tai) Phase 6 +(handlers) (LLM/Agent) │ (device backend) + ▼ │ + Phase 5 ▼ + (yao-grpc) Phase 7 + (CUI auth page) + +--- V2 --- + +Phase 8 (gou/stream) + │ + ▼ +Phase 9 (Stream, ShellStream) +``` diff --git a/grpc/TEST.md b/grpc/TEST.md new file mode 100644 index 00000000..3f947753 --- /dev/null +++ b/grpc/TEST.md @@ -0,0 +1,426 @@ +# Yao gRPC Server — Test Specification + +Design: [DESIGN.md](./DESIGN.md) | Implementation: [IMPL.md](./IMPL.md) + +## Principles + +- **Black-box testing**: all `*_test.go` files use `package xxx_test` — tests only access exported API via gRPC client +- **Tests follow implementation**: `*_test.go` lives next to the code it tests (`grpc/auth/guard_test.go` beside `grpc/auth/guard.go`) +- **Real server**: every test starts a real gRPC server on a random TCP port, exercises the full interceptor → handler chain +- **Coverage > 80%**: per sub-package and overall + +## Prerequisites + +```bash +source $YAO_SOURCE_ROOT/env.local.sh +``` + +Required environment variables (same as existing Yao tests): + +| Variable | Purpose | +|----------|---------| +| `YAO_TEST_APPLICATION` | Path to `yao-dev-app` | +| `YAO_DB_DRIVER` / `YAO_DB_PRIMARY` | Database connection | +| `YAO_JWT_SECRET` / `YAO_DB_AESKEY` | Crypto keys | +| `OPENAI_TEST_KEY` | LLM streaming tests | +| `ANTHROPIC_API_KEY` | LLM streaming tests (Anthropic) | + +## Directory Structure + +``` +grpc/ +├── grpc.go +├── tests/ +│ └── testutils/ +│ └── testutils.go # shared test utilities +├── auth/ +│ ├── guard.go +│ ├── guard_test.go # package auth_test +│ ├── endpoint.go +│ ├── endpoint_test.go # package auth_test +│ └── scope.go +├── run/ +│ ├── run.go +│ └── run_test.go # package run_test +├── shell/ +│ ├── shell.go +│ └── shell_test.go # package shell_test +├── api/ +│ ├── api.go +│ └── api_test.go # package api_test +├── mcp/ +│ ├── mcp.go +│ └── mcp_test.go # package mcp_test +├── llm/ +│ ├── llm.go +│ └── llm_test.go # package llm_test +├── agent/ +│ ├── agent.go +│ └── agent_test.go # package agent_test +└── health/ + ├── health.go + └── health_test.go # package health_test +``` + +Tests live beside the code they verify. `grpc/tests/testutils/` is shared infrastructure only. + +## testutils API + +`grpc/tests/testutils/testutils.go` provides the test harness used by all sub-packages. + +```go +package testutils + +// Prepare initializes the full Yao runtime (DB, V8, models, scripts, etc.) +// then starts a real gRPC server on :0 (random port). +// Returns a connected grpc.ClientConn ready to create service clients. +// +// Internally calls: +// test.Prepare(t, config.Conf) — Yao runtime +// grpc.StartServer(cfg{Port:0}) — gRPC server +// grpc.Dial("127.0.0.1:port") — client connection +func Prepare(t *testing.T) *grpc.ClientConn + +// Clean gracefully stops the gRPC server and tears down the Yao runtime. +// Always use with defer: +// conn := testutils.Prepare(t) +// defer testutils.Clean() +func Clean() + +// Addr returns the gRPC server address "127.0.0.1:xxxxx". +func Addr() string + +// ObtainAccessToken mints a token with the given scopes. +// Calls oauth.MakeAccessToken directly — no HTTP round-trip. +func ObtainAccessToken(t *testing.T, scopes ...string) string + +// ObtainAccessTokenForUser mints a token for a specific user ID. +func ObtainAccessTokenForUser(t *testing.T, userID string, scopes ...string) string + +// WithToken returns ctx with Bearer token in gRPC metadata. +func WithToken(ctx context.Context, token string) context.Context + +// WithRefreshToken returns ctx with both Bearer and x-refresh-token metadata. +func WithRefreshToken(ctx context.Context, token, refreshToken string) context.Context + +// WithSandboxMetadata returns ctx with x-sandbox-id and x-grpc-upstream metadata. +func WithSandboxMetadata(ctx context.Context, sandboxID, upstream string) context.Context + +// NewClient creates a pb.YaoServiceClient from a connection. +func NewClient(conn *grpc.ClientConn) pb.YaoServiceClient +``` + +## How to Write a Test + +### Standard pattern + +Every test file follows this structure: + +```go +// grpc/run/run_test.go +package run_test + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/yaoapp/yao/grpc/pb" + "github.com/yaoapp/yao/grpc/tests/testutils" +) + +func TestRun_ProcessExec(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.Run(ctx, &pb.RunRequest{ + Process: "utils.app.Ping", + }) + assert.NoError(t, err) + assert.NotNil(t, resp.Data) +} + +func TestRun_InvalidProcess(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.Run(ctx, &pb.RunRequest{Process: "nonexistent.process"}) + assert.Error(t, err) +} +``` + +### Auth tests + +Auth tests verify the interceptor chain through the gRPC client: + +```go +// grpc/auth/guard_test.go +package auth_test + +func TestAuth_NoToken_Rejected(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + client := testutils.NewClient(conn) + + _, err := client.Run(context.Background(), &pb.RunRequest{Process: "utils.app.Ping"}) + st, _ := status.FromError(err) + assert.Equal(t, codes.Unauthenticated, st.Code()) +} + +func TestAuth_WrongScope_Denied(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + client := testutils.NewClient(conn) + + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.Run(ctx, &pb.RunRequest{Process: "utils.app.Ping"}) + st, _ := status.FromError(err) + assert.Equal(t, codes.PermissionDenied, st.Code()) +} + +func TestAuth_TokenRefresh(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + client := testutils.NewClient(conn) + + // Mint an expired token + valid refresh token, + // send request with x-refresh-token metadata, + // verify response header contains x-new-access-token. +} + +func TestHealthz_Public(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + client := testutils.NewClient(conn) + + resp, err := client.Healthz(context.Background(), &pb.Empty{}) + assert.NoError(t, err) + assert.Equal(t, "ok", resp.Status) +} +``` + +### Streaming tests + +```go +// grpc/llm/llm_test.go +package llm_test + +func TestChatCompletionsStream(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + client := testutils.NewClient(conn) + + token := testutils.ObtainAccessToken(t, "grpc:llm") + ctx := testutils.WithToken(context.Background(), token) + + stream, err := client.ChatCompletionsStream(ctx, &pb.ChatRequest{ + // ... model, messages, etc. + }) + assert.NoError(t, err) + + var chunks int + for { + chunk, err := stream.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + chunks++ + assert.NotEmpty(t, chunk.Data) + } + assert.Greater(t, chunks, 0) +} +``` + +```go +// grpc/agent/agent_test.go +package agent_test + +func TestAgentStream(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + client := testutils.NewClient(conn) + + token := testutils.ObtainAccessToken(t, "grpc:agent") + ctx := testutils.WithToken(context.Background(), token) + + stream, err := client.AgentStream(ctx, &pb.AgentRequest{ + RobotID: "test-robot", + // ... + }) + assert.NoError(t, err) + + var chunks int + for { + chunk, err := stream.Recv() + if err == io.EOF { + break + } + assert.NoError(t, err) + chunks++ + // Each chunk carries JSON-serialized agent/output/message.Message + } + assert.Greater(t, chunks, 0) +} + +func TestAgentStream_InvalidRobot(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + client := testutils.NewClient(conn) + + token := testutils.ObtainAccessToken(t, "grpc:agent") + ctx := testutils.WithToken(context.Background(), token) + + stream, err := client.AgentStream(ctx, &pb.AgentRequest{ + RobotID: "nonexistent-robot", + }) + // Either err on open or first Recv returns error + if err == nil { + _, err = stream.Recv() + } + assert.Error(t, err) +} +``` + +## Required Test Cases + +Each sub-package must cover at minimum: + +| Sub-package | Required cases | +|-------------|----------------| +| `auth` | valid token / no token (Unauthenticated) / expired token + refresh / wrong scope (PermissionDenied) / Healthz skips auth | +| `health` | Healthz returns ok without token | +| `run` | valid process / nonexistent process / bad arguments | +| `shell` | valid command / command not found / timeout | +| `api` | valid proxy / 404 endpoint | +| `mcp` | MCPListTools / MCPCallTool / MCPListResources / MCPReadResource | +| `llm` | ChatCompletions (unary) / ChatCompletionsStream (multiple chunks) / invalid model | +| `agent` | AgentStream (receives message chunks) / nonexistent robot ID | + +## Makefile + +Add to [Makefile](../Makefile): + +```makefile +TESTFOLDER_GRPC := $(shell $(GO) list ./grpc/...) + +.PHONY: unit-test-grpc +unit-test-grpc: + echo "mode: count" > coverage.out + for d in $(TESTFOLDER_GRPC); do \ + $(GO) test -tags $(TESTTAGS) -v -timeout=10m \ + -covermode=count -coverprofile=profile.out \ + -coverpkg=$$(echo $$d | sed "s/\/test$$//g") \ + -skip='TestMemoryLeak|TestIsolateDisposal|TestLeak_|TestScenario_' \ + $$d > tmp.out; \ + cat tmp.out; \ + if grep -q "^--- FAIL" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "build failed" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "setup failed" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + elif grep -q "runtime error" tmp.out; then \ + rm tmp.out; \ + exit 1; \ + fi; \ + if [ -f profile.out ]; then \ + cat profile.out | grep -v "mode:" >> coverage.out; \ + rm profile.out; \ + fi; \ + done +``` + +Also add `|grpc` to the `TESTFOLDER_CORE` exclude pattern so core-test does not duplicate gRPC tests. + +## CI Integration + +Add `grpc-test` job to `unit-test.yml` and `pr-test.yml`: + +```yaml +grpc-test: + runs-on: ubuntu-latest + services: + mongodb: + image: mongo:6.0 + ports: + - 27017:27017 + env: + MONGO_INITDB_ROOT_USERNAME: root + MONGO_INITDB_ROOT_PASSWORD: "123456" + MONGO_INITDB_DATABASE: test + strategy: + matrix: + go: ["1.25"] + steps: + # ... standard checkout + setup (same as core-test) ... + + - name: Setup ENV (SQLite) + run: | + mkdir -p ${{ github.WORKSPACE }}/../app/db + echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV + echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV + + - name: Run gRPC Tests + run: make unit-test-grpc + + - name: Codecov Report + uses: codecov/codecov-action@v4 + with: + token: ${{ secrets.CODECOV_TOKEN }} +``` + +Key decisions: +- SQLite only — gRPC is a transport layer, no need for MySQL matrix +- No Qdrant/Neo4j/MCP-everything services needed +- LLM/Agent streaming uses real `OPENAI_TEST_KEY` + `ANTHROPIC_API_KEY` (same secrets as agent-test job) + +## Coverage + +- Target: >80% per sub-package, >80% overall +- `grpc.go` (server lifecycle) covered indirectly via testutils.Prepare/Clean +- Coverage collected via `-coverprofile`, reported to Codecov + +## Phase Test Schedule + +Tests are written alongside implementation, not after: + +| Phase | Test files | Repo | +|-------|------------|------| +| Phase 1 (auth + server) | `auth/guard_test.go`, `health/health_test.go` | yao | +| Phase 2 (handlers) | `run/run_test.go`, `shell/shell_test.go`, `api/api_test.go`, `mcp/mcp_test.go` | yao | +| Phase 3 (LLM + Agent) | `llm/llm_test.go`, `agent/agent_test.go` | yao | +| Phase 4 (Tai gateway) | Tai repo tests — gateway forwards `x-grpc-upstream`, conn cache reuse, missing metadata rejected | tai | +| Phase 5 (yao-grpc client) | `tai/grpc/grpc_test.go` — dial, method wrappers, token refresh via response metadata, `x-grpc-upstream` attachment | yao | +| Phase 6 (Device Flow) | `openapi/oauth/*_test.go` — DeviceAuthorization, device_code grant, poll pending/approved/expired | yao | + +Each Phase PR must include tests for all new code. Coverage must meet threshold before merge. + +## Running Tests + +```bash +# All gRPC tests +make unit-test-grpc + +# Single sub-package +go test -v ./grpc/auth/ + +# Single test +go test -v -run TestAuth_NoToken_Rejected ./grpc/auth/ +``` diff --git a/grpc/agent/agent.go b/grpc/agent/agent.go new file mode 100644 index 00000000..bd29cf3f --- /dev/null +++ b/grpc/agent/agent.go @@ -0,0 +1,93 @@ +package agent + +import ( + "net/http" + "strings" + + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/yaoapp/yao/agent" + "github.com/yaoapp/yao/agent/assistant" + agentContext "github.com/yaoapp/yao/agent/context" + "github.com/yaoapp/yao/grpc/auth" + "github.com/yaoapp/yao/grpc/pb" +) + +// Handler implements the AgentStream gRPC method. +type Handler struct{} + +// AgentStream resolves an assistant by ID and streams agent output as AgentChunk messages. +// Mirrors openapi/chat/completions.go GinCreateCompletions flow via context.GetGRPCAgentRequest. +func (h *Handler) AgentStream(req *pb.AgentRequest, stream grpc.ServerStreamingServer[pb.AgentChunk]) error { + ctx := stream.Context() + + if req.AssistantId == "" { + return status.Error(codes.InvalidArgument, "assistant_id is required") + } + + agentDSL := agent.GetAgent() + if agentDSL == nil { + return status.Error(codes.Internal, "agent DSL not initialized") + } + + cache, err := agentDSL.GetCacheStore() + if err != nil { + return status.Errorf(codes.Internal, "failed to get cache store: %v", err) + } + + messages, agentCtx, opts, err := agentContext.GetGRPCAgentRequest(ctx, agentContext.GRPCAgentInput{ + AssistantID: req.AssistantId, + Messages: req.Messages, + Options: req.Options, + AuthInfo: auth.GetAuthorizedInfo(ctx), + Cache: cache, + Writer: &grpcStreamWriter{stream: stream, header: make(http.Header)}, + }) + if err != nil { + return toGRPCError(err) + } + defer agentCtx.Release() + + ast, err := assistant.Get(agentCtx.AssistantID) + if err != nil { + return status.Errorf(codes.NotFound, "assistant not found: %v", err) + } + + _, err = ast.Stream(agentCtx, messages, opts) + if err != nil { + return status.Errorf(codes.Internal, "agent stream failed: %v", err) + } + + return stream.Send(&pb.AgentChunk{Done: true}) +} + +func toGRPCError(err error) error { + msg := err.Error() + if strings.Contains(msg, "is required") || + strings.Contains(msg, "must not be empty") || + strings.Contains(msg, "invalid") { + return status.Error(codes.InvalidArgument, msg) + } + return status.Error(codes.Internal, msg) +} + +// grpcStreamWriter bridges agent/context.Writer (http.ResponseWriter) to gRPC stream. +type grpcStreamWriter struct { + stream grpc.ServerStreamingServer[pb.AgentChunk] + header http.Header + code int +} + +func (w *grpcStreamWriter) Header() http.Header { return w.header } +func (w *grpcStreamWriter) WriteHeader(statusCode int) { w.code = statusCode } +func (w *grpcStreamWriter) Write(data []byte) (int, error) { + if err := w.stream.Send(&pb.AgentChunk{Data: data}); err != nil { + return 0, err + } + return len(data), nil +} + +// Flush implements http.Flusher for streaming compatibility. +func (w *grpcStreamWriter) Flush() {} diff --git a/grpc/agent/agent_test.go b/grpc/agent/agent_test.go new file mode 100644 index 00000000..d06edeea --- /dev/null +++ b/grpc/agent/agent_test.go @@ -0,0 +1,206 @@ +package agent_test + +import ( + "context" + "encoding/json" + "io" + "os" + "testing" + + "github.com/stretchr/testify/assert" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/yaoapp/yao/grpc/pb" + "github.com/yaoapp/yao/grpc/tests/testutils" +) + +func TestAgentStream_InvalidAssistant(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:agent") + ctx := testutils.WithToken(context.Background(), token) + + msgs, _ := json.Marshal([]map[string]interface{}{ + {"role": "user", "content": "hello"}, + }) + + stream, err := client.AgentStream(ctx, &pb.AgentRequest{ + AssistantId: "nonexistent-assistant-id", + Messages: msgs, + }) + if err != nil { + st, _ := status.FromError(err) + assert.Equal(t, codes.NotFound, st.Code()) + return + } + _, err = stream.Recv() + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.NotFound, st.Code()) +} + +func TestAgentStream_EmptyAssistantID(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:agent") + ctx := testutils.WithToken(context.Background(), token) + + stream, err := client.AgentStream(ctx, &pb.AgentRequest{ + AssistantId: "", + }) + if err != nil { + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) + return + } + _, err = stream.Recv() + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestAgentStream_EmptyMessages(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:agent") + ctx := testutils.WithToken(context.Background(), token) + + msgs, _ := json.Marshal([]map[string]interface{}{}) + + stream, err := client.AgentStream(ctx, &pb.AgentRequest{ + AssistantId: "some-assistant", + Messages: msgs, + }) + if err != nil { + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) + return + } + _, err = stream.Recv() + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestAgentStream_NilMessages(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:agent") + ctx := testutils.WithToken(context.Background(), token) + + stream, err := client.AgentStream(ctx, &pb.AgentRequest{ + AssistantId: "some-assistant", + Messages: nil, + }) + if err != nil { + st, _ := status.FromError(err) + assert.NotEqual(t, codes.OK, st.Code()) + return + } + _, err = stream.Recv() + assert.Error(t, err) + st, _ := status.FromError(err) + assert.NotEqual(t, codes.OK, st.Code()) +} + +func TestAgentStream_BadMessagesJSON(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:agent") + ctx := testutils.WithToken(context.Background(), token) + + stream, err := client.AgentStream(ctx, &pb.AgentRequest{ + AssistantId: "some-assistant", + Messages: []byte("{bad-json"), + }) + if err != nil { + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) + return + } + _, err = stream.Recv() + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestAgentStream_BadOptionsJSON(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:agent") + ctx := testutils.WithToken(context.Background(), token) + + msgs, _ := json.Marshal([]map[string]interface{}{ + {"role": "user", "content": "hello"}, + }) + + stream, err := client.AgentStream(ctx, &pb.AgentRequest{ + AssistantId: "some-assistant", + Messages: msgs, + Options: []byte("{bad-options"), + }) + if err != nil { + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) + return + } + _, err = stream.Recv() + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestAgentStream_RealAgent(t *testing.T) { + if os.Getenv("OPENAI_TEST_KEY") == "" { + t.Skip("OPENAI_TEST_KEY not set, skipping real agent test") + } + + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:agent") + ctx := testutils.WithToken(context.Background(), token) + + msgs, _ := json.Marshal([]map[string]interface{}{ + {"role": "user", "content": "Say hello in one word."}, + }) + + stream, err := client.AgentStream(ctx, &pb.AgentRequest{ + AssistantId: "tests.nested.demo", + Messages: msgs, + }) + if !assert.NoError(t, err) { + return + } + + var chunks int + for { + chunk, err := stream.Recv() + if err == io.EOF { + break + } + if !assert.NoError(t, err) { + break + } + chunks++ + if chunk.Done { + break + } + assert.NotEmpty(t, chunk.Data) + } + assert.Greater(t, chunks, 0) +} diff --git a/grpc/api/api.go b/grpc/api/api.go new file mode 100644 index 00000000..fcf5d91f --- /dev/null +++ b/grpc/api/api.go @@ -0,0 +1,69 @@ +package api + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + + "github.com/yaoapp/yao/grpc/pb" + "github.com/yaoapp/yao/service" +) + +// Handler implements the API gRPC method (internal HTTP proxy). +type Handler struct{} + +// API proxies a gRPC request to the internal openapi HTTP router. +func (h *Handler) API(ctx context.Context, req *pb.APIRequest) (*pb.APIResponse, error) { + router := service.Router + if router == nil { + return nil, status.Error(codes.Unavailable, "HTTP router not initialized") + } + + if req.Method == "" { + return nil, status.Error(codes.InvalidArgument, "method is required") + } + if req.Path == "" { + return nil, status.Error(codes.InvalidArgument, "path is required") + } + + httpReq, err := http.NewRequestWithContext(ctx, req.Method, req.Path, bytes.NewReader(req.Body)) + if err != nil { + return nil, status.Errorf(codes.Internal, "failed to build HTTP request: %v", err) + } + + for k, v := range req.Headers { + httpReq.Header.Set(k, v) + } + + // Forward Bearer token from gRPC metadata to HTTP Authorization header + // when the caller didn't explicitly set it. + if httpReq.Header.Get("Authorization") == "" { + if md, ok := metadata.FromIncomingContext(ctx); ok { + if vals := md.Get("authorization"); len(vals) > 0 { + httpReq.Header.Set("Authorization", vals[0]) + } + } + } + + w := httptest.NewRecorder() + router.ServeHTTP(w, httpReq) + + result := w.Result() + defer result.Body.Close() + + respHeaders := make(map[string]string, len(result.Header)) + for k := range result.Header { + respHeaders[k] = result.Header.Get(k) + } + + return &pb.APIResponse{ + Status: int32(result.StatusCode), + Headers: respHeaders, + Body: w.Body.Bytes(), + }, nil +} diff --git a/grpc/api/api_test.go b/grpc/api/api_test.go new file mode 100644 index 00000000..cd03df17 --- /dev/null +++ b/grpc/api/api_test.go @@ -0,0 +1,135 @@ +package api_test + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/yaoapp/yao/grpc/pb" + "github.com/yaoapp/yao/grpc/tests/testutils" +) + +func TestAPI_Proxy(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + + // The API method's ACL check uses the actual openapi path, so we grant all gRPC scopes. + // The openapi guard inside the HTTP router handles further auth via the forwarded Authorization header. + token := testutils.ObtainAccessToken(t, + "grpc:run", "grpc:stream", "grpc:shell", "grpc:mcp", "grpc:llm", "grpc:agent", + ) + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.API(ctx, &pb.APIRequest{ + Method: "GET", + Path: "/api/__yao/app/setting", + }) + + // The proxy itself should succeed (no gRPC error), even if the HTTP response + // is a non-200 status (e.g. 401 from openapi's own guard). + assert.NoError(t, err) + if assert.NotNil(t, resp) { + assert.Greater(t, resp.Status, int32(0)) + assert.NotNil(t, resp.Body) + } +} + +func TestAPI_NotFoundEndpoint(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, + "grpc:run", "grpc:stream", "grpc:shell", "grpc:mcp", "grpc:llm", "grpc:agent", + ) + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.API(ctx, &pb.APIRequest{ + Method: "GET", + Path: "/api/this/does/not/exist", + }) + assert.NoError(t, err) + if assert.NotNil(t, resp) { + assert.Equal(t, int32(404), resp.Status) + } +} + +func TestAPI_MissingMethod(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.API(ctx, &pb.APIRequest{ + Method: "", + Path: "/api/test", + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestAPI_MissingPath(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.API(ctx, &pb.APIRequest{ + Method: "GET", + Path: "", + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestAPI_WithHeaders(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, + "grpc:run", "grpc:stream", "grpc:shell", "grpc:mcp", "grpc:llm", "grpc:agent", + ) + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.API(ctx, &pb.APIRequest{ + Method: "GET", + Path: "/api/__yao/app/setting", + Headers: map[string]string{"X-Custom-Header": "test-value"}, + }) + assert.NoError(t, err) + assert.NotNil(t, resp) + assert.Greater(t, resp.Status, int32(0)) +} + +func TestAPI_PostWithBody(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, + "grpc:run", "grpc:stream", "grpc:shell", "grpc:mcp", "grpc:llm", "grpc:agent", + ) + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.API(ctx, &pb.APIRequest{ + Method: "POST", + Path: "/api/this/does/not/exist", + Body: []byte(`{"key":"value"}`), + }) + assert.NoError(t, err) + if assert.NotNil(t, resp) { + assert.Equal(t, int32(404), resp.Status) + } +} diff --git a/grpc/auth/endpoint.go b/grpc/auth/endpoint.go new file mode 100644 index 00000000..67fcaf6e --- /dev/null +++ b/grpc/auth/endpoint.go @@ -0,0 +1,66 @@ +package auth + +import ( + "fmt" + "strings" + + "github.com/yaoapp/yao/grpc/pb" +) + +// VirtualEndpoint maps a gRPC full method + request to a virtual HTTP endpoint for ACL. +// Returns the HTTP method and path used for scope-based access control. +func VirtualEndpoint(fullMethod string, req interface{}) (method string, path string) { + switch fullMethod { + case "/yao.Yao/Run": + if r, ok := req.(*pb.RunRequest); ok && r.Process != "" { + return "POST", "/grpc/run/" + r.Process + } + return "POST", "/grpc/run/" + + case "/yao.Yao/Stream": + if r, ok := req.(*pb.RunRequest); ok && r.Process != "" { + return "POST", "/grpc/stream/" + r.Process + } + return "POST", "/grpc/stream/" + + case "/yao.Yao/Shell", "/yao.Yao/ShellStream": + return "POST", "/grpc/shell" + + case "/yao.Yao/API": + if r, ok := req.(*pb.APIRequest); ok && r.Path != "" { + m := strings.ToUpper(r.Method) + if m == "" { + m = "POST" + } + return m, r.Path + } + return "POST", "/" + + case "/yao.Yao/MCPListTools": + return "GET", "/grpc/mcp/tools" + + case "/yao.Yao/MCPCallTool": + if r, ok := req.(*pb.MCPCallRequest); ok && r.Tool != "" { + return "POST", "/grpc/mcp/call/" + r.Tool + } + return "POST", "/grpc/mcp/call/" + + case "/yao.Yao/MCPListResources": + return "GET", "/grpc/mcp/resources" + + case "/yao.Yao/MCPReadResource": + return "GET", "/grpc/mcp/resources/read" + + case "/yao.Yao/ChatCompletions", "/yao.Yao/ChatCompletionsStream": + return "POST", "/grpc/llm/completions" + + case "/yao.Yao/AgentStream": + if r, ok := req.(*pb.AgentRequest); ok && r.AssistantId != "" { + return "POST", fmt.Sprintf("/grpc/agent/%s", r.AssistantId) + } + return "POST", "/grpc/agent/" + + default: + return "POST", "/grpc/unknown" + } +} diff --git a/grpc/auth/endpoint_test.go b/grpc/auth/endpoint_test.go new file mode 100644 index 00000000..edbe3c1e --- /dev/null +++ b/grpc/auth/endpoint_test.go @@ -0,0 +1,136 @@ +package auth_test + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/yaoapp/yao/grpc/auth" + "github.com/yaoapp/yao/grpc/pb" +) + +func TestVirtualEndpoint_Run(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/Run", &pb.RunRequest{Process: "models.user.Find"}) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/run/models.user.Find", path) +} + +func TestVirtualEndpoint_Stream(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/Stream", &pb.RunRequest{Process: "flows.report"}) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/stream/flows.report", path) +} + +func TestVirtualEndpoint_Shell(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/Shell", &pb.ShellRequest{Command: "ls"}) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/shell", path) +} + +func TestVirtualEndpoint_ShellStream(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/ShellStream", &pb.ShellRequest{Command: "ls"}) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/shell", path) +} + +func TestVirtualEndpoint_API(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/API", &pb.APIRequest{Method: "GET", Path: "/kb/collections"}) + assert.Equal(t, "GET", method) + assert.Equal(t, "/kb/collections", path) +} + +func TestVirtualEndpoint_MCPListTools(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/MCPListTools", &pb.MCPListRequest{SessionId: "abc"}) + assert.Equal(t, "GET", method) + assert.Equal(t, "/grpc/mcp/tools", path) +} + +func TestVirtualEndpoint_MCPCallTool(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/MCPCallTool", &pb.MCPCallRequest{Tool: "search"}) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/mcp/call/search", path) +} + +func TestVirtualEndpoint_MCPListResources(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/MCPListResources", &pb.MCPListRequest{}) + assert.Equal(t, "GET", method) + assert.Equal(t, "/grpc/mcp/resources", path) +} + +func TestVirtualEndpoint_MCPReadResource(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/MCPReadResource", &pb.MCPResourceRequest{Uri: "file://test"}) + assert.Equal(t, "GET", method) + assert.Equal(t, "/grpc/mcp/resources/read", path) +} + +func TestVirtualEndpoint_ChatCompletions(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/ChatCompletions", &pb.ChatRequest{Connector: "openai"}) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/llm/completions", path) +} + +func TestVirtualEndpoint_ChatCompletionsStream(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/ChatCompletionsStream", &pb.ChatRequest{}) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/llm/completions", path) +} + +func TestVirtualEndpoint_AgentStream(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/AgentStream", &pb.AgentRequest{AssistantId: "my-robot"}) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/agent/my-robot", path) +} + +func TestVirtualEndpoint_Unknown(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/NonExistent", nil) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/unknown", path) +} + +func TestVirtualEndpoint_RunNilReq(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/Run", nil) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/run/", path) +} + +func TestVirtualEndpoint_RunEmptyProcess(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/Run", &pb.RunRequest{Process: ""}) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/run/", path) +} + +func TestVirtualEndpoint_StreamNilReq(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/Stream", nil) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/stream/", path) +} + +func TestVirtualEndpoint_APINilReq(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/API", nil) + assert.Equal(t, "POST", method) + assert.Equal(t, "/", path) +} + +func TestVirtualEndpoint_APIEmptyMethod(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/API", &pb.APIRequest{Method: "", Path: "/test"}) + assert.Equal(t, "POST", method) + assert.Equal(t, "/test", path) +} + +func TestVirtualEndpoint_MCPCallToolNilReq(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/MCPCallTool", nil) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/mcp/call/", path) +} + +func TestVirtualEndpoint_AgentStreamNilReq(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/AgentStream", nil) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/agent/", path) +} + +func TestVirtualEndpoint_AgentStreamEmptyID(t *testing.T) { + method, path := auth.VirtualEndpoint("/yao.Yao/AgentStream", &pb.AgentRequest{AssistantId: ""}) + assert.Equal(t, "POST", method) + assert.Equal(t, "/grpc/agent/", path) +} diff --git a/grpc/auth/guard.go b/grpc/auth/guard.go new file mode 100644 index 00000000..c2abccf3 --- /dev/null +++ b/grpc/auth/guard.go @@ -0,0 +1,169 @@ +package auth + +import ( + "context" + "strings" + + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + + "github.com/yaoapp/yao/openapi/oauth" + "github.com/yaoapp/yao/openapi/oauth/acl" + "github.com/yaoapp/yao/openapi/oauth/types" +) + +const ( + healthzMethod = "/yao.Yao/Healthz" + apiMethod = "/yao.Yao/API" + + metaAuthorization = "authorization" + metaRefreshToken = "x-refresh-token" + metaAccessToken = "x-access-token" + metaSandboxID = "x-sandbox-id" + metaSessionID = "x-session-id" +) + +type authCtxKey struct{} + +// WithAuthorizedInfo stores AuthorizedInfo in context for downstream handlers. +func WithAuthorizedInfo(ctx context.Context, info *types.AuthorizedInfo) context.Context { + return context.WithValue(ctx, authCtxKey{}, info) +} + +// GetAuthorizedInfo retrieves AuthorizedInfo from context (set by the interceptor). +func GetAuthorizedInfo(ctx context.Context) *types.AuthorizedInfo { + info, _ := ctx.Value(authCtxKey{}).(*types.AuthorizedInfo) + return info +} + +// UnaryInterceptor is the gRPC unary server interceptor for authentication and authorization. +func UnaryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) { + if info.FullMethod == healthzMethod { + return handler(ctx, req) + } + + ctx, err := authenticate(ctx, info.FullMethod, req) + if err != nil { + return nil, err + } + return handler(ctx, req) +} + +// StreamInterceptor is the gRPC stream server interceptor for authentication and authorization. +// For streaming RPCs, the request object is not available at intercept time, +// so ACL scope check uses the method-level virtual path (without request-specific IDs). +func StreamInterceptor(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { + if info.FullMethod == healthzMethod { + return handler(srv, ss) + } + + ctx, err := authenticate(ss.Context(), info.FullMethod, nil) + if err != nil { + return err + } + + return handler(srv, &wrappedStream{ServerStream: ss, ctx: ctx}) +} + +// authenticate calls oauth.Service.AuthenticateToken directly — no gin/HTTP shim. +func authenticate(ctx context.Context, fullMethod string, req interface{}) (context.Context, error) { + md, ok := metadata.FromIncomingContext(ctx) + if !ok { + return ctx, status.Error(codes.Unauthenticated, "missing metadata") + } + + svc := oauth.OAuth + if svc == nil { + return ctx, status.Error(codes.Internal, "oauth service not initialized") + } + + bearer := extractBearer(md) + if bearer == "" { + return ctx, status.Error(codes.Unauthenticated, "missing authorization token") + } + + result, err := svc.AuthenticateToken(oauth.AuthInput{ + AccessToken: bearer, + RefreshToken: extractMeta(md, metaRefreshToken), + SessionID: extractMeta(md, metaSessionID), + }) + if err != nil { + return ctx, status.Error(codes.Unauthenticated, err.Error()) + } + + ctx = WithAuthorizedInfo(ctx, result.Info) + + if result.NewAccessToken != "" { + _ = grpc.SendHeader(ctx, metadata.Pairs( + metaAccessToken, result.NewAccessToken, + metaRefreshToken, result.NewRefreshToken, + )) + } + + // ACL scope check — skip for API proxy (the openapi router does its own auth). + if fullMethod != apiMethod { + httpMethod, httpPath := VirtualEndpoint(fullMethod, req) + scopes := strings.Fields(result.Info.Scope) + + enforcer := getACLEnforcer() + if enforcer != nil && enforcer.Scope != nil { + decision := enforcer.Scope.Check(&acl.AccessRequest{ + Method: httpMethod, + Path: httpPath, + Scopes: scopes, + }) + if !decision.Allowed { + return ctx, status.Errorf(codes.PermissionDenied, "insufficient scope: %s", decision.Reason) + } + } + } + + return ctx, nil +} + +// getACLEnforcer returns the ACL enforcer if available and enabled. +func getACLEnforcer() *acl.ACL { + if acl.Global == nil { + return nil + } + enforcer, ok := acl.Global.(*acl.ACL) + if !ok || enforcer == nil { + return nil + } + if !enforcer.Config.Enabled { + return nil + } + return enforcer +} + +func extractBearer(md metadata.MD) string { + vals := md.Get(metaAuthorization) + if len(vals) == 0 { + return "" + } + parts := strings.SplitN(vals[0], " ", 2) + if len(parts) == 2 && strings.EqualFold(parts[0], "bearer") { + return parts[1] + } + return vals[0] +} + +func extractMeta(md metadata.MD, key string) string { + vals := md.Get(key) + if len(vals) == 0 { + return "" + } + return vals[0] +} + +// wrappedStream wraps grpc.ServerStream with a custom context. +type wrappedStream struct { + grpc.ServerStream + ctx context.Context +} + +func (w *wrappedStream) Context() context.Context { + return w.ctx +} diff --git a/grpc/auth/guard_test.go b/grpc/auth/guard_test.go new file mode 100644 index 00000000..093d03d6 --- /dev/null +++ b/grpc/auth/guard_test.go @@ -0,0 +1,169 @@ +package auth_test + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" + + "github.com/yaoapp/yao/grpc/pb" + "github.com/yaoapp/yao/grpc/tests/testutils" +) + +func TestAuth_NoToken_Rejected(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + _, err := client.Run(context.Background(), &pb.RunRequest{Process: "utils.app.Ping"}) + assert.Error(t, err) + + st, ok := status.FromError(err) + assert.True(t, ok) + assert.Equal(t, codes.Unauthenticated, st.Code()) +} + +func TestAuth_ValidToken(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.Run(ctx, &pb.RunRequest{Process: "utils.app.Ping"}) + // Run returns Unimplemented (handler stub), not an auth error + st, ok := status.FromError(err) + assert.True(t, ok) + assert.NotEqual(t, codes.Unauthenticated, st.Code()) + assert.NotEqual(t, codes.PermissionDenied, st.Code()) +} + +func TestAuth_WrongScope_Denied(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.Run(ctx, &pb.RunRequest{Process: "utils.app.Ping"}) + assert.Error(t, err) + + st, ok := status.FromError(err) + assert.True(t, ok) + assert.Equal(t, codes.PermissionDenied, st.Code()) +} + +func TestAuth_TokenRefresh(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + + expiredToken := testutils.ObtainExpiredAccessToken(t, "grpc:run") + refreshToken := testutils.ObtainRefreshToken(t, "grpc:run") + ctx := testutils.WithRefreshToken(context.Background(), expiredToken, refreshToken) + + // The call should succeed (auth interceptor refreshes the token) + _, err := client.Run(ctx, &pb.RunRequest{Process: "utils.app.Ping"}) + st, ok := status.FromError(err) + assert.True(t, ok) + // Should not be an auth error — either Unimplemented (handler stub) or OK + assert.NotEqual(t, codes.Unauthenticated, st.Code()) + assert.NotEqual(t, codes.PermissionDenied, st.Code()) +} + +func TestHealthz_Public(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + resp, err := client.Healthz(context.Background(), &pb.Empty{}) + assert.NoError(t, err) + assert.NotNil(t, resp) + assert.Equal(t, "ok", resp.Status) +} + +func TestAuth_InvalidBearerFormat(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + ctx := metadata.AppendToOutgoingContext(context.Background(), "authorization", "not-a-valid-token-at-all") + + _, err := client.Run(ctx, &pb.RunRequest{Process: "utils.app.Ping"}) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.Unauthenticated, st.Code()) +} + +func TestAuth_StreamInterceptor_NoToken(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + stream, err := client.ChatCompletionsStream(context.Background(), &pb.ChatRequest{ + Connector: "openai", + }) + if err != nil { + st, _ := status.FromError(err) + assert.Equal(t, codes.Unauthenticated, st.Code()) + return + } + _, err = stream.Recv() + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.Unauthenticated, st.Code()) +} + +func TestAuth_StreamInterceptor_ValidToken(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:agent") + ctx := testutils.WithToken(context.Background(), token) + + stream, err := client.AgentStream(ctx, &pb.AgentRequest{ + AssistantId: "nonexistent", + Messages: []byte(`[{"role":"user","content":"hi"}]`), + }) + if err != nil { + st, _ := status.FromError(err) + assert.NotEqual(t, codes.Unauthenticated, st.Code()) + assert.NotEqual(t, codes.PermissionDenied, st.Code()) + return + } + _, err = stream.Recv() + assert.Error(t, err) + st, _ := status.FromError(err) + assert.NotEqual(t, codes.Unauthenticated, st.Code()) + assert.NotEqual(t, codes.PermissionDenied, st.Code()) +} + +func TestAuth_StreamInterceptor_WrongScope(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + stream, err := client.AgentStream(ctx, &pb.AgentRequest{ + AssistantId: "test", + Messages: []byte(`[{"role":"user","content":"hi"}]`), + }) + if err != nil { + st, _ := status.FromError(err) + assert.Equal(t, codes.PermissionDenied, st.Code()) + return + } + _, err = stream.Recv() + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.PermissionDenied, st.Code()) +} diff --git a/grpc/auth/scope.go b/grpc/auth/scope.go new file mode 100644 index 00000000..4957c482 --- /dev/null +++ b/grpc/auth/scope.go @@ -0,0 +1,14 @@ +package auth + +import "github.com/yaoapp/yao/openapi/oauth/acl" + +func init() { + acl.Register( + &acl.ScopeDefinition{Name: "grpc:run", Endpoints: []string{"POST /grpc/run/*", "POST /grpc/run/"}}, + &acl.ScopeDefinition{Name: "grpc:stream", Endpoints: []string{"POST /grpc/stream/*", "POST /grpc/stream/"}}, + &acl.ScopeDefinition{Name: "grpc:shell", Endpoints: []string{"POST /grpc/shell"}}, + &acl.ScopeDefinition{Name: "grpc:mcp", Endpoints: []string{"GET /grpc/mcp/tools", "POST /grpc/mcp/call/*", "POST /grpc/mcp/call/", "GET /grpc/mcp/resources", "GET /grpc/mcp/resources/read"}}, + &acl.ScopeDefinition{Name: "grpc:llm", Endpoints: []string{"POST /grpc/llm/completions"}}, + &acl.ScopeDefinition{Name: "grpc:agent", Endpoints: []string{"POST /grpc/agent/*", "POST /grpc/agent/"}}, + ) +} diff --git a/grpc/grpc.go b/grpc/grpc.go new file mode 100644 index 00000000..22755275 --- /dev/null +++ b/grpc/grpc.go @@ -0,0 +1,173 @@ +package grpc + +import ( + "context" + "net" + "strconv" + "strings" + "sync" + + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/config" + agenthandler "github.com/yaoapp/yao/grpc/agent" + apihandler "github.com/yaoapp/yao/grpc/api" + "github.com/yaoapp/yao/grpc/auth" + "github.com/yaoapp/yao/grpc/health" + llmhandler "github.com/yaoapp/yao/grpc/llm" + mcphandler "github.com/yaoapp/yao/grpc/mcp" + "github.com/yaoapp/yao/grpc/pb" + runhandler "github.com/yaoapp/yao/grpc/run" + shellhandler "github.com/yaoapp/yao/grpc/shell" +) + +var ( + mu sync.Mutex + server *grpc.Server + listeners []net.Listener + addrs []string +) + +type yaoServer struct { + pb.UnimplementedYaoServer + health health.Handler + run runhandler.Handler + shell shellhandler.Handler + api apihandler.Handler + mcp mcphandler.Handler + llm llmhandler.Handler + agent agenthandler.Handler +} + +// ── Health ─────────────────────────────────────────────────────────────────── + +func (s *yaoServer) Healthz(ctx context.Context, req *pb.Empty) (*pb.HealthzResponse, error) { + return s.health.Healthz(ctx, req) +} + +// ── Base ───────────────────────────────────────────────────────────────────── + +func (s *yaoServer) Run(ctx context.Context, req *pb.RunRequest) (*pb.RunResponse, error) { + return s.run.Run(ctx, req) +} + +func (s *yaoServer) Shell(ctx context.Context, req *pb.ShellRequest) (*pb.ShellResponse, error) { + return s.shell.Shell(ctx, req) +} + +// V2 stubs — Stream and ShellStream depend on gou/stream package. +func (s *yaoServer) Stream(req *pb.RunRequest, stream grpc.ServerStreamingServer[pb.Chunk]) error { + return status.Error(codes.Unimplemented, "Stream not implemented (V2)") +} + +func (s *yaoServer) ShellStream(req *pb.ShellRequest, stream grpc.ServerStreamingServer[pb.Chunk]) error { + return status.Error(codes.Unimplemented, "ShellStream not implemented (V2)") +} + +// ── API ────────────────────────────────────────────────────────────────────── + +func (s *yaoServer) API(ctx context.Context, req *pb.APIRequest) (*pb.APIResponse, error) { + return s.api.API(ctx, req) +} + +// ── MCP ────────────────────────────────────────────────────────────────────── + +func (s *yaoServer) MCPListTools(ctx context.Context, req *pb.MCPListRequest) (*pb.MCPListResponse, error) { + return s.mcp.MCPListTools(ctx, req) +} + +func (s *yaoServer) MCPCallTool(ctx context.Context, req *pb.MCPCallRequest) (*pb.MCPCallResponse, error) { + return s.mcp.MCPCallTool(ctx, req) +} + +func (s *yaoServer) MCPListResources(ctx context.Context, req *pb.MCPListRequest) (*pb.MCPResourcesResponse, error) { + return s.mcp.MCPListResources(ctx, req) +} + +func (s *yaoServer) MCPReadResource(ctx context.Context, req *pb.MCPResourceRequest) (*pb.MCPResourceResponse, error) { + return s.mcp.MCPReadResource(ctx, req) +} + +// ── LLM ────────────────────────────────────────────────────────────────────── + +func (s *yaoServer) ChatCompletions(ctx context.Context, req *pb.ChatRequest) (*pb.ChatResponse, error) { + return s.llm.ChatCompletions(ctx, req) +} + +func (s *yaoServer) ChatCompletionsStream(req *pb.ChatRequest, stream grpc.ServerStreamingServer[pb.ChatChunk]) error { + return s.llm.ChatCompletionsStream(req, stream) +} + +// ── Agent ──────────────────────────────────────────────────────────────────── + +func (s *yaoServer) AgentStream(req *pb.AgentRequest, stream grpc.ServerStreamingServer[pb.AgentChunk]) error { + return s.agent.AgentStream(req, stream) +} + +// ── Server lifecycle ───────────────────────────────────────────────────────── + +// StartServer initializes and starts the gRPC server based on config. +// It supports multiple bind addresses and returns immediately (listeners run in goroutines). +func StartServer(cfg config.Config) error { + if strings.ToLower(cfg.GRPC.Enabled) == "off" { + log.Info("gRPC server disabled (YAO_GRPC=off)") + return nil + } + + mu.Lock() + defer mu.Unlock() + + server = grpc.NewServer( + grpc.ChainUnaryInterceptor(auth.UnaryInterceptor), + grpc.ChainStreamInterceptor(auth.StreamInterceptor), + ) + pb.RegisterYaoServer(server, &yaoServer{}) + + hosts := strings.Split(cfg.GRPC.Host, ",") + port := strconv.Itoa(cfg.GRPC.Port) + + for _, h := range hosts { + addr := net.JoinHostPort(strings.TrimSpace(h), port) + lis, err := net.Listen("tcp", addr) + if err != nil { + Stop() + return err + } + listeners = append(listeners, lis) + addrs = append(addrs, lis.Addr().String()) + log.Info("gRPC server listening on %s", lis.Addr().String()) + + go func(l net.Listener) { + if err := server.Serve(l); err != nil { + log.Error("gRPC server error on %s: %s", l.Addr().String(), err.Error()) + } + }(lis) + } + + return nil +} + +// Stop gracefully stops the gRPC server. Safe to call if server was never started. +func Stop() { + mu.Lock() + defer mu.Unlock() + + if server != nil { + server.GracefulStop() + server = nil + } + listeners = nil + addrs = nil +} + +// Addr returns all addresses the gRPC server is listening on. +func Addr() []string { + mu.Lock() + defer mu.Unlock() + result := make([]string, len(addrs)) + copy(result, addrs) + return result +} diff --git a/grpc/health/health.go b/grpc/health/health.go new file mode 100644 index 00000000..fadf1ce6 --- /dev/null +++ b/grpc/health/health.go @@ -0,0 +1,15 @@ +package health + +import ( + "context" + + "github.com/yaoapp/yao/grpc/pb" +) + +// Handler implements the Healthz RPC. +type Handler struct{} + +// Healthz returns server health status. This method is public (no auth required). +func (h *Handler) Healthz(ctx context.Context, req *pb.Empty) (*pb.HealthzResponse, error) { + return &pb.HealthzResponse{Status: "ok"}, nil +} diff --git a/grpc/health/health_test.go b/grpc/health/health_test.go new file mode 100644 index 00000000..740e83d9 --- /dev/null +++ b/grpc/health/health_test.go @@ -0,0 +1,22 @@ +package health_test + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/yaoapp/yao/grpc/pb" + "github.com/yaoapp/yao/grpc/tests/testutils" +) + +func TestHealthz_ReturnsOk(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + resp, err := client.Healthz(context.Background(), &pb.Empty{}) + assert.NoError(t, err) + assert.NotNil(t, resp) + assert.Equal(t, "ok", resp.Status) +} diff --git a/grpc/llm/llm.go b/grpc/llm/llm.go new file mode 100644 index 00000000..93461b59 --- /dev/null +++ b/grpc/llm/llm.go @@ -0,0 +1,165 @@ +package llm + +import ( + "context" + "encoding/json" + + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/yaoapp/gou/connector" + agentContext "github.com/yaoapp/yao/agent/context" + agentLLM "github.com/yaoapp/yao/agent/llm" + "github.com/yaoapp/yao/agent/output/message" + "github.com/yaoapp/yao/grpc/auth" + "github.com/yaoapp/yao/grpc/pb" +) + +// Handler implements the LLM gRPC methods. +type Handler struct{} + +// ChatCompletions sends messages to an LLM connector and returns the full response (unary). +func (h *Handler) ChatCompletions(ctx context.Context, req *pb.ChatRequest) (*pb.ChatResponse, error) { + if req.Connector == "" { + return nil, status.Error(codes.InvalidArgument, "connector is required") + } + + llmInstance, completionOpts, ctxMessages, agentCtx, err := prepareLLMCall(ctx, req) + if err != nil { + return nil, err + } + defer agentCtx.Release() + + noopHandler := func(chunkType message.StreamChunkType, data []byte) int { return 0 } + response, err := llmInstance.Stream(agentCtx, ctxMessages, completionOpts, noopHandler) + if err != nil { + return nil, status.Errorf(codes.Internal, "LLM call failed: %v", err) + } + + data, err := json.Marshal(toOpenAIFormat(response)) + if err != nil { + return nil, status.Errorf(codes.Internal, "failed to marshal LLM response: %v", err) + } + + return &pb.ChatResponse{Data: data}, nil +} + +// ChatCompletionsStream sends messages to an LLM connector and streams response chunks. +func (h *Handler) ChatCompletionsStream(req *pb.ChatRequest, stream grpc.ServerStreamingServer[pb.ChatChunk]) error { + ctx := stream.Context() + + if req.Connector == "" { + return status.Error(codes.InvalidArgument, "connector is required") + } + + llmInstance, completionOpts, ctxMessages, agentCtx, err := prepareLLMCall(ctx, req) + if err != nil { + return err + } + defer agentCtx.Release() + + streamHandler := func(chunkType message.StreamChunkType, data []byte) int { + if ctx.Err() != nil { + return 1 + } + if chunkType == message.ChunkText || chunkType == message.ChunkThinking { + if sendErr := stream.Send(&pb.ChatChunk{Data: data}); sendErr != nil { + return 1 + } + } + return 0 + } + + _, err = llmInstance.Stream(agentCtx, ctxMessages, completionOpts, streamHandler) + if err != nil { + return status.Errorf(codes.Internal, "LLM stream failed: %v", err) + } + + return stream.Send(&pb.ChatChunk{Done: true}) +} + +// prepareLLMCall builds the LLM instance, messages, and agent context from the gRPC request. +// Mirrors agent/llm/process.go ProcessChatCompletions logic without the process wrapper. +func prepareLLMCall(ctx context.Context, req *pb.ChatRequest) (agentLLM.LLM, *agentContext.CompletionOptions, []agentContext.Message, *agentContext.Context, error) { + ctxMessages, err := parseMessagesToContext(req.Messages) + if err != nil { + return nil, nil, nil, nil, err + } + + var opts map[string]interface{} + if len(req.Options) > 0 { + if err := json.Unmarshal(req.Options, &opts); err != nil { + return nil, nil, nil, nil, status.Errorf(codes.InvalidArgument, "invalid options JSON: %v", err) + } + } + + conn, err := connector.Select(req.Connector) + if err != nil { + return nil, nil, nil, nil, status.Errorf(codes.NotFound, "connector %s not found: %v", req.Connector, err) + } + + completionOpts := agentLLM.BuildCompletionOptions(conn, opts) + + llmInstance, err := agentLLM.New(conn, completionOpts) + if err != nil { + return nil, nil, nil, nil, status.Errorf(codes.Internal, "failed to create LLM: %v", err) + } + + authInfo := auth.GetAuthorizedInfo(ctx) + chatID := agentContext.GenChatID() + agentCtx := agentContext.New(ctx, authInfo, chatID) + + return llmInstance, completionOpts, ctxMessages, agentCtx, nil +} + +// parseMessagesToContext converts raw JSON message bytes to []agentContext.Message via JSON round-trip. +func parseMessagesToContext(raw []byte) ([]agentContext.Message, error) { + if len(raw) == 0 { + return nil, status.Error(codes.InvalidArgument, "messages are required") + } + + var messages []agentContext.Message + if err := json.Unmarshal(raw, &messages); err != nil { + return nil, status.Errorf(codes.InvalidArgument, "invalid messages JSON: %v", err) + } + if len(messages) == 0 { + return nil, status.Error(codes.InvalidArgument, "messages must not be empty") + } + + return messages, nil +} + +// toOpenAIFormat converts CompletionResponse to OpenAI chat.completions format. +func toOpenAIFormat(resp *agentContext.CompletionResponse) map[string]interface{} { + if resp == nil { + return map[string]interface{}{"choices": []interface{}{}} + } + + msgMap := map[string]interface{}{ + "role": resp.Role, + "content": resp.Content, + } + if len(resp.ToolCalls) > 0 { + msgMap["tool_calls"] = resp.ToolCalls + } + + choice := map[string]interface{}{ + "index": 0, + "message": msgMap, + "finish_reason": "stop", + } + + result := map[string]interface{}{ + "id": resp.ID, + "object": "chat.completion", + "created": resp.Created, + "model": resp.Model, + "choices": []interface{}{choice}, + } + if resp.Usage != nil { + result["usage"] = resp.Usage + } + + return result +} diff --git a/grpc/llm/llm_test.go b/grpc/llm/llm_test.go new file mode 100644 index 00000000..0106adc6 --- /dev/null +++ b/grpc/llm/llm_test.go @@ -0,0 +1,265 @@ +package llm_test + +import ( + "context" + "encoding/json" + "io" + "os" + "testing" + + "github.com/stretchr/testify/assert" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/yaoapp/yao/grpc/pb" + "github.com/yaoapp/yao/grpc/tests/testutils" +) + +func TestChatCompletions_InvalidConnector(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:llm") + ctx := testutils.WithToken(context.Background(), token) + + msgs, _ := json.Marshal([]map[string]interface{}{ + {"role": "user", "content": "hello"}, + }) + + _, err := client.ChatCompletions(ctx, &pb.ChatRequest{ + Connector: "nonexistent-connector", + Messages: msgs, + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.NotFound, st.Code()) +} + +func TestChatCompletions_EmptyConnector(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:llm") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.ChatCompletions(ctx, &pb.ChatRequest{ + Connector: "", + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestChatCompletionsStream_EmptyConnector(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:llm") + ctx := testutils.WithToken(context.Background(), token) + + stream, err := client.ChatCompletionsStream(ctx, &pb.ChatRequest{ + Connector: "", + }) + if err != nil { + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) + return + } + _, err = stream.Recv() + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestChatCompletions_BadMessagesJSON(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:llm") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.ChatCompletions(ctx, &pb.ChatRequest{ + Connector: "openai", + Messages: []byte("{bad-json"), + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestChatCompletions_EmptyMessages(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:llm") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.ChatCompletions(ctx, &pb.ChatRequest{ + Connector: "openai", + Messages: nil, + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestChatCompletions_EmptyMessageArray(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:llm") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.ChatCompletions(ctx, &pb.ChatRequest{ + Connector: "openai", + Messages: []byte("[]"), + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestChatCompletions_BadOptionsJSON(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:llm") + ctx := testutils.WithToken(context.Background(), token) + + msgs, _ := json.Marshal([]map[string]interface{}{ + {"role": "user", "content": "hello"}, + }) + _, err := client.ChatCompletions(ctx, &pb.ChatRequest{ + Connector: "openai", + Messages: msgs, + Options: []byte("{bad-options"), + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestChatCompletionsStream_InvalidConnector(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:llm") + ctx := testutils.WithToken(context.Background(), token) + + msgs, _ := json.Marshal([]map[string]interface{}{ + {"role": "user", "content": "hello"}, + }) + + stream, err := client.ChatCompletionsStream(ctx, &pb.ChatRequest{ + Connector: "nonexistent-connector", + Messages: msgs, + }) + if err != nil { + st, _ := status.FromError(err) + assert.Equal(t, codes.NotFound, st.Code()) + return + } + _, err = stream.Recv() + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.NotFound, st.Code()) +} + +func TestChatCompletionsStream_BadMessages(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:llm") + ctx := testutils.WithToken(context.Background(), token) + + stream, err := client.ChatCompletionsStream(ctx, &pb.ChatRequest{ + Connector: "openai", + Messages: []byte("{bad-json"), + }) + if err != nil { + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) + return + } + _, err = stream.Recv() + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +// TestChatCompletions_RealLLM tests against a real LLM if OPENAI_TEST_KEY is set. +func TestChatCompletions_RealLLM(t *testing.T) { + if os.Getenv("OPENAI_TEST_KEY") == "" { + t.Skip("OPENAI_TEST_KEY not set, skipping real LLM test") + } + + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:llm") + ctx := testutils.WithToken(context.Background(), token) + + msgs, _ := json.Marshal([]map[string]interface{}{ + {"role": "user", "content": "Say hello in one word."}, + }) + + resp, err := client.ChatCompletions(ctx, &pb.ChatRequest{ + Connector: "gpt-4o-mini", + Messages: msgs, + }) + assert.NoError(t, err) + if assert.NotNil(t, resp) { + assert.NotEmpty(t, resp.Data) + } +} + +// TestChatCompletionsStream_RealLLM tests streaming against a real LLM if OPENAI_TEST_KEY is set. +func TestChatCompletionsStream_RealLLM(t *testing.T) { + if os.Getenv("OPENAI_TEST_KEY") == "" { + t.Skip("OPENAI_TEST_KEY not set, skipping real LLM stream test") + } + + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:llm") + ctx := testutils.WithToken(context.Background(), token) + + msgs, _ := json.Marshal([]map[string]interface{}{ + {"role": "user", "content": "Count from 1 to 3."}, + }) + + stream, err := client.ChatCompletionsStream(ctx, &pb.ChatRequest{ + Connector: "gpt-4o-mini", + Messages: msgs, + }) + assert.NoError(t, err) + + var chunks int + for { + chunk, err := stream.Recv() + if err == io.EOF { + break + } + if !assert.NoError(t, err) { + break + } + chunks++ + if chunk.Done { + break + } + assert.NotEmpty(t, chunk.Data) + } + assert.Greater(t, chunks, 0) +} diff --git a/grpc/mcp/mcp.go b/grpc/mcp/mcp.go new file mode 100644 index 00000000..62ac1a08 --- /dev/null +++ b/grpc/mcp/mcp.go @@ -0,0 +1,102 @@ +package mcp + +import ( + "context" + "encoding/json" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + goumcp "github.com/yaoapp/gou/mcp" + "github.com/yaoapp/yao/grpc/pb" +) + +// Handler implements the MCP gRPC methods. +type Handler struct{} + +// MCPListTools lists all available MCP tools for a given session. +func (h *Handler) MCPListTools(ctx context.Context, req *pb.MCPListRequest) (*pb.MCPListResponse, error) { + client, err := goumcp.Select(req.SessionId) + if err != nil { + return nil, status.Errorf(codes.NotFound, "MCP client not found: %v", err) + } + + resp, err := client.ListTools(ctx, "") + if err != nil { + return nil, status.Errorf(codes.Internal, "ListTools failed: %v", err) + } + + data, err := json.Marshal(resp.Tools) + if err != nil { + return nil, status.Errorf(codes.Internal, "failed to marshal tools: %v", err) + } + + return &pb.MCPListResponse{Tools: data}, nil +} + +// MCPCallTool calls an MCP tool by name with the provided arguments. +func (h *Handler) MCPCallTool(ctx context.Context, req *pb.MCPCallRequest) (*pb.MCPCallResponse, error) { + client, err := goumcp.Select(req.SessionId) + if err != nil { + return nil, status.Errorf(codes.NotFound, "MCP client not found: %v", err) + } + + var args interface{} + if len(req.Arguments) > 0 { + if err := json.Unmarshal(req.Arguments, &args); err != nil { + return nil, status.Errorf(codes.InvalidArgument, "invalid arguments JSON: %v", err) + } + } + + resp, err := client.CallTool(ctx, req.Tool, args) + if err != nil { + return nil, status.Errorf(codes.Internal, "CallTool failed: %v", err) + } + + data, err := json.Marshal(resp) + if err != nil { + return nil, status.Errorf(codes.Internal, "failed to marshal result: %v", err) + } + + return &pb.MCPCallResponse{Result: data}, nil +} + +// MCPListResources lists all available MCP resources for a given session. +func (h *Handler) MCPListResources(ctx context.Context, req *pb.MCPListRequest) (*pb.MCPResourcesResponse, error) { + client, err := goumcp.Select(req.SessionId) + if err != nil { + return nil, status.Errorf(codes.NotFound, "MCP client not found: %v", err) + } + + resp, err := client.ListResources(ctx, "") + if err != nil { + return nil, status.Errorf(codes.Internal, "ListResources failed: %v", err) + } + + data, err := json.Marshal(resp.Resources) + if err != nil { + return nil, status.Errorf(codes.Internal, "failed to marshal resources: %v", err) + } + + return &pb.MCPResourcesResponse{Resources: data}, nil +} + +// MCPReadResource reads a specific MCP resource by URI. +func (h *Handler) MCPReadResource(ctx context.Context, req *pb.MCPResourceRequest) (*pb.MCPResourceResponse, error) { + client, err := goumcp.Select(req.SessionId) + if err != nil { + return nil, status.Errorf(codes.NotFound, "MCP client not found: %v", err) + } + + resp, err := client.ReadResource(ctx, req.Uri) + if err != nil { + return nil, status.Errorf(codes.Internal, "ReadResource failed: %v", err) + } + + data, err := json.Marshal(resp.Contents) + if err != nil { + return nil, status.Errorf(codes.Internal, "failed to marshal contents: %v", err) + } + + return &pb.MCPResourceResponse{Contents: data}, nil +} diff --git a/grpc/mcp/mcp_test.go b/grpc/mcp/mcp_test.go new file mode 100644 index 00000000..2b6a8e9a --- /dev/null +++ b/grpc/mcp/mcp_test.go @@ -0,0 +1,264 @@ +package mcp_test + +import ( + "context" + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/yaoapp/yao/grpc/pb" + "github.com/yaoapp/yao/grpc/tests/testutils" +) + +const echoSession = "echo" + +// --- MCPListTools --- + +func TestMCPListTools_Success(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.MCPListTools(ctx, &pb.MCPListRequest{ + SessionId: echoSession, + }) + assert.NoError(t, err) + if assert.NotNil(t, resp) { + assert.NotEmpty(t, resp.Tools) + + var tools []map[string]interface{} + err := json.Unmarshal(resp.Tools, &tools) + assert.NoError(t, err) + assert.GreaterOrEqual(t, len(tools), 3, "echo MCP defines ping, status, echo") + + names := make(map[string]bool) + for _, tool := range tools { + if n, ok := tool["name"].(string); ok { + names[n] = true + } + } + assert.True(t, names["ping"], "should contain ping tool") + assert.True(t, names["status"], "should contain status tool") + assert.True(t, names["echo"], "should contain echo tool") + } +} + +func TestMCPListTools_InvalidSession(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.MCPListTools(ctx, &pb.MCPListRequest{ + SessionId: "nonexistent-session", + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.NotFound, st.Code()) +} + +// --- MCPCallTool --- + +func TestMCPCallTool_Ping(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + args, _ := json.Marshal(map[string]interface{}{"count": 2, "message": "ping"}) + resp, err := client.MCPCallTool(ctx, &pb.MCPCallRequest{ + SessionId: echoSession, + Tool: "ping", + Arguments: args, + }) + assert.NoError(t, err) + if assert.NotNil(t, resp) { + assert.NotEmpty(t, resp.Result) + var result map[string]interface{} + err := json.Unmarshal(resp.Result, &result) + assert.NoError(t, err) + } +} + +func TestMCPCallTool_Echo(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + args, _ := json.Marshal(map[string]interface{}{"message": "hello", "uppercase": true}) + resp, err := client.MCPCallTool(ctx, &pb.MCPCallRequest{ + SessionId: echoSession, + Tool: "echo", + Arguments: args, + }) + assert.NoError(t, err) + if assert.NotNil(t, resp) { + assert.NotEmpty(t, resp.Result) + } +} + +func TestMCPCallTool_NilArgs(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.MCPCallTool(ctx, &pb.MCPCallRequest{ + SessionId: echoSession, + Tool: "ping", + Arguments: nil, + }) + assert.NoError(t, err) + if assert.NotNil(t, resp) { + assert.NotEmpty(t, resp.Result) + } +} + +func TestMCPCallTool_InvalidSession(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.MCPCallTool(ctx, &pb.MCPCallRequest{ + SessionId: "nonexistent-session", + Tool: "some-tool", + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.NotFound, st.Code()) +} + +func TestMCPCallTool_BadArgs(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.MCPCallTool(ctx, &pb.MCPCallRequest{ + SessionId: echoSession, + Tool: "ping", + Arguments: []byte("{not-json"), + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +// --- MCPListResources --- + +func TestMCPListResources_Success(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.MCPListResources(ctx, &pb.MCPListRequest{ + SessionId: echoSession, + }) + assert.NoError(t, err) + if assert.NotNil(t, resp) { + assert.NotEmpty(t, resp.Resources) + + var resources []map[string]interface{} + err := json.Unmarshal(resp.Resources, &resources) + assert.NoError(t, err) + assert.GreaterOrEqual(t, len(resources), 2, "echo MCP defines info and health resources") + } +} + +func TestMCPListResources_InvalidSession(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.MCPListResources(ctx, &pb.MCPListRequest{ + SessionId: "nonexistent-session", + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.NotFound, st.Code()) +} + +// --- MCPReadResource --- + +func TestMCPReadResource_Success(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.MCPReadResource(ctx, &pb.MCPResourceRequest{ + SessionId: echoSession, + Uri: "echo://info", + }) + assert.NoError(t, err) + if assert.NotNil(t, resp) { + assert.NotEmpty(t, resp.Contents) + + var contents []map[string]interface{} + err := json.Unmarshal(resp.Contents, &contents) + assert.NoError(t, err) + assert.Greater(t, len(contents), 0) + } +} + +func TestMCPReadResource_InvalidSession(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.MCPReadResource(ctx, &pb.MCPResourceRequest{ + SessionId: "nonexistent-session", + Uri: "echo://info", + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.NotFound, st.Code()) +} + +func TestMCPReadResource_NotFoundURI(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:mcp") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.MCPReadResource(ctx, &pb.MCPResourceRequest{ + SessionId: echoSession, + Uri: "echo://nonexistent", + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.Internal, st.Code()) +} diff --git a/grpc/pb/yao.pb.go b/grpc/pb/yao.pb.go new file mode 100644 index 00000000..cea122a1 --- /dev/null +++ b/grpc/pb/yao.pb.go @@ -0,0 +1,1316 @@ +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.36.11 +// protoc v4.25.0 +// source: yao.proto + +package pb + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + reflect "reflect" + sync "sync" + unsafe "unsafe" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +type RunRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Process string `protobuf:"bytes,1,opt,name=process,proto3" json:"process,omitempty"` + Args []byte `protobuf:"bytes,2,opt,name=args,proto3" json:"args,omitempty"` // JSON-encoded argument array + Timeout int32 `protobuf:"varint,3,opt,name=timeout,proto3" json:"timeout,omitempty"` // seconds, 0 = server default + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RunRequest) Reset() { + *x = RunRequest{} + mi := &file_yao_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RunRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RunRequest) ProtoMessage() {} + +func (x *RunRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[0] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RunRequest.ProtoReflect.Descriptor instead. +func (*RunRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{0} +} + +func (x *RunRequest) GetProcess() string { + if x != nil { + return x.Process + } + return "" +} + +func (x *RunRequest) GetArgs() []byte { + if x != nil { + return x.Args + } + return nil +} + +func (x *RunRequest) GetTimeout() int32 { + if x != nil { + return x.Timeout + } + return 0 +} + +type RunResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Data []byte `protobuf:"bytes,1,opt,name=data,proto3" json:"data,omitempty"` // JSON-encoded result + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *RunResponse) Reset() { + *x = RunResponse{} + mi := &file_yao_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *RunResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RunResponse) ProtoMessage() {} + +func (x *RunResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[1] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RunResponse.ProtoReflect.Descriptor instead. +func (*RunResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{1} +} + +func (x *RunResponse) GetData() []byte { + if x != nil { + return x.Data + } + return nil +} + +type Chunk struct { + state protoimpl.MessageState `protogen:"open.v1"` + Data []byte `protobuf:"bytes,1,opt,name=data,proto3" json:"data,omitempty"` + Done bool `protobuf:"varint,2,opt,name=done,proto3" json:"done,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *Chunk) Reset() { + *x = Chunk{} + mi := &file_yao_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *Chunk) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*Chunk) ProtoMessage() {} + +func (x *Chunk) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[2] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use Chunk.ProtoReflect.Descriptor instead. +func (*Chunk) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{2} +} + +func (x *Chunk) GetData() []byte { + if x != nil { + return x.Data + } + return nil +} + +func (x *Chunk) GetDone() bool { + if x != nil { + return x.Done + } + return false +} + +type ShellRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Command string `protobuf:"bytes,1,opt,name=command,proto3" json:"command,omitempty"` + Args []string `protobuf:"bytes,2,rep,name=args,proto3" json:"args,omitempty"` + Env map[string]string `protobuf:"bytes,3,rep,name=env,proto3" json:"env,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + Timeout int32 `protobuf:"varint,4,opt,name=timeout,proto3" json:"timeout,omitempty"` // seconds, 0 = default 30s + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ShellRequest) Reset() { + *x = ShellRequest{} + mi := &file_yao_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ShellRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ShellRequest) ProtoMessage() {} + +func (x *ShellRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[3] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ShellRequest.ProtoReflect.Descriptor instead. +func (*ShellRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{3} +} + +func (x *ShellRequest) GetCommand() string { + if x != nil { + return x.Command + } + return "" +} + +func (x *ShellRequest) GetArgs() []string { + if x != nil { + return x.Args + } + return nil +} + +func (x *ShellRequest) GetEnv() map[string]string { + if x != nil { + return x.Env + } + return nil +} + +func (x *ShellRequest) GetTimeout() int32 { + if x != nil { + return x.Timeout + } + return 0 +} + +type ShellResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Stdout []byte `protobuf:"bytes,1,opt,name=stdout,proto3" json:"stdout,omitempty"` + Stderr []byte `protobuf:"bytes,2,opt,name=stderr,proto3" json:"stderr,omitempty"` + ExitCode int32 `protobuf:"varint,3,opt,name=exit_code,json=exitCode,proto3" json:"exit_code,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ShellResponse) Reset() { + *x = ShellResponse{} + mi := &file_yao_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ShellResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ShellResponse) ProtoMessage() {} + +func (x *ShellResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[4] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ShellResponse.ProtoReflect.Descriptor instead. +func (*ShellResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{4} +} + +func (x *ShellResponse) GetStdout() []byte { + if x != nil { + return x.Stdout + } + return nil +} + +func (x *ShellResponse) GetStderr() []byte { + if x != nil { + return x.Stderr + } + return nil +} + +func (x *ShellResponse) GetExitCode() int32 { + if x != nil { + return x.ExitCode + } + return 0 +} + +type APIRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Method string `protobuf:"bytes,1,opt,name=method,proto3" json:"method,omitempty"` // HTTP method + Path string `protobuf:"bytes,2,opt,name=path,proto3" json:"path,omitempty"` // openapi path + Headers map[string]string `protobuf:"bytes,3,rep,name=headers,proto3" json:"headers,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + Body []byte `protobuf:"bytes,4,opt,name=body,proto3" json:"body,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *APIRequest) Reset() { + *x = APIRequest{} + mi := &file_yao_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *APIRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*APIRequest) ProtoMessage() {} + +func (x *APIRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[5] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use APIRequest.ProtoReflect.Descriptor instead. +func (*APIRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{5} +} + +func (x *APIRequest) GetMethod() string { + if x != nil { + return x.Method + } + return "" +} + +func (x *APIRequest) GetPath() string { + if x != nil { + return x.Path + } + return "" +} + +func (x *APIRequest) GetHeaders() map[string]string { + if x != nil { + return x.Headers + } + return nil +} + +func (x *APIRequest) GetBody() []byte { + if x != nil { + return x.Body + } + return nil +} + +type APIResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Status int32 `protobuf:"varint,1,opt,name=status,proto3" json:"status,omitempty"` // HTTP status code + Headers map[string]string `protobuf:"bytes,2,rep,name=headers,proto3" json:"headers,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + Body []byte `protobuf:"bytes,3,opt,name=body,proto3" json:"body,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *APIResponse) Reset() { + *x = APIResponse{} + mi := &file_yao_proto_msgTypes[6] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *APIResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*APIResponse) ProtoMessage() {} + +func (x *APIResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[6] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use APIResponse.ProtoReflect.Descriptor instead. +func (*APIResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{6} +} + +func (x *APIResponse) GetStatus() int32 { + if x != nil { + return x.Status + } + return 0 +} + +func (x *APIResponse) GetHeaders() map[string]string { + if x != nil { + return x.Headers + } + return nil +} + +func (x *APIResponse) GetBody() []byte { + if x != nil { + return x.Body + } + return nil +} + +type MCPListRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPListRequest) Reset() { + *x = MCPListRequest{} + mi := &file_yao_proto_msgTypes[7] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPListRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPListRequest) ProtoMessage() {} + +func (x *MCPListRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[7] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPListRequest.ProtoReflect.Descriptor instead. +func (*MCPListRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{7} +} + +func (x *MCPListRequest) GetSessionId() string { + if x != nil { + return x.SessionId + } + return "" +} + +type MCPListResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Tools []byte `protobuf:"bytes,1,opt,name=tools,proto3" json:"tools,omitempty"` // JSON array of tool definitions + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPListResponse) Reset() { + *x = MCPListResponse{} + mi := &file_yao_proto_msgTypes[8] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPListResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPListResponse) ProtoMessage() {} + +func (x *MCPListResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[8] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPListResponse.ProtoReflect.Descriptor instead. +func (*MCPListResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{8} +} + +func (x *MCPListResponse) GetTools() []byte { + if x != nil { + return x.Tools + } + return nil +} + +type MCPCallRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` + Tool string `protobuf:"bytes,2,opt,name=tool,proto3" json:"tool,omitempty"` + Arguments []byte `protobuf:"bytes,3,opt,name=arguments,proto3" json:"arguments,omitempty"` // JSON-encoded arguments + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPCallRequest) Reset() { + *x = MCPCallRequest{} + mi := &file_yao_proto_msgTypes[9] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPCallRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPCallRequest) ProtoMessage() {} + +func (x *MCPCallRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[9] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPCallRequest.ProtoReflect.Descriptor instead. +func (*MCPCallRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{9} +} + +func (x *MCPCallRequest) GetSessionId() string { + if x != nil { + return x.SessionId + } + return "" +} + +func (x *MCPCallRequest) GetTool() string { + if x != nil { + return x.Tool + } + return "" +} + +func (x *MCPCallRequest) GetArguments() []byte { + if x != nil { + return x.Arguments + } + return nil +} + +type MCPCallResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Result []byte `protobuf:"bytes,1,opt,name=result,proto3" json:"result,omitempty"` // JSON-encoded result + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPCallResponse) Reset() { + *x = MCPCallResponse{} + mi := &file_yao_proto_msgTypes[10] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPCallResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPCallResponse) ProtoMessage() {} + +func (x *MCPCallResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[10] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPCallResponse.ProtoReflect.Descriptor instead. +func (*MCPCallResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{10} +} + +func (x *MCPCallResponse) GetResult() []byte { + if x != nil { + return x.Result + } + return nil +} + +type MCPResourcesResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Resources []byte `protobuf:"bytes,1,opt,name=resources,proto3" json:"resources,omitempty"` // JSON array of resource definitions + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPResourcesResponse) Reset() { + *x = MCPResourcesResponse{} + mi := &file_yao_proto_msgTypes[11] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPResourcesResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPResourcesResponse) ProtoMessage() {} + +func (x *MCPResourcesResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[11] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPResourcesResponse.ProtoReflect.Descriptor instead. +func (*MCPResourcesResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{11} +} + +func (x *MCPResourcesResponse) GetResources() []byte { + if x != nil { + return x.Resources + } + return nil +} + +type MCPResourceRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` + Uri string `protobuf:"bytes,2,opt,name=uri,proto3" json:"uri,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPResourceRequest) Reset() { + *x = MCPResourceRequest{} + mi := &file_yao_proto_msgTypes[12] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPResourceRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPResourceRequest) ProtoMessage() {} + +func (x *MCPResourceRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[12] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPResourceRequest.ProtoReflect.Descriptor instead. +func (*MCPResourceRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{12} +} + +func (x *MCPResourceRequest) GetSessionId() string { + if x != nil { + return x.SessionId + } + return "" +} + +func (x *MCPResourceRequest) GetUri() string { + if x != nil { + return x.Uri + } + return "" +} + +type MCPResourceResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Contents []byte `protobuf:"bytes,1,opt,name=contents,proto3" json:"contents,omitempty"` // JSON-encoded resource contents + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *MCPResourceResponse) Reset() { + *x = MCPResourceResponse{} + mi := &file_yao_proto_msgTypes[13] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *MCPResourceResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*MCPResourceResponse) ProtoMessage() {} + +func (x *MCPResourceResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[13] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use MCPResourceResponse.ProtoReflect.Descriptor instead. +func (*MCPResourceResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{13} +} + +func (x *MCPResourceResponse) GetContents() []byte { + if x != nil { + return x.Contents + } + return nil +} + +type ChatRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + Connector string `protobuf:"bytes,1,opt,name=connector,proto3" json:"connector,omitempty"` // connector ID + Messages []byte `protobuf:"bytes,2,opt,name=messages,proto3" json:"messages,omitempty"` // JSON-encoded message array + Options []byte `protobuf:"bytes,3,opt,name=options,proto3" json:"options,omitempty"` // JSON-encoded options + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ChatRequest) Reset() { + *x = ChatRequest{} + mi := &file_yao_proto_msgTypes[14] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ChatRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ChatRequest) ProtoMessage() {} + +func (x *ChatRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[14] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ChatRequest.ProtoReflect.Descriptor instead. +func (*ChatRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{14} +} + +func (x *ChatRequest) GetConnector() string { + if x != nil { + return x.Connector + } + return "" +} + +func (x *ChatRequest) GetMessages() []byte { + if x != nil { + return x.Messages + } + return nil +} + +func (x *ChatRequest) GetOptions() []byte { + if x != nil { + return x.Options + } + return nil +} + +type ChatResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Data []byte `protobuf:"bytes,1,opt,name=data,proto3" json:"data,omitempty"` // JSON-encoded completion result + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ChatResponse) Reset() { + *x = ChatResponse{} + mi := &file_yao_proto_msgTypes[15] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ChatResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ChatResponse) ProtoMessage() {} + +func (x *ChatResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[15] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ChatResponse.ProtoReflect.Descriptor instead. +func (*ChatResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{15} +} + +func (x *ChatResponse) GetData() []byte { + if x != nil { + return x.Data + } + return nil +} + +type ChatChunk struct { + state protoimpl.MessageState `protogen:"open.v1"` + Data []byte `protobuf:"bytes,1,opt,name=data,proto3" json:"data,omitempty"` // JSON-encoded chunk + Done bool `protobuf:"varint,2,opt,name=done,proto3" json:"done,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ChatChunk) Reset() { + *x = ChatChunk{} + mi := &file_yao_proto_msgTypes[16] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ChatChunk) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ChatChunk) ProtoMessage() {} + +func (x *ChatChunk) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[16] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use ChatChunk.ProtoReflect.Descriptor instead. +func (*ChatChunk) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{16} +} + +func (x *ChatChunk) GetData() []byte { + if x != nil { + return x.Data + } + return nil +} + +func (x *ChatChunk) GetDone() bool { + if x != nil { + return x.Done + } + return false +} + +type AgentRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + AssistantId string `protobuf:"bytes,1,opt,name=assistant_id,json=assistantId,proto3" json:"assistant_id,omitempty"` + Messages []byte `protobuf:"bytes,2,opt,name=messages,proto3" json:"messages,omitempty"` // JSON-encoded message array + Options []byte `protobuf:"bytes,3,opt,name=options,proto3" json:"options,omitempty"` // JSON-encoded options + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *AgentRequest) Reset() { + *x = AgentRequest{} + mi := &file_yao_proto_msgTypes[17] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *AgentRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*AgentRequest) ProtoMessage() {} + +func (x *AgentRequest) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[17] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use AgentRequest.ProtoReflect.Descriptor instead. +func (*AgentRequest) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{17} +} + +func (x *AgentRequest) GetAssistantId() string { + if x != nil { + return x.AssistantId + } + return "" +} + +func (x *AgentRequest) GetMessages() []byte { + if x != nil { + return x.Messages + } + return nil +} + +func (x *AgentRequest) GetOptions() []byte { + if x != nil { + return x.Options + } + return nil +} + +// Each chunk carries JSON-serialized agent/output/message.Message. +type AgentChunk struct { + state protoimpl.MessageState `protogen:"open.v1"` + Data []byte `protobuf:"bytes,1,opt,name=data,proto3" json:"data,omitempty"` + Done bool `protobuf:"varint,2,opt,name=done,proto3" json:"done,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *AgentChunk) Reset() { + *x = AgentChunk{} + mi := &file_yao_proto_msgTypes[18] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *AgentChunk) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*AgentChunk) ProtoMessage() {} + +func (x *AgentChunk) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[18] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use AgentChunk.ProtoReflect.Descriptor instead. +func (*AgentChunk) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{18} +} + +func (x *AgentChunk) GetData() []byte { + if x != nil { + return x.Data + } + return nil +} + +func (x *AgentChunk) GetDone() bool { + if x != nil { + return x.Done + } + return false +} + +type Empty struct { + state protoimpl.MessageState `protogen:"open.v1"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *Empty) Reset() { + *x = Empty{} + mi := &file_yao_proto_msgTypes[19] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *Empty) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*Empty) ProtoMessage() {} + +func (x *Empty) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[19] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use Empty.ProtoReflect.Descriptor instead. +func (*Empty) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{19} +} + +type HealthzResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + Status string `protobuf:"bytes,1,opt,name=status,proto3" json:"status,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *HealthzResponse) Reset() { + *x = HealthzResponse{} + mi := &file_yao_proto_msgTypes[20] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *HealthzResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*HealthzResponse) ProtoMessage() {} + +func (x *HealthzResponse) ProtoReflect() protoreflect.Message { + mi := &file_yao_proto_msgTypes[20] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use HealthzResponse.ProtoReflect.Descriptor instead. +func (*HealthzResponse) Descriptor() ([]byte, []int) { + return file_yao_proto_rawDescGZIP(), []int{20} +} + +func (x *HealthzResponse) GetStatus() string { + if x != nil { + return x.Status + } + return "" +} + +var File_yao_proto protoreflect.FileDescriptor + +const file_yao_proto_rawDesc = "" + + "\n" + + "\tyao.proto\x12\x03yao\"T\n" + + "\n" + + "RunRequest\x12\x18\n" + + "\aprocess\x18\x01 \x01(\tR\aprocess\x12\x12\n" + + "\x04args\x18\x02 \x01(\fR\x04args\x12\x18\n" + + "\atimeout\x18\x03 \x01(\x05R\atimeout\"!\n" + + "\vRunResponse\x12\x12\n" + + "\x04data\x18\x01 \x01(\fR\x04data\"/\n" + + "\x05Chunk\x12\x12\n" + + "\x04data\x18\x01 \x01(\fR\x04data\x12\x12\n" + + "\x04done\x18\x02 \x01(\bR\x04done\"\xbc\x01\n" + + "\fShellRequest\x12\x18\n" + + "\acommand\x18\x01 \x01(\tR\acommand\x12\x12\n" + + "\x04args\x18\x02 \x03(\tR\x04args\x12,\n" + + "\x03env\x18\x03 \x03(\v2\x1a.yao.ShellRequest.EnvEntryR\x03env\x12\x18\n" + + "\atimeout\x18\x04 \x01(\x05R\atimeout\x1a6\n" + + "\bEnvEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\\\n" + + "\rShellResponse\x12\x16\n" + + "\x06stdout\x18\x01 \x01(\fR\x06stdout\x12\x16\n" + + "\x06stderr\x18\x02 \x01(\fR\x06stderr\x12\x1b\n" + + "\texit_code\x18\x03 \x01(\x05R\bexitCode\"\xc0\x01\n" + + "\n" + + "APIRequest\x12\x16\n" + + "\x06method\x18\x01 \x01(\tR\x06method\x12\x12\n" + + "\x04path\x18\x02 \x01(\tR\x04path\x126\n" + + "\aheaders\x18\x03 \x03(\v2\x1c.yao.APIRequest.HeadersEntryR\aheaders\x12\x12\n" + + "\x04body\x18\x04 \x01(\fR\x04body\x1a:\n" + + "\fHeadersEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\xae\x01\n" + + "\vAPIResponse\x12\x16\n" + + "\x06status\x18\x01 \x01(\x05R\x06status\x127\n" + + "\aheaders\x18\x02 \x03(\v2\x1d.yao.APIResponse.HeadersEntryR\aheaders\x12\x12\n" + + "\x04body\x18\x03 \x01(\fR\x04body\x1a:\n" + + "\fHeadersEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"/\n" + + "\x0eMCPListRequest\x12\x1d\n" + + "\n" + + "session_id\x18\x01 \x01(\tR\tsessionId\"'\n" + + "\x0fMCPListResponse\x12\x14\n" + + "\x05tools\x18\x01 \x01(\fR\x05tools\"a\n" + + "\x0eMCPCallRequest\x12\x1d\n" + + "\n" + + "session_id\x18\x01 \x01(\tR\tsessionId\x12\x12\n" + + "\x04tool\x18\x02 \x01(\tR\x04tool\x12\x1c\n" + + "\targuments\x18\x03 \x01(\fR\targuments\")\n" + + "\x0fMCPCallResponse\x12\x16\n" + + "\x06result\x18\x01 \x01(\fR\x06result\"4\n" + + "\x14MCPResourcesResponse\x12\x1c\n" + + "\tresources\x18\x01 \x01(\fR\tresources\"E\n" + + "\x12MCPResourceRequest\x12\x1d\n" + + "\n" + + "session_id\x18\x01 \x01(\tR\tsessionId\x12\x10\n" + + "\x03uri\x18\x02 \x01(\tR\x03uri\"1\n" + + "\x13MCPResourceResponse\x12\x1a\n" + + "\bcontents\x18\x01 \x01(\fR\bcontents\"a\n" + + "\vChatRequest\x12\x1c\n" + + "\tconnector\x18\x01 \x01(\tR\tconnector\x12\x1a\n" + + "\bmessages\x18\x02 \x01(\fR\bmessages\x12\x18\n" + + "\aoptions\x18\x03 \x01(\fR\aoptions\"\"\n" + + "\fChatResponse\x12\x12\n" + + "\x04data\x18\x01 \x01(\fR\x04data\"3\n" + + "\tChatChunk\x12\x12\n" + + "\x04data\x18\x01 \x01(\fR\x04data\x12\x12\n" + + "\x04done\x18\x02 \x01(\bR\x04done\"g\n" + + "\fAgentRequest\x12!\n" + + "\fassistant_id\x18\x01 \x01(\tR\vassistantId\x12\x1a\n" + + "\bmessages\x18\x02 \x01(\fR\bmessages\x12\x18\n" + + "\aoptions\x18\x03 \x01(\fR\aoptions\"4\n" + + "\n" + + "AgentChunk\x12\x12\n" + + "\x04data\x18\x01 \x01(\fR\x04data\x12\x12\n" + + "\x04done\x18\x02 \x01(\bR\x04done\"\a\n" + + "\x05Empty\")\n" + + "\x0fHealthzResponse\x12\x16\n" + + "\x06status\x18\x01 \x01(\tR\x06status2\xb8\x05\n" + + "\x03Yao\x12(\n" + + "\x03Run\x12\x0f.yao.RunRequest\x1a\x10.yao.RunResponse\x12'\n" + + "\x06Stream\x12\x0f.yao.RunRequest\x1a\n" + + ".yao.Chunk0\x01\x12.\n" + + "\x05Shell\x12\x11.yao.ShellRequest\x1a\x12.yao.ShellResponse\x12.\n" + + "\vShellStream\x12\x11.yao.ShellRequest\x1a\n" + + ".yao.Chunk0\x01\x12(\n" + + "\x03API\x12\x0f.yao.APIRequest\x1a\x10.yao.APIResponse\x129\n" + + "\fMCPListTools\x12\x13.yao.MCPListRequest\x1a\x14.yao.MCPListResponse\x128\n" + + "\vMCPCallTool\x12\x13.yao.MCPCallRequest\x1a\x14.yao.MCPCallResponse\x12B\n" + + "\x10MCPListResources\x12\x13.yao.MCPListRequest\x1a\x19.yao.MCPResourcesResponse\x12D\n" + + "\x0fMCPReadResource\x12\x17.yao.MCPResourceRequest\x1a\x18.yao.MCPResourceResponse\x126\n" + + "\x0fChatCompletions\x12\x10.yao.ChatRequest\x1a\x11.yao.ChatResponse\x12;\n" + + "\x15ChatCompletionsStream\x12\x10.yao.ChatRequest\x1a\x0e.yao.ChatChunk0\x01\x123\n" + + "\vAgentStream\x12\x11.yao.AgentRequest\x1a\x0f.yao.AgentChunk0\x01\x12+\n" + + "\aHealthz\x12\n" + + ".yao.Empty\x1a\x14.yao.HealthzResponseB\x1fZ\x1dgithub.com/yaoapp/yao/grpc/pbb\x06proto3" + +var ( + file_yao_proto_rawDescOnce sync.Once + file_yao_proto_rawDescData []byte +) + +func file_yao_proto_rawDescGZIP() []byte { + file_yao_proto_rawDescOnce.Do(func() { + file_yao_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_yao_proto_rawDesc), len(file_yao_proto_rawDesc))) + }) + return file_yao_proto_rawDescData +} + +var file_yao_proto_msgTypes = make([]protoimpl.MessageInfo, 24) +var file_yao_proto_goTypes = []any{ + (*RunRequest)(nil), // 0: yao.RunRequest + (*RunResponse)(nil), // 1: yao.RunResponse + (*Chunk)(nil), // 2: yao.Chunk + (*ShellRequest)(nil), // 3: yao.ShellRequest + (*ShellResponse)(nil), // 4: yao.ShellResponse + (*APIRequest)(nil), // 5: yao.APIRequest + (*APIResponse)(nil), // 6: yao.APIResponse + (*MCPListRequest)(nil), // 7: yao.MCPListRequest + (*MCPListResponse)(nil), // 8: yao.MCPListResponse + (*MCPCallRequest)(nil), // 9: yao.MCPCallRequest + (*MCPCallResponse)(nil), // 10: yao.MCPCallResponse + (*MCPResourcesResponse)(nil), // 11: yao.MCPResourcesResponse + (*MCPResourceRequest)(nil), // 12: yao.MCPResourceRequest + (*MCPResourceResponse)(nil), // 13: yao.MCPResourceResponse + (*ChatRequest)(nil), // 14: yao.ChatRequest + (*ChatResponse)(nil), // 15: yao.ChatResponse + (*ChatChunk)(nil), // 16: yao.ChatChunk + (*AgentRequest)(nil), // 17: yao.AgentRequest + (*AgentChunk)(nil), // 18: yao.AgentChunk + (*Empty)(nil), // 19: yao.Empty + (*HealthzResponse)(nil), // 20: yao.HealthzResponse + nil, // 21: yao.ShellRequest.EnvEntry + nil, // 22: yao.APIRequest.HeadersEntry + nil, // 23: yao.APIResponse.HeadersEntry +} +var file_yao_proto_depIdxs = []int32{ + 21, // 0: yao.ShellRequest.env:type_name -> yao.ShellRequest.EnvEntry + 22, // 1: yao.APIRequest.headers:type_name -> yao.APIRequest.HeadersEntry + 23, // 2: yao.APIResponse.headers:type_name -> yao.APIResponse.HeadersEntry + 0, // 3: yao.Yao.Run:input_type -> yao.RunRequest + 0, // 4: yao.Yao.Stream:input_type -> yao.RunRequest + 3, // 5: yao.Yao.Shell:input_type -> yao.ShellRequest + 3, // 6: yao.Yao.ShellStream:input_type -> yao.ShellRequest + 5, // 7: yao.Yao.API:input_type -> yao.APIRequest + 7, // 8: yao.Yao.MCPListTools:input_type -> yao.MCPListRequest + 9, // 9: yao.Yao.MCPCallTool:input_type -> yao.MCPCallRequest + 7, // 10: yao.Yao.MCPListResources:input_type -> yao.MCPListRequest + 12, // 11: yao.Yao.MCPReadResource:input_type -> yao.MCPResourceRequest + 14, // 12: yao.Yao.ChatCompletions:input_type -> yao.ChatRequest + 14, // 13: yao.Yao.ChatCompletionsStream:input_type -> yao.ChatRequest + 17, // 14: yao.Yao.AgentStream:input_type -> yao.AgentRequest + 19, // 15: yao.Yao.Healthz:input_type -> yao.Empty + 1, // 16: yao.Yao.Run:output_type -> yao.RunResponse + 2, // 17: yao.Yao.Stream:output_type -> yao.Chunk + 4, // 18: yao.Yao.Shell:output_type -> yao.ShellResponse + 2, // 19: yao.Yao.ShellStream:output_type -> yao.Chunk + 6, // 20: yao.Yao.API:output_type -> yao.APIResponse + 8, // 21: yao.Yao.MCPListTools:output_type -> yao.MCPListResponse + 10, // 22: yao.Yao.MCPCallTool:output_type -> yao.MCPCallResponse + 11, // 23: yao.Yao.MCPListResources:output_type -> yao.MCPResourcesResponse + 13, // 24: yao.Yao.MCPReadResource:output_type -> yao.MCPResourceResponse + 15, // 25: yao.Yao.ChatCompletions:output_type -> yao.ChatResponse + 16, // 26: yao.Yao.ChatCompletionsStream:output_type -> yao.ChatChunk + 18, // 27: yao.Yao.AgentStream:output_type -> yao.AgentChunk + 20, // 28: yao.Yao.Healthz:output_type -> yao.HealthzResponse + 16, // [16:29] is the sub-list for method output_type + 3, // [3:16] is the sub-list for method input_type + 3, // [3:3] is the sub-list for extension type_name + 3, // [3:3] is the sub-list for extension extendee + 0, // [0:3] is the sub-list for field type_name +} + +func init() { file_yao_proto_init() } +func file_yao_proto_init() { + if File_yao_proto != nil { + return + } + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + GoPackagePath: reflect.TypeOf(x{}).PkgPath(), + RawDescriptor: unsafe.Slice(unsafe.StringData(file_yao_proto_rawDesc), len(file_yao_proto_rawDesc)), + NumEnums: 0, + NumMessages: 24, + NumExtensions: 0, + NumServices: 1, + }, + GoTypes: file_yao_proto_goTypes, + DependencyIndexes: file_yao_proto_depIdxs, + MessageInfos: file_yao_proto_msgTypes, + }.Build() + File_yao_proto = out.File + file_yao_proto_goTypes = nil + file_yao_proto_depIdxs = nil +} diff --git a/grpc/pb/yao.proto b/grpc/pb/yao.proto new file mode 100644 index 00000000..f120151f --- /dev/null +++ b/grpc/pb/yao.proto @@ -0,0 +1,149 @@ +syntax = "proto3"; +package yao; +option go_package = "github.com/yaoapp/yao/grpc/pb"; + +// Yao gRPC gateway. Shares OAuth + ACL scope system with openapi. +service Yao { + + // Base + rpc Run(RunRequest) returns (RunResponse); + rpc Stream(RunRequest) returns (stream Chunk); + rpc Shell(ShellRequest) returns (ShellResponse); + rpc ShellStream(ShellRequest) returns (stream Chunk); + + // API gateway + rpc API(APIRequest) returns (APIResponse); + + // MCP + rpc MCPListTools(MCPListRequest) returns (MCPListResponse); + rpc MCPCallTool(MCPCallRequest) returns (MCPCallResponse); + rpc MCPListResources(MCPListRequest) returns (MCPResourcesResponse); + rpc MCPReadResource(MCPResourceRequest) returns (MCPResourceResponse); + + // AI - LLM + rpc ChatCompletions(ChatRequest) returns (ChatResponse); + rpc ChatCompletionsStream(ChatRequest) returns (stream ChatChunk); + + // AI - Agent + rpc AgentStream(AgentRequest) returns (stream AgentChunk); + + // Health + rpc Healthz(Empty) returns (HealthzResponse); +} + +// ── Base ───────────────────────────────────────────────────────────────────── + +message RunRequest { + string process = 1; + bytes args = 2; // JSON-encoded argument array + int32 timeout = 3; // seconds, 0 = server default +} + +message RunResponse { + bytes data = 1; // JSON-encoded result +} + +message Chunk { + bytes data = 1; + bool done = 2; +} + +message ShellRequest { + string command = 1; + repeated string args = 2; + map env = 3; + int32 timeout = 4; // seconds, 0 = default 30s +} + +message ShellResponse { + bytes stdout = 1; + bytes stderr = 2; + int32 exit_code = 3; +} + +// ── API gateway ────────────────────────────────────────────────────────────── + +message APIRequest { + string method = 1; // HTTP method + string path = 2; // openapi path + map headers = 3; + bytes body = 4; +} + +message APIResponse { + int32 status = 1; // HTTP status code + map headers = 2; + bytes body = 3; +} + +// ── MCP ────────────────────────────────────────────────────────────────────── + +message MCPListRequest { + string session_id = 1; +} + +message MCPListResponse { + bytes tools = 1; // JSON array of tool definitions +} + +message MCPCallRequest { + string session_id = 1; + string tool = 2; + bytes arguments = 3; // JSON-encoded arguments +} + +message MCPCallResponse { + bytes result = 1; // JSON-encoded result +} + +message MCPResourcesResponse { + bytes resources = 1; // JSON array of resource definitions +} + +message MCPResourceRequest { + string session_id = 1; + string uri = 2; +} + +message MCPResourceResponse { + bytes contents = 1; // JSON-encoded resource contents +} + +// ── LLM ────────────────────────────────────────────────────────────────────── + +message ChatRequest { + string connector = 1; // connector ID + bytes messages = 2; // JSON-encoded message array + bytes options = 3; // JSON-encoded options +} + +message ChatResponse { + bytes data = 1; // JSON-encoded completion result +} + +message ChatChunk { + bytes data = 1; // JSON-encoded chunk + bool done = 2; +} + +// ── Agent ──────────────────────────────────────────────────────────────────── + +message AgentRequest { + string assistant_id = 1; + bytes messages = 2; // JSON-encoded message array + bytes options = 3; // JSON-encoded options +} + +// Each chunk carries JSON-serialized agent/output/message.Message. +message AgentChunk { + bytes data = 1; + bool done = 2; +} + +// ── Health ─────────────────────────────────────────────────────────────────── + +message Empty {} + +message HealthzResponse { + string status = 1; +} diff --git a/grpc/pb/yao_grpc.pb.go b/grpc/pb/yao_grpc.pb.go new file mode 100644 index 00000000..db5b07d6 --- /dev/null +++ b/grpc/pb/yao_grpc.pb.go @@ -0,0 +1,606 @@ +// Code generated by protoc-gen-go-grpc. DO NOT EDIT. +// versions: +// - protoc-gen-go-grpc v1.6.1 +// - protoc v4.25.0 +// source: yao.proto + +package pb + +import ( + context "context" + grpc "google.golang.org/grpc" + codes "google.golang.org/grpc/codes" + status "google.golang.org/grpc/status" +) + +// This is a compile-time assertion to ensure that this generated file +// is compatible with the grpc package it is being compiled against. +// Requires gRPC-Go v1.64.0 or later. +const _ = grpc.SupportPackageIsVersion9 + +const ( + Yao_Run_FullMethodName = "/yao.Yao/Run" + Yao_Stream_FullMethodName = "/yao.Yao/Stream" + Yao_Shell_FullMethodName = "/yao.Yao/Shell" + Yao_ShellStream_FullMethodName = "/yao.Yao/ShellStream" + Yao_API_FullMethodName = "/yao.Yao/API" + Yao_MCPListTools_FullMethodName = "/yao.Yao/MCPListTools" + Yao_MCPCallTool_FullMethodName = "/yao.Yao/MCPCallTool" + Yao_MCPListResources_FullMethodName = "/yao.Yao/MCPListResources" + Yao_MCPReadResource_FullMethodName = "/yao.Yao/MCPReadResource" + Yao_ChatCompletions_FullMethodName = "/yao.Yao/ChatCompletions" + Yao_ChatCompletionsStream_FullMethodName = "/yao.Yao/ChatCompletionsStream" + Yao_AgentStream_FullMethodName = "/yao.Yao/AgentStream" + Yao_Healthz_FullMethodName = "/yao.Yao/Healthz" +) + +// YaoClient is the client API for Yao service. +// +// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream. +// +// Yao gRPC gateway. Shares OAuth + ACL scope system with openapi. +type YaoClient interface { + // Base + Run(ctx context.Context, in *RunRequest, opts ...grpc.CallOption) (*RunResponse, error) + Stream(ctx context.Context, in *RunRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[Chunk], error) + Shell(ctx context.Context, in *ShellRequest, opts ...grpc.CallOption) (*ShellResponse, error) + ShellStream(ctx context.Context, in *ShellRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[Chunk], error) + // API gateway + API(ctx context.Context, in *APIRequest, opts ...grpc.CallOption) (*APIResponse, error) + // MCP + MCPListTools(ctx context.Context, in *MCPListRequest, opts ...grpc.CallOption) (*MCPListResponse, error) + MCPCallTool(ctx context.Context, in *MCPCallRequest, opts ...grpc.CallOption) (*MCPCallResponse, error) + MCPListResources(ctx context.Context, in *MCPListRequest, opts ...grpc.CallOption) (*MCPResourcesResponse, error) + MCPReadResource(ctx context.Context, in *MCPResourceRequest, opts ...grpc.CallOption) (*MCPResourceResponse, error) + // AI - LLM + ChatCompletions(ctx context.Context, in *ChatRequest, opts ...grpc.CallOption) (*ChatResponse, error) + ChatCompletionsStream(ctx context.Context, in *ChatRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ChatChunk], error) + // AI - Agent + AgentStream(ctx context.Context, in *AgentRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[AgentChunk], error) + // Health + Healthz(ctx context.Context, in *Empty, opts ...grpc.CallOption) (*HealthzResponse, error) +} + +type yaoClient struct { + cc grpc.ClientConnInterface +} + +func NewYaoClient(cc grpc.ClientConnInterface) YaoClient { + return &yaoClient{cc} +} + +func (c *yaoClient) Run(ctx context.Context, in *RunRequest, opts ...grpc.CallOption) (*RunResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(RunResponse) + err := c.cc.Invoke(ctx, Yao_Run_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *yaoClient) Stream(ctx context.Context, in *RunRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[Chunk], error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + stream, err := c.cc.NewStream(ctx, &Yao_ServiceDesc.Streams[0], Yao_Stream_FullMethodName, cOpts...) + if err != nil { + return nil, err + } + x := &grpc.GenericClientStream[RunRequest, Chunk]{ClientStream: stream} + if err := x.ClientStream.SendMsg(in); err != nil { + return nil, err + } + if err := x.ClientStream.CloseSend(); err != nil { + return nil, err + } + return x, nil +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type Yao_StreamClient = grpc.ServerStreamingClient[Chunk] + +func (c *yaoClient) Shell(ctx context.Context, in *ShellRequest, opts ...grpc.CallOption) (*ShellResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(ShellResponse) + err := c.cc.Invoke(ctx, Yao_Shell_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *yaoClient) ShellStream(ctx context.Context, in *ShellRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[Chunk], error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + stream, err := c.cc.NewStream(ctx, &Yao_ServiceDesc.Streams[1], Yao_ShellStream_FullMethodName, cOpts...) + if err != nil { + return nil, err + } + x := &grpc.GenericClientStream[ShellRequest, Chunk]{ClientStream: stream} + if err := x.ClientStream.SendMsg(in); err != nil { + return nil, err + } + if err := x.ClientStream.CloseSend(); err != nil { + return nil, err + } + return x, nil +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type Yao_ShellStreamClient = grpc.ServerStreamingClient[Chunk] + +func (c *yaoClient) API(ctx context.Context, in *APIRequest, opts ...grpc.CallOption) (*APIResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(APIResponse) + err := c.cc.Invoke(ctx, Yao_API_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *yaoClient) MCPListTools(ctx context.Context, in *MCPListRequest, opts ...grpc.CallOption) (*MCPListResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(MCPListResponse) + err := c.cc.Invoke(ctx, Yao_MCPListTools_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *yaoClient) MCPCallTool(ctx context.Context, in *MCPCallRequest, opts ...grpc.CallOption) (*MCPCallResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(MCPCallResponse) + err := c.cc.Invoke(ctx, Yao_MCPCallTool_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *yaoClient) MCPListResources(ctx context.Context, in *MCPListRequest, opts ...grpc.CallOption) (*MCPResourcesResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(MCPResourcesResponse) + err := c.cc.Invoke(ctx, Yao_MCPListResources_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *yaoClient) MCPReadResource(ctx context.Context, in *MCPResourceRequest, opts ...grpc.CallOption) (*MCPResourceResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(MCPResourceResponse) + err := c.cc.Invoke(ctx, Yao_MCPReadResource_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *yaoClient) ChatCompletions(ctx context.Context, in *ChatRequest, opts ...grpc.CallOption) (*ChatResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(ChatResponse) + err := c.cc.Invoke(ctx, Yao_ChatCompletions_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *yaoClient) ChatCompletionsStream(ctx context.Context, in *ChatRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ChatChunk], error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + stream, err := c.cc.NewStream(ctx, &Yao_ServiceDesc.Streams[2], Yao_ChatCompletionsStream_FullMethodName, cOpts...) + if err != nil { + return nil, err + } + x := &grpc.GenericClientStream[ChatRequest, ChatChunk]{ClientStream: stream} + if err := x.ClientStream.SendMsg(in); err != nil { + return nil, err + } + if err := x.ClientStream.CloseSend(); err != nil { + return nil, err + } + return x, nil +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type Yao_ChatCompletionsStreamClient = grpc.ServerStreamingClient[ChatChunk] + +func (c *yaoClient) AgentStream(ctx context.Context, in *AgentRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[AgentChunk], error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + stream, err := c.cc.NewStream(ctx, &Yao_ServiceDesc.Streams[3], Yao_AgentStream_FullMethodName, cOpts...) + if err != nil { + return nil, err + } + x := &grpc.GenericClientStream[AgentRequest, AgentChunk]{ClientStream: stream} + if err := x.ClientStream.SendMsg(in); err != nil { + return nil, err + } + if err := x.ClientStream.CloseSend(); err != nil { + return nil, err + } + return x, nil +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type Yao_AgentStreamClient = grpc.ServerStreamingClient[AgentChunk] + +func (c *yaoClient) Healthz(ctx context.Context, in *Empty, opts ...grpc.CallOption) (*HealthzResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(HealthzResponse) + err := c.cc.Invoke(ctx, Yao_Healthz_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + +// YaoServer is the server API for Yao service. +// All implementations must embed UnimplementedYaoServer +// for forward compatibility. +// +// Yao gRPC gateway. Shares OAuth + ACL scope system with openapi. +type YaoServer interface { + // Base + Run(context.Context, *RunRequest) (*RunResponse, error) + Stream(*RunRequest, grpc.ServerStreamingServer[Chunk]) error + Shell(context.Context, *ShellRequest) (*ShellResponse, error) + ShellStream(*ShellRequest, grpc.ServerStreamingServer[Chunk]) error + // API gateway + API(context.Context, *APIRequest) (*APIResponse, error) + // MCP + MCPListTools(context.Context, *MCPListRequest) (*MCPListResponse, error) + MCPCallTool(context.Context, *MCPCallRequest) (*MCPCallResponse, error) + MCPListResources(context.Context, *MCPListRequest) (*MCPResourcesResponse, error) + MCPReadResource(context.Context, *MCPResourceRequest) (*MCPResourceResponse, error) + // AI - LLM + ChatCompletions(context.Context, *ChatRequest) (*ChatResponse, error) + ChatCompletionsStream(*ChatRequest, grpc.ServerStreamingServer[ChatChunk]) error + // AI - Agent + AgentStream(*AgentRequest, grpc.ServerStreamingServer[AgentChunk]) error + // Health + Healthz(context.Context, *Empty) (*HealthzResponse, error) + mustEmbedUnimplementedYaoServer() +} + +// UnimplementedYaoServer must be embedded to have +// forward compatible implementations. +// +// NOTE: this should be embedded by value instead of pointer to avoid a nil +// pointer dereference when methods are called. +type UnimplementedYaoServer struct{} + +func (UnimplementedYaoServer) Run(context.Context, *RunRequest) (*RunResponse, error) { + return nil, status.Error(codes.Unimplemented, "method Run not implemented") +} +func (UnimplementedYaoServer) Stream(*RunRequest, grpc.ServerStreamingServer[Chunk]) error { + return status.Error(codes.Unimplemented, "method Stream not implemented") +} +func (UnimplementedYaoServer) Shell(context.Context, *ShellRequest) (*ShellResponse, error) { + return nil, status.Error(codes.Unimplemented, "method Shell not implemented") +} +func (UnimplementedYaoServer) ShellStream(*ShellRequest, grpc.ServerStreamingServer[Chunk]) error { + return status.Error(codes.Unimplemented, "method ShellStream not implemented") +} +func (UnimplementedYaoServer) API(context.Context, *APIRequest) (*APIResponse, error) { + return nil, status.Error(codes.Unimplemented, "method API not implemented") +} +func (UnimplementedYaoServer) MCPListTools(context.Context, *MCPListRequest) (*MCPListResponse, error) { + return nil, status.Error(codes.Unimplemented, "method MCPListTools not implemented") +} +func (UnimplementedYaoServer) MCPCallTool(context.Context, *MCPCallRequest) (*MCPCallResponse, error) { + return nil, status.Error(codes.Unimplemented, "method MCPCallTool not implemented") +} +func (UnimplementedYaoServer) MCPListResources(context.Context, *MCPListRequest) (*MCPResourcesResponse, error) { + return nil, status.Error(codes.Unimplemented, "method MCPListResources not implemented") +} +func (UnimplementedYaoServer) MCPReadResource(context.Context, *MCPResourceRequest) (*MCPResourceResponse, error) { + return nil, status.Error(codes.Unimplemented, "method MCPReadResource not implemented") +} +func (UnimplementedYaoServer) ChatCompletions(context.Context, *ChatRequest) (*ChatResponse, error) { + return nil, status.Error(codes.Unimplemented, "method ChatCompletions not implemented") +} +func (UnimplementedYaoServer) ChatCompletionsStream(*ChatRequest, grpc.ServerStreamingServer[ChatChunk]) error { + return status.Error(codes.Unimplemented, "method ChatCompletionsStream not implemented") +} +func (UnimplementedYaoServer) AgentStream(*AgentRequest, grpc.ServerStreamingServer[AgentChunk]) error { + return status.Error(codes.Unimplemented, "method AgentStream not implemented") +} +func (UnimplementedYaoServer) Healthz(context.Context, *Empty) (*HealthzResponse, error) { + return nil, status.Error(codes.Unimplemented, "method Healthz not implemented") +} +func (UnimplementedYaoServer) mustEmbedUnimplementedYaoServer() {} +func (UnimplementedYaoServer) testEmbeddedByValue() {} + +// UnsafeYaoServer may be embedded to opt out of forward compatibility for this service. +// Use of this interface is not recommended, as added methods to YaoServer will +// result in compilation errors. +type UnsafeYaoServer interface { + mustEmbedUnimplementedYaoServer() +} + +func RegisterYaoServer(s grpc.ServiceRegistrar, srv YaoServer) { + // If the following call panics, it indicates UnimplementedYaoServer was + // embedded by pointer and is nil. This will cause panics if an + // unimplemented method is ever invoked, so we test this at initialization + // time to prevent it from happening at runtime later due to I/O. + if t, ok := srv.(interface{ testEmbeddedByValue() }); ok { + t.testEmbeddedByValue() + } + s.RegisterService(&Yao_ServiceDesc, srv) +} + +func _Yao_Run_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(RunRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(YaoServer).Run(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Yao_Run_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(YaoServer).Run(ctx, req.(*RunRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _Yao_Stream_Handler(srv interface{}, stream grpc.ServerStream) error { + m := new(RunRequest) + if err := stream.RecvMsg(m); err != nil { + return err + } + return srv.(YaoServer).Stream(m, &grpc.GenericServerStream[RunRequest, Chunk]{ServerStream: stream}) +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type Yao_StreamServer = grpc.ServerStreamingServer[Chunk] + +func _Yao_Shell_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(ShellRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(YaoServer).Shell(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Yao_Shell_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(YaoServer).Shell(ctx, req.(*ShellRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _Yao_ShellStream_Handler(srv interface{}, stream grpc.ServerStream) error { + m := new(ShellRequest) + if err := stream.RecvMsg(m); err != nil { + return err + } + return srv.(YaoServer).ShellStream(m, &grpc.GenericServerStream[ShellRequest, Chunk]{ServerStream: stream}) +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type Yao_ShellStreamServer = grpc.ServerStreamingServer[Chunk] + +func _Yao_API_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(APIRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(YaoServer).API(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Yao_API_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(YaoServer).API(ctx, req.(*APIRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _Yao_MCPListTools_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(MCPListRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(YaoServer).MCPListTools(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Yao_MCPListTools_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(YaoServer).MCPListTools(ctx, req.(*MCPListRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _Yao_MCPCallTool_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(MCPCallRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(YaoServer).MCPCallTool(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Yao_MCPCallTool_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(YaoServer).MCPCallTool(ctx, req.(*MCPCallRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _Yao_MCPListResources_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(MCPListRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(YaoServer).MCPListResources(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Yao_MCPListResources_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(YaoServer).MCPListResources(ctx, req.(*MCPListRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _Yao_MCPReadResource_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(MCPResourceRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(YaoServer).MCPReadResource(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Yao_MCPReadResource_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(YaoServer).MCPReadResource(ctx, req.(*MCPResourceRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _Yao_ChatCompletions_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(ChatRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(YaoServer).ChatCompletions(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Yao_ChatCompletions_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(YaoServer).ChatCompletions(ctx, req.(*ChatRequest)) + } + return interceptor(ctx, in, info, handler) +} + +func _Yao_ChatCompletionsStream_Handler(srv interface{}, stream grpc.ServerStream) error { + m := new(ChatRequest) + if err := stream.RecvMsg(m); err != nil { + return err + } + return srv.(YaoServer).ChatCompletionsStream(m, &grpc.GenericServerStream[ChatRequest, ChatChunk]{ServerStream: stream}) +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type Yao_ChatCompletionsStreamServer = grpc.ServerStreamingServer[ChatChunk] + +func _Yao_AgentStream_Handler(srv interface{}, stream grpc.ServerStream) error { + m := new(AgentRequest) + if err := stream.RecvMsg(m); err != nil { + return err + } + return srv.(YaoServer).AgentStream(m, &grpc.GenericServerStream[AgentRequest, AgentChunk]{ServerStream: stream}) +} + +// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name. +type Yao_AgentStreamServer = grpc.ServerStreamingServer[AgentChunk] + +func _Yao_Healthz_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(Empty) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(YaoServer).Healthz(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Yao_Healthz_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(YaoServer).Healthz(ctx, req.(*Empty)) + } + return interceptor(ctx, in, info, handler) +} + +// Yao_ServiceDesc is the grpc.ServiceDesc for Yao service. +// It's only intended for direct use with grpc.RegisterService, +// and not to be introspected or modified (even as a copy) +var Yao_ServiceDesc = grpc.ServiceDesc{ + ServiceName: "yao.Yao", + HandlerType: (*YaoServer)(nil), + Methods: []grpc.MethodDesc{ + { + MethodName: "Run", + Handler: _Yao_Run_Handler, + }, + { + MethodName: "Shell", + Handler: _Yao_Shell_Handler, + }, + { + MethodName: "API", + Handler: _Yao_API_Handler, + }, + { + MethodName: "MCPListTools", + Handler: _Yao_MCPListTools_Handler, + }, + { + MethodName: "MCPCallTool", + Handler: _Yao_MCPCallTool_Handler, + }, + { + MethodName: "MCPListResources", + Handler: _Yao_MCPListResources_Handler, + }, + { + MethodName: "MCPReadResource", + Handler: _Yao_MCPReadResource_Handler, + }, + { + MethodName: "ChatCompletions", + Handler: _Yao_ChatCompletions_Handler, + }, + { + MethodName: "Healthz", + Handler: _Yao_Healthz_Handler, + }, + }, + Streams: []grpc.StreamDesc{ + { + StreamName: "Stream", + Handler: _Yao_Stream_Handler, + ServerStreams: true, + }, + { + StreamName: "ShellStream", + Handler: _Yao_ShellStream_Handler, + ServerStreams: true, + }, + { + StreamName: "ChatCompletionsStream", + Handler: _Yao_ChatCompletionsStream_Handler, + ServerStreams: true, + }, + { + StreamName: "AgentStream", + Handler: _Yao_AgentStream_Handler, + ServerStreams: true, + }, + }, + Metadata: "yao.proto", +} diff --git a/grpc/run/run.go b/grpc/run/run.go new file mode 100644 index 00000000..fc5b0820 --- /dev/null +++ b/grpc/run/run.go @@ -0,0 +1,79 @@ +package run + +import ( + "context" + "encoding/json" + "time" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/yaoapp/gou/process" + "github.com/yaoapp/yao/grpc/auth" + "github.com/yaoapp/yao/grpc/pb" +) + +// Handler implements the Run gRPC method. +type Handler struct{} + +// Run executes a Yao process by name and returns the JSON-encoded result. +func (h *Handler) Run(ctx context.Context, req *pb.RunRequest) (*pb.RunResponse, error) { + if req.Process == "" { + return nil, status.Error(codes.InvalidArgument, "process name is required") + } + + if req.Timeout > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, time.Duration(req.Timeout)*time.Second) + defer cancel() + } + + var args []interface{} + if len(req.Args) > 0 { + if err := json.Unmarshal(req.Args, &args); err != nil { + return nil, status.Errorf(codes.InvalidArgument, "invalid args JSON: %v", err) + } + } + + p, err := process.Of(req.Process, args...) + if err != nil { + return nil, status.Errorf(codes.Internal, "process error: %v", err) + } + + p.WithContext(ctx) + injectAuth(p, ctx) + + if err := p.Execute(); err != nil { + if ctx.Err() == context.DeadlineExceeded { + return nil, status.Error(codes.DeadlineExceeded, "process execution timed out") + } + return nil, status.Errorf(codes.Internal, "process execution failed: %v", err) + } + defer p.Release() + + val := p.Value() + data, err := json.Marshal(val) + if err != nil { + return nil, status.Errorf(codes.Internal, "failed to marshal result: %v", err) + } + + return &pb.RunResponse{Data: data}, nil +} + +// injectAuth propagates AuthorizedInfo from the gRPC context into the Process. +func injectAuth(p *process.Process, ctx context.Context) { + authInfo := auth.GetAuthorizedInfo(ctx) + if authInfo == nil { + return + } + p.WithSID(authInfo.SessionID) + p.WithAuthorized(&process.AuthorizedInfo{ + Subject: authInfo.Subject, + ClientID: authInfo.ClientID, + Scope: authInfo.Scope, + SessionID: authInfo.SessionID, + UserID: authInfo.UserID, + TeamID: authInfo.TeamID, + TenantID: authInfo.TenantID, + }) +} diff --git a/grpc/run/run_test.go b/grpc/run/run_test.go new file mode 100644 index 00000000..fa64bfd2 --- /dev/null +++ b/grpc/run/run_test.go @@ -0,0 +1,155 @@ +package run_test + +import ( + "context" + "encoding/json" + "testing" + + "github.com/stretchr/testify/assert" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/yaoapp/yao/grpc/pb" + "github.com/yaoapp/yao/grpc/tests/testutils" +) + +func TestRun_ProcessExec(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.Run(ctx, &pb.RunRequest{ + Process: "utils.app.Ping", + }) + assert.NoError(t, err) + assert.NotNil(t, resp) + assert.NotEmpty(t, resp.Data) +} + +func TestRun_WithArgs(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + args, _ := json.Marshal([]interface{}{"hello", " world"}) + resp, err := client.Run(ctx, &pb.RunRequest{ + Process: "utils.str.Concat", + Args: args, + }) + assert.NoError(t, err) + if assert.NotNil(t, resp) { + assert.NotEmpty(t, resp.Data) + + var result string + err = json.Unmarshal(resp.Data, &result) + assert.NoError(t, err) + assert.Equal(t, "hello world", result) + } +} + +func TestRun_InvalidProcess(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.Run(ctx, &pb.RunRequest{Process: "nonexistent.process.here"}) + assert.Error(t, err) +} + +func TestRun_EmptyProcessName(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.Run(ctx, &pb.RunRequest{Process: ""}) + assert.Error(t, err) +} + +func TestRun_BadArgsJSON(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.Run(ctx, &pb.RunRequest{ + Process: "utils.app.Ping", + Args: []byte("{not-json"), + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestRun_WithTimeout(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.Run(ctx, &pb.RunRequest{ + Process: "utils.app.Ping", + Timeout: 30, + }) + assert.NoError(t, err) + assert.NotNil(t, resp) +} + +func TestRun_EmptyProcessName_StatusCode(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.Run(ctx, &pb.RunRequest{Process: ""}) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestRun_InvalidProcess_StatusCode(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.Run(ctx, &pb.RunRequest{Process: "nonexistent.process.here"}) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.Internal, st.Code()) +} + +func TestRun_NilArgs(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:run") + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.Run(ctx, &pb.RunRequest{ + Process: "utils.app.Ping", + Args: nil, + }) + assert.NoError(t, err) + assert.NotNil(t, resp) +} diff --git a/grpc/shell/shell.go b/grpc/shell/shell.go new file mode 100644 index 00000000..fe51e0bb --- /dev/null +++ b/grpc/shell/shell.go @@ -0,0 +1,91 @@ +package shell + +import ( + "bytes" + "context" + "errors" + "os" + "os/exec" + "syscall" + "time" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/yaoapp/yao/grpc/pb" +) + +const ( + defaultTimeout = 30 * time.Second + maxTimeout = 300 * time.Second +) + +// Handler implements the Shell gRPC method. +type Handler struct{} + +// Shell executes a system command in the host process and returns stdout/stderr/exit code. +func (h *Handler) Shell(ctx context.Context, req *pb.ShellRequest) (*pb.ShellResponse, error) { + if os.Getuid() == 0 { + return nil, status.Error(codes.PermissionDenied, "shell execution refused when running as root") + } + + if req.Command == "" { + return nil, status.Error(codes.InvalidArgument, "command is required") + } + + timeout := defaultTimeout + if req.Timeout > 0 { + timeout = time.Duration(req.Timeout) * time.Second + if timeout > maxTimeout { + timeout = maxTimeout + } + } + + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + cmd := exec.CommandContext(ctx, req.Command, req.Args...) + + if len(req.Env) > 0 { + env := os.Environ() + for k, v := range req.Env { + env = append(env, k+"="+v) + } + cmd.Env = env + } + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + err := cmd.Run() + + resp := &pb.ShellResponse{ + Stdout: stdout.Bytes(), + Stderr: stderr.Bytes(), + ExitCode: 0, + } + + if err != nil { + if ctx.Err() == context.DeadlineExceeded { + return nil, status.Error(codes.DeadlineExceeded, "command timed out") + } + + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { + if ws, ok := exitErr.Sys().(syscall.WaitStatus); ok { + resp.ExitCode = int32(ws.ExitStatus()) + } else { + resp.ExitCode = int32(exitErr.ExitCode()) + } + return resp, nil + } + + if errors.Is(err, exec.ErrNotFound) { + return nil, status.Errorf(codes.NotFound, "command not found: %s", req.Command) + } + return nil, status.Errorf(codes.Internal, "command execution failed: %v", err) + } + + return resp, nil +} diff --git a/grpc/shell/shell_test.go b/grpc/shell/shell_test.go new file mode 100644 index 00000000..c87c2aff --- /dev/null +++ b/grpc/shell/shell_test.go @@ -0,0 +1,166 @@ +package shell_test + +import ( + "context" + "runtime" + "testing" + + "github.com/stretchr/testify/assert" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/yaoapp/yao/grpc/pb" + "github.com/yaoapp/yao/grpc/tests/testutils" +) + +func TestShell_Echo(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:shell") + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.Shell(ctx, &pb.ShellRequest{ + Command: "echo", + Args: []string{"hello"}, + }) + assert.NoError(t, err) + assert.NotNil(t, resp) + assert.Contains(t, string(resp.Stdout), "hello") + assert.Equal(t, int32(0), resp.ExitCode) +} + +func TestShell_CommandNotFound(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:shell") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.Shell(ctx, &pb.ShellRequest{ + Command: "this_command_does_not_exist_xyz", + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.NotFound, st.Code()) +} + +func TestShell_Timeout(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("sleep command not available on Windows") + } + + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:shell") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.Shell(ctx, &pb.ShellRequest{ + Command: "sleep", + Args: []string{"10"}, + Timeout: 1, + }) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.DeadlineExceeded, st.Code()) +} + +func TestShell_EmptyCommand(t *testing.T) { + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:shell") + ctx := testutils.WithToken(context.Background(), token) + + _, err := client.Shell(ctx, &pb.ShellRequest{Command: ""}) + assert.Error(t, err) + st, _ := status.FromError(err) + assert.Equal(t, codes.InvalidArgument, st.Code()) +} + +func TestShell_NonZeroExit(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("false command not available on Windows") + } + + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:shell") + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.Shell(ctx, &pb.ShellRequest{ + Command: "false", + }) + assert.NoError(t, err) + assert.NotEqual(t, int32(0), resp.ExitCode) +} + +func TestShell_WithEnv(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("printenv not available on Windows") + } + + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:shell") + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.Shell(ctx, &pb.ShellRequest{ + Command: "printenv", + Args: []string{"TEST_GRPC_VAR"}, + Env: map[string]string{"TEST_GRPC_VAR": "grpc_value"}, + }) + assert.NoError(t, err) + assert.Contains(t, string(resp.Stdout), "grpc_value") +} + +func TestShell_MaxTimeoutCapped(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("echo not available on Windows") + } + + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:shell") + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.Shell(ctx, &pb.ShellRequest{ + Command: "echo", + Args: []string{"ok"}, + Timeout: 9999, + }) + assert.NoError(t, err) + assert.Contains(t, string(resp.Stdout), "ok") +} + +func TestShell_Stderr(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("bash not available on Windows") + } + + conn := testutils.Prepare(t) + defer testutils.Clean() + + client := testutils.NewClient(conn) + token := testutils.ObtainAccessToken(t, "grpc:shell") + ctx := testutils.WithToken(context.Background(), token) + + resp, err := client.Shell(ctx, &pb.ShellRequest{ + Command: "bash", + Args: []string{"-c", "echo error_msg >&2"}, + }) + assert.NoError(t, err) + assert.Contains(t, string(resp.Stderr), "error_msg") + assert.Equal(t, int32(0), resp.ExitCode) +} diff --git a/grpc/tests/testutils/testutils.go b/grpc/tests/testutils/testutils.go new file mode 100644 index 00000000..d7abab80 --- /dev/null +++ b/grpc/tests/testutils/testutils.go @@ -0,0 +1,220 @@ +package testutils + +import ( + "context" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/metadata" + + gouapi "github.com/yaoapp/gou/api" + "github.com/yaoapp/gou/model" + "github.com/yaoapp/gou/query" + "github.com/yaoapp/gou/query/gou" + "github.com/yaoapp/xun/capsule" + yaoagent "github.com/yaoapp/yao/agent" + "github.com/yaoapp/yao/agent/caller" + agentllm "github.com/yaoapp/yao/agent/llm" + "github.com/yaoapp/yao/config" + yaogrpc "github.com/yaoapp/yao/grpc" + _ "github.com/yaoapp/yao/grpc/auth" + "github.com/yaoapp/yao/grpc/pb" + "github.com/yaoapp/yao/kb" + "github.com/yaoapp/yao/openapi" + "github.com/yaoapp/yao/openapi/oauth" + "github.com/yaoapp/yao/service" + "github.com/yaoapp/yao/test" + + _ "github.com/yaoapp/gou/encoding" + _ "github.com/yaoapp/gou/text" + _ "github.com/yaoapp/yao/agent/assistant" +) + +// Prepare initializes the Yao runtime (DB, V8, models, stores, scripts), +// loads the OpenAPI server (which bootstraps oauth.OAuth and acl.Global), +// sets up the HTTP router for API proxy tests, +// then starts a real gRPC server on a random port. +// Returns a connected grpc.ClientConn ready to create service clients. +func Prepare(t *testing.T) *grpc.ClientConn { + t.Helper() + + cfg := config.Conf + cfg.GRPC.Port = 0 + cfg.GRPC.Host = "127.0.0.1" + cfg.GRPC.Enabled = "" + + test.Prepare(t, config.Conf) + + if openapi.Server == nil { + if _, err := openapi.Load(config.Conf); err != nil { + t.Fatalf("failed to load OpenAPI server: %v", err) + } + } + + // Load KB (required for agent KB features). + if _, err := kb.Load(config.Conf); err != nil { + t.Logf("warning: failed to load KB: %v", err) + } + + // Load agent DSL (required for AgentStream handler). + if yaoagent.GetAgent() == nil { + if err := yaoagent.Load(config.Conf); err != nil { + t.Logf("warning: failed to load agent DSL: %v", err) + } + } + + // Register JSAPI factories (idempotent, needed because Go init order is not guaranteed). + caller.SetJSAPIFactory() + agentllm.SetJSAPIFactory() + + // Register default query engine (required for DB search). + if _, has := query.Engines["default"]; !has && capsule.Global != nil { + query.Register("default", &gou.Query{ + Query: capsule.Query(), + GetTableName: func(s string) string { + if mod, has := model.Models[s]; has { + return mod.MetaData.Table.Name + } + return s + }, + AESKey: config.Conf.DB.AESKey, + }) + } + + // Set up the HTTP router so grpc/api can proxy requests internally. + if service.Router == nil { + router := gin.New() + if openapi.Server != nil { + gouapi.SetRoutes(router, openapi.Server.Config.BaseURL) + gouapi.BuildRouteTable() + openapi.Server.Attach(router) + } + service.Router = router + } + + if err := yaogrpc.StartServer(cfg); err != nil { + t.Fatalf("failed to start gRPC server: %v", err) + } + + addrs := yaogrpc.Addr() + if len(addrs) == 0 { + t.Fatal("gRPC server has no listen address") + } + + conn, err := grpc.NewClient(addrs[0], grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("failed to dial gRPC server: %v", err) + } + + return conn +} + +// Clean stops the gRPC server and tears down the Yao runtime. +func Clean() { + yaogrpc.Stop() + service.Router = nil + openapi.Server = nil + test.Clean() +} + +// Addr returns the gRPC server listen address. +func Addr() string { + addrs := yaogrpc.Addr() + if len(addrs) == 0 { + return "" + } + return addrs[0] +} + +// ObtainAccessToken mints a token with the given scopes via oauth.MakeAccessToken. +func ObtainAccessToken(t *testing.T, scopes ...string) string { + t.Helper() + svc := oauth.OAuth + if svc == nil { + t.Fatal("oauth service not initialized") + } + + scope := strings.Join(scopes, " ") + token, err := svc.MakeAccessToken("grpc-test", scope, "test-user", 3600) + if err != nil { + t.Fatalf("failed to make access token: %v", err) + } + return token +} + +// ObtainAccessTokenForUser mints a token for a specific user ID. +func ObtainAccessTokenForUser(t *testing.T, userID string, scopes ...string) string { + t.Helper() + svc := oauth.OAuth + if svc == nil { + t.Fatal("oauth service not initialized") + } + + scope := strings.Join(scopes, " ") + token, err := svc.MakeAccessToken("grpc-test", scope, userID, 3600) + if err != nil { + t.Fatalf("failed to make access token: %v", err) + } + return token +} + +// ObtainExpiredAccessToken mints an already-expired token (TTL=1s already elapsed). +func ObtainExpiredAccessToken(t *testing.T, scopes ...string) string { + t.Helper() + svc := oauth.OAuth + if svc == nil { + t.Fatal("oauth service not initialized") + } + + scope := strings.Join(scopes, " ") + token, err := svc.MakeAccessToken("grpc-test", scope, "test-user", -1) + if err != nil { + t.Fatalf("failed to make expired access token: %v", err) + } + return token +} + +// ObtainRefreshToken mints a refresh token. +func ObtainRefreshToken(t *testing.T, scopes ...string) string { + t.Helper() + svc := oauth.OAuth + if svc == nil { + t.Fatal("oauth service not initialized") + } + + scope := strings.Join(scopes, " ") + token, err := svc.MakeRefreshToken("grpc-test", scope, "test-user", 0) + if err != nil { + t.Fatalf("failed to make refresh token: %v", err) + } + return token +} + +// WithToken attaches a Bearer token to the context via gRPC metadata. +func WithToken(ctx context.Context, token string) context.Context { + return metadata.AppendToOutgoingContext(ctx, "authorization", "Bearer "+token) +} + +// WithRefreshToken attaches both Bearer and x-refresh-token to the context. +func WithRefreshToken(ctx context.Context, token, refreshToken string) context.Context { + return metadata.AppendToOutgoingContext(ctx, + "authorization", "Bearer "+token, + "x-refresh-token", refreshToken, + ) +} + +// WithSandboxMetadata attaches x-sandbox-id and x-grpc-upstream metadata. +func WithSandboxMetadata(ctx context.Context, sandboxID, upstream string) context.Context { + return metadata.AppendToOutgoingContext(ctx, + "x-sandbox-id", sandboxID, + "x-grpc-upstream", upstream, + ) +} + +// NewClient creates a pb.YaoClient from a connection. +func NewClient(conn *grpc.ClientConn) pb.YaoClient { + return pb.NewYaoClient(conn) +} diff --git a/openapi/oauth/authenticate.go b/openapi/oauth/authenticate.go new file mode 100644 index 00000000..288ff2ad --- /dev/null +++ b/openapi/oauth/authenticate.go @@ -0,0 +1,207 @@ +package oauth + +import ( + "errors" + "fmt" + "strings" + "time" + + "github.com/yaoapp/kun/log" + "github.com/yaoapp/yao/openapi/oauth/types" +) + +// AuthInput contains the raw tokens extracted from the transport layer +// (HTTP headers/cookies or gRPC metadata). No framework dependency. +type AuthInput struct { + AccessToken string + RefreshToken string + SessionID string +} + +// AuthResult holds the outcome of a successful authentication. +type AuthResult struct { + Claims *types.TokenClaims + Info *types.AuthorizedInfo + NewAccessToken string // non-empty when token refresh occurred + NewRefreshToken string // non-empty when token refresh occurred +} + +// AuthenticateToken performs token verification and optional refresh +// without any gin/HTTP dependency. The caller is responsible for +// extracting tokens from the transport and delivering refreshed tokens +// back to the client. +func (s *Service) AuthenticateToken(input AuthInput) (*AuthResult, error) { + token := input.AccessToken + + // API Key resolution (same as getAccessToken in guard.go) + if s.isAPIKey(token) { + token = s.getAccessTokenFromAPIKey(token) + } + token = strings.TrimPrefix(token, "Bearer ") + + if token == "" { + return nil, fmt.Errorf("%s", types.ErrTokenMissing.Error()) + } + + var newAccessToken, newRefreshToken string + + claims, err := s.VerifyToken(token) + if err != nil { + expiredClaims, expErr := s.VerifyTokenAllowExpired(token) + if expErr != nil || expiredClaims == nil { + return nil, fmt.Errorf("%s", types.ErrInvalidToken.Error()) + } + + if !expiredClaims.ExpiresAt.IsZero() && expiredClaims.ExpiresAt.Before(time.Now()) { + newClaims, access, refresh, refreshErr := s.refreshTokenDirect(input.RefreshToken, expiredClaims) + if refreshErr != nil { + if errors.Is(refreshErr, errRefreshInProgress) || errors.Is(refreshErr, errRefreshAlreadyDone) { + claims = expiredClaims + } else { + log.Error("[OAuth] Token refresh failed: %v", refreshErr) + return nil, fmt.Errorf("%s", types.ErrInvalidRefreshToken.Error()) + } + } else { + claims = newClaims + newAccessToken = access + newRefreshToken = refresh + } + } else { + return nil, fmt.Errorf("%s", types.ErrInvalidToken.Error()) + } + } + + info := s.buildAuthInfo(claims, input.SessionID) + + return &AuthResult{ + Claims: claims, + Info: info, + NewAccessToken: newAccessToken, + NewRefreshToken: newRefreshToken, + }, nil +} + +// refreshTokenDirect performs token rotation without any gin/HTTP dependency. +// It shares the same refreshGates concurrency control as TryRefreshToken. +// Returns (newClaims, newAccessToken, newRefreshToken, error). +func (s *Service) refreshTokenDirect(refreshToken string, expiredClaims *types.TokenClaims) (*types.TokenClaims, string, string, error) { + if refreshToken == "" { + return nil, "", "", fmt.Errorf("refresh token missing") + } + + gate := &refreshGate{done: make(chan struct{})} + if actual, loaded := refreshGates.LoadOrStore(refreshToken, gate); loaded { + existing := actual.(*refreshGate) + select { + case <-existing.done: + return nil, "", "", errRefreshAlreadyDone + default: + return nil, "", "", errRefreshInProgress + } + } + + defer func() { + close(gate.done) + time.AfterFunc(30*time.Second, func() { + refreshGates.CompareAndDelete(refreshToken, gate) + }) + }() + + refreshClaims, err := s.VerifyRefreshToken(refreshToken) + if err != nil { + return nil, "", "", fmt.Errorf("invalid or expired refresh token: %w", err) + } + + var accessTTL time.Duration + if expiredClaims != nil && !expiredClaims.IssuedAt.IsZero() && !expiredClaims.ExpiresAt.IsZero() { + accessTTL = expiredClaims.ExpiresAt.Sub(expiredClaims.IssuedAt) + } + if accessTTL <= 0 { + accessTTL = s.config.Token.AccessTokenLifetime + } + if accessTTL <= 0 { + accessTTL = time.Hour + } + + sourceClaims := expiredClaims + if sourceClaims == nil { + sourceClaims = refreshClaims + } + + extraClaims := sourceClaims.Extra + if extraClaims == nil { + extraClaims = make(map[string]interface{}) + } + if sourceClaims.TeamID != "" { + extraClaims["team_id"] = sourceClaims.TeamID + } + if sourceClaims.TenantID != "" { + extraClaims["tenant_id"] = sourceClaims.TenantID + } + + s.revokeRefreshToken(refreshToken) + + var refreshRemainingSeconds int + if !refreshClaims.ExpiresAt.IsZero() { + refreshRemainingSeconds = int(time.Until(refreshClaims.ExpiresAt).Seconds()) + if refreshRemainingSeconds <= 0 { + return nil, "", "", fmt.Errorf("refresh token already expired after revocation") + } + } else { + refreshTTL := s.config.Token.RefreshTokenLifetime + if refreshTTL == 0 { + refreshTTL = 24 * time.Hour + } + refreshRemainingSeconds = int(refreshTTL.Seconds()) + } + + newRefreshToken, err := s.MakeRefreshToken( + sourceClaims.ClientID, + sourceClaims.Scope, + sourceClaims.Subject, + refreshRemainingSeconds, + extraClaims, + ) + if err != nil { + return nil, "", "", fmt.Errorf("failed to issue new refresh token: %w", err) + } + + newTokenStr, err := s.MakeAccessToken( + sourceClaims.ClientID, + sourceClaims.Scope, + sourceClaims.Subject, + int(accessTTL.Seconds()), + extraClaims, + ) + if err != nil { + return nil, "", "", fmt.Errorf("failed to issue access token: %w", err) + } + + newClaims, err := s.VerifyToken(newTokenStr) + if err != nil { + return nil, "", "", fmt.Errorf("failed to verify refreshed token: %w", err) + } + + log.Info("[OAuth] Token rotated for subject %s (access + refresh)", sourceClaims.Subject) + return newClaims, newTokenStr, newRefreshToken, nil +} + +// buildAuthInfo constructs AuthorizedInfo directly from token claims, +// equivalent to the SetInfo+GetInfo round-trip through gin.Context. +func (s *Service) buildAuthInfo(claims *types.TokenClaims, sessionID string) *types.AuthorizedInfo { + info := &types.AuthorizedInfo{ + Subject: claims.Subject, + ClientID: claims.ClientID, + Scope: claims.Scope, + SessionID: sessionID, + TeamID: claims.TeamID, + TenantID: claims.TenantID, + } + + userID, err := s.UserID(claims.ClientID, claims.Subject) + if err == nil && userID != "" { + info.UserID = userID + } + + return info +} diff --git a/sandbox/DESIGN.md b/sandbox/DESIGN.md index 500a9f45..19fc29fd 100644 --- a/sandbox/DESIGN.md +++ b/sandbox/DESIGN.md @@ -1,1379 +1,305 @@ -# Sandbox Design +# Sandbox Refactoring Design -## 1. Overview +## Background -Sandbox provides **persistent Docker containers** as isolated execution environments for external CLI agents like Claude Code. +The current `sandbox.Manager` was built as a quick prototype for the Claude coding agent. It directly depends on the local Docker client, uses bind mounts for file IO, and Unix sockets for IPC. This limits it to single-node, local-only operation. -### Why Docker? +This document outlines the refactoring plan to make sandbox a production-grade, multi-node capable system built on top of the Tai SDK (`yao/tai`). -- **Persistence**: Claude installs dependencies (npm, pip, apt), which must persist across sessions -- **Cross-platform**: Works on Linux, macOS, and Windows -- **Strong isolation**: Process, filesystem, and network isolation -- **Mature ecosystem**: Well-documented, easy to maintain - -### Architecture +## Architecture ``` -┌─────────────────────────────────────────────────────────────────────┐ -│ Yao Server │ -│ │ -│ ┌────────────────────────────────────────────────────────────────┐ │ -│ │ Sandbox Manager │ │ -│ │ │ │ -│ │ containers: map[containerName]*Container │ │ -│ │ │ │ -│ │ - GetOrCreate(userID, chatID) → get or create container │ │ -│ │ - Exec(containerName, cmd) → execute command in container │ │ -│ │ - Stop(containerName) → stop container (preserve data) │ │ -│ │ - Remove(containerName) → delete container │ │ -│ │ │ │ -│ └──────────────────────────┬─────────────────────────────────────┘ │ -│ │ │ -│ ┌──────────────┼──────────────┐ │ -│ │ │ │ │ -│ ▼ ▼ ▼ │ -│ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │ -│ │ Container-A │ │ Container-B │ │ Container-C │ │ -│ │ (user1-chat1)│ │ (user1-chat2)│ │ (user2-chat1)│ │ -│ │ │ │ │ │ │ │ -│ │ - Claude CLI │ │ - Claude CLI │ │ - Claude CLI │ │ -│ │ - Node.js │ │ - Python │ │ - Go │ │ -│ │ - User code │ │ - User code │ │ - User code │ │ -│ └──────┬───────┘ └──────┬───────┘ └──────┬───────┘ │ -│ │ │ │ │ -│ ───────┴────────────────┴────────────────┴─────── │ -│ Unix Socket IPC │ -│ (one socket per container) │ -│ │ -└─────────────────────────────────────────────────────────────────────┘ +tai.Client Single connection to a Tai endpoint (or local Docker) + │ sandbox / volume / proxy / vnc low-level APIs + │ +sandbox.Manager Business layer + Lifecycle, user isolation, TTL, cleanup, IPC ``` ---- +`sandbox.Manager` takes a single `tai.Client` at construction time. Scaling is handled externally by the container runtime — K8s scheduler for pod placement and node scaling, Docker for single-host. The SDK does not manage multiple endpoints or do any scheduling. -## 2. Container Lifecycle +### Why Tai stays in the K8s path +K8s handles pod scheduling and container lifecycle, but it does **not** provide: + +| Capability | K8s native? | What you'd need without Tai | +|-----------|-------------|---------------------------| +| File sync to/from container | No | PVC + init container or sidecar | +| HTTP preview proxy | No | Ingress + Service per sandbox | +| VNC access | No | VNC sidecar + Service + Ingress | +| gRPC IPC relay (container → Yao) | No | Pod must reach Yao directly (network policy, Service) | + +Tai bundles all four behind a single endpoint. Bypassing Tai to "direct-connect" K8s only covers pod CRUD and exec — you'd still need to solve file IO, preview, VNC, and IPC separately, which means either deploying Tai anyway or assembling equivalent infrastructure from K8s primitives. + +The SDK's `NewK8s()` already supports direct kube-apiserver connection (pass empty `addr`), but this is only useful for bare compute scenarios with no file sync or web preview requirements. + +### sandbox.Manager (yao/sandbox) + +High-level business layer on top of `tai.Client`. Manages container lifecycle, user/session isolation, file operations, and IPC. + +**Responsibilities:** +- Create / get / start / stop / remove sandboxes +- Lifecycle policies: one-shot, session-bound, long-running, persistent +- Per-user and global container limits +- Idle timeout and cleanup +- File operations (via `tai.Client.Volume()` for remote, bind mount for local) +- IPC relay to Yao gRPC server + +### Yao gRPC Server (yao/grpc) + +General-purpose gRPC service exposed by the Yao process. Not limited to sandbox IPC — it exposes Yao's process execution capability to any gRPC client. + +**Clients:** +- Container-internal MCP tools (via Tai Gateway relay) +- `yao run --remote` CLI +- Other Yao instances (future node-to-node) + +**IPC path (replacing Unix socket):** ``` -Create ──────────► Running ──────────► Stopped ──────────► Removed -(docker create) (docker start) (docker stop) (docker rm) - │ │ │ │ - │ │ │ │ - ▼ ▼ ▼ ▼ -First request Execute tasks Idle timeout Cleanup policy - (persistent) (data preserved) (manual/scheduled) +Container process → yao-bridge (tai/bridge/) → Tai Gateway (:9100 gRPC) → Yao gRPC Server (:9099) + │ + process.Run(...) ``` -### Container States +Tai does **not** know Yao gRPC address at startup. The upstream is passed per-container via `CreateRequest.GRPCUpstream` — Tai records the mapping and routes relay traffic by source container. This keeps Tai stateless and allows one Tai to serve multiple Yao instances. -| State | Description | -| --------- | ---------------------------------------------- | -| `created` | Container created, not started | -| `running` | Container running, can execute commands | -| `stopped` | Container stopped, data preserved, can restart | -| `removed` | Container deleted | +## Authentication -### Naming Convention +The gRPC server reuses the existing `openapi/oauth` service — no new auth system needed. -``` -yao-sandbox-{userID}-{chatID} +### What already exists -Example: yao-sandbox-u123-c456 -``` +| Capability | Module | Reuse | Needs changes | +|------------|--------|-------|---------------| +| JWT sign (RS256) | `oauth.MakeAccessToken()` | Issue tokens for gRPC clients | None — supports custom scope/subject/extraClaims | +| JWT verify | `oauth.VerifyToken(token string)` | Validate Bearer token in interceptor | None — pure string input, no Gin dependency | +| Signing certs | `oauth.SigningCertificates` | Same keypair for HTTP and gRPC | None | +| Identity | `TokenClaims` (Subject/ClientID/Scope) | gRPC request context | None | +| Scope/ACL | `acl.Scope.Check(*AccessRequest)` | Method-level access control | None — only needs `(Method, Path, Scopes)`, no Gin dependency | +| Scope registration | `acl.Register(...)` | gRPC scopes via same pattern | None — add `grpc:*` scope definitions in `init()` | +| Client auth | `ClientProvider` | `client_credentials` grant for CLI/containers | None | +| Token revocation | `oauth.Revoke(ctx, token, hint)` | Container token cleanup | None | +| Device Flow scaffolding | `types.DeviceAuthorizationResponse`, `GrantTypeDeviceCode`, error codes | CLI `yao login` | Implement `DeviceAuthorization()` (currently stub) | ---- +**Key insight**: `authorized.SetInfo/GetInfo` are Gin-bound, but gRPC does NOT need them. The gRPC interceptor builds `AccessRequest` directly from JWT claims and calls `ScopeManager.Check` — bypasses the full `Enforce` chain (client/team/member), which is HTTP multi-tenant only. -## 3. IPC Communication +**Impact on existing code: zero.** All gRPC auth is purely additive (~80 lines interceptor + scope registration). Device Flow adds ~190 lines new code + ~10 lines to existing `Token()` switch. -### Problem - -Claude CLI runs inside the sandbox but needs to call Yao's MCP Tools (Yao Processes) which run outside. - -### Solution: Unix Socket + MCP JSON-RPC - -``` -┌────────────────────────────────────────────────────────────────────┐ -│ Docker Container │ -│ │ -│ ┌──────────────────────────────────────────────────────────────┐ │ -│ │ Claude CLI │ │ -│ │ │ │ │ -│ │ │ .mcp.json: "yao" → stdio │ │ -│ │ ▼ │ │ -│ │ ┌──────────────────────────────────────────────────────────┐ │ │ -│ │ │ yao-bridge (lightweight binary) │ │ │ -│ │ │ stdin/stdout ↔ /tmp/yao.sock │ │ │ -│ │ └──────────────────────────────────────────────────────────┘ │ │ -│ └──────────────────────────────────────────────────────────────┘ │ -│ │ │ -│ /tmp/yao.sock │ -│ │ │ -└──────────────────────────┼─────────────────────────────────────────┘ - │ - Unix Socket (bind mount) - │ -┌──────────────────────────┼─────────────────────────────────────────┐ -│ │ │ -│ {YAO_DATA_ROOT}/sandbox/ipc/{sessionID}.sock │ -│ │ │ -│ ┌───────────────────────▼──────────────────────────────────────┐ │ -│ │ IPC Server (goroutine) │ │ -│ │ │ │ -│ │ MCP JSON-RPC Methods: │ │ -│ │ - initialize → handshake │ │ -│ │ - tools/list → return authorized Yao MCP tools │ │ -│ │ - tools/call → execute process.New(name, args...) │ │ -│ │ - resources/list → list Yao resources │ │ -│ │ - resources/read → read Yao resource │ │ -│ │ │ │ -│ └───────────────────────────────────────────────────────────────┘ │ -│ │ -│ Yao Server │ -└────────────────────────────────────────────────────────────────────┘ -``` - -### Protocol - -- **Format**: MCP standard JSON-RPC 2.0 over NDJSON (newline-delimited JSON) -- **Transport**: Unix Socket -- **Bridge**: `yao-bridge` binary converts stdio ↔ socket - ---- - -## 4. Core Interfaces - -### 4.1 Sandbox Manager +### gRPC interceptor ```go -// sandbox/manager.go -package sandbox - -type Manager struct { - mu sync.Mutex // Protects creation - containers sync.Map // containerName → *Container - running int32 // Running container count - ipcManager *ipc.Manager - dockerClient *docker.Client - config *Config -} - -var ErrTooManyContainers = errors.New("sandbox: too many running containers, please try again later") - -type Config struct { - Image string // Docker image, default: yao/sandbox:latest - WorkspaceRoot string // Host workspace root directory - IPCDir string // IPC socket directory - MaxContainers int // Maximum concurrent containers - IdleTimeout time.Duration // Idle timeout before stopping container - MaxMemory string // Memory limit, e.g., "2g" - MaxCPU float64 // CPU limit, e.g., 1.0 -} - -type Container struct { - ID string - Name string // yao-sandbox-{userID}-{chatID} - UserID string - ChatID string - Status string // created, running, stopped - CreatedAt time.Time - LastUsedAt time.Time - IPCSession *ipc.Session +func authInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { + md, _ := metadata.FromIncomingContext(ctx) + token := extractBearer(md) + claims, err := oauth.OAuth.VerifyToken(token) + if err != nil { + return nil, status.Errorf(codes.Unauthenticated, "invalid token") + } + ctx = withClaims(ctx, claims) + return handler(ctx, req) } ``` -### 4.2 Manager Methods +`oauth.OAuth` is a global singleton initialized at Yao startup. The gRPC server simply references it — same signing keys, same token format, same user/client model. -```go -// GetOrCreate returns existing container or creates new one -// Returns ErrTooManyContainers if limit exceeded -func (m *Manager) GetOrCreate(ctx context.Context, userID, chatID string) (*Container, error) +### Token flow by client type -// Stream executes command and returns stdout reader -func (m *Manager) Stream(ctx context.Context, containerName string, cmd []string, opts *ExecOptions) (io.ReadCloser, error) +| Client | How it gets a token | +|--------|-------------------| +| Container MCP tool | `oauth.MakeAccessToken(clientID, "grpc:mcp grpc:run", userID, 900)` injected as `YAO_TOKEN` env var. `yao-bridge` auto-refreshes via `YAO_REFRESH_TOKEN`. Manager revokes refresh token on container Remove. | +| `yao run` CLI | `yao login` → OAuth Device Authorization Grant → token saved to `~/.yao/credentials`. Logged in = gRPC, not logged in = local. | +| Yao-to-Yao | Pre-shared service token or `client_credentials` | -// Exec executes command and waits for completion -func (m *Manager) Exec(ctx context.Context, containerName string, cmd []string, opts *ExecOptions) (*ExecResult, error) +## Network Security -// Stop stops container but preserves data -func (m *Manager) Stop(ctx context.Context, containerName string) error +### Yao ↔ Tai communication -// Start starts a stopped container -func (m *Manager) Start(ctx context.Context, containerName string) error +Tai exposes Docker Engine API (:2375), K8s API (:6443), gRPC Volume (:9100), HTTP proxy (:8080), and VNC (:6080). These are raw protocol proxies with **no built-in auth** — security is handled at the network layer. -// Remove deletes container and its data -func (m *Manager) Remove(ctx context.Context, containerName string) error +| Deployment | Strategy | +|-----------|----------| +| Same host (local) | Bind to `127.0.0.1` or Unix socket, no exposure | +| Same VPC / LAN | Firewall rules / security groups, private subnet only | +| Cross-network | VPN / WireGuard tunnel, or mTLS termination at Tai | -// List returns all containers for a user -func (m *Manager) List(ctx context.Context, userID string) ([]*Container, error) +### gRPC Server listen policy -// Cleanup stops idle containers -func (m *Manager) Cleanup(ctx context.Context) error +The Yao gRPC server (:9099) supports configurable listen address: + +| Scenario | Listen | Why | +|----------|--------|-----| +| Local dev | `127.0.0.1:9099` | Only local containers reach it | +| Production (same host) | `127.0.0.1:9099` | Tai on same machine forwards via loopback | +| Production (multi-node) | `0.0.0.0:9099` + IP allowlist | Remote Tai nodes need access | + +### IP allowlist (gRPC server) + +For multi-node deployment where gRPC must listen on `0.0.0.0`, the server should support an IP/CIDR allowlist: + +``` +YAO_GRPC_LISTEN=0.0.0.0:9099 +YAO_GRPC_ALLOW=10.0.0.0/8,172.16.0.0/12,192.168.0.0/16 ``` -### 4.3 Filesystem Methods +Enforcement is a simple gRPC interceptor that runs **before** the auth interceptor: ```go -// WriteFile writes content to a file in container -func (m *Manager) WriteFile(ctx context.Context, containerName, path string, content []byte) error - -// ReadFile reads content from a file in container -func (m *Manager) ReadFile(ctx context.Context, containerName, path string) ([]byte, error) - -// ListDir lists directory contents in container -func (m *Manager) ListDir(ctx context.Context, containerName, path string) ([]FileInfo, error) - -// Stat returns file info -func (m *Manager) Stat(ctx context.Context, containerName, path string) (*FileInfo, error) - -// MkDir creates directory in container -func (m *Manager) MkDir(ctx context.Context, containerName, path string) error - -// Remove removes file or directory in container -func (m *Manager) RemoveFile(ctx context.Context, containerName, path string) error - -// CopyToContainer copies file/directory from host to container -func (m *Manager) CopyToContainer(ctx context.Context, containerName, hostPath, containerPath string) error - -// CopyFromContainer copies file/directory from container to host -func (m *Manager) CopyFromContainer(ctx context.Context, containerName, containerPath, hostPath string) error - -// FileInfo represents file metadata -type FileInfo struct { - Name string - Path string - Size int64 - Mode os.FileMode - ModTime time.Time - IsDir bool -} -``` - -### 4.4 ExecOptions - -```go -type ExecOptions struct { - WorkDir string // Working directory inside container - Env map[string]string // Environment variables - Stdin io.Reader // Standard input - Timeout time.Duration // Execution timeout -} - -type ExecResult struct { - ExitCode int - Stdout string - Stderr string -} -``` - ---- - -## 5. IPC System - -### 5.1 IPC Session - -```go -// ipc/session.go -package ipc - -type Session struct { - ID string // Usually equals chatID - SocketPath string // {IPCDir}/{id}.sock - Listener net.Listener - Conn net.Conn - Context *AgentContext - MCPTools map[string]*MCPTool - cancel context.CancelFunc -} - -type AgentContext struct { - UserID string - ChatID string - Locale string -} - -type MCPTool struct { - Name string - Description string - Process string // Yao process name - InputSchema json.RawMessage // JSON Schema -} -``` - -### 5.2 IPC Manager - -```go -// ipc/manager.go -type Manager struct { - sessions sync.Map // sessionID → *Session - sockDir string // {YAO_DATA_ROOT}/sandbox/ipc/ -} - -// Create creates new IPC session -func (m *Manager) Create(ctx context.Context, sessionID string, agentCtx *AgentContext, mcpTools map[string]*MCPTool) (*Session, error) - -// Close closes IPC session and cleans up -func (m *Manager) Close(sessionID string) error - -// Get returns existing session -func (m *Manager) Get(sessionID string) (*Session, bool) -``` - -### 5.3 JSON-RPC Message Handling - -```go -// JSON-RPC request structure -type JSONRPCRequest struct { - JSONRPC string `json:"jsonrpc"` - ID interface{} `json:"id,omitempty"` - Method string `json:"method"` - Params json.RawMessage `json:"params,omitempty"` -} - -// JSON-RPC response structure -type JSONRPCResponse struct { - JSONRPC string `json:"jsonrpc"` - ID interface{} `json:"id,omitempty"` - Result interface{} `json:"result,omitempty"` - Error *JSONRPCError `json:"error,omitempty"` -} - -type JSONRPCError struct { - Code int `json:"code"` - Message string `json:"message"` - Data interface{} `json:"data,omitempty"` -} -``` - -### 5.4 Session Message Loop - -```go -func (s *Session) serve(ctx context.Context) { - defer s.cleanup() - - for { - select { - case <-ctx.Done(): - return - default: +func ipAllowInterceptor(allowedCIDRs []*net.IPNet) grpc.UnaryServerInterceptor { + return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { + peer, _ := peer.FromContext(ctx) + if !isAllowed(peer.Addr, allowedCIDRs) { + return nil, status.Errorf(codes.PermissionDenied, "ip not allowed") } - - conn, err := s.Listener.Accept() - if err != nil { - continue - } - s.Conn = conn - s.handleConnection(ctx, conn) - } -} - -func (s *Session) handleConnection(ctx context.Context, conn net.Conn) { - defer conn.Close() - - scanner := bufio.NewScanner(conn) - for scanner.Scan() { - select { - case <-ctx.Done(): - return - default: - } - - line := scanner.Text() - response := s.handleMessage(line) - if response != "" { - conn.Write([]byte(response + "\n")) - } - } -} - -func (s *Session) handleMessage(line string) string { - var req JSONRPCRequest - if err := json.Unmarshal([]byte(line), &req); err != nil { - return s.errorResponse(nil, -32700, "Parse error") - } - - switch req.Method { - case "initialize": - return s.handleInitialize(req) - case "initialized": - return "" // notification, no response - case "tools/list": - return s.handleListTools(req) - case "tools/call": - return s.handleCallTool(req) - case "resources/list": - return s.handleListResources(req) - case "resources/read": - return s.handleReadResource(req) - default: - return s.errorResponse(req.ID, -32601, "Method not found") + return handler(ctx, req) } } ``` -### 5.5 Tool Call Handler +Defense in depth: IP allowlist is the first gate, OAuth token is the second. Both must pass. + +### Tai side security + +Tai itself does not need auth — it trusts its network boundary. Recommended: +- Docker: Tai container runs on a private network, ports not exposed to public +- K8s: Tai runs as a DaemonSet or Deployment, service only accessible within cluster +- If Tai must be exposed, put it behind a reverse proxy (nginx/envoy) with mTLS or VPN + +### Tai high availability + +Tai is a single endpoint, but all its services except VNC WebSocket are stateless. Avoiding single-point-of-failure is a deployment concern, not an SDK concern. + +| Deployment | HA strategy | +|-----------|-------------| +| Docker single-host | Docker restart policy (`--restart=always`), Tai failure = transient | +| K8s Deployment | `replicas: N` + K8s Service load balancing, liveness probe on `/healthz` | +| K8s DaemonSet | One Tai per node, pod talks to local Tai via node-local Service | + +VNC uses WebSocket long connections — if Tai restarts, active VNC sessions drop and the client reconnects. Stateless services (K8s proxy, Docker proxy, Volume gRPC, HTTP proxy) recover transparently behind a Service. + +The SDK `tai.Client` connects to a single address. In K8s this address is a Service VIP — Tai replicas behind it are invisible to the SDK. + +## Container Lifecycle + +Lifecycle is managed by `sandbox.Manager`, not by tai.Client. + +| Policy | TTL | Behavior | +|--------|-----|----------| +| One-shot | 0 | Destroyed immediately after execution | +| Session | Minutes | Alive while user is active, cleaned up on idle timeout | +| Long-running | Hours/Days | User workspace, recoverable, cleaned up after extended idle | +| Persistent | None | User-managed, never auto-cleaned | + +## File Operations + +| Mode | tai.Client | File IO | +|------|-----------|---------| +| Local | `tai.New("")` | Bind mount, direct host filesystem | +| Remote | `tai.New("tai://host")` | `tai.Client.Volume()` via gRPC | + +Local mode preserves bind mount for performance. Remote mode uses `tai/volume` (gRPC + lz4 compression). `sandbox.Manager` routes based on `client.IsLocal()`. + +## Agent Layer Adaptation + +The agent layer (`agent/assistant`, `agent/sandbox`, `agent/context`) currently hardcodes local-only assumptions. It needs to be adapted to work with the new `sandbox.Manager` backed by `tai.Client`. + +### Current coupling + +``` +agent/assistant/sandbox.go + │ + ├─ GetSandboxManager() Global singleton, local Docker only + ├─ initSandbox() Creates executor, calls manager.GetOrCreate() + ├─ BuildMCPConfigForSandbox() Hardcodes /tmp/yao.sock for yao-bridge + └─ loadMCPToolsForIPC() Loads MCP tools, injects into IPC session + +agent/sandbox/claude/executor.go + │ + ├─ manager.GetOrCreate() Direct Docker container creation + ├─ manager.Stream() Docker exec + attach + └─ manager.Remove() Docker container removal + +agent/context/jsapi_sandbox.go + │ + ├─ ReadFile() Host filesystem via bind mount path translation + ├─ WriteFile() Docker CopyToContainer + └─ Exec() Docker exec +``` + +### What changes + +| Component | Before | After | +|-----------|--------|-------| +| `GetSandboxManager()` | Global singleton, `docker.NewClientWithOpts(FromEnv)` | Initialized with a `tai.Client` from Yao config | +| Container creation | `dockerClient.ContainerCreate()` | `tai.Client.Sandbox().Create()` | +| Container exec | `dockerClient.ContainerExecCreate/Start/Attach` | `tai.Client.Sandbox().Exec()` | +| File read | Host path via bind mount (`containerPathToHost`) | Local: bind mount (same). Remote: `tai.Client.Volume().Read()` | +| File write | `dockerClient.CopyToContainer` | Local: bind mount. Remote: `tai.Client.Volume().Write()` | +| IPC | Unix socket bind mount + yao-bridge | Local: Unix socket (same). Remote: Tai gRPC relay → Yao gRPC server | +| MCP config | `{args: ["/tmp/yao.sock"]}` hardcoded | Local: socket path from config. Remote: gRPC endpoint injected as env var | +| VNC | `vncproxy.NewProxy(nil)` local assumption | `tai.Client.VNC().URL()` | +| Cleanup | `dockerClient.ContainerRemove` | `tai.Client.Sandbox().Remove()` | + +### IPC migration detail + +**Local mode** (same host): Unix socket preserved — zero overhead, no change needed. + +**Remote mode** (via Tai): +``` +Container process → yao-bridge (tai/bridge/) → Tai relay (:9100 gRPC) → Yao gRPC Server +``` + +`yao-bridge` source lives in `yao/tai/bridge/` — it's a Tai SDK client (consumes Tai relay), shares gRPC deps with `tai/`, and is version-locked with the Tai protocol. Built via `go build ./tai/bridge/cmd/yao-bridge`. + +Bridge mode determined by env var: + +``` +YAO_IPC_MODE=socket YAO_IPC_ADDR=/tmp/yao.sock # local +YAO_IPC_MODE=grpc YAO_IPC_ADDR=tai-host:9100 # remote +``` + +In gRPC mode, bridge also reads `YAO_TOKEN` / `YAO_REFRESH_TOKEN` and handles automatic token refresh (see grpc/DESIGN.md Container token section). + +Tai relay upstream is NOT configured at Tai startup. Manager passes `GRPCUpstream` per-container in `CreateRequest` — Tai records the mapping and routes by source container. One Tai can serve containers from different Yao instances. + +`BuildMCPConfigForSandbox()` sets the env vars based on `client.IsLocal()`. + +### SandboxExecutor interface + +The `agent/context/jsapi_sandbox.go` `SandboxExecutor` interface stays the same — it's already abstract. Implementation behind it changes: ```go -func (s *Session) handleCallTool(req JSONRPCRequest) string { - var params struct { - Name string `json:"name"` - Arguments map[string]interface{} `json:"arguments"` - } - json.Unmarshal(req.Params, ¶ms) - - // Check authorization - tool, ok := s.MCPTools[params.Name] - if !ok { - return s.errorResponse(req.ID, -32602, "Tool not found or not authorized") - } - - // Execute Yao Process - proc := process.New(tool.Process, params.Arguments) - proc.WithContext(s.Context) - - if err := proc.Execute(); err != nil { - return s.toolErrorResponse(req.ID, params.Name, err) - } - defer proc.Release() - - result := proc.Value() - return s.toolSuccessResponse(req.ID, result) +type SandboxExecutor interface { + ReadFile(path string) ([]byte, error) + WriteFile(path string, data []byte) error + ListDir(path string) ([]FileInfo, error) + Exec(cmd string, args ...string) (string, error) + GetWorkDir() string + GetSandboxID() string + GetVNCUrl() (string, error) } ``` ---- +Hooks (`ctx.sandbox.ReadFile()`, etc.) work unchanged. The executor routes to bind mount or `tai.Client.Volume()` internally. -## 6. Docker Container Management +### Agent lifecycle policy -### 6.1 NewManager Constructor +Currently: sandbox created on chat start, removed on chat end (`defer sandboxCleanup`). -```go -func NewManager(config *Config) (*Manager, error) { - // Initialize Docker client - cli, err := docker.NewClientWithOpts(docker.FromEnv, docker.WithAPIVersionNegotiation()) - if err != nil { - return nil, fmt.Errorf("failed to create Docker client: %w", err) - } - - // Ping Docker to verify connection - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - if _, err := cli.Ping(ctx); err != nil { - return nil, fmt.Errorf("Docker not available: %w", err) - } - - // Ensure directories exist - os.MkdirAll(config.WorkspaceRoot, 0755) - os.MkdirAll(config.IPCDir, 0755) - - m := &Manager{ - dockerClient: cli, - config: config, - ipcManager: ipc.NewManager(config.IPCDir), - } - - // Start cleanup loop - go m.startCleanupLoop(context.Background()) - - return m, nil -} -``` - -### 6.2 GetOrCreate with Limit Check - -```go -func (m *Manager) GetOrCreate(ctx context.Context, userID, chatID string) (*Container, error) { - containerName := fmt.Sprintf("yao-sandbox-%s-%s", userID, chatID) - - // Check if container already exists (fast path) - if c, ok := m.containers.Load(containerName); ok { - container := c.(*Container) - container.LastUsedAt = time.Now() - return container, nil - } - - // Use mutex for creation to avoid race condition - m.mu.Lock() - defer m.mu.Unlock() - - // Double-check after acquiring lock - if c, ok := m.containers.Load(containerName); ok { - container := c.(*Container) - container.LastUsedAt = time.Now() - return container, nil - } - - // Check running container limit - if m.running >= int32(m.config.MaxContainers) { - return nil, ErrTooManyContainers - } - - // Create new container - container, err := m.createContainer(ctx, userID, chatID) - if err != nil { - return nil, err - } - - // Store and increment counter - m.containers.Store(containerName, container) - m.running++ - - return container, nil -} -``` - -### 6.3 Create Container (internal) - -```go -func (m *Manager) createContainer(ctx context.Context, userID, chatID string) (*Container, error) { - containerName := fmt.Sprintf("yao-sandbox-%s-%s", userID, chatID) - - // Ensure image exists, auto-pull if not - if err := m.ensureImage(ctx, m.config.Image); err != nil { - return nil, err - } - - // Workspace directory - workspaceHost := filepath.Join(m.config.WorkspaceRoot, userID, chatID) - os.MkdirAll(workspaceHost, 0755) - - // IPC socket path - sessionID := chatID - ipcSocketHost := filepath.Join(m.config.IPCDir, sessionID+".sock") - - // Create container - resp, err := m.dockerClient.ContainerCreate(ctx, - &container.Config{ - Image: m.config.Image, - Cmd: []string{"sleep", "infinity"}, // Keep running - WorkingDir: "/workspace", - Env: []string{ - "YAO_IPC_SOCKET=/tmp/yao.sock", - }, - }, - &container.HostConfig{ - Binds: []string{ - workspaceHost + ":/workspace", - ipcSocketHost + ":/tmp/yao.sock", - }, - Resources: container.Resources{ - Memory: parseMemory(m.config.MaxMemory), - NanoCPUs: int64(m.config.MaxCPU * 1e9), - }, - SecurityOpt: []string{"no-new-privileges"}, - CapDrop: []string{"ALL"}, - }, - nil, nil, containerName, - ) - - if err != nil { - return nil, err - } - - return &Container{ - ID: resp.ID, - Name: containerName, - UserID: userID, - ChatID: chatID, - Status: "created", - CreatedAt: time.Now(), - }, nil -} - -// ensureImage ensures the image exists locally, pulls if not -func (m *Manager) ensureImage(ctx context.Context, imageName string) error { - // Check if image exists locally - _, _, err := m.dockerClient.ImageInspectWithRaw(ctx, imageName) - if err == nil { - return nil // Image exists - } - - // Image not found, pull it - reader, err := m.dockerClient.ImagePull(ctx, imageName, image.PullOptions{}) - if err != nil { - return fmt.Errorf("failed to pull image %s: %w", imageName, err) - } - defer reader.Close() - - // Wait for pull to complete - io.Copy(io.Discard, reader) - return nil -} -``` - -### 6.4 Ensure Running - -```go -func (m *Manager) ensureRunning(ctx context.Context, containerName string) error { - c, ok := m.containers.Load(containerName) - if !ok { - return fmt.Errorf("container not found: %s", containerName) - } - cont := c.(*Container) - - if cont.Status == "running" { - return nil - } - - // Start the container - if err := m.dockerClient.ContainerStart(ctx, cont.ID, container.StartOptions{}); err != nil { - return err - } - - m.mu.Lock() - cont.Status = "running" - cont.LastUsedAt = time.Now() - m.mu.Unlock() - - return nil -} -``` - -### 6.5 Execute Command (Streaming) - -```go -func (m *Manager) Stream(ctx context.Context, containerName string, cmd []string, opts *ExecOptions) (io.ReadCloser, error) { - // Ensure container is running - if err := m.ensureRunning(ctx, containerName); err != nil { - return nil, err - } - - // Get container - c, _ := m.containers.Load(containerName) - cont := c.(*Container) - - // Create exec instance - execConfig := container.ExecOptions{ - Cmd: cmd, - WorkingDir: opts.WorkDir, - Env: mapToSlice(opts.Env), - AttachStdout: true, - AttachStderr: true, - } - - execResp, err := m.dockerClient.ContainerExecCreate(ctx, cont.ID, execConfig) - if err != nil { - return nil, err - } - - // Attach to exec - attachResp, err := m.dockerClient.ContainerExecAttach(ctx, execResp.ID, container.ExecAttachOptions{}) - if err != nil { - return nil, err - } - - return attachResp.Reader, nil -} - -// Exec executes command and waits for completion -func (m *Manager) Exec(ctx context.Context, containerName string, cmd []string, opts *ExecOptions) (*ExecResult, error) { - if opts == nil { - opts = &ExecOptions{} - } - - reader, err := m.Stream(ctx, containerName, cmd, opts) - if err != nil { - return nil, err - } - defer reader.Close() - - // Read all output - output, err := io.ReadAll(reader) - if err != nil { - return nil, err - } - - // TODO: Parse stdout/stderr from Docker multiplexed stream - // TODO: Get exit code from ContainerExecInspect - - return &ExecResult{ - ExitCode: 0, - Stdout: string(output), - Stderr: "", - }, nil -} -``` - -### 6.6 Filesystem Operations - -```go -// WriteFile writes content to a file in container using docker cp -func (m *Manager) WriteFile(ctx context.Context, containerName, path string, content []byte) error { - c, ok := m.containers.Load(containerName) - if !ok { - return fmt.Errorf("container not found: %s", containerName) - } - cont := c.(*Container) - // Create a tar archive with the file - var buf bytes.Buffer - tw := tar.NewWriter(&buf) - - hdr := &tar.Header{ - Name: filepath.Base(path), - Mode: 0644, - Size: int64(len(content)), - } - tw.WriteHeader(hdr) - tw.Write(content) - tw.Close() - - // Copy to container - return m.dockerClient.CopyToContainer(ctx, cont.ID, filepath.Dir(path), &buf, container.CopyToContainerOptions{}) -} - -// ReadFile reads content from a file in container -func (m *Manager) ReadFile(ctx context.Context, containerName, path string) ([]byte, error) { - c, ok := m.containers.Load(containerName) - if !ok { - return nil, fmt.Errorf("container not found: %s", containerName) - } - cont := c.(*Container) - - reader, _, err := m.dockerClient.CopyFromContainer(ctx, cont.ID, path) - if err != nil { - return nil, err - } - defer reader.Close() - - // Extract from tar - tr := tar.NewReader(reader) - _, err = tr.Next() - if err != nil { - return nil, err - } - - return io.ReadAll(tr) -} - -// ListDir lists directory contents -func (m *Manager) ListDir(ctx context.Context, containerName, path string) ([]FileInfo, error) { - result, err := m.Exec(ctx, containerName, []string{"ls", "-la", "--time-style=+%s", path}, nil) - if err != nil { - return nil, err - } - - return parseLS(result.Stdout), nil -} - -// Stat returns file info -func (m *Manager) Stat(ctx context.Context, containerName, path string) (*FileInfo, error) { - result, err := m.Exec(ctx, containerName, []string{"stat", "--format=%n|%s|%f|%Y|%F", path}, nil) - if err != nil { - return nil, err - } - return parseStat(result.Stdout), nil -} - -// MkDir creates directory in container -func (m *Manager) MkDir(ctx context.Context, containerName, path string) error { - _, err := m.Exec(ctx, containerName, []string{"mkdir", "-p", path}, nil) - return err -} - -// RemoveFile removes file or directory in container -func (m *Manager) RemoveFile(ctx context.Context, containerName, path string) error { - _, err := m.Exec(ctx, containerName, []string{"rm", "-rf", path}, nil) - return err -} - -// CopyToContainer copies from host to container -func (m *Manager) CopyToContainer(ctx context.Context, containerName, hostPath, containerPath string) error { - c, ok := m.containers.Load(containerName) - if !ok { - return fmt.Errorf("container not found: %s", containerName) - } - cont := c.(*Container) - - // Create tar archive from host path - archive, err := createTarFromPath(hostPath) - if err != nil { - return err - } - defer archive.Close() - - return m.dockerClient.CopyToContainer(ctx, cont.ID, containerPath, archive, container.CopyToContainerOptions{}) -} - -// CopyFromContainer copies from container to host -func (m *Manager) CopyFromContainer(ctx context.Context, containerName, containerPath, hostPath string) error { - c, ok := m.containers.Load(containerName) - if !ok { - return fmt.Errorf("container not found: %s", containerName) - } - cont := c.(*Container) - - reader, _, err := m.dockerClient.CopyFromContainer(ctx, cont.ID, containerPath) - if err != nil { - return err - } - defer reader.Close() - - return extractTarToPath(reader, hostPath) -} -``` - -### 6.7 Cleanup Strategy - -```go -func (m *Manager) startCleanupLoop(ctx context.Context) { - ticker := time.NewTicker(5 * time.Minute) - defer ticker.Stop() - - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - m.Cleanup(ctx) - } - } -} - -func (m *Manager) Cleanup(ctx context.Context) error { - now := time.Now() - - m.containers.Range(func(key, value interface{}) bool { - containerName := key.(string) - c := value.(*Container) - - // Stop idle containers - if c.Status == "running" && now.Sub(c.LastUsedAt) > m.config.IdleTimeout { - m.Stop(ctx, containerName) - } - - return true - }) - - return nil -} - -func (m *Manager) Start(ctx context.Context, containerName string) error { - return m.ensureRunning(ctx, containerName) -} - -func (m *Manager) Stop(ctx context.Context, containerName string) error { - c, ok := m.containers.Load(containerName) - if !ok { - return nil - } - cont := c.(*Container) - - if err := m.dockerClient.ContainerStop(ctx, cont.ID, container.StopOptions{}); err != nil { - return err - } - - // Update status, decrement running count - m.mu.Lock() - if cont.Status == "running" { - cont.Status = "stopped" - m.running-- - } - m.mu.Unlock() - - return nil -} - -func (m *Manager) Remove(ctx context.Context, containerName string) error { - // Stop first if running - m.Stop(ctx, containerName) - - c, ok := m.containers.Load(containerName) - if !ok { - return nil - } - cont := c.(*Container) - - if err := m.dockerClient.ContainerRemove(ctx, cont.ID, container.RemoveOptions{}); err != nil { - return err - } - - // Remove from map - m.containers.Delete(containerName) - - return nil -} - -func (m *Manager) List(ctx context.Context, userID string) ([]*Container, error) { - var result []*Container - prefix := fmt.Sprintf("yao-sandbox-%s-", userID) - - m.containers.Range(func(key, value interface{}) bool { - containerName := key.(string) - if strings.HasPrefix(containerName, prefix) { - result = append(result, value.(*Container)) - } - return true - }) - - return result, nil -} -``` - -### 6.8 Helper Functions - -```go -// mapToSlice converts map to []string for env vars -func mapToSlice(m map[string]string) []string { - if m == nil { - return nil - } - result := make([]string, 0, len(m)) - for k, v := range m { - result = append(result, k+"="+v) - } - return result -} - -// parseMemory converts string like "2g" to bytes -func parseMemory(s string) int64 { - // Implementation: parse "2g" → 2*1024*1024*1024 - // Use Docker's units package or manual parsing - return 0 // placeholder -} - -// parseLS parses ls -la output to []FileInfo -func parseLS(output string) []FileInfo { - // Implementation: parse ls output lines - return nil // placeholder -} - -// parseStat parses stat output to *FileInfo -func parseStat(output string) *FileInfo { - // Implementation: parse stat --format output - return nil // placeholder -} - -// createTarFromPath creates a tar archive from a host path -func createTarFromPath(hostPath string) (io.ReadCloser, error) { - // Implementation: walk directory, create tar entries - return nil, nil // placeholder -} - -// extractTarToPath extracts a tar archive to a host path -func extractTarToPath(reader io.Reader, hostPath string) error { - // Implementation: read tar entries, write to disk - return nil // placeholder -} -``` - ---- - -## 7. yao-bridge - -Lightweight binary inside container that bridges stdio to Unix socket. - -```go -// cmd/yao-bridge/main.go -package main - -import ( - "io" - "net" - "os" -) - -func main() { - if len(os.Args) < 2 { - os.Exit(1) - } - - sockPath := os.Args[1] - - // Connect to Unix socket - conn, err := net.Dial("unix", sockPath) - if err != nil { - os.Exit(1) - } - defer conn.Close() - - // stdin → socket - go func() { - io.Copy(conn, os.Stdin) - conn.(*net.UnixConn).CloseWrite() - }() - - // socket → stdout - io.Copy(os.Stdout, conn) -} -``` - -Build as static binary and include in Docker image. - ---- - -## 8. Docker Image - -### 8.1 Image Naming Convention - -``` -yao/sandbox-{tool}:{variant} - -Examples: - yao/sandbox-claude:latest # Claude CLI + Node.js + Python (default) - yao/sandbox-claude:full # + Go - yao/sandbox-cursor:latest # Cursor CLI + Node.js + Python (future) -``` - -### 8.2 Source Directory Structure - -``` -sandbox/ -├── docker/ -│ ├── base/ -│ │ └── Dockerfile.base # Common base image -│ ├── claude/ -│ │ ├── Dockerfile # Default: Claude + Node + Python -│ │ └── Dockerfile.full # + Go -│ ├── cursor/ # Future -│ │ └── Dockerfile -│ ├── build.sh -│ └── scripts/ -│ └── entrypoint.sh -├── bridge/ -│ └── main.go # yao-bridge source -├── ipc/ -│ ├── manager.go -│ └── session.go -├── manager.go -├── config.go -└── types.go -``` - -### 8.3 Base Image - -```dockerfile -# sandbox/docker/base/Dockerfile.base -FROM ubuntu:22.04 - -# Base tools -RUN apt-get update && apt-get install -y \ - curl \ - git \ - ca-certificates \ - && rm -rf /var/lib/apt/lists/* - -# yao-bridge (common to all tools) -COPY yao-bridge /usr/local/bin/yao-bridge -RUN chmod +x /usr/local/bin/yao-bridge - -# Working directory -WORKDIR /workspace - -# Non-root user -RUN useradd -m -s /bin/bash sandbox -USER sandbox - -CMD ["sleep", "infinity"] -``` - -### 8.4 Claude Tool Images - -```dockerfile -# sandbox/docker/claude/Dockerfile -# Default image: Claude CLI + Node.js + Python -FROM yao/sandbox-base:latest - -USER root - -# Claude CLI -RUN curl -fsSL https://claude.ai/install.sh | sh - -# Node.js 20 -RUN curl -fsSL https://deb.nodesource.com/setup_20.x | bash - \ - && apt-get install -y nodejs \ - && rm -rf /var/lib/apt/lists/* - -# Python 3.11 -RUN apt-get update && apt-get install -y \ - python3.11 \ - python3-pip \ - && rm -rf /var/lib/apt/lists/* - -USER sandbox -``` - -```dockerfile -# sandbox/docker/claude/Dockerfile.full -# Full image: + Go -FROM yao/sandbox-claude:latest - -USER root - -# Go 1.23 -RUN curl -fsSL https://go.dev/dl/go1.23.linux-amd64.tar.gz | tar -C /usr/local -xzf - \ - && ln -s /usr/local/go/bin/go /usr/local/bin/go - -USER sandbox -``` - -### 8.5 Build Script - -```bash -#!/bin/bash -# sandbox/docker/build.sh - -set -e - -TOOL=${1:-claude} - -# Build yao-bridge -cd ../bridge -CGO_ENABLED=0 go build -o ../docker/yao-bridge . -cd ../docker - -# Build base image -docker build -t yao/sandbox-base:latest -f base/Dockerfile.base . - -# Build tool-specific images -case $TOOL in - claude) - docker build -t yao/sandbox-claude:latest -f claude/Dockerfile . - docker build -t yao/sandbox-claude:full -f claude/Dockerfile.full . - ;; - cursor) - docker build -t yao/sandbox-cursor:latest -f cursor/Dockerfile . - ;; - all) - $0 claude - $0 cursor - ;; -esac - -echo "Images built for tool: $TOOL" -``` - -### 8.6 Image Variants - -| Image | Tool | Size | Pre-installed | -| --------------------------- | ------ | ------ | ----------------------------------- | -| `yao/sandbox-base:latest` | - | ~200MB | git, curl, yao-bridge | -| `yao/sandbox-claude:latest` | Claude | ~700MB | Claude CLI, Node.js 20, Python 3.11 | -| `yao/sandbox-claude:full` | Claude | ~1.3GB | + Go 1.23 | -| `yao/sandbox-cursor:latest` | Cursor | ~700MB | Cursor CLI, Node.js 20, Python 3.11 | - -Default: `yao/sandbox-claude:latest` (includes Node + Python) - ---- - -## 9. Configuration - -### 9.1 Environment Variables - -| Env Variable | Default | Description | -| -------------------------- | ----------------------------------- | ------------------------- | -| `YAO_SANDBOX_IMAGE` | `yao/sandbox-claude:latest` | Default Docker image | -| `YAO_SANDBOX_WORKSPACE` | `{YAO_DATA_ROOT}/sandbox/workspace` | Workspace root directory | -| `YAO_SANDBOX_IPC` | `{YAO_DATA_ROOT}/sandbox/ipc` | IPC socket directory | -| `YAO_SANDBOX_MAX` | `100` | Max concurrent containers | -| `YAO_SANDBOX_IDLE_TIMEOUT` | `30m` | Idle timeout | -| `YAO_SANDBOX_MEMORY` | `2g` | Default memory limit | -| `YAO_SANDBOX_CPU` | `1.0` | Default CPU limit | - -### 9.2 Go Config Struct - -```go -// sandbox/config.go -type Config struct { - Image string `json:"image,omitempty" env:"YAO_SANDBOX_IMAGE" envDefault:"yao/sandbox-claude:latest"` - WorkspaceRoot string `json:"workspace_root,omitempty" env:"YAO_SANDBOX_WORKSPACE"` - IPCDir string `json:"ipc_dir,omitempty" env:"YAO_SANDBOX_IPC"` - MaxContainers int `json:"max_containers,omitempty" env:"YAO_SANDBOX_MAX" envDefault:"100"` - IdleTimeout time.Duration `json:"idle_timeout,omitempty" env:"YAO_SANDBOX_IDLE_TIMEOUT" envDefault:"30m"` - MaxMemory string `json:"max_memory,omitempty" env:"YAO_SANDBOX_MEMORY" envDefault:"2g"` - MaxCPU float64 `json:"max_cpu,omitempty" env:"YAO_SANDBOX_CPU" envDefault:"1.0"` -} - -// Init sets defaults based on Yao config -func (c *Config) Init(dataRoot string) { - if c.WorkspaceRoot == "" { - c.WorkspaceRoot = filepath.Join(dataRoot, "sandbox", "workspace") - } - if c.IPCDir == "" { - c.IPCDir = filepath.Join(dataRoot, "sandbox", "ipc") - } -} -``` - -### 9.3 app.yao (optional override) +New: lifecycle policy set per-assistant config: ```yaml sandbox: - image: "yao/sandbox-claude:full" - max_memory: "4g" + lifecycle: session # one-shot | session | long-running | persistent + idle_timeout: 30m + image: yaoapp/workspace:latest ``` -### 9.4 Assistant-level Configuration (package.yao) +`initSandbox()` passes the policy to `sandbox.Manager`, which enforces TTL and cleanup. `sandboxCleanup()` only disconnects the executor — the Manager decides whether to actually remove the container based on policy. -```yaml -name: "My Coder" -type: claude +## Migration Path -sandbox: - image: "yao/sandbox-claude:full" # Override image - max_memory: "4g" # Override memory limit -``` - -### 9.5 Image Resolution - -``` -1. If package.yao sandbox.image is set → use it -2. Else if type is set → use yao/sandbox-{type}:latest -3. Else → use YAO_SANDBOX_IMAGE (or app.yao sandbox.image) -``` - ---- - -## 10. Data Persistence - -### Directory Structure - -``` -{YAO_DATA_ROOT}/sandbox/ -├── workspace/ -│ └── {userID}/ -│ ├── {chatID-1}/ # Mounted as /workspace in container -│ │ ├── .mcp.json # MCP configuration -│ │ ├── .claude/ # Claude configuration -│ │ │ └── skills/ # Skills symlink -│ │ ├── project/ # User project code -│ │ └── node_modules/ # Installed dependencies -│ │ -│ └── {chatID-2}/ -│ └── ... -│ -└── ipc/ - ├── {chatID-1}.sock # IPC socket - └── {chatID-2}.sock -``` - -### What Persists - -| Item | Location | Persists | -| ------------------ | --------------------- | ------------------------------ | -| User code | `/workspace/` | ✅ Yes (host mount) | -| Installed packages | Container filesystem | ✅ Yes (container persists) | -| Claude config | `/workspace/.claude/` | ✅ Yes | -| IPC socket | `/tmp/yao.sock` | ❌ No (recreated each session) | - ---- - -## 11. Stability Guarantees - -| Concern | Solution | -| ----------------------- | ----------------------------------------------------------------- | -| **Container isolation** | One container per user+chat | -| **IPC isolation** | One socket per session | -| **Resource limits** | Docker memory/CPU limits | -| **Idle cleanup** | Auto-stop after timeout (preserve data) | -| **Data persistence** | Workspace directory mount, container preserves installed packages | -| **Connection handling** | Goroutine detects EOF, auto-cleanup | -| **Concurrency safety** | sync.Map + dedicated goroutines | - ---- - -## 12. Claude Executor Integration - -### Execution Flow - -```go -func (e *ClaudeExecutor) Stream(ctx *context.Context, messages []context.Message, opts ...*context.Options) (*context.Response, error) { - - // 1. Get or create sandbox container - container, err := e.SandboxManager.GetOrCreate(ctx, ctx.User.ID, ctx.ChatID) - if err != nil { - return nil, fmt.Errorf("failed to get sandbox: %w", err) - } - - // 2. Create IPC session - mcpTools := e.getMCPTools() - ipcSession, err := e.IPCManager.Create(ctx, ctx.ChatID, &AgentContext{ - UserID: ctx.User.ID, - ChatID: ctx.ChatID, - Locale: ctx.Locale, - }, mcpTools) - if err != nil { - return nil, fmt.Errorf("failed to create IPC session: %w", err) - } - defer e.IPCManager.Close(ctx.ChatID) - - // 3. Generate .mcp.json - if err := e.writeMCPConfig(ctx, container); err != nil { - return nil, err - } - - // 4. Setup skills - if err := e.setupSkills(container); err != nil { - return nil, err - } - - // 5. Build Claude CLI arguments - args := e.buildArgs(ctx, messages) - - // 6. Execute Claude CLI in container - stdout, err := e.SandboxManager.Stream(ctx, container.Name, - append([]string{"claude"}, args...), - &ExecOptions{ - WorkDir: "/workspace", - Env: e.buildEnvMap(ctx), - Timeout: e.getTimeout(), - }, - ) - if err != nil { - return nil, err - } - defer stdout.Close() - - // 7. Parse stream-json output - return e.parseClaudeOutput(ctx, stdout) -} -``` - -### MCP Configuration Generation - -```go -func (e *ClaudeExecutor) writeMCPConfig(ctx context.Context, container *Container) error { - config := map[string]interface{}{ - "mcpServers": map[string]interface{}{ - "yao": map[string]interface{}{ - "command": "yao-bridge", - "args": []string{"/tmp/yao.sock"}, - }, - }, - } - - // Add other MCP servers (external stdio/sse) - for _, server := range e.Assistant.MCP.Servers { - if server.Transport != "process" { - config["mcpServers"].(map[string]interface{})[server.Name] = server.ToClaudeConfig() - } - } - - data, _ := json.MarshalIndent(config, "", " ") - return e.SandboxManager.WriteFile(ctx, container.Name, "/workspace/.mcp.json", data) -} -``` - ---- - -## 13. Security Considerations - -| Layer | Measures | -| -------------- | ------------------------------------------------------ | -| **Filesystem** | Only workspace mounted, host filesystem not accessible | -| **Network** | Can be restricted with `--network none` if needed | -| **Privileges** | `--cap-drop ALL`, `no-new-privileges` | -| **Resources** | Memory and CPU limits | -| **User** | Non-root user inside container | -| **IPC** | Per-session socket, authorized tools only | - ---- - -## 14. Summary - -| Aspect | Description | -| -------------------------- | ----------------------------------------- | -| **Core approach** | Persistent Docker containers | -| **Communication** | Unix Socket + MCP JSON-RPC | -| **Container granularity** | One container per user+chat | -| **Data persistence** | Workspace mount + container filesystem | -| **Dependency persistence** | npm/pip packages persist in container | -| **Security isolation** | Full isolation between users and sessions | -| **Resource control** | Memory, CPU, idle timeout | -| **Estimated code** | ~1200 lines | +1. **Phase 1:** Yao gRPC server — expose process execution, replace Unix socket IPC +2. **Phase 2:** `sandbox.Manager` refactoring — replace Docker client with `tai.Client`, unified file ops, new lifecycle model +3. **Phase 3:** Agent layer adaptation — executor uses new Manager, IPC mode switch, lifecycle policy +4. **Phase 4:** `yao run --remote` — CLI calls remote Yao via gRPC +5. **Phase 5:** Workspace persistence — browser preview, service exposure, delivery diff --git a/sandbox/SPEC.md b/sandbox/SPEC.md new file mode 100644 index 00000000..c7613eda --- /dev/null +++ b/sandbox/SPEC.md @@ -0,0 +1,579 @@ +# Sandbox Functional Specification + +Detailed interfaces, types, and behavior for the sandbox refactoring. +Architecture and rationale: see [DESIGN.md](./DESIGN.md). + +--- + +## 1. sandbox.Manager + +Replaces the current Docker-only Manager. Backed by a single `tai.Client`. + +### Config + +```go +type Config struct { + Image string // container image, default "yaoapp/workspace:latest" + MaxContainers int // global limit, default 100 + IdleTimeout time.Duration // default cleanup interval, default 30m + MaxMemory string // per-container, e.g. "2g" + MaxCPU float64 // per-container, e.g. 1.0 + ContainerWorkDir string // mount target inside container, default "/workspace" + ContainerUser string // empty = image default +} +``` + +Environment variable overrides remain the same (`YAO_SANDBOX_IMAGE`, etc.). `WorkspaceRoot` and `IPCDir` are removed — local paths derived from `tai.Client.IsLocal()` at runtime; remote mode uses `tai.Client.Volume()`. + +### Constructor + +```go +func NewManager(client *tai.Client, cfg *Config) (*Manager, error) +``` + +- Validates `client` is non-nil and healthy (calls `client.Sandbox().List()` as connectivity check) +- Starts background cleanup goroutine +- Returns ready Manager + +### Manager struct + +```go +type Manager struct { + client *tai.Client + config *Config + sandboxes sync.Map // name → *Sandbox + running atomic.Int32 + ipc *IPCRouter // local: Unix socket manager, remote: gRPC stub + cleanup *time.Ticker + done chan struct{} +} +``` + +### Public methods + +```go +// Lifecycle +func (m *Manager) GetOrCreate(ctx context.Context, opts GetOrCreateOptions) (*Sandbox, error) +func (m *Manager) Get(ctx context.Context, name string) (*Sandbox, error) +func (m *Manager) Start(ctx context.Context, name string) error +func (m *Manager) Stop(ctx context.Context, name string, timeout time.Duration) error +func (m *Manager) Remove(ctx context.Context, name string) error +func (m *Manager) List(ctx context.Context, filter ListFilter) ([]*Sandbox, error) + +// Execution +func (m *Manager) Exec(ctx context.Context, name string, cmd []string, opts ExecOptions) (*ExecResult, error) +func (m *Manager) Stream(ctx context.Context, name string, cmd []string, opts ExecOptions) (io.ReadCloser, error) +func (m *Manager) KillProcess(ctx context.Context, name string, pattern string) error + +// File operations (routes local/remote internally) +func (m *Manager) ReadFile(ctx context.Context, name string, path string) ([]byte, error) +func (m *Manager) WriteFile(ctx context.Context, name string, path string, data []byte) error +func (m *Manager) ListDir(ctx context.Context, name string, path string) ([]FileInfo, error) +func (m *Manager) Stat(ctx context.Context, name string, path string) (*FileInfo, error) +func (m *Manager) MkDir(ctx context.Context, name string, path string) error +func (m *Manager) RemoveFile(ctx context.Context, name string, path string) error +func (m *Manager) CopyToContainer(ctx context.Context, name string, hostPath, containerPath string) error +func (m *Manager) CopyFromContainer(ctx context.Context, name string, containerPath, hostPath string) error + +// Info +func (m *Manager) IsLocal() bool +func (m *Manager) Close() error +``` + +--- + +## 2. Types + +### Sandbox + +```go +type Sandbox struct { + Name string + UserID string + ChatID string + Image string + Status Status + Lifecycle Lifecycle + CreatedAt time.Time + LastUsedAt time.Time + IP string +} +``` + +### Status + +```go +type Status string + +const ( + StatusCreated Status = "created" + StatusRunning Status = "running" + StatusStopped Status = "stopped" +) +``` + +### Lifecycle + +```go +type Lifecycle string + +const ( + LifecycleOneShot Lifecycle = "one-shot" // destroyed after execution + LifecycleSession Lifecycle = "session" // alive while user active, idle timeout + LifecycleLongRunning Lifecycle = "long-running" // hours/days, recoverable + LifecyclePersistent Lifecycle = "persistent" // never auto-cleaned +) +``` + +### GetOrCreateOptions + +```go +type GetOrCreateOptions struct { + UserID string + ChatID string + Image string // override Config.Image + Lifecycle Lifecycle // default: LifecycleSession + Env map[string]string // injected into container + Cmd []string // override entrypoint + Memory string // override Config.MaxMemory + CPU float64 // override Config.MaxCPU +} +``` + +### ExecOptions / ExecResult + +```go +type ExecOptions struct { + WorkDir string + Env map[string]string + Stdin io.Reader + Timeout time.Duration +} + +type ExecResult struct { + ExitCode int + Stdout string + Stderr string +} +``` + +### ListFilter + +```go +type ListFilter struct { + UserID string // empty = all users + Status Status // empty = all statuses + Lifecycle Lifecycle // empty = all policies +} +``` + +### FileInfo + +```go +type FileInfo struct { + Name string + Path string + Size int64 + Mode os.FileMode + ModTime time.Time + IsDir bool +} +``` + +### Errors + +```go +var ( + ErrTooManyContainers = errors.New("sandbox: container limit reached") + ErrNotFound = errors.New("sandbox: not found") + ErrNotRunning = errors.New("sandbox: not running") + ErrAlreadyExists = errors.New("sandbox: already exists") +) +``` + +--- + +## 3. Lifecycle State Machine + +``` + GetOrCreate() + │ + ▼ + ┌─────────┐ + ┌────────│ Created │ + │ └────┬────┘ + │ Start() │ + │ ▼ + │ ┌─────────┐ idle timeout / Stop() + │ │ Running │──────────────────┐ + │ └────┬────┘ │ + │ │ ▼ + │ │ ┌─────────┐ + │ │ │ Stopped │ + │ │ └────┬────┘ + │ │ Start() │ + │ │ ┌───────────────────┘ + │ │ │ + │ ▼ ▼ + │ Remove() from any state + │ │ + │ ▼ + │ [Destroyed] + │ + └── one-shot: auto Remove() after Exec/Stream returns +``` + +### Cleanup rules + +| Lifecycle | Trigger | Action | +|-----------|---------|--------| +| one-shot | Exec/Stream completes | Manager.Remove() immediately | +| session | `IdleTimeout` since `LastUsedAt` | Manager.Stop() then Remove() | +| long-running | `IdleTimeout * 24` since `LastUsedAt` | Manager.Stop() (not removed, can restart) | +| persistent | Never | No automatic action | + +Background goroutine runs every `Config.IdleTimeout / 2`, scans `sandboxes`, applies rules. + +### Touch + +Every `Exec`, `Stream`, `ReadFile`, `WriteFile`, `ListDir` call updates `LastUsedAt`. + +--- + +## 4. File Operations Routing + +```go +func (m *Manager) ReadFile(ctx context.Context, name string, path string) ([]byte, error) { + if m.client.IsLocal() { + hostPath := m.hostPath(name, path) + return os.ReadFile(hostPath) + } + sessionID := m.sessionID(name) + ws := m.client.Workspace(sessionID) + f, err := ws.Open(path) + // ... read and return +} +``` + +| Operation | Local | Remote | +|-----------|-------|--------| +| ReadFile | `os.ReadFile(hostPath)` | `client.Workspace(session).Open(path)` | +| WriteFile | `os.WriteFile(hostPath)` | `client.Volume().Write(session, path, data)` | +| ListDir | `os.ReadDir(hostPath)` | `client.Volume().ReadDir(session, path)` | +| Stat | `os.Stat(hostPath)` | `client.Volume().Stat(session, path)` | +| MkDir | `os.MkdirAll(hostPath)` | `client.Volume().MkDir(session, path)` | +| RemoveFile | `os.RemoveAll(hostPath)` | `client.Volume().Remove(session, path)` | +| CopyToContainer | bind mount (noop, already on host) | `client.Volume().Write()` streamed | +| CopyFromContainer | bind mount (direct read) | `client.Volume().Read()` streamed | + +`hostPath` = `dataDir/{userID}/{chatID}/{containerRelativePath}` + +`sessionID` = `{userID}/{chatID}` (maps to volume session on Tai) + +--- + +## 5. IPC Router + +Abstracts local Unix socket vs remote gRPC relay. Manager creates the right one based on `client.IsLocal()`. + +### Interface + +```go +type IPCRouter interface { + Create(sessionID string, tools []MCPTool) (IPCSession, error) + Get(sessionID string) (IPCSession, error) + Close(sessionID string) error + CloseAll() error +} + +type IPCSession interface { + SetTools(tools []MCPTool) + SetContext(ctx *AgentContext) + SocketPath() string // local only, empty for remote + GRPCAddr() string // remote only, empty for local + Close() error +} +``` + +### Local implementation + +Same as current `ipc.Manager` — creates Unix socket per session, bind-mounts into container, yao-bridge connects to it. + +### Remote implementation + +No socket. Container receives `YAO_IPC_MODE=grpc` and `YAO_IPC_ADDR=tai-host:9100`. `yao-bridge` (`yao/tai/bridge/`) connects to Tai's gRPC relay, which forwards to Yao gRPC Server. Tai relay upstream is per-container via `CreateRequest.GRPCUpstream`, not a Tai startup parameter. + +Tool registration: remote IPCSession sends tool list to Yao gRPC Server via a registration RPC at session creation. + +### Container env injection + +```go +func (m *Manager) buildContainerEnv(session IPCSession, userEnv map[string]string) map[string]string { + env := maps.Clone(userEnv) + if m.client.IsLocal() { + env["YAO_IPC_MODE"] = "socket" + env["YAO_IPC_ADDR"] = session.SocketPath() + } else { + env["YAO_IPC_MODE"] = "grpc" + env["YAO_IPC_ADDR"] = session.GRPCAddr() + env["YAO_TOKEN"] = m.issueAccessToken(session) + env["YAO_REFRESH_TOKEN"] = m.issueRefreshToken(session) + } + return env +} + +// CreateRequest also carries GRPCUpstream for Tai relay routing (per-container, not per-Tai) + +``` + +--- + +## 6. Yao gRPC Server + +### Proto definition + +```protobuf +syntax = "proto3"; +package yao.v1; + +service Yao { + rpc Exec(ExecRequest) returns (ExecResponse); + rpc StreamExec(ExecRequest) returns (stream ExecChunk); + + // MCP tool registration (called by remote IPC sessions) + rpc RegisterTools(RegisterToolsRequest) returns (RegisterToolsResponse); + + // Health + rpc Healthz(HealthzRequest) returns (HealthzResponse); +} + +message ExecRequest { + string process = 1; // e.g. "models.user.Find" + bytes args = 2; // JSON-encoded arguments + string session = 3; // sandbox session ID for context +} + +message ExecResponse { + bytes result = 1; // JSON-encoded result + string error = 2; +} + +message ExecChunk { + bytes data = 1; + bool done = 2; +} + +message RegisterToolsRequest { + string session = 1; + repeated MCPToolDef tools = 2; +} + +message MCPToolDef { + string name = 1; + string description = 2; + string process = 3; // Yao process to call + bytes input_schema = 4; // JSON Schema +} + +message RegisterToolsResponse {} + +message HealthzRequest {} +message HealthzResponse { + string status = 1; +} +``` + +### Server startup + +```go +func StartGRPCServer(cfg GRPCConfig) (*grpc.Server, error) + +type GRPCConfig struct { + Listen string // "127.0.0.1:9099" or "0.0.0.0:9099" + AllowCIDR []string // IP allowlist, empty = no restriction +} +``` + +Interceptor chain: `ipAllowInterceptor` → `authInterceptor` → handler. + +### Exec handler + +```go +func (s *yaoServer) Exec(ctx context.Context, req *pb.ExecRequest) (*pb.ExecResponse, error) { + claims := claimsFromContext(ctx) + // ACL check: does this token have permission to call this process? + + p := process.New(req.Process) + var args []interface{} + json.Unmarshal(req.Args, &args) + + result, err := p.Exec(args...) + if err != nil { + return &pb.ExecResponse{Error: err.Error()}, nil + } + + data, _ := json.Marshal(result) + return &pb.ExecResponse{Result: data}, nil +} +``` + +--- + +## 7. Agent Layer + +### Assistant sandbox config + +```yaml +# assistants/coder.yao +sandbox: + enabled: true + lifecycle: session + idle_timeout: 30m + image: yaoapp/workspace:latest + command: claude + memory: "4g" + cpu: 2.0 +``` + +### Parsed config type + +```go +type AssistantSandboxConfig struct { + Enabled bool `json:"enabled"` + Lifecycle Lifecycle `json:"lifecycle"` + IdleTimeout time.Duration `json:"idle_timeout"` + Image string `json:"image"` + Command string `json:"command"` + Memory string `json:"memory"` + CPU float64 `json:"cpu"` +} +``` + +### Init flow (new) + +```go +func (a *Assistant) initSandbox(ctx context.Context) (*agentsandbox.Executor, error) { + mgr := GetSandboxManager() // global, initialized with tai.Client at Yao startup + + sb, err := mgr.GetOrCreate(ctx, sandbox.GetOrCreateOptions{ + UserID: a.userID, + ChatID: a.chatID, + Image: a.config.Sandbox.Image, + Lifecycle: a.config.Sandbox.Lifecycle, + Memory: a.config.Sandbox.Memory, + CPU: a.config.Sandbox.CPU, + }) + // ... + executor := agentsandbox.New(mgr, sb, a.config.Sandbox.Command) + return executor, nil +} +``` + +### Cleanup (new) + +```go +func (a *Assistant) sandboxCleanup(executor *agentsandbox.Executor) { + executor.Disconnect() + // Manager handles actual removal based on lifecycle policy. + // one-shot: already removed after Exec. + // session: will be cleaned up by background goroutine after idle timeout. + // long-running/persistent: stays. +} +``` + +### GetSandboxManager (new) + +```go +var ( + managerOnce sync.Once + manager *sandbox.Manager +) + +func GetSandboxManager() *sandbox.Manager { + managerOnce.Do(func() { + client := config.GetTaiClient() // initialized at Yao startup from env/config + mgr, err := sandbox.NewManager(client, loadSandboxConfig()) + if err != nil { + log.Fatal("sandbox manager init: %v", err) + } + manager = mgr + }) + return manager +} +``` + +### Executor factory + +```go +// agent/sandbox/executor.go +func New(mgr *sandbox.Manager, sb *sandbox.Sandbox, command string) Executor { + switch command { + case "claude": + return claude.NewExecutor(mgr, sb) + default: + return generic.NewExecutor(mgr, sb) + } +} +``` + +### Executor interface (unchanged) + +```go +type Executor interface { + Stream(ctx context.Context, opts StreamOptions) (io.ReadCloser, error) + Disconnect() error + + // Delegated to Manager internally + ReadFile(ctx context.Context, path string) ([]byte, error) + WriteFile(ctx context.Context, path string, data []byte) error + ListDir(ctx context.Context, path string) ([]FileInfo, error) + Exec(ctx context.Context, cmd []string) (string, error) + GetWorkDir() string + GetSandboxID() string + GetVNCUrl() string +} +``` + +Each method delegates to `mgr.ReadFile(ctx, sb.Name, path)` etc. The executor is a thin wrapper that knows the sandbox name. + +--- + +## 8. Naming Convention + +| Entity | Pattern | Example | +|--------|---------|---------| +| Container/Pod name | `yao-sb-{userID}-{chatID}` | `yao-sb-u123-c456` | +| Volume session | `{userID}/{chatID}` | `u123/c456` | +| IPC session | `{chatID}` | `c456` | +| Host workspace (local) | `{dataDir}/{userID}/{chatID}/` | `/data/u123/c456/` | + +Prefix shortened from `yao-sandbox-` to `yao-sb-` for K8s DNS name length limit (63 chars). + +--- + +## 9. Environment Variables + +### Yao process + +| Variable | Purpose | Default | +|----------|---------|---------| +| `YAO_TAI_ADDR` | Tai endpoint, e.g. `tai://10.0.0.1` or empty for local Docker | `""` (local) | +| `YAO_TAI_RUNTIME` | `docker` or `k8s` | `docker` | +| `YAO_TAI_KUBECONFIG` | Path to kubeconfig (K8s only) | | +| `YAO_TAI_NAMESPACE` | K8s namespace | `default` | +| `YAO_GRPC_LISTEN` | gRPC server listen address | `127.0.0.1:9099` | +| `YAO_GRPC_ALLOW` | CIDR allowlist, comma-separated | (empty = no filter) | +| `YAO_SANDBOX_IMAGE` | Default container image | `yaoapp/workspace:latest` | +| `YAO_SANDBOX_MAX` | Max containers | `100` | +| `YAO_SANDBOX_IDLE_TIMEOUT` | Idle timeout duration | `30m` | +| `YAO_SANDBOX_MEMORY` | Memory limit | `2g` | +| `YAO_SANDBOX_CPU` | CPU limit | `1.0` | + +### Container-internal + +| Variable | Purpose | Set by | +|----------|---------|--------| +| `YAO_IPC_MODE` | `socket` or `grpc` | Manager at creation | +| `YAO_IPC_ADDR` | Socket path or gRPC host:port | Manager at creation | +| `YAO_TOKEN` | JWT access token for gRPC auth (remote only, short TTL 15m) | Manager at creation | +| `YAO_REFRESH_TOKEN` | JWT refresh token (remote only, no expiry, revoked on Remove) | Manager at creation | diff --git a/service/service.go b/service/service.go index 3af013c1..59f21df2 100644 --- a/service/service.go +++ b/service/service.go @@ -11,6 +11,10 @@ import ( "github.com/yaoapp/yao/share" ) +// Router holds the active gin.Engine so the gRPC API proxy can forward +// requests internally without an HTTP round-trip. +var Router *gin.Engine + // Start the yao service func Start(cfg config.Config) (*http.Server, error) { @@ -24,6 +28,7 @@ func Start(cfg config.Config) (*http.Server, error) { } router := gin.New() + Router = router router.Use(Middlewares...) var apiRoot string @@ -68,6 +73,7 @@ func Start(cfg config.Config) (*http.Server, error) { // Restart the yao service func Restart(srv *http.Server, cfg config.Config) error { router := gin.New() + Router = router router.Use(Middlewares...) if openapi.Server != nil { From daa4da763bcd601a6e410bafe1979ac918140a3d Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 4 Mar 2026 14:58:03 +0800 Subject: [PATCH 3/7] 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]) + } +} From 1c79908649d60453146e023d57df2dcc676ed74d Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 4 Mar 2026 15:19:48 +0800 Subject: [PATCH 4/7] Enhance OAuth Device Flow implementation - Add support for the OAuth Device Authorization Flow (RFC 8628) in the OpenAPI service, allowing devices with limited input capabilities to obtain authorization. - Implement `DeviceAuthorization()` and `AuthorizeDevice()` methods to handle device and user code generation, storage, and authorization. - Update the OAuth endpoints to include `/device/authorize` for user code authorization and fix the discovery endpoint path for device authorization. - Introduce MongoDB service in CI workflows for testing and enhance the unit test workflow with Redis setup. - Update Go module dependencies to include necessary packages for the new features. This commit significantly advances the OAuth capabilities of the application, enabling a more flexible authorization process for devices. --- .github/workflows/pr-test.yml | 27 +++ .github/workflows/unit-test.yml | 27 +++ engine/machine.go | 88 +++++++++ engine/machine_darwin.go | 24 +++ engine/machine_linux.go | 21 +++ engine/machine_test.go | 57 ++++++ engine/machine_windows.go | 21 +++ go.mod | 2 +- grpc/IMPL.md | 16 +- openapi/oauth.go | 85 ++++++++- openapi/oauth/core.go | 87 +++++++++ openapi/oauth/device.go | 112 ++++++++++- openapi/oauth/discovery.go | 2 +- openapi/oauth/oauth.go | 11 ++ openapi/oauth/token.go | 118 ++++++++++++ openapi/tests/oauth/device_test.go | 292 +++++++++++++++++++++++++++++ 16 files changed, 968 insertions(+), 22 deletions(-) create mode 100644 engine/machine.go create mode 100644 engine/machine_darwin.go create mode 100644 engine/machine_linux.go create mode 100644 engine/machine_test.go create mode 100644 engine/machine_windows.go create mode 100644 openapi/tests/oauth/device_test.go diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index c3c8ab5f..c8203e46 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -1537,6 +1537,16 @@ jobs: # ============================================================================= TaiTest: runs-on: ubuntu-latest + services: + mongodb: + image: mongo:6.0 + ports: + - 27017:27017 + env: + MONGO_INITDB_ROOT_USERNAME: root + MONGO_INITDB_ROOT_PASSWORD: 123456 + MONGO_INITDB_DATABASE: test + strategy: matrix: go: ["1.25"] @@ -1647,11 +1657,28 @@ jobs: with: ref: ${{ env.HEAD }} + - name: Setup Apple Private Key + run: | + mkdir -p ../app/openapi/certs/apple + echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8 + - name: Setup Go ${{ matrix.go }} uses: actions/setup-go@v5 with: go-version: ${{ matrix.go }} + - name: Start Redis + run: docker run --name redis --publish 6379:6379 --detach redis:6 + + - name: Setup Go Tools + run: make tools + + - name: Setup ENV (SQLite) + run: | + mkdir -p ${{ github.WORKSPACE }}/../app/db + echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV + echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV + - name: Pull Tai & Test Images run: | docker pull yaoapp/tai:latest diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index 1c5b36a4..e7432c52 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -1139,6 +1139,16 @@ jobs: # ============================================================================= tai-test: runs-on: ubuntu-latest + services: + mongodb: + image: mongo:6.0 + ports: + - 27017:27017 + env: + MONGO_INITDB_ROOT_USERNAME: root + MONGO_INITDB_ROOT_PASSWORD: 123456 + MONGO_INITDB_DATABASE: test + strategy: matrix: go: ["1.25"] @@ -1201,11 +1211,28 @@ jobs: - name: Checkout Code uses: actions/checkout@v4 + - name: Setup Apple Private Key + run: | + mkdir -p ../app/openapi/certs/apple + echo "${{ secrets.APPLE_PRIVATE_KEY_USER }}" > ../app/openapi/certs/apple/signin_client_secret_key.p8 + - name: Setup Go ${{ matrix.go }} uses: actions/setup-go@v5 with: go-version: ${{ matrix.go }} + - name: Start Redis + run: docker run --name redis --publish 6379:6379 --detach redis:6 + + - name: Setup Go Tools + run: make tools + + - name: Setup ENV (SQLite) + run: | + mkdir -p ${{ github.WORKSPACE }}/../app/db + echo "YAO_DB_DRIVER=sqlite3" >> $GITHUB_ENV + echo "YAO_DB_PRIMARY=${{ github.WORKSPACE }}/../app/db/yao.db" >> $GITHUB_ENV + - name: Pull Tai & Test Images run: | docker pull yaoapp/tai:latest diff --git a/engine/machine.go b/engine/machine.go new file mode 100644 index 00000000..7d205de4 --- /dev/null +++ b/engine/machine.go @@ -0,0 +1,88 @@ +package engine + +import ( + "crypto/sha256" + "fmt" + "net" + "os" + "runtime" + "strings" + "sync" + + "github.com/yaoapp/gou/process" +) + +// MachineInfo contains deterministic machine identification. +type MachineInfo struct { + ID string `json:"id"` // "yao-cli-{hash32}" deterministic client ID + Hostname string `json:"hostname"` // OS hostname + Platform string `json:"platform"` // runtime.GOOS: "darwin", "linux", "windows" +} + +var ( + cachedMachineInfo *MachineInfo + machineOnce sync.Once + machineErr error +) + +func init() { + process.Register("utils.app.MachineID", processMachineID) +} + +// GetMachineID returns a deterministic machine fingerprint. +// The result is cached after the first call. +func GetMachineID() (*MachineInfo, error) { + machineOnce.Do(func() { + cachedMachineInfo, machineErr = computeMachineID() + }) + return cachedMachineInfo, machineErr +} + +func computeMachineID() (*MachineInfo, error) { + hostname, _ := os.Hostname() + + raw, err := platformMachineID() + if err != nil || strings.TrimSpace(raw) == "" { + raw = fallbackMachineID(hostname) + } + + hash := sha256.Sum256([]byte(raw)) + id := fmt.Sprintf("yao-cli-%x", hash[:16]) // 32 hex chars + + return &MachineInfo{ + ID: id, + Hostname: hostname, + Platform: runtime.GOOS, + }, nil +} + +func fallbackMachineID(hostname string) string { + mac := firstHardwareAddr() + return hostname + ":" + mac +} + +func firstHardwareAddr() string { + ifaces, err := net.Interfaces() + if err != nil { + return "unknown" + } + for _, iface := range ifaces { + if iface.Flags&net.FlagLoopback != 0 || len(iface.HardwareAddr) == 0 { + continue + } + return iface.HardwareAddr.String() + } + return "unknown" +} + +func processMachineID(p *process.Process) interface{} { + info, err := GetMachineID() + if err != nil { + return map[string]interface{}{"error": err.Error()} + } + return map[string]interface{}{ + "id": info.ID, + "hostname": info.Hostname, + "platform": info.Platform, + } +} diff --git a/engine/machine_darwin.go b/engine/machine_darwin.go new file mode 100644 index 00000000..62535e4a --- /dev/null +++ b/engine/machine_darwin.go @@ -0,0 +1,24 @@ +//go:build darwin + +package engine + +import ( + "os/exec" + "strings" +) + +func platformMachineID() (string, error) { + out, err := exec.Command("ioreg", "-rd1", "-c", "IOPlatformExpertDevice").Output() + if err != nil { + return "", err + } + for _, line := range strings.Split(string(out), "\n") { + if strings.Contains(line, "IOPlatformUUID") { + parts := strings.SplitN(line, `"`, 4) + if len(parts) >= 4 { + return strings.TrimSpace(parts[3]), nil + } + } + } + return "", nil +} diff --git a/engine/machine_linux.go b/engine/machine_linux.go new file mode 100644 index 00000000..08520e3a --- /dev/null +++ b/engine/machine_linux.go @@ -0,0 +1,21 @@ +//go:build linux + +package engine + +import ( + "os" + "strings" +) + +func platformMachineID() (string, error) { + for _, path := range []string{"/etc/machine-id", "/var/lib/dbus/machine-id"} { + data, err := os.ReadFile(path) + if err == nil { + id := strings.TrimSpace(string(data)) + if id != "" { + return id, nil + } + } + } + return "", nil +} diff --git a/engine/machine_test.go b/engine/machine_test.go new file mode 100644 index 00000000..3e008925 --- /dev/null +++ b/engine/machine_test.go @@ -0,0 +1,57 @@ +package engine + +import ( + "strings" + "testing" +) + +func TestGetMachineID_Deterministic(t *testing.T) { + info1, err := GetMachineID() + if err != nil { + t.Fatalf("GetMachineID() returned error: %v", err) + } + + info2, err := GetMachineID() + if err != nil { + t.Fatalf("GetMachineID() second call returned error: %v", err) + } + + if info1.ID != info2.ID { + t.Errorf("GetMachineID() not deterministic: %q != %q", info1.ID, info2.ID) + } +} + +func TestGetMachineID_Format(t *testing.T) { + info, err := GetMachineID() + if err != nil { + t.Fatalf("GetMachineID() returned error: %v", err) + } + + if !strings.HasPrefix(info.ID, "yao-cli-") { + t.Errorf("ID should have prefix 'yao-cli-', got %q", info.ID) + } + + // "yao-cli-" (8) + 32 hex chars = 40 + if len(info.ID) != 40 { + t.Errorf("ID should be 40 chars, got %d: %q", len(info.ID), info.ID) + } + + if info.Hostname == "" { + t.Error("Hostname should not be empty") + } + + if info.Platform == "" { + t.Error("Platform should not be empty") + } +} + +func TestGetMachineID_NonEmpty(t *testing.T) { + info, err := GetMachineID() + if err != nil { + t.Fatalf("GetMachineID() returned error: %v", err) + } + + if info.ID == "" { + t.Error("ID should not be empty") + } +} diff --git a/engine/machine_windows.go b/engine/machine_windows.go new file mode 100644 index 00000000..6490f215 --- /dev/null +++ b/engine/machine_windows.go @@ -0,0 +1,21 @@ +//go:build windows + +package engine + +import ( + "golang.org/x/sys/windows/registry" +) + +func platformMachineID() (string, error) { + k, err := registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\Microsoft\Cryptography`, registry.READ|registry.WOW64_64KEY) + if err != nil { + return "", err + } + defer k.Close() + + val, _, err := k.GetStringValue("MachineGuid") + if err != nil { + return "", err + } + return val, nil +} diff --git a/go.mod b/go.mod index 9d69fa70..8290777e 100644 --- a/go.mod +++ b/go.mod @@ -51,6 +51,7 @@ require ( go.mongodb.org/mongo-driver v1.17.3 golang.org/x/crypto v0.48.0 golang.org/x/net v0.50.0 + golang.org/x/sys v0.41.0 golang.org/x/text v0.34.0 google.golang.org/grpc v1.78.0 google.golang.org/protobuf v1.36.11 @@ -235,7 +236,6 @@ require ( golang.org/x/mod v0.33.0 // indirect golang.org/x/oauth2 v0.32.0 // indirect golang.org/x/sync v0.19.0 // indirect - golang.org/x/sys v0.41.0 // indirect golang.org/x/term v0.40.0 // indirect golang.org/x/time v0.14.0 // indirect golang.org/x/tools v0.42.0 // indirect diff --git a/grpc/IMPL.md b/grpc/IMPL.md index adad4644..578d2553 100644 --- a/grpc/IMPL.md +++ b/grpc/IMPL.md @@ -197,18 +197,22 @@ Deliverable: `go build -o yao-grpc ./tai/grpc/cmd`. Depends on: Phase 1. Three sub-phases with sequential dependency: 6.1 → 6.2 → 6.3. -#### Phase 6.1: OAuth Device Flow backend ⏳ +#### 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/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 | +| `engine/machine.go` + platform files | `GetMachineID()` Go API + `utils.app.MachineID` process — cross-platform (macOS/Linux/Windows) deterministic machine fingerprint | ✅ Done | +| `oauth/token.go` | `deviceCodeKey`, `userCodeKey`, `storeDeviceCode`, `getDeviceCodeData`, `authorizeDeviceCode`, `consumeDeviceCode` — device_code + user_code storage/retrieval/consumption helpers | ✅ Done | +| `oauth/device.go` | Implement `DeviceAuthorization()` + `AuthorizeDevice()` + `generateUserCode()` — generate codes (gonanoid, XXXX-XXXX format), validate client + grant type, store, return `DeviceAuthorizationResponse` | ✅ Done | +| `oauth/core.go` | Add `case types.GrantTypeDeviceCode` → `handleDeviceCodeGrant()` — poll returns `authorization_pending` / `expired_token` / token | ✅ Done | +| `openapi/oauth.go` | Replace stub `oauthDeviceAuthorization` handler → call `DeviceAuthorization()`. Add `POST /oauth/device/authorize` → `oauthDeviceAuthorize` (bearer token + user_code → authorize device). | ✅ Done | +| `oauth/discovery.go` | Fix path: `/oauth/device` → `/oauth/device_authorization` | ✅ Done | +| `oauth/oauth.go` | Config defaults: `DeviceCodeLength=8`, `UserCodeLength=8`, `DeviceCodeInterval=5s`, `DeviceFlowEnabled=true`, `DynamicClientRegistrationEnabled=true` | ✅ Done | +| `openapi/tests/oauth/device_test.go` | Full test suite: device auth success/error, token polling (pending/invalid), end-to-end flow | ✅ Done | -Deliverable: Device flow endpoints functional — `POST /oauth/device_authorization` issues codes, `POST /oauth/token` with `grant_type=device_code` polls status. +Deliverable: Device flow endpoints functional — `POST /oauth/device_authorization` issues codes, `POST /oauth/token` with `grant_type=device_code` polls status. `POST /oauth/device/authorize` allows authenticated user to authorize device. #### Phase 6.2: CUI auth/device page (frontend) ⏳ diff --git a/openapi/oauth.go b/openapi/oauth.go index f10e466e..f306f6c9 100644 --- a/openapi/oauth.go +++ b/openapi/oauth.go @@ -58,6 +58,7 @@ func (openapi *OpenAPI) attachOAuth(base *gin.RouterGroup) { // Device Authorization Flow - RFC 8628 oauth.POST("/device_authorization", openapi.oauthDeviceAuthorization) + oauth.POST("/device/authorize", openapi.oauthDeviceAuthorize) // Pushed Authorization Request - RFC 9126 oauth.POST("/par", openapi.oauthPushedAuthorizationRequest) @@ -523,22 +524,86 @@ func (openapi *OpenAPI) oauthDeleteClient(c *gin.Context) { // oauthDeviceAuthorization handles device authorization - RFC 8628 func (openapi *OpenAPI) oauthDeviceAuthorization(c *gin.Context) { clientID := c.PostForm("client_id") - if clientID == "" { - response.RespondWithError(c, response.StatusBadRequest, response.ErrInvalidRequest) + response.RespondWithSecureError(c, response.StatusBadRequest, response.ErrInvalidRequest) return } - // TODO: Implement device authorization logic - deviceResponse := &response.DeviceAuthorizationResponse{ - DeviceCode: "generated-device-code", - UserCode: "USER-CODE", - VerificationURI: "https://example.com/device", - ExpiresIn: 900, // 15 minutes - Interval: 5, // 5 seconds + scope := c.PostForm("scope") + oauthService := openapi.OAuth + + res, err := oauthService.DeviceAuthorization(c, clientID, scope) + if err != nil { + if oauthErr, ok := err.(*response.ErrorResponse); ok { + response.RespondWithSecureError(c, response.StatusBadRequest, oauthErr) + } else { + response.RespondWithSecureError(c, response.StatusBadRequest, response.ErrInvalidRequest) + } + return } - response.RespondWithSuccess(c, response.StatusOK, deviceResponse) + response.RespondWithSecureSuccess(c, response.StatusOK, res) +} + +// oauthDeviceAuthorize allows an authenticated user to authorize a pending device code. +func (openapi *OpenAPI) oauthDeviceAuthorize(c *gin.Context) { + authHeader := c.GetHeader("Authorization") + if authHeader == "" || !strings.HasPrefix(authHeader, "Bearer ") { + response.RespondWithSecureError(c, response.StatusUnauthorized, &response.ErrorResponse{ + Code: types.ErrorInvalidGrant, + ErrorDescription: "Bearer token required", + }) + return + } + + tokenStr := strings.TrimPrefix(authHeader, "Bearer ") + oauthService := openapi.OAuth + introspection, err := oauthService.Introspect(c, tokenStr) + if err != nil || introspection == nil || !introspection.Active { + response.RespondWithSecureError(c, response.StatusUnauthorized, &response.ErrorResponse{ + Code: types.ErrorInvalidGrant, + ErrorDescription: "Invalid or expired token", + }) + return + } + + subject := introspection.Subject + if subject == "" { + response.RespondWithSecureError(c, response.StatusUnauthorized, &response.ErrorResponse{ + Code: types.ErrorInvalidGrant, + ErrorDescription: "Token has no subject", + }) + return + } + + userCode := c.PostForm("user_code") + if userCode == "" { + userCode = c.Query("user_code") + } + if userCode == "" { + response.RespondWithSecureError(c, response.StatusBadRequest, response.ErrInvalidRequest) + return + } + + svc, ok := oauthService.(*oauth.Service) + if !ok { + response.RespondWithSecureError(c, response.StatusInternalServerError, &response.ErrorResponse{ + Code: types.ErrorServerError, + ErrorDescription: "OAuth service unavailable", + }) + return + } + + if err := svc.AuthorizeDevice(c, userCode, subject); err != nil { + if oauthErr, ok := err.(*response.ErrorResponse); ok { + response.RespondWithSecureError(c, response.StatusBadRequest, oauthErr) + } else { + response.RespondWithSecureError(c, response.StatusBadRequest, response.ErrInvalidGrant) + } + return + } + + response.RespondWithSecureSuccess(c, response.StatusOK, map[string]string{"status": "authorized"}) } // oauthPushedAuthorizationRequest handles PAR - RFC 9126 diff --git a/openapi/oauth/core.go b/openapi/oauth/core.go index 045ad4d2..3cbb8411 100644 --- a/openapi/oauth/core.go +++ b/openapi/oauth/core.go @@ -131,6 +131,8 @@ func (s *Service) Token(ctx context.Context, grantType string, code string, clie return s.handleClientCredentialsGrant(ctx, client) case types.GrantTypeRefreshToken: return s.handleRefreshTokenGrant(ctx, client, code) // code is refresh token in this case + case types.GrantTypeDeviceCode: + return s.handleDeviceCodeGrant(ctx, client, code) // code is device_code in this case default: return nil, &types.ErrorResponse{ Code: types.ErrorUnsupportedGrantType, @@ -615,3 +617,88 @@ func (s *Service) validatePKCE(ctx context.Context, client *types.ClientInfo, co return nil } + +// handleDeviceCodeGrant handles the device_code grant type (RFC 8628 Section 3.4). +func (s *Service) handleDeviceCodeGrant(ctx context.Context, client *types.ClientInfo, deviceCode string) (*types.Token, error) { + if !s.config.Features.DeviceFlowEnabled { + return nil, &types.ErrorResponse{ + Code: types.ErrorUnsupportedGrantType, + ErrorDescription: "Device flow is not enabled", + } + } + + codeData, err := s.getDeviceCodeData(deviceCode) + if err != nil { + return nil, err + } + + storedClientID, _ := codeData["client_id"].(string) + if storedClientID != client.ClientID { + return nil, &types.ErrorResponse{ + Code: types.ErrorInvalidGrant, + ErrorDescription: "Device code was issued to a different client", + } + } + + expiresAt, _ := codeData["expires_at"].(int64) + if expiresAt == 0 { + if f, ok := codeData["expires_at"].(float64); ok { + expiresAt = int64(f) + } + } + if expiresAt > 0 && time.Now().Unix() > expiresAt { + s.consumeDeviceCode(deviceCode) + return nil, &types.ErrorResponse{ + Code: types.ErrorExpiredToken, + ErrorDescription: "Device code has expired", + } + } + + status, _ := codeData["status"].(string) + switch status { + case "pending": + return nil, &types.ErrorResponse{ + Code: types.ErrorAuthorizationPending, + ErrorDescription: "The authorization request is still pending", + } + + case "authorized": + scope, _ := codeData["scope"].(string) + subject, _ := codeData["subject"].(string) + s.consumeDeviceCode(deviceCode) + + expiresIn := int(s.config.Token.AccessTokenLifetime.Seconds()) + accessToken, err := s.generateAccessTokenWithScope(client.ClientID, scope, subject, expiresIn, nil) + if err != nil { + return nil, &types.ErrorResponse{ + Code: types.ErrorServerError, + ErrorDescription: "Failed to generate access token", + } + } + + token := &types.Token{ + AccessToken: accessToken, + TokenType: "Bearer", + ExpiresIn: expiresIn, + } + + if types.Contains(client.GrantTypes, types.GrantTypeRefreshToken) { + refreshToken, err := s.generateRefreshToken(client.ClientID, scope, subject, 0, nil) + if err != nil { + return nil, &types.ErrorResponse{ + Code: types.ErrorServerError, + ErrorDescription: "Failed to generate refresh token", + } + } + token.RefreshToken = refreshToken + } + + return token, nil + + default: + return nil, &types.ErrorResponse{ + Code: types.ErrorInvalidGrant, + ErrorDescription: "Invalid device code status", + } + } +} diff --git a/openapi/oauth/device.go b/openapi/oauth/device.go index 124aa09f..ed567513 100644 --- a/openapi/oauth/device.go +++ b/openapi/oauth/device.go @@ -2,13 +2,117 @@ package oauth import ( "context" + "fmt" + "strings" + gonanoid "github.com/matoous/go-nanoid/v2" "github.com/yaoapp/yao/openapi/oauth/types" ) -// DeviceAuthorization initiates the device authorization flow -// This is used for devices with limited input capabilities +const userCodeAlphabet = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789" + +// DeviceAuthorization initiates the device authorization flow (RFC 8628). func (s *Service) DeviceAuthorization(ctx context.Context, clientID string, scope string) (*types.DeviceAuthorizationResponse, error) { - // TODO: Implement device authorization flow - return nil, nil + if !s.config.Features.DeviceFlowEnabled { + return nil, &types.ErrorResponse{ + Code: types.ErrorUnsupportedGrantType, + ErrorDescription: "Device flow is not enabled", + } + } + + client, err := s.clientProvider.GetClientByID(ctx, clientID) + if err != nil || client == nil { + return nil, &types.ErrorResponse{ + Code: types.ErrorInvalidClient, + ErrorDescription: "Invalid client", + } + } + + if !clientSupportsGrantType(client, types.GrantTypeDeviceCode) { + return nil, &types.ErrorResponse{ + Code: types.ErrorUnauthorizedClient, + ErrorDescription: "Client does not support device code grant", + } + } + + deviceCode, err := s.generateToken("dc", clientID) + if err != nil { + return nil, &types.ErrorResponse{ + Code: types.ErrorServerError, + ErrorDescription: "Failed to generate device code", + } + } + + userCode, err := s.generateUserCode() + if err != nil { + return nil, &types.ErrorResponse{ + Code: types.ErrorServerError, + ErrorDescription: "Failed to generate user code", + } + } + + if err := s.storeDeviceCode(deviceCode, userCode, clientID, scope); err != nil { + return nil, &types.ErrorResponse{ + Code: types.ErrorServerError, + ErrorDescription: "Failed to store device code", + } + } + + verificationURI := fmt.Sprintf("%s/auth/device", s.config.IssuerURL) + verificationURIComplete := fmt.Sprintf("%s?user_code=%s", verificationURI, userCode) + + return &types.DeviceAuthorizationResponse{ + DeviceCode: deviceCode, + UserCode: userCode, + VerificationURI: verificationURI, + VerificationURIComplete: verificationURIComplete, + ExpiresIn: int(s.config.Token.DeviceCodeLifetime.Seconds()), + Interval: int(s.config.Token.DeviceCodeInterval.Seconds()), + }, nil +} + +// AuthorizeDevice allows an authenticated user to authorize a device code via user_code. +func (s *Service) AuthorizeDevice(ctx context.Context, userCode string, subject string) error { + if !s.config.Features.DeviceFlowEnabled { + return &types.ErrorResponse{ + Code: types.ErrorUnsupportedGrantType, + ErrorDescription: "Device flow is not enabled", + } + } + + normalized := strings.ToUpper(strings.ReplaceAll(userCode, "-", "")) + formatted := normalized + if len(normalized) == 8 { + formatted = normalized[:4] + "-" + normalized[4:] + } + + return s.authorizeDeviceCode(formatted, subject) +} + +// generateUserCode generates a user-friendly code formatted as XXXX-XXXX. +func (s *Service) generateUserCode() (string, error) { + length := s.config.Token.UserCodeLength + if length <= 0 { + length = 8 + } + raw, err := gonanoid.Generate(userCodeAlphabet, length) + if err != nil { + return "", err + } + if len(raw) == 8 { + return raw[:4] + "-" + raw[4:], nil + } + return raw, nil +} + +func clientSupportsGrantType(client *types.ClientInfo, grantType string) bool { + if client == nil || len(client.GrantTypes) == 0 { + return false + } + for _, gt := range client.GrantTypes { + if gt == grantType { + return true + } + } + return false } diff --git a/openapi/oauth/discovery.go b/openapi/oauth/discovery.go index ad0cb2bc..b7f44dd7 100644 --- a/openapi/oauth/discovery.go +++ b/openapi/oauth/discovery.go @@ -63,7 +63,7 @@ func (s *Service) Endpoints(ctx context.Context) (map[string]string, error) { "registration_endpoint": fmt.Sprintf("%s/oauth/register", baseURL), "introspection_endpoint": fmt.Sprintf("%s/oauth/introspect", baseURL), "revocation_endpoint": fmt.Sprintf("%s/oauth/revoke", baseURL), - "device_authorization_endpoint": fmt.Sprintf("%s/oauth/device", baseURL), + "device_authorization_endpoint": fmt.Sprintf("%s/oauth/device_authorization", baseURL), "pushed_authorization_request_endpoint": fmt.Sprintf("%s/oauth/par", baseURL), } diff --git a/openapi/oauth/oauth.go b/openapi/oauth/oauth.go index c4c85085..b7012502 100644 --- a/openapi/oauth/oauth.go +++ b/openapi/oauth/oauth.go @@ -212,6 +212,15 @@ func setConfigDefaults(config *Config) error { if config.Token.DeviceCodeLifetime == 0 { config.Token.DeviceCodeLifetime = 15 * time.Minute } + if config.Token.DeviceCodeLength == 0 { + config.Token.DeviceCodeLength = 8 + } + if config.Token.UserCodeLength == 0 { + config.Token.UserCodeLength = 8 + } + if config.Token.DeviceCodeInterval == 0 { + config.Token.DeviceCodeInterval = 5 * time.Second + } if config.Token.AccessTokenFormat == "" { config.Token.AccessTokenFormat = "jwt" } @@ -257,6 +266,8 @@ func setConfigDefaults(config *Config) error { config.Features.OAuth21Enabled = true config.Features.PKCEEnforced = true config.Features.RefreshTokenRotationEnabled = true + config.Features.DeviceFlowEnabled = true + config.Features.DynamicClientRegistrationEnabled = true return nil } diff --git a/openapi/oauth/token.go b/openapi/oauth/token.go index 04a99c7c..586764bb 100644 --- a/openapi/oauth/token.go +++ b/openapi/oauth/token.go @@ -539,6 +539,124 @@ func (s *Service) consumeAuthorizationCode(code string) error { return nil } +// deviceCodeKey generates a key for device code storage +func (s *Service) deviceCodeKey(code string) string { + return fmt.Sprintf("%soauth:device_code:%s", s.prefix, code) +} + +// userCodeKey generates a key for user code storage (reverse mapping) +func (s *Service) userCodeKey(code string) string { + return fmt.Sprintf("%soauth:user_code:%s", s.prefix, code) +} + +// storeDeviceCode stores device code data and user_code -> device_code reverse mapping +func (s *Service) storeDeviceCode(deviceCode, userCode, clientID, scope string) error { + ttl := s.config.Token.DeviceCodeLifetime + + codeData := map[string]interface{}{ + "client_id": clientID, + "user_code": userCode, + "scope": scope, + "status": "pending", + "issued_at": time.Now().Unix(), + "expires_at": time.Now().Add(ttl).Unix(), + } + if err := s.store.Set(s.deviceCodeKey(deviceCode), codeData, ttl); err != nil { + return err + } + + reverseData := map[string]interface{}{ + "device_code": deviceCode, + } + return s.store.Set(s.userCodeKey(userCode), reverseData, ttl) +} + +// getDeviceCodeData retrieves device code data from store +func (s *Service) getDeviceCodeData(deviceCode string) (map[string]interface{}, error) { + data, exists := s.store.Get(s.deviceCodeKey(deviceCode)) + if !exists { + return nil, &types.ErrorResponse{ + Code: types.ErrorExpiredToken, + ErrorDescription: "Device code not found or expired", + } + } + + codeInfo, ok := data.(map[string]interface{}) + if !ok { + if m, ok := data.(primitive.M); ok { + codeInfo = map[string]interface{}(m) + } else { + return nil, &types.ErrorResponse{ + Code: types.ErrorServerError, + ErrorDescription: "Invalid device code data format", + } + } + } + return codeInfo, nil +} + +// authorizeDeviceCode marks a device code as authorized via user_code lookup +func (s *Service) authorizeDeviceCode(userCode, subject string) error { + reverseData, exists := s.store.Get(s.userCodeKey(userCode)) + if !exists { + return &types.ErrorResponse{ + Code: types.ErrorInvalidGrant, + ErrorDescription: "Invalid or expired user code", + } + } + + var deviceCode string + switch v := reverseData.(type) { + case map[string]interface{}: + deviceCode, _ = v["device_code"].(string) + case primitive.M: + deviceCode, _ = v["device_code"].(string) + } + if deviceCode == "" { + return &types.ErrorResponse{ + Code: types.ErrorServerError, + ErrorDescription: "Invalid user code mapping", + } + } + + codeData, err := s.getDeviceCodeData(deviceCode) + if err != nil { + return err + } + + codeData["status"] = "authorized" + codeData["subject"] = subject + + // Re-store with remaining TTL + expiresAt, _ := codeData["expires_at"].(int64) + if expiresAt == 0 { + if f, ok := codeData["expires_at"].(float64); ok { + expiresAt = int64(f) + } + } + remaining := time.Until(time.Unix(expiresAt, 0)) + if remaining <= 0 { + return &types.ErrorResponse{ + Code: types.ErrorExpiredToken, + ErrorDescription: "Device code has expired", + } + } + + return s.store.Set(s.deviceCodeKey(deviceCode), codeData, remaining) +} + +// consumeDeviceCode deletes both device_code and user_code entries +func (s *Service) consumeDeviceCode(deviceCode string) error { + codeData, _ := s.getDeviceCodeData(deviceCode) + if codeData != nil { + if uc, ok := codeData["user_code"].(string); ok && uc != "" { + s.store.Del(s.userCodeKey(uc)) + } + } + s.store.Del(s.deviceCodeKey(deviceCode)) + return nil +} + // storeRefreshToken stores refresh token with metadata func (s *Service) storeRefreshToken(refreshToken, clientID string) error { tokenData := map[string]interface{}{ diff --git a/openapi/tests/oauth/device_test.go b/openapi/tests/oauth/device_test.go new file mode 100644 index 00000000..54cc7185 --- /dev/null +++ b/openapi/tests/oauth/device_test.go @@ -0,0 +1,292 @@ +package openapi_test + +import ( + "bytes" + "encoding/json" + "io" + "net/http" + "net/url" + "regexp" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/openapi" + "github.com/yaoapp/yao/openapi/oauth/types" + "github.com/yaoapp/yao/openapi/tests/testutils" +) + +// registerDeviceClient registers a device-flow-capable public client via HTTP POST to /oauth/register. +// Returns the client ID. +func registerDeviceClient(t *testing.T, serverURL, baseURL string) string { + t.Helper() + + endpoint := serverURL + baseURL + "/oauth/register" + req := types.DynamicClientRegistrationRequest{ + ClientName: "device-test-client", + RedirectURIs: []string{"http://localhost/device-callback"}, + GrantTypes: []string{types.GrantTypeDeviceCode, types.GrantTypeRefreshToken}, + TokenEndpointAuthMethod: types.TokenEndpointAuthNone, + } + + jsonData, err := json.Marshal(req) + assert.NoError(t, err) + + resp, err := http.Post(endpoint, "application/json", bytes.NewBuffer(jsonData)) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusCreated, resp.StatusCode, "device client registration should succeed") + + var regResp types.DynamicClientRegistrationResponse + err = json.NewDecoder(resp.Body).Decode(®Resp) + assert.NoError(t, err) + assert.NotEmpty(t, regResp.ClientID) + + return regResp.ClientID +} + +// registerConfidentialClient registers a confidential client with client_credentials grant. +// Returns clientID and clientSecret. +func registerConfidentialClient(t *testing.T, serverURL, baseURL string) (string, string) { + t.Helper() + + endpoint := serverURL + baseURL + "/oauth/register" + req := types.DynamicClientRegistrationRequest{ + ClientName: "confidential-token-client", + RedirectURIs: []string{"http://localhost/callback"}, + GrantTypes: []string{types.GrantTypeClientCredentials}, + TokenEndpointAuthMethod: types.TokenEndpointAuthBasic, + } + + jsonData, err := json.Marshal(req) + assert.NoError(t, err) + + resp, err := http.Post(endpoint, "application/json", bytes.NewBuffer(jsonData)) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusCreated, resp.StatusCode, "confidential client registration should succeed") + + var regResp types.DynamicClientRegistrationResponse + err = json.NewDecoder(resp.Body).Decode(®Resp) + assert.NoError(t, err) + assert.NotEmpty(t, regResp.ClientID) + assert.NotEmpty(t, regResp.ClientSecret) + + return regResp.ClientID, regResp.ClientSecret +} + +func TestDeviceAuthorization_Success(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := openapi.Server.Config.BaseURL + clientID := registerDeviceClient(t, serverURL, baseURL) + + endpoint := serverURL + baseURL + "/oauth/device_authorization" + form := url.Values{} + form.Set("client_id", clientID) + + resp, err := http.PostForm(endpoint, form) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + + bodyBytes, err := io.ReadAll(resp.Body) + assert.NoError(t, err) + + var devResp types.DeviceAuthorizationResponse + err = json.Unmarshal(bodyBytes, &devResp) + assert.NoError(t, err) + + assert.NotEmpty(t, devResp.DeviceCode) + assert.NotEmpty(t, devResp.UserCode) + // user_code format XXXX-XXXX (9 chars including hyphen) + assert.Len(t, devResp.UserCode, 9) + assert.Regexp(t, regexp.MustCompile(`^[A-Z0-9]{4}-[A-Z0-9]{4}$`), devResp.UserCode) + assert.NotEmpty(t, devResp.VerificationURI) + assert.Greater(t, devResp.ExpiresIn, 0) + assert.Greater(t, devResp.Interval, 0) +} + +func TestDeviceAuthorization_MissingClientID(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := openapi.Server.Config.BaseURL + endpoint := serverURL + baseURL + "/oauth/device_authorization" + + form := url.Values{} + // no client_id + + resp, err := http.PostForm(endpoint, form) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) +} + +func TestDeviceAuthorization_InvalidClient(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := openapi.Server.Config.BaseURL + endpoint := serverURL + baseURL + "/oauth/device_authorization" + + form := url.Values{} + form.Set("client_id", "nonexistent") + + resp, err := http.PostForm(endpoint, form) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) +} + +func TestDeviceToken_AuthorizationPending(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := openapi.Server.Config.BaseURL + clientID := registerDeviceClient(t, serverURL, baseURL) + + // Get device code + devAuthEndpoint := serverURL + baseURL + "/oauth/device_authorization" + form := url.Values{} + form.Set("client_id", clientID) + + resp, err := http.PostForm(devAuthEndpoint, form) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var devResp types.DeviceAuthorizationResponse + err = json.NewDecoder(resp.Body).Decode(&devResp) + assert.NoError(t, err) + assert.NotEmpty(t, devResp.DeviceCode) + + // Poll token endpoint before user authorizes - should get authorization_pending + tokenEndpoint := serverURL + baseURL + "/oauth/token" + tokenForm := url.Values{} + tokenForm.Set("grant_type", types.GrantTypeDeviceCode) + tokenForm.Set("device_code", devResp.DeviceCode) + tokenForm.Set("client_id", clientID) + + tokenResp, err := http.PostForm(tokenEndpoint, tokenForm) + assert.NoError(t, err) + defer tokenResp.Body.Close() + + // RFC 8628: authorization_pending returns 400 with error + assert.Equal(t, http.StatusBadRequest, tokenResp.StatusCode) + + bodyBytes, err := io.ReadAll(tokenResp.Body) + assert.NoError(t, err) + + var errResp types.ErrorResponse + err = json.Unmarshal(bodyBytes, &errResp) + assert.NoError(t, err) + assert.Equal(t, types.ErrorAuthorizationPending, errResp.Code) +} + +func TestDeviceToken_InvalidDeviceCode(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := openapi.Server.Config.BaseURL + clientID := registerDeviceClient(t, serverURL, baseURL) + + tokenEndpoint := serverURL + baseURL + "/oauth/token" + tokenForm := url.Values{} + tokenForm.Set("grant_type", types.GrantTypeDeviceCode) + tokenForm.Set("device_code", "bogus-invalid-device-code") + tokenForm.Set("client_id", clientID) + + resp, err := http.PostForm(tokenEndpoint, tokenForm) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + bodyBytes, err := io.ReadAll(resp.Body) + assert.NoError(t, err) + + var errResp types.ErrorResponse + err = json.Unmarshal(bodyBytes, &errResp) + assert.NoError(t, err) + assert.NotEmpty(t, errResp.Code) +} + +func TestDeviceFlow_EndToEnd(t *testing.T) { + serverURL := testutils.Prepare(t) + defer testutils.Clean() + + baseURL := openapi.Server.Config.BaseURL + + // a. Register device client + deviceClientID := registerDeviceClient(t, serverURL, baseURL) + + // b. POST /oauth/device_authorization -> get device_code + user_code + devAuthEndpoint := serverURL + baseURL + "/oauth/device_authorization" + form := url.Values{} + form.Set("client_id", deviceClientID) + + resp, err := http.PostForm(devAuthEndpoint, form) + assert.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusOK, resp.StatusCode) + + var devResp types.DeviceAuthorizationResponse + err = json.NewDecoder(resp.Body).Decode(&devResp) + assert.NoError(t, err) + assert.NotEmpty(t, devResp.DeviceCode) + assert.NotEmpty(t, devResp.UserCode) + + // c. Get bearer token: register confidential client, get token via client_credentials. + // Device authorize requires a token with subject; client_credentials tokens have no subject. + // Use ObtainAccessTokenWithRootPermission to get a token with subject for device authorize. + confClientID, confClientSecret := registerConfidentialClient(t, serverURL, baseURL) + tokenInfo := testutils.ObtainAccessTokenWithRootPermission(t, serverURL, confClientID, confClientSecret, "http://localhost/callback", "openid profile") + bearerToken := tokenInfo.AccessToken + + tokenEndpoint := serverURL + baseURL + "/oauth/token" + + // d. POST /oauth/device/authorize with bearer + user_code -> assert 200 + deviceAuthorizeEndpoint := serverURL + baseURL + "/oauth/device/authorize" + authForm := url.Values{} + authForm.Set("user_code", devResp.UserCode) + + authReq, err := http.NewRequest("POST", deviceAuthorizeEndpoint, bytes.NewBufferString(authForm.Encode())) + assert.NoError(t, err) + authReq.Header.Set("Content-Type", "application/x-www-form-urlencoded") + authReq.Header.Set("Authorization", "Bearer "+bearerToken) + + authResp, err := http.DefaultClient.Do(authReq) + assert.NoError(t, err) + defer authResp.Body.Close() + + assert.Equal(t, http.StatusOK, authResp.StatusCode, "device authorize should succeed") + + // e. POST /oauth/token with device_code -> assert access_token returned + dcForm := url.Values{} + dcForm.Set("grant_type", types.GrantTypeDeviceCode) + dcForm.Set("device_code", devResp.DeviceCode) + dcForm.Set("client_id", deviceClientID) + + dcResp, err := http.PostForm(tokenEndpoint, dcForm) + assert.NoError(t, err) + defer dcResp.Body.Close() + + assert.Equal(t, http.StatusOK, dcResp.StatusCode, "device token exchange should succeed") + + var finalToken struct { + AccessToken string `json:"access_token"` + TokenType string `json:"token_type"` + } + err = json.NewDecoder(dcResp.Body).Decode(&finalToken) + assert.NoError(t, err) + assert.NotEmpty(t, finalToken.AccessToken) + assert.Equal(t, "Bearer", finalToken.TokenType) +} From 9dc99d8b8edbc9e4e1ae059fddc20331824c1133 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 4 Mar 2026 15:32:10 +0800 Subject: [PATCH 5/7] Enhance health check and integration tests for Tai - Update health check messages in CI workflows to specify HTTP and gRPC readiness for the Tai service. - Refactor integration tests to use `require` assertions instead of `assert`, improving error handling and test reliability. - Ensure that responses from gRPC calls are not nil, enhancing test robustness. These changes improve the clarity of service readiness checks and strengthen the integration test suite. --- .github/workflows/pr-test.yml | 12 ++++++++++-- .github/workflows/unit-test.yml | 12 ++++++++++-- tai/grpc/integration_test.go | 18 +++++++++--------- 3 files changed, 29 insertions(+), 13 deletions(-) diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index c8203e46..db5a324d 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -1706,10 +1706,18 @@ jobs: yaoapp/tai:latest for i in $(seq 1 30); do if curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1; then - echo "Tai is ready" + echo "Tai HTTP is ready" break fi - echo "Waiting for Tai... ($i)" + echo "Waiting for Tai HTTP... ($i)" + sleep 1 + done + for i in $(seq 1 15); do + if nc -z 127.0.0.1 9100 2>/dev/null; then + echo "Tai gRPC is ready" + break + fi + echo "Waiting for Tai gRPC... ($i)" sleep 1 done diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index e7432c52..9bdc380d 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -1260,10 +1260,18 @@ jobs: yaoapp/tai:latest for i in $(seq 1 30); do if curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1; then - echo "Tai is ready" + echo "Tai HTTP is ready" break fi - echo "Waiting for Tai... ($i)" + echo "Waiting for Tai HTTP... ($i)" + sleep 1 + done + for i in $(seq 1 15); do + if nc -z 127.0.0.1 9100 2>/dev/null; then + echo "Tai gRPC is ready" + break + fi + echo "Waiting for Tai gRPC... ($i)" sleep 1 done diff --git a/tai/grpc/integration_test.go b/tai/grpc/integration_test.go index 02a36a51..5b534409 100644 --- a/tai/grpc/integration_test.go +++ b/tai/grpc/integration_test.go @@ -90,7 +90,8 @@ 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) + require.NoError(t, err) + require.NotNil(t, resp) assert.Equal(t, int32(0), resp.ExitCode) assert.Contains(t, string(resp.Stdout), "hello") } @@ -147,10 +148,8 @@ 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. + require.NoError(t, err) + require.NotNil(t, resp) t.Logf("API proxy status: %d", resp.Status) } @@ -313,7 +312,7 @@ func TestRelay_Healthz(t *testing.T) { defer client.Close() status, err := client.Healthz(context.Background()) - assert.NoError(t, err) + require.NoError(t, err) assert.Equal(t, "ok", status) } @@ -321,8 +320,8 @@ 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) + require.NoError(t, err) + require.NotNil(t, data) t.Logf("relay Run result: %s", string(data)) } @@ -337,7 +336,8 @@ 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) + require.NoError(t, err) + require.NotNil(t, resp) assert.Equal(t, int32(0), resp.ExitCode) assert.Contains(t, string(resp.Stdout), "relay-test") } From c6e1c449e1c6f24e946e4957eb66a24b6147d0e6 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 4 Mar 2026 16:55:29 +0800 Subject: [PATCH 6/7] Enhance Tai service readiness checks and OAuth device flow - Improve health check logic in CI workflows for both HTTP and gRPC readiness of the Tai service, ensuring clearer error reporting if the service fails to start. - Update the OAuth Device Flow implementation to support additional claims during device authorization, enhancing the flexibility of the authorization process. - Refactor the `AuthorizeDevice` method to accept extra claims, allowing for more detailed user context during authorization. - Introduce a new utility function to extract bearer tokens from requests, streamlining token handling across the OpenAPI service. These changes enhance the robustness of service readiness checks and improve the OAuth device authorization flow, contributing to a more reliable and flexible authentication mechanism. --- .github/workflows/pr-test.yml | 20 ++ .github/workflows/unit-test.yml | 20 ++ cmd/credential.go | 118 ++++++++++ cmd/login.go | 342 ++++++++++++++++++++++++++++ cmd/logout.go | 80 +++++++ cmd/root.go | 2 + cmd/run.go | 392 ++++++++++++++++++++------------ grpc/IMPL.md | 33 +-- openapi/config.go | 1 + openapi/oauth.go | 96 +++++--- openapi/oauth/core.go | 15 +- openapi/oauth/device.go | 8 +- openapi/oauth/discovery.go | 3 +- openapi/oauth/oauth.go | 1 + openapi/oauth/token.go | 5 +- openapi/well-known.go | 46 +++- 16 files changed, 980 insertions(+), 202 deletions(-) create mode 100644 cmd/credential.go create mode 100644 cmd/login.go create mode 100644 cmd/logout.go diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index db5a324d..3f1bcfa4 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -1704,22 +1704,42 @@ jobs: -p 8080:8080 -p 9100:9100 -p 2375:2375 -p 6080:6080 -p 6443:6443 \ -e TAI_K8S_UPSTREAM="tcp://${K3D_IP}:6443" \ yaoapp/tai:latest + + TAI_HTTP_READY=false for i in $(seq 1 30); do if curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1; then echo "Tai HTTP is ready" + TAI_HTTP_READY=true break fi echo "Waiting for Tai HTTP... ($i)" sleep 1 done + if [ "$TAI_HTTP_READY" != "true" ]; then + echo "::error::Tai HTTP failed to become ready within 30s" + echo "--- Tai container logs ---" + docker logs tai 2>&1 || true + echo "--- Tai container status ---" + docker inspect tai --format='{{.State.Status}} exit={{.State.ExitCode}}' || true + exit 1 + fi + + TAI_GRPC_READY=false for i in $(seq 1 15); do if nc -z 127.0.0.1 9100 2>/dev/null; then echo "Tai gRPC is ready" + TAI_GRPC_READY=true break fi echo "Waiting for Tai gRPC... ($i)" sleep 1 done + if [ "$TAI_GRPC_READY" != "true" ]; then + echo "::error::Tai gRPC failed to become ready within 15s" + echo "--- Tai container logs ---" + docker logs tai 2>&1 || true + exit 1 + fi - name: Generate kubeconfig for Tai K8s proxy run: | diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index 9bdc380d..3e04605e 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -1258,22 +1258,42 @@ jobs: -p 8080:8080 -p 9100:9100 -p 2375:2375 -p 6080:6080 -p 6443:6443 \ -e TAI_K8S_UPSTREAM="tcp://${K3D_IP}:6443" \ yaoapp/tai:latest + + TAI_HTTP_READY=false for i in $(seq 1 30); do if curl -sf http://127.0.0.1:8080/healthz > /dev/null 2>&1; then echo "Tai HTTP is ready" + TAI_HTTP_READY=true break fi echo "Waiting for Tai HTTP... ($i)" sleep 1 done + if [ "$TAI_HTTP_READY" != "true" ]; then + echo "::error::Tai HTTP failed to become ready within 30s" + echo "--- Tai container logs ---" + docker logs tai 2>&1 || true + echo "--- Tai container status ---" + docker inspect tai --format='{{.State.Status}} exit={{.State.ExitCode}}' || true + exit 1 + fi + + TAI_GRPC_READY=false for i in $(seq 1 15); do if nc -z 127.0.0.1 9100 2>/dev/null; then echo "Tai gRPC is ready" + TAI_GRPC_READY=true break fi echo "Waiting for Tai gRPC... ($i)" sleep 1 done + if [ "$TAI_GRPC_READY" != "true" ]; then + echo "::error::Tai gRPC failed to become ready within 15s" + echo "--- Tai container logs ---" + docker logs tai 2>&1 || true + exit 1 + fi - name: Generate kubeconfig for Tai K8s proxy run: | diff --git a/cmd/credential.go b/cmd/credential.go new file mode 100644 index 00000000..9f8f3dc9 --- /dev/null +++ b/cmd/credential.go @@ -0,0 +1,118 @@ +package cmd + +import ( + "encoding/base64" + "encoding/json" + "fmt" + "os" + "path/filepath" + "time" +) + +// Credential represents the stored OAuth credential for gRPC mode. +type Credential struct { + Server string `json:"server"` + GRPCAddr string `json:"grpc_addr,omitempty"` + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token,omitempty"` + Scope string `json:"scope,omitempty"` + User string `json:"user,omitempty"` + ExpiresAt string `json:"expires_at,omitempty"` +} + +// Expired returns true if the credential has an expires_at in the past. +func (c *Credential) Expired() bool { + if c.ExpiresAt == "" { + return false + } + t, err := time.Parse(time.RFC3339, c.ExpiresAt) + if err != nil { + return false + } + return time.Now().After(t) +} + +func credentialPath() (string, error) { + home, err := os.UserHomeDir() + if err != nil { + return "", fmt.Errorf("cannot determine home directory: %w", err) + } + return filepath.Join(home, ".yao", "credentials"), nil +} + +// LoadCredential reads and decodes ~/.yao/credentials. Returns nil if the file +// does not exist. +func LoadCredential() (*Credential, error) { + path, err := credentialPath() + if err != nil { + return nil, err + } + raw, err := os.ReadFile(path) + if os.IsNotExist(err) { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("read credentials: %w", err) + } + + decoded, err := base64.StdEncoding.DecodeString(string(raw)) + if err != nil { + return nil, fmt.Errorf("decode credentials: %w", err) + } + + var cred Credential + if err := json.Unmarshal(decoded, &cred); err != nil { + return nil, fmt.Errorf("unmarshal credentials: %w", err) + } + return &cred, nil +} + +// LoadCredentialFrom reads and decodes a credential file from a custom path. +func LoadCredentialFrom(path string) (*Credential, error) { + raw, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("read credentials from %s: %w", path, err) + } + decoded, err := base64.StdEncoding.DecodeString(string(raw)) + if err != nil { + return nil, fmt.Errorf("decode credentials: %w", err) + } + var cred Credential + if err := json.Unmarshal(decoded, &cred); err != nil { + return nil, fmt.Errorf("unmarshal credentials: %w", err) + } + return &cred, nil +} + +// SaveCredential encodes and writes the credential to ~/.yao/credentials. +func SaveCredential(cred *Credential) error { + path, err := credentialPath() + if err != nil { + return err + } + dir := filepath.Dir(path) + if err := os.MkdirAll(dir, 0700); err != nil { + return fmt.Errorf("create directory %s: %w", dir, err) + } + data, err := json.Marshal(cred) + if err != nil { + return fmt.Errorf("marshal credentials: %w", err) + } + encoded := base64.StdEncoding.EncodeToString(data) + if err := os.WriteFile(path, []byte(encoded), 0600); err != nil { + return fmt.Errorf("write credentials: %w", err) + } + return nil +} + +// RemoveCredential deletes ~/.yao/credentials. +func RemoveCredential() error { + path, err := credentialPath() + if err != nil { + return err + } + if err := os.Remove(path); err != nil && !os.IsNotExist(err) { + return fmt.Errorf("remove credentials: %w", err) + } + return nil +} diff --git a/cmd/login.go b/cmd/login.go new file mode 100644 index 00000000..b2221e6d --- /dev/null +++ b/cmd/login.go @@ -0,0 +1,342 @@ +package cmd + +import ( + "encoding/base64" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "os" + "strings" + "time" + + "github.com/fatih/color" + "github.com/spf13/cobra" + "github.com/yaoapp/yao/engine" +) + +var loginServer string + +var loginCmd = &cobra.Command{ + Use: "login", + Short: L("Login to remote Yao server"), + Long: L("Login to remote Yao server using device authorization flow"), + Run: func(cmd *cobra.Command, args []string) { + if loginServer == "" { + color.Red(L("Missing --server flag\n")) + fmt.Println(" yao login --server https://yaoagents.com") + os.Exit(1) + } + + serverURL := strings.TrimRight(loginServer, "/") + + // 1. Discover OAuth endpoints via well-known metadata + endpoints, err := discoverEndpoints(serverURL) + if err != nil { + color.Red(" %s %s\n", L("Server discovery failed:"), err) + os.Exit(1) + } + + // 2. Compute deterministic client_id from machine fingerprint + machine, err := engine.GetMachineID() + if err != nil { + color.Red("Failed to compute machine ID: %s\n", err) + os.Exit(1) + } + clientID := machine.ID + + // 3. Register the client (idempotent for same client_id) + if endpoints.RegistrationEndpoint != "" { + if err := registerClient(endpoints.RegistrationEndpoint, clientID); err != nil { + color.Red("Client registration failed: %s\n", err) + os.Exit(1) + } + } + + // 4. Start device authorization + deviceResp, err := requestDeviceAuthorization(endpoints.DeviceAuthorizationEndpoint, clientID) + if err != nil { + color.Red("Device authorization failed: %s\n", err) + os.Exit(1) + } + + // 5. Display the code to the user + dashboard := endpoints.Dashboard + if dashboard == "" { + dashboard = "/admin" + } + verifyURI := strings.TrimRight(serverURL, "/") + dashboard + "/auth/device" + verifyURIComplete := verifyURI + "?user_code=" + deviceResp.UserCode + + fmt.Println() + color.White(" %s %s\n", + L("Open:"), + color.CyanString(verifyURIComplete)) + fmt.Println() + color.White(" %s %s\n", + L("Or visit:"), + color.CyanString(verifyURI)) + color.White(" %s %s\n", + L("Enter code:"), + color.YellowString(deviceResp.UserCode)) + fmt.Println() + + // 6. Poll for token + interval := deviceResp.Interval + if interval < 5 { + interval = 5 + } + + color.White(" %s", L("Waiting for authorization...")) + tokenResp, err := pollForToken(endpoints.TokenEndpoint, clientID, deviceResp.DeviceCode, interval, deviceResp.ExpiresIn) + if err != nil { + fmt.Println() + color.Red("\n %s %s\n", L("Login failed:"), err) + os.Exit(1) + } + + // 6. Save credential + expiresAt := "" + if tokenResp.ExpiresIn > 0 { + expiresAt = time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second).UTC().Format(time.RFC3339) + } + + cred := &Credential{ + Server: serverURL, + GRPCAddr: endpoints.GRPCAddr, + AccessToken: tokenResp.AccessToken, + RefreshToken: tokenResp.RefreshToken, + Scope: tokenResp.Scope, + User: parseJWTSubject(tokenResp.AccessToken), + ExpiresAt: expiresAt, + } + + if err := SaveCredential(cred); err != nil { + color.Red("\n Failed to save credentials: %s\n", err) + os.Exit(1) + } + + fmt.Print("\033[2J\033[H") + color.Green(" ✓ %s\n", L("Login successful")) + color.White(" %s %s\n", L("Server:"), serverURL) + if cred.GRPCAddr != "" { + color.White(" %s %s\n", L("gRPC:"), cred.GRPCAddr) + } + if cred.User != "" { + color.White(" %s %s\n", L("User:"), cred.User) + } + if cred.ExpiresAt != "" { + color.White(" %s %s\n", L("Expires:"), cred.ExpiresAt) + } + fmt.Println() + }, +} + +func init() { + loginCmd.PersistentFlags().StringVar(&loginServer, "server", "", L("Remote Yao server URL")) +} + +// --- types --- + +type oauthEndpoints struct { + RegistrationEndpoint string `json:"registration_endpoint"` + DeviceAuthorizationEndpoint string `json:"device_authorization_endpoint"` + TokenEndpoint string `json:"token_endpoint"` + RevocationEndpoint string `json:"revocation_endpoint"` + Dashboard string `json:"-"` + GRPCAddr string `json:"-"` +} + +type deviceAuthResponse struct { + DeviceCode string `json:"device_code"` + UserCode string `json:"user_code"` + VerificationURI string `json:"verification_uri"` + VerificationURIComplete string `json:"verification_uri_complete"` + ExpiresIn int `json:"expires_in"` + Interval int `json:"interval"` +} + +type tokenResponse struct { + AccessToken string `json:"access_token"` + TokenType string `json:"token_type"` + ExpiresIn int `json:"expires_in"` + RefreshToken string `json:"refresh_token"` + Scope string `json:"scope"` +} + +type oauthError struct { + Error string `json:"error"` + ErrorDescription string `json:"error_description"` +} + +// --- HTTP helpers --- + +// discoverEndpoints fetches OAuth endpoint URLs from /.well-known/yao, +// using the openapi base prefix to construct correct API paths. +func discoverEndpoints(serverURL string) (*oauthEndpoints, error) { + return discoverFromYaoMetadata(serverURL) +} + +type yaoMetadataResponse struct { + OpenAPI string `json:"openapi"` + Dashboard string `json:"dashboard"` + GRPC string `json:"grpc"` +} + +func discoverFromYaoMetadata(serverURL string) (*oauthEndpoints, error) { + resp, err := http.Get(serverURL + "/.well-known/yao") + if err != nil { + return nil, fmt.Errorf("network error: %w", err) + } + defer resp.Body.Close() + + body, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("/.well-known/yao returned %d", resp.StatusCode) + } + + var meta yaoMetadataResponse + if err := json.Unmarshal(body, &meta); err != nil { + return nil, fmt.Errorf("invalid /.well-known/yao response: %w", err) + } + + base := strings.TrimRight(serverURL, "/") + meta.OpenAPI + + return &oauthEndpoints{ + RegistrationEndpoint: base + "/oauth/register", + DeviceAuthorizationEndpoint: base + "/oauth/device_authorization", + TokenEndpoint: base + "/oauth/token", + RevocationEndpoint: base + "/oauth/revoke", + Dashboard: meta.Dashboard, + GRPCAddr: meta.GRPC, + }, nil +} + +func registerClient(endpoint, clientID string) error { + body := fmt.Sprintf( + `{"client_id":"%s","client_name":"yao-cli","grant_types":["urn:ietf:params:oauth:grant-type:device_code"],"token_endpoint_auth_method":"none","redirect_uris":["http://localhost"]}`, + clientID, + ) + resp, err := http.Post(endpoint, "application/json", strings.NewReader(body)) + if err != nil { + return fmt.Errorf("network error: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusCreated { + return nil + } + + respBody, _ := io.ReadAll(resp.Body) + var oerr oauthError + if json.Unmarshal(respBody, &oerr) == nil && oerr.Error == "invalid_client_metadata" { + return nil // client already registered, idempotent + } + return fmt.Errorf("registration returned %d: %s", resp.StatusCode, string(respBody)) +} + +func requestDeviceAuthorization(endpoint, clientID string) (*deviceAuthResponse, error) { + data := url.Values{ + "client_id": {clientID}, + "scope": {"grpc:run grpc:stream grpc:shell grpc:mcp grpc:llm grpc:agent"}, + } + resp, err := http.PostForm(endpoint, data) + if err != nil { + return nil, fmt.Errorf("network error: %w", err) + } + defer resp.Body.Close() + + respBody, _ := io.ReadAll(resp.Body) + if resp.StatusCode != http.StatusOK { + var oerr oauthError + json.Unmarshal(respBody, &oerr) + if oerr.ErrorDescription != "" { + return nil, fmt.Errorf("%s", oerr.ErrorDescription) + } + return nil, fmt.Errorf("server returned %d: %s", resp.StatusCode, string(respBody)) + } + + var result deviceAuthResponse + if err := json.Unmarshal(respBody, &result); err != nil { + return nil, fmt.Errorf("invalid response: %w", err) + } + return &result, nil +} + +func pollForToken(endpoint, clientID, deviceCode string, interval, expiresIn int) (*tokenResponse, error) { + deadline := time.Now().Add(time.Duration(expiresIn) * time.Second) + ticker := time.NewTicker(time.Duration(interval) * time.Second) + defer ticker.Stop() + + for range ticker.C { + if time.Now().After(deadline) { + return nil, fmt.Errorf("device code expired") + } + + data := url.Values{ + "grant_type": {"urn:ietf:params:oauth:grant-type:device_code"}, + "client_id": {clientID}, + "device_code": {deviceCode}, + } + + resp, err := http.PostForm(endpoint, data) + if err != nil { + continue + } + + respBody, _ := io.ReadAll(resp.Body) + resp.Body.Close() + + if resp.StatusCode == http.StatusOK { + var tok tokenResponse + if err := json.Unmarshal(respBody, &tok); err != nil { + return nil, fmt.Errorf("invalid token response: %w", err) + } + return &tok, nil + } + + var oerr oauthError + json.Unmarshal(respBody, &oerr) + switch oerr.Error { + case "authorization_pending": + fmt.Print(".") + continue + case "slow_down": + interval += 5 + ticker.Reset(time.Duration(interval) * time.Second) + continue + case "expired_token": + return nil, fmt.Errorf("device code expired") + case "access_denied": + return nil, fmt.Errorf("authorization denied by user") + default: + desc := oerr.ErrorDescription + if desc == "" { + desc = oerr.Error + } + return nil, fmt.Errorf("%s", desc) + } + } + return nil, fmt.Errorf("device code expired") +} + +// parseJWTSubject extracts the "sub" claim from a JWT access token +// without verifying the signature (display-only). +func parseJWTSubject(token string) string { + parts := strings.Split(token, ".") + if len(parts) != 3 { + return "" + } + payload, err := base64.RawURLEncoding.DecodeString(parts[1]) + if err != nil { + return "" + } + var claims struct { + Sub string `json:"sub"` + } + if json.Unmarshal(payload, &claims) != nil { + return "" + } + return claims.Sub +} diff --git a/cmd/logout.go b/cmd/logout.go new file mode 100644 index 00000000..6c56ce2e --- /dev/null +++ b/cmd/logout.go @@ -0,0 +1,80 @@ +package cmd + +import ( + "net/http" + "net/url" + "os" + "strings" + + "github.com/fatih/color" + "github.com/spf13/cobra" +) + +var logoutCmd = &cobra.Command{ + Use: "logout", + Short: L("Logout from remote Yao server"), + Long: L("Revoke token and remove stored credentials"), + Run: func(cmd *cobra.Command, args []string) { + cred, err := LoadCredential() + if err != nil { + color.Red(" %s %s\n", L("Failed to read credentials:"), err) + os.Exit(1) + } + if cred == nil { + color.Yellow(" %s\n", L("Not logged in")) + return + } + + // Best-effort token revocation via discovery + if cred.AccessToken != "" && cred.Server != "" { + if ep, err := discoverEndpoints(cred.Server); err == nil && ep.RevocationEndpoint != "" { + revokeToken(ep.RevocationEndpoint, cred.AccessToken) + } + } + + if err := RemoveCredential(); err != nil { + color.Red(" %s %s\n", L("Failed to remove credentials:"), err) + os.Exit(1) + } + + color.Green(" ✓ %s\n", L("Logged out")) + if cred.Server != "" { + color.White(" %s %s\n", L("Server:"), cred.Server) + } + }, +} + +func revokeToken(endpoint, token string) { + data := url.Values{"token": {token}} + req, err := http.NewRequest("POST", endpoint, strings.NewReader(data.Encode())) + if err != nil { + return + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + http.DefaultClient.Do(req) +} + +func init() { + // Add i18n entries + langs["Login to remote Yao server"] = "登录远程 Yao 服务器" + langs["Login to remote Yao server using device authorization flow"] = "使用设备授权流程登录远程 Yao 服务器" + langs["Remote Yao server URL"] = "远程 Yao 服务器地址" + langs["Logout from remote Yao server"] = "登出远程 Yao 服务器" + langs["Revoke token and remove stored credentials"] = "撤销令牌并移除存储的凭证" + langs["Missing --server flag"] = "缺少 --server 参数" + langs["Open:"] = "打开:" + langs["Or visit:"] = "或访问:" + langs["Enter code:"] = "输入设备码:" + langs["Waiting for authorization..."] = "等待授权..." + langs["Login failed:"] = "登录失败:" + langs["Login successful"] = "登录成功" + langs["Server:"] = "服务器:" + langs["Scope:"] = "授权范围:" + langs["Failed to read credentials:"] = "读取凭证失败:" + langs["Not logged in"] = "未登录" + langs["Failed to remove credentials:"] = "移除凭证失败:" + langs["Logged out"] = "已登出" + langs["Path to credentials file"] = "凭证文件路径" + langs["Failed to load credentials:"] = "加载凭证失败:" + langs["Server discovery failed:"] = "服务发现失败:" +} diff --git a/cmd/root.go b/cmd/root.go index 1a228402..7f32b1f4 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -189,6 +189,8 @@ func init() { inspectCmd, startCmd, runCmd, + loginCmd, + logoutCmd, // getCmd, // dumpCmd, // restoreCmd, diff --git a/cmd/run.go b/cmd/run.go index aa840ae8..f6dcca3b 100644 --- a/cmd/run.go +++ b/cmd/run.go @@ -18,173 +18,271 @@ import ( "github.com/yaoapp/yao/engine" ischedule "github.com/yaoapp/yao/schedule" "github.com/yaoapp/yao/share" + taigrpc "github.com/yaoapp/yao/tai/grpc" itask "github.com/yaoapp/yao/task" ) var runSilent = false +var runAuthPath string var runCmd = &cobra.Command{ Use: "run", Short: L("Execute process"), Long: L("Execute process"), Run: func(cmd *cobra.Command, args []string) { - defer share.SessionStop() - defer plugin.KillAll() - defer func() { - err := exception.Catch(recover()) - if err != nil { - if !runSilent { - color.Red(L("Fatal: %s\n"), err.Error()) - return - } - fmt.Printf("%s\n", err.Error()) - } - }() + // Resolve credential: --auth flag > ~/.yao/credentials > nil (local mode) + cred := resolveCredential() - // Auto-detect app root if not specified - if appPath == "" { - cwd, err := os.Getwd() - if err == nil { - if root, err := findAppRootFromPath(cwd); err == nil { - appPath = root - } - } - } - - Boot() - - // Set Runtime Mode - config.Conf.Runtime.Mode = "standard" - - cfg := config.Conf - cfg.Session.IsCLI = true - if len(args) < 1 { - if !runSilent { - color.Red(L("Not enough arguments\n")) - color.White(share.BUILDNAME + " help\n") - return - } - fmt.Print(L("Not enough arguments\n")) + if cred != nil { + runGRPC(cred, args) return } - loadWarnings, err := engine.Load(cfg, engine.LoadOption{Action: "run"}) - if err != nil { - if !runSilent { - color.Red(L("Engine: %s\n"), err.Error()) - return - } - - fmt.Printf("%s\n", err.Error()) - return - } - - name := args[0] - if !runSilent { - color.Green(L("Run: %s\n"), name) - } - - pargs := []interface{}{} - for i, arg := range args { - if i == 0 { - continue - } - - // Parse the arguments - if strings.HasPrefix(arg, "::") { - arg := strings.TrimPrefix(arg, "::") - var v interface{} - err := jsoniter.Unmarshal([]byte(arg), &v) - if err != nil { - color.Red(L("Arguments: %s\n"), err.Error()) - return - } - pargs = append(pargs, v) - - if !runSilent { - color.White("args[%d]: %s\n", i-1, arg) - } - - } else if strings.HasPrefix(arg, "\\::") { - arg := "::" + strings.TrimPrefix(arg, "\\::") - pargs = append(pargs, arg) - if !runSilent { - color.White("args[%d]: %s\n", i-1, arg) - } - - } else { - pargs = append(pargs, arg) - if !runSilent { - color.White("args[%d]: %s\n", i-1, arg) - } - } - - } - - // Start Tasks - itask.Start() - defer itask.Stop() - - // Start Schedules - ischedule.Start() - defer ischedule.Stop() - - process := process.NewWithContext(context.Background(), name, pargs...) - res, err := process.Exec() - if err != nil { - if !runSilent { - color.Red(L("Process: %s\n"), fmt.Sprintf("%s", strings.TrimPrefix(err.Error(), "Exception|404:"))) - return - } - fmt.Printf("%s\n", err.Error()) - return - } - - if !runSilent { - - if len(loadWarnings) > 0 { - fmt.Println(color.YellowString("---------------------------------")) - fmt.Println(color.YellowString(L("Warnings"))) - fmt.Println(color.YellowString("---------------------------------")) - for _, warning := range loadWarnings { - fmt.Println(color.YellowString("[%s] %s", warning.Widget, warning.Error)) - } - fmt.Printf("\n") - } - - color.White("--------------------------------------\n") - color.White(L("%s Response\n"), name) - color.White("--------------------------------------\n") - helper.Dump(res) - color.White("--------------------------------------\n") - color.Green(L("✨DONE✨\n")) - return - } - - // Silent mode output - switch res.(type) { - - case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64, bool: - fmt.Printf("%v\n", res) - return - - case string, []byte: - fmt.Printf("%s\n", res) - return - - default: - txt, err := jsoniter.Marshal(res) - if err != nil { - fmt.Printf("%s\n", err.Error()) - } - fmt.Printf("%s\n", txt) - } + runLocal(args) }, } func init() { runCmd.PersistentFlags().BoolVarP(&runSilent, "silent", "s", false, L("Silent mode")) + runCmd.PersistentFlags().StringVar(&runAuthPath, "auth", "", L("Path to credentials file")) +} + +// resolveCredential loads credential from --auth flag or default path. +func resolveCredential() *Credential { + if runAuthPath != "" { + cred, err := LoadCredentialFrom(runAuthPath) + if err != nil { + color.Red(" %s %s\n", L("Failed to load credentials:"), err) + os.Exit(1) + } + return cred + } + + cred, _ := LoadCredential() + return cred +} + +// runGRPC executes a process via the remote gRPC server. +func runGRPC(cred *Credential, args []string) { + if len(args) < 1 { + if !runSilent { + color.Red(L("Not enough arguments\n")) + color.White(share.BUILDNAME + " help\n") + } else { + fmt.Print(L("Not enough arguments\n")) + } + os.Exit(1) + } + + if cred.GRPCAddr == "" { + color.Red(" %s\n", L("No gRPC address in credentials. Please re-login.")) + os.Exit(1) + } + + name := args[0] + if !runSilent { + color.Green(L("Run: %s gRPC: %s\n"), name, cred.GRPCAddr) + } + + pargs := parseRunArgs(args[1:]) + + argsJSON, err := jsoniter.Marshal(pargs) + if err != nil { + color.Red(" %s %s\n", L("Arguments:"), err.Error()) + os.Exit(1) + } + + tm := taigrpc.NewTokenManager(cred.AccessToken, cred.RefreshToken, "", "") + client, err := taigrpc.Dial(cred.GRPCAddr, tm) + if err != nil { + color.Red(" %s %s\n", L("gRPC connect failed:"), err.Error()) + os.Exit(1) + } + defer client.Close() + + data, err := client.Run(context.Background(), name, argsJSON, 0) + if err != nil { + if !runSilent { + color.Red(" %s %s\n", L("Process:"), err.Error()) + } else { + fmt.Printf("%s\n", err.Error()) + } + os.Exit(1) + } + + if !runSilent { + color.White("--------------------------------------\n") + color.White(L("%s Response\n"), name) + color.White("--------------------------------------\n") + var res interface{} + if jsoniter.Unmarshal(data, &res) == nil { + helper.Dump(res) + } else { + fmt.Printf("%s\n", data) + } + color.White("--------------------------------------\n") + fmt.Printf("\033[32m✨DONE✨\033[0m \033[90mgRPC: %s\033[0m\n", cred.GRPCAddr) + } else { + fmt.Printf("%s\n", data) + } +} + +// runLocal executes a process locally (existing behavior). +func runLocal(args []string) { + defer share.SessionStop() + defer plugin.KillAll() + + defer func() { + err := exception.Catch(recover()) + if err != nil { + if !runSilent { + color.Red(L("Fatal: %s\n"), err.Error()) + return + } + fmt.Printf("%s\n", err.Error()) + } + }() + + // Auto-detect app root if not specified + if appPath == "" { + cwd, err := os.Getwd() + if err == nil { + if root, err := findAppRootFromPath(cwd); err == nil { + appPath = root + } + } + } + + Boot() + + // Set Runtime Mode + config.Conf.Runtime.Mode = "standard" + + cfg := config.Conf + cfg.Session.IsCLI = true + if len(args) < 1 { + if !runSilent { + color.Red(L("Not enough arguments\n")) + color.White(share.BUILDNAME + " help\n") + return + } + fmt.Print(L("Not enough arguments\n")) + return + } + + loadWarnings, err := engine.Load(cfg, engine.LoadOption{Action: "run"}) + if err != nil { + if !runSilent { + color.Red(L("Engine: %s\n"), err.Error()) + return + } + + fmt.Printf("%s\n", err.Error()) + return + } + + name := args[0] + if !runSilent { + color.Green(L("Run: %s\n"), name) + } + + pargs := parseRunArgs(args) + + // Start Tasks + itask.Start() + defer itask.Stop() + + // Start Schedules + ischedule.Start() + defer ischedule.Stop() + + p := process.NewWithContext(context.Background(), name, pargs...) + res, err := p.Exec() + if err != nil { + if !runSilent { + color.Red(L("Process: %s\n"), fmt.Sprintf("%s", strings.TrimPrefix(err.Error(), "Exception|404:"))) + return + } + fmt.Printf("%s\n", err.Error()) + return + } + + if !runSilent { + + if len(loadWarnings) > 0 { + fmt.Println(color.YellowString("---------------------------------")) + fmt.Println(color.YellowString(L("Warnings"))) + fmt.Println(color.YellowString("---------------------------------")) + for _, warning := range loadWarnings { + fmt.Println(color.YellowString("[%s] %s", warning.Widget, warning.Error)) + } + fmt.Printf("\n") + } + + color.White("--------------------------------------\n") + color.White(L("%s Response\n"), name) + color.White("--------------------------------------\n") + helper.Dump(res) + color.White("--------------------------------------\n") + color.Green(L("✨DONE✨\n")) + return + } + + // Silent mode output + switch res.(type) { + + case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64, bool: + fmt.Printf("%v\n", res) + return + + case string, []byte: + fmt.Printf("%s\n", res) + return + + default: + txt, err := jsoniter.Marshal(res) + if err != nil { + fmt.Printf("%s\n", err.Error()) + } + fmt.Printf("%s\n", txt) + } +} + +// parseRunArgs parses the CLI arguments into process arguments, handling :: prefixed JSON. +func parseRunArgs(args []string) []interface{} { + pargs := []interface{}{} + for i, arg := range args { + if i == 0 { + continue + } + + if strings.HasPrefix(arg, "::") { + raw := strings.TrimPrefix(arg, "::") + var v interface{} + err := jsoniter.Unmarshal([]byte(raw), &v) + if err != nil { + color.Red(L("Arguments: %s\n"), err.Error()) + return pargs + } + pargs = append(pargs, v) + if !runSilent { + color.White("args[%d]: %s\n", i-1, raw) + } + } else if strings.HasPrefix(arg, "\\::") { + cleaned := "::" + strings.TrimPrefix(arg, "\\::") + pargs = append(pargs, cleaned) + if !runSilent { + color.White("args[%d]: %s\n", i-1, cleaned) + } + } else { + pargs = append(pargs, arg) + if !runSilent { + color.White("args[%d]: %s\n", i-1, arg) + } + } + } + return pargs } // findAppRootFromPath finds the Yao application root directory by looking for app.yao diff --git a/grpc/IMPL.md b/grpc/IMPL.md index 578d2553..b3f0869b 100644 --- a/grpc/IMPL.md +++ b/grpc/IMPL.md @@ -193,7 +193,7 @@ Container token issuance uses existing `oauth.MakeAccessToken` / `oauth.MakeRefr Deliverable: `go build -o yao-grpc ./tai/grpc/cmd`. -### Phase 6: Device Flow + CLI auth ⏳ +### Phase 6: Device Flow + CLI auth ✅ Depends on: Phase 1. Three sub-phases with sequential dependency: 6.1 → 6.2 → 6.3. @@ -214,7 +214,7 @@ Backend endpoints for RFC 8628 Device Authorization Grant. Scaffolding already i Deliverable: Device flow endpoints functional — `POST /oauth/device_authorization` issues codes, `POST /oauth/token` with `grant_type=device_code` polls status. `POST /oauth/device/authorize` allows authenticated user to authorize device. -#### Phase 6.2: CUI auth/device page (frontend) ⏳ +#### Phase 6.2: CUI auth/device page (frontend) ✅ Depends on: Phase 6.1 (backend endpoints). Frontend-only task in **CUI repo**. @@ -222,8 +222,10 @@ Route: `/auth/device` (Umi convention-based routing → `pages/auth/device/index | Task | Detail | Status | |------|--------|--------| -| `pages/auth/device/index.tsx` | Device authorization page. User enters `user_code`, clicks Authorize. Uses `AuthLayout` + `AuthInput` + `AuthButton` from existing `pages/auth/components/`. | ⏳ Pending | -| `pages/auth/device/index.less` | Styles, follow `pages/auth/entry/index.less` pattern | ⏳ Pending | +| `pages/auth/device/index.tsx` | Device authorization page. User enters `user_code`, clicks Authorize. Uses `AuthLayout` + `AuthInput` + `AuthButton` from existing `pages/auth/components/`. Three states: input, success, error. i18n (zh/en), light/dark, system CSS variables only. | ✅ Done | +| `pages/auth/device/index.less` | Styles, follow `pages/auth/entry/mfa/index.less` pattern. Full responsive + dark theme. | ✅ Done | +| `openapi/user/auth.ts` | `AuthorizeDevice(userCode)` method — `POST /oauth/device/authorize` | ✅ Done | +| `layouts/index.tsx` | Register `['auth_device', '/auth/device']` in `STANDALONE_PAGES` | ✅ Done | Implementation: @@ -237,7 +239,7 @@ Implementation: Deliverable: `/auth/device` page. User authorizes CLI device login from browser. -#### Phase 6.3: CLI commands + TUI status bar ⏳ +#### 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`). @@ -258,10 +260,13 @@ 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 | +| `cmd/credential.go` | `Credential` struct, `LoadCredential`, `LoadCredentialFrom`, `SaveCredential`, `RemoveCredential` — base64-encoded JSON read/write to `~/.yao/credentials` | ✅ Done | +| `cmd/login.go` | `yao login --server ` — compute machine ID → `POST /oauth/register` (dynamic client) → `POST /oauth/device_authorization` → color-print device code + verification URL → poll `POST /oauth/token` with interval + slow_down handling → save to `~/.yao/credentials` | ✅ Done | +| `cmd/logout.go` | `yao logout` — read credentials, best-effort `POST /oauth/revoke`, delete `~/.yao/credentials` | ✅ Done | +| `cmd/run.go` | Detect credentials → gRPC mode vs local mode. `--auth ` flag loads alternate credentials file. `-s` (silent) mode: no TUI. gRPC mode renders TUI status bar then calls remote (gRPC call wiring pending Phase 4/5 integration). Local mode unchanged. | ✅ Done | +| `cmd/tui_status.go` | lipgloss `RenderStatusBar(cred)` — one-line persistent bar: `user (gRPC) │ scope: run,stream,...`. Rounded border, colored connection info. Hidden in silent mode. | ✅ Done | +| `cmd/root.go` | Register `loginCmd`, `logoutCmd` in root command | ✅ Done | +| i18n | All new strings have zh-CN translations via `langs` map | ✅ Done | **`yao run` behavior matrix:** @@ -314,17 +319,17 @@ Phase 1 (auth + server) ✅ │ ├───────────┬───────────┬──────────────────────┐ ▼ ▼ ▼ ▼ -Phase 2 ✅ Phase 3 ✅ Phase 4 ✅ Phase 6 (device flow + CLI) +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) + Phase 5 ✅ 6.1 ✅ 6.2 ✅ + (yao-grpc) (OAuth backend) (CUI page) │ │ └───────┬───────┘ ▼ - 6.3 CMD + TUI - (login/logout/run) + 6.3 ✅ + (CMD + TUI) --- V2 --- diff --git a/openapi/config.go b/openapi/config.go index b0dff37a..8e982a85 100644 --- a/openapi/config.go +++ b/openapi/config.go @@ -388,6 +388,7 @@ func (config *Config) OAuthConfig(appConfig config.Config) (*oauth.Config, error Cache: cacheStore, Store: dataStore, IssuerURL: config.OAuth.IssuerURL, + BaseURL: config.BaseURL, Signing: signingConfig, // Use the converted signing config Token: config.OAuth.Token, Security: config.OAuth.Security, diff --git a/openapi/oauth.go b/openapi/oauth.go index f306f6c9..52f19b1e 100644 --- a/openapi/oauth.go +++ b/openapi/oauth.go @@ -547,45 +547,15 @@ func (openapi *OpenAPI) oauthDeviceAuthorization(c *gin.Context) { // oauthDeviceAuthorize allows an authenticated user to authorize a pending device code. func (openapi *OpenAPI) oauthDeviceAuthorize(c *gin.Context) { - authHeader := c.GetHeader("Authorization") - if authHeader == "" || !strings.HasPrefix(authHeader, "Bearer ") { + tokenStr := extractBearerToken(c) + if tokenStr == "" { response.RespondWithSecureError(c, response.StatusUnauthorized, &response.ErrorResponse{ Code: types.ErrorInvalidGrant, ErrorDescription: "Bearer token required", }) return } - - tokenStr := strings.TrimPrefix(authHeader, "Bearer ") - oauthService := openapi.OAuth - introspection, err := oauthService.Introspect(c, tokenStr) - if err != nil || introspection == nil || !introspection.Active { - response.RespondWithSecureError(c, response.StatusUnauthorized, &response.ErrorResponse{ - Code: types.ErrorInvalidGrant, - ErrorDescription: "Invalid or expired token", - }) - return - } - - subject := introspection.Subject - if subject == "" { - response.RespondWithSecureError(c, response.StatusUnauthorized, &response.ErrorResponse{ - Code: types.ErrorInvalidGrant, - ErrorDescription: "Token has no subject", - }) - return - } - - userCode := c.PostForm("user_code") - if userCode == "" { - userCode = c.Query("user_code") - } - if userCode == "" { - response.RespondWithSecureError(c, response.StatusBadRequest, response.ErrInvalidRequest) - return - } - - svc, ok := oauthService.(*oauth.Service) + svc, ok := openapi.OAuth.(*oauth.Service) if !ok { response.RespondWithSecureError(c, response.StatusInternalServerError, &response.ErrorResponse{ Code: types.ErrorServerError, @@ -594,7 +564,52 @@ func (openapi *OpenAPI) oauthDeviceAuthorize(c *gin.Context) { return } - if err := svc.AuthorizeDevice(c, userCode, subject); err != nil { + tokenClaims, err := svc.VerifyToken(tokenStr) + if err != nil || tokenClaims == nil { + response.RespondWithSecureError(c, response.StatusUnauthorized, &response.ErrorResponse{ + Code: types.ErrorInvalidGrant, + ErrorDescription: "Invalid or expired token", + }) + return + } + + if tokenClaims.Subject == "" { + response.RespondWithSecureError(c, response.StatusUnauthorized, &response.ErrorResponse{ + Code: types.ErrorInvalidGrant, + ErrorDescription: "Token has no subject", + }) + return + } + + extraClaims := tokenClaims.Extra + if extraClaims == nil { + extraClaims = make(map[string]interface{}) + } + if tokenClaims.TeamID != "" { + extraClaims["team_id"] = tokenClaims.TeamID + } + if tokenClaims.TenantID != "" { + extraClaims["tenant_id"] = tokenClaims.TenantID + } + + userCode := c.PostForm("user_code") + if userCode == "" { + userCode = c.Query("user_code") + } + if userCode == "" { + var body struct { + UserCode string `json:"user_code"` + } + if c.ShouldBindJSON(&body) == nil { + userCode = body.UserCode + } + } + if userCode == "" { + response.RespondWithSecureError(c, response.StatusBadRequest, response.ErrInvalidRequest) + return + } + + if err := svc.AuthorizeDevice(c, userCode, tokenClaims.Subject, extraClaims); err != nil { if oauthErr, ok := err.(*response.ErrorResponse); ok { response.RespondWithSecureError(c, response.StatusBadRequest, oauthErr) } else { @@ -705,3 +720,16 @@ func (openapi *OpenAPI) getParam(c *gin.Context, key string) string { // Then try to get from POST form data (POST request) return c.PostForm(key) } + +// extractBearerToken reads the access token from Authorization header or cookie, +// matching the same logic as guard.getAccessToken. +func extractBearerToken(c *gin.Context) string { + if auth := c.GetHeader("Authorization"); strings.HasPrefix(auth, "Bearer ") { + return strings.TrimPrefix(auth, "Bearer ") + } + cookieName := response.GetCookieName("access_token") + if cookie, err := c.Cookie(cookieName); err == nil && cookie != "" { + return strings.TrimPrefix(cookie, "Bearer ") + } + return "" +} diff --git a/openapi/oauth/core.go b/openapi/oauth/core.go index 3cbb8411..9862a5ee 100644 --- a/openapi/oauth/core.go +++ b/openapi/oauth/core.go @@ -6,6 +6,7 @@ import ( "time" "github.com/yaoapp/yao/openapi/oauth/types" + "go.mongodb.org/mongo-driver/bson/primitive" ) // AuthorizationServer returns the authorization server endpoint URL @@ -667,8 +668,18 @@ func (s *Service) handleDeviceCodeGrant(ctx context.Context, client *types.Clien subject, _ := codeData["subject"].(string) s.consumeDeviceCode(deviceCode) + var extraClaims map[string]interface{} + if ec, ok := codeData["extra_claims"]; ok { + switch v := ec.(type) { + case map[string]interface{}: + extraClaims = v + case primitive.M: + extraClaims = map[string]interface{}(v) + } + } + expiresIn := int(s.config.Token.AccessTokenLifetime.Seconds()) - accessToken, err := s.generateAccessTokenWithScope(client.ClientID, scope, subject, expiresIn, nil) + accessToken, err := s.generateAccessTokenWithScope(client.ClientID, scope, subject, expiresIn, extraClaims) if err != nil { return nil, &types.ErrorResponse{ Code: types.ErrorServerError, @@ -683,7 +694,7 @@ func (s *Service) handleDeviceCodeGrant(ctx context.Context, client *types.Clien } if types.Contains(client.GrantTypes, types.GrantTypeRefreshToken) { - refreshToken, err := s.generateRefreshToken(client.ClientID, scope, subject, 0, nil) + refreshToken, err := s.generateRefreshToken(client.ClientID, scope, subject, 0, extraClaims) if err != nil { return nil, &types.ErrorResponse{ Code: types.ErrorServerError, diff --git a/openapi/oauth/device.go b/openapi/oauth/device.go index ed567513..7c5471c3 100644 --- a/openapi/oauth/device.go +++ b/openapi/oauth/device.go @@ -72,7 +72,7 @@ func (s *Service) DeviceAuthorization(ctx context.Context, clientID string, scop } // AuthorizeDevice allows an authenticated user to authorize a device code via user_code. -func (s *Service) AuthorizeDevice(ctx context.Context, userCode string, subject string) error { +func (s *Service) AuthorizeDevice(ctx context.Context, userCode string, subject string, extraClaims ...map[string]interface{}) error { if !s.config.Features.DeviceFlowEnabled { return &types.ErrorResponse{ Code: types.ErrorUnsupportedGrantType, @@ -86,7 +86,11 @@ func (s *Service) AuthorizeDevice(ctx context.Context, userCode string, subject formatted = normalized[:4] + "-" + normalized[4:] } - return s.authorizeDeviceCode(formatted, subject) + var claims map[string]interface{} + if len(extraClaims) > 0 { + claims = extraClaims[0] + } + return s.authorizeDeviceCode(formatted, subject, claims) } // generateUserCode generates a user-friendly code formatted as XXXX-XXXX. diff --git a/openapi/oauth/discovery.go b/openapi/oauth/discovery.go index b7f44dd7..b577c4e8 100644 --- a/openapi/oauth/discovery.go +++ b/openapi/oauth/discovery.go @@ -6,6 +6,7 @@ import ( "encoding/base64" "fmt" "math/big" + "strings" "github.com/yaoapp/yao/openapi/oauth/types" ) @@ -53,7 +54,7 @@ func (s *Service) JWKS(ctx context.Context) (*types.JWKSResponse, error) { // Endpoints returns a map of all available OAuth endpoints // This provides endpoint discovery for clients func (s *Service) Endpoints(ctx context.Context) (map[string]string, error) { - baseURL := s.config.IssuerURL + baseURL := strings.TrimRight(s.config.IssuerURL, "/") + s.config.BaseURL endpoints := map[string]string{ "authorization_endpoint": fmt.Sprintf("%s/oauth/authorize", baseURL), diff --git a/openapi/oauth/oauth.go b/openapi/oauth/oauth.go index b7012502..f313482e 100644 --- a/openapi/oauth/oauth.go +++ b/openapi/oauth/oauth.go @@ -57,6 +57,7 @@ type Config struct { // OAuth server metadata IssuerURL string `json:"issuer_url"` // JWT token issuer URL + BaseURL string `json:"base_url"` // API route prefix (e.g. "/v1") } // FeatureFlags represents feature toggle configuration diff --git a/openapi/oauth/token.go b/openapi/oauth/token.go index 586764bb..5d8b52d3 100644 --- a/openapi/oauth/token.go +++ b/openapi/oauth/token.go @@ -596,7 +596,7 @@ func (s *Service) getDeviceCodeData(deviceCode string) (map[string]interface{}, } // authorizeDeviceCode marks a device code as authorized via user_code lookup -func (s *Service) authorizeDeviceCode(userCode, subject string) error { +func (s *Service) authorizeDeviceCode(userCode, subject string, extraClaims map[string]interface{}) error { reverseData, exists := s.store.Get(s.userCodeKey(userCode)) if !exists { return &types.ErrorResponse{ @@ -626,6 +626,9 @@ func (s *Service) authorizeDeviceCode(userCode, subject string) error { codeData["status"] = "authorized" codeData["subject"] = subject + if extraClaims != nil { + codeData["extra_claims"] = extraClaims + } // Re-store with remaining TTL expiresAt, _ := codeData["expires_at"].(int64) diff --git a/openapi/well-known.go b/openapi/well-known.go index dd38fe82..85e42377 100644 --- a/openapi/well-known.go +++ b/openapi/well-known.go @@ -1,10 +1,14 @@ package openapi import ( + "fmt" + "net" "os" + "strconv" "strings" "github.com/gin-gonic/gin" + "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/share" ) @@ -45,6 +49,7 @@ type YaoMetadata struct { // Dashboard configuration Dashboard string `json:"dashboard,omitempty"` // Admin dashboard root path + GRPC string `json:"grpc,omitempty"` // gRPC server address (e.g., "127.0.0.1:9099") Optional map[string]interface{} `json:"optional,omitempty"` // Optional settings // Developer information @@ -67,6 +72,7 @@ func (openapi *OpenAPI) yaoMetadata(c *gin.Context) { IssuerURL: openapi.Config.OAuth.IssuerURL, ServerURL: resolveServerURL(openapi.Config.OAuth.IssuerURL), Dashboard: "/" + dashboard, + GRPC: resolveGRPCAddr(c), Optional: share.App.Optional, } @@ -97,8 +103,46 @@ func resolveServerURL(issuerURL string) string { return "" } +// resolveGRPCAddr returns the gRPC server address for client discovery. +// Uses the request Host's IP with the configured gRPC port. +func resolveGRPCAddr(c *gin.Context) string { + cfg := config.Conf.GRPC + if strings.ToLower(cfg.Enabled) == "off" { + return "" + } + port := cfg.Port + if port == 0 { + port = 9099 + } + + host := cfg.Host + if host == "" || host == "0.0.0.0" { + reqHost := c.Request.Host + h, _, err := net.SplitHostPort(reqHost) + if err != nil { + h = reqHost + } + host = h + } else if strings.Contains(host, ",") { + host = strings.TrimSpace(strings.Split(host, ",")[0]) + } + + return fmt.Sprintf("%s:%s", host, strconv.Itoa(port)) +} + // oauthServerMetadata returns authorization server metadata - RFC 8414 -func (openapi *OpenAPI) oauthServerMetadata(c *gin.Context) {} +func (openapi *OpenAPI) oauthServerMetadata(c *gin.Context) { + if openapi.OAuth == nil { + c.JSON(503, gin.H{"error": "OAuth service not available"}) + return + } + metadata, err := openapi.OAuth.GetServerMetadata(c.Request.Context()) + if err != nil { + c.JSON(500, gin.H{"error": err.Error()}) + return + } + c.JSON(200, metadata) +} // oauthOpenIDConfiguration returns OpenID Connect configuration func (openapi *OpenAPI) oauthOpenIDConfiguration(c *gin.Context) {} From df63584a6f42a8cd540b89cfe08b4c8e00380914 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 4 Mar 2026 17:13:06 +0800 Subject: [PATCH 7/7] Enhance testing and configuration for Tai service - Add TAI_TEST_HOST_IP environment variable to CI workflows for unit tests, allowing better connectivity to the gRPC server from Docker containers. - Update the `run.go` file to parse command-line arguments correctly. - Modify the test utility to return the gRPC address reachable from Docker, improving integration test reliability. - Refactor integration tests to utilize the new relay address function, ensuring proper communication with the Yao gRPC server. These changes improve the testing framework and enhance the configuration for better service interaction during CI runs. --- .github/workflows/pr-test.yml | 1 + .github/workflows/unit-test.yml | 1 + cmd/run.go | 2 +- grpc/tests/testutils/testutils.go | 24 +++++++++++++++++++++++- tai/grpc/integration_test.go | 10 +++++----- 5 files changed, 31 insertions(+), 7 deletions(-) diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 3f1bcfa4..ed92566b 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -1757,6 +1757,7 @@ jobs: TAI_TEST_K8S_HOST: "127.0.0.1" TAI_TEST_K8S_PORT: "6443" TAI_TEST_KUBECONFIG: "${{ runner.temp }}/kubeconfig-tai.yml" + TAI_TEST_HOST_IP: "172.17.0.1" run: make unit-test-tai - name: Codecov Report diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index 3e04605e..07d7111c 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -1311,6 +1311,7 @@ jobs: TAI_TEST_K8S_HOST: "127.0.0.1" TAI_TEST_K8S_PORT: "6443" TAI_TEST_KUBECONFIG: "${{ runner.temp }}/kubeconfig-tai.yml" + TAI_TEST_HOST_IP: "172.17.0.1" run: make unit-test-tai - name: Codecov Report diff --git a/cmd/run.go b/cmd/run.go index f6dcca3b..819d2f75 100644 --- a/cmd/run.go +++ b/cmd/run.go @@ -85,7 +85,7 @@ func runGRPC(cred *Credential, args []string) { color.Green(L("Run: %s gRPC: %s\n"), name, cred.GRPCAddr) } - pargs := parseRunArgs(args[1:]) + pargs := parseRunArgs(args) argsJSON, err := jsoniter.Marshal(pargs) if err != nil { diff --git a/grpc/tests/testutils/testutils.go b/grpc/tests/testutils/testutils.go index d7abab80..c384af30 100644 --- a/grpc/tests/testutils/testutils.go +++ b/grpc/tests/testutils/testutils.go @@ -2,6 +2,8 @@ package testutils import ( "context" + "net" + "os" "strings" "testing" @@ -43,7 +45,7 @@ func Prepare(t *testing.T) *grpc.ClientConn { cfg := config.Conf cfg.GRPC.Port = 0 - cfg.GRPC.Host = "127.0.0.1" + cfg.GRPC.Host = "0.0.0.0" cfg.GRPC.Enabled = "" test.Prepare(t, config.Conf) @@ -129,6 +131,26 @@ func Addr() string { return addrs[0] } +// RelayAddr returns the gRPC address reachable from a Docker container. +// When TAI_TEST_HOST_IP is set (e.g. to the docker bridge gateway), +// it replaces the host portion so that the Tai container can reach the +// Yao gRPC server running on the CI host. +func RelayAddr() string { + addr := Addr() + if addr == "" { + return "" + } + hostIP := os.Getenv("TAI_TEST_HOST_IP") + if hostIP == "" { + return addr + } + _, port, err := net.SplitHostPort(addr) + if err != nil { + return addr + } + return hostIP + ":" + port +} + // ObtainAccessToken mints a token with the given scopes via oauth.MakeAccessToken. func ObtainAccessToken(t *testing.T, scopes ...string) string { t.Helper() diff --git a/tai/grpc/integration_test.go b/tai/grpc/integration_test.go index 5b534409..4073e163 100644 --- a/tai/grpc/integration_test.go +++ b/tai/grpc/integration_test.go @@ -280,11 +280,11 @@ func setupRelayClient(t *testing.T, scopes ...string) *yaogrpc.Client { testutils.Clean() }) - yaoAddr := testutils.Addr() + yaoAddr := testutils.RelayAddr() token := testutils.ObtainAccessToken(t, scopes...) refreshToken := testutils.ObtainRefreshToken(t, scopes...) - // upstream = Yao gRPC address; taiMode = true + // upstream = Yao gRPC address reachable from the Tai container tm := yaogrpc.NewTokenManager(token, refreshToken, "relay-sandbox", yaoAddr) client, err := yaogrpc.Dial(taiAddr, tm) require.NoError(t, err) @@ -305,7 +305,7 @@ func TestRelay_Healthz(t *testing.T) { testutils.Clean() }() - yaoAddr := testutils.Addr() + yaoAddr := testutils.RelayAddr() tm := yaogrpc.NewTokenManager("", "", "", yaoAddr) client, err := yaogrpc.Dial(taiAddr, tm) require.NoError(t, err) @@ -375,7 +375,7 @@ func TestRelay_Run_NoToken(t *testing.T) { testutils.Clean() }() - yaoAddr := testutils.Addr() + yaoAddr := testutils.RelayAddr() tm := yaogrpc.NewTokenManager("", "", "", yaoAddr) client, err := yaogrpc.Dial(taiAddr, tm) require.NoError(t, err) @@ -398,7 +398,7 @@ func TestRelay_TokenRefresh(t *testing.T) { testutils.Clean() }() - yaoAddr := testutils.Addr() + yaoAddr := testutils.RelayAddr() scopes := []string{"grpc:run"} expiredToken := testutils.ObtainExpiredAccessToken(t, scopes...)