feat(sandbox): integrate sandbox initialization and heartbeat management

- Added sandbox initialization in the load process to ensure proper setup of the sandbox manager.
- Implemented gRPC heartbeat handling to track container liveness, enhancing the monitoring capabilities of the sandbox.
- Updated related documentation to reflect the new integration and functionality.

Made-with: Cursor

cmd/start.go

@@ -26,12 +26,14 @@ import (
 	"github.com/yaoapp/yao/engine"
 	yaogrpc "github.com/yaoapp/yao/grpc"
 	_ "github.com/yaoapp/yao/grpc/auth"
+	sandboxhandler "github.com/yaoapp/yao/grpc/sandbox"
 	"github.com/yaoapp/yao/openapi"
+	sandbox "github.com/yaoapp/yao/sandbox/v2"
 	ischedule "github.com/yaoapp/yao/schedule"
 	"github.com/yaoapp/yao/service"
 	"github.com/yaoapp/yao/setup"
 	"github.com/yaoapp/yao/share"
-	tairegistry "github.com/yaoapp/yao/tai/registry"
+
 	itask "github.com/yaoapp/yao/task"
 )
 
@@ -176,10 +178,6 @@ var startCmd = &cobra.Command{
 		ischedule.Start()
 		defer ischedule.Stop()
 
-		// Initialize the global Tai registry for tunnel and direct connections
-		// (must happen before HTTP/gRPC start so handlers can access it)
-		tairegistry.Init(nil)
-
 		// Pre-flight: detect port conflicts before attempting to start servers.
 		if occupied, proc := portOccupied(config.Conf.Host, config.Conf.Port); occupied {
 			fmt.Println(color.RedString(L("Fatal: HTTP port %d is already in use%s"), config.Conf.Port, proc))
@@ -194,6 +192,12 @@ var startCmd = &cobra.Command{
 			}
 		}
 
+		// Wire gRPC heartbeat → sandbox Manager so container liveness is tracked.
+		yaogrpc.SetSandboxOnBeat(func(data *sandboxhandler.HeartbeatData) string {
+			sandbox.M().Heartbeat(data.SandboxID, true, int(data.RunningProcs))
+			return "ok"
+		})
+
 		// Start all servers (gRPC + HTTP) as a single unit.
 		// Start() blocks until HTTP port is bound (READY) or returns error.
 		svc, err := service.Start(config.Conf, service.ServerHooks{

engine/load.go

@@ -40,12 +40,14 @@ import (
 	"github.com/yaoapp/yao/plugin"
 	"github.com/yaoapp/yao/query"
 	"github.com/yaoapp/yao/runtime"
+	sandbox "github.com/yaoapp/yao/sandbox/v2"
 	"github.com/yaoapp/yao/schedule"
 	"github.com/yaoapp/yao/script"
 	"github.com/yaoapp/yao/share"
 	"github.com/yaoapp/yao/socket"
 	"github.com/yaoapp/yao/store"
 	sui "github.com/yaoapp/yao/sui/api"
+	tairegistry "github.com/yaoapp/yao/tai/registry"
 	"github.com/yaoapp/yao/task"
 	"github.com/yaoapp/yao/websocket"
 	"github.com/yaoapp/yao/widget"
@@ -130,6 +132,22 @@ func Load(cfg config.Config, options LoadOption, progressCallback ...func(string
 		warnings = append(warnings, Warning{Widget: "DB", Error: err})
 	}
 
+	// Initialize the Tai node registry (idempotent, safe to call early).
+	loadStep("Registry", func() error {
+		tairegistry.InitWithWriter(config.LogOutput, cfg.LogMode)
+		return nil
+	}, callback)
+
+	// Initialize the Sandbox manager and start it (auto-registers local Docker
+	// node if available, recovers existing containers, starts cleanup loop).
+	err = loadStep("Sandbox", func() error {
+		sandbox.Init()
+		return sandbox.M().Start(context.Background())
+	}, callback)
+	if err != nil {
+		warnings = append(warnings, Warning{Widget: "Sandbox", Error: err})
+	}
+
 	// Load Certs
 	err = loadStep("Cert", func() error {
 		return cert.Load(cfg)

sandbox/v2/IMPL.md

@@ -166,12 +166,12 @@ Token provisioning is the **caller's responsibility** via `CreateOptions.Env`:
 - Passes it in `CreateOptions.Env["YAO_TOKEN"]` / `Env["YAO_REFRESH_TOKEN"]`
 - `opts.Env` takes priority over `BuildGRPCEnv` output (caller can override anything)
 
-### Remaining (Startup) — PENDING
+### Remaining (Startup)
 
-| Task | Package | Detail |
-|------|---------|--------|
-| `cmd/start.go` integration | `yao` | Call `sandbox.Init()` + `sandbox.M().Start(ctx)` in startup (no config needed — node discovery via tai/registry) |
-| Heartbeat bridge | `yao/grpc` | Wire gRPC Heartbeat handler → `sandbox.M().Heartbeat()` |
+| Task | Package | Status | Detail |
+|------|---------|--------|--------|
+| `engine/load.go` integration | `yao` | **DONE** | `sandbox.Init()` + `sandbox.M().Start(ctx)` added as a `loadStep("Sandbox", ...)` right after Registry init |
+| Heartbeat bridge | `yao/cmd` | **DONE** | `cmd/start.go` calls `yaogrpc.SetSandboxOnBeat(...)` before `service.Start`, forwarding gRPC heartbeats to `sandbox.M().Heartbeat()` |
 
 ---
 

sandbox/v2/manager.go

@@ -3,9 +3,11 @@ package sandbox
 import (
 	"context"
 	"fmt"
+	"path/filepath"
 	"sync"
 	"time"
 
+	"github.com/yaoapp/yao/config"
 	"github.com/yaoapp/yao/tai"
 	"github.com/yaoapp/yao/tai/registry"
 	taisandbox "github.com/yaoapp/yao/tai/sandbox"
@@ -25,12 +27,16 @@ func newManager() *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.
 func (m *Manager) Start(ctx context.Context) error {
 	reg := registry.Global()
 	if reg == nil {
 		return nil
 	}
 
+	m.ensureLocalNode(reg)
+
 	for _, snap := range reg.List() {
 		client, err := m.getNode(snap.TaiID)
 		if err != nil {
@@ -45,6 +51,15 @@ func (m *Manager) Start(ctx context.Context) error {
 	return nil
 }
 
+// ensureLocalNode delegates to tai.RegisterLocal() which probes the local
+// Docker environment and registers a "local" node in the registry if available.
+// The workspace data directory is derived from config.Conf.DataRoot so that
+// workspace files persist across restarts.
+func (m *Manager) ensureLocalNode(_ *registry.Registry) {
+	dataDir := filepath.Join(config.Conf.DataRoot, "workspaces")
+	tai.RegisterLocal(tai.WithDataDir(dataDir))
+}
+
 // Nodes returns the list of registered Tai nodes from the registry.
 func (m *Manager) Nodes() []registry.NodeSnapshot {
 	reg := registry.Global()

tai/registry/registry.go

@@ -7,6 +7,8 @@ import (
 	"io"
 	"log/slog"
 	"net"
+	"os"
+	"strings"
 	"sync"
 	"time"
 
@@ -145,6 +147,23 @@ func Init(logger *slog.Logger) {
 	})
 }
 
+// InitWithWriter initializes the global registry using the given io.Writer
+// and log format ("JSON" or "TEXT"). If w is nil it falls back to stderr.
+// This is the preferred way to integrate with the application log system.
+func InitWithWriter(w io.Writer, logMode string) {
+	if w == nil {
+		w = os.Stderr
+	}
+	opts := &slog.HandlerOptions{Level: slog.LevelInfo}
+	var handler slog.Handler
+	if strings.EqualFold(logMode, "JSON") {
+		handler = slog.NewJSONHandler(w, opts)
+	} else {
+		handler = slog.NewTextHandler(w, opts)
+	}
+	Init(slog.New(handler))
+}
+
 // Global returns the global registry instance.
 func Global() *Registry {
 	return global

tai/tai.go

@@ -579,6 +579,27 @@ func (c *Client) discoverServerInfo(conn *grpc.ClientConn, cfg *config) (map[str
 	return caps, nil
 }
 
+// RegisterLocal probes the local Docker environment and, if reachable,
+// creates a Client and registers it as the "local" node in the registry.
+// Returns true if a local node was successfully registered.
+// Silently returns false if Docker is not available — this is not an error.
+func RegisterLocal(opts ...Option) bool {
+	reg := registry.Global()
+	if reg == nil {
+		return false
+	}
+	if _, ok := reg.Get("local"); ok {
+		return true
+	}
+
+	c, err := New("local", opts...)
+	if err != nil {
+		return false
+	}
+	_ = c // registered by initLocal → reg.Register + reg.SetClient
+	return true
+}
+
 // GetClient returns a registered *Client by taiID from the global registry.
 func GetClient(taiID string) (*Client, bool) {
 	reg := registry.Global()

tai/tai_test.go

@@ -5,6 +5,8 @@ import (
 	"os"
 	"strconv"
 	"testing"
+
+	"github.com/yaoapp/yao/tai/registry"
 )
 
 func taiTestHost() string {
@@ -342,3 +344,77 @@ func TestDiscoverPortsWithUserOverride(t *testing.T) {
 	t.Logf("ports: GRPC=%d HTTP=%d(user) VNC=%d Docker=%d",
 		c.ports.GRPC, c.ports.HTTP, c.ports.VNC, c.ports.Docker)
 }
+
+func TestRegisterLocal(t *testing.T) {
+	registry.Init(nil)
+	reg := registry.Global()
+
+	dir := t.TempDir()
+	ok := RegisterLocal(WithDataDir(dir))
+	if !ok {
+		t.Skip("Docker not available, skipping RegisterLocal test")
+	}
+
+	snap, found := reg.Get("local")
+	if !found {
+		t.Fatal("expected 'local' node in registry after RegisterLocal")
+	}
+	if snap.Mode != "local" {
+		t.Errorf("mode = %q, want 'local'", snap.Mode)
+	}
+	if snap.Status != "online" {
+		t.Errorf("status = %q, want 'online'", snap.Status)
+	}
+
+	c, got := GetClient("local")
+	if !got {
+		t.Fatal("GetClient('local') returned false after RegisterLocal")
+	}
+	if c.DataDir() != dir {
+		t.Errorf("DataDir = %q, want %q", c.DataDir(), dir)
+	}
+	if c.Sandbox() == nil {
+		t.Error("local client Sandbox should not be nil")
+	}
+
+	// Idempotent: second call should return true without error
+	ok2 := RegisterLocal(WithDataDir(dir))
+	if !ok2 {
+		t.Error("second RegisterLocal should return true (idempotent)")
+	}
+
+	c.Close()
+}
+
+func TestRegisterLocal_NoRegistry(t *testing.T) {
+	// RegisterLocal without a registry should return false, not panic
+	origReg := registry.Global()
+	defer func() {
+		if origReg != nil {
+			registry.Init(nil)
+		}
+	}()
+
+	// registry.Global() returns the singleton; we can't un-init it,
+	// but we can verify RegisterLocal returns true (registry exists from
+	// other tests) or false gracefully.
+	ok := RegisterLocal()
+	// Just verify it doesn't panic; result depends on Docker availability
+	_ = ok
+}
+
+func TestRegisterLocal_NoDocker(t *testing.T) {
+	registry.Init(nil)
+
+	// Use an unreachable Docker socket to ensure failure
+	ok := RegisterLocal(WithDataDir(t.TempDir()))
+	if !ok {
+		// Expected when Docker is not available — just ensure no panic
+		return
+	}
+	// If Docker happens to be available, that's also fine
+	c, _ := GetClient("local")
+	if c != nil {
+		c.Close()
+	}
+}