From 223d02ebfe9085478ab709ce7d4b41ea368db8ae Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 14 Mar 2026 15:50:59 +0800 Subject: [PATCH] feat(sandbox): implement sandbox token handling in stream execution - 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 --- agent/assistant/sandbox_v2.go | 10 ++++ agent/sandbox/v2/claude/runner.go | 9 ++++ agent/sandbox/v2/token.go | 90 +++++++++++++++++++++++++++++++ agent/sandbox/v2/types/runner.go | 1 + agent/sandbox/v2/types/token.go | 9 ++-- openapi/computer/computer.go | 2 + openapi/openapi.go | 2 +- openapi/sandbox/manage.go | 15 +++--- sandbox/v2/manager.go | 13 +++++ tai/dial.go | 2 +- 10 files changed, 138 insertions(+), 15 deletions(-) create mode 100644 agent/sandbox/v2/token.go diff --git a/agent/assistant/sandbox_v2.go b/agent/assistant/sandbox_v2.go index a63ba159..a8026ae9 100644 --- a/agent/assistant/sandbox_v2.go +++ b/agent/assistant/sandbox_v2.go @@ -149,6 +149,15 @@ func (ast *Assistant) executeSandboxV2Stream( // Resolve connector for Stream. conn, _, _ := ast.GetConnector(ctx) + var tok *sandboxTypes.SandboxToken + if ctx.Authorized != nil { + var err error + tok, err = sandboxv2.IssueSandboxToken(ctx.Authorized.TeamID, ctx.Authorized.UserID) + if err != nil { + return nil, fmt.Errorf("issue sandbox token: %w", err) + } + } + streamReq := &sandboxTypes.StreamRequest{ Computer: computer, Config: cfg, @@ -156,6 +165,7 @@ func (ast *Assistant) executeSandboxV2Stream( Messages: completionMessages, SystemPrompt: systemPrompt, ChatID: ctx.ChatID, + Token: tok, } execReq := &sandboxv2.ExecuteRequest{ diff --git a/agent/sandbox/v2/claude/runner.go b/agent/sandbox/v2/claude/runner.go index 886fde38..576933a7 100644 --- a/agent/sandbox/v2/claude/runner.go +++ b/agent/sandbox/v2/claude/runner.go @@ -235,6 +235,15 @@ func (r *ClaudeRunner) buildCLICommand(req *types.StreamRequest, oe *osEnv, isCo } } + if req.Token != nil { + if req.Token.Token != "" { + env["YAO_TOKEN"] = req.Token.Token + } + if req.Token.RefreshToken != "" { + env["YAO_REFRESH_TOKEN"] = req.Token.RefreshToken + } + } + var systemPrompt string envPrompt := buildSandboxEnvPrompt(oe.WorkDir) if !isContinuation && req.SystemPrompt != "" { diff --git a/agent/sandbox/v2/token.go b/agent/sandbox/v2/token.go new file mode 100644 index 00000000..dc47f9d0 --- /dev/null +++ b/agent/sandbox/v2/token.go @@ -0,0 +1,90 @@ +package sandboxv2 + +import ( + "fmt" + "time" + + lrustore "github.com/yaoapp/gou/store/lru" + "github.com/yaoapp/yao/agent/sandbox/v2/types" + "github.com/yaoapp/yao/openapi/oauth" +) + +const ( + accessTokenTTL = 2 * time.Hour + refreshTokenTTL = 30 * 24 * time.Hour // 30 days + tokenCacheSize = 1024 +) + +var tokenCache *lrustore.Cache + +func init() { + c, err := lrustore.New(tokenCacheSize) + if err != nil { + panic("sandbox token cache init failed: " + err.Error()) + } + tokenCache = c +} + +func cacheKey(teamID, userID string) string { + if teamID == "" { + return userID + } + return teamID + "/" + userID +} + +func getToken(teamID, userID string) *types.SandboxToken { + val, ok := tokenCache.Get(cacheKey(teamID, userID)) + if !ok { + return nil + } + tok, _ := val.(*types.SandboxToken) + return tok +} + +func setToken(teamID, userID string, tok *types.SandboxToken, ttl time.Duration) { + tokenCache.Set(cacheKey(teamID, userID), tok, ttl) +} + +// IssueSandboxToken returns a valid identity token for the given user. +// Tokens are cached by (teamID, userID); a new token is only issued on +// cache miss or expiry. Returns nil without error when oauth.OAuth is nil. +func IssueSandboxToken(teamID, userID string) (*types.SandboxToken, error) { + if tok := getToken(teamID, userID); tok != nil { + return tok, nil + } + + svc := oauth.OAuth + if svc == nil { + return nil, nil + } + + subject, err := svc.Subject("__yao.sandbox", userID) + if err != nil { + return nil, fmt.Errorf("sandbox token: derive subject: %w", err) + } + + extraClaims := map[string]interface{}{ + "user_id": userID, + } + if teamID != "" { + extraClaims["team_id"] = teamID + } + + tokenStr, err := svc.MakeAccessToken("__yao.sandbox", "sandbox:mcp", subject, + int(accessTokenTTL.Seconds()), extraClaims) + if err != nil { + return nil, fmt.Errorf("sandbox token: issue access token: %w", err) + } + + tok := &types.SandboxToken{Token: tokenStr} + + refreshStr, err := svc.MakeRefreshToken("__yao.sandbox", "sandbox:mcp", subject, + int(refreshTokenTTL.Seconds()), extraClaims) + if err != nil { + return nil, fmt.Errorf("sandbox token: issue refresh token: %w", err) + } + tok.RefreshToken = refreshStr + + setToken(teamID, userID, tok, accessTokenTTL) + return tok, nil +} diff --git a/agent/sandbox/v2/types/runner.go b/agent/sandbox/v2/types/runner.go index 5a90f80f..8674e2e2 100644 --- a/agent/sandbox/v2/types/runner.go +++ b/agent/sandbox/v2/types/runner.go @@ -50,4 +50,5 @@ type StreamRequest struct { Messages []agentContext.Message SystemPrompt string ChatID string + Token *SandboxToken // current user's sandbox token for MCP callbacks } diff --git a/agent/sandbox/v2/types/token.go b/agent/sandbox/v2/types/token.go index 7f1cf525..4900f6eb 100644 --- a/agent/sandbox/v2/types/token.go +++ b/agent/sandbox/v2/types/token.go @@ -1,9 +1,8 @@ package types -import "time" - -// SandboxToken is a short-lived JWT issued for a sandbox computer. +// SandboxToken holds credentials for a sandbox execution session. +// Expiry is managed by the LRU store TTL, not stored here. type SandboxToken struct { - Token string - ExpiresAt time.Time + Token string // access token → YAO_TOKEN + RefreshToken string // refresh token → YAO_REFRESH_TOKEN } diff --git a/openapi/computer/computer.go b/openapi/computer/computer.go index fcd6420d..e0c80837 100644 --- a/openapi/computer/computer.go +++ b/openapi/computer/computer.go @@ -37,6 +37,7 @@ type computerOption struct { Kind string `json:"kind"` ID string `json:"id"` DisplayName string `json:"display_name"` + ContainerID string `json:"container_id,omitempty"` NodeID string `json:"node_id"` Status string `json:"status"` Mode string `json:"mode,omitempty"` @@ -276,6 +277,7 @@ func boxToOption(b *sandboxv2.Box) computerOption { Kind: "box", ID: snap.ID, DisplayName: displayName, + ContainerID: snap.ContainerID, NodeID: snap.NodeID, Status: snap.Status, Mode: mode, diff --git a/openapi/openapi.go b/openapi/openapi.go index 4a4a0b7c..4c20ccc4 100644 --- a/openapi/openapi.go +++ b/openapi/openapi.go @@ -180,7 +180,7 @@ func (openapi *OpenAPI) Attach(router *gin.Engine) { sandbox.SetPathPrefix(baseURL) sandboxGroup := group.Group("/sandbox") sandbox.Attach(sandboxGroup, openapi.OAuth) - sandbox.AttachManage(sandboxGroup) + sandbox.AttachManage(sandboxGroup, openapi.OAuth) // Computer option handlers (for InputArea selector) openapiComputer.Attach(group.Group("/computer"), openapi.OAuth) diff --git a/openapi/sandbox/manage.go b/openapi/sandbox/manage.go index 71754f94..270c3d62 100644 --- a/openapi/sandbox/manage.go +++ b/openapi/sandbox/manage.go @@ -18,20 +18,19 @@ import ( ) // AttachManage registers sandbox management CRUD routes on the given group. -// oauth.Guard is already applied by the parent Attach call on the same group. // - GET / — list sandboxes (filtered by owner) // - POST / — create sandbox (owner from token) // - GET /:id — get sandbox (owner check) // - DELETE /:id — remove sandbox (owner check) // - POST /:id/exec — execute command (owner check) // - POST /:id/heartbeat — heartbeat (owner check) -func AttachManage(group *gin.RouterGroup) { - group.GET("", handleList) - group.POST("", handleCreate) - group.GET("/:id", handleGet) - group.DELETE("/:id", handleRemove) - group.POST("/:id/exec", handleExec) - group.POST("/:id/heartbeat", handleHeartbeat) +func AttachManage(group *gin.RouterGroup, oauth types.OAuth) { + group.GET("", oauth.Guard, handleList) + group.POST("", oauth.Guard, handleCreate) + group.GET("/:id", oauth.Guard, handleGet) + group.DELETE("/:id", oauth.Guard, handleRemove) + group.POST("/:id/exec", oauth.Guard, handleExec) + group.POST("/:id/heartbeat", oauth.Guard, handleHeartbeat) } // resolveOwner returns TeamID if present, otherwise UserID. diff --git a/sandbox/v2/manager.go b/sandbox/v2/manager.go index 7a271257..3306d95d 100644 --- a/sandbox/v2/manager.go +++ b/sandbox/v2/manager.go @@ -372,6 +372,9 @@ func (m *Manager) buildTaiCreateOptions(opts CreateOptions, nodeID, sandboxID st "sandbox-node-id": nodeID, "sandbox-policy": string(opts.Policy), } + if opts.VNC { + labels["sandbox-vnc"] = "true" + } if opts.WorkspaceID != "" { labels["workspace-id"] = opts.WorkspaceID } @@ -455,6 +458,15 @@ func (m *Manager) recoverBoxes(ctx context.Context, nodeID string, res *tai.Conn if c.Name != "" { cid = c.Name } + hasVNC := c.Labels["sandbox-vnc"] == "true" + if !hasVNC { + for _, p := range c.Ports { + if p.ContainerPort == 5900 || p.ContainerPort == 6080 { + hasVNC = true + break + } + } + } box := &Box{ id: sandboxID, containerID: cid, @@ -465,6 +477,7 @@ func (m *Manager) recoverBoxes(ctx context.Context, nodeID string, res *tai.Conn createdAt: time.Now(), image: c.Image, workspaceID: c.Labels["workspace-id"], + vnc: hasVNC, workDir: "/workspace", manager: m, } diff --git a/tai/dial.go b/tai/dial.go index 5844f25d..0f560cd3 100644 --- a/tai/dial.go +++ b/tai/dial.go @@ -189,8 +189,8 @@ func buildResources(conn *grpc.ClientConn, cfg *dialConfig, env dialEnv) (*ConnR if res.Runtime != nil { res.Proxy = env.newProxy(cfg.ports) - res.VNC = env.newVNC(cfg.ports) } + res.VNC = env.newVNC(cfg.ports) return res, nil }