feat(grpc): refactor BuildGRPCEnv to use port from Tai node and update tests

- Modified BuildGRPCEnv to accept the Tai node's gRPC port directly, enhancing flexibility for different modes.
- Updated test cases to reflect changes in gRPC address construction, ensuring accurate environment variable settings for local, direct, and tunnel modes.
- Added a new test for handling unknown modes, improving test coverage and robustness.
- Adjusted Docker configuration to ensure proper host resolution for gRPC communication.

Made-with: Cursor

sandbox/v2/grpc.go

@@ -2,53 +2,47 @@ package sandbox
 
 import (
 	"fmt"
-	"net/url"
-	"strconv"
 
 	"github.com/yaoapp/yao/config"
 )
 
-// BuildGRPCEnv builds the gRPC environment variables for a sandbox container
-// based on the Tai node's mode and address from the registry.
-//
-// mode is the TaiNode.Mode ("local", "direct", "tunnel").
-// addr is the TaiNode.Addr (e.g. "tai://host:port" for direct mode).
-// sandboxID is the container's sandbox identifier.
-//
-// The Yao gRPC port is read from config.Conf.GRPC.Port.
-func BuildGRPCEnv(mode, addr, sandboxID string) map[string]string {
-	grpcPort := config.Conf.GRPC.Port
-	if grpcPort == 0 {
-		grpcPort = 9099
-	}
-	portStr := strconv.Itoa(grpcPort)
+const taiHost = "host.tai.internal"
 
+// BuildGRPCEnv builds the gRPC environment variables for a sandbox container.
+//
+// All containers reach the host via "host.tai.internal" (injected by Tai at
+// container creation). The port depends on the mode:
+//
+//   - local:          Yao gRPC port (Tai and Yao on the same machine)
+//   - tunnel/direct:  Tai gRPC port (Tai Gateway forwards to Yao)
+//
+// taiGRPCPort is the Tai node's gRPC port from registration (Ports.GRPC).
+func BuildGRPCEnv(mode string, taiGRPCPort int, sandboxID string) map[string]string {
 	env := map[string]string{
 		"YAO_SANDBOX_ID": sandboxID,
 	}
 
 	switch mode {
 	case "local":
-		env["YAO_GRPC_ADDR"] = fmt.Sprintf("host.docker.internal:%s", portStr)
-
-	case "tunnel":
-		env["YAO_GRPC_ADDR"] = fmt.Sprintf("127.0.0.1:%s", portStr)
-
-	case "direct":
-		u, err := url.Parse(addr)
-		if err != nil || u.Hostname() == "" {
-			env["YAO_GRPC_ADDR"] = fmt.Sprintf("host.docker.internal:%s", portStr)
-			return env
+		port := config.Conf.GRPC.Port
+		if port == 0 {
+			port = 9099
 		}
-		taiHost := u.Hostname()
-		taiPort := u.Port()
-		if taiPort == "" {
-			taiPort = "19100"
+		env["YAO_GRPC_ADDR"] = fmt.Sprintf("%s:%d", taiHost, port)
+
+	case "tunnel", "direct":
+		port := taiGRPCPort
+		if port == 0 {
+			port = 19100
 		}
-		env["YAO_GRPC_ADDR"] = fmt.Sprintf("%s:%s", taiHost, taiPort)
+		env["YAO_GRPC_ADDR"] = fmt.Sprintf("%s:%d", taiHost, port)
 
 	default:
-		env["YAO_GRPC_ADDR"] = fmt.Sprintf("host.docker.internal:%s", portStr)
+		port := config.Conf.GRPC.Port
+		if port == 0 {
+			port = 9099
+		}
+		env["YAO_GRPC_ADDR"] = fmt.Sprintf("%s:%d", taiHost, port)
 	}
 	return env
 }

sandbox/v2/grpc_test.go

@@ -9,7 +9,7 @@ import (
 
 func TestBuildGRPCEnvLocal(t *testing.T) {
 	config.Conf.GRPC.Port = 9099
-	env := sandbox.BuildGRPCEnv("local", "", "sb-001")
+	env := sandbox.BuildGRPCEnv("local", 19100, "sb-001")
 
 	if env["YAO_SANDBOX_ID"] != "sb-001" {
 		t.Errorf("YAO_SANDBOX_ID = %q", env["YAO_SANDBOX_ID"])
@@ -17,25 +17,48 @@ func TestBuildGRPCEnvLocal(t *testing.T) {
 	if _, ok := env["YAO_TOKEN"]; ok {
 		t.Error("YAO_TOKEN should not be set by BuildGRPCEnv")
 	}
-	if env["YAO_GRPC_ADDR"] != "host.docker.internal:9099" {
-		t.Errorf("YAO_GRPC_ADDR = %q, want host.docker.internal:9099", env["YAO_GRPC_ADDR"])
+	want := "host.tai.internal:9099"
+	if env["YAO_GRPC_ADDR"] != want {
+		t.Errorf("YAO_GRPC_ADDR = %q, want %q", env["YAO_GRPC_ADDR"], want)
 	}
 }
 
 func TestBuildGRPCEnvDirect(t *testing.T) {
 	config.Conf.GRPC.Port = 9099
-	env := sandbox.BuildGRPCEnv("direct", "tai://gpu-server", "sb-002")
+	env := sandbox.BuildGRPCEnv("direct", 19100, "sb-002")
 
-	if env["YAO_GRPC_ADDR"] != "gpu-server:19100" {
-		t.Errorf("YAO_GRPC_ADDR = %q, want gpu-server:19100", env["YAO_GRPC_ADDR"])
+	want := "host.tai.internal:19100"
+	if env["YAO_GRPC_ADDR"] != want {
+		t.Errorf("YAO_GRPC_ADDR = %q, want %q", env["YAO_GRPC_ADDR"], want)
+	}
+}
+
+func TestBuildGRPCEnvDirectDefaultPort(t *testing.T) {
+	config.Conf.GRPC.Port = 9099
+	env := sandbox.BuildGRPCEnv("direct", 0, "sb-002")
+
+	want := "host.tai.internal:19100"
+	if env["YAO_GRPC_ADDR"] != want {
+		t.Errorf("YAO_GRPC_ADDR = %q, want %q (default tai port)", env["YAO_GRPC_ADDR"], want)
 	}
 }
 
 func TestBuildGRPCEnvTunnel(t *testing.T) {
 	config.Conf.GRPC.Port = 9099
-	env := sandbox.BuildGRPCEnv("tunnel", "tunnel://relay.example.com", "sb-003")
+	env := sandbox.BuildGRPCEnv("tunnel", 19200, "sb-003")
 
-	if env["YAO_GRPC_ADDR"] != "127.0.0.1:9099" {
-		t.Errorf("YAO_GRPC_ADDR = %q, want 127.0.0.1:9099", env["YAO_GRPC_ADDR"])
+	want := "host.tai.internal:19200"
+	if env["YAO_GRPC_ADDR"] != want {
+		t.Errorf("YAO_GRPC_ADDR = %q, want %q", env["YAO_GRPC_ADDR"], want)
+	}
+}
+
+func TestBuildGRPCEnvUnknownMode(t *testing.T) {
+	config.Conf.GRPC.Port = 8888
+	env := sandbox.BuildGRPCEnv("unknown", 19100, "sb-004")
+
+	want := "host.tai.internal:8888"
+	if env["YAO_GRPC_ADDR"] != want {
+		t.Errorf("YAO_GRPC_ADDR = %q, want %q (fallback to yao port)", env["YAO_GRPC_ADDR"], want)
 	}
 }

sandbox/v2/manager.go

@@ -350,7 +350,7 @@ func (m *Manager) buildTaiCreateOptions(opts CreateOptions, nodeID, sandboxID st
 	reg := registry.Global()
 	if reg != nil {
 		if snap, ok := reg.Get(nodeID); ok {
-			grpcEnv := BuildGRPCEnv(snap.Mode, snap.Addr, sandboxID)
+			grpcEnv := BuildGRPCEnv(snap.Mode, snap.Ports.GRPC, sandboxID)
 			for k, v := range grpcEnv {
 				env[k] = v
 			}

tai/runtime/docker_core.go

@@ -31,7 +31,8 @@ func (d *dockerCore) create(ctx context.Context, opts CreateOptions, addVNCPorts
 	}
 
 	hostCfg := &container.HostConfig{
-		Binds: opts.Binds,
+		Binds:      opts.Binds,
+		ExtraHosts: []string{"host.tai.internal:host-gateway"},
 	}
 
 	if opts.Memory > 0 {