From 3797091d374c2a9959db9c7d25a3d201fda90130 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 15 Mar 2026 11:51:09 +0800 Subject: [PATCH] feat(sandbox): enhance lifecycle management and box status handling - 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 --- agent/sandbox/v2/lifecycle.go | 22 ++++--- agent/sandbox/v2/options.go | 8 +++ sandbox/v2/box.go | 34 +++++++++- sandbox/v2/manager.go | 99 +++++++++++----------------- sandbox/v2/manager_lifecycle_test.go | 83 ++++++++++------------- sandbox/v2/types.go | 6 +- sandbox/v2/watcher.go | 87 ++++++++++++++++++++++++ 7 files changed, 219 insertions(+), 120 deletions(-) create mode 100644 sandbox/v2/watcher.go diff --git a/agent/sandbox/v2/lifecycle.go b/agent/sandbox/v2/lifecycle.go index bad0b57c..3bbf0c8d 100644 --- a/agent/sandbox/v2/lifecycle.go +++ b/agent/sandbox/v2/lifecycle.go @@ -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 + } } } diff --git a/agent/sandbox/v2/options.go b/agent/sandbox/v2/options.go index 2ac6ce92..4be74d31 100644 --- a/agent/sandbox/v2/options.go +++ b/agent/sandbox/v2/options.go @@ -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 { diff --git a/sandbox/v2/box.go b/sandbox/v2/box.go index d8412014..9e3a58c5 100644 --- a/sandbox/v2/box.go +++ b/sandbox/v2/box.go @@ -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 +} diff --git a/sandbox/v2/manager.go b/sandbox/v2/manager.go index ff26c2a6..b25f6976 100644 --- a/sandbox/v2/manager.go +++ b/sandbox/v2/manager.go @@ -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) } diff --git a/sandbox/v2/manager_lifecycle_test.go b/sandbox/v2/manager_lifecycle_test.go index 67412792..f4afe12a 100644 --- a/sandbox/v2/manager_lifecycle_test.go +++ b/sandbox/v2/manager_lifecycle_test.go @@ -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) } }) } diff --git a/sandbox/v2/types.go b/sandbox/v2/types.go index 828ef799..109c25b4 100644 --- a/sandbox/v2/types.go +++ b/sandbox/v2/types.go @@ -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 diff --git a/sandbox/v2/watcher.go b/sandbox/v2/watcher.go new file mode 100644 index 00000000..e0c8315f --- /dev/null +++ b/sandbox/v2/watcher.go @@ -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 +}