diff --git a/embed.go b/embed.go index 0005e0c..2d82289 100644 --- a/embed.go +++ b/embed.go @@ -41,6 +41,9 @@ type Instance struct { wg *sync.WaitGroup shutdownOnce sync.Once + lg *zerolog.Logger + profileRegistered bool + pg *embeddedpostgres.EmbeddedPostgres stopPGOnce sync.Once } @@ -57,6 +60,17 @@ func (i *Instance) GRPCAddress() string { return i.grpcAddress } func (i *Instance) Shutdown(ctx context.Context) error { i.shutdownOnce.Do(func() { + // Remove the CLI profile registration first, so tooling stops + // discovering an engine that is about to go away. Token matched: a + // newer instance's registration is never deleted. + if i.profileRegistered { + if removed, path, err := deregisterEmbeddedProfile(i.token); err != nil { + i.lg.Warn().Err(err).Msg("could not remove the embedded CLI profile registration") + } else if removed { + i.lg.Debug().Msgf("removed the %q CLI profile from %s", embeddedProfileName, path) + } + } + i.cancel() close(i.interruptCh) }) @@ -310,15 +324,37 @@ func StartServer(ctx context.Context, opts ...Option) (inst *Instance, err error instanceAPIURL = apiURL } + // Register this instance as the "embedded" CLI profile so the hatchet CLI + // and its MCP server can discover it. Best effort: a failure (for example + // an unwritable home directory) never fails startup. Skipped without the + // API server, since a profile without an API URL is not usable by the CLI. + profileRegistered := false + if startServerAPI { + if profilePath, regErr := registerEmbeddedProfile(embeddedRegistration{ + tenantID: tenantID, + token: tok.Token, + apiURL: apiURL, + grpcAddress: grpcBroadcast, + expiresAt: expiresAt, + }); regErr != nil { + lg.Warn().Err(regErr).Msg("could not register the embedded engine as a CLI profile; continuing without it") + } else { + profileRegistered = true + lg.Info().Msgf("registered the %q CLI profile in %s", embeddedProfileName, profilePath) + } + } + return &Instance{ - token: tok.Token, - tenantID: tenantID, - apiURL: instanceAPIURL, - grpcAddress: grpcBroadcast, - interruptCh: interruptCh, - cancel: cancel, - wg: wg, - pg: pg, + token: tok.Token, + tenantID: tenantID, + apiURL: instanceAPIURL, + grpcAddress: grpcBroadcast, + interruptCh: interruptCh, + cancel: cancel, + wg: wg, + pg: pg, + lg: lg, + profileRegistered: profileRegistered, }, nil } diff --git a/go.mod b/go.mod index 11af1e8..437ee52 100644 --- a/go.mod +++ b/go.mod @@ -5,7 +5,7 @@ go 1.26.0 require ( github.com/fergusstrange/embedded-postgres v1.34.0 github.com/google/uuid v1.6.0 - github.com/hatchet-dev/hatchet v0.106.5 + github.com/hatchet-dev/hatchet v0.106.11 github.com/jackc/pgx/v5 v5.10.0 github.com/rs/zerolog v1.35.1 ) diff --git a/go.sum b/go.sum index d0eb8b0..09770b0 100644 --- a/go.sum +++ b/go.sum @@ -153,8 +153,8 @@ github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+l github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9/fmwbPZ6JB6eMoM= github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= -github.com/hatchet-dev/hatchet v0.106.5 h1:K2pK4ocAbEXjNw1XFMa7D+bY0X9Fn/BBendgh6wSlZw= -github.com/hatchet-dev/hatchet v0.106.5/go.mod h1:Dq1EmU3Ph2A5WzfijdCFf4KadC2mJjzAe5aqaCO6bcY= +github.com/hatchet-dev/hatchet v0.106.11 h1:xmcbTqJDgi4oKs911nFuTY7v32ceZR0cxrVTaFCHAT8= +github.com/hatchet-dev/hatchet v0.106.11/go.mod h1:Dq1EmU3Ph2A5WzfijdCFf4KadC2mJjzAe5aqaCO6bcY= github.com/hatchet-dev/pgoutbox v0.4.0 h1:+q9l14zzl/IVnxQ5PHZUyuZ+5D+leXqPJzJrt2iE+2I= github.com/hatchet-dev/pgoutbox v0.4.0/go.mod h1:x7wEFajIrOJ+goWri49MjzlkdwLHMdr+dZkXjVCm9ho= github.com/hatchet-dev/timediff v0.0.4 h1:RfYX1ehoa/qxHKAGQBMAvmkPx+FRQfUV37tDy/G1pOY= diff --git a/profileregistration.go b/profileregistration.go new file mode 100644 index 0000000..4b8209f --- /dev/null +++ b/profileregistration.go @@ -0,0 +1,91 @@ +package embed + +import ( + "os" + "time" + + "github.com/hatchet-dev/hatchet/pkg/config/cli" + "github.com/hatchet-dev/hatchet/pkg/config/cli/profilestore" +) + +// The embedded engine registers itself as a Hatchet CLI profile named +// "embedded" in ~/.hatchet/profiles.yaml when it becomes ready, and removes +// the registration again on graceful shutdown. This lets the hatchet CLI and +// its MCP server discover a running embedded engine through the ordinary +// profile store instead of environment handshakes or port probes. +// +// The store itself is the CLI's own pkg/config/cli/profilestore package, so +// the file location, lock protocol, and on-disk format always agree with the +// CLI. The extra metadata fields written below (embedded, pid, cwd, +// startedat) are ignored by the CLI's lenient reader; the registration itself +// shows up in `hatchet profile list` and is directly usable. +// +// Registration is best effort: any error (including opening the store) is +// logged as a warning by the caller and never fails engine startup. Removal +// is token matched, so an instance that shut down late never deletes a newer +// instance's registration. Last-writer-wins for the single "embedded" name is +// intended. + +// embeddedProfileName is the reserved CLI profile name for the embedded engine. +const embeddedProfileName = "embedded" + +// embeddedRegistration is the connection info registered for this instance. +type embeddedRegistration struct { + tenantID string + token string + apiURL string + grpcAddress string + expiresAt time.Time +} + +// registerEmbeddedProfile upserts the "embedded" profile into the CLI profile +// store, creating ~/.hatchet and the profiles file if absent. All other +// profiles, the default-profile setting, and any unrecognized file content +// (including comments) are preserved. Returns the path written for logging. +func registerEmbeddedProfile(reg embeddedRegistration) (string, error) { + store, err := profilestore.NewDefaultStore() + if err != nil { + return "", err + } + + // The cwd key is always written, even when empty, so a stale value from a + // previous registration cannot linger across upserts (UpsertProfile merges + // into the existing entry rather than replacing it). + cwd, _ := os.Getwd() + + err = store.UpsertProfile(embeddedProfileName, &cli.Profile{ + TenantId: reg.tenantID, + Name: embeddedProfileName, + Token: reg.token, + ExpiresAt: reg.expiresAt.UTC(), + ApiServerURL: reg.apiURL, + GrpcHostPort: reg.grpcAddress, + TLSStrategy: "none", + }, map[string]any{ + "embedded": true, + "pid": os.Getpid(), + "cwd": cwd, + "startedat": time.Now().UTC(), + }) + if err != nil { + return "", err + } + + return store.Path(), nil +} + +// deregisterEmbeddedProfile removes the "embedded" profile from the CLI +// profile store, but only when its token still matches this instance's token: +// a newer instance's registration (last writer wins) is left alone. The +// user's default-profile setting is intentionally never touched. Reports +// whether an entry was removed and the path acted on. +func deregisterEmbeddedProfile(token string) (bool, string, error) { + store, err := profilestore.NewDefaultStore() + if err != nil { + return false, "", err + } + + removed, err := store.RemoveProfileIfTokenMatches(embeddedProfileName, token) + + return removed, store.Path(), err +} diff --git a/profileregistration_test.go b/profileregistration_test.go new file mode 100644 index 0000000..91c4111 --- /dev/null +++ b/profileregistration_test.go @@ -0,0 +1,184 @@ +package embed + +import ( + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/hatchet-dev/hatchet/pkg/config/cli" + "github.com/hatchet-dev/hatchet/pkg/config/cli/profilestore" +) + +// These tests cover this repo's registration semantics through the exported +// profilestore API: what the embedded engine writes into the profile entry, +// that registration errors surface (the caller downgrades them to warnings), +// and the token-matched deregistration behavior. The store's own suite in +// github.com/hatchet-dev/hatchet covers the file format details (unquoted +// timestamps, comment preservation, flow-style traps, the lock protocol). + +func testRegistration(token string) embeddedRegistration { + return embeddedRegistration{ + tenantID: "707d0855-80ab-4e1f-a156-f1c4546cbf52", + token: token, + apiURL: "http://localhost:28243", + grpcAddress: "127.0.0.1:50051", + expiresAt: time.Now().UTC().Add(90 * 24 * time.Hour), + } +} + +func setTestHome(t *testing.T) string { + t.Helper() + home := t.TempDir() + t.Setenv("HOME", home) + return home +} + +// readEmbeddedProfile loads the profile store the same way the CLI does and +// returns the "embedded" profile. +func readEmbeddedProfile(t *testing.T) *cli.Profile { + t.Helper() + + store, err := profilestore.NewDefaultStore() + if err != nil { + t.Fatalf("could not open the profile store: %v", err) + } + + profile, err := store.GetProfile(embeddedProfileName) + if err != nil { + t.Fatalf("could not read the embedded profile: %v", err) + } + + return profile +} + +func TestRegisterEmbeddedProfileFreshFile(t *testing.T) { + home := setTestHome(t) + + path, err := registerEmbeddedProfile(testRegistration("token-fresh")) + if err != nil { + t.Fatalf("register failed: %v", err) + } + if want := filepath.Join(home, ".hatchet", "profiles.yaml"); path != want { + t.Fatalf("wrote to %s, want %s", path, want) + } + + profile := readEmbeddedProfile(t) + if profile.TenantId != "707d0855-80ab-4e1f-a156-f1c4546cbf52" || + profile.Name != "embedded" || + profile.Token != "token-fresh" || + profile.ApiServerURL != "http://localhost:28243" || + profile.GrpcHostPort != "127.0.0.1:50051" || + profile.TLSStrategy != "none" { + t.Errorf("unexpected embedded profile: %+v", profile) + } + if profile.ExpiresAt.IsZero() { + t.Errorf("expiresat did not round-trip through the store") + } + + // The metadata fields ride along in the same entry (the CLI ignores them). + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("could not read profiles file: %v", err) + } + raw := string(data) + for _, want := range []string{"embedded: true", "pid: ", "cwd: ", "startedat: "} { + if !strings.Contains(raw, want) { + t.Errorf("metadata %q missing from profiles file:\n%s", want, raw) + } + } +} + +func TestRegisterOverwritesPreviousRegistration(t *testing.T) { + setTestHome(t) + + if _, err := registerEmbeddedProfile(testRegistration("token-old")); err != nil { + t.Fatalf("first register failed: %v", err) + } + if _, err := registerEmbeddedProfile(testRegistration("token-new")); err != nil { + t.Fatalf("second register failed: %v", err) + } + + if got := readEmbeddedProfile(t).Token; got != "token-new" { + t.Errorf("embedded token = %q, want token-new (last writer wins)", got) + } +} + +func TestRegisterPermissionDeniedIsAnError(t *testing.T) { + if os.Getuid() == 0 { + t.Skip("running as root; permission checks do not apply") + } + + home := setTestHome(t) + + dir := filepath.Join(home, ".hatchet") + if err := os.MkdirAll(dir, 0o500); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = os.Chmod(dir, 0o700) }) + + // The caller treats this as a warning, never a startup failure; here we + // only assert that it surfaces as an error instead of a panic or a write. + if _, err := registerEmbeddedProfile(testRegistration("token-denied")); err == nil { + t.Fatal("expected an error registering into an unwritable directory") + } +} + +func TestDeregisterRemovesMatchingToken(t *testing.T) { + setTestHome(t) + + if _, err := registerEmbeddedProfile(testRegistration("token-mine")); err != nil { + t.Fatalf("register failed: %v", err) + } + + removed, _, err := deregisterEmbeddedProfile("token-mine") + if err != nil { + t.Fatalf("deregister failed: %v", err) + } + if !removed { + t.Fatal("expected the registration to be removed") + } + + store, err := profilestore.NewDefaultStore() + if err != nil { + t.Fatalf("could not open the profile store: %v", err) + } + if _, err := store.GetProfile(embeddedProfileName); err == nil { + t.Errorf("embedded profile still present after deregistration") + } +} + +func TestDeregisterKeepsNewerRegistration(t *testing.T) { + setTestHome(t) + + // A newer instance overwrote the registration; the older instance's + // shutdown must not delete it. + if _, err := registerEmbeddedProfile(testRegistration("token-newer")); err != nil { + t.Fatalf("register failed: %v", err) + } + + removed, _, err := deregisterEmbeddedProfile("token-older") + if err != nil { + t.Fatalf("deregister failed: %v", err) + } + if removed { + t.Fatal("deregistration removed another instance's registration") + } + + if got := readEmbeddedProfile(t).Token; got != "token-newer" { + t.Errorf("embedded token = %q, want token-newer", got) + } +} + +func TestDeregisterWithoutFileIsANoOp(t *testing.T) { + setTestHome(t) + + removed, _, err := deregisterEmbeddedProfile("token-any") + if err != nil { + t.Fatalf("deregister failed: %v", err) + } + if removed { + t.Fatal("nothing to remove, but removed was reported") + } +}