From 5c2abd2027d59d9212d4a13f2fcbc733f097e049 Mon Sep 17 00:00:00 2001 From: Alexander Belanger Date: Wed, 9 Sep 2026 14:18:08 -0400 Subject: [PATCH 1/2] feat: register a CLI profile for running embedded engines When the engine becomes ready, StartServer registers an embedded profile in the CLI profile store (~/.hatchet/profiles.yaml, created if absent) with the minted token and endpoints, plus metadata fields the released CLI ignores (embedded, pid, cwd, startedat). Registration is always attempted and never fails startup: write and permission errors log a warning and continue. Graceful shutdown removes the entry only when its token still matches this instance, so a newer engine's registration survives. The user's default profile is never changed. The writer honors the CLI's config lock protocol and profile file name overrides, upserts via yaml.Node so other profiles, ordering, and comments are preserved, and writes atomically. Timestamps must be emitted as unquoted YAML timestamps because viper does not decode quoted strings into time.Time; the upsert forces block style and a regression test guards it. Registration is skipped under WithoutAPI since such a profile would be unusable and unverifiable. The sidecar inherits this through StartServer, so TypeScript and Python embedded runs register identically with no SDK changes. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XHyfTNZGxYR249FURndHEX --- embed.go | 52 ++++- go.mod | 4 +- profileregistration.go | 426 ++++++++++++++++++++++++++++++++++++ profileregistration_test.go | 379 ++++++++++++++++++++++++++++++++ 4 files changed, 851 insertions(+), 10 deletions(-) create mode 100644 profileregistration.go create mode 100644 profileregistration_test.go 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..9db2433 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,8 @@ require ( github.com/hatchet-dev/hatchet v0.106.5 github.com/jackc/pgx/v5 v5.10.0 github.com/rs/zerolog v1.35.1 + github.com/spf13/viper v1.21.0 + gopkg.in/yaml.v3 v3.0.1 ) require ( @@ -104,7 +106,6 @@ require ( github.com/spf13/afero v1.15.0 // indirect github.com/spf13/cast v1.10.0 // indirect github.com/spf13/pflag v1.0.10 // indirect - github.com/spf13/viper v1.21.0 // indirect github.com/subosito/gotenv v1.6.0 // indirect github.com/tink-crypto/tink-go v0.0.0-20230613075026-d6de17e3f164 // indirect github.com/tink-crypto/tink-go-gcpkms v0.0.0-20230602082706-31d0d09ccc8d // indirect @@ -140,5 +141,4 @@ require ( google.golang.org/genproto/googleapis/rpc v0.0.0-20260825221802-da73d73af1c5 // indirect google.golang.org/grpc v1.83.2 // indirect google.golang.org/protobuf v1.36.12 // indirect - gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/profileregistration.go b/profileregistration.go new file mode 100644 index 0000000..3d3c91d --- /dev/null +++ b/profileregistration.go @@ -0,0 +1,426 @@ +package embed + +import ( + "bytes" + "fmt" + "os" + "path/filepath" + "strings" + "time" + + "gopkg.in/yaml.v3" +) + +// 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 CLI cannot share its profile-store code with this module (it lives +// under an internal/ path), so the file format is written directly here: +// +// - The CLI loads the file leniently via viper (no strict decoding), so the +// extra metadata fields written below (embedded, pid, cwd, startedat) are +// ignored by released CLI versions; the registration itself shows up in +// `hatchet profile list` and is directly usable. +// - Keys are written lowercased to match what viper's own writer produces; +// viper reads keys case-insensitively either way. +// - expiresat and startedat must serialize as unquoted YAML timestamps: the +// CLI decodes profiles into a struct with time.Time fields, and a quoted +// string there would fail its (otherwise lenient) unmarshalling. +// +// Registration is best effort: any error 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" + +const ( + profilesLockTimeout = 5 * time.Second + profilesLockRetryDelay = 50 * time.Millisecond + profilesLockMaxAttempts = 100 +) + +// embeddedProfileEntry is the profile written to the CLI profile store. The +// first block mirrors the CLI's released profile schema; the second block is +// metadata the CLI ignores. +type embeddedProfileEntry struct { + TenantID string `yaml:"tenantid"` + Name string `yaml:"name"` + Token string `yaml:"token"` + ExpiresAt time.Time `yaml:"expiresat"` + APIServerURL string `yaml:"apiserverurl"` + GrpcHostPort string `yaml:"grpchostport"` + TLSStrategy string `yaml:"tlsstrategy"` + + Embedded bool `yaml:"embedded"` + PID int `yaml:"pid"` + Cwd string `yaml:"cwd,omitempty"` + StartedAt time.Time `yaml:"startedat"` +} + +// 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) { + dir, path, err := profilesFilePath() + if err != nil { + return "", err + } + + if err := os.MkdirAll(dir, 0o700); err != nil { + return "", fmt.Errorf("could not create %s: %w", dir, err) + } + + unlock, err := acquireProfilesLock(dir) + if err != nil { + return "", err + } + defer unlock() + + doc, root, err := loadProfilesDocument(path) + if err != nil { + return "", err + } + + profiles, err := mappingValueOrCreate(root, "profiles") + if err != nil { + return "", err + } + + cwd, _ := os.Getwd() + entryNode, err := yamlNodeFor(embeddedProfileEntry{ + TenantID: reg.tenantID, + Name: embeddedProfileName, + Token: reg.token, + ExpiresAt: reg.expiresAt.UTC(), + APIServerURL: reg.apiURL, + GrpcHostPort: reg.grpcAddress, + TLSStrategy: "none", + Embedded: true, + PID: os.Getpid(), + Cwd: cwd, + StartedAt: time.Now().UTC(), + }) + if err != nil { + return "", err + } + + setMappingValue(profiles, embeddedProfileName, entryNode) + + // Force block style on the mappings we touch: an empty "profiles: {}" + // left behind by a previous deregistration parses as flow style, and flow + // style would make the encoder quote the timestamp values, which the CLI + // cannot decode into its time.Time fields. + root.Style = 0 + profiles.Style = 0 + + // The user's defaultprofile setting is intentionally never touched. + if err := writeYAMLAtomic(path, doc); err != nil { + return "", err + } + + return 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. Reports +// whether an entry was removed and the path acted on. +func deregisterEmbeddedProfile(token string) (bool, string, error) { + dir, path, err := profilesFilePath() + if err != nil { + return false, "", err + } + + if _, err := os.Stat(path); os.IsNotExist(err) { + return false, path, nil + } + + unlock, err := acquireProfilesLock(dir) + if err != nil { + return false, path, err + } + defer unlock() + + doc, root, err := loadProfilesDocument(path) + if err != nil { + return false, path, err + } + + profiles := mappingValue(root, "profiles") + if profiles == nil || profiles.Kind != yaml.MappingNode { + return false, path, nil + } + + entry := mappingValue(profiles, embeddedProfileName) + if entry == nil || entry.Kind != yaml.MappingNode { + return false, path, nil + } + + entryToken := mappingValue(entry, "token") + if entryToken == nil || entryToken.Value != token { + return false, path, nil + } + + deleteMappingKey(profiles, embeddedProfileName) + + // A defaultprofile setting pointing at "embedded" is left as is: changing + // the user's default is out of scope here, and both the CLI and the MCP + // server tolerate a default that names a missing profile. + if err := writeYAMLAtomic(path, doc); err != nil { + return false, path, err + } + + return true, path, nil +} + +// profilesFilePath resolves the CLI profile store location the same way the +// CLI does: ~/.hatchet plus the profile file name from the +// HATCHET_CLI_PROFILE_FILE_NAME env var or the profileFileName key in +// ~/.hatchet/config.yaml, defaulting to profiles.yaml. +func profilesFilePath() (dir string, path string, err error) { + home, err := os.UserHomeDir() + if err != nil { + return "", "", fmt.Errorf("could not determine the home directory: %w", err) + } + + dir = filepath.Join(home, ".hatchet") + + name := os.Getenv("HATCHET_CLI_PROFILE_FILE_NAME") + if name == "" { + name = profileFileNameFromConfig(filepath.Join(dir, "config.yaml")) + } + if name == "" { + name = "profiles.yaml" + } + + return dir, filepath.Join(dir, name), nil +} + +// profileFileNameFromConfig reads the profileFileName key (any casing) from +// the CLI config file, returning "" when the file or key is absent or +// unreadable. +func profileFileNameFromConfig(configPath string) string { + data, err := os.ReadFile(configPath) + if err != nil { + return "" + } + + var conf map[string]any + if err := yaml.Unmarshal(data, &conf); err != nil { + return "" + } + + for key, value := range conf { + if strings.EqualFold(key, "profilefilename") { + if name, ok := value.(string); ok { + return name + } + } + } + + return "" +} + +// acquireProfilesLock takes the CLI's config lock (~/.hatchet/config.lock) +// with the same protocol the CLI uses: exclusive create, retry every 50ms up +// to 100 attempts, and locks older than 5s are considered stale. +func acquireProfilesLock(dir string) (func(), error) { + lockFile := filepath.Join(dir, "config.lock") + + for attempts := 0; attempts < profilesLockMaxAttempts; attempts++ { + f, err := os.OpenFile(lockFile, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o644) + if err == nil { + _, writeErr := f.WriteString(time.Now().Format(time.RFC3339)) + _ = f.Close() + if writeErr != nil { + _ = os.Remove(lockFile) + return nil, fmt.Errorf("could not write lock file: %w", writeErr) + } + return func() { _ = os.Remove(lockFile) }, nil + } + + if os.IsExist(err) { + if stat, statErr := os.Stat(lockFile); statErr == nil && time.Since(stat.ModTime()) > profilesLockTimeout { + _ = os.Remove(lockFile) + continue + } + } else { + return nil, fmt.Errorf("could not create lock file: %w", err) + } + + time.Sleep(profilesLockRetryDelay) + } + + return nil, fmt.Errorf("could not acquire the profile store lock at %s", lockFile) +} + +// loadProfilesDocument parses the profiles file into a yaml document whose +// root is a mapping, creating an empty document when the file is absent or +// empty. Parse failures are returned rather than overwriting the file. +func loadProfilesDocument(path string) (doc *yaml.Node, root *yaml.Node, err error) { + data, err := os.ReadFile(path) + if err != nil && !os.IsNotExist(err) { + return nil, nil, fmt.Errorf("could not read %s: %w", path, err) + } + + if len(bytes.TrimSpace(data)) == 0 { + root = &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} + doc = &yaml.Node{Kind: yaml.DocumentNode, Content: []*yaml.Node{root}} + return doc, root, nil + } + + doc = &yaml.Node{} + if err := yaml.Unmarshal(data, doc); err != nil { + return nil, nil, fmt.Errorf("could not parse %s (leaving it untouched): %w", path, err) + } + + if len(doc.Content) == 0 { + root = &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} + doc = &yaml.Node{Kind: yaml.DocumentNode, Content: []*yaml.Node{root}} + return doc, root, nil + } + + root = doc.Content[0] + if root.Kind == yaml.ScalarNode && root.Tag == "!!null" { + *root = yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} + return doc, root, nil + } + if root.Kind != yaml.MappingNode { + return nil, nil, fmt.Errorf("%s has an unexpected structure (top level is not a mapping); leaving it untouched", path) + } + + return doc, root, nil +} + +// yamlNodeFor round-trips v through the yaml encoder to obtain its node +// representation, so values like time.Time serialize exactly as the encoder +// would write them (unquoted timestamps). +func yamlNodeFor(v any) (*yaml.Node, error) { + data, err := yaml.Marshal(v) + if err != nil { + return nil, err + } + + doc := &yaml.Node{} + if err := yaml.Unmarshal(data, doc); err != nil { + return nil, err + } + if len(doc.Content) == 0 { + return nil, fmt.Errorf("could not build a yaml node") + } + + return doc.Content[0], nil +} + +// mappingValue returns the value node for key in mapping m, matching keys +// case-insensitively (viper treats profile keys case-insensitively). +func mappingValue(m *yaml.Node, key string) *yaml.Node { + if m == nil || m.Kind != yaml.MappingNode { + return nil + } + + for i := 0; i+1 < len(m.Content); i += 2 { + if strings.EqualFold(m.Content[i].Value, key) { + return m.Content[i+1] + } + } + + return nil +} + +// mappingValueOrCreate returns the mapping value for key, creating an empty +// mapping entry when the key is absent or null. +func mappingValueOrCreate(m *yaml.Node, key string) (*yaml.Node, error) { + value := mappingValue(m, key) + if value == nil { + value = &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} + m.Content = append(m.Content, + &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key}, + value, + ) + return value, nil + } + + if value.Kind == yaml.ScalarNode && value.Tag == "!!null" { + *value = yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} + return value, nil + } + if value.Kind != yaml.MappingNode { + return nil, fmt.Errorf("the %q key in the profiles file is not a mapping; leaving the file untouched", key) + } + + return value, nil +} + +// setMappingValue replaces the value for key in mapping m (case-insensitive), +// appending the pair when the key is absent. +func setMappingValue(m *yaml.Node, key string, value *yaml.Node) { + for i := 0; i+1 < len(m.Content); i += 2 { + if strings.EqualFold(m.Content[i].Value, key) { + m.Content[i+1] = value + return + } + } + + m.Content = append(m.Content, + &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key}, + value, + ) +} + +// deleteMappingKey removes key (case-insensitive) and its value from mapping m. +func deleteMappingKey(m *yaml.Node, key string) bool { + for i := 0; i+1 < len(m.Content); i += 2 { + if strings.EqualFold(m.Content[i].Value, key) { + m.Content = append(m.Content[:i], m.Content[i+2:]...) + return true + } + } + + return false +} + +// writeYAMLAtomic writes doc to path via a same-directory temp file and +// rename, so concurrent readers never see a partial file. Two concurrent +// writers still race whole-file (last rename wins); with the config lock held +// this only matters against writers that ignore the lock, and the loser's +// write is a well-formed file. +func writeYAMLAtomic(path string, doc *yaml.Node) error { + var buf bytes.Buffer + enc := yaml.NewEncoder(&buf) + enc.SetIndent(4) + if err := enc.Encode(doc); err != nil { + _ = enc.Close() + return fmt.Errorf("could not encode the profiles file: %w", err) + } + if err := enc.Close(); err != nil { + return fmt.Errorf("could not encode the profiles file: %w", err) + } + + tmp := fmt.Sprintf("%s.%d.tmp", path, os.Getpid()) + if err := os.WriteFile(tmp, buf.Bytes(), 0o600); err != nil { + return fmt.Errorf("could not write %s: %w", tmp, err) + } + + if err := os.Rename(tmp, path); err != nil { + _ = os.Remove(tmp) + return fmt.Errorf("could not replace %s: %w", path, err) + } + + return nil +} diff --git a/profileregistration_test.go b/profileregistration_test.go new file mode 100644 index 0000000..8eb4660 --- /dev/null +++ b/profileregistration_test.go @@ -0,0 +1,379 @@ +package embed + +import ( + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "github.com/spf13/viper" + "gopkg.in/yaml.v3" +) + +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 +} + +func readProfilesFile(t *testing.T, home string) (string, map[string]any) { + t.Helper() + + data, err := os.ReadFile(filepath.Join(home, ".hatchet", "profiles.yaml")) + if err != nil { + t.Fatalf("could not read profiles file: %v", err) + } + + var parsed map[string]any + if err := yaml.Unmarshal(data, &parsed); err != nil { + t.Fatalf("written profiles file does not parse: %v", err) + } + + return string(data), parsed +} + +func profileEntry(t *testing.T, parsed map[string]any, name string) map[string]any { + t.Helper() + + profiles, ok := parsed["profiles"].(map[string]any) + if !ok { + t.Fatalf("profiles key missing or not a mapping: %#v", parsed["profiles"]) + } + entry, ok := profiles[name].(map[string]any) + if !ok { + t.Fatalf("profile %q missing or not a mapping: %#v", name, profiles[name]) + } + + return entry +} + +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) + } + + raw, parsed := readProfilesFile(t, home) + entry := profileEntry(t, parsed, "embedded") + + for key, want := range map[string]any{ + "tenantid": "707d0855-80ab-4e1f-a156-f1c4546cbf52", + "name": "embedded", + "token": "token-fresh", + "apiserverurl": "http://localhost:28243", + "grpchostport": "127.0.0.1:50051", + "tlsstrategy": "none", + "embedded": true, + } { + if entry[key] != want { + t.Errorf("entry[%q] = %#v, want %#v", key, entry[key], want) + } + } + if _, ok := entry["pid"].(int); !ok { + t.Errorf("entry[pid] = %#v, want an int", entry["pid"]) + } + + // The CLI decodes expiresat into a time.Time; a quoted string there would + // break its unmarshalling, so the timestamp must be written unquoted. + if strings.Contains(raw, `expiresat: "`) || strings.Contains(raw, "expiresat: '") { + t.Errorf("expiresat was written quoted:\n%s", raw) + } + if _, ok := entry["expiresat"].(time.Time); !ok { + t.Errorf("expiresat did not parse as a yaml timestamp: %#v", entry["expiresat"]) + } +} + +// TestRegisteredProfileReadableByViper anchors interop with the CLI: the CLI +// loads the profiles file via viper (lenient unmarshal into a struct with +// time.Time fields) and per-key lookups, so both must work on our output. +func TestRegisteredProfileReadableByViper(t *testing.T) { + home := setTestHome(t) + + if _, err := registerEmbeddedProfile(testRegistration("token-viper")); err != nil { + t.Fatalf("register failed: %v", err) + } + + v := viper.New() + v.SetConfigFile(filepath.Join(home, ".hatchet", "profiles.yaml")) + v.SetConfigType("yaml") + if err := v.ReadInConfig(); err != nil { + t.Fatalf("viper could not read the profiles file: %v", err) + } + + if got := v.GetString("profiles.embedded.token"); got != "token-viper" { + t.Errorf("viper token = %q, want token-viper", got) + } + if got := v.GetTime("profiles.embedded.expiresat"); got.IsZero() { + t.Errorf("viper expiresat is zero") + } + + // Mirrors pkg/config/cli.ProfileFile in the CLI: unknown metadata fields + // must not break the struct unmarshal. + type profile struct { + TenantId string `mapstructure:"tenantId"` + Name string `mapstructure:"name"` + Token string `mapstructure:"token"` + ExpiresAt time.Time `mapstructure:"expiresAt"` + ApiServerURL string `mapstructure:"apiServerURL"` + GrpcHostPort string `mapstructure:"grpcHostPort"` + TLSStrategy string `mapstructure:"tlsStrategy"` + } + var file struct { + Profiles map[string]profile `mapstructure:"profiles"` + } + if err := v.Unmarshal(&file); err != nil { + t.Fatalf("viper struct unmarshal failed (the released CLI would fatal on this): %v", err) + } + + got := file.Profiles["embedded"] + if got.Token != "token-viper" || got.TLSStrategy != "none" || got.ExpiresAt.IsZero() { + t.Errorf("unmarshalled profile = %+v", got) + } +} + +func TestRegisterPreservesExistingContent(t *testing.T) { + home := setTestHome(t) + + existing := `# keep this comment +defaultprofile: prod +profiles: + prod: + apiserverurl: https://prod.example.com + expiresat: 2027-01-02T03:04:05Z + grpchostport: prod.example.com:443 + name: prod + tenantid: 11111111-2222-3333-4444-555555555555 + tlsstrategy: tls + token: prod-token +` + dir := filepath.Join(home, ".hatchet") + if err := os.MkdirAll(dir, 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "profiles.yaml"), []byte(existing), 0o600); err != nil { + t.Fatal(err) + } + + if _, err := registerEmbeddedProfile(testRegistration("token-merge")); err != nil { + t.Fatalf("register failed: %v", err) + } + + raw, parsed := readProfilesFile(t, home) + + if parsed["defaultprofile"] != "prod" { + t.Errorf("defaultprofile = %#v, want prod", parsed["defaultprofile"]) + } + if !strings.Contains(raw, "# keep this comment") { + t.Errorf("comment was not preserved:\n%s", raw) + } + + prod := profileEntry(t, parsed, "prod") + for key, want := range map[string]any{ + "apiserverurl": "https://prod.example.com", + "grpchostport": "prod.example.com:443", + "name": "prod", + "tenantid": "11111111-2222-3333-4444-555555555555", + "tlsstrategy": "tls", + "token": "prod-token", + } { + if prod[key] != want { + t.Errorf("prod[%q] = %#v, want %#v", key, prod[key], want) + } + } + + embedded := profileEntry(t, parsed, "embedded") + if embedded["token"] != "token-merge" { + t.Errorf("embedded token = %#v", embedded["token"]) + } +} + +func TestRegisterOverwritesPreviousRegistration(t *testing.T) { + home := 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) + } + + _, parsed := readProfilesFile(t, home) + entry := profileEntry(t, parsed, "embedded") + if entry["token"] != "token-new" { + t.Errorf("embedded token = %#v, want token-new (last writer wins)", entry["token"]) + } +} + +// TestReRegistrationAfterDeregistrationStaysCLIReadable guards against a flow +// style trap: a deregistration that empties the store leaves "profiles: {}" +// behind, and a later registration merged into that flow mapping would be +// emitted in flow style with quoted timestamps, which the CLI cannot decode +// into its time.Time fields. +func TestReRegistrationAfterDeregistrationStaysCLIReadable(t *testing.T) { + home := setTestHome(t) + + if _, err := registerEmbeddedProfile(testRegistration("token-first")); err != nil { + t.Fatalf("register failed: %v", err) + } + if _, _, err := deregisterEmbeddedProfile("token-first"); err != nil { + t.Fatalf("deregister failed: %v", err) + } + if _, err := registerEmbeddedProfile(testRegistration("token-second")); err != nil { + t.Fatalf("re-register failed: %v", err) + } + + raw, _ := readProfilesFile(t, home) + if strings.Contains(raw, "{") { + t.Errorf("profiles file was written in flow style:\n%s", raw) + } + + v := viper.New() + v.SetConfigFile(filepath.Join(home, ".hatchet", "profiles.yaml")) + v.SetConfigType("yaml") + if err := v.ReadInConfig(); err != nil { + t.Fatalf("viper could not read the profiles file: %v", err) + } + var file struct { + Profiles map[string]struct { + ExpiresAt time.Time `mapstructure:"expiresAt"` + Token string `mapstructure:"token"` + } `mapstructure:"profiles"` + } + if err := v.Unmarshal(&file); err != nil { + t.Fatalf("viper struct unmarshal failed (the released CLI would fatal on this): %v", err) + } + if file.Profiles["embedded"].Token != "token-second" || file.Profiles["embedded"].ExpiresAt.IsZero() { + t.Errorf("unexpected re-registered profile: %+v", file.Profiles["embedded"]) + } +} + +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) { + home := 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") + } + + _, parsed := readProfilesFile(t, home) + if profiles, ok := parsed["profiles"].(map[string]any); ok { + if _, still := profiles["embedded"]; still { + t.Errorf("embedded profile still present after deregistration") + } + } +} + +func TestDeregisterKeepsNewerRegistration(t *testing.T) { + home := 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") + } + + _, parsed := readProfilesFile(t, home) + entry := profileEntry(t, parsed, "embedded") + if entry["token"] != "token-newer" { + t.Errorf("embedded token = %#v, want token-newer", entry["token"]) + } +} + +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") + } +} + +func TestConcurrentRegisterDeregisterKeepsFileWellFormed(t *testing.T) { + home := setTestHome(t) + + var wg sync.WaitGroup + for i := 0; i < 8; i++ { + wg.Add(1) + go func(n int) { + defer wg.Done() + token := "token-a" + if n%2 == 1 { + token = "token-b" + } + if _, err := registerEmbeddedProfile(testRegistration(token)); err != nil { + t.Errorf("register failed: %v", err) + } + if _, _, err := deregisterEmbeddedProfile("token-a"); err != nil { + t.Errorf("deregister failed: %v", err) + } + }(i) + } + wg.Wait() + + // Whatever interleaving happened, the file must parse and contain either + // no embedded profile or a well-formed one. + _, parsed := readProfilesFile(t, home) + if profiles, ok := parsed["profiles"].(map[string]any); ok { + if raw, present := profiles["embedded"]; present { + if _, ok := raw.(map[string]any); !ok { + t.Errorf("embedded entry is malformed: %#v", raw) + } + } + } +} From 1fad43d43bf24ac8ee6da7369398faa21ee19e98 Mon Sep 17 00:00:00 2001 From: Alexander Belanger Date: Fri, 11 Sep 2026 13:07:45 -0400 Subject: [PATCH 2/2] refactor: use the public profilestore package from hatchet v0.106.11 Replaces the duplicated profile-writing machinery (path resolution, lock protocol, yaml.Node surgery, atomic writes, roughly 350 lines) with the profilestore package the CLI itself now uses, so all writers share one implementation. Registration upserts through UpsertProfile with the embedded metadata fields, always including cwd since upserts merge rather than replace, and deregistration uses RemoveProfileIfTokenMatches. Store construction happens inside register and deregister so failures stay best-effort warnings that never fail engine startup. Tests are pruned to this repo's semantics; format-fidelity coverage lives in the profilestore suite. Verified live against the released v0.106.11 CLI: registration parses and lists, comments survive, graceful shutdown removes the entry, re-registration into an emptied file stays readable, and a killed engine leaves a stale entry as designed. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01XHyfTNZGxYR249FURndHEX --- go.mod | 6 +- go.sum | 4 +- profileregistration.go | 397 +++--------------------------------- profileregistration_test.go | 293 +++++--------------------- 4 files changed, 85 insertions(+), 615 deletions(-) diff --git a/go.mod b/go.mod index 9db2433..437ee52 100644 --- a/go.mod +++ b/go.mod @@ -5,11 +5,9 @@ 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 - github.com/spf13/viper v1.21.0 - gopkg.in/yaml.v3 v3.0.1 ) require ( @@ -106,6 +104,7 @@ require ( github.com/spf13/afero v1.15.0 // indirect github.com/spf13/cast v1.10.0 // indirect github.com/spf13/pflag v1.0.10 // indirect + github.com/spf13/viper v1.21.0 // indirect github.com/subosito/gotenv v1.6.0 // indirect github.com/tink-crypto/tink-go v0.0.0-20230613075026-d6de17e3f164 // indirect github.com/tink-crypto/tink-go-gcpkms v0.0.0-20230602082706-31d0d09ccc8d // indirect @@ -141,4 +140,5 @@ require ( google.golang.org/genproto/googleapis/rpc v0.0.0-20260825221802-da73d73af1c5 // indirect google.golang.org/grpc v1.83.2 // indirect google.golang.org/protobuf v1.36.12 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect ) 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 index 3d3c91d..4b8209f 100644 --- a/profileregistration.go +++ b/profileregistration.go @@ -1,14 +1,11 @@ package embed import ( - "bytes" - "fmt" "os" - "path/filepath" - "strings" "time" - "gopkg.in/yaml.v3" + "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 @@ -17,51 +14,21 @@ import ( // its MCP server discover a running embedded engine through the ordinary // profile store instead of environment handshakes or port probes. // -// The CLI cannot share its profile-store code with this module (it lives -// under an internal/ path), so the file format is written directly here: +// 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. // -// - The CLI loads the file leniently via viper (no strict decoding), so the -// extra metadata fields written below (embedded, pid, cwd, startedat) are -// ignored by released CLI versions; the registration itself shows up in -// `hatchet profile list` and is directly usable. -// - Keys are written lowercased to match what viper's own writer produces; -// viper reads keys case-insensitively either way. -// - expiresat and startedat must serialize as unquoted YAML timestamps: the -// CLI decodes profiles into a struct with time.Time fields, and a quoted -// string there would fail its (otherwise lenient) unmarshalling. -// -// Registration is best effort: any error 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. +// 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" -const ( - profilesLockTimeout = 5 * time.Second - profilesLockRetryDelay = 50 * time.Millisecond - profilesLockMaxAttempts = 100 -) - -// embeddedProfileEntry is the profile written to the CLI profile store. The -// first block mirrors the CLI's released profile schema; the second block is -// metadata the CLI ignores. -type embeddedProfileEntry struct { - TenantID string `yaml:"tenantid"` - Name string `yaml:"name"` - Token string `yaml:"token"` - ExpiresAt time.Time `yaml:"expiresat"` - APIServerURL string `yaml:"apiserverurl"` - GrpcHostPort string `yaml:"grpchostport"` - TLSStrategy string `yaml:"tlsstrategy"` - - Embedded bool `yaml:"embedded"` - PID int `yaml:"pid"` - Cwd string `yaml:"cwd,omitempty"` - StartedAt time.Time `yaml:"startedat"` -} - // embeddedRegistration is the connection info registered for this instance. type embeddedRegistration struct { tenantID string @@ -76,351 +43,49 @@ type embeddedRegistration struct { // 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) { - dir, path, err := profilesFilePath() - if err != nil { - return "", err - } - - if err := os.MkdirAll(dir, 0o700); err != nil { - return "", fmt.Errorf("could not create %s: %w", dir, err) - } - - unlock, err := acquireProfilesLock(dir) - if err != nil { - return "", err - } - defer unlock() - - doc, root, err := loadProfilesDocument(path) - if err != nil { - return "", err - } - - profiles, err := mappingValueOrCreate(root, "profiles") + 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() - entryNode, err := yamlNodeFor(embeddedProfileEntry{ - TenantID: reg.tenantID, + + err = store.UpsertProfile(embeddedProfileName, &cli.Profile{ + TenantId: reg.tenantID, Name: embeddedProfileName, Token: reg.token, ExpiresAt: reg.expiresAt.UTC(), - APIServerURL: reg.apiURL, + ApiServerURL: reg.apiURL, GrpcHostPort: reg.grpcAddress, TLSStrategy: "none", - Embedded: true, - PID: os.Getpid(), - Cwd: cwd, - StartedAt: time.Now().UTC(), + }, map[string]any{ + "embedded": true, + "pid": os.Getpid(), + "cwd": cwd, + "startedat": time.Now().UTC(), }) if err != nil { return "", err } - setMappingValue(profiles, embeddedProfileName, entryNode) - - // Force block style on the mappings we touch: an empty "profiles: {}" - // left behind by a previous deregistration parses as flow style, and flow - // style would make the encoder quote the timestamp values, which the CLI - // cannot decode into its time.Time fields. - root.Style = 0 - profiles.Style = 0 - - // The user's defaultprofile setting is intentionally never touched. - if err := writeYAMLAtomic(path, doc); err != nil { - return "", err - } - - return path, nil + 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. Reports +// 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) { - dir, path, err := profilesFilePath() + store, err := profilestore.NewDefaultStore() if err != nil { return false, "", err } - if _, err := os.Stat(path); os.IsNotExist(err) { - return false, path, nil - } - - unlock, err := acquireProfilesLock(dir) - if err != nil { - return false, path, err - } - defer unlock() - - doc, root, err := loadProfilesDocument(path) - if err != nil { - return false, path, err - } - - profiles := mappingValue(root, "profiles") - if profiles == nil || profiles.Kind != yaml.MappingNode { - return false, path, nil - } - - entry := mappingValue(profiles, embeddedProfileName) - if entry == nil || entry.Kind != yaml.MappingNode { - return false, path, nil - } - - entryToken := mappingValue(entry, "token") - if entryToken == nil || entryToken.Value != token { - return false, path, nil - } - - deleteMappingKey(profiles, embeddedProfileName) - - // A defaultprofile setting pointing at "embedded" is left as is: changing - // the user's default is out of scope here, and both the CLI and the MCP - // server tolerate a default that names a missing profile. - if err := writeYAMLAtomic(path, doc); err != nil { - return false, path, err - } - - return true, path, nil -} - -// profilesFilePath resolves the CLI profile store location the same way the -// CLI does: ~/.hatchet plus the profile file name from the -// HATCHET_CLI_PROFILE_FILE_NAME env var or the profileFileName key in -// ~/.hatchet/config.yaml, defaulting to profiles.yaml. -func profilesFilePath() (dir string, path string, err error) { - home, err := os.UserHomeDir() - if err != nil { - return "", "", fmt.Errorf("could not determine the home directory: %w", err) - } - - dir = filepath.Join(home, ".hatchet") - - name := os.Getenv("HATCHET_CLI_PROFILE_FILE_NAME") - if name == "" { - name = profileFileNameFromConfig(filepath.Join(dir, "config.yaml")) - } - if name == "" { - name = "profiles.yaml" - } - - return dir, filepath.Join(dir, name), nil -} - -// profileFileNameFromConfig reads the profileFileName key (any casing) from -// the CLI config file, returning "" when the file or key is absent or -// unreadable. -func profileFileNameFromConfig(configPath string) string { - data, err := os.ReadFile(configPath) - if err != nil { - return "" - } - - var conf map[string]any - if err := yaml.Unmarshal(data, &conf); err != nil { - return "" - } - - for key, value := range conf { - if strings.EqualFold(key, "profilefilename") { - if name, ok := value.(string); ok { - return name - } - } - } - - return "" -} - -// acquireProfilesLock takes the CLI's config lock (~/.hatchet/config.lock) -// with the same protocol the CLI uses: exclusive create, retry every 50ms up -// to 100 attempts, and locks older than 5s are considered stale. -func acquireProfilesLock(dir string) (func(), error) { - lockFile := filepath.Join(dir, "config.lock") - - for attempts := 0; attempts < profilesLockMaxAttempts; attempts++ { - f, err := os.OpenFile(lockFile, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o644) - if err == nil { - _, writeErr := f.WriteString(time.Now().Format(time.RFC3339)) - _ = f.Close() - if writeErr != nil { - _ = os.Remove(lockFile) - return nil, fmt.Errorf("could not write lock file: %w", writeErr) - } - return func() { _ = os.Remove(lockFile) }, nil - } - - if os.IsExist(err) { - if stat, statErr := os.Stat(lockFile); statErr == nil && time.Since(stat.ModTime()) > profilesLockTimeout { - _ = os.Remove(lockFile) - continue - } - } else { - return nil, fmt.Errorf("could not create lock file: %w", err) - } - - time.Sleep(profilesLockRetryDelay) - } - - return nil, fmt.Errorf("could not acquire the profile store lock at %s", lockFile) -} - -// loadProfilesDocument parses the profiles file into a yaml document whose -// root is a mapping, creating an empty document when the file is absent or -// empty. Parse failures are returned rather than overwriting the file. -func loadProfilesDocument(path string) (doc *yaml.Node, root *yaml.Node, err error) { - data, err := os.ReadFile(path) - if err != nil && !os.IsNotExist(err) { - return nil, nil, fmt.Errorf("could not read %s: %w", path, err) - } - - if len(bytes.TrimSpace(data)) == 0 { - root = &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} - doc = &yaml.Node{Kind: yaml.DocumentNode, Content: []*yaml.Node{root}} - return doc, root, nil - } - - doc = &yaml.Node{} - if err := yaml.Unmarshal(data, doc); err != nil { - return nil, nil, fmt.Errorf("could not parse %s (leaving it untouched): %w", path, err) - } - - if len(doc.Content) == 0 { - root = &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} - doc = &yaml.Node{Kind: yaml.DocumentNode, Content: []*yaml.Node{root}} - return doc, root, nil - } - - root = doc.Content[0] - if root.Kind == yaml.ScalarNode && root.Tag == "!!null" { - *root = yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} - return doc, root, nil - } - if root.Kind != yaml.MappingNode { - return nil, nil, fmt.Errorf("%s has an unexpected structure (top level is not a mapping); leaving it untouched", path) - } - - return doc, root, nil -} - -// yamlNodeFor round-trips v through the yaml encoder to obtain its node -// representation, so values like time.Time serialize exactly as the encoder -// would write them (unquoted timestamps). -func yamlNodeFor(v any) (*yaml.Node, error) { - data, err := yaml.Marshal(v) - if err != nil { - return nil, err - } - - doc := &yaml.Node{} - if err := yaml.Unmarshal(data, doc); err != nil { - return nil, err - } - if len(doc.Content) == 0 { - return nil, fmt.Errorf("could not build a yaml node") - } - - return doc.Content[0], nil -} - -// mappingValue returns the value node for key in mapping m, matching keys -// case-insensitively (viper treats profile keys case-insensitively). -func mappingValue(m *yaml.Node, key string) *yaml.Node { - if m == nil || m.Kind != yaml.MappingNode { - return nil - } - - for i := 0; i+1 < len(m.Content); i += 2 { - if strings.EqualFold(m.Content[i].Value, key) { - return m.Content[i+1] - } - } - - return nil -} - -// mappingValueOrCreate returns the mapping value for key, creating an empty -// mapping entry when the key is absent or null. -func mappingValueOrCreate(m *yaml.Node, key string) (*yaml.Node, error) { - value := mappingValue(m, key) - if value == nil { - value = &yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} - m.Content = append(m.Content, - &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key}, - value, - ) - return value, nil - } - - if value.Kind == yaml.ScalarNode && value.Tag == "!!null" { - *value = yaml.Node{Kind: yaml.MappingNode, Tag: "!!map"} - return value, nil - } - if value.Kind != yaml.MappingNode { - return nil, fmt.Errorf("the %q key in the profiles file is not a mapping; leaving the file untouched", key) - } - - return value, nil -} - -// setMappingValue replaces the value for key in mapping m (case-insensitive), -// appending the pair when the key is absent. -func setMappingValue(m *yaml.Node, key string, value *yaml.Node) { - for i := 0; i+1 < len(m.Content); i += 2 { - if strings.EqualFold(m.Content[i].Value, key) { - m.Content[i+1] = value - return - } - } - - m.Content = append(m.Content, - &yaml.Node{Kind: yaml.ScalarNode, Tag: "!!str", Value: key}, - value, - ) -} - -// deleteMappingKey removes key (case-insensitive) and its value from mapping m. -func deleteMappingKey(m *yaml.Node, key string) bool { - for i := 0; i+1 < len(m.Content); i += 2 { - if strings.EqualFold(m.Content[i].Value, key) { - m.Content = append(m.Content[:i], m.Content[i+2:]...) - return true - } - } - - return false -} - -// writeYAMLAtomic writes doc to path via a same-directory temp file and -// rename, so concurrent readers never see a partial file. Two concurrent -// writers still race whole-file (last rename wins); with the config lock held -// this only matters against writers that ignore the lock, and the loser's -// write is a well-formed file. -func writeYAMLAtomic(path string, doc *yaml.Node) error { - var buf bytes.Buffer - enc := yaml.NewEncoder(&buf) - enc.SetIndent(4) - if err := enc.Encode(doc); err != nil { - _ = enc.Close() - return fmt.Errorf("could not encode the profiles file: %w", err) - } - if err := enc.Close(); err != nil { - return fmt.Errorf("could not encode the profiles file: %w", err) - } - - tmp := fmt.Sprintf("%s.%d.tmp", path, os.Getpid()) - if err := os.WriteFile(tmp, buf.Bytes(), 0o600); err != nil { - return fmt.Errorf("could not write %s: %w", tmp, err) - } - - if err := os.Rename(tmp, path); err != nil { - _ = os.Remove(tmp) - return fmt.Errorf("could not replace %s: %w", path, err) - } + removed, err := store.RemoveProfileIfTokenMatches(embeddedProfileName, token) - return nil + return removed, store.Path(), err } diff --git a/profileregistration_test.go b/profileregistration_test.go index 8eb4660..91c4111 100644 --- a/profileregistration_test.go +++ b/profileregistration_test.go @@ -4,14 +4,20 @@ import ( "os" "path/filepath" "strings" - "sync" "testing" "time" - "github.com/spf13/viper" - "gopkg.in/yaml.v3" + "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", @@ -29,35 +35,22 @@ func setTestHome(t *testing.T) string { return home } -func readProfilesFile(t *testing.T, home string) (string, map[string]any) { +// 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() - data, err := os.ReadFile(filepath.Join(home, ".hatchet", "profiles.yaml")) + store, err := profilestore.NewDefaultStore() if err != nil { - t.Fatalf("could not read profiles file: %v", err) - } - - var parsed map[string]any - if err := yaml.Unmarshal(data, &parsed); err != nil { - t.Fatalf("written profiles file does not parse: %v", err) + t.Fatalf("could not open the profile store: %v", err) } - return string(data), parsed -} - -func profileEntry(t *testing.T, parsed map[string]any, name string) map[string]any { - t.Helper() - - profiles, ok := parsed["profiles"].(map[string]any) - if !ok { - t.Fatalf("profiles key missing or not a mapping: %#v", parsed["profiles"]) - } - entry, ok := profiles[name].(map[string]any) - if !ok { - t.Fatalf("profile %q missing or not a mapping: %#v", name, profiles[name]) + profile, err := store.GetProfile(embeddedProfileName) + if err != nil { + t.Fatalf("could not read the embedded profile: %v", err) } - return entry + return profile } func TestRegisterEmbeddedProfileFreshFile(t *testing.T) { @@ -71,142 +64,34 @@ func TestRegisterEmbeddedProfileFreshFile(t *testing.T) { t.Fatalf("wrote to %s, want %s", path, want) } - raw, parsed := readProfilesFile(t, home) - entry := profileEntry(t, parsed, "embedded") - - for key, want := range map[string]any{ - "tenantid": "707d0855-80ab-4e1f-a156-f1c4546cbf52", - "name": "embedded", - "token": "token-fresh", - "apiserverurl": "http://localhost:28243", - "grpchostport": "127.0.0.1:50051", - "tlsstrategy": "none", - "embedded": true, - } { - if entry[key] != want { - t.Errorf("entry[%q] = %#v, want %#v", key, entry[key], want) - } - } - if _, ok := entry["pid"].(int); !ok { - t.Errorf("entry[pid] = %#v, want an int", entry["pid"]) - } - - // The CLI decodes expiresat into a time.Time; a quoted string there would - // break its unmarshalling, so the timestamp must be written unquoted. - if strings.Contains(raw, `expiresat: "`) || strings.Contains(raw, "expiresat: '") { - t.Errorf("expiresat was written quoted:\n%s", raw) - } - if _, ok := entry["expiresat"].(time.Time); !ok { - t.Errorf("expiresat did not parse as a yaml timestamp: %#v", entry["expiresat"]) - } -} - -// TestRegisteredProfileReadableByViper anchors interop with the CLI: the CLI -// loads the profiles file via viper (lenient unmarshal into a struct with -// time.Time fields) and per-key lookups, so both must work on our output. -func TestRegisteredProfileReadableByViper(t *testing.T) { - home := setTestHome(t) - - if _, err := registerEmbeddedProfile(testRegistration("token-viper")); err != nil { - t.Fatalf("register failed: %v", err) - } - - v := viper.New() - v.SetConfigFile(filepath.Join(home, ".hatchet", "profiles.yaml")) - v.SetConfigType("yaml") - if err := v.ReadInConfig(); err != nil { - t.Fatalf("viper could not read the profiles file: %v", err) - } - - if got := v.GetString("profiles.embedded.token"); got != "token-viper" { - t.Errorf("viper token = %q, want token-viper", got) - } - if got := v.GetTime("profiles.embedded.expiresat"); got.IsZero() { - t.Errorf("viper expiresat is zero") - } - - // Mirrors pkg/config/cli.ProfileFile in the CLI: unknown metadata fields - // must not break the struct unmarshal. - type profile struct { - TenantId string `mapstructure:"tenantId"` - Name string `mapstructure:"name"` - Token string `mapstructure:"token"` - ExpiresAt time.Time `mapstructure:"expiresAt"` - ApiServerURL string `mapstructure:"apiServerURL"` - GrpcHostPort string `mapstructure:"grpcHostPort"` - TLSStrategy string `mapstructure:"tlsStrategy"` - } - var file struct { - Profiles map[string]profile `mapstructure:"profiles"` - } - if err := v.Unmarshal(&file); err != nil { - t.Fatalf("viper struct unmarshal failed (the released CLI would fatal on this): %v", err) - } - - got := file.Profiles["embedded"] - if got.Token != "token-viper" || got.TLSStrategy != "none" || got.ExpiresAt.IsZero() { - t.Errorf("unmarshalled profile = %+v", got) - } -} - -func TestRegisterPreservesExistingContent(t *testing.T) { - home := setTestHome(t) - - existing := `# keep this comment -defaultprofile: prod -profiles: - prod: - apiserverurl: https://prod.example.com - expiresat: 2027-01-02T03:04:05Z - grpchostport: prod.example.com:443 - name: prod - tenantid: 11111111-2222-3333-4444-555555555555 - tlsstrategy: tls - token: prod-token -` - dir := filepath.Join(home, ".hatchet") - if err := os.MkdirAll(dir, 0o700); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(filepath.Join(dir, "profiles.yaml"), []byte(existing), 0o600); err != nil { - t.Fatal(err) + 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 _, err := registerEmbeddedProfile(testRegistration("token-merge")); err != nil { - t.Fatalf("register failed: %v", err) + if profile.ExpiresAt.IsZero() { + t.Errorf("expiresat did not round-trip through the store") } - raw, parsed := readProfilesFile(t, home) - - if parsed["defaultprofile"] != "prod" { - t.Errorf("defaultprofile = %#v, want prod", parsed["defaultprofile"]) - } - if !strings.Contains(raw, "# keep this comment") { - t.Errorf("comment was not preserved:\n%s", raw) + // 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) } - - prod := profileEntry(t, parsed, "prod") - for key, want := range map[string]any{ - "apiserverurl": "https://prod.example.com", - "grpchostport": "prod.example.com:443", - "name": "prod", - "tenantid": "11111111-2222-3333-4444-555555555555", - "tlsstrategy": "tls", - "token": "prod-token", - } { - if prod[key] != want { - t.Errorf("prod[%q] = %#v, want %#v", key, prod[key], want) + 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) } } - - embedded := profileEntry(t, parsed, "embedded") - if embedded["token"] != "token-merge" { - t.Errorf("embedded token = %#v", embedded["token"]) - } } func TestRegisterOverwritesPreviousRegistration(t *testing.T) { - home := setTestHome(t) + setTestHome(t) if _, err := registerEmbeddedProfile(testRegistration("token-old")); err != nil { t.Fatalf("first register failed: %v", err) @@ -215,53 +100,8 @@ func TestRegisterOverwritesPreviousRegistration(t *testing.T) { t.Fatalf("second register failed: %v", err) } - _, parsed := readProfilesFile(t, home) - entry := profileEntry(t, parsed, "embedded") - if entry["token"] != "token-new" { - t.Errorf("embedded token = %#v, want token-new (last writer wins)", entry["token"]) - } -} - -// TestReRegistrationAfterDeregistrationStaysCLIReadable guards against a flow -// style trap: a deregistration that empties the store leaves "profiles: {}" -// behind, and a later registration merged into that flow mapping would be -// emitted in flow style with quoted timestamps, which the CLI cannot decode -// into its time.Time fields. -func TestReRegistrationAfterDeregistrationStaysCLIReadable(t *testing.T) { - home := setTestHome(t) - - if _, err := registerEmbeddedProfile(testRegistration("token-first")); err != nil { - t.Fatalf("register failed: %v", err) - } - if _, _, err := deregisterEmbeddedProfile("token-first"); err != nil { - t.Fatalf("deregister failed: %v", err) - } - if _, err := registerEmbeddedProfile(testRegistration("token-second")); err != nil { - t.Fatalf("re-register failed: %v", err) - } - - raw, _ := readProfilesFile(t, home) - if strings.Contains(raw, "{") { - t.Errorf("profiles file was written in flow style:\n%s", raw) - } - - v := viper.New() - v.SetConfigFile(filepath.Join(home, ".hatchet", "profiles.yaml")) - v.SetConfigType("yaml") - if err := v.ReadInConfig(); err != nil { - t.Fatalf("viper could not read the profiles file: %v", err) - } - var file struct { - Profiles map[string]struct { - ExpiresAt time.Time `mapstructure:"expiresAt"` - Token string `mapstructure:"token"` - } `mapstructure:"profiles"` - } - if err := v.Unmarshal(&file); err != nil { - t.Fatalf("viper struct unmarshal failed (the released CLI would fatal on this): %v", err) - } - if file.Profiles["embedded"].Token != "token-second" || file.Profiles["embedded"].ExpiresAt.IsZero() { - t.Errorf("unexpected re-registered profile: %+v", file.Profiles["embedded"]) + if got := readEmbeddedProfile(t).Token; got != "token-new" { + t.Errorf("embedded token = %q, want token-new (last writer wins)", got) } } @@ -286,7 +126,7 @@ func TestRegisterPermissionDeniedIsAnError(t *testing.T) { } func TestDeregisterRemovesMatchingToken(t *testing.T) { - home := setTestHome(t) + setTestHome(t) if _, err := registerEmbeddedProfile(testRegistration("token-mine")); err != nil { t.Fatalf("register failed: %v", err) @@ -300,16 +140,17 @@ func TestDeregisterRemovesMatchingToken(t *testing.T) { t.Fatal("expected the registration to be removed") } - _, parsed := readProfilesFile(t, home) - if profiles, ok := parsed["profiles"].(map[string]any); ok { - if _, still := profiles["embedded"]; still { - t.Errorf("embedded profile still present after deregistration") - } + 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) { - home := setTestHome(t) + setTestHome(t) // A newer instance overwrote the registration; the older instance's // shutdown must not delete it. @@ -325,10 +166,8 @@ func TestDeregisterKeepsNewerRegistration(t *testing.T) { t.Fatal("deregistration removed another instance's registration") } - _, parsed := readProfilesFile(t, home) - entry := profileEntry(t, parsed, "embedded") - if entry["token"] != "token-newer" { - t.Errorf("embedded token = %#v, want token-newer", entry["token"]) + if got := readEmbeddedProfile(t).Token; got != "token-newer" { + t.Errorf("embedded token = %q, want token-newer", got) } } @@ -343,37 +182,3 @@ func TestDeregisterWithoutFileIsANoOp(t *testing.T) { t.Fatal("nothing to remove, but removed was reported") } } - -func TestConcurrentRegisterDeregisterKeepsFileWellFormed(t *testing.T) { - home := setTestHome(t) - - var wg sync.WaitGroup - for i := 0; i < 8; i++ { - wg.Add(1) - go func(n int) { - defer wg.Done() - token := "token-a" - if n%2 == 1 { - token = "token-b" - } - if _, err := registerEmbeddedProfile(testRegistration(token)); err != nil { - t.Errorf("register failed: %v", err) - } - if _, _, err := deregisterEmbeddedProfile("token-a"); err != nil { - t.Errorf("deregister failed: %v", err) - } - }(i) - } - wg.Wait() - - // Whatever interleaving happened, the file must parse and contain either - // no embedded profile or a well-formed one. - _, parsed := readProfilesFile(t, home) - if profiles, ok := parsed["profiles"].(map[string]any); ok { - if raw, present := profiles["embedded"]; present { - if _, ok := raw.(map[string]any); !ok { - t.Errorf("embedded entry is malformed: %#v", raw) - } - } - } -}