refactor(runtime): use the SDK host lifecycle

go.mod

@@ -4,7 +4,7 @@ go 1.26.0
 
 require (
 	github.com/NeKiro-project/NeKiro v0.0.0-20260810043416-3e815b89cb37
-	github.com/NeKiro-project/nekiro-sdk-go v0.0.0-20260804145402-39c37c8929b5
+	github.com/NeKiro-project/nekiro-sdk-go v0.0.0-20260811153316-3f4c32d6e895
 	github.com/a2aproject/a2a-go v0.3.15
 	github.com/golang-jwt/jwt/v5 v5.3.1
 	trpc.group/trpc-go/trpc-agent-go v1.10.0

go.sum

@@ -2,8 +2,8 @@ github.com/Masterminds/semver/v3 v3.5.0 h1:kQceYJfbupGfZOKZQg0kou0DgAKhzDg2NZPAw
 github.com/Masterminds/semver/v3 v3.5.0/go.mod h1:4V+yj/TJE1HU9XfppCwVMZq3I84lprf4nC11bSS5beM=
 github.com/NeKiro-project/NeKiro v0.0.0-20260810043416-3e815b89cb37 h1:ai0eN+G2k6rwtF9h1dDhDKzloueH6AyuJp2UOIoyQec=
 github.com/NeKiro-project/NeKiro v0.0.0-20260810043416-3e815b89cb37/go.mod h1:JCIEeiLu52WC/Q5QlcAKmWKRtW7CNLkZ3lV3BAn92Oo=
-github.com/NeKiro-project/nekiro-sdk-go v0.0.0-20260804145402-39c37c8929b5 h1:wx8nuBluyNb7MhMOURtX/9/Q+FBNcniJqoCMpxBnrwE=
-github.com/NeKiro-project/nekiro-sdk-go v0.0.0-20260804145402-39c37c8929b5/go.mod h1:lxAQLsSVXE3lmWoz+SMDKsDOc1b0Vj7okAGNkE7KFRI=
+github.com/NeKiro-project/nekiro-sdk-go v0.0.0-20260811153316-3f4c32d6e895 h1:In02gBJq2ZHUdfjcPYrz2BePtJOfdscNBxQOiL8mu14=
+github.com/NeKiro-project/nekiro-sdk-go v0.0.0-20260811153316-3f4c32d6e895/go.mod h1:lxAQLsSVXE3lmWoz+SMDKsDOc1b0Vj7okAGNkE7KFRI=
 github.com/a2aproject/a2a-go v0.3.15 h1:h5YpCiPq3jxQ5rIns7oDjPag3ivP8u817AzdA4F+NiI=
 github.com/a2aproject/a2a-go v0.3.15/go.mod h1:I7Cm+a1oL+UT6zMoP+roaRE5vdfUa1iQGVN8aSOuZ0I=
 github.com/bmatcuk/doublestar/v4 v4.9.1 h1:X8jg9rRZmJd4yRy7ZeNDRnM+T3ZfHv15JiBJ/avrEXE=

runtime-a/cmd/runtime-a/main.go

@@ -2,18 +2,16 @@ package main
 
 import (
 	"context"
-	"errors"
-	"fmt"
 	"log"
 	"net/http"
 	"os"
-	"os/signal"
 	"syscall"
 	"time"
 
 	"github.com/NeKiro-project/NeKiro-Samples/internal/challengeproof"
 	"github.com/NeKiro-project/NeKiro-Samples/internal/nacosregistration"
 	runtimea "github.com/NeKiro-project/NeKiro-Samples/runtime-a"
+	agenthost "github.com/NeKiro-project/nekiro-sdk-go/agent/host"
 )
 
 func main() {
@@ -23,86 +21,55 @@ func main() {
 }
 
 func run() error {
-	config, err := runtimea.LoadConfig(os.LookupEnv)
+	return runWithLookup(os.LookupEnv)
+}
+
+func runWithLookup(lookup func(string) (string, bool)) error {
+	config, err := runtimea.LoadConfig(lookup)
 	if err != nil {
-		return err
+		return agenthost.Wrap(agenthost.StageConfig, "load Runtime A configuration", err)
 	}
-	registrationConfig, err := nacosregistration.Load(os.LookupEnv, "RUNTIME_A", config.AgentID, config.InstanceID)
+	registrationConfig, err := nacosregistration.Load(lookup, "RUNTIME_A", config.AgentID, config.InstanceID)
 	if err != nil {
-		return err
+		return agenthost.Wrap(agenthost.StageConfig, "load Runtime A registration configuration", err)
 	}
-	var registration *nacosregistration.Registration
+	var registration agenthost.Registration
 	var readiness runtimea.Readiness = ready(true)
 	if registrationConfig.Mode == nacosregistration.ModeNacos {
 		registrationClient, clientErr := nacosregistration.NewHTTPClient(registrationConfig)
 		if clientErr != nil {
-			return fmt.Errorf("runtime-a Nacos registration transport: %w", clientErr)
+			return agenthost.Wrap(agenthost.StageRegistration, "create Runtime A Nacos transport", clientErr)
 		}
-		registration, err = nacosregistration.New(registrationConfig, registrationClient)
+		runtimeRegistration, err := nacosregistration.New(registrationConfig, registrationClient)
 		if err != nil {
-			return fmt.Errorf("runtime-a Nacos registration config: %w", err)
+			return agenthost.Wrap(agenthost.StageRegistration, "create Runtime A registration", err)
 		}
-		readiness = registration
+		registration = runtimeRegistration
+		readiness = runtimeRegistration
 	}
 	handler, err := runtimea.NewHandler(config, http.DefaultClient)
 	if err != nil {
-		return fmt.Errorf("runtime-a initialize: %w", err)
+		return agenthost.Wrap(agenthost.StageHandler, "create Runtime A handler", err)
 	}
-	application, err := challengeproof.NewHandler(runtimea.NewHTTPHandlerWithReadiness(handler, readiness), os.LookupEnv)
+	application, err := challengeproof.NewHandler(runtimea.NewHTTPHandlerWithReadiness(handler, readiness), lookup)
 	if err != nil {
-		return fmt.Errorf("runtime-a challenge proof: %w", err)
+		return agenthost.Wrap(agenthost.StageHandler, "configure Runtime A endpoint challenge", err)
 	}
-	if registration != nil {
-		if err := registration.Register(context.Background()); err != nil {
-			return err
-		}
-	}
-	server := &http.Server{Addr: config.ListenAddress, Handler: application}
-	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
-	defer stop()
-	serverErrors := make(chan error, 1)
-	go func() {
-		err := server.ListenAndServe()
-		if err != nil && !errors.Is(err, http.ErrServerClosed) {
-			serverErrors <- fmt.Errorf("runtime-a serve: %w", err)
-			return
-		}
-		serverErrors <- nil
-	}()
-	var registrationErrors chan error
-	if registration != nil {
-		registrationErrors = make(chan error, 1)
-		go func() { registrationErrors <- registration.Run(ctx) }()
-	}
-	var runErr error
-	registrationStopped := registration == nil
-	select {
-	case <-ctx.Done():
-	case runErr = <-serverErrors:
-	case runErr = <-registrationErrors:
-		registrationStopped = true
-	}
-	stop()
 	shutdownTimeout := 5 * time.Second
 	if registrationConfig.RequestTimeout > 0 {
 		shutdownTimeout = registrationConfig.RequestTimeout
 	}
-	shutdownContext, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
-	defer cancel()
-	shutdownErr := server.Shutdown(shutdownContext)
-	if !registrationStopped {
-		select {
-		case registrationErr := <-registrationErrors:
-			runErr = errors.Join(runErr, registrationErr)
-		case <-shutdownContext.Done():
-			runErr = errors.Join(runErr, errors.New("Runtime A Nacos heartbeat did not stop before shutdown"))
-		}
-	}
-	var deregisterErr error
-	if registration != nil {
-		deregisterErr = registration.Deregister(shutdownContext)
+	runtimeHost, err := agenthost.New(agenthost.Config{
+		Address:         config.ListenAddress,
+		Handler:         application,
+		Registration:    registration,
+		ShutdownTimeout: shutdownTimeout,
+		Signals:         []os.Signal{os.Interrupt, syscall.SIGTERM},
+	})
+	if err != nil {
+		return err
 	}
-	return errors.Join(runErr, shutdownErr, deregisterErr)
+	return runtimeHost.Run(context.Background())
 }
 
 type ready bool

runtime-a/cmd/runtime-a/main_test.go

@@ -0,0 +1,15 @@
+package main
+
+import (
+	"testing"
+
+	agenthost "github.com/NeKiro-project/nekiro-sdk-go/agent/host"
+)
+
+func TestRunStagesMissingRuntimeAConfiguration(t *testing.T) {
+	err := runWithLookup(func(string) (string, bool) { return "", false })
+	stage, ok := agenthost.StageOf(err)
+	if !ok || stage != agenthost.StageConfig {
+		t.Fatalf("StageOf(run error) = %q, %v; error=%v", stage, ok, err)
+	}
+}

runtime-b/cmd/runtime-b/main.go

@@ -2,18 +2,16 @@ package main
 
 import (
 	"context"
-	"errors"
-	"fmt"
 	"log"
 	"net/http"
 	"os"
-	"os/signal"
 	"syscall"
 	"time"
 
 	"github.com/NeKiro-project/NeKiro-Samples/internal/challengeproof"
 	"github.com/NeKiro-project/NeKiro-Samples/internal/nacosregistration"
 	runtimeb "github.com/NeKiro-project/NeKiro-Samples/runtime-b"
+	agenthost "github.com/NeKiro-project/nekiro-sdk-go/agent/host"
 	"github.com/NeKiro-project/nekiro-sdk-go/agent/routerauth"
 )
 
@@ -24,98 +22,67 @@ func main() {
 }
 
 func run() error {
-	address, err := runtimeb.ListenAddressFromEnvironment(os.LookupEnv)
+	return runWithLookup(os.LookupEnv)
+}
+
+func runWithLookup(lookup func(string) (string, bool)) error {
+	address, err := runtimeb.ListenAddressFromEnvironment(lookup)
 	if err != nil {
-		return err
+		return agenthost.Wrap(agenthost.StageConfig, "load Runtime B listen address", err)
 	}
-	authenticationConfig, err := routerauth.LoadConfig(os.LookupEnv)
+	authenticationConfig, err := routerauth.LoadConfig(lookup)
 	if err != nil {
-		return fmt.Errorf("runtime-b authentication config: %w", err)
+		return agenthost.Wrap(agenthost.StageConfig, "load Runtime B authentication configuration", err)
 	}
-	config, err := runtimeb.LoadConfig(os.LookupEnv)
+	config, err := runtimeb.LoadConfig(lookup)
 	if err != nil {
-		return err
+		return agenthost.Wrap(agenthost.StageConfig, "load Runtime B configuration", err)
 	}
-	registrationConfig, err := runtimeb.LoadRegistrationConfig(os.LookupEnv, config.AgentID, config.InstanceID)
+	registrationConfig, err := runtimeb.LoadRegistrationConfig(lookup, config.AgentID, config.InstanceID)
 	if err != nil {
-		return err
+		return agenthost.Wrap(agenthost.StageConfig, "load Runtime B registration configuration", err)
 	}
-	var registration *runtimeb.NacosRegistration
+	var registration agenthost.Registration
 	var readiness runtimeb.Readiness = ready(true)
 	if registrationConfig.Mode == runtimeb.RegistrationModeNacos {
 		registrationClient, clientErr := nacosregistration.NewHTTPClient(registrationConfig)
 		if clientErr != nil {
-			return fmt.Errorf("runtime-b Nacos registration transport: %w", clientErr)
+			return agenthost.Wrap(agenthost.StageRegistration, "create Runtime B Nacos transport", clientErr)
 		}
-		registration, err = runtimeb.NewNacosRegistration(registrationConfig, registrationClient)
+		runtimeRegistration, err := runtimeb.NewNacosRegistration(registrationConfig, registrationClient)
 		if err != nil {
-			return fmt.Errorf("runtime-b Nacos registration config: %w", err)
+			return agenthost.Wrap(agenthost.StageRegistration, "create Runtime B registration", err)
 		}
-		readiness = registration
+		registration = runtimeRegistration
+		readiness = runtimeRegistration
 	}
 	handler, err := runtimeb.NewConfiguredHandler(config, http.DefaultClient)
 	if err != nil {
-		return fmt.Errorf("runtime-b initialize: %w", err)
+		return agenthost.Wrap(agenthost.StageHandler, "create Runtime B handler", err)
 	}
 	execution, err := runtimeb.NewHTTPHandlerWithAuthAndReadiness(handler, authenticationConfig, readiness)
 	if err != nil {
-		return fmt.Errorf("runtime-b authentication: %w", err)
+		return agenthost.Wrap(agenthost.StageHandler, "configure Runtime B authentication", err)
 	}
-	application, err := challengeproof.NewHandler(execution, os.LookupEnv)
+	application, err := challengeproof.NewHandler(execution, lookup)
 	if err != nil {
-		return fmt.Errorf("runtime-b challenge proof: %w", err)
-	}
-	if registration != nil {
-		if err := registration.Register(context.Background()); err != nil {
-			return err
-		}
+		return agenthost.Wrap(agenthost.StageHandler, "configure Runtime B endpoint challenge", err)
 	}
-	server := &http.Server{Addr: address, Handler: application}
-	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
-	defer stop()
-	serverErrors := make(chan error, 1)
-	go func() {
-		err := server.ListenAndServe()
-		if err != nil && !errors.Is(err, http.ErrServerClosed) {
-			serverErrors <- fmt.Errorf("runtime-b serve: %w", err)
-			return
-		}
-		serverErrors <- nil
-	}()
-	var registrationErrors chan error
-	if registration != nil {
-		registrationErrors = make(chan error, 1)
-		go func() { registrationErrors <- registration.Run(ctx) }()
-	}
-	var runErr error
-	registrationStopped := registration == nil
-	select {
-	case <-ctx.Done():
-	case runErr = <-serverErrors:
-	case runErr = <-registrationErrors:
-		registrationStopped = true
-	}
-	stop()
 	shutdownTimeout := 5 * time.Second
 	if registrationConfig.RequestTimeout > 0 {
 		shutdownTimeout = registrationConfig.RequestTimeout
 	}
-	shutdownContext, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
-	defer cancel()
-	shutdownErr := server.Shutdown(shutdownContext)
-	if !registrationStopped {
-		select {
-		case registrationErr := <-registrationErrors:
-			runErr = errors.Join(runErr, registrationErr)
-		case <-shutdownContext.Done():
-			runErr = errors.Join(runErr, errors.New("Runtime B Nacos heartbeat did not stop before shutdown"))
-		}
-	}
-	var deregisterErr error
-	if registration != nil {
-		deregisterErr = registration.Deregister(shutdownContext)
+	runtimeHost, err := agenthost.New(agenthost.Config{
+		Address:         address,
+		Handler:         application,
+		Registration:    registration,
+		ShutdownTimeout: shutdownTimeout,
+		Signals:         []os.Signal{os.Interrupt, syscall.SIGTERM},
+	})
+	if err != nil {
+		return err
 	}
-	return errors.Join(runErr, shutdownErr, deregisterErr)
+	return runtimeHost.Run(context.Background())
 }
 
 type ready bool

runtime-b/cmd/runtime-b/main_test.go

@@ -0,0 +1,15 @@
+package main
+
+import (
+	"testing"
+
+	agenthost "github.com/NeKiro-project/nekiro-sdk-go/agent/host"
+)
+
+func TestRunStagesMissingRuntimeBConfiguration(t *testing.T) {
+	err := runWithLookup(func(string) (string, bool) { return "", false })
+	stage, ok := agenthost.StageOf(err)
+	if !ok || stage != agenthost.StageConfig {
+		t.Fatalf("StageOf(run error) = %q, %v; error=%v", stage, ok, err)
+	}
+}