opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(sandbox): enhance lifecycle management and box status handling

Browse source 
- Updated BuildIdentifier to include assistant ID in session identifiers for better uniqueness.
- Implemented automatic starting of stopped boxes in resolveBox, improving recovery processes.
- Introduced status management for boxes, allowing for accurate tracking of their state (running, exited, stopped).
- Added idle timeout defaults based on lifecycle policy in BuildCreateOptions, enhancing configuration flexibility.
- Refactored tests to validate new box status behavior and lifecycle management improvements.

Made-with: Cursor
This commit is contained in:
Max 2026-03-15 11:51:09 +08:00
parent 6d5ae17d2f
commit 3797091d37
7 changed files with 219 additions and 120 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
agent/sandbox/v2/lifecycle.go
View file

@ -22,16 +22,9 @@ func BuildIdentifier(cfg *types.SandboxConfig, ownerID, chatID, assistantID, wor
return ""
}
// Custom identifier from metadata takes precedence.
if metadata != nil {
if cid, ok := metadata["computer_id"].(string); ok && cid != "" {
return fmt.Sprintf("%s-%s.%s", ownerID, cid, workspaceID)
}
}
switch cfg.Lifecycle {
case "session":
return fmt.Sprintf("%s-%s", ownerID, chatID)
return fmt.Sprintf("%s-%s-%s", ownerID, assistantID, chatID)
case "longrunning", "persistent":
return fmt.Sprintf("%s-%s.%s", ownerID, assistantID, workspaceID)
default:
@ -184,8 +177,17 @@ func resolveBox(
if identifier != "" {
box, err := manager.Get(context.Background(), identifier)
if err == nil && box != nil {
box.BindWorkplace(workspaceID)
return box, identifier, nil
if box.IsStopped() {
if startErr := manager.StartBox(context.Background(), identifier); startErr != nil {
log.Printf("[sandbox/v2] auto-start stopped box %s failed: %v, creating new", identifier, startErr)
} else {
box.BindWorkplace(workspaceID)
return box, identifier, nil
}
} else {
box.BindWorkplace(workspaceID)
return box, identifier, nil
}
}
}

8
agent/sandbox/v2/options.go
View file

@ -63,6 +63,14 @@ func BuildCreateOptions(cfg *types.SandboxConfig, identifier, ownerID, workspace
}
opts.IdleTimeout = d
}
if opts.IdleTimeout == 0 {
switch opts.Policy {
case infra.Session:
opts.IdleTimeout = infra.DefaultSessionIdleTimeout
case infra.LongRunning:
opts.IdleTimeout = infra.DefaultLongRunningIdleTimeout
}
}
if cfg.MaxLifetime != "" {
d, err := time.ParseDuration(cfg.MaxLifetime)
if err != nil {

34
sandbox/v2/box.go
View file

@ -22,6 +22,7 @@ type Box struct {
lastCall atomic.Int64
lastHeartbeat atomic.Int64
processCount atomic.Int32
status atomic.Value // string: "running", "exited", "stopped", "created", "unknown"
idleTimeoutD time.Duration
maxLifetimeD time.Duration
stopTimeoutD time.Duration
@ -210,14 +211,18 @@ func (b *Box) GetWorkDir() string {
func (b *Box) WorkspaceID() string { return b.workspaceID }
// Snapshot returns a local-only BoxInfo snapshot without any remote calls.
// Status is inferred from local state (not from the container runtime).
// Status is maintained by the sandbox watcher (see watcher.go).
func (b *Box) Snapshot() BoxInfo {
s, _ := b.status.Load().(string)
if s == "" {
s = "unknown"
}
return BoxInfo{
ID: b.id,
ContainerID: b.containerID,
NodeID: b.nodeID,
Owner: b.owner,
Status: "running",
Status: s,
Policy: b.policy,
Labels: b.labels,
Image: b.image,
@ -313,6 +318,12 @@ func (b *Box) lastActiveTime() time.Time {
return time.UnixMilli(ts)
}
// idleSince returns the timestamp of the last business call (Exec/Stream/VNC/etc).
// Unlike lastActiveTime, heartbeats do NOT reset this — only real user activity does.
func (b *Box) idleSince() time.Time {
return time.UnixMilli(b.lastCall.Load())
}
func (b *Box) idleTimeout() time.Duration {
return b.idleTimeoutD
}
@ -327,3 +338,22 @@ func (b *Box) stopTimeout() time.Duration {
}
return DefaultStopTimeout
}
// IsStopped reports whether the box's last known status indicates a non-running container.
func (b *Box) IsStopped() bool {
s, _ := b.status.Load().(string)
return s == "exited" || s == "stopped"
}
// inspectStatus queries the container runtime for the real container state.
func (b *Box) inspectStatus(ctx context.Context) string {
res, err := b.manager.getNode(b.nodeID)
if err != nil || res.Runtime == nil {
return "unknown"
}
info, err := res.Runtime.Inspect(ctx, b.containerID)
if err != nil {
return "unknown"
}
return info.Status
}

99
sandbox/v2/manager.go
View file

@ -19,20 +19,18 @@ import (
)
// Manager manages sandbox lifecycle. Node connections are delegated to tai/registry.
// Idle-timeout and lifecycle enforcement is handled by the sandbox watcher (watcher.go).
type Manager struct {
boxes sync.Map
mu sync.Mutex
cancel context.CancelFunc
boxes sync.Map
}
func newManager() *Manager {
return &Manager{}
}
// Start discovers existing containers from all registered nodes, rebuilds
// the boxes map, and starts the cleanup loop.
// If no "local" node is registered yet, it probes the local Docker environment
// and auto-registers one when available.
// Start discovers existing containers from all registered nodes and rebuilds
// the boxes map. If no "local" node is registered yet, it probes the local
// Docker environment and auto-registers one when available.
func (m *Manager) Start(ctx context.Context) error {
reg := registry.Global()
if reg == nil {
@ -49,9 +47,6 @@ func (m *Manager) Start(ctx context.Context) error {
m.recoverBoxes(ctx, snap.TaiID, res)
}
loopCtx, cancel := context.WithCancel(ctx)
m.cancel = cancel
go m.cleanupLoop(loopCtx)
return nil
}
@ -218,6 +213,7 @@ func (m *Manager) Create(ctx context.Context, opts CreateOptions) (*Box, error)
displayName: opts.DisplayName,
system: sys,
}
box.status.Store("running")
box.lastCall.Store(time.Now().UnixMilli())
m.boxes.Store(id, box)
@ -267,6 +263,31 @@ func (m *Manager) List(_ context.Context, opts ListOptions) ([]*Box, error) {
return result, nil
}
// StartBox starts a stopped sandbox and updates its lastCall timestamp.
func (m *Manager) StartBox(ctx context.Context, id string) error {
v, ok := m.boxes.Load(id)
if !ok {
return ErrNotFound
}
b := v.(*Box)
res, err := m.getNode(b.nodeID)
if err != nil {
return err
}
if res.Runtime == nil {
return fmt.Errorf("sandbox: node %q has no container runtime", b.nodeID)
}
if err := res.Runtime.Start(ctx, b.containerID); err != nil {
return fmt.Errorf("sandbox: start container %s: %w", b.containerID, err)
}
b.status.Store("running")
b.touch()
return nil
}
// Remove force-removes a sandbox (SIGKILL + delete).
func (m *Manager) Remove(ctx context.Context, id string) error {
v, ok := m.boxes.Load(id)
@ -284,58 +305,11 @@ func (m *Manager) Remove(ctx context.Context, id string) error {
return nil
}
// Cleanup removes idle/expired sandboxes.
func (m *Manager) Cleanup(ctx context.Context) error {
now := time.Now()
m.boxes.Range(func(key, value any) bool {
b := value.(*Box)
idle := now.Sub(b.lastActiveTime())
switch b.policy {
case OneShot:
// handled after Exec
case Session:
if timeout := b.idleTimeout(); timeout > 0 && idle > timeout {
m.Remove(ctx, b.id)
}
case LongRunning:
if timeout := b.idleTimeout(); timeout > 0 && idle > timeout {
if res, err := m.getNode(b.nodeID); err == nil && res.Runtime != nil {
res.Runtime.Stop(ctx, b.containerID, b.stopTimeout())
}
}
if lifetime := b.maxLifetime(); lifetime > 0 && now.Sub(b.createdAt) > lifetime {
m.Remove(ctx, b.id)
}
case Persistent:
// never auto-cleaned
}
return true
})
return nil
}
// Close stops the cleanup loop. Node connections are managed by the registry.
// Close is a no-op; lifecycle management is handled by the sandbox watcher.
func (m *Manager) Close() error {
if m.cancel != nil {
m.cancel()
}
return nil
}
func (m *Manager) cleanupLoop(ctx context.Context) {
ticker := time.NewTicker(1 * time.Minute)
defer ticker.Stop()
for {
select {
case <-ticker.C:
m.Cleanup(ctx)
case <-ctx.Done():
return
}
}
}
func (m *Manager) getNode(name string) (*tai.ConnResources, error) {
res, ok := tai.GetResources(name)
if !ok {
@ -488,12 +462,13 @@ func (m *Manager) recoverBoxes(ctx context.Context, nodeID string, res *tai.Conn
if sys.OS == "" {
sys = inferSystemInfo(ctx, res, c.Image)
}
policy := LifecyclePolicy(c.Labels["sandbox-policy"])
box := &Box{
id: sandboxID,
containerID: cid,
nodeID: c.Labels["sandbox-node-id"],
owner: c.Labels["sandbox-owner"],
policy: LifecyclePolicy(c.Labels["sandbox-policy"]),
policy: policy,
labels: c.Labels,
createdAt: time.Now(),
image: c.Image,
@ -504,6 +479,12 @@ func (m *Manager) recoverBoxes(ctx context.Context, nodeID string, res *tai.Conn
system: sys,
manager: m,
}
switch policy {
case Session:
box.idleTimeoutD = DefaultSessionIdleTimeout
case LongRunning:
box.idleTimeoutD = DefaultLongRunningIdleTimeout
}
box.lastCall.Store(time.Now().UnixMilli())
m.boxes.Store(sandboxID, box)
}

83
sandbox/v2/manager_lifecycle_test.go
View file

@ -46,42 +46,6 @@ func TestHeartbeatUnknownBox(t *testing.T) {
}
}
func TestIdleCleanup(t *testing.T) {
skipIfNoDocker(t)
for _, pc := range testNodes() {
pc := pc
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForNode(t, &pc)
ensureTestImage(t, m, pc.TaiID)
ctx := context.Background()
box, err := m.Create(ctx, sandbox.CreateOptions{
Image: testImage(),
Owner: "test-user",
NodeID: pc.TaiID,
Policy: sandbox.Session,
IdleTimeout: 1 * time.Second,
})
if err != nil {
t.Fatalf("Create: %v", err)
}
boxID := box.ID()
time.Sleep(2 * time.Second)
if err := m.Cleanup(ctx); err != nil {
t.Fatalf("Cleanup: %v", err)
}
_, err = m.Get(ctx, boxID)
if err != sandbox.ErrNotFound {
t.Errorf("after idle cleanup, Get err = %v, want ErrNotFound", err)
}
})
}
}
func TestStartRecovery(t *testing.T) {
skipIfNoDocker(t)
@ -114,27 +78,50 @@ func TestStartRecovery(t *testing.T) {
}
}
func TestPersistentNotCleaned(t *testing.T) {
func TestStartBox(t *testing.T) {
skipIfNoDocker(t)
for _, pc := range testNodes() {
pc := pc
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForNode(t, &pc)
box := createTestBox(t, m, pc, func(co *sandbox.CreateOptions) {
co.Policy = sandbox.Persistent
co.IdleTimeout = 1 * time.Second
})
time.Sleep(2 * time.Second)
box := createTestBox(t, m, pc)
boxID := box.ID()
ctx := context.Background()
m.Cleanup(ctx)
_, err := m.Get(ctx, box.ID())
if err := box.Stop(ctx); err != nil {
t.Fatalf("Stop: %v", err)
}
time.Sleep(500 * time.Millisecond)
if err := m.StartBox(ctx, boxID); err != nil {
t.Fatalf("StartBox: %v", err)
}
info, err := box.Info(ctx)
if err != nil {
t.Errorf("persistent box should not be cleaned: %v", err)
t.Fatalf("Info after StartBox: %v", err)
}
if info.Status != "running" {
t.Errorf("status = %q after StartBox, want running", info.Status)
}
})
}
}
func TestSnapshotReadsStatus(t *testing.T) {
skipIfNoDocker(t)
for _, pc := range testNodes() {
pc := pc
t.Run(pc.Name, func(t *testing.T) {
m := setupManagerForNode(t, &pc)
box := createTestBox(t, m, pc)
snap := box.Snapshot()
if snap.Status != "running" {
t.Errorf("initial snapshot status = %q, want running", snap.Status)
}
})
}

6
sandbox/v2/types.go
View file

@ -71,7 +71,11 @@ const (
Persistent LifecyclePolicy = "persistent"
)
const DefaultStopTimeout = 2 * time.Second
const (
DefaultStopTimeout = 2 * time.Second
DefaultSessionIdleTimeout = 30 * time.Minute
DefaultLongRunningIdleTimeout = 2 * time.Hour
)
// ---------------------------------------------------------------------------
// Create / List options

87
sandbox/v2/watcher.go Normal file
View file

@ -0,0 +1,87 @@
package sandbox
import (
"context"
"fmt"
"time"
"github.com/yaoapp/yao/monitor"
)
func init() {
monitor.Register(&sandboxWatcher{})
}
type sandboxWatcher struct{}
func (w *sandboxWatcher) Name() string { return "sandbox" }
func (w *sandboxWatcher) Interval() time.Duration { return 30 * time.Second }
func (w *sandboxWatcher) Check(ctx context.Context) []monitor.Alert {
if mgr == nil {
return nil
}
var alerts []monitor.Alert
mgr.boxes.Range(func(_, v any) bool {
b := v.(*Box)
status := b.inspectStatus(ctx)
old, _ := b.status.Swap(status).(string)
if old != "" && old != status {
alerts = append(alerts, monitor.Alert{
Level: monitor.Info,
Target: "box:" + b.id,
Message: fmt.Sprintf("status %s → %s", old, status),
})
}
if status != "running" {
return true
}
idle := time.Since(b.idleSince())
timeout := b.idleTimeout()
if timeout <= 0 || idle <= timeout {
return true
}
switch b.policy {
case Session:
alerts = append(alerts, monitor.Alert{
Level: monitor.Warn,
Target: "box:" + b.id,
Message: fmt.Sprintf("session idle expired (idle=%s, timeout=%s), removing", idle.Round(time.Second), timeout),
Action: func(ctx context.Context) {
mgr.Remove(ctx, b.id)
},
})
case LongRunning:
alerts = append(alerts, monitor.Alert{
Level: monitor.Warn,
Target: "box:" + b.id,
Message: fmt.Sprintf("longrunning idle expired (idle=%s, timeout=%s), stopping", idle.Round(time.Second), timeout),
Action: func(ctx context.Context) {
b.Stop(ctx)
},
})
}
if b.policy == LongRunning {
if lifetime := b.maxLifetime(); lifetime > 0 && time.Since(b.createdAt) > lifetime {
alerts = append(alerts, monitor.Alert{
Level: monitor.Warn,
Target: "box:" + b.id,
Message: fmt.Sprintf("lifetime expired (%s), removing", lifetime),
Action: func(ctx context.Context) {
mgr.Remove(ctx, b.id)
},
})
}
}
return true
})
return alerts
}