diff --git a/sandbox/v2/grpc.go b/sandbox/v2/grpc.go index e3d85399..9dc949ef 100644 --- a/sandbox/v2/grpc.go +++ b/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 } diff --git a/sandbox/v2/grpc_test.go b/sandbox/v2/grpc_test.go index 800a2cae..4da6f0f0 100644 --- a/sandbox/v2/grpc_test.go +++ b/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) } } diff --git a/sandbox/v2/manager.go b/sandbox/v2/manager.go index 4f0fe5aa..ff26c2a6 100644 --- a/sandbox/v2/manager.go +++ b/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 } diff --git a/tai/runtime/docker_core.go b/tai/runtime/docker_core.go index 26efe21d..4394f7c7 100644 --- a/tai/runtime/docker_core.go +++ b/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 {