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

@@ -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{

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 != "" {

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
+}

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
 }

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
 }

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,

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)

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.

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,
 		}

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
 }