223d02ebfe
- Added support for issuing and passing a sandbox token during stream execution in the Assistant. - Updated StreamRequest to include a Token field for managing user authentication. - Enhanced ClaudeRunner to set environment variables for the sandbox token and refresh token. - Refactored SandboxToken structure to clarify its purpose and manage token credentials effectively. - Modified sandbox management routes to enforce OAuth guard for improved security. Made-with: Cursor
387 lines
10 KiB
Go
387 lines
10 KiB
Go
package tai
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net"
|
|
"net/http"
|
|
"time"
|
|
|
|
hepb "github.com/yaoapp/yao/tai/hostexec/pb"
|
|
"github.com/yaoapp/yao/tai/proxy"
|
|
"github.com/yaoapp/yao/tai/registry"
|
|
"github.com/yaoapp/yao/tai/runtime"
|
|
sipb "github.com/yaoapp/yao/tai/serverinfo/pb"
|
|
"github.com/yaoapp/yao/tai/types"
|
|
"github.com/yaoapp/yao/tai/vnc"
|
|
"github.com/yaoapp/yao/tai/volume"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/credentials/insecure"
|
|
"google.golang.org/grpc/keepalive"
|
|
)
|
|
|
|
// DialRemote establishes connections to a remote Tai node via gRPC (direct mode).
|
|
// Does NOT interact with the registry. Caller must call ConnResources.Close().
|
|
func DialRemote(host string, ports types.Ports, opts ...DialOption) (*ConnResources, error) {
|
|
cfg := &dialConfig{ports: mergedPorts(ports)}
|
|
for _, o := range opts {
|
|
o.applyDial(cfg)
|
|
}
|
|
|
|
grpcAddr := fmt.Sprintf("%s:%d", host, cfg.ports.GRPC)
|
|
conn, err := dialGRPC(grpcAddr)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("grpc dial %s: %w", grpcAddr, err)
|
|
}
|
|
|
|
return buildResources(conn, cfg, &remoteEnv{host: host, httpClient: cfg.httpClient})
|
|
}
|
|
|
|
// DialTunnel establishes connections to a Tai node through the WebSocket tunnel.
|
|
// Requires the node to already be registered in the registry (online).
|
|
// Does NOT call registry.SetResources. Caller must call ConnResources.Close().
|
|
func DialTunnel(taiID string, reg *registry.Registry, opts ...DialOption) (*ConnResources, error) {
|
|
node, ok := reg.Get(taiID)
|
|
if !ok || node.Status != "online" {
|
|
return nil, fmt.Errorf("tai node %s not online", taiID)
|
|
}
|
|
|
|
cfg := &dialConfig{
|
|
ports: types.Ports{
|
|
GRPC: intOr(node.Ports.GRPC, 19100),
|
|
HTTP: intOr(node.Ports.HTTP, 8099),
|
|
VNC: intOr(node.Ports.VNC, 16080),
|
|
Docker: intOr(node.Ports.Docker, 12375),
|
|
K8s: intOr(node.Ports.K8s, 16443),
|
|
},
|
|
}
|
|
for _, o := range opts {
|
|
o.applyDial(cfg)
|
|
}
|
|
|
|
grpcLn, err := reg.OpenLocalListener(taiID, cfg.ports.GRPC)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("open grpc tunnel listener: %w", err)
|
|
}
|
|
|
|
conn, err := dialGRPC("passthrough:///" + grpcLn.Addr().String())
|
|
if err != nil {
|
|
grpcLn.Close()
|
|
return nil, fmt.Errorf("grpc dial tunnel %s: %w", grpcLn.Addr(), err)
|
|
}
|
|
|
|
env := &tunnelEnv{
|
|
taiID: taiID,
|
|
yaoBase: node.YaoBase,
|
|
reg: reg,
|
|
regCaps: node.Capabilities,
|
|
listeners: []net.Listener{grpcLn},
|
|
}
|
|
|
|
res, err := buildResources(conn, cfg, env)
|
|
if err != nil {
|
|
grpcLn.Close()
|
|
conn.Close()
|
|
return nil, err
|
|
}
|
|
res.Listeners = env.listeners
|
|
return res, nil
|
|
}
|
|
|
|
// DialLocal establishes connections to the local Docker daemon.
|
|
// Does NOT interact with the registry. Caller must call ConnResources.Close().
|
|
func DialLocal(addr string, dataDir string, vol volume.Volume) (*ConnResources, error) {
|
|
sb, err := runtime.NewLocal(addr)
|
|
if err != nil && vol == nil {
|
|
return nil, err
|
|
}
|
|
|
|
res := &ConnResources{DataDir: dataDir}
|
|
|
|
if sb != nil {
|
|
res.Runtime = sb
|
|
res.Image = runtime.NewDockerImage(runtime.DockerCli(sb))
|
|
res.Proxy = proxy.NewLocal(sb)
|
|
res.VNC = vnc.NewLocal(sb)
|
|
}
|
|
|
|
if vol != nil {
|
|
res.Volume = vol
|
|
} else {
|
|
if dataDir == "" {
|
|
dataDir = "/tmp/tai-volumes"
|
|
}
|
|
res.DataDir = dataDir
|
|
res.Volume = volume.NewLocal(dataDir)
|
|
}
|
|
|
|
return res, nil
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Shared build logic
|
|
// ---------------------------------------------------------------------------
|
|
|
|
// dialEnv abstracts the mode-specific differences (remote vs tunnel) that
|
|
// buildResources needs.
|
|
type dialEnv interface {
|
|
fallbackCaps() map[string]bool
|
|
mergeCaps(discovered map[string]bool) types.Capabilities
|
|
// listenAddr opens or formats a host:port address for the given port.
|
|
// Tunnel mode opens a local listener; remote mode formats host:port.
|
|
listenAddr(port int) (string, error)
|
|
newProxy(ports types.Ports) proxy.Proxy
|
|
newVNC(ports types.Ports) vnc.VNC
|
|
}
|
|
|
|
// buildResources constructs a ConnResources from an established gRPC
|
|
// connection. Shared by DialRemote and DialTunnel.
|
|
func buildResources(conn *grpc.ClientConn, cfg *dialConfig, env dialEnv) (*ConnResources, error) {
|
|
info, err := discoverInfo(conn, cfg)
|
|
if err != nil {
|
|
info = &discoveredInfo{Capabilities: env.fallbackCaps()}
|
|
}
|
|
|
|
caps := env.mergeCaps(info.Capabilities)
|
|
|
|
res := &ConnResources{
|
|
GRPCConn: conn,
|
|
HostExec: hepb.NewHostExecClient(conn),
|
|
Volume: volume.NewRemote(conn),
|
|
Caps: caps,
|
|
System: info.System,
|
|
Ports: cfg.ports,
|
|
Version: info.Version,
|
|
}
|
|
|
|
if cfg.runtime == types.K8s || (!caps.Docker && caps.K8s) {
|
|
if cfg.kubeConfig != "" {
|
|
k8sPort := cfg.ports.K8s
|
|
if k8sPort == 0 {
|
|
k8sPort = 16443
|
|
}
|
|
addr, err := env.listenAddr(k8sPort)
|
|
if err == nil {
|
|
sb, err := runtime.NewK8s(addr, runtime.K8sOption{
|
|
Namespace: cfg.namespace,
|
|
KubeConfig: cfg.kubeConfig,
|
|
})
|
|
if err == nil {
|
|
res.Runtime = sb
|
|
res.Image = runtime.NewK8sImage()
|
|
}
|
|
}
|
|
}
|
|
} else if caps.Docker {
|
|
dockerPort := cfg.ports.Docker
|
|
if dockerPort == 0 {
|
|
dockerPort = 12375
|
|
}
|
|
addr, err := env.listenAddr(dockerPort)
|
|
if err == nil {
|
|
sb, err := runtime.NewDocker("tcp://" + addr)
|
|
if err == nil {
|
|
res.Runtime = sb
|
|
res.Image = runtime.NewDockerImage(runtime.DockerCli(sb))
|
|
}
|
|
}
|
|
}
|
|
|
|
if res.Runtime != nil {
|
|
res.Proxy = env.newProxy(cfg.ports)
|
|
}
|
|
res.VNC = env.newVNC(cfg.ports)
|
|
|
|
return res, nil
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// remoteEnv — direct TCP connections
|
|
// ---------------------------------------------------------------------------
|
|
|
|
type remoteEnv struct {
|
|
host string
|
|
httpClient *http.Client
|
|
}
|
|
|
|
func (e *remoteEnv) fallbackCaps() map[string]bool {
|
|
return map[string]bool{"docker": true}
|
|
}
|
|
|
|
func (e *remoteEnv) mergeCaps(discovered map[string]bool) types.Capabilities {
|
|
return types.Capabilities{
|
|
Docker: discovered["docker"],
|
|
K8s: discovered["k8s"],
|
|
HostExec: discovered["host_exec"],
|
|
}
|
|
}
|
|
|
|
func (e *remoteEnv) listenAddr(port int) (string, error) {
|
|
return fmt.Sprintf("%s:%d", e.host, port), nil
|
|
}
|
|
|
|
func (e *remoteEnv) newProxy(ports types.Ports) proxy.Proxy {
|
|
return proxy.NewRemote(e.host, ports.HTTP, e.httpClient)
|
|
}
|
|
|
|
func (e *remoteEnv) newVNC(ports types.Ports) vnc.VNC {
|
|
return vnc.NewRemote(e.host, ports.VNC, e.httpClient)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// tunnelEnv — connections via WebSocket tunnel
|
|
// ---------------------------------------------------------------------------
|
|
|
|
type tunnelEnv struct {
|
|
taiID string
|
|
yaoBase string
|
|
reg *registry.Registry
|
|
regCaps types.Capabilities
|
|
listeners []net.Listener
|
|
}
|
|
|
|
func (e *tunnelEnv) fallbackCaps() map[string]bool {
|
|
return make(map[string]bool)
|
|
}
|
|
|
|
func (e *tunnelEnv) mergeCaps(discovered map[string]bool) types.Capabilities {
|
|
return types.Capabilities{
|
|
Docker: discovered["docker"] || e.regCaps.Docker,
|
|
K8s: discovered["k8s"] || e.regCaps.K8s,
|
|
HostExec: discovered["host_exec"] || e.regCaps.HostExec,
|
|
}
|
|
}
|
|
|
|
func (e *tunnelEnv) listenAddr(port int) (string, error) {
|
|
ln, err := e.reg.OpenLocalListener(e.taiID, port)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
e.listeners = append(e.listeners, ln)
|
|
return ln.Addr().String(), nil
|
|
}
|
|
|
|
func (e *tunnelEnv) newProxy(_ types.Ports) proxy.Proxy {
|
|
return proxy.NewTunnel(e.taiID, e.yaoBase)
|
|
}
|
|
|
|
func (e *tunnelEnv) newVNC(_ types.Ports) vnc.VNC {
|
|
return vnc.NewTunnel(e.taiID, e.yaoBase)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Dial options
|
|
// ---------------------------------------------------------------------------
|
|
|
|
// DialOption configures a Dial* call.
|
|
type DialOption interface {
|
|
applyDial(*dialConfig)
|
|
}
|
|
|
|
type dialOptionFunc func(*dialConfig)
|
|
|
|
func (f dialOptionFunc) applyDial(c *dialConfig) { f(c) }
|
|
|
|
// WithDialRuntime selects the container runtime for the dial call.
|
|
func WithDialRuntime(rt types.Runtime) DialOption {
|
|
return dialOptionFunc(func(c *dialConfig) { c.runtime = rt })
|
|
}
|
|
|
|
// WithDialKubeConfig sets the kubeconfig for K8s runtime.
|
|
func WithDialKubeConfig(path string) DialOption {
|
|
return dialOptionFunc(func(c *dialConfig) { c.kubeConfig = path })
|
|
}
|
|
|
|
// WithDialNamespace sets the K8s namespace.
|
|
func WithDialNamespace(ns string) DialOption {
|
|
return dialOptionFunc(func(c *dialConfig) { c.namespace = ns })
|
|
}
|
|
|
|
// WithDialHTTPClient sets a custom HTTP client for proxy/VNC.
|
|
func WithDialHTTPClient(hc *http.Client) DialOption {
|
|
return dialOptionFunc(func(c *dialConfig) { c.httpClient = hc })
|
|
}
|
|
|
|
type dialConfig struct {
|
|
runtime types.Runtime
|
|
ports types.Ports
|
|
kubeConfig string
|
|
namespace string
|
|
httpClient *http.Client
|
|
userPorts types.Ports
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Helpers
|
|
// ---------------------------------------------------------------------------
|
|
|
|
func dialGRPC(target string) (*grpc.ClientConn, error) {
|
|
return grpc.NewClient(target,
|
|
grpc.WithTransportCredentials(insecure.NewCredentials()),
|
|
grpc.WithKeepaliveParams(keepalive.ClientParameters{
|
|
Time: 20 * time.Second,
|
|
Timeout: 5 * time.Second,
|
|
PermitWithoutStream: true,
|
|
}),
|
|
)
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// ServerInfo discovery (shared by DialRemote / DialTunnel)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
type discoveredInfo struct {
|
|
Capabilities map[string]bool
|
|
System types.SystemInfo
|
|
Version string
|
|
}
|
|
|
|
func discoverInfo(conn *grpc.ClientConn, cfg *dialConfig) (*discoveredInfo, error) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
client := sipb.NewServerInfoClient(conn)
|
|
resp, err := client.GetInfo(ctx, &sipb.GetInfoRequest{})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
up := cfg.userPorts
|
|
|
|
if p := int(resp.Ports["http"]); p > 0 && up.HTTP == 0 {
|
|
cfg.ports.HTTP = p
|
|
}
|
|
if p := int(resp.Ports["docker"]); p > 0 && up.Docker == 0 {
|
|
cfg.ports.Docker = p
|
|
}
|
|
if p := int(resp.Ports["vnc"]); p > 0 && up.VNC == 0 {
|
|
cfg.ports.VNC = p
|
|
}
|
|
if p := int(resp.Ports["k8s"]); p > 0 && up.K8s == 0 {
|
|
cfg.ports.K8s = p
|
|
}
|
|
|
|
caps := resp.Capabilities
|
|
if caps == nil {
|
|
caps = make(map[string]bool)
|
|
}
|
|
|
|
var sys types.SystemInfo
|
|
if s := resp.System; s != nil {
|
|
sys = types.SystemInfo{
|
|
OS: s.Os,
|
|
Arch: s.Arch,
|
|
Hostname: s.Hostname,
|
|
NumCPU: int(s.NumCpu),
|
|
TotalMem: s.TotalMem,
|
|
Shell: s.Shell,
|
|
TempDir: s.TempDir,
|
|
}
|
|
}
|
|
|
|
return &discoveredInfo{
|
|
Capabilities: caps,
|
|
System: sys,
|
|
Version: resp.Version,
|
|
}, nil
|
|
}
|