diff --git a/internal/config/agent_id_test.go b/internal/config/agent_id_test.go new file mode 100644 index 0000000000000000000000000000000000000000..74bad7f563dd1aa4c5f535f43c2204aacb1930b0 --- /dev/null +++ b/internal/config/agent_id_test.go @@ -0,0 +1,29 @@ +package config + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestConfig_AgentIDs(t *testing.T) { + cfg := &Config{ + Options: &Options{ + DisabledTools: []string{}, + }, + } + cfg.SetupAgents() + + t.Run("Coder agent should have correct ID", func(t *testing.T) { + coderAgent, ok := cfg.Agents[AgentCoder] + require.True(t, ok) + assert.Equal(t, AgentCoder, coderAgent.ID, "Coder agent ID should be '%s'", AgentCoder) + }) + + t.Run("Task agent should have correct ID", func(t *testing.T) { + taskAgent, ok := cfg.Agents[AgentTask] + require.True(t, ok) + assert.Equal(t, AgentTask, taskAgent.ID, "Task agent ID should be '%s'", AgentTask) + }) +} diff --git a/internal/config/atomicwrite.go b/internal/config/atomicwrite.go new file mode 100644 index 0000000000000000000000000000000000000000..6dc6d142cad22ce6c9090e4c7db5f39a1e968fd5 --- /dev/null +++ b/internal/config/atomicwrite.go @@ -0,0 +1,63 @@ +package config + +import ( + "os" + "path/filepath" + "time" +) + +// atomicWriteFile writes data to a file atomically by writing to a unique +// temporary file in the same directory and renaming it into place. This +// prevents concurrent readers from observing a partially-written file. +func atomicWriteFile(path string, data []byte, perm os.FileMode) error { + path = filepath.Clean(path) + dir := filepath.Dir(path) + f, err := os.CreateTemp(dir, filepath.Base(path)+".*.tmp") + if err != nil { + return err + } + tmp := f.Name() + if _, err := f.Write(data); err != nil { + f.Close() + os.Remove(tmp) + return err + } + if err := f.Chmod(perm); err != nil { + f.Close() + os.Remove(tmp) + return err + } + if err := f.Close(); err != nil { + os.Remove(tmp) + return err + } + if err := renameFile(tmp, path); err != nil { + os.Remove(tmp) + return err + } + return nil +} + +// renameRetryBudget bounds how long renameFile keeps retrying transient +// failures before giving up and returning the error. +const renameRetryBudget = 2 * time.Second + +// renameFile renames tmp over path. On Windows the rename fails with +// ERROR_ACCESS_DENIED or ERROR_SHARING_VIOLATION while another process +// (antivirus, search indexer) or a concurrent reader briefly holds a +// handle on the destination, so transient failures are retried with +// backoff. On other platforms isTransientRenameError is always false +// and this is a plain os.Rename. +func renameFile(tmp, path string) error { + var slept time.Duration + delay := time.Millisecond + for { + err := os.Rename(tmp, path) + if err == nil || !isTransientRenameError(err) || slept >= renameRetryBudget { + return err + } + time.Sleep(delay) + slept += delay + delay = min(delay*2, 50*time.Millisecond) + } +} diff --git a/internal/config/atomicwrite_test.go b/internal/config/atomicwrite_test.go new file mode 100644 index 0000000000000000000000000000000000000000..089ee5c78eae7fecdb5bc3cd3d8f74713098fc38 --- /dev/null +++ b/internal/config/atomicwrite_test.go @@ -0,0 +1,43 @@ +package config + +import ( + "os" + "path/filepath" + "runtime" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAtomicWriteFile(t *testing.T) { + t.Parallel() + dir := t.TempDir() + path := filepath.Join(dir, "test.json") + + require.NoError(t, atomicWriteFile(path, []byte(`{"key":"value"}`), 0o600)) + + data, err := os.ReadFile(path) + require.NoError(t, err) + require.Equal(t, `{"key":"value"}`, string(data)) + + // No temp files should linger. + entries, err := os.ReadDir(dir) + require.NoError(t, err) + require.Len(t, entries, 1) + require.Equal(t, "test.json", entries[0].Name()) +} + +func TestAtomicWriteFile_PermissionsApplied(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("Windows does not support Unix file permissions") + } + t.Parallel() + dir := t.TempDir() + path := filepath.Join(dir, "test.json") + + require.NoError(t, atomicWriteFile(path, []byte(`{}`), 0o600)) + + info, err := os.Stat(path) + require.NoError(t, err) + require.Equal(t, os.FileMode(0o600), info.Mode().Perm()) +} diff --git a/internal/config/atomicwrite_unix.go b/internal/config/atomicwrite_unix.go new file mode 100644 index 0000000000000000000000000000000000000000..da7c7645c32939bc15d0a88a840ee17519627240 --- /dev/null +++ b/internal/config/atomicwrite_unix.go @@ -0,0 +1,7 @@ +//go:build !windows + +package config + +// isTransientRenameError reports whether err is a rename failure that +// can resolve on its own. Only Windows has such failures. +func isTransientRenameError(error) bool { return false } diff --git a/internal/config/atomicwrite_windows_test.go b/internal/config/atomicwrite_windows_test.go new file mode 100644 index 0000000000000000000000000000000000000000..14d2e2d2f2426a76aa1fedc6fbdaf2e8c528b0d5 --- /dev/null +++ b/internal/config/atomicwrite_windows_test.go @@ -0,0 +1,45 @@ +//go:build windows + +package config + +import ( + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" + "golang.org/x/sys/windows" +) + +// Regression test for the windows-latest CI flake where +// TestConfigStore_SetConfigFields_concurrentInProcess failed with +// "rename ...crush.json..tmp ...crush.json: Access is denied": +// renaming over a destination while another handle is open on it +// without sharing must be retried until the handle closes, not +// surfaced as an error. +func TestAtomicWriteFile_RetriesWhileDestinationHandleOpen(t *testing.T) { + t.Parallel() + dir := t.TempDir() + path := filepath.Join(dir, "test.json") + require.NoError(t, atomicWriteFile(path, []byte(`{"v":1}`), 0o600)) + + // Hold the destination open with no sharing flags so the rename + // fails with ERROR_ACCESS_DENIED/ERROR_SHARING_VIOLATION until the + // handle closes, mimicking an antivirus scan or exclusive reader. + p, err := windows.UTF16PtrFromString(path) + require.NoError(t, err) + h, err := windows.CreateFile(p, windows.GENERIC_READ, 0, nil, + windows.OPEN_EXISTING, windows.FILE_ATTRIBUTE_NORMAL, 0) + require.NoError(t, err) + go func() { + time.Sleep(200 * time.Millisecond) + windows.CloseHandle(h) + }() + + require.NoError(t, atomicWriteFile(path, []byte(`{"v":2}`), 0o600)) + + data, err := os.ReadFile(path) + require.NoError(t, err) + require.JSONEq(t, `{"v":2}`, string(data)) +} diff --git a/internal/config/attribution_migration_test.go b/internal/config/attribution_migration_test.go new file mode 100644 index 0000000000000000000000000000000000000000..cc8c9e2e278b89ee2f86996975774abe18413843 --- /dev/null +++ b/internal/config/attribution_migration_test.go @@ -0,0 +1,93 @@ +package config + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAttributionMigration(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + configJSON string + expectedTrailer TrailerStyle + expectedGenerate bool + }{ + { + name: "old setting co_authored_by=true migrates to co-authored-by", + configJSON: `{ + "options": { + "attribution": { + "co_authored_by": true, + "generated_with": false + } + } + }`, + expectedTrailer: TrailerStyleCoAuthoredBy, + expectedGenerate: false, + }, + { + name: "old setting co_authored_by=false migrates to none", + configJSON: `{ + "options": { + "attribution": { + "co_authored_by": false, + "generated_with": true + } + } + }`, + expectedTrailer: TrailerStyleNone, + expectedGenerate: true, + }, + { + name: "new setting takes precedence over old setting", + configJSON: `{ + "options": { + "attribution": { + "trailer_style": "assisted-by", + "co_authored_by": true, + "generated_with": false + } + } + }`, + expectedTrailer: TrailerStyleAssistedBy, + expectedGenerate: false, + }, + { + name: "default when neither setting present", + configJSON: `{ + "options": { + "attribution": { + "generated_with": true + } + } + }`, + expectedTrailer: TrailerStyleAssistedBy, + expectedGenerate: true, + }, + { + name: "default when attribution is null", + configJSON: `{ + "options": {} + }`, + expectedTrailer: TrailerStyleAssistedBy, + expectedGenerate: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + cfg, err := loadFromBytes([][]byte{[]byte(tt.configJSON)}) + require.NoError(t, err) + + cfg.setDefaults(t.TempDir(), "") + + require.Equal(t, tt.expectedTrailer, cfg.Options.Attribution.TrailerStyle) + require.Equal(t, tt.expectedGenerate, cfg.Options.Attribution.GeneratedWith) + }) + } +} diff --git a/internal/config/config_unix.go b/internal/config/config_unix.go new file mode 100644 index 0000000000000000000000000000000000000000..b6d372f609c8529bcd5bc2792eca0194b199a25a --- /dev/null +++ b/internal/config/config_unix.go @@ -0,0 +1,7 @@ +//go:build !windows + +package config + +// systemConfigPath is the system-wide configuration file path. It is +// loaded at the lowest priority so user and project configs override it. +const systemConfigPath = "/etc/crush/crush.json" diff --git a/internal/config/copilot.go b/internal/config/copilot.go new file mode 100644 index 0000000000000000000000000000000000000000..d912156bec00a9f00850ab2ec3a3baf1016c2141 --- /dev/null +++ b/internal/config/copilot.go @@ -0,0 +1 @@ +package config diff --git a/internal/config/docker_mcp.go b/internal/config/docker_mcp.go new file mode 100644 index 0000000000000000000000000000000000000000..3da11ff86a4a8689a10edef37fbfadb354dc8d0d --- /dev/null +++ b/internal/config/docker_mcp.go @@ -0,0 +1,141 @@ +package config + +import ( + "context" + "fmt" + "os/exec" + "sync" + "time" +) + +var dockerMCPVersionRunner = func(ctx context.Context) error { + cmd := exec.CommandContext(ctx, "docker", "mcp", "version") + return cmd.Run() +} + +const dockerMCPAvailabilityTTL = 10 * time.Second + +var dockerMCPAvailabilityCache struct { + mu sync.Mutex + available bool + checkedAt time.Time + known bool +} + +// DockerMCPName is the name of the Docker MCP configuration. +const DockerMCPName = "docker" + +// IsDockerMCPAvailable checks if Docker MCP is available by running +// 'docker mcp version'. +func IsDockerMCPAvailable() bool { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + err := dockerMCPVersionRunner(ctx) + return err == nil +} + +// DockerMCPAvailabilityCached returns the cached Docker MCP availability and +// whether the cached value is still fresh. +func DockerMCPAvailabilityCached() (available bool, known bool) { + dockerMCPAvailabilityCache.mu.Lock() + defer dockerMCPAvailabilityCache.mu.Unlock() + + if !dockerMCPAvailabilityCache.known { + return false, false + } + if time.Since(dockerMCPAvailabilityCache.checkedAt) > dockerMCPAvailabilityTTL { + return dockerMCPAvailabilityCache.available, false + } + return dockerMCPAvailabilityCache.available, true +} + +// RefreshDockerMCPAvailability refreshes and caches Docker MCP availability. +func RefreshDockerMCPAvailability() bool { + available := IsDockerMCPAvailable() + dockerMCPAvailabilityCache.mu.Lock() + dockerMCPAvailabilityCache.available = available + dockerMCPAvailabilityCache.checkedAt = time.Now() + dockerMCPAvailabilityCache.known = true + dockerMCPAvailabilityCache.mu.Unlock() + return available +} + +// IsDockerMCPEnabled checks if Docker MCP is already configured. +func (c *Config) IsDockerMCPEnabled() bool { + if c.MCP == nil { + return false + } + _, exists := c.MCP[DockerMCPName] + return exists +} + +// DockerMCPConfig returns the default Docker MCP stdio configuration. +func DockerMCPConfig() MCPConfig { + return MCPConfig{ + Type: MCPStdio, + Command: "docker", + Args: []string{"mcp", "gateway", "run"}, + Disabled: false, + } +} + +// PrepareDockerMCPConfig validates Docker MCP availability and stages the +// Docker MCP configuration in memory. +func (s *ConfigStore) PrepareDockerMCPConfig() (MCPConfig, error) { + if !IsDockerMCPAvailable() { + return MCPConfig{}, fmt.Errorf("docker mcp is not available, please ensure docker is installed and 'docker mcp version' succeeds") + } + + mcpConfig := DockerMCPConfig() + // In-memory only; persistence happens in PersistDockerMCPConfig. + s.mutateInMemory(func(c *Config) { + if c.MCP == nil { + c.MCP = make(map[string]MCPConfig) + } + c.MCP[DockerMCPName] = mcpConfig + }) + return mcpConfig, nil +} + +// PersistDockerMCPConfig persists a previously prepared Docker MCP +// configuration to the global config file. +func (s *ConfigStore) PersistDockerMCPConfig(mcpConfig MCPConfig) error { + if err := s.SetConfigField(ScopeGlobal, "mcp."+DockerMCPName, mcpConfig); err != nil { + return fmt.Errorf("failed to persist docker mcp configuration: %w", err) + } + return nil +} + +// EnableDockerMCP adds Docker MCP configuration and persists it. +func (s *ConfigStore) EnableDockerMCP() error { + mcpConfig, err := s.PrepareDockerMCPConfig() + if err != nil { + return err + } + if err := s.PersistDockerMCPConfig(mcpConfig); err != nil { + return err + } + return nil +} + +// DisableDockerMCP removes Docker MCP configuration and persists the change. +func (s *ConfigStore) DisableDockerMCP() error { + return s.update(ScopeGlobal, func(c *Config) map[string]any { + if c.MCP == nil { + return nil + } + delete(c.MCP, DockerMCPName) + return map[string]any{"mcp": c.MCP} + }) +} + +// RemoveDockerMCPInMemory removes the Docker MCP entry from the in-memory +// config via copy-on-write, without persisting to disk. It rolls back a +// staged PrepareDockerMCPConfig when starting or persisting the server +// fails, so callers must not mutate Config().MCP directly. +func (s *ConfigStore) RemoveDockerMCPInMemory() { + s.mutateInMemory(func(c *Config) { + delete(c.MCP, DockerMCPName) + }) +} diff --git a/internal/config/docker_mcp_test.go b/internal/config/docker_mcp_test.go new file mode 100644 index 0000000000000000000000000000000000000000..54690a8a7517740fc58fcdea87708a27964072ff --- /dev/null +++ b/internal/config/docker_mcp_test.go @@ -0,0 +1,195 @@ +package config + +import ( + "context" + "errors" + "os" + "path/filepath" + "testing" + + "github.com/charmbracelet/crush/internal/env" + "github.com/stretchr/testify/require" +) + +var errDockerUnavailable = errors.New("docker unavailable") + +func setDockerMCPVersionRunner(t *testing.T, runner func(context.Context) error) { + t.Helper() + orig := dockerMCPVersionRunner + dockerMCPVersionRunner = runner + t.Cleanup(func() { + dockerMCPVersionRunner = orig + }) +} + +func TestIsDockerMCPEnabled(t *testing.T) { + t.Parallel() + + t.Run("returns false when MCP is nil", func(t *testing.T) { + t.Parallel() + cfg := &Config{ + MCP: nil, + } + require.False(t, cfg.IsDockerMCPEnabled()) + }) + + t.Run("returns false when docker mcp not configured", func(t *testing.T) { + t.Parallel() + cfg := &Config{ + MCP: make(map[string]MCPConfig), + } + require.False(t, cfg.IsDockerMCPEnabled()) + }) + + t.Run("returns true when docker mcp is configured", func(t *testing.T) { + t.Parallel() + cfg := &Config{ + MCP: map[string]MCPConfig{ + DockerMCPName: { + Type: MCPStdio, + Command: "docker", + }, + }, + } + require.True(t, cfg.IsDockerMCPEnabled()) + }) +} + +func TestEnableDockerMCP(t *testing.T) { + t.Run("adds docker mcp to config", func(t *testing.T) { + setDockerMCPVersionRunner(t, func(context.Context) error { return nil }) + + // Create a temporary directory for config. + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "crush.json") + + cfg := &Config{ + MCP: make(map[string]MCPConfig), + } + store := &ConfigStore{ + config: cfg, + globalDataPath: configPath, + resolver: NewShellVariableResolver(env.New()), + } + + err := store.EnableDockerMCP() + require.NoError(t, err) + + // Check in-memory config via the store (copy-on-write publishes + // a new Config; the captured cfg pointer stays unchanged). + require.True(t, store.Config().IsDockerMCPEnabled()) + mcpConfig, exists := store.Config().MCP[DockerMCPName] + require.True(t, exists) + require.Equal(t, MCPStdio, mcpConfig.Type) + require.Equal(t, "docker", mcpConfig.Command) + require.Equal(t, []string{"mcp", "gateway", "run"}, mcpConfig.Args) + require.False(t, mcpConfig.Disabled) + + // Check persisted config. + data, err := os.ReadFile(configPath) + require.NoError(t, err) + require.Contains(t, string(data), "docker") + require.Contains(t, string(data), "gateway") + }) + + t.Run("fails when docker mcp not available", func(t *testing.T) { + setDockerMCPVersionRunner(t, func(context.Context) error { return errDockerUnavailable }) + + // Create a temporary directory for config. + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "crush.json") + + cfg := &Config{ + MCP: make(map[string]MCPConfig), + } + store := &ConfigStore{ + config: cfg, + globalDataPath: configPath, + resolver: NewShellVariableResolver(env.New()), + } + + err := store.EnableDockerMCP() + require.Error(t, err) + require.Contains(t, err.Error(), "docker mcp is not available") + }) +} + +func TestDisableDockerMCP(t *testing.T) { + t.Parallel() + + t.Run("removes docker mcp from config", func(t *testing.T) { + t.Parallel() + + // Create a temporary directory for config. + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "crush.json") + + cfg := &Config{ + MCP: map[string]MCPConfig{ + DockerMCPName: { + Type: MCPStdio, + Command: "docker", + Args: []string{"mcp", "gateway", "run"}, + Disabled: false, + }, + }, + } + store := &ConfigStore{ + config: cfg, + globalDataPath: configPath, + resolver: NewShellVariableResolver(env.New()), + } + + // Verify it's enabled first. + require.True(t, cfg.IsDockerMCPEnabled()) + + err := store.DisableDockerMCP() + require.NoError(t, err) + + // Check in-memory config via the store (copy-on-write publishes a + // new Config; the original cfg pointer is intentionally unchanged). + require.False(t, store.Config().IsDockerMCPEnabled()) + _, exists := store.Config().MCP[DockerMCPName] + require.False(t, exists) + }) + + t.Run("does nothing when MCP is nil", func(t *testing.T) { + t.Parallel() + + cfg := &Config{ + MCP: nil, + } + store := &ConfigStore{ + config: cfg, + globalDataPath: filepath.Join(t.TempDir(), "crush.json"), + resolver: NewShellVariableResolver(env.New()), + } + + err := store.DisableDockerMCP() + require.NoError(t, err) + }) +} + +func TestEnableDockerMCPWithRealDockerWhenAvailable(t *testing.T) { + t.Parallel() + + if !IsDockerMCPAvailable() { + t.Skip("docker mcp not available on this machine") + } + + tmpDir := t.TempDir() + configPath := filepath.Join(tmpDir, "crush.json") + + cfg := &Config{ + MCP: make(map[string]MCPConfig), + } + store := &ConfigStore{ + config: cfg, + globalDataPath: configPath, + resolver: NewShellVariableResolver(env.New()), + } + + err := store.EnableDockerMCP() + require.NoError(t, err) + require.True(t, store.Config().IsDockerMCPEnabled()) +} diff --git a/internal/config/init.go b/internal/config/init.go new file mode 100644 index 0000000000000000000000000000000000000000..6138c49d496b16d054ea9ff3f6c49906b3f433ca --- /dev/null +++ b/internal/config/init.go @@ -0,0 +1,127 @@ +package config + +import ( + "fmt" + "os" + "path/filepath" + "slices" + "strings" + + "github.com/charmbracelet/crush/internal/fsext" +) + +const ( + InitFlagFilename = "init" +) + +type ProjectInitFlag struct { + Initialized bool `json:"initialized"` +} + +func Init(workingDir, dataDir string, debug bool) (*ConfigStore, error) { + store, err := Load(workingDir, dataDir, debug) + if err != nil { + return nil, err + } + return store, nil +} + +func ProjectNeedsInitialization(store *ConfigStore) (bool, error) { + if store == nil { + return false, fmt.Errorf("config not loaded") + } + + cfg := store.Config() + flagFilePath := filepath.Join(cfg.Options.DataDirectory, InitFlagFilename) + + _, err := os.Stat(flagFilePath) + if err == nil { + return false, nil + } + + if !os.IsNotExist(err) { + return false, fmt.Errorf("failed to check init flag file: %w", err) + } + + someContextFileExists, err := contextPathsExist(store.WorkingDir()) + if err != nil { + return false, fmt.Errorf("failed to check for context files: %w", err) + } + if someContextFileExists { + return false, nil + } + + // If the working directory has no non-ignored files, skip initialization step + empty, err := dirHasNoVisibleFiles(store.WorkingDir()) + if err != nil { + return false, fmt.Errorf("failed to check if directory is empty: %w", err) + } + if empty { + return false, nil + } + + return true, nil +} + +func contextPathsExist(dir string) (bool, error) { + entries, err := os.ReadDir(dir) + if err != nil { + return false, err + } + + // Create a slice of lowercase filenames for lookup with slices.Contains + var files []string + for _, entry := range entries { + if !entry.IsDir() { + files = append(files, strings.ToLower(entry.Name())) + } + } + + // Check if any of the default context paths exist in the directory + for _, path := range defaultContextPaths { + // Extract just the filename from the path + _, filename := filepath.Split(path) + filename = strings.ToLower(filename) + + if slices.Contains(files, filename) { + return true, nil + } + } + + return false, nil +} + +// dirHasNoVisibleFiles returns true if the directory has no files/dirs after applying ignore rules. +func dirHasNoVisibleFiles(dir string) (bool, error) { + files, _, err := fsext.ListDirectory(dir, nil, 1, 1) + if err != nil { + return false, err + } + return len(files) == 0, nil +} + +func MarkProjectInitialized(store *ConfigStore) error { + if store == nil { + return fmt.Errorf("config not loaded") + } + flagFilePath := filepath.Join(store.Config().Options.DataDirectory, InitFlagFilename) + + file, err := os.Create(flagFilePath) + if err != nil { + return fmt.Errorf("failed to create init flag file: %w", err) + } + defer file.Close() + + return nil +} + +func HasInitialDataConfig(store *ConfigStore) bool { + if store == nil { + return false + } + cfgPath := GlobalConfigData() + if _, err := os.Stat(cfgPath); err != nil { + return false + } + return store.Config().IsConfigured() +} diff --git a/internal/config/load.go b/internal/config/load.go new file mode 100644 index 0000000000000000000000000000000000000000..3b4b59a770f07445dedd89c0f86778ef7115cab6 --- /dev/null +++ b/internal/config/load.go @@ -0,0 +1,1454 @@ +package config + +import ( + "cmp" + "context" + "encoding/json" + "fmt" + "log/slog" + "maps" + "os" + "os/exec" + "path/filepath" + "regexp" + "runtime" + "slices" + "strconv" + "strings" + "sync" + "testing" + "time" + + "charm.land/catwalk/pkg/catwalk" + "github.com/charmbracelet/crush/internal/agent/hyper" + "github.com/charmbracelet/crush/internal/csync" + "github.com/charmbracelet/crush/internal/discover" + "github.com/charmbracelet/crush/internal/env" + "github.com/charmbracelet/crush/internal/filepathext" + "github.com/charmbracelet/crush/internal/fsext" + "github.com/charmbracelet/crush/internal/home" + "github.com/charmbracelet/crush/internal/shellconfig" + powernapConfig "github.com/charmbracelet/x/powernap/pkg/config" + "github.com/qjebbs/go-jsons" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +const defaultCatwalkURL = "https://catwalk.charm.land" + +// Load loads the configuration from the default paths and returns a +// ConfigStore that owns both the pure-data Config and all runtime state. +func Load(workingDir, dataDir string, debug bool) (*ConfigStore, error) { + // Migrate deprecated disable_notifications before loading config. + migrateDisableNotifications() + + configPaths := lookupConfigs(workingDir) + + cfg, loadedPaths, err := loadFromConfigPaths(context.Background(), configPaths) + if err != nil { + return nil, fmt.Errorf("failed to load config from paths %v: %w", configPaths, err) + } + + cfg.setDefaults(workingDir, dataDir) + + store := &ConfigStore{ + config: cfg, + workingDir: workingDir, + globalDataPath: GlobalConfigData(), + workspacePath: filepath.Join(cfg.Options.DataDirectory, fmt.Sprintf("%s.json", appName)), + loadedPaths: loadedPaths, + } + + if debug { + cfg.Options.Debug = true + } + + // Load workspace config last so it has highest priority. + if wsData, err := os.ReadFile(store.workspacePath); err == nil && len(wsData) > 0 { + if !json.Valid(wsData) { + return nil, fmt.Errorf("invalid JSON in config file %s", store.workspacePath) + } + merged, mergeErr := loadFromBytes(append([][]byte{mustMarshalConfig(cfg)}, wsData)) + if mergeErr == nil { + // Preserve defaults that setDefaults already applied. + dataDir := cfg.Options.DataDirectory + *cfg = *merged + cfg.setDefaults(workingDir, dataDir) + store.config = cfg + store.loadedPaths = append(store.loadedPaths, store.workspacePath) + } + } + + // Validate hooks after all config merging is complete so workspace + // hooks also get their matcher regexes compiled. + if err := cfg.ValidateHooks(); err != nil { + return nil, fmt.Errorf("invalid hook configuration: %w", err) + } + + if !isInsideWorktree() { + const depth = 2 + const items = 100 + slog.Warn("No git repository detected in working directory, will limit file walk operations", "depth", depth, "items", items) + assignIfNil(&cfg.Tools.Ls.MaxDepth, depth) + assignIfNil(&cfg.Tools.Ls.MaxItems, items) + assignIfNil(&cfg.Options.TUI.Completions.MaxDepth, depth) + assignIfNil(&cfg.Options.TUI.Completions.MaxItems, items) + } + + if isAppleTerminal() { + slog.Warn("Detected Apple Terminal, enabling transparent mode") + assignIfNil(&cfg.Options.TUI.Transparent, true) + } + + // Load known providers, this loads the config from catwalk. A failed + // refresh still yields the cached or embedded catalog, so only an empty + // list is fatal: starting up without providers is worse than starting + // up with slightly stale ones. Pass a Hyper token refresher so the + // catalog fetch can retry on 401. + providers, err := Providers(cfg, func(ctx context.Context) error { + return store.RefreshOAuthToken(ctx, ScopeGlobal, "hyper") + }) + if err != nil { + if len(providers) == 0 { + return nil, err + } + slog.Warn("Continuing with the previously known providers", "error", err) + } + store.knownProviders = providers + + env := env.New() + // Configure providers + valueResolver := NewShellVariableResolver(env) + store.resolver = valueResolver + + // Hold writeMu during initial load to prevent configureProviders + // from triggering auto-reload via RemoveConfigField. + store.writeMu.Lock() + defer store.writeMu.Unlock() + + // Apply top-level env vars before configuring providers so variables + // like AWS_PROFILE are visible to the AWS SDK credential chain. + cfg.applyEnv(valueResolver) + + if err := cfg.configureProviders(context.Background(), store, env, valueResolver, store.knownProviders); err != nil { + return nil, fmt.Errorf("failed to configure providers: %w", err) + } + + if !cfg.IsConfigured() { + slog.Warn("No providers configured") + return store, nil + } + + resolved, err := resolveSelectedModels(cfg, store.knownProviders) + if err != nil { + return nil, fmt.Errorf("failed to configure selected models: %w", err) + } + cfg.Models[SelectedModelTypeLarge] = resolved.Large + cfg.Models[SelectedModelTypeSmall] = resolved.Small + + // Persist any fallback corrections while we still hold writeMu. + if resolved.LargeFallback { + if err := store.updateLocked(ScopeGlobal, func(c *Config) map[string]any { + return store.updatePreferredModelFields(c, SelectedModelTypeLarge, resolved.Large) + }); err != nil { + return nil, fmt.Errorf("failed to update preferred large model: %w", err) + } + } + if resolved.SmallFallback { + if err := store.updateLocked(ScopeGlobal, func(c *Config) map[string]any { + return store.updatePreferredModelFields(c, SelectedModelTypeSmall, resolved.Small) + }); err != nil { + return nil, fmt.Errorf("failed to update preferred small model: %w", err) + } + } + store.SetupAgents() + + // Capture initial staleness snapshot + // Capture initial staleness snapshot. Track every discovered config path, + // not just the ones that loaded, so a config file created after startup + // (e.g. a crushrc added mid-session) is detected as a change. + store.captureStalenessSnapshot(append(slices.Clone(configPaths), loadedPaths...)) + + return store, nil +} + +// mustMarshalConfig marshals the config to JSON bytes, returning empty JSON on +// error. +func mustMarshalConfig(cfg *Config) []byte { + data, err := json.Marshal(cfg) + if err != nil { + return []byte("{}") + } + return data +} + +func PushPopCrushEnv() func() { + var found []string + for _, ev := range os.Environ() { + if strings.HasPrefix(ev, "CRUSH_") { + pair := strings.SplitN(ev, "=", 2) + if len(pair) != 2 { + continue + } + found = append(found, strings.TrimPrefix(pair[0], "CRUSH_")) + } + } + backups := make(map[string]string) + for _, ev := range found { + backups[ev] = os.Getenv(ev) + } + + for _, ev := range found { + os.Setenv(ev, os.Getenv("CRUSH_"+ev)) + } + + restore := func() { + for k, v := range backups { + os.Setenv(k, v) + } + } + return restore +} + +func (c *Config) configureProviders(ctx context.Context, store *ConfigStore, env env.Env, resolver VariableResolver, knownProviders []catwalk.Provider) error { + knownProviderNames := make(map[string]bool) + restore := PushPopCrushEnv() + defer restore() + + // When disable_default_providers is enabled, skip all default/embedded + // providers entirely. Users must fully specify any providers they want. + // We skip to the custom provider validation loop which handles all + // user-configured providers uniformly. + if c.Options.DisableDefaultProviders { + knownProviders = nil + } + + for _, p := range knownProviders { + knownProviderNames[string(p.ID)] = true + config, configExists := c.Providers.Get(string(p.ID)) + // if the user configured a known provider we need to allow it to override a couple of parameters + if configExists { + if config.BaseURL != "" { + p.APIEndpoint = config.BaseURL + } + if config.APIKey != "" { + p.APIKey = config.APIKey + } + if len(config.Models) > 0 { + models := []catwalk.Model{} + seen := make(map[string]bool) + + for _, model := range config.Models { + if seen[model.ID] { + continue + } + seen[model.ID] = true + if model.Name == "" { + model.Name = model.ID + } + models = append(models, model) + } + for _, model := range p.Models { + if seen[model.ID] { + continue + } + seen[model.ID] = true + if model.Name == "" { + model.Name = model.ID + } + models = append(models, model) + } + + p.Models = models + } + } + + headers := map[string]string{} + if len(p.DefaultHeaders) > 0 { + maps.Copy(headers, p.DefaultHeaders) + } + if len(config.ExtraHeaders) > 0 { + maps.Copy(headers, config.ExtraHeaders) + } + // Provider headers use the same error contract as MCP headers: + // a failing $(...) aborts the provider load with a clear + // message, and a header that resolves to the empty string + // (unset bare $VAR under lenient nounset, $(echo), or literal + // "") is dropped from the outgoing request. + for k, v := range headers { + resolved, err := resolver.ResolveValue(v) + if err != nil { + return fmt.Errorf("resolving provider %s header %q: %w", p.ID, k, err) + } + if resolved == "" { + delete(headers, k) + continue + } + headers[k] = resolved + } + // Start from user config so all user fields survive without + // explicit copying. Overlay catwalk identity/endpoint fields + // (already merged with user overrides above). + prepared := config + prepared.ID = string(p.ID) + prepared.Name = p.Name + prepared.BaseURL = p.APIEndpoint + prepared.APIKey = p.APIKey + prepared.APIKeyTemplate = p.APIKey // Store original template for re-resolution + prepared.Type = p.Type + prepared.Models = p.Models + prepared.ExtraHeaders = headers + if prepared.ExtraParams == nil { + prepared.ExtraParams = make(map[string]string) + } + + switch { + case p.ID == catwalk.InferenceProviderAnthropic && config.OAuthToken != nil: + // Claude Code subscription is not supported anymore. Remove to show onboarding. + // RemoveConfigField persists the deletion to disk. The in-memory + // state is kept consistent by the Providers.Del call below; any + // concurrent reload that races with this write will also see the + // removal because it re-reads from disk. + store.RemoveConfigField(ScopeGlobal, "providers.anthropic") + c.Providers.Del(string(p.ID)) + continue + case p.ID == catwalk.InferenceProviderCopilot && config.OAuthToken != nil: + prepared.SetupGitHubCopilot() + } + + switch p.ID { + // Handle specific providers that require additional configuration + case catwalk.InferenceProviderVertexAI: + var ( + project = env.Get("VERTEXAI_PROJECT") + location = env.Get("VERTEXAI_LOCATION") + ) + if project == "" || location == "" { + if configExists { + slog.Warn("Skipping Vertex AI provider due to missing credentials") + c.Providers.Del(string(p.ID)) + } + continue + } + prepared.ExtraParams["project"] = project + prepared.ExtraParams["location"] = location + case catwalk.InferenceProviderAzure: + endpoint, err := resolver.ResolveValue(p.APIEndpoint) + if err != nil || endpoint == "" { + if configExists { + slog.Warn("Skipping Azure provider due to missing API endpoint", "provider", p.ID, "error", err) + c.Providers.Del(string(p.ID)) + } + continue + } + prepared.BaseURL = endpoint + prepared.ExtraParams["apiVersion"] = env.Get("AZURE_OPENAI_API_VERSION") + case catwalk.InferenceProviderBedrock, catwalk.InferenceProviderBedrockEurope: + if p.APIKey == "" && !hasAWSCredentials(env) { + if configExists { + slog.Warn("Skipping Bedrock provider due to missing AWS credentials") + c.Providers.Del(string(p.ID)) + } + continue + } + case catwalk.InferenceProvider("hyper"): + if apiKey := env.Get("HYPER_API_KEY"); apiKey != "" { + prepared.APIKey = apiKey + prepared.APIKeyTemplate = apiKey + } else { + v, err := resolver.ResolveValue(p.APIKey) + if v == "" || err != nil { + if configExists { + slog.Warn("Skipping Hyper provider due to missing API key", "provider", p.ID) + c.Providers.Del(string(p.ID)) + } + continue + } + } + default: + // An OAuth login is a credential too: providers signed in + // through OAuth (e.g. OpenAI with a ChatGPT account) are + // configured even when no API key is present. + if config.OAuthToken != nil { + break + } + // if the provider api or endpoint are missing we skip them + v, err := resolver.ResolveValue(p.APIKey) + if v == "" || err != nil { + if configExists { + slog.Warn("Skipping provider due to missing API key", "provider", p.ID) + c.Providers.Del(string(p.ID)) + } + continue + } + } + c.Providers.Set(string(p.ID), prepared) + } + + // Discover models concurrently for custom providers that need it. + // A provider needs discovery when discover_models is explicitly true, + // or when the models list is empty (auto-trigger, unless opted out). + type discoveryResult struct { + models []catwalk.Model + err error + } + + discoveryResults := make(map[string]discoveryResult) + var mu sync.Mutex + var wg sync.WaitGroup + + discoverCtx, discoverCancel := context.WithTimeout(ctx, 3*time.Second) + for id, pc := range c.Providers.Seq2() { + if knownProviderNames[id] { + continue + } + if pc.Disable || pc.BaseURL == "" { + continue + } + wantsDiscovery := pc.AutoDiscoverModels != nil && *pc.AutoDiscoverModels + autoTrigger := len(pc.Models) == 0 && (pc.AutoDiscoverModels == nil || *pc.AutoDiscoverModels) + if !wantsDiscovery && !autoTrigger { + continue + } + providerID := cmp.Or(pc.ID, id) + cfg := discover.Config{ + ID: providerID, + BaseURL: pc.BaseURL, + APIKey: pc.APIKey, + ExtraHeaders: pc.ExtraHeaders, + ExistingModels: pc.Models, + } + providerType := cmp.Or(pc.Type, catwalk.TypeOpenAICompat) + wg.Go(func() { + models, err := discover.DiscoverModels(discoverCtx, cfg, resolver) + if err == nil && len(models) > 0 { + if enricher := discover.GetEnricher(string(providerType)); enricher != nil { + models, _ = enricher.EnrichModels(discoverCtx, cfg, resolver, models) + } + } + mu.Lock() + discoveryResults[id] = discoveryResult{models: models, err: err} + mu.Unlock() + }) + } + wg.Wait() + discoverCancel() + + // Validate the custom providers. + for id, providerConfig := range c.Providers.Seq2() { + if knownProviderNames[id] { + continue + } + + // Make sure the provider ID is set. + providerConfig.ID = id + providerConfig.Name = cmp.Or(providerConfig.Name, id) // Use ID as name if not set + // Default to OpenAI if not set. + providerConfig.Type = cmp.Or(providerConfig.Type, catwalk.TypeOpenAICompat) + if !slices.Contains(catwalk.KnownProviderTypes(), providerConfig.Type) && + providerConfig.Type != hyper.Name && + !discover.IsKnownCustomProvider(string(providerConfig.Type)) { + slog.Warn("Skipping custom provider due to unsupported provider type", "provider", id) + c.Providers.Del(id) + continue + } + + if providerConfig.Disable { + slog.Debug("Skipping custom provider due to disable flag", "provider", id) + c.Providers.Del(id) + continue + } + apiKey, err := resolver.ResolveValue(providerConfig.APIKey) + if apiKey == "" || err != nil { + slog.Warn("Provider is missing API key, this might be OK for local providers", "provider", id) + } + baseURL, err := resolver.ResolveValue(providerConfig.BaseURL) + if baseURL == "" || err != nil { + slog.Warn("Skipping custom provider due to missing API endpoint", "provider", id, "error", err) + c.Providers.Del(id) + continue + } + + // Apply discovery results if available. + if result, ok := discoveryResults[id]; ok { + if result.err != nil { + slog.Warn("Model discovery failed", "provider", id, "error", result.err) + if len(providerConfig.Models) == 0 { + slog.Warn("Skipping provider with no models after failed discovery", "provider", id) + c.Providers.Del(id) + continue + } + } else if len(result.models) > 0 { + providerConfig.Models = result.models + slog.Info("Discovered models for provider", "provider", id, "count", len(result.models)) + } + } + + if len(providerConfig.Models) == 0 { + slog.Warn("Skipping custom provider because the provider has no models", "provider", id) + c.Providers.Del(id) + continue + } + + // Custom-provider headers share the MCP error contract; see + // the known-provider loop above. + for k, v := range providerConfig.ExtraHeaders { + resolved, err := resolver.ResolveValue(v) + if err != nil { + return fmt.Errorf("resolving provider %s header %q: %w", id, k, err) + } + if resolved == "" { + delete(providerConfig.ExtraHeaders, k) + continue + } + providerConfig.ExtraHeaders[k] = resolved + } + + c.Providers.Set(id, providerConfig) + } + + if c.Providers.Len() == 0 && c.Options.DisableDefaultProviders { + return fmt.Errorf("default providers are disabled and there are no custom providers are configured") + } + + return nil +} + +// applyEnv sets top-level env vars from the config. Keys are sorted for +// deterministic ordering so that vars referencing other vars via the +// value resolver produce consistent results. +func (c *Config) applyEnv(resolver VariableResolver) { + keys := make([]string, 0, len(c.Env)) + for k := range c.Env { + keys = append(keys, k) + } + slices.Sort(keys) + for _, k := range keys { + resolved, err := resolver.ResolveValue(c.Env[k]) + if err != nil { + slog.Warn("Skipping env var due to resolution failure.", "key", k, "value", c.Env[k], "error", err) + continue + } + os.Setenv(k, resolved) + } +} + +// NormalizeOptions allocates Options and Options.TUI and fills in the option +// defaults the UI relies on, so readers can dereference them without guarding. +// Configs loaded from disk get this via setDefaults; configs arriving over the +// wire from a Crush server need the same treatment before the UI reads them. +// +// DiffMode is deliberately left alone: the permissions dialog reads its zero +// value as "choose split or unified from the terminal width". +func (c *Config) NormalizeOptions() { + if c.Options == nil { + c.Options = &Options{} + } + if c.Options.TUI == nil { + c.Options.TUI = &TUIOptions{} + } + if c.Options.TUI.Scrollbar == "" { + c.Options.TUI.Scrollbar = ScrollbarDefault + } + if c.Options.TUI.ExitBanner == "" { + c.Options.TUI.ExitBanner = ExitBannerDefault + } +} + +func (c *Config) setDefaults(workingDir, dataDir string) { + c.NormalizeOptions() + if len(c.Options.GlobalContextPaths) == 0 { + crushConfigDir := filepath.Dir(GlobalConfig()) + c.Options.GlobalContextPaths = []string{ + filepath.Join(crushConfigDir, "CRUSH.md"), + filepath.Join(filepath.Dir(crushConfigDir), "AGENTS.md"), + } + } + slices.Sort(c.Options.GlobalContextPaths) + c.Options.GlobalContextPaths = slices.Compact(c.Options.GlobalContextPaths) + + if dataDir != "" { + c.Options.DataDirectory = dataDir + } else if c.Options.DataDirectory == "" { + if path, ok := fsext.LookupClosestBounded(workingDir, projectBoundary(workingDir), defaultDataDirectory); ok { + c.Options.DataDirectory = path + } else { + c.Options.DataDirectory = filepath.Join(workingDir, defaultDataDirectory) + } + } + c.Options.DataDirectory = filepath.Clean(filepathext.SmartJoin(workingDir, c.Options.DataDirectory)) + if c.Providers == nil { + c.Providers = csync.NewMap[string, ProviderConfig]() + } + if c.Models == nil { + c.Models = make(map[SelectedModelType]SelectedModel) + } + if c.RecentModels == nil { + c.RecentModels = make(map[SelectedModelType][]SelectedModel) + } + if c.MCP == nil { + c.MCP = make(map[string]MCPConfig) + } + // Drop orphaned OAuth token entries left behind when a user removes + // an MCP from crush.json. See MCPConfig.isOrphanedToken. + for name, m := range c.MCP { + if m.isOrphanedToken() { + delete(c.MCP, name) + } + } + if c.LSP == nil { + c.LSP = make(map[string]LSPConfig) + } + + // Apply defaults to LSP configurations + c.applyLSPDefaults() + + // Add the default context paths if they are not already present + c.Options.ContextPaths = append(slices.Clone(defaultContextPaths), c.Options.ContextPaths...) + + slices.Sort(c.Options.ContextPaths) + c.Options.ContextPaths = slices.Compact(c.Options.ContextPaths) + + // Add the default skills directories if not already present. + for _, dir := range GlobalSkillsDirs() { + if !slices.Contains(c.Options.SkillsPaths, dir) { + c.Options.SkillsPaths = append(c.Options.SkillsPaths, dir) + } + } + + // Project specific skills dirs. + c.Options.SkillsPaths = append(c.Options.SkillsPaths, ProjectSkillsDir(workingDir)...) + + if str, ok := os.LookupEnv("CRUSH_DISABLE_PROVIDER_AUTO_UPDATE"); ok { + c.Options.DisableProviderAutoUpdate, _ = strconv.ParseBool(str) + } + + if str, ok := os.LookupEnv("CRUSH_DISABLE_DEFAULT_PROVIDERS"); ok { + c.Options.DisableDefaultProviders, _ = strconv.ParseBool(str) + } + + if c.Options.Attribution == nil { + c.Options.Attribution = &Attribution{ + TrailerStyle: TrailerStyleAssistedBy, + GeneratedWith: true, + } + } else if c.Options.Attribution.TrailerStyle == "" { + // Migrate deprecated co_authored_by or apply default + if c.Options.Attribution.CoAuthoredBy != nil { + if *c.Options.Attribution.CoAuthoredBy { + c.Options.Attribution.TrailerStyle = TrailerStyleCoAuthoredBy + } else { + c.Options.Attribution.TrailerStyle = TrailerStyleNone + } + } else { + c.Options.Attribution.TrailerStyle = TrailerStyleAssistedBy + } + } + + c.Options.InitializeAs = cmp.Or(c.Options.InitializeAs, defaultInitializeAs) +} + +// powernapDefaults caches the powernap default LSP server catalog. The +// catalog is static and immutable for the life of the process, but +// building it (NewManager + LoadDefaults) is expensive and was previously +// repeated on every config reload. We load it once and only ever read from +// it via GetServer, so a shared instance is safe. +var ( + powernapDefaultsOnce sync.Once + powernapDefaults *powernapConfig.Manager +) + +func lspDefaultsManager() *powernapConfig.Manager { + powernapDefaultsOnce.Do(func() { + m := powernapConfig.NewManager() + // LoadDefaults only fails on malformed embedded defaults, which + // would be a build-time bug; treat the manager as usable either + // way so a transient error never wedges config loading. + _ = m.LoadDefaults() + powernapDefaults = m + }) + return powernapDefaults +} + +// applyLSPDefaults applies default values from powernap to LSP configurations +func (c *Config) applyLSPDefaults() { + // Reuse the process-wide default catalog; building it per reload was a + // significant chunk of reload latency. + configManager := lspDefaultsManager() + + // Apply defaults to each LSP configuration + for name, cfg := range c.LSP { + // Try to get defaults from powernap based on name or command name. + base, ok := configManager.GetServer(name) + if !ok { + base, ok = configManager.GetServer(cfg.Command) + if !ok { + continue + } + } + if cfg.Options == nil { + cfg.Options = base.Settings + } + if cfg.InitOptions == nil { + cfg.InitOptions = base.InitOptions + } + if len(cfg.FileTypes) == 0 { + cfg.FileTypes = base.FileTypes + } + if len(cfg.RootMarkers) == 0 { + cfg.RootMarkers = base.RootMarkers + } + cfg.Command = cmp.Or(cfg.Command, base.Command) + if len(cfg.Args) == 0 { + cfg.Args = base.Args + } + if len(cfg.Env) == 0 { + cfg.Env = base.Environment + } + // Update the config in the map + c.LSP[name] = cfg + } +} + +func (c *Config) defaultModelSelection(knownProviders []catwalk.Provider) (largeModel SelectedModel, smallModel SelectedModel, err error) { + if len(knownProviders) == 0 && c.Providers.Len() == 0 { + err = fmt.Errorf("no providers configured, please configure at least one provider") + return largeModel, smallModel, err + } + + // Use the first provider enabled based on the known providers order + // if no provider found that is known use the first provider configured + for _, p := range knownProviders { + providerConfig, ok := c.Providers.Get(string(p.ID)) + if !ok || providerConfig.Disable { + continue + } + defaultLargeModel := c.GetModel(string(p.ID), p.DefaultLargeModelID) + if defaultLargeModel == nil { + slog.Warn("Default large model %s not found for provider %s", p.DefaultLargeModelID, p.ID) + if len(providerConfig.Models) == 0 { + return largeModel, smallModel, fmt.Errorf("default large model %s not found for provider %s", p.DefaultLargeModelID, p.ID) + } + defaultLargeModel = &providerConfig.Models[0] + } + largeModel = SelectedModel{ + Provider: string(p.ID), + Model: defaultLargeModel.ID, + MaxTokens: defaultLargeModel.DefaultMaxTokens, + ReasoningEffort: defaultLargeModel.DefaultReasoningEffort, + } + + defaultSmallModel := c.GetModel(string(p.ID), p.DefaultSmallModelID) + if defaultSmallModel == nil { + slog.Warn("Default small model %s not found for provider %s", p.DefaultSmallModelID, p.ID) + if len(providerConfig.Models) == 0 { + return largeModel, smallModel, fmt.Errorf("default small model %s not found for provider %s", p.DefaultSmallModelID, p.ID) + } + defaultSmallModel = &providerConfig.Models[0] + } + smallModel = SelectedModel{ + Provider: string(p.ID), + Model: defaultSmallModel.ID, + MaxTokens: defaultSmallModel.DefaultMaxTokens, + ReasoningEffort: defaultSmallModel.DefaultReasoningEffort, + } + return largeModel, smallModel, err + } + + enabledProviders := c.EnabledProviders() + slices.SortFunc(enabledProviders, func(a, b ProviderConfig) int { + return strings.Compare(a.ID, b.ID) + }) + + if len(enabledProviders) == 0 { + err = fmt.Errorf("no providers configured, please configure at least one provider") + return largeModel, smallModel, err + } + + providerConfig := enabledProviders[0] + if len(providerConfig.Models) == 0 { + err = fmt.Errorf("provider %s has no models configured", providerConfig.ID) + return largeModel, smallModel, err + } + defaultLargeModel := c.GetModel(providerConfig.ID, providerConfig.Models[0].ID) + largeModel = SelectedModel{ + Provider: providerConfig.ID, + Model: defaultLargeModel.ID, + MaxTokens: defaultLargeModel.DefaultMaxTokens, + } + defaultSmallModel := c.GetModel(providerConfig.ID, providerConfig.Models[0].ID) + smallModel = SelectedModel{ + Provider: providerConfig.ID, + Model: defaultSmallModel.ID, + MaxTokens: defaultSmallModel.DefaultMaxTokens, + } + return largeModel, smallModel, err +} + +// resolvedModels holds the result of resolving user-configured model +// selections against the provider catalog. +type resolvedModels struct { + Large SelectedModel + Small SelectedModel + LargeFallback bool // true if Large was corrected to a default + SmallFallback bool // true if Small was corrected to a default +} + +// resolveSelectedModels validates the user's configured model selections +// against the provider catalog, falling back to defaults when a model ID is +// invalid. It is pure resolution logic: it does not mutate the store or +// touch disk. The caller assigns the results to c.Models and persists any +// fallback corrections as appropriate. +func resolveSelectedModels(cfg *Config, knownProviders []catwalk.Provider) (resolvedModels, error) { + var result resolvedModels + defaultLarge, defaultSmall, err := cfg.defaultModelSelection(knownProviders) + if err != nil { + return result, fmt.Errorf("failed to select default models: %w", err) + } + large, small := defaultLarge, defaultSmall + + largeModelSelected, largeModelConfigured := cfg.Models[SelectedModelTypeLarge] + if largeModelConfigured { + if largeModelSelected.Model != "" { + large.Model = largeModelSelected.Model + } + if largeModelSelected.Provider != "" { + large.Provider = largeModelSelected.Provider + } + model := cfg.GetModel(large.Provider, large.Model) + if model == nil { + large = defaultLarge + result.LargeFallback = true + } else { + if largeModelSelected.MaxTokens > 0 { + large.MaxTokens = largeModelSelected.MaxTokens + } else { + large.MaxTokens = model.DefaultMaxTokens + } + if largeModelSelected.ReasoningEffort != "" { + large.ReasoningEffort = largeModelSelected.ReasoningEffort + } else { + large.ReasoningEffort = model.DefaultReasoningEffort + } + large.Think = largeModelSelected.Think + if largeModelSelected.Temperature != nil { + large.Temperature = largeModelSelected.Temperature + } + if largeModelSelected.TopP != nil { + large.TopP = largeModelSelected.TopP + } + if largeModelSelected.TopK != nil { + large.TopK = largeModelSelected.TopK + } + if largeModelSelected.FrequencyPenalty != nil { + large.FrequencyPenalty = largeModelSelected.FrequencyPenalty + } + if largeModelSelected.PresencePenalty != nil { + large.PresencePenalty = largeModelSelected.PresencePenalty + } + if largeModelSelected.ProviderOptions != nil { + large.ProviderOptions = maps.Clone(largeModelSelected.ProviderOptions) + } + } + } + smallModelSelected, smallModelConfigured := cfg.Models[SelectedModelTypeSmall] + if smallModelConfigured { + if smallModelSelected.Model != "" { + small.Model = smallModelSelected.Model + } + if smallModelSelected.Provider != "" { + small.Provider = smallModelSelected.Provider + } + + model := cfg.GetModel(small.Provider, small.Model) + if model == nil { + small = defaultSmall + result.SmallFallback = true + } else { + if smallModelSelected.MaxTokens > 0 { + small.MaxTokens = smallModelSelected.MaxTokens + } else { + small.MaxTokens = model.DefaultMaxTokens + } + if smallModelSelected.ReasoningEffort != "" { + small.ReasoningEffort = smallModelSelected.ReasoningEffort + } else { + small.ReasoningEffort = model.DefaultReasoningEffort + } + if smallModelSelected.Temperature != nil { + small.Temperature = smallModelSelected.Temperature + } + if smallModelSelected.TopP != nil { + small.TopP = smallModelSelected.TopP + } + if smallModelSelected.TopK != nil { + small.TopK = smallModelSelected.TopK + } + if smallModelSelected.FrequencyPenalty != nil { + small.FrequencyPenalty = smallModelSelected.FrequencyPenalty + } + if smallModelSelected.PresencePenalty != nil { + small.PresencePenalty = smallModelSelected.PresencePenalty + } + if smallModelSelected.ProviderOptions != nil { + small.ProviderOptions = maps.Clone(smallModelSelected.ProviderOptions) + } + small.Think = smallModelSelected.Think + } + } + + // When small isn't explicitly configured and the provider isn't a + // known built-in, use the large model as the small model. This + // prevents two different models from being requested concurrently + // for local/openai-compat providers. + if !smallModelConfigured { + isKnownProvider := false + for _, kp := range knownProviders { + if string(kp.ID) == small.Provider { + isKnownProvider = true + break + } + } + if !isKnownProvider { + slog.Warn("Using large model as small model for unknown provider", "provider", large.Provider, "model", large.Model) + small = large + } + } + + result.Large = large + result.Small = small + return result, nil +} + +// lookupConfigs searches config files starting at cwd and walking up +// through the current project. The upward walk stops at the git +// working tree root when one can be detected, otherwise at cwd itself, +// so an unrelated crush.json placed above the project is never picked +// up. Global user-level config locations are always included +// regardless of the boundary. +func lookupConfigs(cwd string) []string { + // Prepend global user config and machine-owned data JSON. Only the user + // config directory contributes a crushrc; the data directory is writable + // machine state and must never be executed as Bash. Missing files are + // skipped when loaded. + configPaths := []string{ + systemConfigPath, + GlobalConfig(), + shellConfigSibling(GlobalConfig()), + GlobalConfigData(), + } + + // Ordered high-to-low priority within a directory. LookupBounded returns + // matches in this order, and the later reverse + merge make the earliest + // listed name win on conflict. So: .crushrc beats crushrc, both beat the + // JSON configs, and .crush.json beats crush.json. + configNames := []string{ + "." + appName + "rc", + appName + "rc", + "." + appName + ".json", + appName + ".json", + } + + foundConfigs, err := fsext.LookupBounded(cwd, projectBoundary(cwd), configNames...) + if err != nil { + // returns at least default configs + return configPaths + } + + // reverse order so last config has more priority + slices.Reverse(foundConfigs) + + return append(configPaths, foundConfigs...) +} + +func loadFromConfigPaths(ctx context.Context, configPaths []string) (*Config, []string, error) { + var configs [][]byte + var loaded []string + + // Track directories that have both crush.json and crushrc to warn + // about potential confusion, along with the top-level keys each + // defines so we can report conflicts. + jsonDirKeys := make(map[string]map[string]bool) + shDirKeys := make(map[string]map[string]bool) + + for _, path := range configPaths { + if path == "" { + continue + } + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + continue + } + return nil, nil, fmt.Errorf("failed to open config file %s: %w", path, err) + } + if len(data) == 0 { + continue + } + + dir := filepath.Dir(path) + if isShellConfig(path) { + jsonBytes, err := shellconfig.LoadShellConfig(ctx, path, data) + if err != nil { + return nil, nil, fmt.Errorf("failed to load shell config %s: %w", path, err) + } + if len(jsonBytes) > 0 { + if !json.Valid(jsonBytes) { + return nil, nil, fmt.Errorf("shell config %s produced invalid JSON", path) + } + addTopLevelKeys(shDirKeys, dir, jsonBytes) + configs = append(configs, jsonBytes) + loaded = append(loaded, path) + } + } else { + if !json.Valid(data) { + return nil, nil, fmt.Errorf("invalid JSON in config file %s", path) + } + addTopLevelKeys(jsonDirKeys, dir, data) + configs = append(configs, data) + loaded = append(loaded, path) + } + } + + // Warn if both a JSON config and a crushrc exist in the same directory + // and define overlapping top-level keys. Disjoint coexistence is + // intentional and not worth warning about. + for dir, jKeys := range jsonDirKeys { + sKeys, ok := shDirKeys[dir] + if !ok { + continue + } + var conflicts []string + for k := range jKeys { + if sKeys[k] { + conflicts = append(conflicts, k) + } + } + if len(conflicts) > 0 { + slices.Sort(conflicts) + slog.Warn("Found both a JSON config and a crushrc in the same directory; merging with crushrc taking precedence", + "dir", dir, "conflicting_keys", strings.Join(conflicts, ", ")) + } + } + + cfg, err := loadFromBytes(configs) + if err != nil { + return nil, nil, err + } + return cfg, loaded, nil +} + +// addTopLevelKeys records the top-level JSON keys present in data into the +// set for dir. +func addTopLevelKeys(m map[string]map[string]bool, dir string, data []byte) { + keys := m[dir] + if keys == nil { + keys = make(map[string]bool) + m[dir] = keys + } + gjson.ParseBytes(data).ForEach(func(key, _ gjson.Result) bool { + keys[key.String()] = true + return true + }) +} + +func loadFromBytes(configs [][]byte) (*Config, error) { + if len(configs) == 0 { + return &Config{}, nil + } + + data, err := jsons.Merge(configs) + if err != nil { + return nil, err + } + var config Config + if err := json.Unmarshal(data, &config); err != nil { + return nil, err + } + return &config, nil +} + +func hasAWSCredentials(env env.Env) bool { + if env.Get("AWS_BEARER_TOKEN_BEDROCK") != "" { + return true + } + + if env.Get("AWS_ACCESS_KEY_ID") != "" && env.Get("AWS_SECRET_ACCESS_KEY") != "" { + return true + } + + if env.Get("AWS_PROFILE") != "" || env.Get("AWS_DEFAULT_PROFILE") != "" { + return true + } + + if env.Get("AWS_REGION") != "" || env.Get("AWS_DEFAULT_REGION") != "" { + return true + } + + if env.Get("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI") != "" || + env.Get("AWS_CONTAINER_CREDENTIALS_FULL_URI") != "" { + return true + } + + // File-based credential discovery requires filesystem stats, so do it + // last and skip it under test. Checking testing.Testing() before the + // os.Stat call (rather than after, in the && tail) ensures the syscall + // is never issued during tests, where it otherwise ran unconditionally + // and only had its result discarded. + if testing.Testing() { + return false + } + if _, err := os.Stat(filepath.Join(home.Dir(), ".aws/credentials")); err == nil { + return true + } + if _, err := os.Stat(filepath.Join(home.Dir(), ".aws/login")); err == nil { + return true + } + + return false +} + +// migrateDisableNotifications migrates the deprecated disable_notifications +// and notification_style fields to the unified notifications field. It checks +// both the user config (~/.config) and data config (~/.local) files. If +// disable_notifications is true, it sets notifications to "disabled" in the +// data file. If notification_style is set, it moves the value to notifications. +// Regardless of value, it removes the deprecated fields from any file that +// contains them. +func migrateDisableNotifications() { + globalConfig := GlobalConfig() + dataConfig := GlobalConfigData() + + var wasDisabled bool + var styleValue string + filesToClean := []string{} + + for _, path := range []string{globalConfig, dataConfig} { + data, err := os.ReadFile(path) + if err != nil { + continue + } + needsClean := false + if gjson.Get(string(data), "options.disable_notifications").Exists() { + needsClean = true + if gjson.Get(string(data), "options.disable_notifications").Bool() { + wasDisabled = true + } + } + if v := gjson.Get(string(data), "options.notification_style"); v.Exists() { + needsClean = true + if styleValue == "" { + styleValue = v.String() + } + } + if needsClean { + filesToClean = append(filesToClean, path) + } + } + + if len(filesToClean) == 0 { + return + } + + // Determine the value to persist: notification_style takes precedence, + // then disable_notifications: true maps to "disabled". + migratedValue := styleValue + if migratedValue == "" && wasDisabled { + migratedValue = "disabled" + } + + if migratedValue != "" { + data, err := os.ReadFile(dataConfig) + if err == nil { + if !gjson.Get(string(data), "options.notifications").Exists() { + updated, err := sjson.Set(string(data), "options.notifications", migratedValue) + if err == nil { + if err := atomicWriteFile(dataConfig, []byte(updated), 0o600); err != nil { + slog.Warn("Failed to migrate to notifications field", "error", err) + } else { + slog.Info("Migrated notification settings to notifications field", "value", migratedValue) + } + } + } + } + } + + // Remove deprecated fields from all files that contain them. + for _, path := range filesToClean { + data, err := os.ReadFile(path) + if err != nil { + continue + } + updated := string(data) + updated, _ = sjson.Delete(updated, "options.disable_notifications") + updated, _ = sjson.Delete(updated, "options.notification_style") + if updated == string(data) { + continue + } + if err := atomicWriteFile(path, []byte(updated), 0o600); err != nil { + slog.Warn("Failed to write migrated config", "path", path, "error", err) + } + } +} + +// GlobalConfig returns the global configuration file path for the application. +func GlobalConfig() string { + if crushGlobal := os.Getenv("CRUSH_GLOBAL_CONFIG"); crushGlobal != "" { + return filepath.Join(crushGlobal, fmt.Sprintf("%s.json", appName)) + } + return filepath.Join(home.Config(), appName, fmt.Sprintf("%s.json", appName)) +} + +// shellConfigSibling returns the crushrc path that sits alongside a given +// crush.json path (same directory). Used so global config locations pick up a +// shell config, not just JSON. +func shellConfigSibling(jsonPath string) string { + return filepath.Join(filepath.Dir(jsonPath), appName+"rc") +} + +// isShellConfig reports whether a config path is a shell config (crushrc or +// the hidden .crushrc), as opposed to a JSON config. +func isShellConfig(path string) bool { + base := filepath.Base(path) + return base == appName+"rc" || base == "."+appName+"rc" +} + +// GlobalCacheDir returns the path to the global cache directory for the +// application. +func GlobalCacheDir() string { + if crushCache := os.Getenv("CRUSH_CACHE_DIR"); crushCache != "" { + return crushCache + } + if xdgCacheHome := os.Getenv("XDG_CACHE_HOME"); xdgCacheHome != "" { + return filepath.Join(xdgCacheHome, appName) + } + if runtime.GOOS == "windows" { + localAppData := cmp.Or( + os.Getenv("LOCALAPPDATA"), + filepath.Join(os.Getenv("USERPROFILE"), "AppData", "Local"), + ) + return filepath.Join(localAppData, appName, "cache") + } + return filepath.Join(home.Dir(), ".cache", appName) +} + +// ProjectConfigs returns list of current project configs paths. +func ProjectConfigs(cwd string) []string { + return lookupConfigs(cwd) +} + +// GlobalConfigData returns the path to the main data directory for the application. +// this config is used when the app overrides configurations instead of updating the global config. +func GlobalConfigData() string { + if crushData := os.Getenv("CRUSH_GLOBAL_DATA"); crushData != "" { + return filepath.Join(crushData, fmt.Sprintf("%s.json", appName)) + } + if xdgDataHome := os.Getenv("XDG_DATA_HOME"); xdgDataHome != "" { + return filepath.Join(xdgDataHome, appName, fmt.Sprintf("%s.json", appName)) + } + + // return the path to the main data directory + // for windows, it should be in `%LOCALAPPDATA%/crush/` + // for linux and macOS, it should be in `$HOME/.local/share/crush/` + if runtime.GOOS == "windows" { + localAppData := cmp.Or( + os.Getenv("LOCALAPPDATA"), + filepath.Join(os.Getenv("USERPROFILE"), "AppData", "Local"), + ) + return filepath.Join(localAppData, appName, fmt.Sprintf("%s.json", appName)) + } + + return filepath.Join(home.Dir(), ".local", "share", appName, fmt.Sprintf("%s.json", appName)) +} + +// GlobalWorkspaceDir returns the path to the global server workspace +// directory. This directory acts as a meta-workspace for the server +// process, giving it a real workingDir so that config loading, scoped +// writes, and provider resolution behave identically to project +// workspaces. +func GlobalWorkspaceDir() string { + return filepath.Dir(GlobalConfigData()) +} + +func assignIfNil[T any](ptr **T, val T) { + if *ptr == nil { + *ptr = &val + } +} + +func isInsideWorktree() bool { + bts, err := exec.CommandContext( + context.Background(), + "git", "rev-parse", + "--is-inside-work-tree", + ).CombinedOutput() + return err == nil && strings.TrimSpace(string(bts)) == "true" +} + +// worktreeRoot returns the absolute path of the git working tree root for +// dir, or the empty string if dir is not inside a working tree (bare +// repositories, missing git binary, plain directories, or any other +// failure mode). Linked worktrees and submodules each report their own +// top-level, which is what callers want when bounding lookups. +// worktreeRootCache memoizes the git worktree root per directory. The root +// is stable for the life of the process, so we avoid re-shelling out to +// "git rev-parse" on every config reload. Keyed by the requested dir; the +// value is the resolved root ("" when dir is not in a git worktree). +var worktreeRootCache sync.Map // map[string]string + +func worktreeRoot(dir string) string { + if cached, ok := worktreeRootCache.Load(dir); ok { + return cached.(string) + } + root := computeWorktreeRoot(dir) + worktreeRootCache.Store(dir, root) + return root +} + +func computeWorktreeRoot(dir string) string { + cmd := exec.CommandContext( + context.Background(), + "git", "rev-parse", "--show-toplevel", + ) + cmd.Dir = dir + out, err := cmd.Output() + if err != nil { + return "" + } + root := strings.TrimSpace(string(out)) + if root == "" { + return "" + } + abs, err := filepath.Abs(root) + if err != nil { + return "" + } + return abs +} + +// projectBoundary returns the directory at which an upward configuration +// search rooted at dir should stop. It is the git working tree root when +// one can be detected, otherwise dir itself. Returning dir as a +// fallback keeps Crush from silently adopting state files placed above +// the current project. +func projectBoundary(dir string) string { + if root := worktreeRoot(dir); root != "" { + return root + } + abs, err := filepath.Abs(dir) + if err != nil { + return dir + } + return abs +} + +// GlobalSkillsDirs returns the default directories for Agent Skills. +// Skills in these directories are auto-discovered and their files can be read +// without permission prompts. +func GlobalSkillsDirs() []string { + if crushSkills := os.Getenv("CRUSH_SKILLS_DIR"); crushSkills != "" { + return []string{crushSkills} + } + + paths := []string{ + filepath.Join(home.Config(), appName, "skills"), + filepath.Join(home.Config(), "agents", "skills"), + // Per the Agent Skills spec, scan ~/.agents/skills + filepath.Join(home.Dir(), ".agents", "skills"), + filepath.Join(home.Dir(), ".claude", "skills"), + } + + // On Windows, also load from app data on top of `$HOME/.config/crush`. + // This is here mostly for backwards compatibility. + if runtime.GOOS == "windows" { + appData := cmp.Or( + os.Getenv("LOCALAPPDATA"), + filepath.Join(os.Getenv("USERPROFILE"), "AppData", "Local"), + ) + paths = append( + paths, + filepath.Join(appData, appName, "skills"), + filepath.Join(appData, "agents", "skills"), + ) + } + + return paths +} + +// projectSkillSubdirs lists the conventional subdirectories where +// project-level skills are discovered. Shared across working-dir and +// git-root lookups to prevent drift when a new convention is added. +var projectSkillSubdirs = []string{ + ".agents/skills", + ".crush/skills", + ".claude/skills", + ".cursor/skills", +} + +// ProjectSkillsDir returns the default project directories for which Crush +// will look for skills. In addition to the working directory, it also +// checks the git working tree root so that monorepo-level skills are +// discovered when the user is inside a subdirectory. +// Working-directory paths come first so local skills take precedence +// over monorepo-level ones. +func ProjectSkillsDir(workingDir string) []string { + dirs := make([]string, 0, len(projectSkillSubdirs)*2) + for _, sub := range projectSkillSubdirs { + dirs = append(dirs, filepath.Join(workingDir, sub)) + } + + // When the working directory is inside a git repository, also look at + // the repository root so monorepo-level .agents/skills are found. + if root := worktreeRoot(workingDir); root != "" && root != workingDir { + for _, sub := range projectSkillSubdirs { + dirs = append(dirs, filepath.Join(root, sub)) + } + } + + return dirs +} + +func isAppleTerminal() bool { return os.Getenv("TERM_PROGRAM") == "Apple_Terminal" } + +// normalizeHookEvent maps user-provided event names to their canonical +// form. Matching is case-insensitive and accepts snake_case variants +// (e.g. "pre_tool_use" → "PreToolUse"). +func normalizeHookEvent(name string) string { + switch strings.ToLower(strings.ReplaceAll(name, "_", "")) { + case "pretooluse": + return "PreToolUse" + default: + return name + } +} + +// ValidateHooks normalizes event names and checks that every configured +// hook has a command and a syntactically valid matcher regex. Matcher +// compilation used for matching is owned by hooks.Runner; this function +// only validates up front so the user sees config errors at load time +// rather than on the first tool call. +func (c *Config) ValidateHooks() error { + // Normalize event name keys. + for event, eventHooks := range c.Hooks { + canonical := normalizeHookEvent(event) + if canonical != event { + c.Hooks[canonical] = append(c.Hooks[canonical], eventHooks...) + delete(c.Hooks, event) + } + } + + for event, eventHooks := range c.Hooks { + for i, h := range eventHooks { + if h.Command == "" { + return fmt.Errorf("hook %s[%d]: command is required", event, i) + } + if h.Matcher == "" { + continue + } + if _, err := regexp.Compile(h.Matcher); err != nil { + return fmt.Errorf("hook %s[%d]: invalid matcher regex %q: %w", event, i, h.Matcher, err) + } + } + } + return nil +} diff --git a/internal/config/load_test.go b/internal/config/load_test.go new file mode 100644 index 0000000000000000000000000000000000000000..417fc779d8f83d3962681a4addc8a28dd8e9e4b2 --- /dev/null +++ b/internal/config/load_test.go @@ -0,0 +1,2536 @@ +package config + +import ( + "context" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "testing" + "time" + + "charm.land/catwalk/pkg/catwalk" + "github.com/charmbracelet/crush/internal/csync" + "github.com/charmbracelet/crush/internal/env" + "github.com/charmbracelet/crush/internal/oauth" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestMain(m *testing.M) { + slog.SetDefault(slog.New(slog.NewTextHandler(io.Discard, nil))) + + exitVal := m.Run() + os.Exit(exitVal) +} + +func TestConfig_LoadFromBytes(t *testing.T) { + data1 := []byte(`{"providers": {"openai": {"api_key": "key1", "base_url": "https://api.openai.com/v1"}}}`) + data2 := []byte(`{"providers": {"openai": {"api_key": "key2", "base_url": "https://api.openai.com/v2"}}}`) + data3 := []byte(`{"providers": {"openai": {}}}`) + + loadedConfig, err := loadFromBytes([][]byte{data1, data2, data3}) + + require.NoError(t, err) + require.NotNil(t, loadedConfig) + require.Equal(t, 1, loadedConfig.Providers.Len()) + pc, _ := loadedConfig.Providers.Get("openai") + require.Equal(t, "key2", pc.APIKey) + require.Equal(t, "https://api.openai.com/v2", pc.BaseURL) +} + +func TestLookupConfigs_BoundedByProject(t *testing.T) { + // Force GlobalConfig and GlobalConfigData to point at locations we + // control so they can be present in the result without polluting + // the developer's real config. + globalDir := t.TempDir() + dataDir := t.TempDir() + t.Setenv("CRUSH_GLOBAL_CONFIG", globalDir) + t.Setenv("CRUSH_GLOBAL_DATA", dataDir) + + t.Run("does not pick up crush.json above non-git project", func(t *testing.T) { + parent := t.TempDir() + + // crush.json above the project must not be adopted. + require.NoError(t, os.WriteFile( + filepath.Join(parent, "crush.json"), + []byte(`{}`), + 0o644, + )) + + project := filepath.Join(parent, "project") + require.NoError(t, os.Mkdir(project, 0o755)) + + got := lookupConfigs(project) + for _, p := range got { + require.NotEqual(t, filepath.Join(parent, "crush.json"), p) + } + }) + + t.Run("does not climb out of git worktree to find crush.json", func(t *testing.T) { + if _, err := exec.LookPath("git"); err != nil { + t.Skip("git not available") + } + + parent := t.TempDir() + + require.NoError(t, os.WriteFile( + filepath.Join(parent, "crush.json"), + []byte(`{}`), + 0o644, + )) + + worktree := filepath.Join(parent, "worktree") + require.NoError(t, os.Mkdir(worktree, 0o755)) + gitInit := exec.CommandContext(t.Context(), "git", "init", "-q") + gitInit.Dir = worktree + require.NoError(t, gitInit.Run()) + + got := lookupConfigs(worktree) + strayEval, err := filepath.EvalSymlinks(filepath.Join(parent, "crush.json")) + require.NoError(t, err) + for _, p := range got { + pEval, err := filepath.EvalSymlinks(p) + if err != nil { + continue + } + require.NotEqual(t, strayEval, pEval, "must not adopt parent crush.json") + } + }) + + t.Run("picks up crush.json inside the project", func(t *testing.T) { + project := t.TempDir() + local := filepath.Join(project, "crush.json") + require.NoError(t, os.WriteFile(local, []byte(`{}`), 0o644)) + + got := lookupConfigs(project) + + localEval, err := filepath.EvalSymlinks(local) + require.NoError(t, err) + var foundLocal bool + for _, p := range got { + pEval, err := filepath.EvalSymlinks(p) + if err != nil { + continue + } + if pEval == localEval { + foundLocal = true + break + } + } + require.True(t, foundLocal, "expected project crush.json to be in lookup result: %v", got) + }) + + t.Run("global config is always included regardless of boundary", func(t *testing.T) { + project := t.TempDir() + + got := lookupConfigs(project) + // Global config and global data path are always prepended, + // even when no project file exists. + require.Contains(t, got, GlobalConfig()) + require.Contains(t, got, GlobalConfigData()) + }) + + t.Run("global shell config (crushrc) is included", func(t *testing.T) { + project := t.TempDir() + + got := lookupConfigs(project) + // A global crushrc is discovered only beside the user config. The data + // directory is machine-owned state and must never execute a crushrc. + require.Contains(t, got, shellConfigSibling(GlobalConfig())) + require.NotContains(t, got, shellConfigSibling(GlobalConfigData())) + }) + + t.Run("project crushrc and .crushrc are discovered", func(t *testing.T) { + project := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(project, "crushrc"), []byte(""), 0o644)) + require.NoError(t, os.WriteFile(filepath.Join(project, ".crushrc"), []byte(""), 0o644)) + + got := lookupConfigs(project) + require.Contains(t, got, filepath.Join(project, "crushrc")) + require.Contains(t, got, filepath.Join(project, ".crushrc")) + }) + + t.Run("system config is loaded first", func(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("system config not supported on Windows") + } + + got := lookupConfigs(t.TempDir()) + require.NotEmpty(t, got) + // The system-wide config must be first so it has the lowest + // priority when configs are merged. + require.Equal(t, "/etc/crush/crush.json", got[0]) + }) +} + +func TestLoadFromConfigPaths_InvalidJSON(t *testing.T) { + t.Parallel() + + t.Run("identifies the offending file", func(t *testing.T) { + t.Parallel() + tmpDir := t.TempDir() + good := filepath.Join(tmpDir, "good.json") + bad := filepath.Join(tmpDir, "bad.json") + require.NoError(t, os.WriteFile(good, []byte(`{"providers":{}}`), 0o644)) + require.NoError(t, os.WriteFile(bad, []byte(`{not valid json}`), 0o644)) + + _, _, err := loadFromConfigPaths(context.Background(), []string{good, bad}) + require.Error(t, err) + require.Contains(t, err.Error(), "invalid JSON in config file") + require.Contains(t, err.Error(), "bad.json") + }) + + t.Run("skips missing and empty files", func(t *testing.T) { + t.Parallel() + tmpDir := t.TempDir() + empty := filepath.Join(tmpDir, "empty.json") + require.NoError(t, os.WriteFile(empty, []byte(""), 0o644)) + + cfg, _, err := loadFromConfigPaths(context.Background(), []string{ + filepath.Join(tmpDir, "nonexistent.json"), + empty, + }) + require.NoError(t, err) + require.NotNil(t, cfg) + }) +} + +// TestLoadFromConfigPaths_ConflictWarningNamesKeys verifies that when a JSON +// config and a crushrc coexist in the same directory, the merge warning names +// the overlapping top-level keys so incremental migrations can spot stale +// duplicates. +func TestLoadFromConfigPaths_ConflictWarningNamesKeys(t *testing.T) { + capture := func(t *testing.T) *strings.Builder { + t.Helper() + var buf strings.Builder + prev := slog.Default() + slog.SetDefault(slog.New(slog.NewJSONHandler(&buf, &slog.HandlerOptions{Level: slog.LevelWarn}))) + t.Cleanup(func() { slog.SetDefault(prev) }) + return &buf + } + + t.Run("names overlapping keys", func(t *testing.T) { + buf := capture(t) + tmpDir := t.TempDir() + jsonPath := filepath.Join(tmpDir, "crush.json") + rcPath := filepath.Join(tmpDir, "crushrc") + require.NoError(t, os.WriteFile(jsonPath, []byte(`{"options":{"debug":true},"providers":{}}`), 0o644)) + require.NoError(t, os.WriteFile(rcPath, []byte("option debug true\n"), 0o644)) + + _, _, err := loadFromConfigPaths(context.Background(), []string{jsonPath, rcPath}) + require.NoError(t, err) + require.Contains(t, buf.String(), "crushrc taking precedence") + require.Contains(t, buf.String(), `"conflicting_keys":"options"`) + }) + + t.Run("no warning when nothing overlaps", func(t *testing.T) { + buf := capture(t) + tmpDir := t.TempDir() + jsonPath := filepath.Join(tmpDir, "crush.json") + rcPath := filepath.Join(tmpDir, "crushrc") + require.NoError(t, os.WriteFile(jsonPath, []byte(`{"providers":{}}`), 0o644)) + require.NoError(t, os.WriteFile(rcPath, []byte("option debug true\n"), 0o644)) + + _, _, err := loadFromConfigPaths(context.Background(), []string{jsonPath, rcPath}) + require.NoError(t, err) + require.NotContains(t, buf.String(), "crushrc taking precedence", + "disjoint coexistence should not warn") + }) +} + +// testStore wraps a Config in a minimal ConfigStore for testing. +func testStore(cfg *Config) *ConfigStore { + return &ConfigStore{config: cfg} +} + +func TestConfig_setDefaults(t *testing.T) { + t.Run("sets default data directory", func(t *testing.T) { + cfg := &Config{} + workingDir := t.TempDir() + + cfg.setDefaults(workingDir, "") + + require.NotNil(t, cfg.Options) + require.NotNil(t, cfg.Options.TUI) + require.NotNil(t, cfg.Options.ContextPaths) + require.NotNil(t, cfg.Providers) + require.NotNil(t, cfg.Models) + require.NotNil(t, cfg.LSP) + require.NotNil(t, cfg.MCP) + require.Equal(t, filepath.Join(workingDir, ".crush"), cfg.Options.DataDirectory) + require.Equal(t, "AGENTS.md", cfg.Options.InitializeAs) + // DiffMode is deliberately left empty: the permissions dialog treats + // the zero value as "pick split or unified based on terminal width". + require.Empty(t, cfg.Options.TUI.DiffMode) + require.Equal(t, ScrollbarDefault, cfg.Options.TUI.Scrollbar) + require.Equal(t, ExitBannerDefault, cfg.Options.TUI.ExitBanner) + for _, path := range defaultContextPaths { + require.Contains(t, cfg.Options.ContextPaths, path) + } + }) + + t.Run("sets TUI defaults only when unset", func(t *testing.T) { + cfg := &Config{} + workingDir := t.TempDir() + + cfg.setDefaults(workingDir, "") + + require.Empty(t, cfg.Options.TUI.DiffMode) + require.Equal(t, ScrollbarDefault, cfg.Options.TUI.Scrollbar) + require.Equal(t, ExitBannerDefault, cfg.Options.TUI.ExitBanner) + + cfg.Options.TUI.DiffMode = DiffModeSplit + cfg.Options.TUI.Scrollbar = ScrollbarNever + cfg.Options.TUI.ExitBanner = ExitBannerCompact + cfg.setDefaults(workingDir, "") + + require.Equal(t, DiffModeSplit, cfg.Options.TUI.DiffMode) + require.Equal(t, ScrollbarNever, cfg.Options.TUI.Scrollbar) + require.Equal(t, ExitBannerCompact, cfg.Options.TUI.ExitBanner) + }) + + t.Run("prunes orphaned OAuth token MCP entries but keeps real ones", func(t *testing.T) { + cfg := &Config{ + MCP: map[string]MCPConfig{ + "orphan": {OAuthToken: &oauth.Token{AccessToken: "stale"}}, + "real-http": {Type: MCPHttp, URL: "https://example.com/mcp", OAuthToken: &oauth.Token{AccessToken: "live"}}, + "real-stdio": {Type: MCPStdio, Command: "npx"}, + "malformed": {Command: "npx"}, // missing type but has a command: surface the error, don't prune + }, + } + + cfg.setDefaults(t.TempDir(), "") + + require.NotContains(t, cfg.MCP, "orphan", "orphaned token entry should be pruned") + require.Contains(t, cfg.MCP, "real-http") + require.Contains(t, cfg.MCP, "real-stdio") + require.Contains(t, cfg.MCP, "malformed", "malformed entry should survive so its error surfaces") + }) + + t.Run("resolves relative configured data directory from working directory", func(t *testing.T) { + cfg := &Config{Options: &Options{DataDirectory: "."}} + workingDir := filepath.Join(t.TempDir(), "worktree") + + cfg.setDefaults(workingDir, "") + + require.Equal(t, workingDir, cfg.Options.DataDirectory) + }) + + t.Run("resolves relative flag data directory from working directory", func(t *testing.T) { + cfg := &Config{} + workingDir := filepath.Join(t.TempDir(), "worktree") + + cfg.setDefaults(workingDir, "./state") + + require.Equal(t, filepath.Join(workingDir, "state"), cfg.Options.DataDirectory) + }) + + t.Run("preserves absolute configured data directory", func(t *testing.T) { + // Use a platform-appropriate absolute path so the test runs + // the same way on POSIX and Windows. + absDir := filepath.Join(t.TempDir(), "data") + cfg := &Config{Options: &Options{DataDirectory: absDir}} + + cfg.setDefaults(filepath.Join(t.TempDir(), "worktree"), "") + + require.Equal(t, absDir, cfg.Options.DataDirectory) + }) + + t.Run("workspace merge re-entry keeps an absolute data directory", func(t *testing.T) { + // Simulate the load and reload paths: defaults are applied + // twice with the data directory potentially carried through + // from an earlier merge as a relative string. + workingDir := filepath.Join(t.TempDir(), "worktree") + cfg := &Config{} + cfg.setDefaults(workingDir, "") + + // Workspace JSON sets data_directory to a relative value; the + // merge replaces the struct, then setDefaults runs again. + cfg.Options.DataDirectory = "./state" + cfg.setDefaults(workingDir, "") + + require.True(t, filepath.IsAbs(cfg.Options.DataDirectory), + "data directory must remain absolute after re-merge, got %q", + cfg.Options.DataDirectory) + require.Equal(t, filepath.Join(workingDir, "state"), cfg.Options.DataDirectory) + }) + + t.Run("does not adopt .crush from a parent project", func(t *testing.T) { + parent := t.TempDir() + + // .crush in the parent: it should not be reused by the child + // because there is no git context joining them. + require.NoError(t, os.Mkdir(filepath.Join(parent, defaultDataDirectory), 0o755)) + + child := filepath.Join(parent, "child") + require.NoError(t, os.Mkdir(child, 0o755)) + + cfg := &Config{} + cfg.setDefaults(child, "") + + require.Equal( + t, + filepath.Clean(filepath.Join(child, defaultDataDirectory)), + filepath.Clean(cfg.Options.DataDirectory), + ) + }) + + t.Run("does not climb out of git worktree to find .crush", func(t *testing.T) { + if _, err := exec.LookPath("git"); err != nil { + t.Skip("git not available") + } + + parent := t.TempDir() + + // Stray .crush above the worktree root. + require.NoError(t, os.Mkdir(filepath.Join(parent, defaultDataDirectory), 0o755)) + + worktree := filepath.Join(parent, "worktree") + require.NoError(t, os.Mkdir(worktree, 0o755)) + + sub := filepath.Join(worktree, "pkg") + require.NoError(t, os.Mkdir(sub, 0o755)) + + // Make worktree a real git repo so the boundary detection + // resolves to it, mirroring what happens with linked worktrees + // in real usage. + gitInit := exec.CommandContext(t.Context(), "git", "init", "-q") + gitInit.Dir = worktree + require.NoError(t, gitInit.Run()) + + cfg := &Config{} + cfg.setDefaults(sub, "") + + // Resolve symlinks because TempDir on macOS sits under /var + // which is a symlink to /private/var. The data directory has + // not been created yet, so resolve its parent and join. + gotDir, gotName := filepath.Split(cfg.Options.DataDirectory) + gotEvalDir, err := filepath.EvalSymlinks(filepath.Clean(gotDir)) + require.NoError(t, err) + gotEval := filepath.Join(gotEvalDir, gotName) + + strayEval, err := filepath.EvalSymlinks(filepath.Join(parent, defaultDataDirectory)) + require.NoError(t, err) + require.NotEqual(t, strayEval, gotEval, "must not adopt parent .crush") + + subEval, err := filepath.EvalSymlinks(sub) + require.NoError(t, err) + require.Equal(t, filepath.Join(subEval, defaultDataDirectory), gotEval) + }) +} + +func TestConfig_configureProviders(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$OPENAI_API_KEY", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + } + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{ + "OPENAI_API_KEY": "test-key", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + require.Equal(t, 1, cfg.Providers.Len()) + + // We want to make sure that we keep the configured API key as a placeholder + pc, _ := cfg.Providers.Get("openai") + require.Equal(t, "$OPENAI_API_KEY", pc.APIKey) +} + +func TestConfig_configureProvidersWithOverride(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$OPENAI_API_KEY", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + } + + cfg := &Config{ + Providers: csync.NewMap[string, ProviderConfig](), + } + cfg.Providers.Set("openai", ProviderConfig{ + APIKey: "xyz", + BaseURL: "https://api.openai.com/v2", + Models: []catwalk.Model{ + { + ID: "test-model", + Name: "Updated", + }, + { + ID: "another-model", + }, + }, + }) + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{ + "OPENAI_API_KEY": "test-key", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + require.Equal(t, 1, cfg.Providers.Len()) + + // We want to make sure that we keep the configured API key as a placeholder + pc, _ := cfg.Providers.Get("openai") + require.Equal(t, "xyz", pc.APIKey) + require.Equal(t, "https://api.openai.com/v2", pc.BaseURL) + require.Len(t, pc.Models, 2) + require.Equal(t, "Updated", pc.Models[0].Name) +} + +func TestConfig_configureProvidersWithNewProvider(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$OPENAI_API_KEY", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + } + + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "xyz", + BaseURL: "https://api.someendpoint.com/v2", + Models: []catwalk.Model{ + { + ID: "test-model", + }, + }, + }, + }), + } + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{ + "OPENAI_API_KEY": "test-key", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + // Should be to because of the env variable + require.Equal(t, cfg.Providers.Len(), 2) + + // We want to make sure that we keep the configured API key as a placeholder + pc, _ := cfg.Providers.Get("custom") + require.Equal(t, "xyz", pc.APIKey) + // Make sure we set the ID correctly + require.Equal(t, "custom", pc.ID) + require.Equal(t, "https://api.someendpoint.com/v2", pc.BaseURL) + require.Len(t, pc.Models, 1) + + _, ok := cfg.Providers.Get("openai") + require.True(t, ok, "OpenAI provider should still be present") +} + +func TestConfig_configureProvidersBedrockWithCredentials(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: catwalk.InferenceProviderBedrock, + APIKey: "", + APIEndpoint: "", + Models: []catwalk.Model{{ + ID: "anthropic.claude-sonnet-4-20250514-v1:0", + }}, + }, + } + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{ + "AWS_ACCESS_KEY_ID": "test-key-id", + "AWS_SECRET_ACCESS_KEY": "test-secret-key", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + require.Equal(t, cfg.Providers.Len(), 1) + + bedrockProvider, ok := cfg.Providers.Get("bedrock") + require.True(t, ok, "Bedrock provider should be present") + require.Len(t, bedrockProvider.Models, 1) + require.Equal(t, "anthropic.claude-sonnet-4-20250514-v1:0", bedrockProvider.Models[0].ID) +} + +func TestConfig_configureProvidersBedrockWithoutCredentials(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: catwalk.InferenceProviderBedrock, + APIKey: "", + APIEndpoint: "", + Models: []catwalk.Model{{ + ID: "anthropic.claude-sonnet-4-20250514-v1:0", + }}, + }, + } + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + // Provider should not be configured without credentials + require.Equal(t, cfg.Providers.Len(), 0) +} + +func TestConfig_configureProvidersVertexAIWithCredentials(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: catwalk.InferenceProviderVertexAI, + APIKey: "", + APIEndpoint: "", + Models: []catwalk.Model{{ + ID: "gemini-pro", + }}, + }, + } + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{ + "VERTEXAI_PROJECT": "test-project", + "VERTEXAI_LOCATION": "us-central1", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + require.Equal(t, cfg.Providers.Len(), 1) + + vertexProvider, ok := cfg.Providers.Get("vertexai") + require.True(t, ok, "VertexAI provider should be present") + require.Len(t, vertexProvider.Models, 1) + require.Equal(t, "gemini-pro", vertexProvider.Models[0].ID) + require.Equal(t, "test-project", vertexProvider.ExtraParams["project"]) + require.Equal(t, "us-central1", vertexProvider.ExtraParams["location"]) +} + +func TestConfig_configureProvidersVertexAIWithoutCredentials(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: catwalk.InferenceProviderVertexAI, + APIKey: "", + APIEndpoint: "", + Models: []catwalk.Model{{ + ID: "gemini-pro", + }}, + }, + } + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{ + "GOOGLE_GENAI_USE_VERTEXAI": "false", + "GOOGLE_CLOUD_PROJECT": "test-project", + "GOOGLE_CLOUD_LOCATION": "us-central1", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + // Provider should not be configured without proper credentials + require.Equal(t, cfg.Providers.Len(), 0) +} + +func TestConfig_configureProvidersVertexAIMissingProject(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: catwalk.InferenceProviderVertexAI, + APIKey: "", + APIEndpoint: "", + Models: []catwalk.Model{{ + ID: "gemini-pro", + }}, + }, + } + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{ + "GOOGLE_GENAI_USE_VERTEXAI": "true", + "GOOGLE_CLOUD_LOCATION": "us-central1", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + // Provider should not be configured without project + require.Equal(t, cfg.Providers.Len(), 0) +} + +func TestConfig_configureProvidersSetProviderID(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$OPENAI_API_KEY", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + } + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{ + "OPENAI_API_KEY": "test-key", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + require.Equal(t, cfg.Providers.Len(), 1) + + // Provider ID should be set + pc, _ := cfg.Providers.Get("openai") + require.Equal(t, "openai", pc.ID) +} + +func TestConfig_EnabledProviders(t *testing.T) { + t.Run("all providers enabled", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "openai": { + ID: "openai", + APIKey: "key1", + Disable: false, + }, + "anthropic": { + ID: "anthropic", + APIKey: "key2", + Disable: false, + }, + }), + } + + enabled := cfg.EnabledProviders() + require.Len(t, enabled, 2) + }) + + t.Run("some providers disabled", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "openai": { + ID: "openai", + APIKey: "key1", + Disable: false, + }, + "anthropic": { + ID: "anthropic", + APIKey: "key2", + Disable: true, + }, + }), + } + + enabled := cfg.EnabledProviders() + require.Len(t, enabled, 1) + require.Equal(t, "openai", enabled[0].ID) + }) + + t.Run("empty providers map", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMap[string, ProviderConfig](), + } + + enabled := cfg.EnabledProviders() + require.Len(t, enabled, 0) + }) +} + +func TestConfig_IsConfigured(t *testing.T) { + t.Run("returns true when at least one provider is enabled", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "openai": { + ID: "openai", + APIKey: "key1", + Disable: false, + }, + }), + } + + require.True(t, cfg.IsConfigured()) + }) + + t.Run("returns false when no providers are configured", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMap[string, ProviderConfig](), + } + + require.False(t, cfg.IsConfigured()) + }) + + t.Run("returns false when all providers are disabled", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "openai": { + ID: "openai", + APIKey: "key1", + Disable: true, + }, + "anthropic": { + ID: "anthropic", + APIKey: "key2", + Disable: true, + }, + }), + } + + require.False(t, cfg.IsConfigured()) + }) +} + +func TestConfig_setupAgentsWithNoDisabledTools(t *testing.T) { + cfg := &Config{ + Options: &Options{ + DisabledTools: []string{}, + }, + } + + cfg.SetupAgents() + coderAgent, ok := cfg.Agents[AgentCoder] + require.True(t, ok) + assert.Equal(t, allToolNames(), coderAgent.AllowedTools) + + taskAgent, ok := cfg.Agents[AgentTask] + require.True(t, ok) + assert.Equal(t, []string{"lsp_symbols", "lsp_definition", "lsp_call_hierarchy", "glob", "grep", "ls", "sourcegraph", "view"}, taskAgent.AllowedTools) +} + +func TestConfig_setupAgentsWithDisabledTools(t *testing.T) { + cfg := &Config{ + Options: &Options{ + DisabledTools: []string{ + "edit", + "download", + "grep", + }, + }, + } + + cfg.SetupAgents() + coderAgent, ok := cfg.Agents[AgentCoder] + require.True(t, ok) + + assert.Equal(t, []string{"agent", "bash", "crush_info", "crush_logs", "job_output", "job_kill", "multiedit", "lsp_diagnostics", "lsp_references", "lsp_restart", "lsp_symbols", "lsp_definition", "lsp_call_hierarchy", "lsp_rename", "lsp_replace_symbol", "fetch", "agentic_fetch", "glob", "ls", "question", "sourcegraph", "todos", "view", "write", "list_mcp_resources", "read_mcp_resource"}, coderAgent.AllowedTools) + + taskAgent, ok := cfg.Agents[AgentTask] + require.True(t, ok) + assert.Equal(t, []string{"lsp_symbols", "lsp_definition", "lsp_call_hierarchy", "glob", "ls", "sourcegraph", "view"}, taskAgent.AllowedTools) +} + +func TestConfig_setupAgentsWithEveryReadOnlyToolDisabled(t *testing.T) { + cfg := &Config{ + Options: &Options{ + DisabledTools: []string{ + "glob", + "grep", + "ls", + "lsp_call_hierarchy", + "lsp_definition", + "lsp_symbols", + "sourcegraph", + "view", + }, + }, + } + + cfg.SetupAgents() + coderAgent, ok := cfg.Agents[AgentCoder] + require.True(t, ok) + assert.Equal(t, []string{"agent", "bash", "crush_info", "crush_logs", "job_output", "job_kill", "download", "edit", "multiedit", "lsp_diagnostics", "lsp_references", "lsp_restart", "lsp_rename", "lsp_replace_symbol", "fetch", "agentic_fetch", "question", "todos", "write", "list_mcp_resources", "read_mcp_resource"}, coderAgent.AllowedTools) + + taskAgent, ok := cfg.Agents[AgentTask] + require.True(t, ok) + assert.Len(t, taskAgent.AllowedTools, 0) +} + +func TestConfig_configureProvidersWithDisabledProvider(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$OPENAI_API_KEY", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + } + + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "openai": { + Disable: true, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{ + "OPENAI_API_KEY": "test-key", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + + require.Equal(t, cfg.Providers.Len(), 1) + prov, exists := cfg.Providers.Get("openai") + require.True(t, exists) + require.True(t, prov.Disable) +} + +func TestConfig_configureProvidersCustomProviderValidation(t *testing.T) { + t.Run("custom provider with missing API key is allowed, but not known providers", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + BaseURL: "https://api.custom.com/v1", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + "openai": { + APIKey: "$MISSING", + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + require.Equal(t, cfg.Providers.Len(), 1) + _, exists := cfg.Providers.Get("custom") + require.True(t, exists) + }) + + t.Run("custom provider with missing BaseURL is removed", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + require.Equal(t, cfg.Providers.Len(), 0) + _, exists := cfg.Providers.Get("custom") + require.False(t, exists) + }) + + t.Run("custom provider with shell-expanded BaseURL is kept", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: "$(echo https://api.custom.com/v1)", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + require.Equal(t, 1, cfg.Providers.Len()) + _, exists := cfg.Providers.Get("custom") + require.True(t, exists, "a BaseURL supplied via shell expansion must pass endpoint validation") + }) + + t.Run("custom provider with env var BaseURL is kept", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: "$CUSTOM_API_URL", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{ + "CUSTOM_API_URL": "https://api.custom.com/v1", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + require.Equal(t, 1, cfg.Providers.Len()) + _, exists := cfg.Providers.Get("custom") + require.True(t, exists, "a BaseURL supplied via an env variable must pass endpoint validation") + }) + + t.Run("custom provider whose BaseURL resolves to empty is removed", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: "$MISSING_API_URL", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + require.Equal(t, 0, cfg.Providers.Len()) + _, exists := cfg.Providers.Get("custom") + require.False(t, exists, "a raw-non-empty BaseURL that resolves to empty must be rejected") + }) + + t.Run("BaseURL resolving to empty with disable_default_providers returns an error", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: "$MISSING_API_URL", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + }), + } + cfg.setDefaults("/tmp", "") + cfg.Options.DisableDefaultProviders = true + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.Error(t, err, "the sole custom provider is invalid after resolution, so configuration must fail") + }) + + t.Run("custom provider with no models attempts discovery and is removed on failure", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: "https://api.custom.com/v1", + Models: []catwalk.Model{}, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + // Discovery fails (unreachable URL) so provider is removed. + require.Equal(t, 0, cfg.Providers.Len()) + _, exists := cfg.Providers.Get("custom") + require.False(t, exists) + }) + + t.Run("custom provider with no models and discover_models:false is removed", func(t *testing.T) { + discoverFalse := false + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: "https://api.custom.com/v1", + Models: []catwalk.Model{}, + AutoDiscoverModels: &discoverFalse, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + require.Equal(t, 0, cfg.Providers.Len()) + _, exists := cfg.Providers.Get("custom") + require.False(t, exists) + }) + + t.Run("custom provider with models and discover_models:true merges discovered models", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data": [ + {"id": "existing-model", "object": "model"}, + {"id": "discovered-model", "object": "model"} + ]}`)) + })) + defer server.Close() + + discoverTrue := true + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: server.URL + "/v1", + Models: []catwalk.Model{ + {ID: "existing-model", Name: "My Custom Name", ContextWindow: 200000}, + }, + AutoDiscoverModels: &discoverTrue, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + require.Equal(t, 1, cfg.Providers.Len()) + p, exists := cfg.Providers.Get("custom") + require.True(t, exists) + require.Len(t, p.Models, 2) + + // User-specified model keeps its custom fields. + require.Equal(t, "existing-model", p.Models[0].ID) + require.Equal(t, "My Custom Name", p.Models[0].Name) + require.Equal(t, int64(200000), p.Models[0].ContextWindow) + + // Discovered model is appended. + require.Equal(t, "discovered-model", p.Models[1].ID) + }) + + t.Run("custom provider with models and no discover_models uses only listed models", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: "https://api.custom.com/v1", + Models: []catwalk.Model{ + {ID: "my-model", Name: "My Model"}, + }, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + require.Equal(t, 1, cfg.Providers.Len()) + p, exists := cfg.Providers.Get("custom") + require.True(t, exists) + require.Len(t, p.Models, 1) + require.Equal(t, "my-model", p.Models[0].ID) + }) + + t.Run("custom provider with no models auto-discovers successfully", func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data": [ + {"id": "auto-model-a", "object": "model"}, + {"id": "auto-model-b", "object": "model"} + ]}`)) + })) + defer server.Close() + + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: server.URL + "/v1", + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + require.Equal(t, 1, cfg.Providers.Len()) + p, exists := cfg.Providers.Get("custom") + require.True(t, exists) + require.Len(t, p.Models, 2) + require.Equal(t, "auto-model-a", p.Models[0].ID) + require.Equal(t, "auto-model-b", p.Models[1].ID) + }) + + t.Run("custom provider with unsupported type is removed", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: "https://api.custom.com/v1", + Type: "unsupported", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + require.Equal(t, cfg.Providers.Len(), 0) + _, exists := cfg.Providers.Get("custom") + require.False(t, exists) + }) + + t.Run("valid custom provider is kept and ID is set", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: "https://api.custom.com/v1", + Type: catwalk.TypeOpenAI, + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + require.Equal(t, cfg.Providers.Len(), 1) + customProvider, exists := cfg.Providers.Get("custom") + require.True(t, exists) + require.Equal(t, "custom", customProvider.ID) + require.Equal(t, "test-key", customProvider.APIKey) + require.Equal(t, "https://api.custom.com/v1", customProvider.BaseURL) + }) + + t.Run("custom anthropic provider is supported", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom-anthropic": { + APIKey: "test-key", + BaseURL: "https://api.anthropic.com/v1", + Type: catwalk.TypeAnthropic, + Models: []catwalk.Model{{ + ID: "claude-3-sonnet", + }}, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + require.Equal(t, cfg.Providers.Len(), 1) + customProvider, exists := cfg.Providers.Get("custom-anthropic") + require.True(t, exists) + require.Equal(t, "custom-anthropic", customProvider.ID) + require.Equal(t, "test-key", customProvider.APIKey) + require.Equal(t, "https://api.anthropic.com/v1", customProvider.BaseURL) + require.Equal(t, catwalk.TypeAnthropic, customProvider.Type) + }) + + t.Run("disabled custom provider is removed", func(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: "https://api.custom.com/v1", + Type: catwalk.TypeOpenAI, + Disable: true, + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + require.Equal(t, cfg.Providers.Len(), 0) + _, exists := cfg.Providers.Get("custom") + require.False(t, exists) + }) +} + +func TestConfig_configureProvidersEnhancedCredentialValidation(t *testing.T) { + t.Run("VertexAI provider removed when credentials missing with existing config", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: catwalk.InferenceProviderVertexAI, + APIKey: "", + APIEndpoint: "", + Models: []catwalk.Model{{ + ID: "gemini-pro", + }}, + }, + } + + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "vertexai": { + BaseURL: "custom-url", + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{ + "GOOGLE_GENAI_USE_VERTEXAI": "false", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + + require.Equal(t, cfg.Providers.Len(), 0) + _, exists := cfg.Providers.Get("vertexai") + require.False(t, exists) + }) + + t.Run("Bedrock provider removed when AWS credentials missing with existing config", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: catwalk.InferenceProviderBedrock, + APIKey: "", + APIEndpoint: "", + Models: []catwalk.Model{{ + ID: "anthropic.claude-sonnet-4-20250514-v1:0", + }}, + }, + } + + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "bedrock": { + BaseURL: "custom-url", + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + + require.Equal(t, cfg.Providers.Len(), 0) + _, exists := cfg.Providers.Get("bedrock") + require.False(t, exists) + }) + + t.Run("provider removed when API key missing with existing config", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$MISSING_API_KEY", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + } + + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "openai": { + BaseURL: "custom-url", + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + + require.Equal(t, cfg.Providers.Len(), 0) + _, exists := cfg.Providers.Get("openai") + require.False(t, exists) + }) + + t.Run("known provider should still be added if the endpoint is missing the client will use default endpoints", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$OPENAI_API_KEY", + APIEndpoint: "$MISSING_ENDPOINT", + Models: []catwalk.Model{{ + ID: "test-model", + }}, + }, + } + + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "openai": { + APIKey: "test-key", + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{ + "OPENAI_API_KEY": "test-key", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + + require.Equal(t, cfg.Providers.Len(), 1) + _, exists := cfg.Providers.Get("openai") + require.True(t, exists) + }) +} + +func TestConfig_defaultModelSelection(t *testing.T) { + t.Run("default behavior uses the default models for given provider", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "abc", + DefaultLargeModelID: "large-model", + DefaultSmallModelID: "small-model", + Models: []catwalk.Model{ + { + ID: "large-model", + DefaultMaxTokens: 1000, + }, + { + ID: "small-model", + DefaultMaxTokens: 500, + }, + }, + }, + } + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + + large, small, err := cfg.defaultModelSelection(knownProviders) + require.NoError(t, err) + require.Equal(t, "large-model", large.Model) + require.Equal(t, "openai", large.Provider) + require.Equal(t, int64(1000), large.MaxTokens) + require.Equal(t, "small-model", small.Model) + require.Equal(t, "openai", small.Provider) + require.Equal(t, int64(500), small.MaxTokens) + }) + t.Run("should error if no providers configured", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$MISSING_KEY", + DefaultLargeModelID: "large-model", + DefaultSmallModelID: "small-model", + Models: []catwalk.Model{ + { + ID: "large-model", + DefaultMaxTokens: 1000, + }, + { + ID: "small-model", + DefaultMaxTokens: 500, + }, + }, + }, + } + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + + _, _, err = cfg.defaultModelSelection(knownProviders) + require.Error(t, err) + }) + t.Run("should not error if model is missing", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "abc", + DefaultLargeModelID: "large-model", + DefaultSmallModelID: "small-model", + Models: []catwalk.Model{ + { + ID: "not-large-model", + DefaultMaxTokens: 1000, + }, + { + ID: "small-model", + DefaultMaxTokens: 500, + }, + }, + }, + } + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + _, _, err = cfg.defaultModelSelection(knownProviders) + require.NoError(t, err) + }) + + t.Run("should configure the default models with a custom provider", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$MISSING", // will not be included in the config + DefaultLargeModelID: "large-model", + DefaultSmallModelID: "small-model", + Models: []catwalk.Model{ + { + ID: "not-large-model", + DefaultMaxTokens: 1000, + }, + { + ID: "small-model", + DefaultMaxTokens: 500, + }, + }, + }, + } + + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: "https://api.custom.com/v1", + Models: []catwalk.Model{ + { + ID: "model", + DefaultMaxTokens: 600, + }, + }, + }, + }), + } + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + large, small, err := cfg.defaultModelSelection(knownProviders) + require.NoError(t, err) + require.Equal(t, "model", large.Model) + require.Equal(t, "custom", large.Provider) + require.Equal(t, int64(600), large.MaxTokens) + require.Equal(t, "model", small.Model) + require.Equal(t, "custom", small.Provider) + require.Equal(t, int64(600), small.MaxTokens) + }) + + t.Run("should fail if no model configured", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$MISSING", // will not be included in the config + DefaultLargeModelID: "large-model", + DefaultSmallModelID: "small-model", + Models: []catwalk.Model{ + { + ID: "not-large-model", + DefaultMaxTokens: 1000, + }, + { + ID: "small-model", + DefaultMaxTokens: 500, + }, + }, + }, + } + + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: "https://api.custom.com/v1", + Models: []catwalk.Model{}, + }, + }), + } + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + _, _, err = cfg.defaultModelSelection(knownProviders) + require.Error(t, err) + }) + t.Run("should use the default provider first", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "set", + DefaultLargeModelID: "large-model", + DefaultSmallModelID: "small-model", + Models: []catwalk.Model{ + { + ID: "large-model", + DefaultMaxTokens: 1000, + }, + { + ID: "small-model", + DefaultMaxTokens: 500, + }, + }, + }, + } + + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "custom": { + APIKey: "test-key", + BaseURL: "https://api.custom.com/v1", + Models: []catwalk.Model{ + { + ID: "large-model", + DefaultMaxTokens: 1000, + }, + }, + }, + }), + } + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + large, small, err := cfg.defaultModelSelection(knownProviders) + require.NoError(t, err) + require.Equal(t, "large-model", large.Model) + require.Equal(t, "openai", large.Provider) + require.Equal(t, int64(1000), large.MaxTokens) + require.Equal(t, "small-model", small.Model) + require.Equal(t, "openai", small.Provider) + require.Equal(t, int64(500), small.MaxTokens) + }) +} + +func TestConfig_configureProvidersDisableDefaultProviders(t *testing.T) { + t.Run("when enabled, ignores all default providers and requires full specification", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$OPENAI_API_KEY", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ + ID: "gpt-4", + }}, + }, + } + + // User references openai but doesn't fully specify it (no base_url, no + // models). This should be rejected because disable_default_providers + // treats all providers as custom. + cfg := &Config{ + Options: &Options{ + DisableDefaultProviders: true, + }, + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "openai": { + APIKey: "$OPENAI_API_KEY", + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{ + "OPENAI_API_KEY": "test-key", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.ErrorContains(t, err, "no custom providers") + + // openai should NOT be present because it lacks base_url and models. + require.Equal(t, 0, cfg.Providers.Len()) + _, exists := cfg.Providers.Get("openai") + require.False(t, exists, "openai should not be present without full specification") + }) + + t.Run("when enabled, fully specified providers work", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$OPENAI_API_KEY", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ + ID: "gpt-4", + }}, + }, + } + + // User fully specifies their provider. + cfg := &Config{ + Options: &Options{ + DisableDefaultProviders: true, + }, + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "my-llm": { + APIKey: "$MY_API_KEY", + BaseURL: "https://my-llm.example.com/v1", + Models: []catwalk.Model{{ + ID: "my-model", + }}, + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{ + "MY_API_KEY": "test-key", + "OPENAI_API_KEY": "test-key", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + + // Only fully specified provider should be present. + require.Equal(t, 1, cfg.Providers.Len()) + provider, exists := cfg.Providers.Get("my-llm") + require.True(t, exists, "my-llm should be present") + require.Equal(t, "https://my-llm.example.com/v1", provider.BaseURL) + require.Len(t, provider.Models, 1) + + // Default openai should NOT be present. + _, exists = cfg.Providers.Get("openai") + require.False(t, exists, "openai should not be present") + }) + + t.Run("when disabled, includes all known providers with valid credentials", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$OPENAI_API_KEY", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ + ID: "gpt-4", + }}, + }, + { + ID: "anthropic", + APIKey: "$ANTHROPIC_API_KEY", + APIEndpoint: "https://api.anthropic.com/v1", + Models: []catwalk.Model{{ + ID: "claude-3", + }}, + }, + } + + // User only configures openai, both API keys are available, but option + // is disabled. + cfg := &Config{ + Options: &Options{ + DisableDefaultProviders: false, + }, + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "openai": { + APIKey: "$OPENAI_API_KEY", + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{ + "OPENAI_API_KEY": "test-key", + "ANTHROPIC_API_KEY": "test-key", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + + // Both providers should be present. + require.Equal(t, 2, cfg.Providers.Len()) + _, exists := cfg.Providers.Get("openai") + require.True(t, exists, "openai should be present") + _, exists = cfg.Providers.Get("anthropic") + require.True(t, exists, "anthropic should be present") + }) + + t.Run("when enabled, provider missing models attempts discovery but still triggers no-custom-providers error", func(t *testing.T) { + cfg := &Config{ + Options: &Options{ + DisableDefaultProviders: true, + }, + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "my-llm": { + APIKey: "test-key", + BaseURL: "https://my-llm.example.com/v1", + Models: []catwalk.Model{}, // No models. + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.ErrorContains(t, err, "no custom providers") + + // Discovery fails (unreachable URL) so provider is removed. + require.Equal(t, 0, cfg.Providers.Len()) + }) + + t.Run("when enabled, provider missing base_url is rejected", func(t *testing.T) { + cfg := &Config{ + Options: &Options{ + DisableDefaultProviders: true, + }, + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "my-llm": { + APIKey: "test-key", + Models: []catwalk.Model{{ID: "model"}}, + // No BaseURL. + }, + }), + } + cfg.setDefaults("/tmp", "") + + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, []catwalk.Provider{}) + require.ErrorContains(t, err, "no custom providers") + + // Provider should be rejected for missing base_url. + require.Equal(t, 0, cfg.Providers.Len()) + }) +} + +func TestConfig_setDefaultsDisableDefaultProvidersEnvVar(t *testing.T) { + t.Run("sets option from environment variable", func(t *testing.T) { + t.Setenv("CRUSH_DISABLE_DEFAULT_PROVIDERS", "true") + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + + require.True(t, cfg.Options.DisableDefaultProviders) + }) + + t.Run("does not override when env var is not set", func(t *testing.T) { + cfg := &Config{ + Options: &Options{ + DisableDefaultProviders: true, + }, + } + cfg.setDefaults("/tmp", "") + + require.True(t, cfg.Options.DisableDefaultProviders) + }) +} + +func TestConfig_configureSelectedModels(t *testing.T) { + t.Run("reload mode should not persist fallback defaults", func(t *testing.T) { + dir := t.TempDir() + globalPath := filepath.Join(dir, "crush.json") + require.NoError(t, os.WriteFile(globalPath, []byte(`{"models":{"large":{"provider":"ghost","model":"missing"}}}`), 0o600)) + + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "abc", + DefaultLargeModelID: "large-model", + DefaultSmallModelID: "small-model", + Models: []catwalk.Model{ + {ID: "large-model", DefaultMaxTokens: 1000}, + {ID: "small-model", DefaultMaxTokens: 500}, + }, + }, + } + + cfg := &Config{ + Models: map[SelectedModelType]SelectedModel{ + SelectedModelTypeLarge: {Provider: "ghost", Model: "missing"}, + }, + } + cfg.setDefaults(dir, "") + store := &ConfigStore{config: cfg, globalDataPath: globalPath} + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), store, env, resolver, knownProviders) + require.NoError(t, err) + + resolved, resolveErr := resolveSelectedModels(cfg, knownProviders) + require.NoError(t, resolveErr) + cfg.Models[SelectedModelTypeLarge] = resolved.Large + cfg.Models[SelectedModelTypeSmall] = resolved.Small + + // In-memory falls back to default. + require.True(t, resolved.LargeFallback) + require.Equal(t, "openai", cfg.Models[SelectedModelTypeLarge].Provider) + require.Equal(t, "large-model", cfg.Models[SelectedModelTypeLarge].Model) + + // Disk remains unchanged (resolveSelectedModels never persists). + data, readErr := os.ReadFile(globalPath) + require.NoError(t, readErr) + require.Contains(t, string(data), `"provider":"ghost"`) + require.Contains(t, string(data), `"model":"missing"`) + }) + t.Run("should override defaults", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "abc", + DefaultLargeModelID: "large-model", + DefaultSmallModelID: "small-model", + Models: []catwalk.Model{ + { + ID: "larger-model", + DefaultMaxTokens: 2000, + }, + { + ID: "large-model", + DefaultMaxTokens: 1000, + }, + { + ID: "small-model", + DefaultMaxTokens: 500, + }, + }, + }, + } + + cfg := &Config{ + Models: map[SelectedModelType]SelectedModel{ + "large": { + Model: "larger-model", + }, + }, + } + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + + resolved, resolveErr := resolveSelectedModels(cfg, knownProviders) + require.NoError(t, resolveErr) + cfg.Models[SelectedModelTypeLarge] = resolved.Large + cfg.Models[SelectedModelTypeSmall] = resolved.Small + large := cfg.Models[SelectedModelTypeLarge] + small := cfg.Models[SelectedModelTypeSmall] + require.Equal(t, "larger-model", large.Model) + require.Equal(t, "openai", large.Provider) + require.Equal(t, int64(2000), large.MaxTokens) + require.Equal(t, "small-model", small.Model) + require.Equal(t, "openai", small.Provider) + require.Equal(t, int64(500), small.MaxTokens) + }) + t.Run("should be possible to use multiple providers", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "abc", + DefaultLargeModelID: "large-model", + DefaultSmallModelID: "small-model", + Models: []catwalk.Model{ + { + ID: "large-model", + DefaultMaxTokens: 1000, + }, + { + ID: "small-model", + DefaultMaxTokens: 500, + }, + }, + }, + { + ID: "anthropic", + APIKey: "abc", + DefaultLargeModelID: "a-large-model", + DefaultSmallModelID: "a-small-model", + Models: []catwalk.Model{ + { + ID: "a-large-model", + DefaultMaxTokens: 1000, + }, + { + ID: "a-small-model", + DefaultMaxTokens: 200, + }, + }, + }, + } + + cfg := &Config{ + Models: map[SelectedModelType]SelectedModel{ + "small": { + Model: "a-small-model", + Provider: "anthropic", + MaxTokens: 300, + }, + }, + } + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + + resolved, resolveErr := resolveSelectedModels(cfg, knownProviders) + require.NoError(t, resolveErr) + cfg.Models[SelectedModelTypeLarge] = resolved.Large + cfg.Models[SelectedModelTypeSmall] = resolved.Small + large := cfg.Models[SelectedModelTypeLarge] + small := cfg.Models[SelectedModelTypeSmall] + require.Equal(t, "large-model", large.Model) + require.Equal(t, "openai", large.Provider) + require.Equal(t, int64(1000), large.MaxTokens) + require.Equal(t, "a-small-model", small.Model) + require.Equal(t, "anthropic", small.Provider) + require.Equal(t, int64(300), small.MaxTokens) + }) + + t.Run("should override the max tokens only", func(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "abc", + DefaultLargeModelID: "large-model", + DefaultSmallModelID: "small-model", + Models: []catwalk.Model{ + { + ID: "large-model", + DefaultMaxTokens: 1000, + }, + { + ID: "small-model", + DefaultMaxTokens: 500, + }, + }, + }, + } + + cfg := &Config{ + Models: map[SelectedModelType]SelectedModel{ + "large": { + MaxTokens: 100, + }, + }, + } + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + + resolved, resolveErr := resolveSelectedModels(cfg, knownProviders) + require.NoError(t, resolveErr) + cfg.Models[SelectedModelTypeLarge] = resolved.Large + cfg.Models[SelectedModelTypeSmall] = resolved.Small + large := cfg.Models[SelectedModelTypeLarge] + require.Equal(t, "large-model", large.Model) + require.Equal(t, "openai", large.Provider) + require.Equal(t, int64(100), large.MaxTokens) + }) + t.Run("resolve and persist fallback under writeMu does not deadlock", func(t *testing.T) { + dir := t.TempDir() + globalPath := filepath.Join(dir, "crush.json") + require.NoError(t, os.WriteFile(globalPath, []byte(`{}`), 0o600)) + + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "abc", + DefaultLargeModelID: "large-model", + DefaultSmallModelID: "small-model", + Models: []catwalk.Model{ + {ID: "large-model", DefaultMaxTokens: 1000}, + {ID: "small-model", DefaultMaxTokens: 500}, + }, + }, + } + + cfg := &Config{ + Models: map[SelectedModelType]SelectedModel{ + SelectedModelTypeLarge: {Provider: "openai", Model: "this-model-does-not-exist"}, + SelectedModelTypeSmall: {Provider: "openai", Model: "also-does-not-exist"}, + }, + } + cfg.setDefaults(dir, "") + store := &ConfigStore{config: cfg, globalDataPath: globalPath} + env := env.NewFromMap(map[string]string{}) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), store, env, resolver, knownProviders) + require.NoError(t, err) + + // Simulate the Load path: resolve (pure), then persist fallbacks + // under writeMu using updateLocked. Before the refactor, the + // combined configureSelectedModels(persist=true) self-deadlocked + // because UpdatePreferredModel re-acquired writeMu. + done := make(chan error, 1) + go func() { + resolved, resolveErr := resolveSelectedModels(cfg, knownProviders) + if resolveErr != nil { + done <- resolveErr + return + } + cfg.Models[SelectedModelTypeLarge] = resolved.Large + cfg.Models[SelectedModelTypeSmall] = resolved.Small + + store.writeMu.Lock() + defer store.writeMu.Unlock() + if resolved.LargeFallback { + if err := store.updateLocked(ScopeGlobal, func(c *Config) map[string]any { + return store.updatePreferredModelFields(c, SelectedModelTypeLarge, resolved.Large) + }); err != nil { + done <- err + return + } + } + if resolved.SmallFallback { + if err := store.updateLocked(ScopeGlobal, func(c *Config) map[string]any { + return store.updatePreferredModelFields(c, SelectedModelTypeSmall, resolved.Small) + }); err != nil { + done <- err + return + } + } + done <- nil + }() + + select { + case err := <-done: + require.NoError(t, err) + // Should have fallen back to defaults. + require.Equal(t, "large-model", cfg.Models[SelectedModelTypeLarge].Model) + require.Equal(t, "small-model", cfg.Models[SelectedModelTypeSmall].Model) + case <-time.After(5 * time.Second): + t.Fatal("resolve + persist deadlocked under writeMu") + } + }) +} + +func TestConfig_configureProviders_HyperAPIKeyFromEnv(t *testing.T) { + // Test that HYPER_API_KEY environment variable works without config + knownProviders := []catwalk.Provider{ + { + ID: "hyper", + APIKey: "", // No API key in provider definition + DefaultLargeModelID: "large-model", + DefaultSmallModelID: "small-model", + Models: []catwalk.Model{ + { + ID: "large-model", + DefaultMaxTokens: 1000, + }, + { + ID: "small-model", + DefaultMaxTokens: 500, + }, + }, + }, + } + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + env := env.NewFromMap(map[string]string{ + "HYPER_API_KEY": "env-api-key", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + require.Equal(t, 1, cfg.Providers.Len()) + + // Verify Hyper provider is configured with the env var API key + pc, ok := cfg.Providers.Get("hyper") + require.True(t, ok, "Hyper provider should be configured") + require.Equal(t, "env-api-key", pc.APIKey) + require.Equal(t, "env-api-key", pc.APIKeyTemplate) +} + +func TestConfig_configureProviders_HyperAPIKeyFromConfigOverrides(t *testing.T) { + // Test that config API key takes precedence when HYPER_API_KEY is also set + knownProviders := []catwalk.Provider{ + { + ID: "hyper", + APIKey: "provider-api-key", + DefaultLargeModelID: "large-model", + DefaultSmallModelID: "small-model", + Models: []catwalk.Model{ + { + ID: "large-model", + DefaultMaxTokens: 1000, + }, + { + ID: "small-model", + DefaultMaxTokens: 500, + }, + }, + }, + } + + // User has Hyper configured with an API key + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "hyper": { + APIKey: "config-api-key", + }, + }), + } + cfg.setDefaults("/tmp", "") + + // But they also have HYPER_API_KEY set - env var should take precedence + env := env.NewFromMap(map[string]string{ + "HYPER_API_KEY": "env-api-key", + }) + resolver := NewShellVariableResolver(env) + err := cfg.configureProviders(context.Background(), testStore(cfg), env, resolver, knownProviders) + require.NoError(t, err) + require.Equal(t, 1, cfg.Providers.Len()) + + // Verify env var takes precedence (as per requirements) + pc, ok := cfg.Providers.Get("hyper") + require.True(t, ok, "Hyper provider should be configured") + require.Equal(t, "env-api-key", pc.APIKey) +} + +// TestConfig_configureProviders_ProviderHeaderResolveError verifies +// that a failing $(cmd) in a provider header fails the provider load +// with a clear message that names the offending header. Provider +// headers share the MCP error contract. +func TestConfig_configureProviders_ProviderHeaderResolveError(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$OPENAI_API_KEY", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ID: "test-model"}}, + }, + } + + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "openai": { + ExtraHeaders: map[string]string{ + // Failing $(...) — inner command exits 1. Must + // propagate as an error, not a silent truncation. + "X-Broken": "$(false)", + }, + }, + }), + } + cfg.setDefaults("/tmp", "") + + testEnv := env.NewFromMap(map[string]string{ + "OPENAI_API_KEY": "test-key", + "PATH": os.Getenv("PATH"), + }) + resolver := NewShellVariableResolver(testEnv) + + err := cfg.configureProviders(context.Background(), testStore(cfg), testEnv, resolver, knownProviders) + require.Error(t, err, "failing $(cmd) in a header must fail the provider load") + require.Contains(t, err.Error(), "X-Broken", "error must name the offending header") +} + +// TestConfig_configureProviders_CatwalkDefaultWithUnsetVarLoads +// verifies that a Catwalk-style default header like +// "OpenAI-Organization": "$OPENAI_ORG_ID" loads cleanly under lenient +// nounset (unset → "" → header dropped), and does not fail the load +// or leave the literal template on the wire. +func TestConfig_configureProviders_CatwalkDefaultWithUnsetVarLoads(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$OPENAI_API_KEY", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ID: "test-model"}}, + DefaultHeaders: map[string]string{ + "OpenAI-Organization": "$OPENAI_ORG_ID", + }, + }, + } + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + + testEnv := env.NewFromMap(map[string]string{ + "OPENAI_API_KEY": "test-key", + "PATH": os.Getenv("PATH"), + }) + resolver := NewShellVariableResolver(testEnv) + + err := cfg.configureProviders(context.Background(), testStore(cfg), testEnv, resolver, knownProviders) + require.NoError(t, err, "optional env-gated header must not fail the load") + + pc, ok := cfg.Providers.Get("openai") + require.True(t, ok, "openai provider must still be configured") + _, present := pc.ExtraHeaders["OpenAI-Organization"] + require.False(t, present, "header whose value resolves to empty must be absent") +} + +// TestConfig_configureProviders_LiteralEmptyHeaderDropped pins design +// decision #18 for the literal case: a user-authored +// "X-Custom": "" in extra_headers is absent from the resolved map. +// Applies to both known- and custom-provider paths; this test +// exercises the custom-provider loop. +func TestConfig_configureProviders_LiteralEmptyHeaderDropped(t *testing.T) { + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "my-llm": { + APIKey: "test-key", + BaseURL: "https://my-llm.example.com/v1", + Type: catwalk.TypeOpenAI, + Models: []catwalk.Model{{ID: "m"}}, + ExtraHeaders: map[string]string{ + "X-Custom": "", + "X-Kept": "present", + }, + }, + }), + } + cfg.setDefaults("/tmp", "") + + testEnv := env.NewFromMap(map[string]string{ + "PATH": os.Getenv("PATH"), + }) + resolver := NewShellVariableResolver(testEnv) + + err := cfg.configureProviders(context.Background(), testStore(cfg), testEnv, resolver, []catwalk.Provider{}) + require.NoError(t, err) + + pc, ok := cfg.Providers.Get("my-llm") + require.True(t, ok) + _, present := pc.ExtraHeaders["X-Custom"] + require.False(t, present, "literal empty-string header must be dropped") + require.Equal(t, "present", pc.ExtraHeaders["X-Kept"]) +} + +// TestConfig_configureProviders_EchoEmptyHeaderDropped pins design +// decision #18 for the non-failing empty case: $(echo) exits 0 with +// empty output, resolves cleanly to "", and must be dropped the same +// way an unset bare $VAR is. Exercises the known-provider loop. +func TestConfig_configureProviders_EchoEmptyHeaderDropped(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$OPENAI_API_KEY", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ID: "test-model"}}, + DefaultHeaders: map[string]string{ + "X-Empty": "$(echo)", + "X-Kept": "present", + }, + }, + } + + cfg := &Config{} + cfg.setDefaults("/tmp", "") + + testEnv := env.NewFromMap(map[string]string{ + "OPENAI_API_KEY": "test-key", + "PATH": os.Getenv("PATH"), + }) + resolver := NewShellVariableResolver(testEnv) + + err := cfg.configureProviders(context.Background(), testStore(cfg), testEnv, resolver, knownProviders) + require.NoError(t, err) + + pc, ok := cfg.Providers.Get("openai") + require.True(t, ok) + _, present := pc.ExtraHeaders["X-Empty"] + require.False(t, present, "$(echo) → empty → header must be dropped") + require.Equal(t, "present", pc.ExtraHeaders["X-Kept"]) +} + +// TestConfig_configureProviders_UnsetAPIKeySkipsProvider verifies that +// under the lenient-nounset shell resolver, $UNSET_API_KEY expands to +// ("", nil) rather than ("", err), and the existing +// `v == "" || err != nil` skip path at load.go:331 still drops the +// provider. The slog.Warn line is emitted on the same +// path but is not asserted here — internal/config/load_test.go's +// TestMain replaces the default slog handler with an io.Discard +// writer, so capturing that log line would require mid-test handler +// swapping and a sync.Mutex dance that adds more flake surface than +// signal. The observable outcome (provider absent from the map) is +// what downstream code — model picker, agent wiring — actually reads, +// so that's what we pin. +func TestConfig_configureProviders_UnsetAPIKeySkipsProvider(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$SOMETHING_UNSET", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ID: "test-model"}}, + }, + } + + // Existing user config for this known provider so the load.go:332 + // `if configExists` branch fires and actually calls Providers.Del. + // Without it the provider was never in the map to begin with and + // the test would pass trivially. + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "openai": {BaseURL: "custom-url"}, + }), + } + cfg.setDefaults("/tmp", "") + + testEnv := env.NewFromMap(map[string]string{ + "PATH": os.Getenv("PATH"), + }) + resolver := NewShellVariableResolver(testEnv) + + err := cfg.configureProviders(context.Background(), testStore(cfg), testEnv, resolver, knownProviders) + require.NoError(t, err, "skip path must not surface as a load error") + + require.Equal(t, 0, cfg.Providers.Len(), "provider with unset API key must be skipped") + _, exists := cfg.Providers.Get("openai") + require.False(t, exists) +} + +// TestConfig_configureProviders_FailingAPIKeyCmdSkipsProvider pins +// that the two failure modes for APIKey — ("", nil) from an unset var +// under lenient nounset and ("", err) from a failing $(cmd) — are +// equivalent for the skip outcome at load.go:331. The `v == "" || +// err != nil` check fires on either branch; this test locks in that +// equivalence so a future refactor that splits the check into two +// paths doesn't accidentally start propagating $(false) as a load +// error while keeping unset-var as a silent skip (or vice versa). +func TestConfig_configureProviders_FailingAPIKeyCmdSkipsProvider(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: "openai", + APIKey: "$(false)", + APIEndpoint: "https://api.openai.com/v1", + Models: []catwalk.Model{{ID: "test-model"}}, + }, + } + + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "openai": {BaseURL: "custom-url"}, + }), + } + cfg.setDefaults("/tmp", "") + + testEnv := env.NewFromMap(map[string]string{ + "PATH": os.Getenv("PATH"), + }) + resolver := NewShellVariableResolver(testEnv) + + err := cfg.configureProviders(context.Background(), testStore(cfg), testEnv, resolver, knownProviders) + require.NoError(t, err, "failing $(cmd) in API key must skip provider, not fail load") + + require.Equal(t, 0, cfg.Providers.Len(), "provider with failing $(cmd) API key must be skipped") + _, exists := cfg.Providers.Get("openai") + require.False(t, exists) +} + +// TestConfig_configureProviders_UnsetAzureEndpointSkipsProvider pins +// the same contract on the Azure path at load.go:287 — APIEndpoint is +// the field that gates Azure and goes through the same +// `v == "" || err != nil` skip check. Covered here so both branches +// of the shared skip pattern (APIKey default path and APIEndpoint +// Azure path) are tested; a future refactor that unifies them can +// rely on these two tests to catch drift. +func TestConfig_configureProviders_UnsetAzureEndpointSkipsProvider(t *testing.T) { + knownProviders := []catwalk.Provider{ + { + ID: catwalk.InferenceProviderAzure, + APIKey: "test-key", + APIEndpoint: "$UNSET_AZURE_ENDPOINT", + Models: []catwalk.Model{{ID: "test-model"}}, + }, + } + + cfg := &Config{ + Providers: csync.NewMapFrom(map[string]ProviderConfig{ + "azure": {BaseURL: ""}, + }), + } + cfg.setDefaults("/tmp", "") + + testEnv := env.NewFromMap(map[string]string{ + "PATH": os.Getenv("PATH"), + }) + resolver := NewShellVariableResolver(testEnv) + + err := cfg.configureProviders(context.Background(), testStore(cfg), testEnv, resolver, knownProviders) + require.NoError(t, err) + + require.Equal(t, 0, cfg.Providers.Len(), "azure provider with unset endpoint must be skipped") + _, exists := cfg.Providers.Get("azure") + require.False(t, exists) +} + +func TestConfig_LoadFromBytes_Env(t *testing.T) { + data := []byte(`{"env": {"AWS_PROFILE": "my-profile", "AWS_REGION": "us-west-2"}}`) + + loadedConfig, err := loadFromBytes([][]byte{data}) + + require.NoError(t, err) + require.NotNil(t, loadedConfig.Env) + require.Equal(t, "my-profile", loadedConfig.Env["AWS_PROFILE"]) + require.Equal(t, "us-west-2", loadedConfig.Env["AWS_REGION"]) +} + +func TestConfig_LoadFromBytes_EnvMerge(t *testing.T) { + data1 := []byte(`{"env": {"AWS_PROFILE": "first", "AWS_REGION": "us-east-1"}}`) + data2 := []byte(`{"env": {"AWS_PROFILE": "second"}}`) + + loadedConfig, err := loadFromBytes([][]byte{data1, data2}) + + require.NoError(t, err) + require.NotNil(t, loadedConfig.Env) + require.Equal(t, "second", loadedConfig.Env["AWS_PROFILE"]) + require.Equal(t, "us-east-1", loadedConfig.Env["AWS_REGION"]) +} diff --git a/internal/config/lsp_defaults_test.go b/internal/config/lsp_defaults_test.go new file mode 100644 index 0000000000000000000000000000000000000000..f04666597803f390fe25699a3a8bf9aba44b68be --- /dev/null +++ b/internal/config/lsp_defaults_test.go @@ -0,0 +1,35 @@ +package config + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestApplyLSPDefaults(t *testing.T) { + t.Parallel() + + // Create a config with an LSP that should get defaults + config := &Config{ + LSP: map[string]LSPConfig{ + "gopls": { + Command: "gopls", // This should get defaults from powernap + }, + "custom": { + Command: "custom-lsp", + RootMarkers: []string{"custom.toml"}, // This should keep its explicit config + }, + }, + } + + // Apply defaults + config.applyLSPDefaults() + + // Check that gopls got defaults (it should have some root markers now) + goplsConfig := config.LSP["gopls"] + require.NotEmpty(t, goplsConfig.RootMarkers, "gopls should have received default root markers") + + // Check that custom LSP kept its explicit config + customConfig := config.LSP["custom"] + require.Equal(t, []string{"custom.toml"}, customConfig.RootMarkers, "custom LSP should keep its explicit root markers") +} diff --git a/internal/config/openai_test.go b/internal/config/openai_test.go new file mode 100644 index 0000000000000000000000000000000000000000..b7930bffebd3bac8b720f26365a0cec28c9cdae4 --- /dev/null +++ b/internal/config/openai_test.go @@ -0,0 +1,220 @@ +package config + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "charm.land/catwalk/pkg/catwalk" + "github.com/charmbracelet/crush/internal/csync" + "github.com/charmbracelet/crush/internal/oauth" + "github.com/stretchr/testify/require" +) + +func TestGetModelIncludesChatGPTModels(t *testing.T) { + t.Parallel() + + providers := csync.NewMap[string, ProviderConfig]() + providers.Set("openai", ProviderConfig{ + Models: []catwalk.Model{{ID: "gpt-5.1"}}, + ChatGPTModels: []catwalk.Model{{ID: "gpt-5.1-codex", Name: "GPT-5.1 Codex"}}, + }) + cfg := &Config{Providers: providers} + + model := cfg.GetModel("openai", "gpt-5.1-codex") + require.NotNil(t, model) + require.Equal(t, "GPT-5.1 Codex", model.Name) + + require.NotNil(t, cfg.GetModel("openai", "gpt-5.1")) + require.Nil(t, cfg.GetModel("openai", "missing")) +} + +func TestTokenFields(t *testing.T) { + t.Parallel() + + t.Run("openai keeps the api key", func(t *testing.T) { + t.Parallel() + + fields := tokenFields("openai", &oauth.Token{AccessToken: "at"}) + require.Contains(t, fields, "providers.openai.oauth") + require.NotContains(t, fields, "providers.openai.api_key") + }) + + t.Run("other providers mirror the access token", func(t *testing.T) { + t.Parallel() + + fields := tokenFields("hyper", &oauth.Token{AccessToken: "at"}) + require.Contains(t, fields, "providers.hyper.oauth") + require.Equal(t, "at", fields["providers.hyper.api_key"]) + }) +} + +func TestTokenAccountIDRoundTrip(t *testing.T) { + t.Parallel() + + // The account ID is persisted with the token in crush.json; verify + // it survives the JSON round trip loadTokenFromDisk performs. + token := &oauth.Token{ + AccessToken: "at", + RefreshToken: "rt", + IDToken: "idt", + AccountID: "acct-42", + ExpiresIn: 3600, + ExpiresAt: 4102444800, + } + + data, err := json.Marshal(token) + require.NoError(t, err) + + var restored oauth.Token + require.NoError(t, json.Unmarshal(data, &restored)) + require.Equal(t, "acct-42", restored.AccountID) + require.Equal(t, "idt", restored.IDToken) + require.Equal(t, "rt", restored.RefreshToken) +} + +// testResolver mimics the shell resolver for templates: variables +// resolve to empty when unset, literals pass through. +type testResolver struct{} + +func (testResolver) ResolveValue(v string) (string, error) { + if strings.HasPrefix(v, "$") { + return "", nil + } + return v, nil +} + +func TestHasAPIKey(t *testing.T) { + t.Parallel() + + t.Run("empty", func(t *testing.T) { + t.Parallel() + + require.False(t, (&ProviderConfig{}).HasAPIKey(testResolver{})) + }) + + t.Run("literal key", func(t *testing.T) { + t.Parallel() + + require.True(t, (&ProviderConfig{APIKey: "sk-key"}).HasAPIKey(testResolver{})) + }) + + t.Run("unresolved template", func(t *testing.T) { + t.Parallel() + + require.False(t, (&ProviderConfig{APIKey: "$OPENAI_API_KEY"}).HasAPIKey(testResolver{})) + }) +} + +// TestSetProviderAPIKeyOpenAIIsEitherOr proves the OpenAI provider holds +// exactly one credential: a ChatGPT login retires a previously entered +// API key, and entering an API key retires a previous ChatGPT login. +// Copilot keeps mirroring the access token into api_key as before. +func TestSetProviderAPIKeyOpenAIIsEitherOr(t *testing.T) { + // Not parallel: t.Setenv below. + + // Point config discovery at the test sandbox: the write below + // triggers an auto-reload, which must not pick up the developer's + // real crush.json. + t.Setenv("CRUSH_GLOBAL_CONFIG", t.TempDir()) + t.Setenv("CRUSH_GLOBAL_DATA", t.TempDir()) + t.Setenv("XDG_DATA_HOME", t.TempDir()) + + token := &oauth.Token{ + AccessToken: "chatgpt-at", + RefreshToken: "chatgpt-rt", + ExpiresIn: 3600, + ExpiresAt: time.Now().Add(time.Hour).Unix(), + } + + newStore := func(t *testing.T, providerID string, initialConfig string) *ConfigStore { + t.Helper() + dir := t.TempDir() + configPath := filepath.Join(dir, "crush.json") + require.NoError(t, os.WriteFile(configPath, []byte(initialConfig), 0o600)) + + // The in-memory config mirrors what a real load would produce. + // SetConfigFields auto-reloads from configPath, so any state the + // assertions expect to survive must be on disk. + providers := csync.NewMap[string, ProviderConfig]() + require.NoError(t, json.Unmarshal([]byte(initialConfig), &struct { + Providers *csync.Map[string, ProviderConfig] `json:"providers"` + }{Providers: providers})) + + return &ConfigStore{ + config: &Config{Providers: providers}, + globalDataPath: configPath, + workingDir: dir, + fetchOpenAIModels: func(context.Context, *oauth.Token) ([]catwalk.Model, error) { + return []catwalk.Model{{ID: "gpt-5.1-codex"}}, nil + }, + } + } + + t.Run("chatgpt login retires the api key", func(t *testing.T) { + store := newStore(t, "openai", `{ + "providers": { + "openai": { + "id": "openai", + "api_key": "sk-keep", + "models": [{"id": "gpt-5.1", "name": "GPT-5.1"}] + } + } + }`) + + require.NoError(t, store.SetProviderAPIKey(ScopeGlobal, "openai", token)) + + pc, ok := store.Config().Providers.Get("openai") + require.True(t, ok) + require.Empty(t, pc.APIKey, "the ChatGPT login replaces the API key") + require.Equal(t, token, pc.OAuthToken) + require.Equal(t, "gpt-5.1-codex", pc.ChatGPTModels[0].ID, "the subscription catalog lands in its own field") + require.Equal(t, "gpt-5.1", pc.Models[0].ID, "the API catalog is untouched") + + disk, err := os.ReadFile(store.globalDataPath) + require.NoError(t, err) + require.NotContains(t, string(disk), "sk-keep", "the retired key is gone from the config file") + require.Contains(t, string(disk), "chatgpt-rt", "the login is persisted") + }) + + t.Run("api key retires the chatgpt login", func(t *testing.T) { + store := newStore(t, "openai", `{ + "providers": { + "openai": { + "id": "openai", + "api_key": "", + "oauth": {"access_token": "old-at", "refresh_token": "old-rt"}, + "chatgpt_models": [{"id": "gpt-5.6-luna"}] + } + } + }`) + + require.NoError(t, store.SetProviderAPIKey(ScopeGlobal, "openai", "sk-new")) + + pc, ok := store.Config().Providers.Get("openai") + require.True(t, ok) + require.Equal(t, "sk-new", pc.APIKey) + require.Nil(t, pc.OAuthToken, "the API key replaces the ChatGPT login") + require.Empty(t, pc.ChatGPTModels, "the subscription catalog goes with it") + + disk, err := os.ReadFile(store.globalDataPath) + require.NoError(t, err) + require.NotContains(t, string(disk), "old-rt", "the retired login is gone from the config file") + require.NotContains(t, string(disk), "gpt-5.6-luna") + }) + + t.Run("copilot", func(t *testing.T) { + store := newStore(t, "copilot", `{"providers":{"copilot":{"id":"copilot"}}}`) + + require.NoError(t, store.SetProviderAPIKey(ScopeGlobal, "copilot", token)) + + pc, ok := store.Config().Providers.Get("copilot") + require.True(t, ok) + require.Equal(t, "chatgpt-at", pc.APIKey, "copilot still mirrors the access token into api_key") + require.Equal(t, token, pc.OAuthToken) + }) +} diff --git a/internal/config/options_request_timeout_test.go b/internal/config/options_request_timeout_test.go new file mode 100644 index 0000000000000000000000000000000000000000..f40619678caaf208827d5ec4da95a7f83e75921a --- /dev/null +++ b/internal/config/options_request_timeout_test.go @@ -0,0 +1,63 @@ +package config + +import ( + "encoding/json" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func ptr[T any](v T) *T { return &v } + +func TestOptionsGetRequestTimeout(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + options *Options + expected time.Duration + }{ + { + name: "nil options", + options: nil, + expected: DefaultRequestTimeout, + }, + { + name: "unset field", + options: &Options{}, + expected: DefaultRequestTimeout, + }, + { + name: "disabled", + options: &Options{RequestTimeout: ptr(0)}, + expected: 0, + }, + { + name: "negative disables too", + options: &Options{RequestTimeout: ptr(-1)}, + expected: 0, + }, + { + name: "seconds to duration", + options: &Options{RequestTimeout: ptr(300)}, + expected: 5 * time.Minute, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tt.expected, tt.options.GetRequestTimeout()) + }) + } +} + +func TestOptionsRequestTimeoutFromJSON(t *testing.T) { + t.Parallel() + + var cfg Config + require.NoError(t, json.Unmarshal([]byte(`{"options":{"request_timeout":300}}`), &cfg)) + require.Equal(t, 300, *cfg.Options.RequestTimeout) + require.Equal(t, 5*time.Minute, cfg.Options.GetRequestTimeout()) +} diff --git a/internal/config/provider_empty_test.go b/internal/config/provider_empty_test.go new file mode 100644 index 0000000000000000000000000000000000000000..9bc62f5c3141d239aaadc3947dce539a4dcf4810 --- /dev/null +++ b/internal/config/provider_empty_test.go @@ -0,0 +1,39 @@ +package config + +import ( + "context" + "os" + "testing" + + "charm.land/catwalk/pkg/catwalk" + "github.com/stretchr/testify/require" +) + +type emptyProviderClient struct{} + +func (m *emptyProviderClient) GetProviders(context.Context, string) ([]catwalk.Provider, error) { + return []catwalk.Provider{}, nil +} + +// TestCatwalkSync_GetEmptyResultFromClient tests that when the client returns +// an empty list, we fall back to cached providers and return an error. +func TestCatwalkSync_GetEmptyResultFromClient(t *testing.T) { + t.Parallel() + + tmpDir := t.TempDir() + path := tmpDir + "/providers.json" + + syncer := &catwalkSync{} + client := &emptyProviderClient{} + + syncer.Init(client, path, true) + + providers, err := syncer.Get(t.Context()) + require.Error(t, err) + require.Contains(t, err.Error(), "empty providers list from catwalk") + require.NotEmpty(t, providers) // Should have embedded providers as fallback. + + // Check that no cache file was created for empty results. + _, statErr := os.Stat(path) + require.True(t, os.IsNotExist(statErr), "Cache file should not exist for empty results") +} diff --git a/internal/config/recent_models_test.go b/internal/config/recent_models_test.go new file mode 100644 index 0000000000000000000000000000000000000000..4506cd336f3b1e130f844e46614ede3d23cdb112 --- /dev/null +++ b/internal/config/recent_models_test.go @@ -0,0 +1,158 @@ +package config + +import ( + "encoding/json" + "io/fs" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +// readConfigJSON reads and unmarshals the JSON config file at path. +func readConfigJSON(t *testing.T, path string) map[string]any { + t.Helper() + baseDir := filepath.Dir(path) + fileName := filepath.Base(path) + b, err := fs.ReadFile(os.DirFS(baseDir), fileName) + require.NoError(t, err) + var out map[string]any + require.NoError(t, json.Unmarshal(b, &out)) + return out +} + +// readRecentModels reads the recent_models section from the config file. +func readRecentModels(t *testing.T, path string) map[string]any { + t.Helper() + out := readConfigJSON(t, path) + rm, ok := out["recent_models"].(map[string]any) + require.True(t, ok) + return rm +} + +// testStoreWithPath creates a ConfigStore backed by a Config for recent model tests. +func testStoreWithPath(cfg *Config, dir string) *ConfigStore { + return &ConfigStore{ + config: cfg, + globalDataPath: filepath.Join(dir, "config.json"), + } +} + +// configWithRecents builds a Config seeded with the given recent models for +// the large type, for exercising the pure nextRecentModels helper. +func configWithRecents(recents ...SelectedModel) *Config { + return &Config{ + RecentModels: map[SelectedModelType][]SelectedModel{ + SelectedModelTypeLarge: recents, + }, + } +} + +func TestNextRecentModels_AddsToFront(t *testing.T) { + t.Parallel() + + cfg := configWithRecents() + updated, changed := nextRecentModels(cfg, SelectedModelTypeLarge, SelectedModel{Provider: "openai", Model: "gpt-4o"}) + require.True(t, changed) + require.Equal(t, []SelectedModel{{Provider: "openai", Model: "gpt-4o"}}, updated) +} + +func TestNextRecentModels_DedupeAndMoveToFront(t *testing.T) { + t.Parallel() + + cfg := configWithRecents( + SelectedModel{Provider: "anthropic", Model: "claude"}, + SelectedModel{Provider: "openai", Model: "gpt-4o"}, + ) + updated, changed := nextRecentModels(cfg, SelectedModelTypeLarge, SelectedModel{Provider: "openai", Model: "gpt-4o"}) + require.True(t, changed) + require.Equal(t, []SelectedModel{ + {Provider: "openai", Model: "gpt-4o"}, + {Provider: "anthropic", Model: "claude"}, + }, updated) +} + +func TestNextRecentModels_TrimsToMax(t *testing.T) { + t.Parallel() + + var seed []SelectedModel + for _, id := range []string{"m5", "m4", "m3", "m2", "m1"} { + seed = append(seed, SelectedModel{Provider: "p", Model: id}) + } + cfg := configWithRecents(seed...) + + updated, changed := nextRecentModels(cfg, SelectedModelTypeLarge, SelectedModel{Provider: "p", Model: "m6"}) + require.True(t, changed) + require.Len(t, updated, maxRecentModelsPerType) + require.Equal(t, SelectedModel{Provider: "p", Model: "m6"}, updated[0]) + require.Equal(t, SelectedModel{Provider: "p", Model: "m2"}, updated[maxRecentModelsPerType-1]) +} + +func TestNextRecentModels_SkipsEmptyValues(t *testing.T) { + t.Parallel() + + cfg := configWithRecents() + _, changed := nextRecentModels(cfg, SelectedModelTypeLarge, SelectedModel{Provider: "", Model: "m"}) + require.False(t, changed) + _, changed = nextRecentModels(cfg, SelectedModelTypeLarge, SelectedModel{Provider: "p", Model: ""}) + require.False(t, changed) +} + +func TestNextRecentModels_NoChangeWhenAlreadyFront(t *testing.T) { + t.Parallel() + + entry := SelectedModel{Provider: "openai", Model: "gpt-4o"} + cfg := configWithRecents(entry) + _, changed := nextRecentModels(cfg, SelectedModelTypeLarge, entry) + require.False(t, changed) +} + +func TestUpdatePreferredModel_PersistsModelAndRecents(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + cfg := &Config{} + cfg.setDefaults(dir, "") + store := testStoreWithPath(cfg, dir) + + sel := SelectedModel{Provider: "openai", Model: "gpt-4o"} + require.NoError(t, store.UpdatePreferredModel(ScopeGlobal, SelectedModelTypeLarge, sel)) + + // in-memory state (read through the store; copy-on-write publishes a + // new Config, so the seed cfg pointer is intentionally unchanged). + require.Equal(t, sel, store.Config().Models[SelectedModelTypeLarge]) + require.Len(t, store.Config().RecentModels[SelectedModelTypeLarge], 1) + + // persisted state + rm := readRecentModels(t, store.globalDataPath) + large, ok := rm[string(SelectedModelTypeLarge)].([]any) + require.True(t, ok) + require.Len(t, large, 1) + item := large[0].(map[string]any) + require.Equal(t, "openai", item["provider"]) + require.Equal(t, "gpt-4o", item["model"]) +} + +func TestUpdatePreferredModel_TypeIsolation(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + cfg := &Config{} + cfg.setDefaults(dir, "") + store := testStoreWithPath(cfg, dir) + + largeModel := SelectedModel{Provider: "openai", Model: "gpt-4o"} + smallModel := SelectedModel{Provider: "anthropic", Model: "claude"} + require.NoError(t, store.UpdatePreferredModel(ScopeGlobal, SelectedModelTypeLarge, largeModel)) + require.NoError(t, store.UpdatePreferredModel(ScopeGlobal, SelectedModelTypeSmall, smallModel)) + + // Adding to large leaves small untouched. + anotherLarge := SelectedModel{Provider: "google", Model: "gemini"} + require.NoError(t, store.UpdatePreferredModel(ScopeGlobal, SelectedModelTypeLarge, anotherLarge)) + + require.Len(t, store.Config().RecentModels[SelectedModelTypeLarge], 2) + require.Equal(t, anotherLarge, store.Config().RecentModels[SelectedModelTypeLarge][0]) + require.Len(t, store.Config().RecentModels[SelectedModelTypeSmall], 1) + require.Equal(t, smallModel, store.Config().RecentModels[SelectedModelTypeSmall][0]) +} diff --git a/internal/config/reload_crushrc_test.go b/internal/config/reload_crushrc_test.go new file mode 100644 index 0000000000000000000000000000000000000000..f4a5f0b187a73a8fc93fda2bb5f3a398f8a7a778 --- /dev/null +++ b/internal/config/reload_crushrc_test.go @@ -0,0 +1,137 @@ +package config_test + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "github.com/charmbracelet/crush/internal/config" + "github.com/stretchr/testify/require" +) + +// isolateReloadEnv points HOME/XDG/CRUSH_* at a throwaway directory so a +// reload only sees the config files the test writes. No t.Parallel(): these +// tests Setenv. +func isolateReloadEnv(t *testing.T) (workDir, dataDir string) { + t.Helper() + isolated := t.TempDir() + t.Setenv("HOME", isolated) + t.Setenv("XDG_CONFIG_HOME", filepath.Join(isolated, ".config")) + t.Setenv("XDG_DATA_HOME", filepath.Join(isolated, ".local", "share")) + t.Setenv("CRUSH_GLOBAL_CONFIG", filepath.Join(isolated, ".config", "crush")) + t.Setenv("CRUSH_GLOBAL_DATA", filepath.Join(isolated, ".local", "share", "crush")) + return t.TempDir(), t.TempDir() +} + +// TestReloadFromDisk_PicksUpEditedCrushrc verifies that a crushrc edit on +// disk is reflected in memory after a reload, matching how JSON config +// already behaves. +func TestReloadFromDisk_PicksUpEditedCrushrc(t *testing.T) { + workDir, dataDir := isolateReloadEnv(t) + rcPath := filepath.Join(workDir, "crushrc") + require.NoError(t, os.WriteFile(rcPath, []byte("option notifications bell\n"), 0o644)) + + store, err := config.Load(workDir, dataDir, false) + require.NoError(t, err) + require.Equal(t, "bell", store.Config().Options.Notifications) + + require.NoError(t, os.WriteFile(rcPath, []byte("option notifications osc\n"), 0o644)) + require.NoError(t, store.ReloadFromDisk(context.Background())) + require.Equal(t, "osc", store.Config().Options.Notifications, + "reload must re-execute the crushrc and pick up the new value") +} + +// TestReloadFromDisk_FailingCrushrcKeepsOldConfig verifies that a crushrc +// which fails to load during a reload does not clobber the in-memory config. +// The reload must return an error and leave the previously-loaded config +// intact (the swap happens only after a successful parse). +func TestReloadFromDisk_FailingCrushrcKeepsOldConfig(t *testing.T) { + workDir, dataDir := isolateReloadEnv(t) + rcPath := filepath.Join(workDir, "crushrc") + require.NoError(t, os.WriteFile(rcPath, []byte("option notifications bell\n"), 0o644)) + + store, err := config.Load(workDir, dataDir, false) + require.NoError(t, err) + require.Equal(t, "bell", store.Config().Options.Notifications) + + // Overwrite with a crushrc that errors at load time (unknown option key). + require.NoError(t, os.WriteFile(rcPath, []byte("option totally-bogus-key value\n"), 0o644)) + + err = store.ReloadFromDisk(context.Background()) + require.Error(t, err, "a failing crushrc must surface a reload error") + require.Equal(t, "bell", store.Config().Options.Notifications, + "in-memory config must be preserved when a reload fails") +} + +// TestReloadFromDisk_HangingCrushrcIsInterruptible is the core regression test +// for the reload-path blocker: config reloads run while holding the store's +// write lock, so a crushrc that blocks (busy loop, hung command substitution) +// must be interruptible via the reload context rather than wedging the store +// forever. A cancelled reload must return an error and leave the old config +// in place. The test bounds its own wait so a regression can't hang CI. +func TestReloadFromDisk_HangingCrushrcIsInterruptible(t *testing.T) { + workDir, dataDir := isolateReloadEnv(t) + rcPath := filepath.Join(workDir, "crushrc") + require.NoError(t, os.WriteFile(rcPath, []byte("option notifications bell\n"), 0o644)) + + store, err := config.Load(workDir, dataDir, false) + require.NoError(t, err) + require.Equal(t, "bell", store.Config().Options.Notifications) + + // Overwrite with a crushrc that never terminates. + require.NoError(t, os.WriteFile(rcPath, []byte("while true; do :; done\n"), 0o644)) + + ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond) + t.Cleanup(cancel) + + done := make(chan error, 1) + go func() { + done <- store.ReloadFromDisk(ctx) + }() + + select { + case err := <-done: + require.Error(t, err, "a cancelled reload must fail, not succeed") + case <-time.After(3 * time.Second): + t.Fatal("ReloadFromDisk did not return after context cancellation; the store would be wedged") + } + + require.Equal(t, "bell", store.Config().Options.Notifications, + "in-memory config must be preserved when a reload is cancelled") +} + +// TestLoad_TracksNotYetCreatedGlobalCrushrc verifies that config.Load tracks +// the global crushrc path even when the file does not exist yet, so a crushrc +// created after startup is detected as a staleness change. Previously only +// successfully-loaded paths were tracked, so a mid-session global crushrc went +// unnoticed until something else triggered a reload. +// +// Project-level crushrc files are discovered by walking the tree for existing +// files, so a not-yet-created project crushrc is still only picked up on the +// next reload; the global path is the common case this covers. +func TestLoad_TracksNotYetCreatedGlobalCrushrc(t *testing.T) { + workDir, dataDir := isolateReloadEnv(t) + globalRC := filepath.Join(t.TempDir(), "crushrc") + t.Setenv("CRUSH_GLOBAL_CONFIG", filepath.Dir(globalRC)) + + // A provider must be configured so Load runs past its early + // "not configured" return and reaches the staleness snapshot capture. + require.NoError(t, os.WriteFile( + filepath.Join(workDir, "crushrc"), + []byte("provider add openai --api-key k\n"), 0o644, + )) + + // Load with no global crushrc present. + store, err := config.Load(workDir, dataDir, false) + require.NoError(t, err) + require.False(t, store.ConfigStaleness().Dirty, "fresh load should be clean") + + // Create the global crushrc after startup. + require.NoError(t, os.WriteFile(globalRC, []byte("option debug true\n"), 0o644)) + + staleness := store.ConfigStaleness() + require.True(t, staleness.Dirty, "creating a global crushrc must be detected") + require.Contains(t, staleness.Changed, globalRC) +} diff --git a/internal/config/reload_hooks_test.go b/internal/config/reload_hooks_test.go new file mode 100644 index 0000000000000000000000000000000000000000..a21dfa89a0b7c61260a96f800acc34a18a2b086b --- /dev/null +++ b/internal/config/reload_hooks_test.go @@ -0,0 +1,102 @@ +package config_test + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/hooks" + "github.com/stretchr/testify/require" +) + +// TestReloadFromDisk_CompilesHookMatchers is a regression test for a bug +// where ReloadFromDisk dropped the compiled matcher regex on every hook, +// causing a matcher like "^bash$" to match every tool call after any +// SetConfigField-triggered reload. +// +// The assertion is phrased in terms of observable Runner behavior (not +// internal field presence) so it stays valid if the Runner later owns +// matcher compilation itself. +func TestReloadFromDisk_CompilesHookMatchers(t *testing.T) { + // No t.Parallel(): we Setenv HOME/XDG_CONFIG_HOME to isolate from the + // developer's real global config, which may define its own hooks. + isolated := t.TempDir() + t.Setenv("HOME", isolated) + t.Setenv("XDG_CONFIG_HOME", filepath.Join(isolated, ".config")) + t.Setenv("XDG_DATA_HOME", filepath.Join(isolated, ".local", "share")) + + workDir := t.TempDir() + dataDir := t.TempDir() + configPath := filepath.Join(workDir, "crush.json") + cfgJSON := `{ + "hooks": { + "PreToolUse": [ + {"matcher": "^bash$", "command": "exit 0"} + ] + } + }` + require.NoError(t, os.WriteFile(configPath, []byte(cfgJSON), 0o600)) + + store, err := config.Load(workDir, dataDir, false) + require.NoError(t, err) + + // Sanity: hook filtering works immediately after Load. + assertHookFilters(t, store) + + require.NoError(t, store.ReloadFromDisk(context.Background())) + + // The actual regression check: filtering must still work after a + // reload, not silently collapse to match-everything. + assertHookFilters(t, store) +} + +// assertHookFilters builds a Runner from the store's current hooks and +// verifies the "^bash$" matcher rejects a non-bash tool while accepting +// bash. +func assertHookFilters(t *testing.T, store *config.ConfigStore) { + t.Helper() + preHooks := store.Config().Hooks[hooks.EventPreToolUse] + require.Len(t, preHooks, 1) + + runner := hooks.NewRunner(preHooks, t.TempDir(), t.TempDir()) + + nonMatch, err := runner.Run(context.Background(), hooks.EventPreToolUse, "sess", "view", `{}`) + require.NoError(t, err) + require.Equal(t, 0, nonMatch.HookCount, "view must not match ^bash$ matcher") + + match, err := runner.Run(context.Background(), hooks.EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, 1, match.HookCount, "bash must match ^bash$ matcher") +} + +// TestSetConfigField_AutoReload_PreservesHookMatcherFiltering verifies the +// dominant real-world trigger path: config writes call autoReload, +// autoReload calls ReloadFromDisk, and hook matching must remain correct. +func TestSetConfigField_AutoReload_PreservesHookMatcherFiltering(t *testing.T) { + isolated := t.TempDir() + t.Setenv("HOME", isolated) + t.Setenv("XDG_CONFIG_HOME", filepath.Join(isolated, ".config")) + t.Setenv("XDG_DATA_HOME", filepath.Join(isolated, ".local", "share")) + + workDir := t.TempDir() + dataDir := t.TempDir() + configPath := filepath.Join(workDir, "crush.json") + cfgJSON := `{ + "hooks": { + "PreToolUse": [ + {"matcher": "^bash$", "command": "exit 0"} + ] + } + }` + require.NoError(t, os.WriteFile(configPath, []byte(cfgJSON), 0o600)) + + store, err := config.Load(workDir, dataDir, false) + require.NoError(t, err) + assertHookFilters(t, store) + + require.NoError(t, store.SetConfigField(config.ScopeGlobal, "options.debug", true)) + + assertHookFilters(t, store) +} diff --git a/internal/config/resolve_real_test.go b/internal/config/resolve_real_test.go new file mode 100644 index 0000000000000000000000000000000000000000..ab107c0e22f88e0155d0a3da9c2e5736d0403e89 --- /dev/null +++ b/internal/config/resolve_real_test.go @@ -0,0 +1,512 @@ +package config + +import ( + "fmt" + "maps" + "os" + "path/filepath" + "slices" + "strings" + "testing" + + "github.com/charmbracelet/crush/internal/env" + "github.com/stretchr/testify/require" +) + +// These tests exercise the full shell-expansion path (no mocks, +// no injected Expander) to catch regressions that only surface when +// internal/shell actually runs the value. Table-level unit tests with +// fake expanders live in resolve_test.go. + +// realShellResolver builds a resolver backed by a shell env that +// contains PATH + the caller-supplied overrides. Production callers +// get PATH for free via env.New(); these tests need it so $(cat ...) +// and similar inner commands can resolve. +func realShellResolver(vars map[string]string) VariableResolver { + m := map[string]string{"PATH": os.Getenv("PATH")} + maps.Copy(m, vars) + return NewShellVariableResolver(env.NewFromMap(m)) +} + +func writeTempFile(t *testing.T, content string) string { + t.Helper() + dir := t.TempDir() + p := filepath.Join(dir, "secret") + require.NoError(t, os.WriteFile(p, []byte(content), 0o600)) + return p +} + +// TestResolvedEnv_RealShell_Success covers the shell constructs the +// PLAN calls out: $(cat tempfile) with and without trailing newline, +// ${VAR:-default} for unset vars, literal-space preservation around +// $(...), nested parens, quoted args inside $(echo ...), and a +// glob-like literal round-tripping unchanged. +func TestResolvedEnv_RealShell_Success(t *testing.T) { + t.Parallel() + + // filepath.ToSlash so Windows temp paths (C:\Users\...) survive + // being injected into a shell command string — the embedded shell + // treats backslashes as escapes, forward slashes work on every OS. + withNL := filepath.ToSlash(writeTempFile(t, "token-with-nl\n")) + noNL := filepath.ToSlash(writeTempFile(t, "token-no-nl")) + + m := MCPConfig{ + Env: map[string]string{ + // POSIX strips exactly one trailing newline from $(...) + // output, so both forms land on the same value. + "TOK_NL": fmt.Sprintf("$(cat %s)", withNL), + "TOK_NO": fmt.Sprintf("$(cat %s)", noNL), + + // ${VAR:-default} must not error on unset: this is the + // opt-in escape hatch for "empty is fine". + "FALLBACK": "${MCP_MISSING:-fallback}", + + // Leading/trailing literal spaces around $(...) must be + // preserved — single-value contract, no field splitting. + "PADDED": " $(echo v) ", + + // ")" inside a quoted arg to echo is a regression guard + // for the old hand-rolled paren matcher. + "PAREN": `$(echo ")")`, + + // Embedded space inside a quoted arg must survive + // verbatim; no word-splitting side effect. + "SPACEY": `$(echo "a b")`, + + // Glob-like literals must not expand. + "GLOB": "*.go", + }, + } + + got, err := m.ResolvedEnv(realShellResolver(nil)) + require.NoError(t, err) + + // ResolvedEnv returns "KEY=value" sorted by key. + want := []string{ + "FALLBACK=fallback", + "GLOB=*.go", + "PADDED= v ", + "PAREN=)", + "SPACEY=a b", + "TOK_NL=token-with-nl", + "TOK_NO=token-no-nl", + } + require.Equal(t, want, got) +} + +// TestResolvedEnv_RealShell_DoesNotMutate pins that both success and +// failure paths leave m.Env untouched. Prior behaviour rewrote the +// value in place on error; that was the exact mechanism that shipped +// empty credentials to MCP servers. +func TestResolvedEnv_RealShell_DoesNotMutate(t *testing.T) { + t.Parallel() + + t.Run("success path leaves Env untouched", func(t *testing.T) { + t.Parallel() + m := MCPConfig{Env: map[string]string{"TOKEN": "$(echo shh)"}} + orig := maps.Clone(m.Env) + + _, err := m.ResolvedEnv(realShellResolver(nil)) + require.NoError(t, err) + require.Equal(t, orig, m.Env) + }) + + t.Run("failure path leaves Env untouched", func(t *testing.T) { + t.Parallel() + m := MCPConfig{Env: map[string]string{"BROKEN": "$(false)"}} + orig := maps.Clone(m.Env) + + _, err := m.ResolvedEnv(realShellResolver(nil)) + require.Error(t, err) + require.Equal(t, orig, m.Env, "map must be preserved on error") + }) +} + +// TestResolvedEnv_RealShell_Idempotent pins the pure-function contract: +// two calls on the same config return deeply-equal slices. +func TestResolvedEnv_RealShell_Idempotent(t *testing.T) { + t.Parallel() + + m := MCPConfig{ + Env: map[string]string{ + "A": "$(echo one)", + "B": "$(echo two)", + "C": "literal", + }, + } + r := realShellResolver(nil) + + first, err := m.ResolvedEnv(r) + require.NoError(t, err) + second, err := m.ResolvedEnv(r) + require.NoError(t, err) + require.Equal(t, first, second) +} + +// TestResolvedEnv_RealShell_Deterministic guards against Go's +// randomized map iteration leaking into the returned slice order. +func TestResolvedEnv_RealShell_Deterministic(t *testing.T) { + t.Parallel() + + m := MCPConfig{Env: map[string]string{ + "Z": "z", + "A": "a", + "M": "m", + }} + + got, err := m.ResolvedEnv(realShellResolver(nil)) + require.NoError(t, err) + require.True(t, slices.IsSorted(got), "env slice must be sorted; got %v", got) +} + +// TestResolvedEnv_RealShell_UnsetExpandsEmpty pins the lenient +// default: an unset bare $VAR expands to the empty string, matching +// bash. The silent-empty-credential class of bug is prevented by the +// pure-function error-returning contract of ResolvedEnv, so we don't +// rely on nounset to catch typo'd variable names. Users who want +// strict behaviour for a required credential opt in per-reference +// with ${VAR:?msg}; see TestResolvedEnv_RealShell_ColonQuestionIsStrict. +func TestResolvedEnv_RealShell_UnsetExpandsEmpty(t *testing.T) { + t.Parallel() + + m := MCPConfig{Env: map[string]string{ + // Intentional typo: user meant $FORGEJO_TOKEN. Under the + // lenient default this expands to "" rather than erroring, + // matching bash's behaviour on bare $VAR. + "FORGEJO_ACCESS_TOKEN": "$FORGJO_TOKEN", + }} + got, err := m.ResolvedEnv(realShellResolver(nil)) + require.NoError(t, err, "unset var must expand to empty, not error") + require.Equal(t, []string{"FORGEJO_ACCESS_TOKEN="}, got) +} + +// TestResolvedEnv_RealShell_ColonQuestionIsStrict pins the opt-in +// strictness contract: ${VAR:?msg} must hard-error when VAR is unset, +// regardless of the global NoUnset toggle. This is the recommended +// mechanism for required credentials under the lenient default, so a +// future refactor that accidentally swallows ${VAR:?...} errors would +// silently ship empty tokens to MCP servers again. +func TestResolvedEnv_RealShell_ColonQuestionIsStrict(t *testing.T) { + t.Parallel() + + m := MCPConfig{Env: map[string]string{ + "FORGEJO_ACCESS_TOKEN": "${FORGJO_TOKEN:?set FORGJO_TOKEN}", + }} + got, err := m.ResolvedEnv(realShellResolver(nil)) + require.Error(t, err, "${VAR:?msg} must error when VAR is unset") + require.Nil(t, got) + // The resolver wraps with the env key and the user-written + // template; the inner shell error carries the :? message so + // users learn which credential is missing and why. + msg := err.Error() + require.Contains(t, msg, "FORGEJO_ACCESS_TOKEN") + require.Contains(t, msg, "${FORGJO_TOKEN:?set FORGJO_TOKEN}") + require.Contains(t, msg, "set FORGJO_TOKEN") +} + +// TestResolvedEnv_RealShell_FailureDetail pins that a failing inner +// command surfaces enough detail (exit code + stderr on POSIX, the +// underlying OS error on Windows where coreutils runs in-process) to +// diagnose without forcing the user to re-run the command by hand. +// Also verifies the template is included so they know which Env +// entry blew up. +func TestResolvedEnv_RealShell_FailureDetail(t *testing.T) { + t.Parallel() + + // Forward slashes so the path survives shell-string injection on + // Windows; see TestResolvedEnv_RealShell_Success for the same note. + missing := filepath.ToSlash(filepath.Join(t.TempDir(), "definitely-not-here")) + m := MCPConfig{Env: map[string]string{ + "FORGEJO_ACCESS_TOKEN": fmt.Sprintf("$(cat %s)", missing), + }} + + _, err := m.ResolvedEnv(realShellResolver(nil)) + require.Error(t, err) + msg := err.Error() + require.Contains(t, msg, "FORGEJO_ACCESS_TOKEN", "must identify the failing env var") + require.Contains(t, msg, missing, "must include the template so users see what failed") + + // Inner diagnostic detail must survive. POSIX surfaces "exit + // status N" + stderr; Windows' in-process coreutils surfaces the + // Go OS error instead. Accept either shape so the test is + // portable without weakening the intent. + lower := strings.ToLower(msg) + hasDetail := strings.Contains(lower, "exit status") || + strings.Contains(lower, "no such file") || + strings.Contains(lower, "cannot find") + require.True(t, hasDetail, "must surface inner error detail: %s", msg) +} + +// TestResolvedHeaders_RealShell_FailurePreservesOriginal pins two +// invariants simultaneously: on failure the returned map is nil (not +// a partially-populated map) and the receiver's Headers map is +// unchanged. A test that only asserted on the returned value could +// hide an in-place mutation regression. +func TestResolvedHeaders_RealShell_FailurePreservesOriginal(t *testing.T) { + t.Parallel() + + m := MCPConfig{Headers: map[string]string{ + "Authorization": "Bearer $(false)", + "X-Static": "kept", + }} + orig := maps.Clone(m.Headers) + + got, err := m.ResolvedHeaders(realShellResolver(nil)) + require.Error(t, err) + require.Nil(t, got, "headers map must be nil on failure") + require.Contains(t, err.Error(), "header Authorization") + require.Equal(t, orig, m.Headers, "receiver Headers must be preserved") +} + +// TestResolvedHeaders_RealShell_DropEmpty verifies that a header +// whose value resolves to the empty string is omitted from the +// returned map. Covers the three ways a value can legitimately land +// on empty — unset bare $VAR under lenient nounset, a literal "", +// and a non-failing command whose stdout is empty — and also pins +// that a failing $(cmd) still errors rather than silently dropping. +func TestResolvedHeaders_RealShell_DropEmpty(t *testing.T) { + t.Parallel() + + t.Run("unset $VAR is absent", func(t *testing.T) { + t.Parallel() + m := MCPConfig{Headers: map[string]string{ + "X-Missing": "$MCP_HEADER_NEVER_SET", + "X-Kept": "present", + }} + got, err := m.ResolvedHeaders(realShellResolver(nil)) + require.NoError(t, err) + _, present := got["X-Missing"] + require.False(t, present, "unset bare $VAR → empty → header dropped") + require.Equal(t, "present", got["X-Kept"]) + }) + + t.Run("literal empty string is absent", func(t *testing.T) { + t.Parallel() + m := MCPConfig{Headers: map[string]string{ + "X-Custom": "", + "X-Kept": "present", + }} + got, err := m.ResolvedHeaders(realShellResolver(nil)) + require.NoError(t, err) + _, present := got["X-Custom"] + require.False(t, present, "literal empty-string header must be dropped") + require.Equal(t, "present", got["X-Kept"]) + }) + + t.Run("$(echo) is absent", func(t *testing.T) { + t.Parallel() + m := MCPConfig{Headers: map[string]string{ + "X-Empty": "$(echo)", + "X-Kept": "present", + }} + got, err := m.ResolvedHeaders(realShellResolver(nil)) + require.NoError(t, err) + _, present := got["X-Empty"] + require.False(t, present, "$(echo) → empty → header dropped") + require.Equal(t, "present", got["X-Kept"]) + }) + + t.Run("$(false) errors and does not mutate", func(t *testing.T) { + t.Parallel() + m := MCPConfig{Headers: map[string]string{ + "X-Broken": "$(false)", + "X-Kept": "present", + }} + orig := maps.Clone(m.Headers) + + got, err := m.ResolvedHeaders(realShellResolver(nil)) + require.Error(t, err) + require.Empty(t, got, "map must be nil/empty on failure, not a partial") + require.Contains(t, err.Error(), "header X-Broken") + require.Equal(t, orig, m.Headers, "receiver Headers must be preserved") + }) +} + +// TestResolvedArgs_RealShell exercises both success and failure for +// m.Args symmetrically with Env. Args are ordered so error messages +// must identify a positional index, not a key. +func TestResolvedArgs_RealShell(t *testing.T) { + t.Parallel() + + t.Run("success expands each element", func(t *testing.T) { + t.Parallel() + m := MCPConfig{Args: []string{"--token", "$(echo shh)", "--host", "example.com"}} + got, err := m.ResolvedArgs(realShellResolver(nil)) + require.NoError(t, err) + require.Equal(t, []string{"--token", "shh", "--host", "example.com"}, got) + }) + + t.Run("failure identifies offending index", func(t *testing.T) { + t.Parallel() + m := MCPConfig{Args: []string{"--token", "$(false)"}} + orig := slices.Clone(m.Args) + + got, err := m.ResolvedArgs(realShellResolver(nil)) + require.Error(t, err) + require.Nil(t, got) + require.Contains(t, err.Error(), "arg 1") + require.Equal(t, orig, m.Args, "receiver Args must be preserved") + }) + + t.Run("nil args returns nil, no error", func(t *testing.T) { + t.Parallel() + m := MCPConfig{} + got, err := m.ResolvedArgs(realShellResolver(nil)) + require.NoError(t, err) + require.Nil(t, got) + }) +} + +// TestLSPConfig_ResolvedArgs_RealShell exercises both success and +// failure for l.Args symmetrically with MCP args. Args are ordered so +// error messages must identify a positional index, not a key. +func TestLSPConfig_ResolvedArgs_RealShell(t *testing.T) { + t.Parallel() + + t.Run("success expands $VAR in each element", func(t *testing.T) { + t.Parallel() + l := LSPConfig{Args: []string{"--root", "$HOME", "--flag", "literal"}} + r := realShellResolver(map[string]string{"HOME": "/home/tester"}) + got, err := l.ResolvedArgs(r) + require.NoError(t, err) + require.Equal(t, []string{"--root", "/home/tester", "--flag", "literal"}, got) + }) + + t.Run("failure identifies offending index", func(t *testing.T) { + t.Parallel() + l := LSPConfig{Args: []string{"--root", "$(false)"}} + orig := slices.Clone(l.Args) + + got, err := l.ResolvedArgs(realShellResolver(nil)) + require.Error(t, err) + require.Nil(t, got) + require.Contains(t, err.Error(), "arg 1") + require.Equal(t, orig, l.Args, "receiver Args must be preserved") + }) + + t.Run("nil args returns nil, no error", func(t *testing.T) { + t.Parallel() + l := LSPConfig{} + got, err := l.ResolvedArgs(realShellResolver(nil)) + require.NoError(t, err) + require.Nil(t, got) + }) +} + +// TestLSPConfig_ResolvedEnv_RealShell pins the LSP env contract: +// success expands $VAR, failure wraps with the key name, and the +// receiver map is never mutated. The shape is map[string]string +// (not the MCP []string form) because powernap.ClientConfig.Environment +// takes a map directly. +func TestLSPConfig_ResolvedEnv_RealShell(t *testing.T) { + t.Parallel() + + t.Run("success expands $VAR", func(t *testing.T) { + t.Parallel() + l := LSPConfig{Env: map[string]string{"GOPATH": "$HOME/go"}} + r := realShellResolver(map[string]string{"HOME": "/home/tester"}) + got, err := l.ResolvedEnv(r) + require.NoError(t, err) + require.Equal(t, map[string]string{"GOPATH": "/home/tester/go"}, got) + }) + + t.Run("failure identifies offending key", func(t *testing.T) { + t.Parallel() + l := LSPConfig{Env: map[string]string{ + "GOPATH": "$(false)", + "OTHER": "literal", + }} + orig := maps.Clone(l.Env) + + got, err := l.ResolvedEnv(realShellResolver(nil)) + require.Error(t, err) + require.Nil(t, got) + require.Contains(t, err.Error(), `env "GOPATH"`) + require.Equal(t, orig, l.Env, "receiver Env must be preserved") + }) + + t.Run("idempotent and non-mutating", func(t *testing.T) { + t.Parallel() + l := LSPConfig{Env: map[string]string{ + "A": "$(echo one)", + "B": "literal", + }} + orig := maps.Clone(l.Env) + r := realShellResolver(nil) + + first, err := l.ResolvedEnv(r) + require.NoError(t, err) + second, err := l.ResolvedEnv(r) + require.NoError(t, err) + require.Equal(t, first, second) + require.Equal(t, orig, l.Env, "receiver Env must be preserved") + }) +} + +// TestLSPConfig_IdentityResolver pins the client-mode contract: both +// ResolvedArgs and ResolvedEnv round-trip the template verbatim under +// IdentityResolver and never error on unset variables. Local +// expansion would double-expand when the server does its own — this +// has to stay a pure pass-through. +func TestLSPConfig_IdentityResolver(t *testing.T) { + t.Parallel() + + l := LSPConfig{ + Args: []string{"--root", "$LSP_ROOT", "$(vault read -f lsp)"}, + Env: map[string]string{ + "GOPATH": "$HOME/go", + "TOKEN": "$(cat /run/secrets/x)", + }, + } + r := IdentityResolver() + + args, err := l.ResolvedArgs(r) + require.NoError(t, err) + require.Equal(t, l.Args, args) + + envs, err := l.ResolvedEnv(r) + require.NoError(t, err) + require.Equal(t, l.Env, envs) +} + +// TestMCPConfig_IdentityResolver pins the client-mode contract: every +// Resolved* method round-trips the template verbatim and never errors +// on unset variables. Local expansion would double-expand when the +// server does its own — this has to stay a pure pass-through. +func TestMCPConfig_IdentityResolver(t *testing.T) { + t.Parallel() + + m := MCPConfig{ + Command: "$CMD", + Args: []string{"--token", "$MCP_MISSING_TOKEN", "$(vault read -f secret)"}, + Env: map[string]string{ + "TOKEN": "$(cat /run/secrets/x)", + "HOST": "$MCP_MISSING_HOST", + }, + Headers: map[string]string{ + "Authorization": "Bearer $(vault read -f token)", + }, + URL: "https://$MCP_HOST/$(vault read -f path)", + } + r := IdentityResolver() + + args, err := m.ResolvedArgs(r) + require.NoError(t, err) + require.Equal(t, m.Args, args) + + envs, err := m.ResolvedEnv(r) + require.NoError(t, err) + // Sorted "KEY=value". + require.Equal(t, []string{ + "HOST=$MCP_MISSING_HOST", + "TOKEN=$(cat /run/secrets/x)", + }, envs) + + headers, err := m.ResolvedHeaders(r) + require.NoError(t, err) + require.Equal(t, m.Headers, headers) + + u, err := m.ResolvedURL(r) + require.NoError(t, err) + require.Equal(t, m.URL, u) +} diff --git a/internal/config/resolve_test.go b/internal/config/resolve_test.go new file mode 100644 index 0000000000000000000000000000000000000000..3d364e523ae4738d23888941bc429d2c00241fa2 --- /dev/null +++ b/internal/config/resolve_test.go @@ -0,0 +1,229 @@ +package config + +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/charmbracelet/crush/internal/env" + "github.com/stretchr/testify/require" +) + +// fakeExpander returns a canned value/error for the last passed value and +// records the context, raw value, and env slice it was called with. It +// lets the config-layer tests assert on delegation behaviour without +// spinning up a real interpreter — real-shell coverage lives in +// internal/shell/expand_test.go and resolve_real_test.go. +type fakeExpander struct { + expand func(ctx context.Context, value string, env []string) (string, error) + lastValue string + lastEnv []string + calls int +} + +func (f *fakeExpander) Expand(ctx context.Context, value string, env []string) (string, error) { + f.calls++ + f.lastValue = value + f.lastEnv = env + if f.expand == nil { + return value, nil + } + return f.expand(ctx, value, env) +} + +func TestShellVariableResolver_DelegatesToExpander(t *testing.T) { + t.Parallel() + + fe := &fakeExpander{ + expand: func(_ context.Context, value string, _ []string) (string, error) { + if value == "hello $FOO" { + return "hello bar", nil + } + return value, nil + }, + } + + e := env.NewFromMap(map[string]string{"FOO": "bar"}) + r := NewShellVariableResolver(e, WithExpander(fe.Expand)) + + got, err := r.ResolveValue("hello $FOO") + require.NoError(t, err) + require.Equal(t, "hello bar", got) + require.Equal(t, 1, fe.calls) + require.Equal(t, "hello $FOO", fe.lastValue) + require.Contains(t, fe.lastEnv, "FOO=bar") +} + +func TestShellVariableResolver_LoneDollarIsError(t *testing.T) { + t.Parallel() + + // Lone "$" must short-circuit before reaching the expander: the + // underlying shell parser would accept it as a literal, but this + // resolver has historically rejected it and callers depend on + // that early-fail behaviour. + fe := &fakeExpander{} + r := NewShellVariableResolver(env.NewFromMap(nil), WithExpander(fe.Expand)) + + _, err := r.ResolveValue("$") + require.Error(t, err) + require.Equal(t, 0, fe.calls, "expander must not be called for lone $") +} + +func TestShellVariableResolver_PassesThroughLiterals(t *testing.T) { + t.Parallel() + + fe := &fakeExpander{ + expand: func(_ context.Context, value string, _ []string) (string, error) { + return value, nil + }, + } + r := NewShellVariableResolver(env.NewFromMap(nil), WithExpander(fe.Expand)) + + got, err := r.ResolveValue("plain-string") + require.NoError(t, err) + require.Equal(t, "plain-string", got) +} + +func TestShellVariableResolver_WrapsErrorsWithTemplate(t *testing.T) { + t.Parallel() + + inner := errors.New("cat: /run/secrets/x: permission denied") + fe := &fakeExpander{ + expand: func(_ context.Context, _ string, _ []string) (string, error) { + return "", inner + }, + } + r := NewShellVariableResolver(env.NewFromMap(nil), WithExpander(fe.Expand)) + + _, err := r.ResolveValue("$(cat /run/secrets/x)") + require.Error(t, err) + require.ErrorIs(t, err, inner) + require.Contains(t, err.Error(), "$(cat /run/secrets/x)") + require.Contains(t, err.Error(), "permission denied") +} + +func TestSanitizeResolveError(t *testing.T) { + t.Parallel() + + t.Run("nil passes through", func(t *testing.T) { + t.Parallel() + require.NoError(t, sanitizeResolveError("anything", nil)) + }) + + t.Run("includes template and wraps inner", func(t *testing.T) { + t.Parallel() + inner := errors.New("cat: /run/secrets/x: permission denied") + got := sanitizeResolveError("$(cat /run/secrets/x)", inner) + require.Error(t, got) + require.ErrorIs(t, got, inner) + require.Contains(t, got.Error(), "$(cat /run/secrets/x)") + require.Contains(t, got.Error(), "permission denied") + }) + + t.Run("unwrap preserves original for errors.Is", func(t *testing.T) { + t.Parallel() + inner := errors.New("sentinel") + got := sanitizeResolveError("$FOO", inner) + require.ErrorIs(t, got, inner) + }) + + t.Run("truncates over-budget inner message", func(t *testing.T) { + t.Parallel() + // Inner message holds far more than the budget. After + // sanitization the rendered inner portion must not exceed + // maxResolveErrBytes, and the characters beyond the budget + // (marked by a distinct tail sentinel) must be gone. + const tailSentinel = "TAIL_SENTINEL_BEYOND_BUDGET" + body := strings.Repeat("x", maxResolveErrBytes) + inner := errors.New(body + tailSentinel) + + got := sanitizeResolveError("$TEMPLATE", inner) + require.Error(t, got) + + prefix := `resolving "$TEMPLATE": ` + rendered := got.Error() + require.True( + t, + strings.HasPrefix(rendered, prefix), + "rendered error must start with template prefix", + ) + innerRendered := strings.TrimPrefix(rendered, prefix) + require.LessOrEqual( + t, + len(innerRendered), + maxResolveErrBytes, + "inner message must be bounded to maxResolveErrBytes", + ) + require.NotContains( + t, + rendered, + tailSentinel, + "content past the budget must not leak", + ) + }) + + t.Run("replaces non-printable bytes", func(t *testing.T) { + t.Parallel() + // NUL, BEL, ESC, DEL, and a UTF-8 high byte should all be + // scrubbed to '?'. Tab and newline are preserved because + // they show up legitimately in command stderr. + inner := errors.New("ok\x00bad\x07\x1b\x7f\xffend\ttab\nline") + got := sanitizeResolveError("$T", inner) + rendered := got.Error() + + require.NotContains(t, rendered, "\x00") + require.NotContains(t, rendered, "\x07") + require.NotContains(t, rendered, "\x1b") + require.NotContains(t, rendered, "\x7f") + require.NotContains(t, rendered, "\xff") + require.Contains(t, rendered, "ok?bad????end\ttab\nline") + }) + + t.Run("scrubbing does not depend on shell.ExpandValue upstream", func(t *testing.T) { + t.Parallel() + // A custom Expander can inject arbitrary error text. The + // config-layer helper is the single chokepoint; it must + // bound + scrub regardless of the error source. + nasty := strings.Repeat("A", maxResolveErrBytes+64) + "\x00BEYOND" + fe := &fakeExpander{ + expand: func(_ context.Context, _ string, _ []string) (string, error) { + return "", errors.New(nasty) + }, + } + r := NewShellVariableResolver(env.NewFromMap(nil), WithExpander(fe.Expand)) + + _, err := r.ResolveValue("$T") + require.Error(t, err) + require.NotContains(t, err.Error(), "BEYOND", "over-budget tail must not leak") + require.NotContains(t, err.Error(), "\x00", "non-printables must be scrubbed") + }) +} + +func TestScrubErrorMessage(t *testing.T) { + t.Parallel() + + t.Run("bounds output to maxResolveErrBytes", func(t *testing.T) { + t.Parallel() + got := scrubErrorMessage(strings.Repeat("a", maxResolveErrBytes*3)) + require.Len(t, got, maxResolveErrBytes) + }) + + t.Run("preserves printable ASCII tab and newline", func(t *testing.T) { + t.Parallel() + require.Equal(t, "a\tb\nc d!", scrubErrorMessage("a\tb\nc d!")) + }) + + t.Run("replaces control and non-ASCII bytes", func(t *testing.T) { + t.Parallel() + require.Equal(t, "a?b??c", scrubErrorMessage("a\x01b\x1b\xe2c")) + }) +} + +func TestNewShellVariableResolver(t *testing.T) { + testEnv := env.NewFromMap(map[string]string{"TEST": "value"}) + resolver := NewShellVariableResolver(testEnv) + + require.NotNil(t, resolver) + require.Implements(t, (*VariableResolver)(nil), resolver) +} diff --git a/internal/config/scope.go b/internal/config/scope.go new file mode 100644 index 0000000000000000000000000000000000000000..89f7cb181f2750995450a3384132252ccf9603cf --- /dev/null +++ b/internal/config/scope.go @@ -0,0 +1,29 @@ +package config + +import "fmt" + +// Scope determines which config file is targeted for read/write operations. +type Scope int + +const ( + // ScopeGlobal targets the global data config (~/.local/share/crush/crush.json). + ScopeGlobal Scope = iota + // ScopeWorkspace targets the workspace config (.crush/crush.json). + ScopeWorkspace +) + +// String returns a human-readable label for the scope. +func (s Scope) String() string { + switch s { + case ScopeGlobal: + return "global" + case ScopeWorkspace: + return "workspace" + default: + return fmt.Sprintf("Scope(%d)", int(s)) + } +} + +// ErrNoWorkspaceConfig is returned when a workspace-scoped write is +// attempted on a ConfigStore that has no workspace config path. +var ErrNoWorkspaceConfig = fmt.Errorf("no workspace config path configured") diff --git a/internal/config/scopeb_race_test.go b/internal/config/scopeb_race_test.go new file mode 100644 index 0000000000000000000000000000000000000000..36ad3466ca9923c25d6a2fcb075f640c41589535 --- /dev/null +++ b/internal/config/scopeb_race_test.go @@ -0,0 +1,74 @@ +package config + +import ( + "os" + "path/filepath" + "sync" + "testing" +) + +// TestScopeB_InPlaceMutationRace probes whether the store's in-place field +// mutators (e.g. SetCompactMode) race with concurrent Config() readers that +// walk the same field. This is the "Scope B" race left open after the +// pointer-swap fix. Run with -race. +// TestScopeB_InPlaceMutationRace verifies that the store's typed field +// mutators (e.g. SetCompactMode) no longer race concurrent Config() +// readers walking the same field. Copy-on-write publishing means the +// mutator swaps in a fresh Config rather than writing through the live +// pointer, so a reader always sees an immutable snapshot. Run with -race. +func TestScopeB_InPlaceMutationRace(t *testing.T) { + dir := t.TempDir() + configPath := filepath.Join(dir, "crush.json") + + t.Setenv("CRUSH_GLOBAL_CONFIG", dir) + t.Setenv("CRUSH_GLOBAL_DATA", dir) + resetProviderState() + t.Cleanup(resetProviderState) + + cfg := `{ + "models": {"large": {"provider": "openai", "model": "gpt-4"}}, + "providers": {"openai": {"api_key": "k", "models": [{"id": "gpt-4", "name": "GPT-4"}]}}, + "options": {"tui": {"compact_mode": false}} + }` + if err := os.WriteFile(configPath, []byte(cfg), 0o600); err != nil { + t.Fatal(err) + } + + store, err := Load(dir, dir, false) + if err != nil { + t.Fatal(err) + } + store.globalDataPath = configPath + + var wg sync.WaitGroup + stop := make(chan struct{}) + + // Reader: walks Options.TUI.CompactMode, the field SetCompactMode writes. + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + c := store.Config() + if c != nil && c.Options != nil { + _ = c.Options.TUI.CompactMode + } + } + } + }() + + // Writer: flips compact mode, mutating Options in place. + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 50; i++ { + _ = store.SetCompactMode(ScopeGlobal, i%2 == 0) + } + close(stop) + }() + + wg.Wait() +} diff --git a/internal/config/shellconfig_hook_test.go b/internal/config/shellconfig_hook_test.go new file mode 100644 index 0000000000000000000000000000000000000000..013407109ff703e6e2fcbcb856daafe11b702087 --- /dev/null +++ b/internal/config/shellconfig_hook_test.go @@ -0,0 +1,37 @@ +package config_test + +import ( + "testing" + + "github.com/charmbracelet/crush/internal/hooks" + "github.com/stretchr/testify/require" +) + +func TestShellConfigHookAdd(t *testing.T) { + store := loadCrushSh(t, `hook add PreToolUse --matcher "^bash$" --command "echo hi" --name greet --timeout 10`) + + hs := store.Config().Hooks[hooks.EventPreToolUse] + require.Len(t, hs, 1) + require.Equal(t, "greet", hs[0].Name) + require.Equal(t, "^bash$", hs[0].Matcher) + require.Equal(t, "echo hi", hs[0].Command) + require.Equal(t, 10, hs[0].Timeout) +} + +func TestShellConfigHookRemoveByName(t *testing.T) { + store := loadCrushSh(t, `hook add PreToolUse --command a --name keep +hook add PreToolUse --command b --name drop +hook remove PreToolUse --name drop`) + + hs := store.Config().Hooks[hooks.EventPreToolUse] + require.Len(t, hs, 1) + require.Equal(t, "keep", hs[0].Name) +} + +func TestShellConfigHookClearEvent(t *testing.T) { + store := loadCrushSh(t, `hook add PreToolUse --command a --name a +hook add PreToolUse --command b --name b +hook remove PreToolUse`) + + require.Empty(t, store.Config().Hooks[hooks.EventPreToolUse]) +} diff --git a/internal/config/shellconfig_lsp_test.go b/internal/config/shellconfig_lsp_test.go new file mode 100644 index 0000000000000000000000000000000000000000..7c5f87b9ed9e85fea9d7240b9c293a8397dd31e5 --- /dev/null +++ b/internal/config/shellconfig_lsp_test.go @@ -0,0 +1,28 @@ +package config_test + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestShellConfigLSPAdd(t *testing.T) { + store := loadCrushSh(t, `lsp add gopls --command gopls --filetypes go --filetypes mod --root-markers go.mod --timeout 60`) + + l, ok := store.Config().LSP["gopls"] + require.True(t, ok, "gopls LSP should be configured") + require.Equal(t, "gopls", l.Command) + require.Subset(t, l.FileTypes, []string{"go", "mod"}) + require.Contains(t, l.RootMarkers, "go.mod") + require.Equal(t, 60, l.Timeout) +} + +func TestShellConfigLSPRemove(t *testing.T) { + store := loadCrushSh(t, `lsp add keepls --command keep-server +lsp add dropls --command drop-server +lsp remove dropls`) + + lsps := store.Config().LSP + require.Contains(t, lsps, "keepls") + require.NotContains(t, lsps, "dropls") +} diff --git a/internal/config/shellconfig_option_test.go b/internal/config/shellconfig_option_test.go new file mode 100644 index 0000000000000000000000000000000000000000..0bb44da45114e47bfe73b524fec2037ad7d8b118 --- /dev/null +++ b/internal/config/shellconfig_option_test.go @@ -0,0 +1,134 @@ +package config_test + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestShellConfigOptionBooleans(t *testing.T) { + store := loadCrushSh(t, `option debug true +option progress false +option auto-lsp false`) + + opts := store.Config().Options + require.True(t, opts.Debug, "debug should be on") + require.NotNil(t, opts.Progress) + require.False(t, *opts.Progress, "progress should be off") + require.NotNil(t, opts.AutoLSP) + require.False(t, *opts.AutoLSP, "auto-lsp should be off") +} + +// Config phrases this field negatively (disable_metrics) but the command +// exposes it positively. "metrics false" must land as disable_metrics = true. +func TestShellConfigOptionPositiveMetricsFalse(t *testing.T) { + store := loadCrushSh(t, `option metrics false`) + require.True(t, store.Config().Options.DisableMetrics, "metrics off => disable_metrics true") +} + +// The bare positive form defaults to true, which inverts to disable = false. +func TestShellConfigOptionPositiveMetricsBare(t *testing.T) { + store := loadCrushSh(t, `option metrics`) + require.False(t, store.Config().Options.DisableMetrics, "metrics on => disable_metrics false") +} + +func TestShellConfigOptionUI(t *testing.T) { + store := loadCrushSh(t, `option ui compact true +option ui diff split +option ui transparent false +option ui mouse false +option ui scrollbar always +option ui completions-max-depth 4 +option ui completions-max-items 200`) + + ui := store.Config().Options.TUI + require.NotNil(t, ui) + require.True(t, ui.CompactMode) + require.Equal(t, "split", ui.DiffMode) + require.NotNil(t, ui.Transparent) + require.False(t, *ui.Transparent) + require.NotNil(t, ui.Mouse) + require.False(t, *ui.Mouse) + require.Equal(t, "always", ui.Scrollbar) + require.NotNil(t, ui.Completions.MaxDepth) + require.Equal(t, 4, *ui.Completions.MaxDepth) + require.NotNil(t, ui.Completions.MaxItems) + require.Equal(t, 200, *ui.Completions.MaxItems) +} + +func TestShellConfigOptionUIRejectsInvalidValue(t *testing.T) { + _, err := loadCrushShErr(t, `option ui diff side-by-side`) + require.Error(t, err) + require.Contains(t, err.Error(), "expects unified or split") +} + +func TestShellConfigOptionAttribution(t *testing.T) { + store := loadCrushSh(t, `option attribution-trailer-style none +option attribution-generated-with false`) + + attribution := store.Config().Options.Attribution + require.NotNil(t, attribution) + require.Equal(t, "none", string(attribution.TrailerStyle)) + require.False(t, attribution.GeneratedWith) +} + +func TestShellConfigOptionAttributionTrailerStylePreservesGeneratedWithDefault(t *testing.T) { + store := loadCrushSh(t, `option attribution-trailer-style co-authored-by`) + + attribution := store.Config().Options.Attribution + require.NotNil(t, attribution) + require.Equal(t, "co-authored-by", string(attribution.TrailerStyle)) + require.True(t, attribution.GeneratedWith) +} + +func TestShellConfigOptionAttributionGeneratedWithCaseInsensitive(t *testing.T) { + store := loadCrushSh(t, `option attribution-generated-with YES`) + + require.NotNil(t, store.Config().Options.Attribution) + require.True(t, store.Config().Options.Attribution.GeneratedWith) +} + +func TestShellConfigOptionAttributionRejectsInvalidStyle(t *testing.T) { + _, err := loadCrushShErr(t, `option attribution-trailer-style bogus`) + require.Error(t, err) + require.Contains(t, err.Error(), "expects none, co-authored-by, or assisted-by") +} + +func TestShellConfigOptionListAppends(t *testing.T) { + store := loadCrushSh(t, `option disable-skill crush-config +option disable-skill jq`) + + require.Subset(t, store.Config().Options.DisabledSkills, []string{"crush-config", "jq"}) +} + +// reset wipes values added earlier (or via source) while keeping anything +// added after it — observable in the effective config. +func TestShellConfigOptionReset(t *testing.T) { + store := loadCrushSh(t, `option skill-path ./inherited-a +option skill-path ./inherited-b +option reset skill-path +option skill-path ./mine`) + + paths := store.Config().Options.SkillsPaths + require.Contains(t, paths, "./mine") + require.NotContains(t, paths, "./inherited-a") + require.NotContains(t, paths, "./inherited-b") +} + +func TestShellConfigOptionUnknownKeyFails(t *testing.T) { + _, err := loadCrushShErr(t, `option bogus-key value`) + require.Error(t, err) + require.Contains(t, err.Error(), "unknown key") +} + +func TestShellConfigOptionDisableToolRemoved(t *testing.T) { + _, err := loadCrushShErr(t, `option disable-tool bash`) + require.Error(t, err) + require.Contains(t, err.Error(), "unknown key") +} + +func TestShellConfigOptionResetRejectsNonList(t *testing.T) { + _, err := loadCrushShErr(t, `option reset debug`) + require.Error(t, err) + require.Contains(t, err.Error(), "not one") +} diff --git a/internal/config/shellconfig_precedence_test.go b/internal/config/shellconfig_precedence_test.go new file mode 100644 index 0000000000000000000000000000000000000000..6c7adf50b3eb32afc420b0a325263f122c30b281 --- /dev/null +++ b/internal/config/shellconfig_precedence_test.go @@ -0,0 +1,35 @@ +package config_test + +import ( + "os" + "path/filepath" + "testing" + + "github.com/charmbracelet/crush/internal/config" + "github.com/stretchr/testify/require" +) + +// TestShellConfigDotCrushrcTakesPrecedence verifies that a project-local +// .crushrc overrides crushrc in the same directory on conflicting settings. +func TestShellConfigDotCrushrcTakesPrecedence(t *testing.T) { + isolated := t.TempDir() + t.Setenv("HOME", isolated) + t.Setenv("XDG_CONFIG_HOME", filepath.Join(isolated, ".config")) + t.Setenv("XDG_DATA_HOME", filepath.Join(isolated, ".local", "share")) + + workDir := t.TempDir() + dataDir := t.TempDir() + require.NoError(t, os.WriteFile( + filepath.Join(workDir, "crushrc"), + []byte("option notifications bell\n"), 0o644, + )) + require.NoError(t, os.WriteFile( + filepath.Join(workDir, ".crushrc"), + []byte("option notifications osc\n"), 0o644, + )) + + store, err := config.Load(workDir, dataDir, false) + require.NoError(t, err) + require.Equal(t, "osc", store.Config().Options.Notifications, + ".crushrc should win over crushrc") +} diff --git a/internal/config/shellconfig_provider_test.go b/internal/config/shellconfig_provider_test.go new file mode 100644 index 0000000000000000000000000000000000000000..5d1ed8c2732dba3929a2fdd2ae3d19a6c8e8b051 --- /dev/null +++ b/internal/config/shellconfig_provider_test.go @@ -0,0 +1,85 @@ +package config_test + +import ( + "slices" + "testing" + + "charm.land/catwalk/pkg/catwalk" + "github.com/charmbracelet/crush/internal/config" + "github.com/stretchr/testify/require" +) + +func TestShellConfigProviderAddAndModel(t *testing.T) { + store := loadCrushSh(t, `provider add myllm \ + --type openai-compat \ + --base-url "http://localhost:1234/v1" \ + --api-key "sk-test" \ + --discover-models false \ + --extra-body '{"service_tier":"flex"}' \ + --provider-options '{"region":"local"}' +provider add myllm \ + --extra-body '{"stream":true}' \ + --provider-options '{"mode":"test"}' +model add myllm/foo-1 --name "Foo 1" --context-window 8000 \ + --price-input 1.25 --price-output 5 \ + --price-cache-create 2 --price-cache-hit 0.25 +model large myllm/foo-1 \ + --top-p 0.9 --top-k 40 \ + --frequency-penalty 0.2 --presence-penalty 0.1 \ + --provider-options '{"routing":{"tier":"fast"}}' +model large myllm/foo-1 --provider-options '{"timeout":30}'`) + + cfg := store.Config() + + p, ok := cfg.Providers.Get("myllm") + require.True(t, ok, "myllm provider should be configured") + require.Equal(t, "sk-test", p.APIKey) + require.Equal(t, "http://localhost:1234/v1", p.BaseURL) + require.NotNil(t, p.AutoDiscoverModels) + require.False(t, *p.AutoDiscoverModels) + require.Equal(t, "flex", p.ExtraBody["service_tier"]) + require.Equal(t, true, p.ExtraBody["stream"]) + require.Equal(t, "local", p.ProviderOptions["region"]) + require.Equal(t, "test", p.ProviderOptions["mode"]) + require.True( + t, + slices.ContainsFunc(p.Models, func(m catwalk.Model) bool { return m.ID == "foo-1" }), + "custom model foo-1 should be in the provider catalog", + ) + model := p.Models[0] + require.Equal(t, 1.25, model.CostPer1MIn) + require.Equal(t, 5.0, model.CostPer1MOut) + require.Equal(t, 2.0, model.CostPer1MInCached) + require.Equal(t, 0.25, model.CostPer1MOutCached) + + large := cfg.Models[config.SelectedModelTypeLarge] + require.Equal(t, "myllm", large.Provider) + require.Equal(t, "foo-1", large.Model) + require.NotNil(t, large.TopP) + require.Equal(t, 0.9, *large.TopP) + require.NotNil(t, large.TopK) + require.Equal(t, int64(40), *large.TopK) + require.NotNil(t, large.FrequencyPenalty) + require.Equal(t, 0.2, *large.FrequencyPenalty) + require.NotNil(t, large.PresencePenalty) + require.Equal(t, 0.1, *large.PresencePenalty) + require.Equal(t, map[string]any{ + "routing": map[string]any{"tier": "fast"}, + "timeout": float64(30), + }, large.ProviderOptions) +} + +func TestShellConfigProviderRemove(t *testing.T) { + // Both providers get a model so they survive provider configuration + // (model-less providers are dropped); the only difference is the remove. + store := loadCrushSh(t, `provider add keepme --type openai-compat --base-url "http://localhost:1/v1" --api-key k +model add keepme/m1 --name M1 +provider add dropme --type openai-compat --base-url "http://localhost:2/v1" --api-key k +model add dropme/m2 --name M2 +provider remove dropme`) + + _, keep := store.Config().Providers.Get("keepme") + _, drop := store.Config().Providers.Get("dropme") + require.True(t, keep, "keepme should remain") + require.False(t, drop, "dropme should be gone after remove") +} diff --git a/internal/config/store.go b/internal/config/store.go new file mode 100644 index 0000000000000000000000000000000000000000..8a47e97f7e824d3945368e500548cc9280b67c9f --- /dev/null +++ b/internal/config/store.go @@ -0,0 +1,1430 @@ +package config + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "maps" + "os" + "path/filepath" + "slices" + "sync" + "time" + + "charm.land/catwalk/pkg/catwalk" + hyperp "github.com/charmbracelet/crush/internal/agent/hyper" + "github.com/charmbracelet/crush/internal/env" + "github.com/charmbracelet/crush/internal/lock" + "github.com/charmbracelet/crush/internal/oauth" + "github.com/charmbracelet/crush/internal/oauth/copilot" + "github.com/charmbracelet/crush/internal/oauth/hyper" + "github.com/charmbracelet/crush/internal/oauth/openai" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" + "golang.org/x/sync/singleflight" +) + +// configLockDeadline bounds how long lockConfig waits for the +// cross-process flock before giving up. A few seconds is plenty for +// honest contention; longer suggests something is wedged. +const configLockDeadline = 5 * time.Second + +// refreshLockDeadline bounds how long RefreshOAuthToken waits for the +// per-provider cross-process refresh lock. It must exceed the token +// exchange HTTP timeout (30s) so that a peer mid-exchange is given time +// to finish and publish its result, which we then adopt instead of +// running our own exchange. Running our own would reuse an +// already-rotated refresh token and trip the provider's reuse detection, +// revoking the whole token family. +const refreshLockDeadline = 45 * time.Second + +// credentialWriteLockDeadline bounds how long a credential write (e.g. +// storing the token from a fresh interactive login) waits for the +// per-provider refresh lock. It is deliberately shorter than +// refreshLockDeadline because a user is watching: if a peer is wedged we +// would rather write and risk a rare clobber than hang the UI. +const credentialWriteLockDeadline = 10 * time.Second + +// fileSnapshot captures metadata about a config file at a point in time. +type fileSnapshot struct { + Path string + Exists bool + Size int64 + ModTime int64 // UnixNano +} + +// RuntimeOverrides holds per-session settings that are never persisted to +// disk. They are applied on top of the loaded Config and survive only for +// the lifetime of the process (or workspace). +type RuntimeOverrides struct { + SkipPermissionRequests bool + // EnabledChannels lists the MCP servers opted in as channels for this + // session (via the --channels flag). A server present in MCP config only + // pushes channel events when it also appears here. Entries may be written + // as "server:" or as a bare "". + EnabledChannels []string + // Models records the model choices made in this instance, whether + // persisted or not. They are reapplied after a config reload so that a + // selection made here always outranks whatever the shared config file + // happens to hold — see pinPreferredModelLocked. + Models map[SelectedModelType]SelectedModel +} + +// ConfigStore is the single entry point for all config access. It owns the +// pure-data Config, runtime state (working directory, resolver, known +// providers), and persistence to both global and workspace config files. +// +// mu serialises all config file mutations (SetConfigFields, +// RemoveConfigField, RefreshOAuthToken) to prevent both in-process +// goroutine races and, together with the shared lock.File, cross-process +// races on the config file. +// +// writeMu serialises every operation that produces a new in-memory Config: +// the typed copy-on-write mutators (SetCompactMode, UpdatePreferredModel, +// ...) and ReloadFromDisk. Typed mutators take Lock; autoReload takes +// TryLock so a write triggered re-entrantly during a reload (e.g. +// configureProviders calling RemoveConfigField) skips the nested reload +// instead of deadlocking. This is what lets published Configs be treated +// as immutable: a mutator clones, mutates the clone, and swaps it in under +// writeMu rather than mutating the live Config in place. +type ConfigStore struct { + config *Config + workingDir string + resolver VariableResolver + globalDataPath string // ~/.local/share/crush/crush.json + workspacePath string // .crush/crush.json + loadedPaths []string // config files that were successfully loaded + knownProviders []catwalk.Provider + overrides RuntimeOverrides + trackedConfigPaths []string // unique, normalized config file paths + snapshots map[string]fileSnapshot // path -> snapshot at last capture + + // configMu guards the config pointer field against concurrent + // readers (Config) and the writeMu-serialised swap (setConfig). It + // protects the pointer word only; the pointed-to Config is treated + // as immutable once published, since both reloads and typed mutators + // build a fresh Config rather than mutating the live one. + configMu sync.RWMutex + + mu sync.Mutex // serialises config file writes + writeMu sync.RWMutex // serialises in-memory config production (mutators + reload); RLock for readers + + // refreshSF collapses concurrent in-process OAuth refreshes for the + // same provider into a single attempt. Combined with the per-provider + // cross-process refresh lock, it ensures only one token exchange runs + // at a time. See RefreshOAuthToken. + refreshSF singleflight.Group + + // exchangeToken performs the provider-specific OAuth token exchange. + // It is a field so tests can substitute a fake exchange without making + // real network calls. Production code leaves it nil, and exchange falls + // back to the real provider clients. + exchangeToken func(ctx context.Context, providerID, refreshToken string) (*oauth.Token, error) + + // fetchOpenAIModels fetches the ChatGPT model catalog. A field for the + // same reason as exchangeToken: tests stub it to avoid network calls, + // production leaves it nil and refetchOpenAIModels calls the real one. + fetchOpenAIModels func(ctx context.Context, token *oauth.Token) ([]catwalk.Model, error) + + // authSignalMu guards authSignals, which maps provider IDs to + // channels that WaitForTokenChange blocks on. SignalAuthComplete + // closes the channel to unblock waiters; a new channel is created + // on the next wait. + authSignalMu sync.Mutex + authSignals map[string]chan struct{} +} + +// Config returns the pure-data config struct (read-only after load). +// +// The pointer read is guarded by configMu so it can never tear against +// the reload swap in reloadFromDiskLocked. Reloads build a brand-new +// Config and swap it in rather than mutating the live one, so holding the +// returned pointer stays safe even across a concurrent reload — the reader +// keeps reading its (now immutable) snapshot. +func (s *ConfigStore) Config() *Config { + s.configMu.RLock() + defer s.configMu.RUnlock() + return s.config +} + +// setConfig atomically swaps the active config pointer under configMu. +// Used by the reload path; in-place field mutators leave the pointer +// untouched and run under mu instead. +func (s *ConfigStore) setConfig(cfg *Config) { + s.configMu.Lock() + defer s.configMu.Unlock() + s.config = cfg +} + +// WorkingDir returns the current working directory. +func (s *ConfigStore) WorkingDir() string { + return s.workingDir +} + +// Resolver returns the variable resolver. +func (s *ConfigStore) Resolver() VariableResolver { + s.writeMu.RLock() + defer s.writeMu.RUnlock() + return s.resolver +} + +// Resolve resolves a variable reference using the configured resolver. +func (s *ConfigStore) Resolve(key string) (string, error) { + s.writeMu.RLock() + r := s.resolver + s.writeMu.RUnlock() + if r == nil { + return "", fmt.Errorf("no variable resolver configured") + } + return r.ResolveValue(key) +} + +// KnownProviders returns the list of known providers. +func (s *ConfigStore) KnownProviders() []catwalk.Provider { + s.writeMu.RLock() + defer s.writeMu.RUnlock() + return s.knownProviders +} + +// RefetchHyperProvider re-fetches the Hyper provider catalog from the +// remote API and updates the in-memory known providers list and config. +// This is called after OAuth authentication completes so the latest +// models are available without restarting. +func (s *ConfigStore) RefetchHyperProvider(ctx context.Context) error { + // Build a fresh client that reads the API key from the live config, + // not the stale snapshot captured at startup. The syncer's original + // client closes over the startup config and would send an expired + // token after OAuth re-authentication. + freshClient := realHyperClient{ + baseURL: hyperp.BaseURL(), + resolveKey: func() string { return resolveHyperAPIKey(s.Config()) }, + } + hyperSyncer.SetClient(freshClient) + + hyperProvider, err := hyperSyncer.Refetch(ctx) + if err != nil { + return fmt.Errorf("failed to refetch Hyper provider: %w", err) + } + if hyperProvider.ID == "" { + return nil + } + + s.writeMu.Lock() + defer s.writeMu.Unlock() + + // Replace or insert the Hyper entry in knownProviders. + found := false + for i, p := range s.knownProviders { + if string(p.ID) == string(hyperProvider.ID) { + s.knownProviders[i] = hyperProvider + found = true + break + } + } + if !found { + s.knownProviders = append([]catwalk.Provider{hyperProvider}, s.knownProviders...) + } + + // Update the Hyper provider config with the refreshed model list + // and endpoint. Use cloneForWrite so readers always see a consistent + // snapshot (the store's contract forbids in-place config mutation). + nc := s.config.cloneForWrite() + if pc, ok := nc.Providers.Get(string(hyperProvider.ID)); ok { + pc.Models = hyperProvider.Models + if hyperProvider.APIEndpoint != "" { + pc.BaseURL = hyperProvider.APIEndpoint + } + nc.Providers.Set(string(hyperProvider.ID), pc) + } + s.setConfig(nc) + + // Also update the memoized provider list so callers of + // config.Providers() (e.g. the models dialog) see fresh data. + UpdateProviderInList(hyperProvider) + + s.SetupAgents() + return nil +} + +// SetupAgents configures the coder and task agents on the config. +func (s *ConfigStore) SetupAgents() { + s.Config().SetupAgents() +} + +// Overrides returns the runtime overrides for this store. +func (s *ConfigStore) Overrides() *RuntimeOverrides { + s.writeMu.RLock() + defer s.writeMu.RUnlock() + return &s.overrides +} + +// LoadedPaths returns the config file paths that were successfully loaded. +func (s *ConfigStore) LoadedPaths() []string { + s.writeMu.RLock() + defer s.writeMu.RUnlock() + return slices.Clone(s.loadedPaths) +} + +// lockConfig acquires both the in-process mutex and a cross-process flock +// on the config file for the given scope. Callers that need to do I/O +// between reading and writing (e.g. an HTTP token exchange) must use +// lockConfig explicitly rather than atomicWrite. +// +// The returned release function drops both locks. Callers must call it +// as soon as the file access is complete — no I/O should be performed +// while the lock is held. +func (s *ConfigStore) lockConfig(scope Scope) (func(), error) { + s.mu.Lock() + path, err := s.configPath(scope) + if err != nil { + s.mu.Unlock() + return nil, err + } + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + s.mu.Unlock() + return nil, fmt.Errorf("create config directory: %w", err) + } + ctx, cancel := context.WithTimeout(context.Background(), configLockDeadline) + defer cancel() + release, err := lock.File(ctx, path+".lock") + if err != nil { + s.mu.Unlock() + return nil, fmt.Errorf("acquire config lock: %w", err) + } + return func() { + release() + s.mu.Unlock() + }, nil +} + +// atomicWrite handles the lock-read-transform-write-unlock cycle for +// config file mutations. The fn callback receives the current file +// contents (raw bytes, or {} if the file is missing) and must return the +// new contents. fn must be pure — no I/O, no network calls. +func (s *ConfigStore) atomicWrite(scope Scope, fn func(current []byte) ([]byte, error)) error { + unlock, err := s.lockConfig(scope) + if err != nil { + return err + } + defer unlock() + + path, err := s.configPath(scope) + if err != nil { + return err + } + + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + data = []byte("{}") + } else { + return fmt.Errorf("read config file: %w", err) + } + } + + newData, err := fn(data) + if err != nil { + return err + } + + return atomicWriteFile(path, newData, 0o600) +} + +// configPath returns the file path for the given scope. +func (s *ConfigStore) configPath(scope Scope) (string, error) { + switch scope { + case ScopeWorkspace: + if s.workspacePath == "" { + return "", ErrNoWorkspaceConfig + } + return s.workspacePath, nil + default: + return s.globalDataPath, nil + } +} + +// HasConfigField checks whether a key exists in the config file for the given +// scope. +func (s *ConfigStore) HasConfigField(scope Scope, key string) bool { + path, err := s.configPath(scope) + if err != nil { + return false + } + data, err := os.ReadFile(path) + if err != nil { + return false + } + return gjson.Get(string(data), key).Exists() +} + +// SetConfigField sets a key/value pair in the config file for the given scope. +// After a successful write, it automatically reloads config to keep in-memory +// state fresh. +func (s *ConfigStore) SetConfigField(scope Scope, key string, value any) error { + return s.SetConfigFields(scope, map[string]any{key: value}) +} + +// SetConfigFields sets multiple key/value pairs in the config file for the +// given scope in a single write, then reloads in-memory state from disk. +// +// Use this for arbitrary external edits where the in-memory effect of the +// change is not known ahead of time. The typed mutators (which know exactly +// what changed) go through update instead and skip the reload. +// +// The write is protected by an in-process mutex and a cross-process flock +// to prevent races between concurrent writers in different processes. +func (s *ConfigStore) SetConfigFields(scope Scope, kv map[string]any) error { + if err := s.writeConfigFields(scope, kv); err != nil { + return err + } + // Auto-reload to keep in-memory state fresh after config edits. + // We use context.Background() since this is an internal operation that + // shouldn't be cancelled by user context. + if err := s.autoReload(context.Background()); err != nil { + // Log warning but don't fail the write - disk is already updated. + slog.Warn("Config file updated but failed to reload in-memory state", "error", err) + } + return nil +} + +// writeConfigFields persists key/value pairs to the config file. It does not +// touch in-memory config state or the staleness snapshot: callers either +// reload (SetConfigFields, whose reload recaptures the snapshot) or have +// already published an updated clone and capture the snapshot themselves +// (update). Both of those run under writeMu, which is what keeps the +// snapshot map free of concurrent writers. +func (s *ConfigStore) writeConfigFields(scope Scope, kv map[string]any) error { + // Sort keys for deterministic output regardless of map iteration + // order. This also ensures consistent results when callers pass + // overlapping JSONPath keys (e.g. "a" and "a.b"). + keys := make([]string, 0, len(kv)) + for k := range kv { + keys = append(keys, k) + } + slices.Sort(keys) + + return s.atomicWrite(scope, func(data []byte) ([]byte, error) { + v := string(data) + for _, key := range keys { + var sErr error + if v, sErr = sjson.Set(v, key, kv[key]); sErr != nil { + return nil, fmt.Errorf("failed to set config field %s: %w", key, sErr) + } + } + return []byte(v), nil + }) +} + +// mutateInMemory applies a copy-on-write change to the config without +// persisting. Under writeMu it clones the live config, lets mutate edit the +// clone, and publishes it. This is the single primitive every in-memory +// config change goes through, so a published Config is never mutated in +// place and readers always see a consistent snapshot. +func (s *ConfigStore) mutateInMemory(mutate func(*Config)) { + s.writeMu.Lock() + defer s.writeMu.Unlock() + + nc := s.Config().cloneForWrite() + mutate(nc) + s.setConfig(nc) +} + +// update applies a copy-on-write change and persists the reported fields. +// mutate edits the clone and returns the JSON-path fields to write to disk; +// because the clone already reflects the change, no reload is needed. +// Returning an empty map publishes the clone without a disk write. +func (s *ConfigStore) update(scope Scope, mutate func(*Config) map[string]any) error { + s.writeMu.Lock() + defer s.writeMu.Unlock() + return s.updateLocked(scope, mutate) +} + +// updateLocked is the lock-free core of update. Caller must hold writeMu. +func (s *ConfigStore) updateLocked(scope Scope, mutate func(*Config) map[string]any) error { + nc := s.Config().cloneForWrite() + fields := mutate(nc) + s.setConfig(nc) + if len(fields) == 0 { + return nil + } + if err := s.writeConfigFields(scope, fields); err != nil { + return err + } + // Refresh the staleness snapshot so the file watcher does not treat + // our own write as an external change. Safe to touch the snapshot map + // here because we hold writeMu. + if path, err := s.configPath(scope); err == nil { + s.captureStalenessSnapshot(append(slices.Clone(s.loadedPaths), path)) + } + return nil +} + +// OverridePreferredModel sets the preferred model for the given type in +// memory only, without persisting. It is for per-run overrides (such as the +// non-interactive --model flags) that must not be written to the user's +// config file. +func (s *ConfigStore) OverridePreferredModel(modelType SelectedModelType, model SelectedModel) { + s.mutateInMemory(func(c *Config) { + if c.Models == nil { + c.Models = make(map[SelectedModelType]SelectedModel) + } + c.Models[modelType] = model + s.pinPreferredModelLocked(modelType, model) + }) +} + +// pinPreferredModelLocked records a model choice made in this instance so +// that a later config reload cannot replace it with a choice made +// somewhere else. Several Crush instances share one global config file, so +// a reload triggered by an unrelated write (a token refresh, say) would +// otherwise import whichever model a sibling instance last selected and +// switch models out from under the user mid-session. +// +// Caller must hold writeMu. +func (s *ConfigStore) pinPreferredModelLocked(modelType SelectedModelType, model SelectedModel) { + if s.overrides.Models == nil { + s.overrides.Models = make(map[SelectedModelType]SelectedModel) + } + s.overrides.Models[modelType] = model +} + +// RemoveConfigField removes a key from the config file for the given scope. +// After a successful write, it automatically reloads config to keep in-memory +// state fresh. +// +// The write is protected by an in-process mutex and a cross-process flock. +func (s *ConfigStore) RemoveConfigField(scope Scope, key string) error { + err := s.atomicWrite(scope, func(data []byte) ([]byte, error) { + v, sErr := sjson.Delete(string(data), key) + if sErr != nil { + return nil, fmt.Errorf("failed to delete config field %s: %w", key, sErr) + } + return []byte(v), nil + }) + if err != nil { + return err + } + + if err := s.autoReload(context.Background()); err != nil { + slog.Warn("Config file updated but failed to reload in-memory state", "error", err) + } + + return nil +} + +// UpdatePreferredModel updates the preferred model for the given type and +// persists it to the config file at the given scope. The selected model and +// the recent-models list are written together in a single config write. +// +// The write skips the full disk reparse/reload (which would rebuild the +// provider catalog and agents on every model switch and dominate selection +// latency); agents are refreshed separately by the caller (see +// UpdateAgentModel). +func (s *ConfigStore) UpdatePreferredModel(scope Scope, modelType SelectedModelType, model SelectedModel) error { + return s.update(scope, func(c *Config) map[string]any { + return s.updatePreferredModelFields(c, modelType, model) + }) +} + +// updatePreferredModelFields builds the fields map for persisting a preferred +// model change. Shared between UpdatePreferredModel and direct updateLocked +// callers (e.g. Load). Caller must hold writeMu. +func (s *ConfigStore) updatePreferredModelFields(c *Config, modelType SelectedModelType, model SelectedModel) map[string]any { + if c.Models == nil { + c.Models = make(map[SelectedModelType]SelectedModel) + } + c.Models[modelType] = model + s.pinPreferredModelLocked(modelType, model) + + fields := map[string]any{ + fmt.Sprintf("models.%s", modelType): model, + } + if updated, changed := nextRecentModels(c, modelType, model); changed { + if c.RecentModels == nil { + c.RecentModels = make(map[SelectedModelType][]SelectedModel) + } + c.RecentModels[modelType] = updated + fields[fmt.Sprintf("recent_models.%s", modelType)] = updated + } + return fields +} + +// SetCompactMode sets the compact mode setting and persists it. +func (s *ConfigStore) SetCompactMode(scope Scope, enabled bool) error { + return s.update(scope, func(c *Config) map[string]any { + c.ensureTUI().CompactMode = enabled + return map[string]any{"options.tui.compact_mode": enabled} + }) +} + +// SetTransparentBackground sets the transparent background setting and persists it. +func (s *ConfigStore) SetTransparentBackground(scope Scope, enabled bool) error { + return s.update(scope, func(c *Config) map[string]any { + c.ensureTUI().Transparent = &enabled + return map[string]any{"options.tui.transparent": enabled} + }) +} + +// SetProviderAPIKey sets the API key for a provider and persists it. +// The OpenAI provider holds exactly one credential: storing a ChatGPT +// token removes a previously entered API key, and storing an API key +// removes a previous ChatGPT login. +func (s *ConfigStore) SetProviderAPIKey(scope Scope, providerID string, apiKey any) error { + var providerConfig ProviderConfig + var exists bool + var setKeyOrToken func() + isToken := false + + switch v := apiKey.(type) { + case string: + if err := s.SetConfigField(scope, fmt.Sprintf("providers.%s.api_key", providerID), v); err != nil { + return fmt.Errorf("failed to save api key to config file: %w", err) + } + setKeyOrToken = func() { + providerConfig.APIKey = v + if providerID == string(catwalk.InferenceProviderOpenAI) { + // Either OAuth or an API key, never both: the login + // leaves nothing usable on the API-key side behind. + providerConfig.OAuthToken = nil + providerConfig.ChatGPTModels = nil + } + } + if providerID == string(catwalk.InferenceProviderOpenAI) { + // Either OAuth or an API key, never both: the new key + // leaves nothing usable on the ChatGPT side behind. + if err := s.RemoveConfigField(scope, fmt.Sprintf("providers.%s.oauth", providerID)); err != nil { + return err + } + if err := s.RemoveConfigField(scope, fmt.Sprintf("providers.%s.chatgpt_models", providerID)); err != nil { + return err + } + } + case *oauth.Token: + // Hold the refresh lock across the write so a peer's in-flight + // token exchange cannot land on top of a credential the user just + // obtained interactively — which would silently invalidate the + // login they only just completed. + fields := map[string]any{ + fmt.Sprintf("providers.%s.oauth", providerID): v, + } + if providerID != string(catwalk.InferenceProviderOpenAI) { + fields[fmt.Sprintf("providers.%s.api_key", providerID)] = v.AccessToken + } + if err := s.withRefreshLock(providerID, func() error { + return s.SetConfigFields(scope, fields) + }); err != nil { + return err + } + if providerID == string(catwalk.InferenceProviderOpenAI) { + // Either OAuth or an API key, never both: the login retires + // any key that came before it. + if err := s.RemoveConfigField(scope, fmt.Sprintf("providers.%s.api_key", providerID)); err != nil { + return err + } + } + setKeyOrToken = func() { + providerConfig.OAuthToken = v + if providerID == string(catwalk.InferenceProviderOpenAI) { + isToken = true + providerConfig.APIKey = "" + return + } + providerConfig.APIKey = v.AccessToken + if providerID == string(catwalk.InferenceProviderCopilot) { + providerConfig.SetupGitHubCopilot() + } + } + } + + cfg := s.Config() + providerConfig, exists = cfg.Providers.Get(providerID) + if exists { + setKeyOrToken() + cfg.Providers.Set(providerID, providerConfig) + } else { + var foundProvider *catwalk.Provider + for _, p := range s.knownProviders { + if string(p.ID) == providerID { + foundProvider = &p + break + } + } + + if foundProvider != nil { + providerConfig = ProviderConfig{ + ID: providerID, + Name: foundProvider.Name, + BaseURL: foundProvider.APIEndpoint, + Type: foundProvider.Type, + Disable: false, + ExtraHeaders: make(map[string]string), + ExtraParams: make(map[string]string), + Models: foundProvider.Models, + } + setKeyOrToken() + } else { + return fmt.Errorf("provider with ID %s not found in known providers", providerID) + } + cfg.Providers.Set(providerID, providerConfig) + } + + // After authenticating with Hyper, re-fetch the provider catalog so + // the latest models are available without restarting. + if providerID == "hyper" { + if refetchErr := s.RefetchHyperProvider(context.Background()); refetchErr != nil { + slog.Warn("Failed to refetch Hyper provider after auth", "error", refetchErr) + } + } + // After authenticating with a ChatGPT account, fetch the Codex model + // catalog the subscription grants and persist it so the models + // dialog can offer it beside the API-key catalog. + if providerID == string(catwalk.InferenceProviderOpenAI) && isToken { + s.refetchOpenAIModels(context.Background(), scope) + } + return nil +} + +// refetchOpenAIModels stores the Codex model catalog the ChatGPT plan +// grants next to the provider's API-key models. Best effort: a failure +// leaves the existing catalog in place and the login still succeeds. +func (s *ConfigStore) refetchOpenAIModels(ctx context.Context, scope Scope) { + cfg := s.Config() + pc, ok := cfg.Providers.Get(string(catwalk.InferenceProviderOpenAI)) + if !ok || pc.OAuthToken == nil { + return + } + fetchModels := s.fetchOpenAIModels + if fetchModels == nil { + fetchModels = openai.Models + } + models, err := fetchModels(ctx, pc.OAuthToken) + if err != nil { + slog.Warn("Failed to fetch ChatGPT model catalog after auth", "error", err) + return + } + if err := s.update(scope, func(c *Config) map[string]any { + p, ok := c.Providers.Get(string(catwalk.InferenceProviderOpenAI)) + if !ok { + return nil + } + p.ChatGPTModels = models + c.Providers.Set(string(catwalk.InferenceProviderOpenAI), p) + return map[string]any{ + "providers.openai.chatgpt_models": models, + } + }); err != nil { + slog.Warn("Failed to persist ChatGPT model catalog", "error", err) + } +} + +// RefetchOpenAIChatGPTModels fills in the ChatGPT model catalog when the +// OpenAI provider is signed in but has none — because the fetch at login +// time failed, or the credentials predate the catalog. A no-op once the +// catalog exists, so callers can invoke it freely on model updates. +func (s *ConfigStore) RefetchOpenAIChatGPTModels(ctx context.Context) { + cfg := s.Config() + pc, ok := cfg.Providers.Get(string(catwalk.InferenceProviderOpenAI)) + if !ok || pc.OAuthToken == nil || len(pc.ChatGPTModels) > 0 { + return + } + s.refetchOpenAIModels(ctx, ScopeGlobal) +} + +// RefreshOAuthToken refreshes the OAuth token for the given provider. +// +// Providers like Hyper rotate refresh tokens: each exchange consumes the +// caller's refresh token, issues a new pair, and revokes the old one. If +// two crush instances (or two goroutines) refresh concurrently with the +// same stored refresh token, the second exchange reuses an already-revoked +// token, trips the provider's reuse detection, and revokes the entire +// token family — leaving both with dead tokens even though each refresh +// "succeeded". +// +// To prevent that, refreshes are single-flighted at two levels: +// +// - In-process: refreshSF collapses concurrent goroutines for the same +// provider into one attempt. +// - Cross-process: a per-provider advisory lock is held across the whole +// read-decide-exchange-write cycle, so only one process exchanges at a +// time. A process that acquires the lock after a peer rotated finds the +// peer's fresh token on disk and adopts it instead of exchanging. +func (s *ConfigStore) RefreshOAuthToken(ctx context.Context, scope Scope, providerID string) error { + key := fmt.Sprintf("%d\x00%s", scope, providerID) + _, err, _ := s.refreshSF.Do(key, func() (any, error) { + return nil, s.refreshOAuthTokenLocked(ctx, scope, providerID) + }) + return err +} + +// refreshOAuthTokenLocked performs the cross-process single-flighted +// refresh. It is invoked through refreshSF, so at most one goroutine per +// provider runs it at a time within this process. +func (s *ConfigStore) refreshOAuthTokenLocked(ctx context.Context, scope Scope, providerID string) error { + cfg := s.Config() + providerConfig, exists := cfg.Providers.Get(providerID) + if !exists { + return fmt.Errorf("provider %s not found", providerID) + } + if providerConfig.OAuthToken == nil { + return fmt.Errorf("provider %s does not have an OAuth token", providerID) + } + entryToken := providerConfig.OAuthToken + + // Acquire the per-provider cross-process refresh lock. This is a + // dedicated lock file, not the config-write lock, and it does not take + // s.mu — so the network exchange below cannot stall unrelated config + // operations. The deadline exceeds the exchange timeout so that a peer + // mid-exchange has time to publish a token we can adopt. Lock ordering: + // the refresh lock is always taken before the config-write lock (via + // SetConfigFields), never the reverse, so no deadlock is possible. + lockCtx, cancel := context.WithTimeout(ctx, refreshLockDeadline) + defer cancel() + release, lockErr := lock.File(lockCtx, s.refreshLockPath(providerID)) + if lockErr != nil { + // Could not acquire the lock (peer wedged or deadline hit). Prefer a + // usable token already on disk over forcing our own exchange, which + // would risk reusing a rotated refresh token. + if diskToken := s.usableDiskToken(scope, providerID, entryToken); diskToken != nil { + slog.Warn("Refresh lock unavailable; adopting token from disk", "provider", providerID, "error", lockErr) + return s.applyToken(providerConfig, diskToken, providerID) + } + return fmt.Errorf("acquire refresh lock for provider %s: %w", providerID, lockErr) + } + defer release() + + // Now that we hold the lock, disk is the authority on which credential + // is current: a peer may have rotated ours away while we waited. Adopt + // a newer token outright when it is still usable, and otherwise switch + // to its refresh token for the exchange below. Presenting our own + // already-rotated refresh token would trip the provider's reuse + // detection and revoke the whole family, forcing an interactive login. + if diskToken := s.newerDiskToken(scope, providerID, entryToken); diskToken != nil { + if !diskToken.IsExpired() { + slog.Info("Adopting token refreshed by another session", "provider", providerID) + return s.applyToken(providerConfig, diskToken, providerID) + } + slog.Info("Exchanging with refresh token rotated by another session", "provider", providerID) + entryToken = diskToken + } + + // Disk still holds our token (or no newer peer token exists) and we hold + // the lock, so we are the sole exchanger. Perform the exchange. + refreshedToken, refreshErr := s.exchange(ctx, providerID, entryToken.RefreshToken) + if refreshErr != nil { + // The exchange may have failed because a peer rotated the refresh + // token in a window we did not cover. Re-check disk: adopt a usable + // token, or retry once with the peer's newer refresh token. + if diskToken := s.newerDiskToken(scope, providerID, entryToken); diskToken != nil { + if !diskToken.IsExpired() { + slog.Info("Adopting token refreshed by another session after exchange failure", "provider", providerID) + return s.applyToken(providerConfig, diskToken, providerID) + } + slog.Info("Retrying exchange with refresh token rotated by another session", "provider", providerID) + refreshedToken, refreshErr = s.exchange(ctx, providerID, diskToken.RefreshToken) + } + } + if refreshErr != nil { + return fmt.Errorf("failed to refresh OAuth token for provider %s: %w", providerID, refreshErr) + } + + slog.Info("Successfully refreshed OAuth token", "provider", providerID) + if err := s.applyToken(providerConfig, refreshedToken, providerID); err != nil { + return err + } + + if err := s.SetConfigFields(scope, tokenFields(providerID, refreshedToken)); err != nil { + return fmt.Errorf("failed to persist refreshed token: %w", err) + } + return nil +} + +// tokenFields builds the config fields that persist an OAuth token. The +// OpenAI provider does not mirror the access token into api_key so a +// manually entered API key can coexist with the ChatGPT login. +func tokenFields(providerID string, token *oauth.Token) map[string]any { + fields := map[string]any{ + fmt.Sprintf("providers.%s.oauth", providerID): token, + } + if providerID != string(catwalk.InferenceProviderOpenAI) { + fields[fmt.Sprintf("providers.%s.api_key", providerID)] = token.AccessToken + } + return fields +} + +// WaitForTokenChange blocks until SignalAuthComplete is called for the +// given provider or the context is cancelled. It is used by OnAuthRefresh +// callbacks to wait for interactive re-authentication to complete before +// retrying a failed request. The channel is created atomically with the +// wait registration so a concurrent SignalAuthComplete cannot miss it. +func (s *ConfigStore) WaitForTokenChange(ctx context.Context, providerID string) error { + s.authSignalMu.Lock() + ch, ok := s.authSignals[providerID] + if !ok { + ch = make(chan struct{}) + if s.authSignals == nil { + s.authSignals = make(map[string]chan struct{}) + } + s.authSignals[providerID] = ch + } + s.authSignalMu.Unlock() + + select { + case <-ch: + // Remove the consumed signal so a subsequent + // SignalAuthComplete does not close an already-closed + // channel. + s.authSignalMu.Lock() + if s.authSignals[providerID] == ch { + delete(s.authSignals, providerID) + } + s.authSignalMu.Unlock() + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// SignalAuthComplete unblocks any goroutine waiting in WaitForTokenChange +// for the given provider. If no waiter exists yet, it pre-creates and +// immediately closes the channel so a subsequent WaitForTokenChange +// returns without blocking. This eliminates the race where the signal +// fires before the waiter registers. +func (s *ConfigStore) SignalAuthComplete(providerID string) { + s.authSignalMu.Lock() + defer s.authSignalMu.Unlock() + if ch, ok := s.authSignals[providerID]; ok { + delete(s.authSignals, providerID) + select { + case <-ch: + // Already closed by a previous signal; nothing to do. + default: + close(ch) + } + } else { + // No waiter yet. Pre-create a closed channel so the next + // WaitForTokenChange returns immediately. + if s.authSignals == nil { + s.authSignals = make(map[string]chan struct{}) + } + ch := make(chan struct{}) + close(ch) + s.authSignals[providerID] = ch + } +} + +// newerDiskToken returns the on-disk token for the provider when it is +// newer than entryToken — i.e. another session (possibly in another +// process) has already rotated the credential. It returns nil when disk +// holds nothing newer than what we started with. +// +// Newness is judged by expiry as well as identity, so a config file that +// somehow holds an older token cannot drag us backwards. The result may +// itself be expired: providers that rotate refresh tokens invalidate ours +// the moment a peer refreshes, so the peer's refresh token is the only one +// the provider will still accept even after its access token ages out. +// Callers decide whether to adopt the token wholesale or merely borrow its +// refresh token. +func (s *ConfigStore) newerDiskToken(scope Scope, providerID string, entryToken *oauth.Token) *oauth.Token { + diskToken, err := s.loadTokenFromDisk(scope, providerID) + if err != nil { + slog.Warn("Failed to read token from config file", "provider", providerID, "error", err) + return nil + } + if diskToken == nil { + return nil + } + if diskToken.AccessToken == entryToken.AccessToken { + // Same token we started with; nobody rotated since. + return nil + } + if diskToken.RefreshToken == "" && entryToken.RefreshToken != "" { + // Adopting would strand us with no way to refresh later, and + // there is nothing to borrow for an exchange. + return nil + } + if diskToken.ExpiresAt < entryToken.ExpiresAt { + // Older than ours; nothing to gain from adopting it. + return nil + } + return diskToken +} + +// usableDiskToken returns the on-disk token only when it is both newer +// than entryToken and still valid, meaning it can be adopted as-is with +// no exchange at all. +func (s *ConfigStore) usableDiskToken(scope Scope, providerID string, entryToken *oauth.Token) *oauth.Token { + diskToken := s.newerDiskToken(scope, providerID, entryToken) + if diskToken == nil || diskToken.IsExpired() { + return nil + } + return diskToken +} + +// exchange performs the provider-specific OAuth token exchange. Tests may +// override it via the exchangeToken field; production uses the real +// provider clients. +func (s *ConfigStore) exchange(ctx context.Context, providerID, refreshToken string) (*oauth.Token, error) { + if s.exchangeToken != nil { + return s.exchangeToken(ctx, providerID, refreshToken) + } + switch providerID { + case string(catwalk.InferenceProviderCopilot): + return copilot.RefreshToken(ctx, refreshToken) + case string(catwalk.InferenceProviderOpenAI): + return openai.RefreshToken(ctx, refreshToken) + case hyperp.Name: + return hyper.ExchangeToken(ctx, refreshToken) + default: + return nil, fmt.Errorf("OAuth refresh not supported for provider %s", providerID) + } +} + +// withRefreshLock runs fn while holding the per-provider cross-process +// refresh lock, so a credential write cannot interleave with a peer's +// token exchange. Acquisition is best effort: when the lock cannot be +// taken in time, fn runs anyway rather than blocking a write the user is +// waiting on. +func (s *ConfigStore) withRefreshLock(providerID string, fn func() error) error { + ctx, cancel := context.WithTimeout(context.Background(), credentialWriteLockDeadline) + defer cancel() + release, err := lock.File(ctx, s.refreshLockPath(providerID)) + if err != nil { + slog.Warn("Writing credentials without the refresh lock", "provider", providerID, "error", err) + return fn() + } + defer release() + return fn() +} + +// refreshLockPath returns the path to the per-provider cross-process refresh +// lock file. Lock files live under a dedicated locks/ subdirectory of the +// data dir so they do not clutter the config directory. The file is created +// on demand by lock.File and is never removed (flock keys on inode, not +// path). +func (s *ConfigStore) refreshLockPath(providerID string) string { + dir := filepath.Join(filepath.Dir(s.globalDataPath), "locks") + _ = os.MkdirAll(dir, 0o755) + return filepath.Join(dir, fmt.Sprintf("%s.refresh.lock", providerID)) +} + +// applyToken updates the in-memory provider config with the given token. +func (s *ConfigStore) applyToken(providerConfig ProviderConfig, token *oauth.Token, providerID string) error { + providerConfig.OAuthToken = token + // The OpenAI provider holds exactly one credential, so a ChatGPT + // token means there is no API key side by side with it. + if providerID != string(catwalk.InferenceProviderOpenAI) { + providerConfig.APIKey = token.AccessToken + } + if providerID == string(catwalk.InferenceProviderCopilot) { + providerConfig.SetupGitHubCopilot() + } + s.Config().Providers.Set(providerID, providerConfig) + return nil +} + +// loadTokenFromDisk reads the OAuth token for the given provider from the +// config file on disk. Returns nil if the token is not found or matches the +// current in-memory token. +func (s *ConfigStore) loadTokenFromDisk(scope Scope, providerID string) (*oauth.Token, error) { + path, err := s.configPath(scope) + if err != nil { + return nil, err + } + + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return nil, nil + } + return nil, err + } + + oauthKey := fmt.Sprintf("providers.%s.oauth", providerID) + oauthResult := gjson.Get(string(data), oauthKey) + if !oauthResult.Exists() { + return nil, nil + } + + var token oauth.Token + if err := json.Unmarshal([]byte(oauthResult.Raw), &token); err != nil { + return nil, err + } + + if token.AccessToken == "" { + return nil, nil + } + + return &token, nil +} + +// nextRecentModels computes the recent-models list for the given type +// after recording the supplied model at the front, operating on the +// provided config without persisting anything. It returns the new slice +// and whether it differs from cfg's current list. Callers fold the result +// into a clone they are about to publish. +func nextRecentModels(cfg *Config, modelType SelectedModelType, model SelectedModel) ([]SelectedModel, bool) { + if model.Provider == "" || model.Model == "" { + return nil, false + } + + eq := func(a, b SelectedModel) bool { + return a.Provider == b.Provider && a.Model == b.Model + } + + entry := SelectedModel{ + Provider: model.Provider, + Model: model.Model, + } + + current := cfg.RecentModels[modelType] + withoutCurrent := slices.DeleteFunc(slices.Clone(current), func(existing SelectedModel) bool { + return eq(existing, entry) + }) + + updated := append([]SelectedModel{entry}, withoutCurrent...) + if len(updated) > maxRecentModelsPerType { + updated = updated[:maxRecentModelsPerType] + } + + if slices.EqualFunc(current, updated, eq) { + return current, false + } + + return updated, true +} + +// NewTestStore creates a ConfigStore for testing purposes. +func NewTestStore(cfg *Config, loadedPaths ...string) *ConfigStore { + return &ConfigStore{ + config: cfg, + loadedPaths: loadedPaths, + resolver: NewShellVariableResolver(env.New()), + } +} + +// ImportCopilot attempts to import a GitHub Copilot token from disk. +func (s *ConfigStore) ImportCopilot() (*oauth.Token, bool) { + if s.HasConfigField(ScopeGlobal, "providers.copilot.api_key") || s.HasConfigField(ScopeGlobal, "providers.copilot.oauth") { + return nil, false + } + + diskToken, hasDiskToken := copilot.RefreshTokenFromDisk() + if !hasDiskToken { + return nil, false + } + + slog.Info("Found existing GitHub Copilot token on disk. Authenticating...") + token, err := copilot.RefreshToken(context.TODO(), diskToken) + if err != nil { + slog.Error("Unable to import GitHub Copilot token", "error", err) + return nil, false + } + + if err := s.SetProviderAPIKey(ScopeGlobal, string(catwalk.InferenceProviderCopilot), token); err != nil { + return token, false + } + + if err := s.SetConfigFields(ScopeGlobal, map[string]any{ + "providers.copilot.api_key": token.AccessToken, + "providers.copilot.oauth": token, + }); err != nil { + slog.Error("Unable to save GitHub Copilot token to disk", "error", err) + } + + slog.Info("GitHub Copilot successfully imported") + return token, true +} + +// StalenessResult contains the result of a staleness check. +type StalenessResult struct { + Dirty bool + Changed []string + Missing []string + Errors map[string]error // stat errors by path +} + +// ConfigStaleness checks whether any tracked config files have changed on disk +// since the last snapshot. Returns dirty=true if any files changed or went +// missing, along with sorted lists of affected paths. Stat errors are +// captured in Errors map but still treated as non-existence for dirty detection. +func (s *ConfigStore) ConfigStaleness() StalenessResult { + var result StalenessResult + result.Errors = make(map[string]error) + + for _, path := range s.trackedConfigPaths { + snapshot, hadSnapshot := s.snapshots[path] + + info, err := os.Stat(path) + exists := err == nil && !info.IsDir() + + if err != nil && !os.IsNotExist(err) { + // Capture permission/IO errors separately from non-existence + result.Errors[path] = err + result.Dirty = true + } + + if !exists { + if hadSnapshot && snapshot.Exists { + // File existed before but now missing + result.Missing = append(result.Missing, path) + result.Dirty = true + } + continue + } + + // File exists now + if !hadSnapshot || !snapshot.Exists { + // File didn't exist before but does now + result.Changed = append(result.Changed, path) + result.Dirty = true + continue + } + + // Check for content or metadata changes + if snapshot.Size != info.Size() || snapshot.ModTime != info.ModTime().UnixNano() { + result.Changed = append(result.Changed, path) + result.Dirty = true + } + } + + // Sort for deterministic output + slices.Sort(result.Changed) + slices.Sort(result.Missing) + + return result +} + +// RefreshStalenessSnapshot captures fresh snapshots of all tracked config files. +// Call this after reloading config to clear dirty state. +func (s *ConfigStore) RefreshStalenessSnapshot() error { + if s.snapshots == nil { + s.snapshots = make(map[string]fileSnapshot) + } + + for _, path := range s.trackedConfigPaths { + info, err := os.Stat(path) + exists := err == nil && !info.IsDir() + + snapshot := fileSnapshot{ + Path: path, + Exists: exists, + } + + if exists { + snapshot.Size = info.Size() + snapshot.ModTime = info.ModTime().UnixNano() + } + + s.snapshots[path] = snapshot + } + + return nil +} + +// CaptureStalenessSnapshot captures snapshots for the given paths, building the +// tracked config paths list. Paths are deduplicated and normalized. +func (s *ConfigStore) CaptureStalenessSnapshot(paths []string) { + // Build unique set of normalized paths + seen := make(map[string]struct{}) + for _, p := range paths { + if p == "" { + continue + } + // Normalize path + abs, err := filepath.Abs(p) + if err != nil { + abs = p + } + seen[abs] = struct{}{} + } + + // Also track workspace and global config paths if set + if s.workspacePath != "" { + abs, err := filepath.Abs(s.workspacePath) + if err == nil { + seen[abs] = struct{}{} + } + } + if s.globalDataPath != "" { + abs, err := filepath.Abs(s.globalDataPath) + if err == nil { + seen[abs] = struct{}{} + } + } + + // Build sorted list for deterministic ordering + s.trackedConfigPaths = make([]string, 0, len(seen)) + for p := range seen { + s.trackedConfigPaths = append(s.trackedConfigPaths, p) + } + slices.Sort(s.trackedConfigPaths) + + // Capture initial snapshots + s.RefreshStalenessSnapshot() +} + +// captureStalenessSnapshot is an alias for CaptureStalenessSnapshot for internal use. +func (s *ConfigStore) captureStalenessSnapshot(paths []string) { + s.CaptureStalenessSnapshot(paths) +} + +// ReloadFromDisk re-runs the config load/merge flow and updates the in-memory +// config atomically. It rebuilds the staleness snapshot after successful reload. +// On failure, the store state is rolled back to its previous state. +// Concurrent calls are serialised via writeMu. +func (s *ConfigStore) ReloadFromDisk(ctx context.Context) error { + if s.workingDir == "" { + return fmt.Errorf("cannot reload: working directory not set") + } + s.writeMu.Lock() + defer s.writeMu.Unlock() + return s.reloadFromDiskLocked(ctx) +} + +// reloadFromDiskLocked performs the actual reload. Caller must hold writeMu. +func (s *ConfigStore) reloadFromDiskLocked(ctx context.Context) error { + // Migrate deprecated disable_notifications before reloading config. + migrateDisableNotifications() + + configPaths := lookupConfigs(s.workingDir) + cfg, loadedPaths, err := loadFromConfigPaths(ctx, configPaths) + if err != nil { + return fmt.Errorf("failed to reload config: %w", err) + } + + // Apply defaults (using existing data directory if set) + var dataDir string + if cur := s.Config(); cur != nil && cur.Options != nil { + dataDir = cur.Options.DataDirectory + } + cfg.setDefaults(s.workingDir, dataDir) + + // Merge workspace config if present + workspacePath := filepath.Join(cfg.Options.DataDirectory, fmt.Sprintf("%s.json", appName)) + if wsData, err := os.ReadFile(workspacePath); err == nil && len(wsData) > 0 { + if !json.Valid(wsData) { + return fmt.Errorf("invalid JSON in config file %s", workspacePath) + } + merged, mergeErr := loadFromBytes(append([][]byte{mustMarshalConfig(cfg)}, wsData)) + if mergeErr == nil { + dataDir := cfg.Options.DataDirectory + *cfg = *merged + cfg.setDefaults(s.workingDir, dataDir) + loadedPaths = append(loadedPaths, workspacePath) + } + } + + // Validate hooks after all config merging is complete so matcher + // regexes are recompiled on the reloaded config (mirrors Load). + if err := cfg.ValidateHooks(); err != nil { + return fmt.Errorf("invalid hook configuration on reload: %w", err) + } + + // Save current state for potential rollback BEFORE configureProviders, + // which may write to disk via RemoveConfigField (e.g. removing stale + // OAuth providers). Capturing after would snapshot a config that has + // already been mutated, and the rollback would restore corrupted state. + oldConfig := s.Config() + oldLoadedPaths := s.loadedPaths + oldResolver := s.resolver + oldKnownProviders := s.knownProviders + oldOverrides := s.overrides + oldWorkspacePath := s.workspacePath + + // Preserve runtime overrides + overrides := s.overrides + + // Reapply model choices made in this instance. The global config file is + // shared, so it may now name a model a sibling instance selected; a + // reload triggered by an unrelated write must not swap the user's model + // mid-session. An external edit to the config still takes effect for any + // model type this instance never chose. + maps.Copy(cfg.Models, overrides.Models) + + // Reconfigure providers + env := env.New() + resolver := NewShellVariableResolver(env) + + // Apply top-level env vars before configuring providers so variables + // like AWS_PROFILE are visible to the AWS SDK credential chain. + cfg.applyEnv(resolver) + + providers, err := Providers(cfg) + if err != nil { + if len(providers) == 0 { + return fmt.Errorf("failed to load providers during reload: %w", err) + } + slog.Warn("Reload continuing with the previously known providers", "error", err) + } + + if err := cfg.configureProviders(ctx, s, env, resolver, providers); err != nil { + return fmt.Errorf("failed to configure providers during reload: %w", err) + } + + // Update store state BEFORE running model/agent setup (so they see new config) + s.setConfig(cfg) + s.loadedPaths = loadedPaths + s.resolver = resolver + s.knownProviders = providers + s.overrides = overrides + s.workspacePath = workspacePath + + // Mirror startup flow: setup models and agents against NEW config. + var setupErr error + if !cfg.IsConfigured() { + slog.Warn("No providers configured after reload") + } else { + resolved, resolveErr := resolveSelectedModels(cfg, providers) + if resolveErr != nil { + setupErr = fmt.Errorf("failed to configure selected models during reload: %w", resolveErr) + } else { + cfg.Models[SelectedModelTypeLarge] = resolved.Large + cfg.Models[SelectedModelTypeSmall] = resolved.Small + s.SetupAgents() + } + } + + // Rollback on setup failure + if setupErr != nil { + s.setConfig(oldConfig) + s.loadedPaths = oldLoadedPaths + s.resolver = oldResolver + s.knownProviders = oldKnownProviders + s.overrides = oldOverrides + s.workspacePath = oldWorkspacePath + return setupErr + } + + // Rebuild staleness tracking. Track every discovered config path, not + // just the ones that loaded, so a config file created after this reload + // is detected as a change on the next staleness check. + s.captureStalenessSnapshot(append(slices.Clone(configPaths), loadedPaths...)) + + return nil +} + +// autoReload conditionally reloads config from disk after writes. +// It returns nil (no error) for expected skip cases: when auto-reload is +// disabled during load/reload flows, or when working directory is not set +// (e.g., during testing). Only actual reload failures return an error. +func (s *ConfigStore) autoReload(ctx context.Context) error { + if s.workingDir == "" { + return nil // Expected skip: working directory not set + } + // Skip if a reload is already in progress. This handles both + // concurrent auto-reloads after parallel writes and re-entrant + // calls from configureProviders during a reload. + // + // Note: if a write completes after the in-progress reload has + // already read the config file, that write won't be reflected in + // memory until the next reload. This is acceptable because writes + // are rare and the next user action or file-watch tick will pick + // up the change. Callers that need guaranteed fresh state after a + // write should call ReloadFromDisk explicitly. + if !s.writeMu.TryLock() { + return nil + } + defer s.writeMu.Unlock() + return s.reloadFromDiskLocked(ctx) +} diff --git a/internal/config/update_model_bench_test.go b/internal/config/update_model_bench_test.go new file mode 100644 index 0000000000000000000000000000000000000000..1c087ce1864489254966d33614811d76ce8c166a --- /dev/null +++ b/internal/config/update_model_bench_test.go @@ -0,0 +1,108 @@ +package config + +import ( + "context" + "os" + "path/filepath" + "testing" +) + +// BenchmarkUpdatePreferredModel measures the full write cost of a single +// model selection, which is what the user experiences when picking a model +// in the TUI. It guards against regressing back to a full config reload on +// every selection. +func BenchmarkUpdatePreferredModel(b *testing.B) { + dir := b.TempDir() + configPath := filepath.Join(dir, "crush.json") + + b.Setenv("CRUSH_GLOBAL_CONFIG", dir) + b.Setenv("CRUSH_GLOBAL_DATA", dir) + resetProviderState() + b.Cleanup(resetProviderState) + + cfg := `{ + "models": { + "large": {"provider": "openai", "model": "gpt-4"}, + "small": {"provider": "openai", "model": "gpt-4o-mini"} + }, + "providers": { + "openai": { + "api_key": "test-key", + "models": [ + {"id": "gpt-4", "name": "GPT-4"}, + {"id": "gpt-4o", "name": "GPT-4o"}, + {"id": "gpt-4o-mini", "name": "GPT-4o mini"} + ] + }, + "anthropic": { + "api_key": "test-key-2", + "models": [ + {"id": "claude-3", "name": "Claude 3"}, + {"id": "claude-3-5", "name": "Claude 3.5"} + ] + } + } + }` + if err := os.WriteFile(configPath, []byte(cfg), 0o600); err != nil { + b.Fatal(err) + } + + store, err := Load(dir, dir, false) + if err != nil { + b.Fatal(err) + } + store.globalDataPath = configPath + + models := []SelectedModel{ + {Provider: "openai", Model: "gpt-4"}, + {Provider: "anthropic", Model: "claude-3"}, + {Provider: "openai", Model: "gpt-4o"}, + {Provider: "anthropic", Model: "claude-3-5"}, + } + + b.ReportAllocs() + i := 0 + for b.Loop() { + m := models[i%len(models)] + if err := store.UpdatePreferredModel(ScopeGlobal, SelectedModelTypeLarge, m); err != nil { + b.Fatal(err) + } + i++ + } +} + +// BenchmarkReloadFromDisk isolates the full reload cost for comparison with +// BenchmarkUpdatePreferredModel. +func BenchmarkReloadFromDisk(b *testing.B) { + dir := b.TempDir() + configPath := filepath.Join(dir, "crush.json") + + b.Setenv("CRUSH_GLOBAL_CONFIG", dir) + b.Setenv("CRUSH_GLOBAL_DATA", dir) + resetProviderState() + b.Cleanup(resetProviderState) + + cfg := `{ + "models": {"large": {"provider": "openai", "model": "gpt-4"}}, + "providers": { + "openai": {"api_key": "test-key", "models": [{"id": "gpt-4", "name": "GPT-4"}]} + } + }` + if err := os.WriteFile(configPath, []byte(cfg), 0o600); err != nil { + b.Fatal(err) + } + + store, err := Load(dir, dir, false) + if err != nil { + b.Fatal(err) + } + store.globalDataPath = configPath + + ctx := context.Background() + b.ReportAllocs() + for b.Loop() { + if err := store.ReloadFromDisk(ctx); err != nil { + b.Fatal(err) + } + } +} diff --git a/internal/csync/doc.go b/internal/csync/doc.go new file mode 100644 index 0000000000000000000000000000000000000000..b44fa2dac52172a8753e620e19f5c8d0a0481618 --- /dev/null +++ b/internal/csync/doc.go @@ -0,0 +1,3 @@ +// Package csync provides concurrent data structures for safe access in +// multi-threaded environments. +package csync diff --git a/internal/csync/maps.go b/internal/csync/maps.go new file mode 100644 index 0000000000000000000000000000000000000000..9de501f64e80ffc1771440674d5aa35dc638eb4c --- /dev/null +++ b/internal/csync/maps.go @@ -0,0 +1,174 @@ +package csync + +import ( + "encoding/json" + "iter" + "maps" + "sync" +) + +// Map is a concurrent map implementation that provides thread-safe access. +type Map[K comparable, V any] struct { + inner map[K]V + mu sync.RWMutex +} + +// NewMap creates a new thread-safe map with the specified key and value types. +func NewMap[K comparable, V any]() *Map[K, V] { + return &Map[K, V]{ + inner: make(map[K]V), + } +} + +// NewMapFrom creates a new thread-safe map from an existing map. +func NewMapFrom[K comparable, V any](m map[K]V) *Map[K, V] { + return &Map[K, V]{ + inner: m, + } +} + +// NewLazyMap creates a new lazy-loaded map. The provided load function is +// executed in a separate goroutine to populate the map. +func NewLazyMap[K comparable, V any](load func() map[K]V) *Map[K, V] { + m := &Map[K, V]{} + m.mu.Lock() + go func() { + defer m.mu.Unlock() + m.inner = load() + }() + return m +} + +// Reset replaces the inner map with the new one. +func (m *Map[K, V]) Reset(input map[K]V) { + m.mu.Lock() + defer m.mu.Unlock() + m.inner = input +} + +// Set sets the value for the specified key in the map. +func (m *Map[K, V]) Set(key K, value V) { + m.mu.Lock() + defer m.mu.Unlock() + m.inner[key] = value +} + +// Del deletes the specified key from the map. +func (m *Map[K, V]) Del(key K) { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.inner, key) +} + +// CompareAndDelete deletes the key only if the current value matches the +// expected pointer. Returns true if the deletion occurred. This is the +// ABA-safe cleanup primitive: it prevents a deferred cleanup from removing +// a value that was replaced by a newer writer in the window between the +// explicit Del and the deferred Del. +func (m *Map[K, V]) CompareAndDelete(key K, expected any) bool { + m.mu.Lock() + defer m.mu.Unlock() + current, ok := m.inner[key] + if !ok { + return false + } + if any(current) != expected { + return false + } + delete(m.inner, key) + return true +} + +// Get gets the value for the specified key from the map. +func (m *Map[K, V]) Get(key K) (V, bool) { + m.mu.RLock() + defer m.mu.RUnlock() + v, ok := m.inner[key] + return v, ok +} + +// Len returns the number of items in the map. +func (m *Map[K, V]) Len() int { + m.mu.RLock() + defer m.mu.RUnlock() + return len(m.inner) +} + +// GetOrSet gets and returns the key if it exists, otherwise, it executes the +// given function, set its return value for the given key, and returns it. +func (m *Map[K, V]) GetOrSet(key K, fn func() V) V { + got, ok := m.Get(key) + if ok { + return got + } + value := fn() + m.Set(key, value) + return value +} + +// Take gets an item and then deletes it. +func (m *Map[K, V]) Take(key K) (V, bool) { + m.mu.Lock() + defer m.mu.Unlock() + v, ok := m.inner[key] + delete(m.inner, key) + return v, ok +} + +// Copy returns a copy of the inner map. +func (m *Map[K, V]) Copy() map[K]V { + m.mu.RLock() + defer m.mu.RUnlock() + return maps.Clone(m.inner) +} + +// Seq2 returns an iter.Seq2 that yields key-value pairs from the map. +func (m *Map[K, V]) Seq2() iter.Seq2[K, V] { + dst := m.Copy() + return func(yield func(K, V) bool) { + for k, v := range dst { + if !yield(k, v) { + return + } + } + } +} + +// Seq returns an iter.Seq that yields values from the map. +func (m *Map[K, V]) Seq() iter.Seq[V] { + return func(yield func(V) bool) { + for _, v := range m.Seq2() { + if !yield(v) { + return + } + } + } +} + +var ( + _ json.Unmarshaler = &Map[string, any]{} + _ json.Marshaler = &Map[string, any]{} +) + +// JSONSchemaAlias returns the underlying map type for JSON schema generation. +// Value receiver is required because github.com/invopop/jsonschema checks +// interface satisfaction on the non-pointer type after stripping pointers. +func (Map[K, V]) JSONSchemaAlias() any { //nolint + m := map[K]V{} + return m +} + +// UnmarshalJSON implements json.Unmarshaler. +func (m *Map[K, V]) UnmarshalJSON(data []byte) error { + m.mu.Lock() + defer m.mu.Unlock() + m.inner = make(map[K]V) + return json.Unmarshal(data, &m.inner) +} + +// MarshalJSON implements json.Marshaler. +func (m *Map[K, V]) MarshalJSON() ([]byte, error) { + m.mu.RLock() + defer m.mu.RUnlock() + return json.Marshal(m.inner) +} diff --git a/internal/csync/maps_test.go b/internal/csync/maps_test.go new file mode 100644 index 0000000000000000000000000000000000000000..98ef763a7258eb02f285ea986d7d0248733f3701 --- /dev/null +++ b/internal/csync/maps_test.go @@ -0,0 +1,770 @@ +package csync + +import ( + "encoding/json" + "maps" + "sync" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/stretchr/testify/require" +) + +func TestNewMap(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + require.NotNil(t, m) + require.NotNil(t, m.inner) + require.Equal(t, 0, m.Len()) +} + +func TestNewMapFrom(t *testing.T) { + t.Parallel() + + original := map[string]int{ + "key1": 1, + "key2": 2, + } + + m := NewMapFrom(original) + require.NotNil(t, m) + require.Equal(t, original, m.inner) + require.Equal(t, 2, m.Len()) + + value, ok := m.Get("key1") + require.True(t, ok) + require.Equal(t, 1, value) +} + +func TestNewLazyMap(t *testing.T) { + t.Parallel() + + synctest.Test(t, func(t *testing.T) { + t.Helper() + + waiter := sync.Mutex{} + waiter.Lock() + var loadCalled atomic.Bool + + loadFunc := func() map[string]int { + waiter.Lock() + defer waiter.Unlock() + loadCalled.Store(true) + return map[string]int{ + "key1": 1, + "key2": 2, + } + } + + m := NewLazyMap(loadFunc) + require.NotNil(t, m) + + waiter.Unlock() // Allow the load function to proceed + time.Sleep(100 * time.Millisecond) + require.True(t, loadCalled.Load()) + require.Equal(t, 2, m.Len()) + + value, ok := m.Get("key1") + require.True(t, ok) + require.Equal(t, 1, value) + }) +} + +func TestMap_Reset(t *testing.T) { + t.Parallel() + + m := NewMapFrom(map[string]int{ + "a": 10, + }) + + m.Reset(map[string]int{ + "b": 20, + }) + value, ok := m.Get("b") + require.True(t, ok) + require.Equal(t, 20, value) + require.Equal(t, 1, m.Len()) +} + +func TestMap_Set(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + + m.Set("key1", 42) + value, ok := m.Get("key1") + require.True(t, ok) + require.Equal(t, 42, value) + require.Equal(t, 1, m.Len()) + + m.Set("key1", 100) + value, ok = m.Get("key1") + require.True(t, ok) + require.Equal(t, 100, value) + require.Equal(t, 1, m.Len()) +} + +func TestMap_GetOrSet(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + + require.Equal(t, 42, m.GetOrSet("key1", func() int { return 42 })) + require.Equal(t, 42, m.GetOrSet("key1", func() int { return 99999 })) + require.Equal(t, 1, m.Len()) +} + +func TestMap_Get(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + + value, ok := m.Get("nonexistent") + require.False(t, ok) + require.Equal(t, 0, value) + + m.Set("key1", 42) + value, ok = m.Get("key1") + require.True(t, ok) + require.Equal(t, 42, value) +} + +func TestMap_Del(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + m.Set("key1", 42) + m.Set("key2", 100) + + require.Equal(t, 2, m.Len()) + + m.Del("key1") + _, ok := m.Get("key1") + require.False(t, ok) + require.Equal(t, 1, m.Len()) + + value, ok := m.Get("key2") + require.True(t, ok) + require.Equal(t, 100, value) + + m.Del("nonexistent") + require.Equal(t, 1, m.Len()) +} + +func TestMap_Len(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + require.Equal(t, 0, m.Len()) + + m.Set("key1", 1) + require.Equal(t, 1, m.Len()) + + m.Set("key2", 2) + require.Equal(t, 2, m.Len()) + + m.Del("key1") + require.Equal(t, 1, m.Len()) + + m.Del("key2") + require.Equal(t, 0, m.Len()) +} + +// TestMap_CompareAndDelete verifies that CompareAndDelete only removes the +// entry when the stored pointer matches the expected value. This is the +// ABA-safe cleanup primitive used by the dispatch completion path to prevent +// a deferred cleanup from wiping a newer run's entry. +func TestMap_CompareAndDelete(t *testing.T) { + t.Parallel() + + type entry struct{ val int } + m := NewMap[string, *entry]() + + original := &entry{val: 1} + m.Set("k", original) + + // Matching pointer: deletion succeeds. + require.True(t, m.CompareAndDelete("k", any(original))) + _, ok := m.Get("k") + require.False(t, ok) + + // Re-set with a new value, try to delete with the stale pointer. + second := &entry{val: 2} + m.Set("k", second) + require.False(t, m.CompareAndDelete("k", any(original)), "stale pointer must not delete") + got, ok := m.Get("k") + require.True(t, ok) + require.Equal(t, 2, got.val) + + // Missing key: returns false. + require.False(t, m.CompareAndDelete("absent", any(original))) +} + +func TestMap_Take(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + m.Set("key1", 42) + m.Set("key2", 100) + + require.Equal(t, 2, m.Len()) + + value, ok := m.Take("key1") + require.True(t, ok) + require.Equal(t, 42, value) + require.Equal(t, 1, m.Len()) + + _, exists := m.Get("key1") + require.False(t, exists) + + value, ok = m.Get("key2") + require.True(t, ok) + require.Equal(t, 100, value) +} + +func TestMap_Take_NonexistentKey(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + m.Set("key1", 42) + + value, ok := m.Take("nonexistent") + require.False(t, ok) + require.Equal(t, 0, value) + require.Equal(t, 1, m.Len()) + + value, ok = m.Get("key1") + require.True(t, ok) + require.Equal(t, 42, value) +} + +func TestMap_Take_EmptyMap(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + + value, ok := m.Take("key1") + require.False(t, ok) + require.Equal(t, 0, value) + require.Equal(t, 0, m.Len()) +} + +func TestMap_Take_SameKeyTwice(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + m.Set("key1", 42) + + value, ok := m.Take("key1") + require.True(t, ok) + require.Equal(t, 42, value) + require.Equal(t, 0, m.Len()) + + value, ok = m.Take("key1") + require.False(t, ok) + require.Equal(t, 0, value) + require.Equal(t, 0, m.Len()) +} + +func TestMap_Seq2(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + m.Set("key1", 1) + m.Set("key2", 2) + m.Set("key3", 3) + + collected := maps.Collect(m.Seq2()) + + require.Equal(t, 3, len(collected)) + require.Equal(t, 1, collected["key1"]) + require.Equal(t, 2, collected["key2"]) + require.Equal(t, 3, collected["key3"]) +} + +func TestMap_Seq2_EarlyReturn(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + m.Set("key1", 1) + m.Set("key2", 2) + m.Set("key3", 3) + + count := 0 + for range m.Seq2() { + count++ + if count == 2 { + break + } + } + + require.Equal(t, 2, count) +} + +func TestMap_Seq2_EmptyMap(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + + count := 0 + for range m.Seq2() { + count++ + } + + require.Equal(t, 0, count) +} + +func TestMap_Seq(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + m.Set("key1", 1) + m.Set("key2", 2) + m.Set("key3", 3) + + collected := make([]int, 0) + for v := range m.Seq() { + collected = append(collected, v) + } + + require.Equal(t, 3, len(collected)) + require.Contains(t, collected, 1) + require.Contains(t, collected, 2) + require.Contains(t, collected, 3) +} + +func TestMap_Seq_EarlyReturn(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + m.Set("key1", 1) + m.Set("key2", 2) + m.Set("key3", 3) + + count := 0 + for range m.Seq() { + count++ + if count == 2 { + break + } + } + + require.Equal(t, 2, count) +} + +func TestMap_Seq_EmptyMap(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + + count := 0 + for range m.Seq() { + count++ + } + + require.Equal(t, 0, count) +} + +func TestMap_MarshalJSON(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + m.Set("key1", 1) + m.Set("key2", 2) + + data, err := json.Marshal(m) + require.NoError(t, err) + + result := &Map[string, int]{} + err = json.Unmarshal(data, result) + require.NoError(t, err) + require.Equal(t, 2, result.Len()) + v1, _ := result.Get("key1") + v2, _ := result.Get("key2") + require.Equal(t, 1, v1) + require.Equal(t, 2, v2) +} + +func TestMap_MarshalJSON_EmptyMap(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + + data, err := json.Marshal(m) + require.NoError(t, err) + require.Equal(t, "{}", string(data)) +} + +func TestMap_UnmarshalJSON(t *testing.T) { + t.Parallel() + + jsonData := `{"key1": 1, "key2": 2}` + + m := NewMap[string, int]() + err := json.Unmarshal([]byte(jsonData), m) + require.NoError(t, err) + + require.Equal(t, 2, m.Len()) + value, ok := m.Get("key1") + require.True(t, ok) + require.Equal(t, 1, value) + + value, ok = m.Get("key2") + require.True(t, ok) + require.Equal(t, 2, value) +} + +func TestMap_UnmarshalJSON_EmptyJSON(t *testing.T) { + t.Parallel() + + jsonData := `{}` + + m := NewMap[string, int]() + err := json.Unmarshal([]byte(jsonData), m) + require.NoError(t, err) + require.Equal(t, 0, m.Len()) +} + +func TestMap_UnmarshalJSON_InvalidJSON(t *testing.T) { + t.Parallel() + + jsonData := `{"key1": 1, "key2":}` + + m := NewMap[string, int]() + err := json.Unmarshal([]byte(jsonData), m) + require.Error(t, err) +} + +func TestMap_UnmarshalJSON_OverwritesExistingData(t *testing.T) { + t.Parallel() + + m := NewMap[string, int]() + m.Set("existing", 999) + + jsonData := `{"key1": 1, "key2": 2}` + err := json.Unmarshal([]byte(jsonData), m) + require.NoError(t, err) + + require.Equal(t, 2, m.Len()) + _, ok := m.Get("existing") + require.False(t, ok) + + value, ok := m.Get("key1") + require.True(t, ok) + require.Equal(t, 1, value) +} + +func TestMap_JSONRoundTrip(t *testing.T) { + t.Parallel() + + original := NewMap[string, int]() + original.Set("key1", 1) + original.Set("key2", 2) + original.Set("key3", 3) + + data, err := json.Marshal(original) + require.NoError(t, err) + + restored := NewMap[string, int]() + err = json.Unmarshal(data, restored) + require.NoError(t, err) + + require.Equal(t, original.Len(), restored.Len()) + + for k, v := range original.Seq2() { + restoredValue, ok := restored.Get(k) + require.True(t, ok) + require.Equal(t, v, restoredValue) + } +} + +func TestMap_ConcurrentAccess(t *testing.T) { + t.Parallel() + + m := NewMap[int, int]() + const numGoroutines = 100 + const numOperations = 100 + + var wg sync.WaitGroup + wg.Add(numGoroutines) + + for i := range numGoroutines { + go func(id int) { + defer wg.Done() + for j := range numOperations { + key := id*numOperations + j + m.Set(key, key*2) + value, ok := m.Get(key) + require.True(t, ok) + require.Equal(t, key*2, value) + } + }(i) + } + + wg.Wait() + + require.Equal(t, numGoroutines*numOperations, m.Len()) +} + +func TestMap_ConcurrentReadWrite(t *testing.T) { + t.Parallel() + + m := NewMap[int, int]() + const numReaders = 50 + const numWriters = 50 + const numOperations = 100 + + for i := range 1000 { + m.Set(i, i) + } + + var wg sync.WaitGroup + wg.Add(numReaders + numWriters) + + for range numReaders { + go func() { + defer wg.Done() + for j := range numOperations { + key := j % 1000 + value, ok := m.Get(key) + if ok { + require.Equal(t, key, value) + } + _ = m.Len() + } + }() + } + + for i := range numWriters { + go func(id int) { + defer wg.Done() + for j := range numOperations { + key := 1000 + id*numOperations + j + m.Set(key, key) + if j%10 == 0 { + m.Del(key) + } + } + }(i) + } + + wg.Wait() +} + +func TestMap_ConcurrentSeq2(t *testing.T) { + t.Parallel() + + m := NewMap[int, int]() + for i := range 100 { + m.Set(i, i*2) + } + + var wg sync.WaitGroup + const numIterators = 10 + + wg.Add(numIterators) + for range numIterators { + go func() { + defer wg.Done() + count := 0 + for k, v := range m.Seq2() { + require.Equal(t, k*2, v) + count++ + } + require.Equal(t, 100, count) + }() + } + + wg.Wait() +} + +func TestMap_ConcurrentSeq(t *testing.T) { + t.Parallel() + + m := NewMap[int, int]() + for i := range 100 { + m.Set(i, i*2) + } + + var wg sync.WaitGroup + const numIterators = 10 + + wg.Add(numIterators) + for range numIterators { + go func() { + defer wg.Done() + count := 0 + values := make(map[int]bool) + for v := range m.Seq() { + values[v] = true + count++ + } + require.Equal(t, 100, count) + for i := range 100 { + require.True(t, values[i*2]) + } + }() + } + + wg.Wait() +} + +func TestMap_ConcurrentTake(t *testing.T) { + t.Parallel() + + m := NewMap[int, int]() + const numItems = 1000 + + for i := range numItems { + m.Set(i, i*2) + } + + var wg sync.WaitGroup + const numWorkers = 10 + taken := make([][]int, numWorkers) + + wg.Add(numWorkers) + for i := range numWorkers { + go func(workerID int) { + defer wg.Done() + taken[workerID] = make([]int, 0) + for j := workerID; j < numItems; j += numWorkers { + if value, ok := m.Take(j); ok { + taken[workerID] = append(taken[workerID], value) + } + } + }(i) + } + + wg.Wait() + + require.Equal(t, 0, m.Len()) + + allTaken := make(map[int]bool) + for _, workerTaken := range taken { + for _, value := range workerTaken { + require.False(t, allTaken[value], "Value %d was taken multiple times", value) + allTaken[value] = true + } + } + + require.Equal(t, numItems, len(allTaken)) + for i := range numItems { + require.True(t, allTaken[i*2], "Expected value %d to be taken", i*2) + } +} + +func TestMap_TypeSafety(t *testing.T) { + t.Parallel() + + stringIntMap := NewMap[string, int]() + stringIntMap.Set("key", 42) + value, ok := stringIntMap.Get("key") + require.True(t, ok) + require.Equal(t, 42, value) + + intStringMap := NewMap[int, string]() + intStringMap.Set(42, "value") + strValue, ok := intStringMap.Get(42) + require.True(t, ok) + require.Equal(t, "value", strValue) + + structMap := NewMap[string, struct{ Name string }]() + structMap.Set("key", struct{ Name string }{Name: "test"}) + structValue, ok := structMap.Get("key") + require.True(t, ok) + require.Equal(t, "test", structValue.Name) +} + +func TestMap_InterfaceCompliance(t *testing.T) { + t.Parallel() + + var _ json.Marshaler = &Map[string, any]{} + var _ json.Unmarshaler = &Map[string, any]{} +} + +func BenchmarkMap_Set(b *testing.B) { + m := NewMap[int, int]() + + for i := 0; b.Loop(); i++ { + m.Set(i, i*2) + } +} + +func BenchmarkMap_Get(b *testing.B) { + m := NewMap[int, int]() + for i := range 1000 { + m.Set(i, i*2) + } + + for i := 0; b.Loop(); i++ { + m.Get(i % 1000) + } +} + +func BenchmarkMap_Seq2(b *testing.B) { + m := NewMap[int, int]() + for i := range 1000 { + m.Set(i, i*2) + } + + for b.Loop() { + for range m.Seq2() { + } + } +} + +func BenchmarkMap_Seq(b *testing.B) { + m := NewMap[int, int]() + for i := range 1000 { + m.Set(i, i*2) + } + + for b.Loop() { + for range m.Seq() { + } + } +} + +func BenchmarkMap_Take(b *testing.B) { + m := NewMap[int, int]() + for i := range 1000 { + m.Set(i, i*2) + } + + b.ResetTimer() + for i := 0; b.Loop(); i++ { + key := i % 1000 + m.Take(key) + if i%1000 == 999 { + b.StopTimer() + for j := range 1000 { + m.Set(j, j*2) + } + b.StartTimer() + } + } +} + +func BenchmarkMap_ConcurrentReadWrite(b *testing.B) { + m := NewMap[int, int]() + for i := range 1000 { + m.Set(i, i*2) + } + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + i := 0 + for pb.Next() { + if i%2 == 0 { + m.Get(i % 1000) + } else { + m.Set(i+1000, i*2) + } + i++ + } + }) +} diff --git a/internal/csync/slices.go b/internal/csync/slices.go new file mode 100644 index 0000000000000000000000000000000000000000..fcce9881b6e27021adcc9462b123f49d469dcd9f --- /dev/null +++ b/internal/csync/slices.go @@ -0,0 +1,121 @@ +package csync + +import ( + "iter" + "sync" +) + +// LazySlice is a thread-safe lazy-loaded slice. +type LazySlice[K any] struct { + inner []K + wg sync.WaitGroup +} + +// NewLazySlice creates a new slice and runs the [load] function in a goroutine +// to populate it. +func NewLazySlice[K any](load func() []K) *LazySlice[K] { + s := &LazySlice[K]{} + s.wg.Go(func() { + s.inner = load() + }) + return s +} + +// Seq returns an iterator that yields elements from the slice. +func (s *LazySlice[K]) Seq() iter.Seq[K] { + s.wg.Wait() + return func(yield func(K) bool) { + for _, v := range s.inner { + if !yield(v) { + return + } + } + } +} + +// Slice is a thread-safe slice implementation that provides concurrent access. +type Slice[T any] struct { + inner []T + mu sync.RWMutex +} + +// NewSlice creates a new thread-safe slice. +func NewSlice[T any]() *Slice[T] { + return &Slice[T]{ + inner: make([]T, 0), + } +} + +// NewSliceFrom creates a new thread-safe slice from an existing slice. +func NewSliceFrom[T any](s []T) *Slice[T] { + inner := make([]T, len(s)) + copy(inner, s) + return &Slice[T]{ + inner: inner, + } +} + +// Append adds an element to the end of the slice. +func (s *Slice[T]) Append(items ...T) { + s.mu.Lock() + defer s.mu.Unlock() + s.inner = append(s.inner, items...) +} + +// Get returns the element at the specified index. +func (s *Slice[T]) Get(index int) (T, bool) { + s.mu.RLock() + defer s.mu.RUnlock() + var zero T + if index < 0 || index >= len(s.inner) { + return zero, false + } + return s.inner[index], true +} + +// Len returns the number of elements in the slice. +func (s *Slice[T]) Len() int { + s.mu.RLock() + defer s.mu.RUnlock() + return len(s.inner) +} + +// SetSlice replaces the entire slice with a new one. +func (s *Slice[T]) SetSlice(items []T) { + s.mu.Lock() + defer s.mu.Unlock() + s.inner = make([]T, len(items)) + copy(s.inner, items) +} + +// Seq returns an iterator that yields elements from the slice. +func (s *Slice[T]) Seq() iter.Seq[T] { + return func(yield func(T) bool) { + for _, v := range s.Seq2() { + if !yield(v) { + return + } + } + } +} + +// Seq2 returns an iterator that yields index-value pairs from the slice. +func (s *Slice[T]) Seq2() iter.Seq2[int, T] { + items := s.Copy() + return func(yield func(int, T) bool) { + for i, v := range items { + if !yield(i, v) { + return + } + } + } +} + +// Copy returns a copy of the inner slice. +func (s *Slice[T]) Copy() []T { + s.mu.RLock() + defer s.mu.RUnlock() + items := make([]T, len(s.inner)) + copy(items, s.inner) + return items +} diff --git a/internal/csync/slices_test.go b/internal/csync/slices_test.go new file mode 100644 index 0000000000000000000000000000000000000000..c7946ac6f1a84614def05b7b6e7e9b0ed11b3a73 --- /dev/null +++ b/internal/csync/slices_test.go @@ -0,0 +1,214 @@ +package csync + +import ( + "slices" + "sync" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/stretchr/testify/require" +) + +func TestLazySlice_Seq(t *testing.T) { + t.Parallel() + + synctest.Test(t, func(t *testing.T) { + t.Helper() + data := []string{"a", "b", "c"} + s := NewLazySlice(func() []string { + time.Sleep(10 * time.Millisecond) // Small delay to ensure loading happens + return data + }) + require.Equal(t, data, slices.Collect(s.Seq())) + }) +} + +func TestLazySlice_SeqWaitsForLoading(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + t.Helper() + + var loaded atomic.Bool + data := []string{"x", "y", "z"} + + s := NewLazySlice(func() []string { + time.Sleep(100 * time.Millisecond) + loaded.Store(true) + return data + }) + + require.False(t, loaded.Load(), "should not be loaded immediately") + require.Equal(t, data, slices.Collect(s.Seq())) + require.True(t, loaded.Load(), "should be loaded after Seq") + }) +} + +func TestLazySlice_EmptySlice(t *testing.T) { + t.Parallel() + s := NewLazySlice(func() []string { + return []string{} + }) + require.Empty(t, slices.Collect(s.Seq())) +} + +func TestLazySlice_EarlyBreak(t *testing.T) { + t.Parallel() + + synctest.Test(t, func(t *testing.T) { + t.Helper() + data := []string{"a", "b", "c", "d", "e"} + s := NewLazySlice(func() []string { + time.Sleep(10 * time.Millisecond) // Small delay to ensure loading happens + return data + }) + + var result []string + for v := range s.Seq() { + result = append(result, v) + if len(result) == 2 { + break + } + } + + require.Equal(t, []string{"a", "b"}, result) + }) +} + +func TestSlice(t *testing.T) { + t.Run("NewSlice", func(t *testing.T) { + s := NewSlice[int]() + require.Equal(t, 0, s.Len()) + }) + + t.Run("NewSliceFrom", func(t *testing.T) { + original := []int{1, 2, 3} + s := NewSliceFrom(original) + require.Equal(t, 3, s.Len()) + + // Verify it's a copy, not a reference + original[0] = 999 + val, ok := s.Get(0) + require.True(t, ok) + require.Equal(t, 1, val) + }) + + t.Run("Append", func(t *testing.T) { + s := NewSlice[string]() + s.Append("hello") + s.Append("world") + + require.Equal(t, 2, s.Len()) + val, ok := s.Get(0) + require.True(t, ok) + require.Equal(t, "hello", val) + + val, ok = s.Get(1) + require.True(t, ok) + require.Equal(t, "world", val) + }) + + t.Run("Get", func(t *testing.T) { + s := NewSliceFrom([]string{"a", "b", "c"}) + + val, ok := s.Get(1) + require.True(t, ok) + require.Equal(t, "b", val) + + // Out of bounds + _, ok = s.Get(10) + require.False(t, ok) + + // Negative index + _, ok = s.Get(-1) + require.False(t, ok) + }) + + t.Run("SetSlice", func(t *testing.T) { + s := NewSlice[int]() + s.Append(1) + s.Append(2) + + newItems := []int{10, 20, 30} + s.SetSlice(newItems) + + require.Equal(t, 3, s.Len()) + require.Equal(t, newItems, slices.Collect(s.Seq())) + + // Verify it's a copy + newItems[0] = 999 + val, ok := s.Get(0) + require.True(t, ok) + require.Equal(t, 10, val) + }) + + t.Run("Slice", func(t *testing.T) { + original := []int{1, 2, 3} + s := NewSliceFrom(original) + + copied := slices.Collect(s.Seq()) + require.Equal(t, original, copied) + + // Verify it's a copy + copied[0] = 999 + val, ok := s.Get(0) + require.True(t, ok) + require.Equal(t, 1, val) + }) + + t.Run("Seq", func(t *testing.T) { + s := NewSliceFrom([]int{1, 2, 3}) + + var result []int + for v := range s.Seq() { + result = append(result, v) + } + + require.Equal(t, []int{1, 2, 3}, result) + }) + + t.Run("SeqWithIndex", func(t *testing.T) { + s := NewSliceFrom([]string{"a", "b", "c"}) + + var indices []int + var values []string + for i, v := range s.Seq2() { + indices = append(indices, i) + values = append(values, v) + } + + require.Equal(t, []int{0, 1, 2}, indices) + require.Equal(t, []string{"a", "b", "c"}, values) + }) + + t.Run("ConcurrentAccess", func(t *testing.T) { + s := NewSlice[int]() + const numGoroutines = 100 + const itemsPerGoroutine = 10 + + var wg sync.WaitGroup + + // Concurrent appends + for i := range numGoroutines { + wg.Add(2) + go func(start int) { + defer wg.Done() + for j := range itemsPerGoroutine { + s.Append(start*itemsPerGoroutine + j) + } + }(i) + go func() { + defer wg.Done() + for range itemsPerGoroutine { + s.Len() // Just read the length + } + }() + } + + wg.Wait() + + // Should have all items + require.Equal(t, numGoroutines*itemsPerGoroutine, s.Len()) + }) +} diff --git a/internal/csync/value.go b/internal/csync/value.go new file mode 100644 index 0000000000000000000000000000000000000000..17528a281e0d34d49b206a7c3901b892370c18ba --- /dev/null +++ b/internal/csync/value.go @@ -0,0 +1,44 @@ +package csync + +import ( + "reflect" + "sync" +) + +// Value is a generic thread-safe wrapper for any value type. +// +// For slices, use [Slice]. For maps, use [Map]. Pointers are not supported. +type Value[T any] struct { + v T + mu sync.RWMutex +} + +// NewValue creates a new Value with the given initial value. +// +// Panics if t is a pointer, slice, or map. Use the dedicated types for those. +func NewValue[T any](t T) *Value[T] { + v := reflect.ValueOf(t) + switch v.Kind() { + case reflect.Pointer: + panic("csync.Value does not support pointer types") + case reflect.Slice: + panic("csync.Value does not support slice types; use csync.Slice") + case reflect.Map: + panic("csync.Value does not support map types; use csync.Map") + } + return &Value[T]{v: t} +} + +// Get returns the current value. +func (v *Value[T]) Get() T { + v.mu.RLock() + defer v.mu.RUnlock() + return v.v +} + +// Set updates the value. +func (v *Value[T]) Set(t T) { + v.mu.Lock() + defer v.mu.Unlock() + v.v = t +} diff --git a/internal/csync/value_test.go b/internal/csync/value_test.go new file mode 100644 index 0000000000000000000000000000000000000000..2d0243a3b0ae8f71802469496c35b5d4a50d260b --- /dev/null +++ b/internal/csync/value_test.go @@ -0,0 +1,97 @@ +package csync + +import ( + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestValue_GetSet(t *testing.T) { + t.Parallel() + + v := NewValue(42) + require.Equal(t, 42, v.Get()) + + v.Set(100) + require.Equal(t, 100, v.Get()) +} + +func TestValue_ZeroValue(t *testing.T) { + t.Parallel() + + v := NewValue("") + require.Equal(t, "", v.Get()) + + v.Set("hello") + require.Equal(t, "hello", v.Get()) +} + +func TestValue_Struct(t *testing.T) { + t.Parallel() + + type config struct { + Name string + Count int + } + + v := NewValue(config{Name: "test", Count: 1}) + require.Equal(t, config{Name: "test", Count: 1}, v.Get()) + + v.Set(config{Name: "updated", Count: 2}) + require.Equal(t, config{Name: "updated", Count: 2}, v.Get()) +} + +func TestValue_PointerPanics(t *testing.T) { + t.Parallel() + + require.Panics(t, func() { + NewValue(&struct{}{}) + }) +} + +func TestValue_SlicePanics(t *testing.T) { + t.Parallel() + + require.Panics(t, func() { + NewValue([]string{"a", "b"}) + }) +} + +func TestValue_MapPanics(t *testing.T) { + t.Parallel() + + require.Panics(t, func() { + NewValue(map[string]int{"a": 1}) + }) +} + +func TestValue_ConcurrentAccess(t *testing.T) { + t.Parallel() + + v := NewValue(0) + var wg sync.WaitGroup + + // Concurrent writers. + for i := range 100 { + wg.Add(1) + go func(val int) { + defer wg.Done() + v.Set(val) + }(i) + } + + // Concurrent readers. + for range 100 { + wg.Go(func() { + _ = v.Get() + }) + } + + wg.Wait() + + // Value should be one of the set values (0-99). + got := v.Get() + require.GreaterOrEqual(t, got, 0) + require.Less(t, got, 100) +} diff --git a/internal/csync/versionedmap.go b/internal/csync/versionedmap.go new file mode 100644 index 0000000000000000000000000000000000000000..6ed996b2ff8d1380aa7fd22cab57342bf71e4a8f --- /dev/null +++ b/internal/csync/versionedmap.go @@ -0,0 +1,56 @@ +package csync + +import ( + "iter" + "sync/atomic" +) + +// NewVersionedMap creates a new versioned, thread-safe map. +func NewVersionedMap[K comparable, V any]() *VersionedMap[K, V] { + return &VersionedMap[K, V]{ + m: NewMap[K, V](), + } +} + +// VersionedMap is a thread-safe map that keeps track of its version. +type VersionedMap[K comparable, V any] struct { + m *Map[K, V] + v atomic.Uint64 +} + +// Get gets the value for the specified key from the map. +func (m *VersionedMap[K, V]) Get(key K) (V, bool) { + return m.m.Get(key) +} + +// Set sets the value for the specified key in the map and increments the version. +func (m *VersionedMap[K, V]) Set(key K, value V) { + m.m.Set(key, value) + m.v.Add(1) +} + +// Del deletes the specified key from the map and increments the version. +func (m *VersionedMap[K, V]) Del(key K) { + m.m.Del(key) + m.v.Add(1) +} + +// Seq2 returns an iter.Seq2 that yields key-value pairs from the map. +func (m *VersionedMap[K, V]) Seq2() iter.Seq2[K, V] { + return m.m.Seq2() +} + +// Copy returns a copy of the inner map. +func (m *VersionedMap[K, V]) Copy() map[K]V { + return m.m.Copy() +} + +// Len returns the number of items in the map. +func (m *VersionedMap[K, V]) Len() int { + return m.m.Len() +} + +// Version returns the current version of the map. +func (m *VersionedMap[K, V]) Version() uint64 { + return m.v.Load() +} diff --git a/internal/csync/versionedmap_test.go b/internal/csync/versionedmap_test.go new file mode 100644 index 0000000000000000000000000000000000000000..2c32004c5f269b7518999f95be23db95d7b6ec15 --- /dev/null +++ b/internal/csync/versionedmap_test.go @@ -0,0 +1,89 @@ +package csync + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestVersionedMap_Set(t *testing.T) { + t.Parallel() + + vm := NewVersionedMap[string, int]() + require.Equal(t, uint64(0), vm.Version()) + + vm.Set("key1", 42) + require.Equal(t, uint64(1), vm.Version()) + + value, ok := vm.Get("key1") + require.True(t, ok) + require.Equal(t, 42, value) +} + +func TestVersionedMap_Del(t *testing.T) { + t.Parallel() + + vm := NewVersionedMap[string, int]() + vm.Set("key1", 42) + initialVersion := vm.Version() + + vm.Del("key1") + require.Equal(t, initialVersion+1, vm.Version()) + + _, ok := vm.Get("key1") + require.False(t, ok) +} + +func TestVersionedMap_VersionIncrement(t *testing.T) { + t.Parallel() + + vm := NewVersionedMap[string, int]() + initialVersion := vm.Version() + + // Setting a value should increment the version + vm.Set("key1", 42) + require.Equal(t, initialVersion+1, vm.Version()) + + // Deleting a value should increment the version + vm.Del("key1") + require.Equal(t, initialVersion+2, vm.Version()) + + // Deleting a non-existent key should still increment the version + vm.Del("nonexistent") + require.Equal(t, initialVersion+3, vm.Version()) +} + +func TestVersionedMap_ConcurrentAccess(t *testing.T) { + t.Parallel() + + vm := NewVersionedMap[int, int]() + const numGoroutines = 100 + const numOperations = 100 + + // Initial version + initialVersion := vm.Version() + + // Perform concurrent Set and Del operations + for i := range numGoroutines { + go func(id int) { + for j := range numOperations { + key := id*numOperations + j + vm.Set(key, key*2) + vm.Del(key) + } + }(i) + } + + // Wait for operations to complete by checking the version + // This is a simplified check - in a real test you might want to use sync.WaitGroup + expectedMinVersion := initialVersion + uint64(numGoroutines*numOperations*2) + + // Allow some time for operations to complete + for vm.Version() < expectedMinVersion { + // Busy wait - in a real test you'd use proper synchronization + } + + // Final version should be at least the expected minimum + require.GreaterOrEqual(t, vm.Version(), expectedMinVersion) + require.Equal(t, 0, vm.Len()) +} diff --git a/internal/diff/diff.go b/internal/diff/diff.go new file mode 100644 index 0000000000000000000000000000000000000000..9e00c87177d1438c5967797a3b4e6f3e7a467dbe --- /dev/null +++ b/internal/diff/diff.go @@ -0,0 +1,29 @@ +package diff + +import ( + "strings" + + "github.com/aymanbagabas/go-udiff" +) + +// GenerateDiff creates a unified diff from two file contents +func GenerateDiff(beforeContent, afterContent, fileName string) (string, int, int) { + fileName = strings.TrimPrefix(fileName, "/") + + var ( + unified = udiff.Unified("a/"+fileName, "b/"+fileName, beforeContent, afterContent) + additions = 0 + removals = 0 + ) + + lines := strings.SplitSeq(unified, "\n") + for line := range lines { + if strings.HasPrefix(line, "+") && !strings.HasPrefix(line, "+++") { + additions++ + } else if strings.HasPrefix(line, "-") && !strings.HasPrefix(line, "---") { + removals++ + } + } + + return unified, additions, removals +} diff --git a/internal/diffdetect/detect.go b/internal/diffdetect/detect.go new file mode 100644 index 0000000000000000000000000000000000000000..213803e6b3f3754491dd3953ca82b100090def89 --- /dev/null +++ b/internal/diffdetect/detect.go @@ -0,0 +1,36 @@ +package diffdetect + +import "strings" + +// Signal describes which unified-diff markers were found while scanning text. +type Signal struct { + HasHunk bool + HasFileHeader bool + HasGitHeader bool +} + +// Inspect scans content for unified-diff markers. +func Inspect(content string) Signal { + var signal Signal + for line := range strings.SplitSeq(content, "\n") { + if strings.HasPrefix(line, "@@") { + signal.HasHunk = true + } + if strings.HasPrefix(line, "--- ") || strings.HasPrefix(line, "+++ ") { + signal.HasFileHeader = true + } + if strings.HasPrefix(line, "diff --git ") { + signal.HasGitHeader = true + } + } + return signal +} + +// IsUnifiedDiff reports whether content appears to be a unified diff. +func IsUnifiedDiff(content string) bool { + signal := Inspect(content) + if signal.HasGitHeader && signal.HasFileHeader { + return true + } + return signal.HasHunk && signal.HasFileHeader +} diff --git a/internal/diffdetect/detect_test.go b/internal/diffdetect/detect_test.go new file mode 100644 index 0000000000000000000000000000000000000000..5d2cb50a5e50ff93a19596d92639abf44e59f2bd --- /dev/null +++ b/internal/diffdetect/detect_test.go @@ -0,0 +1,152 @@ +package diffdetect + +import "testing" + +func TestInspect(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + content string + want Signal + }{ + { + name: "git unified diff", + content: `diff --git a/main.go b/main.go +--- a/main.go ++++ b/main.go +@@ -1,2 +1,3 @@ + package main ++import "fmt" +`, + want: Signal{HasHunk: true, HasFileHeader: true, HasGitHeader: true}, + }, + { + name: "non-git unified diff", + content: `--- old.c ++++ old.c +@@ -1 +1 @@ +-old ++new +`, + want: Signal{HasHunk: true, HasFileHeader: true, HasGitHeader: false}, + }, + { + name: "plain text", + content: "hello world", + want: Signal{}, + }, + { + name: "hunk only", + content: "@@ -1 +1 @@\n-old\n+new\n", + want: Signal{HasHunk: true}, + }, + { + name: "headers only", + content: "--- a/file\n+++ b/file\n", + want: Signal{HasFileHeader: true}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + got := Inspect(tt.content) + if got != tt.want { + t.Errorf("Inspect() = %+v, want %+v", got, tt.want) + } + }) + } +} + +func TestIsUnifiedDiff(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + content string + want bool + }{ + { + name: "github-style multi-file diff", + content: `diff --git a/one.txt b/one.txt +--- a/one.txt ++++ b/one.txt +@@ -1 +1 @@ +-a ++b +diff --git a/two.txt b/two.txt +--- a/two.txt ++++ b/two.txt +@@ -1 +1 @@ +-c ++d +`, + want: true, + }, + { + name: "non-git unified patch", + content: `--- a/old.c ++++ b/old.c +@@ -1,3 +1,4 @@ + #include +-int main() { ++int main(int argc, char **argv) { + return 0; + } +`, + want: true, + }, + { + name: "new file from dev null", + content: `--- /dev/null ++++ newfile.txt +@@ -0,0 +1,2 @@ ++hello ++world +`, + want: true, + }, + { + name: "markdown false positive candidate", + content: `- Item one +- Item two ++ Bonus item +- Item three +`, + want: false, + }, + { + name: "headers without hunk", + content: `--- a/somefile.txt ++++ b/somefile.txt +Just some content here +No hunk markers at all +`, + want: false, + }, + { + name: "hunk without headers", + content: `@@ -1,3 +1,4 @@ + some line +-another line ++changed line +`, + want: false, + }, + { + name: "empty", + content: "", + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := IsUnifiedDiff(tt.content); got != tt.want { + t.Errorf("IsUnifiedDiff() = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/internal/filetracker/service.go b/internal/filetracker/service.go new file mode 100644 index 0000000000000000000000000000000000000000..5a92d4de1d0c2ac585c25f7f31834b94564a0f5d --- /dev/null +++ b/internal/filetracker/service.go @@ -0,0 +1,93 @@ +// Package filetracker provides functionality to track file reads in sessions. +package filetracker + +import ( + "context" + "fmt" + "log/slog" + "os" + "path/filepath" + "time" + + "github.com/charmbracelet/crush/internal/db" +) + +// Service defines the interface for tracking file reads in sessions. +type Service interface { + // RecordRead records when a file was read. + RecordRead(ctx context.Context, sessionID, path string) + + // LastReadTime returns when a file was last read. + // Returns zero time if never read. + LastReadTime(ctx context.Context, sessionID, path string) time.Time + + // ListReadFiles returns the paths of all files read in a session. + ListReadFiles(ctx context.Context, sessionID string) ([]string, error) +} + +type service struct { + q *db.Queries +} + +// NewService creates a new file tracker service. +func NewService(q *db.Queries) Service { + return &service{q: q} +} + +// RecordRead records when a file was read. +func (s *service) RecordRead(ctx context.Context, sessionID, path string) { + if err := s.q.RecordFileRead(ctx, db.RecordFileReadParams{ + SessionID: sessionID, + Path: relpath(path), + }); err != nil { + slog.Error("Error recording file read", "error", err, "file", path) + } +} + +// LastReadTime returns when a file was last read. +// Returns zero time if never read. +func (s *service) LastReadTime(ctx context.Context, sessionID, path string) time.Time { + readFile, err := s.q.GetFileRead(ctx, db.GetFileReadParams{ + SessionID: sessionID, + Path: relpath(path), + }) + if err != nil { + return time.Time{} + } + + return time.Unix(readFile.ReadAt, 0) +} + +func relpath(path string) string { + path = filepath.Clean(path) + basepath, err := os.Getwd() + if err != nil { + slog.Warn("Error getting basepath", "error", err) + return path + } + relpath, err := filepath.Rel(basepath, path) + if err != nil { + slog.Warn("Error getting relpath", "error", err) + return path + } + return relpath +} + +// ListReadFiles returns the paths of all files read in a session. +func (s *service) ListReadFiles(ctx context.Context, sessionID string) ([]string, error) { + readFiles, err := s.q.ListSessionReadFiles(ctx, sessionID) + if err != nil { + return nil, fmt.Errorf("listing read files: %w", err) + } + + basepath, err := os.Getwd() + if err != nil { + return nil, fmt.Errorf("getting working directory: %w", err) + } + + paths := make([]string, 0, len(readFiles)) + for _, rf := range readFiles { + paths = append(paths, filepath.Join(basepath, rf.Path)) + } + return paths, nil +} diff --git a/internal/filetracker/service_test.go b/internal/filetracker/service_test.go new file mode 100644 index 0000000000000000000000000000000000000000..c7fb15090dd31e9591c5c3b9c2a256c839aea3f6 --- /dev/null +++ b/internal/filetracker/service_test.go @@ -0,0 +1,116 @@ +package filetracker + +import ( + "context" + "testing" + "testing/synctest" + "time" + + "github.com/charmbracelet/crush/internal/db" + "github.com/stretchr/testify/require" +) + +type testEnv struct { + ctx context.Context + q *db.Queries + svc Service +} + +func setupTest(t *testing.T) *testEnv { + t.Helper() + + conn, err := db.Connect(t.Context(), t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { conn.Close() }) + + q := db.New(conn) + return &testEnv{ + ctx: t.Context(), + q: q, + svc: NewService(q), + } +} + +func (e *testEnv) createSession(t *testing.T, sessionID string) { + t.Helper() + _, err := e.q.CreateSession(e.ctx, db.CreateSessionParams{ + ID: sessionID, + Title: "Test Session", + }) + require.NoError(t, err) +} + +func TestService_RecordRead(t *testing.T) { + env := setupTest(t) + + sessionID := "test-session-1" + path := "/path/to/file.go" + env.createSession(t, sessionID) + + env.svc.RecordRead(env.ctx, sessionID, path) + + lastRead := env.svc.LastReadTime(env.ctx, sessionID, path) + require.False(t, lastRead.IsZero(), "expected non-zero time after recording read") + require.WithinDuration(t, time.Now(), lastRead, 2*time.Second) +} + +func TestService_LastReadTime_NotFound(t *testing.T) { + env := setupTest(t) + + lastRead := env.svc.LastReadTime(env.ctx, "nonexistent-session", "/nonexistent/path") + require.True(t, lastRead.IsZero(), "expected zero time for unread file") +} + +func TestService_RecordRead_UpdatesTimestamp(t *testing.T) { + env := setupTest(t) + + sessionID := "test-session-2" + path := "/path/to/file.go" + env.createSession(t, sessionID) + + env.svc.RecordRead(env.ctx, sessionID, path) + firstRead := env.svc.LastReadTime(env.ctx, sessionID, path) + require.False(t, firstRead.IsZero()) + + synctest.Test(t, func(t *testing.T) { + time.Sleep(100 * time.Millisecond) + synctest.Wait() + env.svc.RecordRead(env.ctx, sessionID, path) + secondRead := env.svc.LastReadTime(env.ctx, sessionID, path) + + require.False(t, secondRead.Before(firstRead), "second read time should not be before first") + }) +} + +func TestService_RecordRead_DifferentSessions(t *testing.T) { + env := setupTest(t) + + path := "/shared/file.go" + session1, session2 := "session-1", "session-2" + env.createSession(t, session1) + env.createSession(t, session2) + + env.svc.RecordRead(env.ctx, session1, path) + + lastRead1 := env.svc.LastReadTime(env.ctx, session1, path) + require.False(t, lastRead1.IsZero()) + + lastRead2 := env.svc.LastReadTime(env.ctx, session2, path) + require.True(t, lastRead2.IsZero(), "session 2 should not see session 1's read") +} + +func TestService_RecordRead_DifferentPaths(t *testing.T) { + env := setupTest(t) + + sessionID := "test-session-3" + path1, path2 := "/path/to/file1.go", "/path/to/file2.go" + env.createSession(t, sessionID) + + env.svc.RecordRead(env.ctx, sessionID, path1) + + lastRead1 := env.svc.LastReadTime(env.ctx, sessionID, path1) + require.False(t, lastRead1.IsZero()) + + lastRead2 := env.svc.LastReadTime(env.ctx, sessionID, path2) + require.True(t, lastRead2.IsZero(), "path2 should not be recorded") +} diff --git a/internal/lock/lock.go b/internal/lock/lock.go new file mode 100644 index 0000000000000000000000000000000000000000..dbbbec6a2acb1a01d74d7e4e9019c0c7627b4114 --- /dev/null +++ b/internal/lock/lock.go @@ -0,0 +1,75 @@ +// Package lock provides cross-process advisory file locking. +// +// File acquires an exclusive lock on the file at path, blocking until +// the context is cancelled (or its deadline elapses). TryFile does the +// same but returns ErrContended immediately if the lock is already +// held. In both cases the returned release function drops the lock and +// closes the underlying file descriptor. +// +// The lock is released automatically by the kernel on process +// termination (including crash), so no stale-lock recovery is needed. +// +// The lock file at path is created if it does not exist. It is never +// unlinked — flock is keyed by inode, not path, and unlinking could +// create a window where two processes lock different inodes at the +// same path. +// +// This is the canonical file-locking helper for Crush. Callers should +// prefer it over rolling their own platform-specific code. +package lock + +import ( + "context" + "errors" + "fmt" + "os" +) + +// ErrContended is returned by TryFile when the lock is already held by +// another process. +var ErrContended = errors.New("file lock is held by another process") + +// File acquires an exclusive advisory lock on the file at path, blocking +// until the lock is acquired or ctx is cancelled. It returns a release +// function that drops the lock and closes the underlying file descriptor. +// +// Pass a context with a deadline (e.g. context.WithTimeout) to bound the +// wait. Pass context.Background() to block indefinitely. +func File(ctx context.Context, path string) (func(), error) { + f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0o600) + if err != nil { + return nil, fmt.Errorf("open lock file %q: %w", path, err) + } + + release, err := lockFile(ctx, f) + if err != nil { + f.Close() + return nil, err + } + + return func() { + release() + f.Close() + }, nil +} + +// TryFile is like File but returns ErrContended immediately if the lock +// is already held by another process. Use this when you want to fail +// fast rather than wait. +func TryFile(path string) (func(), error) { + f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0o600) + if err != nil { + return nil, fmt.Errorf("open lock file %q: %w", path, err) + } + + release, err := tryLockFile(f) + if err != nil { + f.Close() + return nil, err + } + + return func() { + release() + f.Close() + }, nil +} diff --git a/internal/lock/lock_test.go b/internal/lock/lock_test.go new file mode 100644 index 0000000000000000000000000000000000000000..6d7b8de56cd7d69c365cb4cee6746135df9e8dfd --- /dev/null +++ b/internal/lock/lock_test.go @@ -0,0 +1,168 @@ +package lock + +import ( + "context" + "errors" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestTryFile_AcquiresWhenFree(t *testing.T) { + t.Parallel() + path := filepath.Join(t.TempDir(), "test.lock") + + release, err := TryFile(path) + require.NoError(t, err) + require.NotNil(t, release) + release() +} + +func TestTryFile_ReturnsErrContendedWhenHeld(t *testing.T) { + t.Parallel() + path := filepath.Join(t.TempDir(), "test.lock") + + release, err := TryFile(path) + require.NoError(t, err) + t.Cleanup(release) + + _, err = TryFile(path) + require.ErrorIs(t, err, ErrContended) +} + +func TestTryFile_ReacquireAfterRelease(t *testing.T) { + t.Parallel() + path := filepath.Join(t.TempDir(), "test.lock") + + release, err := TryFile(path) + require.NoError(t, err) + release() + + release2, err := TryFile(path) + require.NoError(t, err) + t.Cleanup(release2) +} + +func TestFile_AcquiresWhenFree(t *testing.T) { + t.Parallel() + path := filepath.Join(t.TempDir(), "test.lock") + + release, err := File(context.Background(), path) + require.NoError(t, err) + t.Cleanup(release) +} + +func TestFile_BlocksThenSucceeds(t *testing.T) { + t.Parallel() + path := filepath.Join(t.TempDir(), "test.lock") + + release, err := TryFile(path) + require.NoError(t, err) + + // Release the lock after a short delay so the blocking acquirer + // can complete within the test timeout. + go func() { + time.Sleep(150 * time.Millisecond) + release() + }() + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + release2, err := File(ctx, path) + require.NoError(t, err, "should acquire after first releases") + release2() +} + +func TestFile_RespectsContextDeadline(t *testing.T) { + t.Parallel() + path := filepath.Join(t.TempDir(), "test.lock") + + release, err := TryFile(path) + require.NoError(t, err) + t.Cleanup(release) + + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + start := time.Now() + _, err = File(ctx, path) + elapsed := time.Since(start) + require.Error(t, err) + require.True(t, errors.Is(err, context.DeadlineExceeded), "expected deadline exceeded, got %v", err) + require.Less(t, elapsed, 1*time.Second, "should return promptly after deadline") +} + +func TestFile_RespectsContextCancellation(t *testing.T) { + t.Parallel() + path := filepath.Join(t.TempDir(), "test.lock") + + release, err := TryFile(path) + require.NoError(t, err) + t.Cleanup(release) + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + _, err := File(ctx, path) + done <- err + }() + + time.Sleep(50 * time.Millisecond) + cancel() + + select { + case err := <-done: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(2 * time.Second): + t.Fatal("File did not return after context cancellation") + } +} + +// TestFile_ConcurrentAcquirers verifies that multiple blocking acquirers +// queue up correctly: each gets the lock in turn, exactly one at a time. +func TestFile_ConcurrentAcquirers(t *testing.T) { + t.Parallel() + path := filepath.Join(t.TempDir(), "test.lock") + + const n = 5 + var ( + mu sync.Mutex + inside int + maxSeen int + finished int + ) + var wg sync.WaitGroup + wg.Add(n) + for range n { + go func() { + defer wg.Done() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + release, err := File(ctx, path) + if err != nil { + t.Errorf("acquire failed: %v", err) + return + } + mu.Lock() + inside++ + if inside > maxSeen { + maxSeen = inside + } + mu.Unlock() + + time.Sleep(20 * time.Millisecond) + + mu.Lock() + inside-- + finished++ + mu.Unlock() + release() + }() + } + wg.Wait() + + require.Equal(t, n, finished) + require.Equal(t, 1, maxSeen, "lock must be mutually exclusive") +} diff --git a/internal/lock/lock_unix.go b/internal/lock/lock_unix.go new file mode 100644 index 0000000000000000000000000000000000000000..a60ee9b234870bb397530f08396a905e38967cc7 --- /dev/null +++ b/internal/lock/lock_unix.go @@ -0,0 +1,45 @@ +//go:build !windows + +package lock + +import ( + "context" + "errors" + "fmt" + "os" + "time" + + "golang.org/x/sys/unix" +) + +// retrySleep is the interval between non-blocking flock retries in the +// blocking File path. Small enough that contention resolution feels +// snappy; large enough that we don't burn a CPU spinning. +const retrySleep = 100 * time.Millisecond + +func lockFile(ctx context.Context, f *os.File) (func(), error) { + for { + err := unix.Flock(int(f.Fd()), unix.LOCK_EX|unix.LOCK_NB) + if err == nil { + return func() { _ = unix.Flock(int(f.Fd()), unix.LOCK_UN) }, nil + } + if !errors.Is(err, unix.EWOULDBLOCK) { + return nil, fmt.Errorf("flock: %w", err) + } + select { + case <-ctx.Done(): + return nil, fmt.Errorf("acquire lock: %w", ctx.Err()) + case <-time.After(retrySleep): + } + } +} + +func tryLockFile(f *os.File) (func(), error) { + if err := unix.Flock(int(f.Fd()), unix.LOCK_EX|unix.LOCK_NB); err != nil { + if errors.Is(err, unix.EWOULDBLOCK) { + return nil, ErrContended + } + return nil, fmt.Errorf("flock: %w", err) + } + return func() { _ = unix.Flock(int(f.Fd()), unix.LOCK_UN) }, nil +} diff --git a/internal/lock/lock_windows.go b/internal/lock/lock_windows.go new file mode 100644 index 0000000000000000000000000000000000000000..2b7845966ff8e7503bfd06e9dd1bd60f99fbbd99 --- /dev/null +++ b/internal/lock/lock_windows.go @@ -0,0 +1,57 @@ +//go:build windows + +package lock + +import ( + "context" + "errors" + "fmt" + "math" + "os" + "time" + + "golang.org/x/sys/windows" +) + +// retrySleep is the interval between non-blocking lock retries in the +// blocking File path. +const retrySleep = 100 * time.Millisecond + +func lockFile(ctx context.Context, f *os.File) (func(), error) { + h := windows.Handle(f.Fd()) + for { + ol := new(windows.Overlapped) + flags := uint32(windows.LOCKFILE_EXCLUSIVE_LOCK | windows.LOCKFILE_FAIL_IMMEDIATELY) + err := windows.LockFileEx(h, flags, 0, math.MaxUint32, math.MaxUint32, ol) + if err == nil { + return func() { + ol := new(windows.Overlapped) + _ = windows.UnlockFileEx(windows.Handle(f.Fd()), 0, math.MaxUint32, math.MaxUint32, ol) + }, nil + } + if !errors.Is(err, windows.ERROR_LOCK_VIOLATION) && !errors.Is(err, windows.ERROR_IO_PENDING) { + return nil, fmt.Errorf("LockFileEx: %w", err) + } + select { + case <-ctx.Done(): + return nil, fmt.Errorf("acquire lock: %w", ctx.Err()) + case <-time.After(retrySleep): + } + } +} + +func tryLockFile(f *os.File) (func(), error) { + h := windows.Handle(f.Fd()) + ol := new(windows.Overlapped) + flags := uint32(windows.LOCKFILE_EXCLUSIVE_LOCK | windows.LOCKFILE_FAIL_IMMEDIATELY) + if err := windows.LockFileEx(h, flags, 0, math.MaxUint32, math.MaxUint32, ol); err != nil { + if errors.Is(err, windows.ERROR_LOCK_VIOLATION) || errors.Is(err, windows.ERROR_IO_PENDING) { + return nil, ErrContended + } + return nil, fmt.Errorf("LockFileEx: %w", err) + } + return func() { + ol := new(windows.Overlapped) + _ = windows.UnlockFileEx(windows.Handle(f.Fd()), 0, math.MaxUint32, math.MaxUint32, ol) + }, nil +} diff --git a/internal/logout/choose.go b/internal/logout/choose.go new file mode 100644 index 0000000000000000000000000000000000000000..a3515b102e4bfa6bf2c7e3cdf5c77b46617c931b --- /dev/null +++ b/internal/logout/choose.go @@ -0,0 +1,160 @@ +package logout + +import ( + "errors" + "strings" + + "charm.land/bubbles/v2/help" + "charm.land/bubbles/v2/key" + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" +) + +type chooseKeyMap struct { + Up key.Binding + Down key.Binding + Confirm key.Binding + Cancel key.Binding +} + +func defaultChooseKeyMap() chooseKeyMap { + return chooseKeyMap{ + Up: key.NewBinding(key.WithKeys("up", "k"), key.WithHelp("↑/↓", "select")), + Down: key.NewBinding(key.WithKeys("down", "j"), key.WithHelp("", "")), + Confirm: key.NewBinding(key.WithKeys("enter"), key.WithHelp("enter", "confirm")), + Cancel: key.NewBinding(key.WithKeys("ctrl+c", "esc"), key.WithHelp("esc", "cancel")), + } +} + +func (k chooseKeyMap) ShortHelp() []key.Binding { + return []key.Binding{k.Up, k.Confirm, k.Cancel} +} + +func (k chooseKeyMap) FullHelp() [][]key.Binding { + return [][]key.Binding{{k.Up, k.Down, k.Confirm, k.Cancel}} +} + +// chooseModel is the Bubble Tea model for a multiple choice question with +// a ">" cursor. +type chooseModel struct { + question string + options []string + cursor int + chosen int + + help help.Model + keymap chooseKeyMap + + quitting bool + blankLines int + width int +} + +func newChooseModel(question string, options []string) chooseModel { + return chooseModel{ + question: question, + options: options, + chosen: -1, + help: help.New(), + keymap: defaultChooseKeyMap(), + } +} + +// runChoose runs the multiple choice TUI and returns the selected option +// index, or -1 when canceled. +func runChoose(question string, options []string) (int, error) { + p := tea.NewProgram(newChooseModel(question, options)) + final, err := p.Run() + if err != nil { + return -1, err + } + m, ok := final.(chooseModel) + if !ok { + return -1, errors.New("unexpected choose program model") + } + return m.chosen, nil +} + +// Init initializes the choose TUI. +func (m chooseModel) Init() tea.Cmd { + return nil +} + +// Update handles messages for the choose TUI. +func (m chooseModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.WindowSizeMsg: + m.width = msg.Width + m.help.SetWidth(msg.Width) + return m, nil + + case tea.KeyPressMsg: + switch { + case key.Matches(msg, m.keymap.Cancel): + m.quitting = true + m.blankLines = m.contentHeight() + return m, tea.Quit + case key.Matches(msg, m.keymap.Confirm): + m.chosen = m.cursor + m.quitting = true + m.blankLines = m.contentHeight() + return m, tea.Quit + case key.Matches(msg, m.keymap.Up): + m.cursor-- + if m.cursor < 0 { + m.cursor = len(m.options) - 1 + } + return m, nil + case key.Matches(msg, m.keymap.Down): + m.cursor++ + if m.cursor >= len(m.options) { + m.cursor = 0 + } + return m, nil + } + } + return m, nil +} + +// content renders the full choose view. The trailing zero-width space +// keeps the final blank line alive (the renderer erases trailing empty +// lines). +func (m chooseModel) content() string { + var b strings.Builder + + b.WriteString("\n ") + b.WriteString(questionStyle.Render(m.question)) + b.WriteString("\n\n") + for i, option := range m.options { + if i == m.cursor { + b.WriteString(" " + choiceSelectedStyle.Render("> "+option)) + } else { + b.WriteString(" " + option) + } + if i < len(m.options)-1 { + b.WriteString("\n") + } + } + b.WriteString("\n\n ") + b.WriteString(m.help.View(m.keymap)) + b.WriteString("\n") + return b.String() +} + +// contentHeight returns the height of the current view in lines. +func (m chooseModel) contentHeight() int { + return lipgloss.Height(m.content()) +} + +// View renders the choose TUI. +func (m chooseModel) View() tea.View { + if m.quitting { + // Render the same number of blank lines as the last frame so the + // renderer overwrites it and the TUI disappears on exit. A truly + // empty view doesn't work under Fang's wrapped output: the renderer + // drops the touched-lines buffer when the height shrinks to zero + // and leaves the frame on screen. + return tea.NewView("") + } + return tea.NewView(m.content()) +} diff --git a/internal/logout/confirm.go b/internal/logout/confirm.go new file mode 100644 index 0000000000000000000000000000000000000000..19b51a1b563c9ce5f7a8a36d79023d52b237b441 --- /dev/null +++ b/internal/logout/confirm.go @@ -0,0 +1,148 @@ +package logout + +import ( + "errors" + "strings" + + "charm.land/bubbles/v2/help" + "charm.land/bubbles/v2/key" + tea "charm.land/bubbletea/v2" + "charm.land/lipgloss/v2" +) + +type confirmKeyMap struct { + Switch key.Binding + Confirm key.Binding + Yes key.Binding + No key.Binding + Cancel key.Binding +} + +func defaultConfirmKeyMap() confirmKeyMap { + return confirmKeyMap{ + Switch: key.NewBinding(key.WithKeys("left", "right", "h", "l"), key.WithHelp("←/→", "switch")), + Confirm: key.NewBinding(key.WithKeys("enter"), key.WithHelp("enter", "confirm")), + Yes: key.NewBinding(key.WithKeys("y", "Y"), key.WithHelp("y", "yes")), + No: key.NewBinding(key.WithKeys("n", "N"), key.WithHelp("n", "no")), + Cancel: key.NewBinding(key.WithKeys("ctrl+c", "esc"), key.WithHelp("esc", "cancel")), + } +} + +func (k confirmKeyMap) ShortHelp() []key.Binding { + return []key.Binding{k.Switch, k.Confirm, k.Yes, k.No, k.Cancel} +} + +func (k confirmKeyMap) FullHelp() [][]key.Binding { + return [][]key.Binding{{k.Switch, k.Confirm, k.Yes, k.No, k.Cancel}} +} + +// confirmModel is the Bubble Tea model for a yes/no question rendered as +// selectable buttons. No is selected by default for safety. +type confirmModel struct { + question string + answered bool + yes bool + quitting bool + blankLines int + + help help.Model + keymap confirmKeyMap + width int +} + +func newConfirmModel(question string) confirmModel { + return confirmModel{ + question: question, + help: help.New(), + keymap: defaultConfirmKeyMap(), + } +} + +// runConfirm runs the yes/no TUI and reports the user's answer. +func runConfirm(question string) (bool, error) { + p := tea.NewProgram(newConfirmModel(question)) + final, err := p.Run() + if err != nil { + return false, err + } + m, ok := final.(confirmModel) + if !ok { + return false, errors.New("unexpected confirm program model") + } + return m.answered && m.yes, nil +} + +// Init initializes the confirm TUI. +func (m confirmModel) Init() tea.Cmd { + return nil +} + +// Update handles messages for the confirm TUI. +func (m confirmModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.WindowSizeMsg: + m.width = msg.Width + m.help.SetWidth(msg.Width) + return m, nil + + case tea.KeyPressMsg: + switch { + case key.Matches(msg, m.keymap.Cancel): + m.quitting = true + m.blankLines = m.contentHeight() + return m, tea.Quit + case key.Matches(msg, m.keymap.Yes): + m.answered = true + m.yes = true + m.quitting = true + m.blankLines = m.contentHeight() + return m, tea.Quit + case key.Matches(msg, m.keymap.No): + m.answered = true + m.yes = false + m.quitting = true + m.blankLines = m.contentHeight() + return m, tea.Quit + case key.Matches(msg, m.keymap.Confirm): + m.answered = true + m.quitting = true + m.blankLines = m.contentHeight() + return m, tea.Quit + case key.Matches(msg, m.keymap.Switch): + m.yes = !m.yes + return m, nil + } + } + return m, nil +} + +// content renders the full confirm view. The trailing zero-width space +// keeps the final blank line alive (the renderer erases trailing empty +// lines). +func (m confirmModel) content() string { + var b strings.Builder + + b.WriteString("\n ") + b.WriteString(questionStyle.Render(m.question)) + b.WriteString("\n\n ") + b.WriteString(renderButton("Yes", m.yes)) + b.WriteString(" ") + b.WriteString(renderButton("No", !m.yes)) + b.WriteString("\n\n ") + b.WriteString(m.help.View(m.keymap)) + b.WriteString("\n") + return b.String() +} + +// contentHeight returns the height of the current view in lines. +func (m confirmModel) contentHeight() int { + return lipgloss.Height(m.content()) +} + +// View renders the confirm TUI. +func (m confirmModel) View() tea.View { + if m.quitting { + return tea.NewView("") + } + return tea.NewView(m.content()) +} diff --git a/internal/logout/logout.go b/internal/logout/logout.go new file mode 100644 index 0000000000000000000000000000000000000000..1f07919a0c65b6bf9a759a74eb9ff78aba66b754 --- /dev/null +++ b/internal/logout/logout.go @@ -0,0 +1,61 @@ +// Package logout implements the CLI logout prompts: small Bubble Tea +// programs in interactive terminals and plain text prompts in +// non-interactive sessions. +package logout + +import ( + "fmt" + "os" + + "github.com/charmbracelet/x/term" +) + +// interactive reports whether the prompts should run their TUIs: both stdin +// and stdout must be terminals, otherwise the plain text fallbacks are used. +func interactive() bool { + return term.IsTerminal(os.Stdin.Fd()) && term.IsTerminal(os.Stdout.Fd()) +} + +// Confirm asks a yes/no question and reports the answer. In interactive +// terminals it shows selectable Yes/No buttons (No is the default); +// otherwise it falls back to a plain (y/N) prompt. +func Confirm(question string) (bool, error) { + if interactive() { + return runConfirm(question) + } + + fmt.Printf("%s (y/N) ", question) + var response string + _, err := fmt.Scanln(&response) + if err != nil { + return false, nil + } + switch response { + case "y", "Y", "yes", "Yes", "YES": + return true, nil + default: + return false, nil + } +} + +// Choose asks a multiple choice question and returns the selected option +// index, or -1 when canceled. In interactive terminals the options are +// navigable with a ">" cursor; otherwise a numbered list is printed. +func Choose(question string, options []string) (int, error) { + if interactive() { + return runChoose(question, options) + } + + fmt.Println(question) + for i, option := range options { + fmt.Printf(" %d. %s\n", i+1, option) + } + fmt.Printf("Select (1-%d): ", len(options)) + + var choice int + _, err := fmt.Scanln(&choice) + if err != nil || choice < 1 || choice > len(options) { + return -1, nil + } + return choice - 1, nil +} diff --git a/internal/logout/styles.go b/internal/logout/styles.go new file mode 100644 index 0000000000000000000000000000000000000000..15b341fd946808b1ad367b291d30ec1d23d5f811 --- /dev/null +++ b/internal/logout/styles.go @@ -0,0 +1,42 @@ +package logout + +import ( + "charm.land/lipgloss/v2" + "github.com/charmbracelet/x/exp/charmtone" +) + +// buttonPadding is the horizontal padding inside a button. +const buttonPadding = 5 + +var ( + questionStyle = lipgloss.NewStyle(). + Bold(true) + + buttonFocusedStyle = lipgloss.NewStyle(). + Foreground(charmtone.Sash). + Background(charmtone.Blush) + buttonBlurredStyle = lipgloss.NewStyle(). + Foreground(charmtone.Soda). + Background(charmtone.BBQ) + + choiceSelectedStyle = lipgloss.NewStyle(). + Bold(true). + Foreground(charmtone.Blush) +) + +// renderButton renders a selectable button with an underlined accelerator +// key (its first letter). +func renderButton(text string, selected bool) string { + base := buttonBlurredStyle + if selected { + base = buttonFocusedStyle + } + rendered := base.Padding(0, buttonPadding).Render(text) + if text != "" { + // Underline the accelerator key (the first letter). The range style + // must not carry padding, as StyleRanges re-renders the matched text. + rendered = lipgloss.StyleRanges(rendered, + lipgloss.NewRange(buttonPadding, buttonPadding+1, base.Underline(true))) + } + return rendered +} diff --git a/internal/message/attachment.go b/internal/message/attachment.go new file mode 100644 index 0000000000000000000000000000000000000000..0ac83707212e11133925306ff93a4a8504a7ddec --- /dev/null +++ b/internal/message/attachment.go @@ -0,0 +1,24 @@ +package message + +import ( + "slices" + "strings" +) + +type Attachment struct { + FilePath string + FileName string + MimeType string + Content []byte +} + +func (a Attachment) IsText() bool { return strings.HasPrefix(a.MimeType, "text/") } +func (a Attachment) IsImage() bool { return strings.HasPrefix(a.MimeType, "image/") } +func (a Attachment) IsMarkdown() bool { return a.MimeType == "text/markdown" } + +// ContainsTextAttachment returns true if any of the attachments is a text attachment. +func ContainsTextAttachment(attachments []Attachment) bool { + return slices.ContainsFunc(attachments, func(a Attachment) bool { + return a.IsText() + }) +} diff --git a/internal/message/content.go b/internal/message/content.go new file mode 100644 index 0000000000000000000000000000000000000000..ac82b16d1354c7ab9540cd175ebec44a69a7378c --- /dev/null +++ b/internal/message/content.go @@ -0,0 +1,660 @@ +package message + +import ( + "encoding/base64" + "errors" + "fmt" + "slices" + "strings" + "time" + + "charm.land/catwalk/pkg/catwalk" + "charm.land/fantasy" + "charm.land/fantasy/providers/anthropic" + "charm.land/fantasy/providers/google" + "charm.land/fantasy/providers/openai" + "github.com/charmbracelet/crush/internal/stringext" + "github.com/charmbracelet/x/ansi" +) + +type MessageRole string + +const ( + Assistant MessageRole = "assistant" + User MessageRole = "user" + System MessageRole = "system" + Tool MessageRole = "tool" +) + +// mediaLoadFailedPlaceholder is the text substituted for image data that +// cannot be decoded during session replay. +const mediaLoadFailedPlaceholder = "[Image data could not be loaded]" + +type FinishReason string + +const ( + FinishReasonEndTurn FinishReason = "end_turn" + FinishReasonMaxTokens FinishReason = "max_tokens" + FinishReasonToolUse FinishReason = "tool_use" + FinishReasonCanceled FinishReason = "canceled" + FinishReasonError FinishReason = "error" + // FinishReasonContentFilter is a provider safety/refusal stop + // (Anthropic stop_reason=refusal, OpenAI content_filter, etc.). + // The TUI renders this as a REFUSED banner rather than a silent + // empty turn. + FinishReasonContentFilter FinishReason = "content_filter" + + // Should never happen + FinishReasonUnknown FinishReason = "unknown" +) + +type ContentPart interface { + isPart() +} + +type ReasoningContent struct { + Thinking string `json:"thinking"` + Signature string `json:"signature"` + ThoughtSignature string `json:"thought_signature"` // Used for google + ToolID string `json:"tool_id"` // Used for openrouter google models + ResponsesData *openai.ResponsesReasoningMetadata `json:"responses_data"` + StartedAt int64 `json:"started_at,omitempty"` + FinishedAt int64 `json:"finished_at,omitempty"` +} + +func (tc ReasoningContent) String() string { + return tc.Thinking +} +func (ReasoningContent) isPart() {} + +type TextContent struct { + Text string `json:"text"` +} + +func (tc TextContent) String() string { + return tc.Text +} + +func (TextContent) isPart() {} + +type ImageURLContent struct { + URL string `json:"url"` + Detail string `json:"detail,omitempty"` +} + +func (iuc ImageURLContent) String() string { + return iuc.URL +} + +func (ImageURLContent) isPart() {} + +type BinaryContent struct { + Path string + MIMEType string + Data []byte +} + +func (bc BinaryContent) String(p catwalk.InferenceProvider) string { + base64Encoded := base64.StdEncoding.EncodeToString(bc.Data) + if p == catwalk.InferenceProviderOpenAI { + return "data:" + bc.MIMEType + ";base64," + base64Encoded + } + return base64Encoded +} + +func (BinaryContent) isPart() {} + +type ToolCall struct { + ID string `json:"id"` + Name string `json:"name"` + Input string `json:"input"` + ProviderExecuted bool `json:"provider_executed"` + Finished bool `json:"finished"` +} + +func (ToolCall) isPart() {} + +type ToolResult struct { + ToolCallID string `json:"tool_call_id"` + Name string `json:"name"` + Content string `json:"content"` + Data string `json:"data"` + MIMEType string `json:"mime_type"` + Metadata string `json:"metadata"` + IsError bool `json:"is_error"` +} + +func (ToolResult) isPart() {} + +type Finish struct { + Reason FinishReason `json:"reason"` + Time int64 `json:"time"` + Message string `json:"message,omitempty"` + Details string `json:"details,omitempty"` +} + +func (Finish) isPart() {} + +// ShellCommand stores a bang-mode shell command and its output as a +// distinct content part so it can be reconstructed on session restore. +type ShellCommand struct { + Command string `json:"command"` + Output string `json:"output"` + ExitCode int `json:"exit_code"` +} + +func (ShellCommand) isPart() {} + +// HasShellCommand reports whether the message contains any ShellCommand parts. +func (m *Message) HasShellCommand() bool { + for _, part := range m.Parts { + if _, ok := part.(ShellCommand); ok { + return true + } + } + return false +} + +// ShellCommands returns all ShellCommand parts from the message. +func (m *Message) ShellCommands() []ShellCommand { + var cmds []ShellCommand + for _, part := range m.Parts { + if sc, ok := part.(ShellCommand); ok { + cmds = append(cmds, sc) + } + } + return cmds +} + +type Message struct { + ID string + Role MessageRole + SessionID string + Parts []ContentPart + Model string + Provider string + CreatedAt int64 + UpdatedAt int64 + IsSummaryMessage bool + // PrismModelID and PrismModelName identify the model that actually + // served the turn, as reported by the Hyper Prism model router + // headers. Empty when the turn was not routed through Prism. + PrismModelID string + PrismModelName string + // PrismHypercreditSavings and PrismDollarSavings are the savings + // from routing through Prism, as reported by its savings trailers. + // Nil when not reported. When both are present the hypercredit + // figure is the one shown. + PrismHypercreditSavings *float64 + PrismDollarSavings *float64 +} + +func (m *Message) Content() TextContent { + for _, part := range m.Parts { + if c, ok := part.(TextContent); ok { + return c + } + } + return TextContent{} +} + +func (m *Message) ReasoningContent() ReasoningContent { + for _, part := range m.Parts { + if c, ok := part.(ReasoningContent); ok { + return c + } + } + return ReasoningContent{} +} + +func (m *Message) ImageURLContent() []ImageURLContent { + imageURLContents := make([]ImageURLContent, 0) + for _, part := range m.Parts { + if c, ok := part.(ImageURLContent); ok { + imageURLContents = append(imageURLContents, c) + } + } + return imageURLContents +} + +func (m *Message) BinaryContent() []BinaryContent { + binaryContents := make([]BinaryContent, 0) + for _, part := range m.Parts { + if c, ok := part.(BinaryContent); ok { + binaryContents = append(binaryContents, c) + } + } + return binaryContents +} + +func (m *Message) ToolCalls() []ToolCall { + toolCalls := make([]ToolCall, 0) + for _, part := range m.Parts { + if c, ok := part.(ToolCall); ok { + toolCalls = append(toolCalls, c) + } + } + return toolCalls +} + +func (m *Message) ToolResults() []ToolResult { + toolResults := make([]ToolResult, 0) + for _, part := range m.Parts { + if c, ok := part.(ToolResult); ok { + toolResults = append(toolResults, c) + } + } + return toolResults +} + +func (m *Message) IsFinished() bool { + for _, part := range m.Parts { + if _, ok := part.(Finish); ok { + return true + } + } + return false +} + +func (m *Message) FinishPart() *Finish { + for _, part := range m.Parts { + if c, ok := part.(Finish); ok { + return &c + } + } + return nil +} + +func (m *Message) FinishReason() FinishReason { + for _, part := range m.Parts { + if c, ok := part.(Finish); ok { + return c.Reason + } + } + return "" +} + +// IsErrorLike reports whether the message finished with an error-style +// banner (a real error or a provider safety refusal). The TUI renders +// both through the same banner path. +func (m *Message) IsErrorLike() bool { + switch m.FinishReason() { + case FinishReasonError, FinishReasonContentFilter: + return true + } + return false +} + +func (m *Message) IsThinking() bool { + if m.ReasoningContent().Thinking != "" && m.Content().Text == "" && !m.IsFinished() { + return true + } + return false +} + +func (m *Message) AppendContent(delta string) { + found := false + for i, part := range m.Parts { + if c, ok := part.(TextContent); ok { + m.Parts[i] = TextContent{Text: c.Text + delta} + found = true + } + } + if !found { + m.Parts = append(m.Parts, TextContent{Text: delta}) + } +} + +func (m *Message) AppendReasoningContent(delta string) { + found := false + for i, part := range m.Parts { + if c, ok := part.(ReasoningContent); ok { + m.Parts[i] = ReasoningContent{ + Thinking: c.Thinking + delta, + Signature: c.Signature, + StartedAt: c.StartedAt, + FinishedAt: c.FinishedAt, + } + found = true + } + } + if !found { + m.Parts = append(m.Parts, ReasoningContent{ + Thinking: delta, + StartedAt: time.Now().Unix(), + }) + } +} + +func (m *Message) AppendThoughtSignature(signature string, toolCallID string) { + for i, part := range m.Parts { + if c, ok := part.(ReasoningContent); ok { + m.Parts[i] = ReasoningContent{ + Thinking: c.Thinking, + ThoughtSignature: c.ThoughtSignature + signature, + ToolID: toolCallID, + Signature: c.Signature, + StartedAt: c.StartedAt, + FinishedAt: c.FinishedAt, + } + return + } + } + m.Parts = append(m.Parts, ReasoningContent{ThoughtSignature: signature}) +} + +func (m *Message) AppendReasoningSignature(signature string) { + for i, part := range m.Parts { + if c, ok := part.(ReasoningContent); ok { + m.Parts[i] = ReasoningContent{ + Thinking: c.Thinking, + Signature: c.Signature + signature, + StartedAt: c.StartedAt, + FinishedAt: c.FinishedAt, + } + return + } + } + m.Parts = append(m.Parts, ReasoningContent{Signature: signature}) +} + +func (m *Message) SetReasoningResponsesData(data *openai.ResponsesReasoningMetadata) { + for i, part := range m.Parts { + if c, ok := part.(ReasoningContent); ok { + m.Parts[i] = ReasoningContent{ + Thinking: c.Thinking, + ResponsesData: data, + StartedAt: c.StartedAt, + FinishedAt: c.FinishedAt, + } + return + } + } +} + +func (m *Message) FinishThinking() { + for i, part := range m.Parts { + if c, ok := part.(ReasoningContent); ok { + if c.FinishedAt == 0 { + m.Parts[i] = ReasoningContent{ + Thinking: c.Thinking, + Signature: c.Signature, + StartedAt: c.StartedAt, + FinishedAt: time.Now().Unix(), + } + } + return + } + } +} + +func (m *Message) ThinkingDuration() time.Duration { + reasoning := m.ReasoningContent() + if reasoning.StartedAt == 0 { + return 0 + } + + endTime := reasoning.FinishedAt + if endTime == 0 { + endTime = time.Now().Unix() + } + + return time.Duration(endTime-reasoning.StartedAt) * time.Second +} + +func (m *Message) FinishToolCall(toolCallID string) { + for i, part := range m.Parts { + if c, ok := part.(ToolCall); ok { + if c.ID == toolCallID { + m.Parts[i] = ToolCall{ + ID: c.ID, + Name: c.Name, + Input: c.Input, + Finished: true, + } + return + } + } + } +} + +func (m *Message) AppendToolCallInput(toolCallID string, inputDelta string) { + for i, part := range m.Parts { + if c, ok := part.(ToolCall); ok { + if c.ID == toolCallID { + m.Parts[i] = ToolCall{ + ID: c.ID, + Name: c.Name, + Input: c.Input + inputDelta, + Finished: c.Finished, + } + return + } + } + } +} + +func (m *Message) AddToolCall(tc ToolCall) { + for i, part := range m.Parts { + if c, ok := part.(ToolCall); ok { + if c.ID == tc.ID { + m.Parts[i] = tc + return + } + } + } + m.Parts = append(m.Parts, tc) +} + +func (m *Message) SetToolCalls(tc []ToolCall) { + // remove any existing tool call part it could have multiple + parts := make([]ContentPart, 0) + for _, part := range m.Parts { + if _, ok := part.(ToolCall); ok { + continue + } + parts = append(parts, part) + } + m.Parts = parts + for _, toolCall := range tc { + m.Parts = append(m.Parts, toolCall) + } +} + +func (m *Message) AddToolResult(tr ToolResult) { + m.Parts = append(m.Parts, tr) +} + +func (m *Message) SetToolResults(tr []ToolResult) { + for _, toolResult := range tr { + m.Parts = append(m.Parts, toolResult) + } +} + +// Clone returns a deep copy of the message with an independent Parts slice. +// This prevents race conditions when the message is modified concurrently. +func (m *Message) Clone() Message { + clone := *m + clone.Parts = make([]ContentPart, len(m.Parts)) + copy(clone.Parts, m.Parts) + return clone +} + +// ResetStreamedContent removes all parts that were added during streaming +// (text, reasoning, tool calls, finish) so the message is ready for a +// retry. Non-streamed parts (images, binary attachments, tool results, +// shell commands) are preserved. +func (m *Message) ResetStreamedContent() { + kept := m.Parts[:0] + for _, part := range m.Parts { + switch part.(type) { + case TextContent, ReasoningContent, ToolCall, Finish: + // Drop streamed parts. + default: + kept = append(kept, part) + } + } + m.Parts = kept +} + +func (m *Message) AddFinish(reason FinishReason, message, details string) { + // remove any existing finish part + for i, part := range m.Parts { + if _, ok := part.(Finish); ok { + m.Parts = slices.Delete(m.Parts, i, i+1) + break + } + } + m.Parts = append(m.Parts, Finish{Reason: reason, Time: time.Now().Unix(), Message: message, Details: details}) +} + +func (m *Message) AddImageURL(url, detail string) { + m.Parts = append(m.Parts, ImageURLContent{URL: url, Detail: detail}) +} + +func (m *Message) AddBinary(mimeType string, data []byte) { + m.Parts = append(m.Parts, BinaryContent{MIMEType: mimeType, Data: data}) +} + +func PromptWithTextAttachments(prompt string, attachments []Attachment) string { + var sb strings.Builder + sb.WriteString(prompt) + addedAttachments := false + for _, content := range attachments { + if !content.IsText() { + continue + } + if !addedAttachments { + sb.WriteString("\nThe files below have been attached by the user, consider them in your response\n") + addedAttachments = true + } + if content.FilePath != "" { + fmt.Fprintf(&sb, "\n", content.FilePath) + } else { + sb.WriteString("\n") + } + sb.WriteString("\n") + sb.Write(content.Content) + sb.WriteString("\n\n") + } + return sb.String() +} + +func (m *Message) ToAIMessage() []fantasy.Message { + var messages []fantasy.Message + switch m.Role { + case User: + var parts []fantasy.MessagePart + text := strings.TrimSpace(m.Content().Text) + var textAttachments []Attachment + for _, content := range m.BinaryContent() { + if !strings.HasPrefix(content.MIMEType, "text/") { + continue + } + textAttachments = append(textAttachments, Attachment{ + FilePath: content.Path, + MimeType: content.MIMEType, + Content: content.Data, + }) + } + text = PromptWithTextAttachments(text, textAttachments) + // Include bang-mode shell commands as context for the agent. + for _, sc := range m.ShellCommands() { + shellText := fmt.Sprintf("$ %s\n%s\n(exit code %d)", sc.Command, ansi.Strip(sc.Output), sc.ExitCode) + if text != "" { + text += "\n\n" + shellText + } else { + text = shellText + } + } + if text != "" { + parts = append(parts, fantasy.TextPart{Text: text}) + } + for _, content := range m.BinaryContent() { + // skip text attachements + if strings.HasPrefix(content.MIMEType, "text/") { + continue + } + parts = append(parts, fantasy.FilePart{ + Filename: content.Path, + Data: content.Data, + MediaType: content.MIMEType, + }) + } + messages = append(messages, fantasy.Message{ + Role: fantasy.MessageRoleUser, + Content: parts, + }) + case Assistant: + var parts []fantasy.MessagePart + text := strings.TrimSpace(m.Content().Text) + if text != "" { + parts = append(parts, fantasy.TextPart{Text: text}) + } + reasoning := m.ReasoningContent() + if reasoning.Thinking != "" { + reasoningPart := fantasy.ReasoningPart{Text: reasoning.Thinking, ProviderOptions: fantasy.ProviderOptions{}} + if reasoning.Signature != "" { + reasoningPart.ProviderOptions[anthropic.Name] = &anthropic.ReasoningOptionMetadata{ + Signature: reasoning.Signature, + } + } + if reasoning.ResponsesData != nil { + reasoningPart.ProviderOptions[openai.Name] = reasoning.ResponsesData + } + if reasoning.ThoughtSignature != "" { + reasoningPart.ProviderOptions[google.Name] = &google.ReasoningMetadata{ + Signature: reasoning.ThoughtSignature, + ToolID: reasoning.ToolID, + } + } + parts = append(parts, reasoningPart) + } + for _, call := range m.ToolCalls() { + parts = append(parts, fantasy.ToolCallPart{ + ToolCallID: call.ID, + ToolName: call.Name, + Input: call.Input, + ProviderExecuted: call.ProviderExecuted, + }) + } + messages = append(messages, fantasy.Message{ + Role: fantasy.MessageRoleAssistant, + Content: parts, + }) + case Tool: + var parts []fantasy.MessagePart + for _, result := range m.ToolResults() { + var content fantasy.ToolResultOutputContent + if result.IsError { + content = fantasy.ToolResultOutputContentError{ + Error: errors.New(result.Content), + } + } else if result.Data != "" { + if stringext.IsValidBase64(result.Data) { + content = fantasy.ToolResultOutputContentMedia{ + Data: result.Data, + MediaType: result.MIMEType, + } + } else { + content = fantasy.ToolResultOutputContentText{ + Text: mediaLoadFailedPlaceholder, + } + } + } else { + content = fantasy.ToolResultOutputContentText{ + Text: result.Content, + } + } + parts = append(parts, fantasy.ToolResultPart{ + ToolCallID: result.ToolCallID, + Output: content, + }) + } + messages = append(messages, fantasy.Message{ + Role: fantasy.MessageRoleTool, + Content: parts, + }) + } + return messages +} diff --git a/internal/message/content_test.go b/internal/message/content_test.go new file mode 100644 index 0000000000000000000000000000000000000000..a92c2a0b33018d9f64a26f75ee9beb1698ea4de5 --- /dev/null +++ b/internal/message/content_test.go @@ -0,0 +1,175 @@ +package message + +import ( + "encoding/base64" + "fmt" + "strings" + "testing" + + "charm.land/fantasy" + "github.com/stretchr/testify/require" +) + +func makeTestAttachments(n int, contentSize int) []Attachment { + attachments := make([]Attachment, n) + content := []byte(strings.Repeat("x", contentSize)) + for i := range n { + attachments[i] = Attachment{ + FilePath: fmt.Sprintf("/path/to/file%d.txt", i), + MimeType: "text/plain", + Content: content, + } + } + return attachments +} + +func TestToAIMessage_CorruptedMediaData(t *testing.T) { + t.Parallel() + + msg := &Message{ + Role: Tool, + Parts: []ContentPart{ + ToolResult{ + ToolCallID: "call_123", + Name: "screenshot", + Content: "Loaded image/png content", + Data: "abc\x80def", + MIMEType: "image/png", + }, + }, + } + + messages := msg.ToAIMessage() + require.Len(t, messages, 1) + require.Len(t, messages[0].Content, 1) + + part, ok := messages[0].Content[0].(fantasy.ToolResultPart) + require.True(t, ok) + + require.Equal(t, "call_123", part.ToolCallID) + + textContent, ok := part.Output.(fantasy.ToolResultOutputContentText) + require.True(t, ok, "corrupted media should be downgraded to text") + require.Equal(t, mediaLoadFailedPlaceholder, textContent.Text) +} + +func TestToAIMessage_ValidMediaData(t *testing.T) { + t.Parallel() + + validBase64 := base64.StdEncoding.EncodeToString([]byte{0x89, 0x50, 0x4E, 0x47}) + + msg := &Message{ + Role: Tool, + Parts: []ContentPart{ + ToolResult{ + ToolCallID: "call_456", + Name: "screenshot", + Content: "Loaded image/png content", + Data: validBase64, + MIMEType: "image/png", + }, + }, + } + + messages := msg.ToAIMessage() + require.Len(t, messages, 1) + require.Len(t, messages[0].Content, 1) + + part, ok := messages[0].Content[0].(fantasy.ToolResultPart) + require.True(t, ok) + + require.Equal(t, "call_456", part.ToolCallID) + + mediaContent, ok := part.Output.(fantasy.ToolResultOutputContentMedia) + require.True(t, ok, "valid media should remain as media") + require.Equal(t, validBase64, mediaContent.Data) + require.Equal(t, "image/png", mediaContent.MediaType) +} + +func TestToAIMessage_ASCIIButInvalidBase64(t *testing.T) { + t.Parallel() + + msg := &Message{ + Role: Tool, + Parts: []ContentPart{ + ToolResult{ + ToolCallID: "call_789", + Name: "screenshot", + Content: "Loaded image/png content", + Data: "not-valid-base64!!!", + MIMEType: "image/png", + }, + }, + } + + messages := msg.ToAIMessage() + require.Len(t, messages, 1) + require.Len(t, messages[0].Content, 1) + + part, ok := messages[0].Content[0].(fantasy.ToolResultPart) + require.True(t, ok) + + require.Equal(t, "call_789", part.ToolCallID) + + textContent, ok := part.Output.(fantasy.ToolResultOutputContentText) + require.True(t, ok, "ASCII but invalid base64 should be downgraded to text") + require.Equal(t, mediaLoadFailedPlaceholder, textContent.Text) +} + +func BenchmarkPromptWithTextAttachments(b *testing.B) { + cases := []struct { + name string + numFiles int + contentSize int + }{ + {"1file_100bytes", 1, 100}, + {"5files_1KB", 5, 1024}, + {"10files_10KB", 10, 10 * 1024}, + {"20files_50KB", 20, 50 * 1024}, + } + + for _, tc := range cases { + attachments := makeTestAttachments(tc.numFiles, tc.contentSize) + prompt := "Process these files" + + b.Run(tc.name, func(b *testing.B) { + b.ReportAllocs() + for range b.N { + _ = PromptWithTextAttachments(prompt, attachments) + } + }) + } +} + +func TestResetStreamedContent(t *testing.T) { + t.Parallel() + + msg := &Message{} + msg.AddImageURL("https://example.com/img.png", "high") + msg.AppendContent("partial answer") + msg.AppendReasoningContent("thinking...") + msg.AddToolCall(ToolCall{ID: "1", Name: "bash"}) + msg.AddToolResult(ToolResult{ToolCallID: "1", Content: "output"}) + msg.AddFinish(FinishReasonError, "boom", "stream died") + + msg.ResetStreamedContent() + + // Streamed parts are gone. + require.Empty(t, msg.Content().Text, "text should be cleared") + require.Empty(t, msg.ReasoningContent().Thinking, "reasoning should be cleared") + require.Empty(t, msg.ToolCalls(), "tool calls should be cleared") + require.Nil(t, msg.FinishPart(), "finish should be cleared") + + // Non-streamed parts survive. + require.Len(t, msg.ImageURLContent(), 1, "image should survive") + require.Len(t, msg.ToolResults(), 1, "tool results should survive") +} + +func TestResetStreamedContentEmpty(t *testing.T) { + t.Parallel() + + // Reset on an empty message is a no-op and must not panic. + msg := &Message{} + msg.ResetStreamedContent() + require.Empty(t, msg.Parts) +} diff --git a/internal/message/message.go b/internal/message/message.go new file mode 100644 index 0000000000000000000000000000000000000000..18e0d712444178c21480bf8df0637e4434022f6d --- /dev/null +++ b/internal/message/message.go @@ -0,0 +1,711 @@ +package message + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "sync" + "time" + + "github.com/charmbracelet/crush/internal/db" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/google/uuid" +) + +// defaultUpdateDebounce is the default debounce window for [Service.Update]. +// Streaming deltas that arrive within the window are coalesced into a +// single SQL write and a single pubsub event. Terminal updates +// (finish/error/cancel/tool-call structural changes) bypass the +// debounce and flush synchronously. +const defaultUpdateDebounce = 33 * time.Millisecond + +type CreateMessageParams struct { + Role MessageRole + Parts []ContentPart + Model string + Provider string + IsSummaryMessage bool +} + +// Service is the public interface to the message store. +// +// [Service.Update] is eventually consistent: it accepts new state into +// an in-memory buffer and writes it to SQLite plus publishes a +// [pubsub.UpdatedEvent] on the next debounce tick (default +// [defaultUpdateDebounce]) or on the next terminal-state update, +// whichever comes first. Terminal-state updates — those that finish +// the message, add or finish a tool call, or end a reasoning section — +// flush synchronously before [Service.Update] returns. +// +// Callers that need stronger ordering (e.g. tests, shutdown, +// session-switch reads) must use [Service.Flush] or [Service.FlushAll] +// before reading via [Service.Get] / [Service.List]. Without an +// explicit flush, a read can race the debounce timer and miss the +// most recent in-memory state. +type Service interface { + pubsub.Subscriber[Message] + Create(ctx context.Context, sessionID string, params CreateMessageParams) (Message, error) + Update(ctx context.Context, message Message) error + Get(ctx context.Context, id string) (Message, error) + List(ctx context.Context, sessionID string) ([]Message, error) + ListUserMessages(ctx context.Context, sessionID string) ([]Message, error) + ListAllUserMessages(ctx context.Context) ([]Message, error) + GetLastAssistantMessage(ctx context.Context, sessionID string) (Message, error) + Delete(ctx context.Context, id string) error + DeleteSessionMessages(ctx context.Context, sessionID string) error + + // Flush synchronously drains any pending debounced state for the + // given message ID, performs the SQL write, and publishes the + // resulting [pubsub.UpdatedEvent]. Idempotent; cheap no-op if no + // updates are pending. Use this before any read that must observe + // the latest [Service.Update]. + Flush(ctx context.Context, id string) error + + // FlushAll synchronously drains pending debounced state for every + // message known to the service. Intended for shutdown and + // session-switch paths. + FlushAll(ctx context.Context) error +} + +// pendingState holds the in-memory coalescing buffer for a single +// message ID. All fields except where noted are guarded by +// service.mu. The flushing flag serializes concurrent flushers for +// the same ID so SQL writes never reorder. +type pendingState struct { + // latest is the most recent [Message] passed to [Service.Update] + // that has not yet been flushed. + latest Message + + // dirty is true when latest contains state that has not been + // written to SQL since the last successful flush. + dirty bool + + // flushing is true while a goroutine is performing the SQL write + // for this ID. New updates are still accepted (and re-mark dirty) + // but other flushers must back off. + flushing bool + + // timer is the active debounce timer, or nil if no flush is + // scheduled. Stopped and reset when a terminal update preempts + // the debounce window. + timer *time.Timer + + // baseline is the projection of the snapshot most recently + // written to SQL that terminal-state detection needs. Used as the + // baseline for [shouldFlushNow]. + baseline flushBaseline + + // hasFlushed is false until the first successful write for this + // ID; until then baseline is the zero value and must not be + // treated as a real prior state. + hasFlushed bool +} + +// flushBaseline is the compact projection of a flushed [Message] that +// [shouldFlushNow] compares against. Pending state lives for as long +// as the process does, so it retains this instead of the whole +// Message: a finished stream would otherwise pin its parts — reasoning +// text, tool inputs, tool results — on the heap forever. +type flushBaseline struct { + toolCallsFinished []bool + reasoningFinishedAt int64 +} + +// newFlushBaseline projects the fields [shouldFlushNow] reads out of a +// flushed message. +func newFlushBaseline(m *Message) flushBaseline { + calls := m.ToolCalls() + finished := make([]bool, len(calls)) + for i, c := range calls { + finished[i] = c.Finished + } + return flushBaseline{ + toolCallsFinished: finished, + reasoningFinishedAt: m.ReasoningContent().FinishedAt, + } +} + +type service struct { + *pubsub.Broker[Message] + q db.Querier + debounce time.Duration + + mu sync.Mutex + pending map[string]*pendingState +} + +// ServiceOption configures a [Service] at construction. +type ServiceOption func(*service) + +// WithDebounce overrides the debounce window for [Service.Update]. A +// zero or negative value disables debouncing entirely (every update +// flushes synchronously). Intended primarily for tests. +func WithDebounce(d time.Duration) ServiceOption { + return func(s *service) { + s.debounce = d + } +} + +func NewService(q db.Querier, opts ...ServiceOption) Service { + s := &service{ + Broker: pubsub.NewBroker[Message](), + q: q, + debounce: defaultUpdateDebounce, + pending: make(map[string]*pendingState), + } + for _, opt := range opts { + opt(s) + } + return s +} + +func (s *service) Delete(ctx context.Context, id string) error { + message, err := s.Get(ctx, id) + if err != nil { + return err + } + err = s.q.DeleteMessage(ctx, message.ID) + if err != nil { + return err + } + // Drop any pending coalesced state for this ID. We never want to + // flush back over a deleted row. + s.mu.Lock() + if p, ok := s.pending[id]; ok { + if p.timer != nil { + p.timer.Stop() + } + delete(s.pending, id) + } + s.mu.Unlock() + // Clone the message before publishing to avoid race conditions with + // concurrent modifications to the Parts slice. + s.Publish(pubsub.DeletedEvent, message.Clone()) + return nil +} + +func (s *service) Create(ctx context.Context, sessionID string, params CreateMessageParams) (Message, error) { + if params.Role != Assistant { + params.Parts = append(params.Parts, Finish{ + Reason: "stop", + }) + } + partsJSON, err := marshalParts(params.Parts) + if err != nil { + return Message{}, err + } + isSummary := int64(0) + if params.IsSummaryMessage { + isSummary = 1 + } + dbMessage, err := s.q.CreateMessage(ctx, db.CreateMessageParams{ + ID: uuid.New().String(), + SessionID: sessionID, + Role: string(params.Role), + Parts: string(partsJSON), + Model: sql.NullString{String: string(params.Model), Valid: true}, + Provider: sql.NullString{String: params.Provider, Valid: params.Provider != ""}, + IsSummaryMessage: isSummary, + }) + if err != nil { + return Message{}, err + } + message, err := s.fromDBItem(dbMessage) + if err != nil { + return Message{}, err + } + // Clone the message before publishing to avoid race conditions with + // concurrent modifications to the Parts slice. + s.Publish(pubsub.CreatedEvent, message.Clone()) + return message, nil +} + +func (s *service) DeleteSessionMessages(ctx context.Context, sessionID string) error { + messages, err := s.List(ctx, sessionID) + if err != nil { + return err + } + for _, message := range messages { + if message.SessionID == sessionID { + err = s.Delete(ctx, message.ID) + if err != nil { + return err + } + } + } + return nil +} + +// Update accepts a new state for a message and either flushes +// synchronously (terminal updates, debounce <= 0) or buffers it until +// the next debounce tick. See [Service] for the contract. +func (s *service) Update(ctx context.Context, msg Message) error { + cloned := msg.Clone() + + // Zero or negative debounce: flush every update synchronously. This + // preserves the pre-coalescing behaviour for tests and any caller + // that explicitly opted out via [WithDebounce]. + if s.debounce <= 0 { + s.mu.Lock() + p, ok := s.pending[msg.ID] + if !ok { + p = &pendingState{} + s.pending[msg.ID] = p + } + p.latest = cloned + p.dirty = true + s.mu.Unlock() + return s.flushOne(ctx, msg.ID, true) + } + + s.mu.Lock() + p, ok := s.pending[msg.ID] + if !ok { + p = &pendingState{} + s.pending[msg.ID] = p + } + p.latest = cloned + p.dirty = true + + var prev *flushBaseline + if p.hasFlushed { + prev = &p.baseline + } + terminal := shouldFlushNow(prev, &cloned) + + if terminal { + if p.timer != nil { + p.timer.Stop() + p.timer = nil + } + s.mu.Unlock() + return s.flushOne(ctx, msg.ID, true) + } + + // Debounce: schedule a single flush per pending state. If a flush + // is already running we let it finish; the trailing dirty bit will + // be picked up by the next Update or by Flush. + if p.timer == nil && !p.flushing { + id := msg.ID + p.timer = time.AfterFunc(s.debounce, func() { + // Detached from caller ctx so a cancelled stream context + // does not strand the buffered write. + _ = s.flushOne(context.Background(), id, false) + }) + } + s.mu.Unlock() + return nil +} + +// Flush implements [Service.Flush]. +func (s *service) Flush(ctx context.Context, id string) error { + return s.flushOne(ctx, id, true) +} + +// FlushAll implements [Service.FlushAll]. It snapshots every ID with +// outstanding work — either dirty buffered state or a flush already in +// flight — then drains each one. Picking up in-flight IDs ensures +// FlushAll cannot return while a timer-fired write is still mid-SQL, +// which is what shutdown and session-switch callers rely on. +func (s *service) FlushAll(ctx context.Context) error { + s.mu.Lock() + ids := make([]string, 0, len(s.pending)) + for id, p := range s.pending { + if p.dirty || p.flushing { + ids = append(ids, id) + } + } + s.mu.Unlock() + var firstErr error + for _, id := range ids { + if err := s.flushOne(ctx, id, true); err != nil && firstErr == nil { + firstErr = err + } + } + return firstErr +} + +// flushOne drains a single message ID. When syncCaller is true the +// caller is willing to wait through a concurrent in-flight flush so +// that, on return, the flushed state equals latest at the moment of +// return. +// When false (timer-fired path) we bail if another flusher is already +// running; that flusher will pick up the trailing dirty bit. +// +// Order matters: a sync caller must wait for any in-flight flush to +// drain even when the buffer is currently clean — that in-flight +// write has not yet updated the SQL row, so returning early would +// violate the contract that on success the flushed state reflects the +// most recent state. +func (s *service) flushOne(ctx context.Context, id string, syncCaller bool) error { + for { + s.mu.Lock() + p, ok := s.pending[id] + if !ok { + s.mu.Unlock() + return nil + } + if p.flushing { + if !syncCaller { + s.mu.Unlock() + return nil + } + s.mu.Unlock() + // Brief yield; in-flight write should land in <1ms typical. + time.Sleep(time.Millisecond) + continue + } + if !p.dirty { + s.mu.Unlock() + return nil + } + + if p.timer != nil { + p.timer.Stop() + p.timer = nil + } + snap := p.latest + // Decide whether this snapshot represents a terminal event + // against the prior baseline. We must do this before resetting + // dirty/flushing because shouldFlushNow looks at p.baseline + // (which is what was on disk before this write). + var prev *flushBaseline + if p.hasFlushed { + prev = &p.baseline + } + isTerminal := shouldFlushNow(prev, &snap) + p.flushing = true + p.dirty = false + s.mu.Unlock() + + err := s.write(ctx, snap) + + s.mu.Lock() + p.flushing = false + if err == nil { + p.baseline = newFlushBaseline(&snap) + p.hasFlushed = true + } else { + // Restore dirty so the next caller retries. + p.dirty = true + } + // If a delta arrived during the SQL write and we are a sync + // caller, the user expects that delta to land too. + wasDirty := p.dirty + if !wasDirty { + // Nothing left to write, so drop the snapshot. Pending + // entries outlive the streams that created them; holding + // latest here would pin every finished message's parts + // for the life of the process. + p.latest = Message{} + } + s.mu.Unlock() + + if err != nil { + return err + } + + // Terminal events — message finished, tool call added or + // finished, reasoning ended — use the bounded must-deliver + // path so they never get dropped under channel contention. + if isTerminal { + s.PublishMustDeliver(ctx, pubsub.UpdatedEvent, snap) + } else { + s.Publish(pubsub.UpdatedEvent, snap) + } + + if wasDirty && syncCaller { + continue + } + return nil + } +} + +// write performs the unguarded SQL write + UpdatedAt stamp. Caller +// owns publishing. +func (s *service) write(ctx context.Context, msg Message) error { + parts, err := marshalParts(msg.Parts) + if err != nil { + return err + } + finishedAt := sql.NullInt64{} + if f := msg.FinishPart(); f != nil { + finishedAt.Int64 = f.Time + finishedAt.Valid = true + } + if err := s.q.UpdateMessage(ctx, db.UpdateMessageParams{ + ID: msg.ID, + Parts: string(parts), + PrismModelID: sql.NullString{String: msg.PrismModelID, Valid: msg.PrismModelID != ""}, + PrismModelName: sql.NullString{String: msg.PrismModelName, Valid: msg.PrismModelName != ""}, + PrismHypercreditSavings: nullableFloat(msg.PrismHypercreditSavings), + PrismDollarSavings: nullableFloat(msg.PrismDollarSavings), + FinishedAt: finishedAt, + }); err != nil { + return err + } + return nil +} + +func nullableFloat(v *float64) sql.NullFloat64 { + if v == nil { + return sql.NullFloat64{} + } + return sql.NullFloat64{Float64: *v, Valid: true} +} + +func floatPtr(v sql.NullFloat64) *float64 { + if !v.Valid { + return nil + } + value := v.Float64 + return &value +} + +// shouldFlushNow returns true when next represents a structural +// change that must not be silently coalesced: the message just +// finished, the tool-call set grew, a tool call transitioned to +// finished, or reasoning just finished. prev is the last-flushed +// snapshot (or nil if no write has landed yet). +func shouldFlushNow(prev *flushBaseline, next *Message) bool { + if next.IsFinished() { + return true + } + + var prevCalls []bool + var prevReasoningFinishedAt int64 + if prev != nil { + prevCalls = prev.toolCallsFinished + prevReasoningFinishedAt = prev.reasoningFinishedAt + } + nextCalls := next.ToolCalls() + if len(nextCalls) != len(prevCalls) { + return true + } + for i := range nextCalls { + // Bounds-safe: lengths are equal here. + if nextCalls[i].Finished != prevCalls[i] { + return true + } + // A tool call's input only matters once it has landed (Finished + // flips true). Earlier deltas to Input are debounced with the + // rest of the streaming state. + } + if next.ReasoningContent().FinishedAt > 0 && prevReasoningFinishedAt == 0 { + return true + } + return false +} + +func (s *service) Get(ctx context.Context, id string) (Message, error) { + dbMessage, err := s.q.GetMessage(ctx, id) + if err != nil { + return Message{}, err + } + return s.fromDBItem(dbMessage) +} + +func (s *service) List(ctx context.Context, sessionID string) ([]Message, error) { + dbMessages, err := s.q.ListMessagesBySession(ctx, sessionID) + if err != nil { + return nil, err + } + messages := make([]Message, len(dbMessages)) + for i, dbMessage := range dbMessages { + messages[i], err = s.fromDBItem(dbMessage) + if err != nil { + return nil, err + } + } + return messages, nil +} + +func (s *service) ListUserMessages(ctx context.Context, sessionID string) ([]Message, error) { + dbMessages, err := s.q.ListUserMessagesBySession(ctx, sessionID) + if err != nil { + return nil, err + } + messages := make([]Message, len(dbMessages)) + for i, dbMessage := range dbMessages { + messages[i], err = s.fromDBItem(dbMessage) + if err != nil { + return nil, err + } + } + return messages, nil +} + +func (s *service) ListAllUserMessages(ctx context.Context) ([]Message, error) { + dbMessages, err := s.q.ListAllUserMessages(ctx) + if err != nil { + return nil, err + } + messages := make([]Message, len(dbMessages)) + for i, dbMessage := range dbMessages { + messages[i], err = s.fromDBItem(dbMessage) + if err != nil { + return nil, err + } + } + return messages, nil +} + +func (s *service) GetLastAssistantMessage(ctx context.Context, sessionID string) (Message, error) { + dbMessage, err := s.q.GetLastAssistantMessageBySession(ctx, sessionID) + if err != nil { + return Message{}, err + } + return s.fromDBItem(dbMessage) +} + +func (s *service) fromDBItem(item db.Message) (Message, error) { + parts, err := unmarshalParts([]byte(item.Parts)) + if err != nil { + return Message{}, err + } + return Message{ + ID: item.ID, + SessionID: item.SessionID, + Role: MessageRole(item.Role), + Parts: parts, + Model: item.Model.String, + Provider: item.Provider.String, + CreatedAt: item.CreatedAt, + UpdatedAt: item.UpdatedAt, + IsSummaryMessage: item.IsSummaryMessage != 0, + PrismModelID: item.PrismModelID.String, + PrismModelName: item.PrismModelName.String, + PrismHypercreditSavings: floatPtr(item.PrismHypercreditSavings), + PrismDollarSavings: floatPtr(item.PrismDollarSavings), + }, nil +} + +type partType string + +const ( + reasoningType partType = "reasoning" + textType partType = "text" + imageURLType partType = "image_url" + binaryType partType = "binary" + toolCallType partType = "tool_call" + toolResultType partType = "tool_result" + finishType partType = "finish" + shellCommandType partType = "shell_command" +) + +type partWrapper struct { + Type partType `json:"type"` + Data ContentPart `json:"data"` +} + +func marshalParts(parts []ContentPart) ([]byte, error) { + wrappedParts := make([]partWrapper, len(parts)) + + for i, part := range parts { + var typ partType + + switch part.(type) { + case ReasoningContent: + typ = reasoningType + case TextContent: + typ = textType + case ImageURLContent: + typ = imageURLType + case BinaryContent: + typ = binaryType + case ToolCall: + typ = toolCallType + case ToolResult: + typ = toolResultType + case Finish: + typ = finishType + case ShellCommand: + typ = shellCommandType + default: + return nil, fmt.Errorf("unknown part type: %T", part) + } + + wrappedParts[i] = partWrapper{ + Type: typ, + Data: part, + } + } + return json.Marshal(wrappedParts) +} + +func unmarshalParts(data []byte) ([]ContentPart, error) { + temp := []json.RawMessage{} + + if err := json.Unmarshal(data, &temp); err != nil { + return nil, err + } + + parts := make([]ContentPart, 0) + + for _, rawPart := range temp { + var wrapper struct { + Type partType `json:"type"` + Data json.RawMessage `json:"data"` + } + + if err := json.Unmarshal(rawPart, &wrapper); err != nil { + return nil, err + } + + switch wrapper.Type { + case reasoningType: + part := ReasoningContent{} + if err := json.Unmarshal(wrapper.Data, &part); err != nil { + return nil, err + } + parts = append(parts, part) + case textType: + part := TextContent{} + if err := json.Unmarshal(wrapper.Data, &part); err != nil { + return nil, err + } + parts = append(parts, part) + case imageURLType: + part := ImageURLContent{} + if err := json.Unmarshal(wrapper.Data, &part); err != nil { + return nil, err + } + parts = append(parts, part) + case binaryType: + part := BinaryContent{} + if err := json.Unmarshal(wrapper.Data, &part); err != nil { + return nil, err + } + parts = append(parts, part) + case toolCallType: + part := ToolCall{} + if err := json.Unmarshal(wrapper.Data, &part); err != nil { + return nil, err + } + parts = append(parts, part) + case toolResultType: + part := ToolResult{} + if err := json.Unmarshal(wrapper.Data, &part); err != nil { + return nil, err + } + parts = append(parts, part) + case finishType: + part := Finish{} + if err := json.Unmarshal(wrapper.Data, &part); err != nil { + return nil, err + } + parts = append(parts, part) + case shellCommandType: + part := ShellCommand{} + if err := json.Unmarshal(wrapper.Data, &part); err != nil { + return nil, err + } + parts = append(parts, part) + default: + return nil, fmt.Errorf("unknown part type: %s", wrapper.Type) + } + } + + return parts, nil +} diff --git a/internal/message/message_test.go b/internal/message/message_test.go new file mode 100644 index 0000000000000000000000000000000000000000..c75822a6ab29460060179b60f09300c4b93969dc --- /dev/null +++ b/internal/message/message_test.go @@ -0,0 +1,810 @@ +package message + +import ( + "context" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/charmbracelet/crush/internal/db" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/charmbracelet/crush/internal/session" + "github.com/stretchr/testify/require" +) + +// slowUpdateQuerier wraps a [db.Querier] and forces UpdateMessage to +// hang on a release channel. Used to simulate an in-flight SQL write. +type slowUpdateQuerier struct { + db.Querier + release chan struct{} + started chan struct{} + startOnce sync.Once +} + +func (s *slowUpdateQuerier) UpdateMessage(ctx context.Context, arg db.UpdateMessageParams) error { + s.startOnce.Do(func() { close(s.started) }) + select { + case <-s.release: + case <-ctx.Done(): + return ctx.Err() + } + return s.Querier.UpdateMessage(ctx, arg) +} + +// newTestService spins up a fresh in-memory message.Service backed by a +// temporary on-disk SQLite database. Returns the service plus a session +// ID to attach messages to. +func newTestService(t *testing.T, opts ...ServiceOption) (Service, string) { + t.Helper() + conn, err := db.Connect(t.Context(), t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + + q := db.New(conn) + sessions := session.NewService(q, conn) + sess, err := sessions.Create(t.Context(), "test") + require.NoError(t, err) + + svc := NewService(q, opts...) + return svc, sess.ID +} + +// eventCollector consumes broker events into a slice in a goroutine +// and exposes thread-safe Snapshot / Reset helpers for assertions. +type eventCollector struct { + mu sync.Mutex + events []pubsub.Event[Message] +} + +func collect(ctx context.Context, sub <-chan pubsub.Event[Message]) *eventCollector { + c := &eventCollector{} + go func() { + for { + select { + case <-ctx.Done(): + return + case ev, ok := <-sub: + if !ok { + return + } + c.mu.Lock() + c.events = append(c.events, ev) + c.mu.Unlock() + } + } + }() + return c +} + +func (c *eventCollector) snapshot() []pubsub.Event[Message] { + c.mu.Lock() + defer c.mu.Unlock() + out := make([]pubsub.Event[Message], len(c.events)) + copy(out, c.events) + return out +} + +func (c *eventCollector) reset() { + c.mu.Lock() + defer c.mu.Unlock() + c.events = nil +} + +func TestUpdate_DebouncesTextDeltas(t *testing.T) { + t.Parallel() + + // Long-enough debounce that we can verify nothing flushes prematurely. + svc, sessionID := newTestService(t, WithDebounce(50*time.Millisecond)) + + subCtx, cancelSub := context.WithCancel(t.Context()) + defer cancelSub() + sub := svc.Subscribe(subCtx) + collector := collect(subCtx, sub) + + msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{ + Role: Assistant, + }) + require.NoError(t, err) + // Drop the CreatedEvent emitted by Create. + time.Sleep(5 * time.Millisecond) + collector.reset() + + // Push 5 deltas inside a single debounce window. + for i := 0; i < 5; i++ { + msg.AppendContent("a") + require.NoError(t, svc.Update(t.Context(), msg)) + } + + // Before the debounce expires no UpdatedEvent should have landed. + time.Sleep(10 * time.Millisecond) + require.Empty(t, collector.snapshot(), "no events should land before debounce window expires") + + // Wait for the debounce timer to fire. + require.Eventually(t, func() bool { + return len(collector.snapshot()) >= 1 + }, time.Second, 5*time.Millisecond) + events := collector.snapshot() + require.Len(t, events, 1, "5 deltas should coalesce into 1 UpdatedEvent") + require.Equal(t, pubsub.UpdatedEvent, events[0].Type) + require.Equal(t, "aaaaa", events[0].Payload.Content().Text) + + // Final state must be persisted. + got, err := svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.Equal(t, "aaaaa", got.Content().Text) +} + +func TestUpdate_TerminalUpdatesFlushSynchronously(t *testing.T) { + t.Parallel() + + svc, sessionID := newTestService(t, WithDebounce(time.Hour)) + + subCtx, cancelSub := context.WithCancel(t.Context()) + defer cancelSub() + sub := svc.Subscribe(subCtx) + collector := collect(subCtx, sub) + + msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant}) + require.NoError(t, err) + time.Sleep(5 * time.Millisecond) + collector.reset() + + // AddFinish makes the message terminal; Update must flush + // synchronously even with a 1-hour debounce. + msg.AppendContent("done") + msg.AddFinish(FinishReasonEndTurn, "", "") + require.NoError(t, svc.Update(t.Context(), msg)) + + require.Eventually(t, func() bool { + return len(collector.snapshot()) >= 1 + }, time.Second, 5*time.Millisecond, + "terminal update must publish without waiting for debounce") + events := collector.snapshot() + require.Len(t, events, 1) + require.True(t, events[0].Payload.IsFinished()) + + got, err := svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.True(t, got.IsFinished()) +} + +func TestUpdate_ToolCallStructuralChangeFlushes(t *testing.T) { + t.Parallel() + + svc, sessionID := newTestService(t, WithDebounce(time.Hour)) + + msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant}) + require.NoError(t, err) + + // Adding a new tool call is a structural change → sync flush. + msg.AddToolCall(ToolCall{ID: "tc1", Name: "view", Finished: false}) + require.NoError(t, svc.Update(t.Context(), msg)) + + got, err := svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.Len(t, got.ToolCalls(), 1) + require.Equal(t, "tc1", got.ToolCalls()[0].ID) + + // Marking the tool call finished is also a structural change. + msg.AddToolCall(ToolCall{ID: "tc1", Name: "view", Input: "{}", Finished: true}) + require.NoError(t, svc.Update(t.Context(), msg)) + + got, err = svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.True(t, got.ToolCalls()[0].Finished) +} + +func TestUpdate_ReasoningEndFlushes(t *testing.T) { + t.Parallel() + + svc, sessionID := newTestService(t, WithDebounce(time.Hour)) + + msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant}) + require.NoError(t, err) + + // Reasoning deltas alone debounce. + msg.AppendReasoningContent("hmm") + require.NoError(t, svc.Update(t.Context(), msg)) + got, err := svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.Empty(t, got.ReasoningContent().Thinking, "reasoning delta should still be in the debounce buffer") + + // FinishThinking sets FinishedAt → terminal flush. + msg.FinishThinking() + require.NoError(t, svc.Update(t.Context(), msg)) + + got, err = svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.Equal(t, "hmm", got.ReasoningContent().Thinking) + require.NotZero(t, got.ReasoningContent().FinishedAt) +} + +func TestFlush_DrainsPendingDebouncedUpdates(t *testing.T) { + t.Parallel() + + svc, sessionID := newTestService(t, WithDebounce(time.Hour)) + + msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant}) + require.NoError(t, err) + msg.AppendContent("buffered") + require.NoError(t, svc.Update(t.Context(), msg)) + + // Without a flush the SQL row is unchanged from Create. + got, err := svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.Empty(t, got.Content().Text) + + require.NoError(t, svc.Flush(t.Context(), msg.ID)) + + got, err = svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.Equal(t, "buffered", got.Content().Text) + + // Subsequent Flush is a no-op. + require.NoError(t, svc.Flush(t.Context(), msg.ID)) +} + +func TestFlushAll_DrainsAllPending(t *testing.T) { + t.Parallel() + + svc, sessionID := newTestService(t, WithDebounce(time.Hour)) + + const n = 5 + msgs := make([]Message, n) + for i := range msgs { + m, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant}) + require.NoError(t, err) + m.AppendContent("hi") + require.NoError(t, svc.Update(t.Context(), m)) + msgs[i] = m + } + + require.NoError(t, svc.FlushAll(t.Context())) + + for _, m := range msgs { + got, err := svc.Get(t.Context(), m.ID) + require.NoError(t, err) + require.Equal(t, "hi", got.Content().Text, "FlushAll should drain every pending message") + } +} + +func TestUpdate_OrderingMatchesNonCoalesced(t *testing.T) { + t.Parallel() + + // Compare the final state after coalesced vs zero-debounce updates. + // A sequence of interleaved text/reasoning/tool-call updates must + // converge to the same final DB row either way. + build := func(svc Service, sessionID string) Message { + msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant}) + require.NoError(t, err) + msg.AppendReasoningContent("thinking 1 ") + require.NoError(t, svc.Update(t.Context(), msg)) + msg.AppendReasoningContent("thinking 2") + require.NoError(t, svc.Update(t.Context(), msg)) + msg.FinishThinking() + require.NoError(t, svc.Update(t.Context(), msg)) + msg.AppendContent("hello ") + require.NoError(t, svc.Update(t.Context(), msg)) + msg.AppendContent("world") + require.NoError(t, svc.Update(t.Context(), msg)) + msg.AddToolCall(ToolCall{ID: "tc", Name: "x", Finished: false}) + require.NoError(t, svc.Update(t.Context(), msg)) + msg.AddToolCall(ToolCall{ID: "tc", Name: "x", Input: "{}", Finished: true}) + require.NoError(t, svc.Update(t.Context(), msg)) + msg.AddFinish(FinishReasonEndTurn, "", "") + require.NoError(t, svc.Update(t.Context(), msg)) + return msg + } + + coalesced, sid1 := newTestService(t, WithDebounce(20*time.Millisecond)) + a := build(coalesced, sid1) + require.NoError(t, coalesced.FlushAll(t.Context())) + gotA, err := coalesced.Get(t.Context(), a.ID) + require.NoError(t, err) + + immediate, sid2 := newTestService(t, WithDebounce(0)) + b := build(immediate, sid2) + gotB, err := immediate.Get(t.Context(), b.ID) + require.NoError(t, err) + + require.Equal(t, gotA.Content().Text, gotB.Content().Text) + require.Equal(t, gotA.ReasoningContent().Thinking, gotB.ReasoningContent().Thinking) + require.Equal(t, len(gotA.ToolCalls()), len(gotB.ToolCalls())) + require.Equal(t, gotA.IsFinished(), gotB.IsFinished()) +} + +func TestDelete_DropsPendingState(t *testing.T) { + t.Parallel() + + svc, sessionID := newTestService(t, WithDebounce(time.Hour)) + msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant}) + require.NoError(t, err) + msg.AppendContent("dropped") + require.NoError(t, svc.Update(t.Context(), msg)) + + require.NoError(t, svc.Delete(t.Context(), msg.ID)) + + // FlushAll after Delete must not write to the deleted row. + require.NoError(t, svc.FlushAll(t.Context())) + + _, err = svc.Get(t.Context(), msg.ID) + require.Error(t, err, "deleted message must remain deleted") +} + +func TestBroker_PublishLossyDropCounter(t *testing.T) { + t.Parallel() + + // Tiny channel buffer so we can saturate from a single sender. + b := pubsub.NewBrokerWithOptions[int](1) + defer b.Shutdown() + + subCtx, cancel := context.WithCancel(t.Context()) + defer cancel() + sub := b.Subscribe(subCtx) + require.NotNil(t, sub) + + // Don't read from sub. Saturate the buffer. + for range 100 { + b.Publish(pubsub.UpdatedEvent, 1) + } + + require.GreaterOrEqual(t, b.DropCount(), uint64(1), + "lossy Publish must increment the drop counter under contention") +} + +func TestBroker_PublishMustDeliverHonorsTimeout(t *testing.T) { + t.Parallel() + + b := pubsub.NewBrokerWithOptions[int](1) + b.SetMustDeliverTimeout(20 * time.Millisecond) + defer b.Shutdown() + + subCtx, cancel := context.WithCancel(t.Context()) + defer cancel() + sub := b.Subscribe(subCtx) + require.NotNil(t, sub) + + // Saturate: one event sits in the buffer, the second must wait. + b.Publish(pubsub.UpdatedEvent, 1) + + // PublishMustDeliver should block up to 20ms then drop. + start := time.Now() + b.PublishMustDeliver(t.Context(), pubsub.UpdatedEvent, 2) + elapsed := time.Since(start) + + require.GreaterOrEqual(t, elapsed, 20*time.Millisecond, + "PublishMustDeliver should block at least the timeout under contention") + require.Less(t, elapsed, 200*time.Millisecond, + "PublishMustDeliver must not block indefinitely") + require.GreaterOrEqual(t, b.MustDeliverDropCount(), uint64(1), + "timeout must increment the must-deliver drop counter") +} + +func TestBroker_PublishMustDeliverWithReader(t *testing.T) { + t.Parallel() + + b := pubsub.NewBrokerWithOptions[int](1) + b.SetMustDeliverTimeout(50 * time.Millisecond) + defer b.Shutdown() + + subCtx, cancel := context.WithCancel(t.Context()) + defer cancel() + sub := b.Subscribe(subCtx) + + var received atomic.Uint64 + done := make(chan struct{}) + go func() { + defer close(done) + for { + select { + case <-subCtx.Done(): + return + case _, ok := <-sub: + if !ok { + return + } + received.Add(1) + } + } + }() + + for i := range 10 { + b.PublishMustDeliver(t.Context(), pubsub.UpdatedEvent, i) + } + + // All 10 should land within the must-deliver timeout window. + require.Eventually(t, func() bool { return received.Load() == 10 }, + time.Second, 5*time.Millisecond, + "all must-deliver events should reach an active subscriber") + require.Zero(t, b.MustDeliverDropCount(), + "no drops expected when subscriber drains promptly") +} + +func TestUpdate_TerminalEventUsesMustDeliver(t *testing.T) { + t.Parallel() + + svc, sessionID := newTestService(t, WithDebounce(time.Hour)) + + subCtx, cancel := context.WithCancel(t.Context()) + defer cancel() + sub := svc.Subscribe(subCtx) + + var seenFinish atomic.Bool + done := make(chan struct{}) + go func() { + defer close(done) + for { + select { + case <-subCtx.Done(): + return + case ev, ok := <-sub: + if !ok { + return + } + if ev.Type == pubsub.UpdatedEvent && ev.Payload.IsFinished() { + seenFinish.Store(true) + } + } + } + }() + + msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant}) + require.NoError(t, err) + msg.AppendContent("final") + msg.AddFinish(FinishReasonEndTurn, "", "") + require.NoError(t, svc.Update(t.Context(), msg)) + + require.Eventually(t, func() bool { return seenFinish.Load() }, + time.Second, 10*time.Millisecond, + "terminal update must reach subscribers via the must-deliver path") +} + +func TestUpdate_ZeroDebounceFlushesEveryUpdate(t *testing.T) { + t.Parallel() + + svc, sessionID := newTestService(t, WithDebounce(0)) + + msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant}) + require.NoError(t, err) + + for i := 0; i < 3; i++ { + msg.AppendContent("x") + require.NoError(t, svc.Update(t.Context(), msg)) + got, err := svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.Len(t, got.Content().Text, i+1, "every update must land synchronously when debounce is 0") + } +} + +// TestFlush_WaitsForInFlightWrite reproduces the failure where Flush +// or FlushAll could return before a concurrent in-flight SQL write +// completed. We block UpdateMessage on a release channel, fire the +// debounce timer, then call Flush and assert it does not return until +// the in-flight write actually lands. +func TestFlush_WaitsForInFlightWrite(t *testing.T) { + t.Parallel() + + conn, err := db.Connect(t.Context(), t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + + q := db.New(conn) + sessions := session.NewService(q, conn) + sess, err := sessions.Create(t.Context(), "test") + require.NoError(t, err) + + slow := &slowUpdateQuerier{ + Querier: q, + release: make(chan struct{}), + started: make(chan struct{}), + } + // Short debounce so the timer fires quickly. + svc := NewService(slow, WithDebounce(10*time.Millisecond)) + + msg, err := svc.Create(t.Context(), sess.ID, CreateMessageParams{Role: Assistant}) + require.NoError(t, err) + msg.AppendContent("payload") + require.NoError(t, svc.Update(t.Context(), msg)) + + // Wait for the timer-fired flush to enter UpdateMessage. + select { + case <-slow.started: + case <-time.After(time.Second): + t.Fatal("timer-fired flush never reached UpdateMessage") + } + + // At this point the buffer is dirty=false but flushing=true. A + // naive Flush would early-return on !dirty. Spawn Flush in a + // goroutine and assert it has not returned while the write is + // still blocked. + flushDone := make(chan error, 1) + go func() { flushDone <- svc.Flush(t.Context(), msg.ID) }() + + select { + case err := <-flushDone: + t.Fatalf("Flush returned %v while in-flight write was still blocked", err) + case <-time.After(50 * time.Millisecond): + // Expected: Flush is correctly waiting. + } + + // Release the slow write; Flush must now return cleanly. + close(slow.release) + select { + case err := <-flushDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("Flush did not return after in-flight write completed") + } + + // The SQL row should now reflect the buffered payload. + got, err := svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.Equal(t, "payload", got.Content().Text) +} + +// TestFlushAll_WaitsForInFlightWrite asserts FlushAll picks up IDs +// whose buffer is currently flushing (dirty=false) so shutdown and +// session-switch callers don't return while an SQL write is mid-flight. +func TestFlushAll_WaitsForInFlightWrite(t *testing.T) { + t.Parallel() + + conn, err := db.Connect(t.Context(), t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + + q := db.New(conn) + sessions := session.NewService(q, conn) + sess, err := sessions.Create(t.Context(), "test") + require.NoError(t, err) + + slow := &slowUpdateQuerier{ + Querier: q, + release: make(chan struct{}), + started: make(chan struct{}), + } + svc := NewService(slow, WithDebounce(10*time.Millisecond)) + + msg, err := svc.Create(t.Context(), sess.ID, CreateMessageParams{Role: Assistant}) + require.NoError(t, err) + msg.AppendContent("payload") + require.NoError(t, svc.Update(t.Context(), msg)) + + select { + case <-slow.started: + case <-time.After(time.Second): + t.Fatal("timer-fired flush never reached UpdateMessage") + } + + flushDone := make(chan error, 1) + go func() { flushDone <- svc.FlushAll(t.Context()) }() + + select { + case err := <-flushDone: + t.Fatalf("FlushAll returned %v while in-flight write was still blocked", err) + case <-time.After(50 * time.Millisecond): + } + + close(slow.release) + select { + case err := <-flushDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("FlushAll did not return after in-flight write completed") + } + + got, err := svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.Equal(t, "payload", got.Content().Text) +} + +// TestUpdate_StructuralFlushUsesMustDeliver covers the second review +// finding: structural terminal events (tool-call add, tool-call +// finish, reasoning end) must publish via the must-deliver path even +// when the message itself is not yet IsFinished. +// +// We detect which path was taken by saturating a subscriber's channel +// buffer with no reader. With a short must-deliver timeout, the +// must-deliver path increments [pubsub.Broker.MustDeliverDropCount] +// after the timeout expires; the lossy path increments +// [pubsub.Broker.DropCount] immediately. The two counters are +// disjoint, so they precisely identify which call site published the +// event. +func TestUpdate_StructuralFlushUsesMustDeliver(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + mut func(*Message) + }{ + { + name: "tool call add", + mut: func(m *Message) { + m.AddToolCall(ToolCall{ID: "tc1", Name: "view"}) + }, + }, + { + name: "tool call finish", + mut: func(m *Message) { + m.AddToolCall(ToolCall{ID: "tc1", Name: "view", Input: "{}", Finished: true}) + }, + }, + { + name: "reasoning end", + mut: func(m *Message) { + m.AppendReasoningContent("hmm") + m.FinishThinking() + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + conn, err := db.Connect(t.Context(), t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + + q := db.New(conn) + sessions := session.NewService(q, conn) + sess, err := sessions.Create(t.Context(), "test") + require.NoError(t, err) + + // Replace the default broker with a tiny buffer + short + // must-deliver timeout so we can fully saturate from a + // single sender and observe drops without long waits. + svc := NewService(q, WithDebounce(time.Hour)) + impl := svc.(*service) + impl.Shutdown() + impl.Broker = pubsub.NewBrokerWithOptions[Message](1) + impl.SetMustDeliverTimeout(40 * time.Millisecond) + + subCtx, cancel := context.WithCancel(t.Context()) + defer cancel() + sub := svc.Subscribe(subCtx) + + msg, err := svc.Create(t.Context(), sess.ID, CreateMessageParams{Role: Assistant}) + require.NoError(t, err) + + // Saturate the subscriber's buffer (capacity 1). The + // CreatedEvent from Create above already left one event + // in the buffer; we never read sub, so the next publish + // has nowhere to go. + _ = sub // intentionally not drained. + + // Drive the structural change. With debounce=1h, Update + // flushes synchronously and routes through whichever + // publish path the service chose for structural events. + tc.mut(&msg) + require.NoError(t, svc.Update(t.Context(), msg)) + + // Must-deliver timeout (40ms) should have expired with + // no drain. If structural events are routed through + // must-deliver: MustDeliverDropCount > 0, DropCount + // unchanged. If routed through lossy Publish: + // DropCount > 0, MustDeliverDropCount == 0. + require.Eventually(t, func() bool { + return impl.MustDeliverDropCount() >= 1 + }, time.Second, 5*time.Millisecond, + "structural terminal event should publish via must-deliver, not lossy Publish") + require.Zero(t, impl.DropCount(), + "structural terminal event must not be silently dropped via lossy Publish") + }) + } +} + +// TestPendingStateDoesNotRetainPayload guards the fix for the +// unbounded heap growth reported in #3162: pending entries live for the +// life of the service, so once a message is fully flushed they must +// hold bookkeeping only, never the message parts. +func TestPendingStateDoesNotRetainPayload(t *testing.T) { + t.Parallel() + + svc, sessionID := newTestService(t, WithDebounce(time.Millisecond)) + s := svc.(*service) + + msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{ + Role: Assistant, + Parts: []ContentPart{}, + }) + require.NoError(t, err) + + msg.AppendReasoningContent(strings.Repeat("x", 64*1024)) + msg.AddToolCall(ToolCall{ID: "tc-1", Name: "bash", Finished: true}) + msg.AddFinish(FinishReasonEndTurn, "", "") + require.NoError(t, svc.Update(t.Context(), msg)) + require.NoError(t, svc.Flush(t.Context(), msg.ID)) + + s.mu.Lock() + defer s.mu.Unlock() + p := s.pending[msg.ID] + require.NotNil(t, p) + require.False(t, p.dirty) + require.Empty(t, p.latest.Parts, "flushed pending state still pins the message parts") + require.True(t, p.hasFlushed) + require.Equal(t, []bool{true}, p.baseline.toolCallsFinished) +} + +// TestShouldFlushNow_AgainstCompactBaseline pins terminal detection to the +// projection stored in flushBaseline rather than to a full Message snapshot. +// +// The existing structural tests all assert the positive direction, which a +// broken projection would still satisfy: an empty baseline yields a tool-call +// length mismatch, so it flushes for the wrong reason. The load-bearing case +// is the negative one — after a flush, a delta that changes nothing structural +// must stay buffered. That only holds if the baseline carries the prior +// finished flags and reasoning finish time accurately. +func TestShouldFlushNow_AgainstCompactBaseline(t *testing.T) { + t.Parallel() + + setup := func(t *testing.T) (Service, Message) { + t.Helper() + svc, sessionID := newTestService(t, WithDebounce(time.Hour)) + msg, err := svc.Create(t.Context(), sessionID, CreateMessageParams{Role: Assistant}) + require.NoError(t, err) + + // Land a baseline: one unfinished tool call, reasoning still open. + msg.AppendReasoningContent("thinking") + msg.AddToolCall(ToolCall{ID: "tc1", Name: "view", Finished: false}) + require.NoError(t, svc.Update(t.Context(), msg)) + got, err := svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.Len(t, got.ToolCalls(), 1) + require.False(t, got.ToolCalls()[0].Finished) + return svc, msg + } + + t.Run("non-structural delta stays buffered", func(t *testing.T) { + t.Parallel() + svc, msg := setup(t) + + // Same tool call, same finished flag, reasoning still open. A + // baseline that lost the projection would see a length or flag + // mismatch here and flush. + msg.AppendReasoningContent(" some more") + require.NoError(t, svc.Update(t.Context(), msg)) + + got, err := svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.Equal(t, "thinking", got.ReasoningContent().Thinking, + "a non-structural delta must not trigger a sync flush against the compact baseline") + }) + + t.Run("finished flip flushes", func(t *testing.T) { + t.Parallel() + svc, msg := setup(t) + + msg.AddToolCall(ToolCall{ID: "tc1", Name: "view", Input: "{}", Finished: true}) + require.NoError(t, svc.Update(t.Context(), msg)) + + got, err := svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.True(t, got.ToolCalls()[0].Finished) + }) + + t.Run("new tool call flushes", func(t *testing.T) { + t.Parallel() + svc, msg := setup(t) + + msg.AddToolCall(ToolCall{ID: "tc2", Name: "bash", Finished: false}) + require.NoError(t, svc.Update(t.Context(), msg)) + + got, err := svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.Len(t, got.ToolCalls(), 2) + }) + + t.Run("reasoning end flushes", func(t *testing.T) { + t.Parallel() + svc, msg := setup(t) + + msg.FinishThinking() + require.NoError(t, svc.Update(t.Context(), msg)) + + got, err := svc.Get(t.Context(), msg.ID) + require.NoError(t, err) + require.NotZero(t, got.ReasoningContent().FinishedAt) + }) +} diff --git a/internal/projects/projects.go b/internal/projects/projects.go new file mode 100644 index 0000000000000000000000000000000000000000..f909fc2cf6e3e13030746cc8804ebc4af2514ba5 --- /dev/null +++ b/internal/projects/projects.go @@ -0,0 +1,126 @@ +package projects + +import ( + "encoding/json" + "os" + "path/filepath" + "slices" + "sync" + "time" + + "github.com/charmbracelet/crush/internal/config" +) + +const projectsFileName = "projects.json" + +// Project represents a tracked project directory. +type Project struct { + Path string `json:"path"` + DataDir string `json:"data_dir"` + LastAccessed time.Time `json:"last_accessed"` +} + +// ProjectList holds the list of tracked projects. +type ProjectList struct { + Projects []Project `json:"projects"` +} + +var mu sync.Mutex + +// projectsFilePath returns the path to the projects.json file. +func projectsFilePath() string { + return filepath.Join(filepath.Dir(config.GlobalConfigData()), projectsFileName) +} + +// Load reads the projects list from disk. +func Load() (*ProjectList, error) { + mu.Lock() + defer mu.Unlock() + + path := projectsFilePath() + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return &ProjectList{Projects: []Project{}}, nil + } + return nil, err + } + + var list ProjectList + if err := json.Unmarshal(data, &list); err != nil { + return nil, err + } + + return &list, nil +} + +// Save writes the projects list to disk. +func Save(list *ProjectList) error { + mu.Lock() + defer mu.Unlock() + + path := projectsFilePath() + + // Ensure directory exists + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return err + } + + data, err := json.MarshalIndent(list, "", " ") + if err != nil { + return err + } + + return os.WriteFile(path, data, 0o600) +} + +// Register adds or updates a project in the list. +func Register(workingDir, dataDir string) error { + list, err := Load() + if err != nil { + return err + } + + now := time.Now().UTC() + + // Check if project already exists + found := false + for i, p := range list.Projects { + if p.Path == workingDir { + list.Projects[i].DataDir = dataDir + list.Projects[i].LastAccessed = now + found = true + break + } + } + + if !found { + list.Projects = append(list.Projects, Project{ + Path: workingDir, + DataDir: dataDir, + LastAccessed: now, + }) + } + + // Sort by last accessed (most recent first) + slices.SortFunc(list.Projects, func(a, b Project) int { + if a.LastAccessed.After(b.LastAccessed) { + return -1 + } + if a.LastAccessed.Before(b.LastAccessed) { + return 1 + } + return 0 + }) + + return Save(list) +} + +// List returns all tracked projects sorted by last accessed. +func List() ([]Project, error) { + list, err := Load() + if err != nil { + return nil, err + } + return list.Projects, nil +} diff --git a/internal/projects/projects_test.go b/internal/projects/projects_test.go new file mode 100644 index 0000000000000000000000000000000000000000..e41ffca74040648315a369f451140aec57bdfb40 --- /dev/null +++ b/internal/projects/projects_test.go @@ -0,0 +1,186 @@ +package projects + +import ( + "path/filepath" + "testing" + "time" +) + +func TestRegisterAndList(t *testing.T) { + // Create a temporary directory for the test + tmpDir := t.TempDir() + + // Override the projects file path for testing + t.Setenv("XDG_DATA_HOME", tmpDir) + t.Setenv("CRUSH_GLOBAL_DATA", filepath.Join(tmpDir, "crush")) + + // Test registering a project + err := Register("/home/user/project1", "/home/user/project1/.crush") + if err != nil { + t.Fatalf("Register failed: %v", err) + } + + // List projects + projects, err := List() + if err != nil { + t.Fatalf("List failed: %v", err) + } + + if len(projects) != 1 { + t.Fatalf("Expected 1 project, got %d", len(projects)) + } + + if projects[0].Path != "/home/user/project1" { + t.Errorf("Expected path /home/user/project1, got %s", projects[0].Path) + } + + if projects[0].DataDir != "/home/user/project1/.crush" { + t.Errorf("Expected data_dir /home/user/project1/.crush, got %s", projects[0].DataDir) + } + + // Register another project + err = Register("/home/user/project2", "/home/user/project2/.crush") + if err != nil { + t.Fatalf("Register failed: %v", err) + } + + projects, err = List() + if err != nil { + t.Fatalf("List failed: %v", err) + } + + if len(projects) != 2 { + t.Fatalf("Expected 2 projects, got %d", len(projects)) + } + + // Most recent should be first + if projects[0].Path != "/home/user/project2" { + t.Errorf("Expected most recent project first, got %s", projects[0].Path) + } +} + +func TestRegisterUpdatesExisting(t *testing.T) { + tmpDir := t.TempDir() + t.Setenv("XDG_DATA_HOME", tmpDir) + t.Setenv("CRUSH_GLOBAL_DATA", filepath.Join(tmpDir, "crush")) + + // Register a project + err := Register("/home/user/project1", "/home/user/project1/.crush") + if err != nil { + t.Fatalf("Register failed: %v", err) + } + + projects, _ := List() + firstAccess := projects[0].LastAccessed + + // Wait a bit and re-register + time.Sleep(10 * time.Millisecond) + + err = Register("/home/user/project1", "/home/user/project1/.crush-new") + if err != nil { + t.Fatalf("Register failed: %v", err) + } + + projects, _ = List() + + if len(projects) != 1 { + t.Fatalf("Expected 1 project after update, got %d", len(projects)) + } + + if projects[0].DataDir != "/home/user/project1/.crush-new" { + t.Errorf("Expected updated data_dir, got %s", projects[0].DataDir) + } + + if !projects[0].LastAccessed.After(firstAccess) { + t.Error("Expected LastAccessed to be updated") + } +} + +func TestLoadEmptyFile(t *testing.T) { + tmpDir := t.TempDir() + t.Setenv("XDG_DATA_HOME", tmpDir) + t.Setenv("CRUSH_GLOBAL_DATA", filepath.Join(tmpDir, "crush")) + + // List before any projects exist + projects, err := List() + if err != nil { + t.Fatalf("List failed: %v", err) + } + + if len(projects) != 0 { + t.Errorf("Expected 0 projects, got %d", len(projects)) + } +} + +func TestProjectsFilePath(t *testing.T) { + tmpDir := t.TempDir() + t.Setenv("XDG_DATA_HOME", tmpDir) + t.Setenv("CRUSH_GLOBAL_DATA", filepath.Join(tmpDir, "crush")) + + expected := filepath.Join(tmpDir, "crush", "projects.json") + actual := projectsFilePath() + + if actual != expected { + t.Errorf("Expected %s, got %s", expected, actual) + } +} + +func TestRegisterWithParentDataDir(t *testing.T) { + tmpDir := t.TempDir() + t.Setenv("XDG_DATA_HOME", tmpDir) + t.Setenv("CRUSH_GLOBAL_DATA", filepath.Join(tmpDir, "crush")) + + // Register a project where .crush is in a parent directory. + // e.g., working in /home/user/monorepo/packages/app but .crush is at /home/user/monorepo/.crush + err := Register("/home/user/monorepo/packages/app", "/home/user/monorepo/.crush") + if err != nil { + t.Fatalf("Register failed: %v", err) + } + + projects, err := List() + if err != nil { + t.Fatalf("List failed: %v", err) + } + + if len(projects) != 1 { + t.Fatalf("Expected 1 project, got %d", len(projects)) + } + + if projects[0].Path != "/home/user/monorepo/packages/app" { + t.Errorf("Expected path /home/user/monorepo/packages/app, got %s", projects[0].Path) + } + + if projects[0].DataDir != "/home/user/monorepo/.crush" { + t.Errorf("Expected data_dir /home/user/monorepo/.crush, got %s", projects[0].DataDir) + } +} + +func TestRegisterWithExternalDataDir(t *testing.T) { + tmpDir := t.TempDir() + t.Setenv("XDG_DATA_HOME", tmpDir) + t.Setenv("CRUSH_GLOBAL_DATA", filepath.Join(tmpDir, "crush")) + + // Register a project where .crush is in a completely different location. + // e.g., project at /home/user/project but data stored at /var/data/crush/myproject + err := Register("/home/user/project", "/var/data/crush/myproject") + if err != nil { + t.Fatalf("Register failed: %v", err) + } + + projects, err := List() + if err != nil { + t.Fatalf("List failed: %v", err) + } + + if len(projects) != 1 { + t.Fatalf("Expected 1 project, got %d", len(projects)) + } + + if projects[0].Path != "/home/user/project" { + t.Errorf("Expected path /home/user/project, got %s", projects[0].Path) + } + + if projects[0].DataDir != "/var/data/crush/myproject" { + t.Errorf("Expected data_dir /var/data/crush/myproject, got %s", projects[0].DataDir) + } +} diff --git a/internal/server/agent_cancel_test.go b/internal/server/agent_cancel_test.go new file mode 100644 index 0000000000000000000000000000000000000000..ea24bfc94b88d27f5513e7a6b3a3b7e69b4a86ac --- /dev/null +++ b/internal/server/agent_cancel_test.go @@ -0,0 +1,227 @@ +package server + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + "charm.land/fantasy" + "github.com/charmbracelet/crush/internal/agent" + "github.com/charmbracelet/crush/internal/app" + "github.com/charmbracelet/crush/internal/backend" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/proto" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +// runCoordinator is a configurable agent.Coordinator stub for the +// cancel/drop tests. Run blocks until either ctx is canceled (so it +// can observe explicit Cancel paths) or release fires (so the test +// can let a "still running" turn finish on its own). The most recent +// ctx and the error returned to the caller are recorded for +// assertions. +type runCoordinator struct { + release chan struct{} + returnFn func(ctx context.Context) error + + mu sync.Mutex + gotCtx context.Context + ranCount atomic.Int32 + entered chan struct{} // closed exactly once when Run is first entered. + enteredOne sync.Once +} + +func newRunCoordinator(returnFn func(ctx context.Context) error) *runCoordinator { + return &runCoordinator{ + release: make(chan struct{}), + returnFn: returnFn, + entered: make(chan struct{}), + } +} + +func (s *runCoordinator) Run(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + s.mu.Lock() + s.gotCtx = ctx + s.mu.Unlock() + s.ranCount.Add(1) + s.enteredOne.Do(func() { close(s.entered) }) + select { + case <-s.release: + case <-ctx.Done(): + // Only fires if the run is actually cancellable. + } + return nil, s.returnFn(ctx) +} + +func (s *runCoordinator) RunAccepted(ctx context.Context, accept *agent.AcceptedRun, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + return s.Run(ctx, sessionID, prompt, attachments...) +} + +func (s *runCoordinator) BeginAccepted(sessionID string) *agent.AcceptedRun { + return nil +} +func (s *runCoordinator) Cancel(string) {} +func (s *runCoordinator) CancelAll() {} +func (s *runCoordinator) IsBusy() bool { return false } +func (s *runCoordinator) IsSessionBusy(string) bool { + return false +} +func (s *runCoordinator) QueuedPrompts(string) int { return 0 } +func (s *runCoordinator) QueuedPromptsList(string) []string { return nil } +func (s *runCoordinator) ClearQueue(string) {} +func (s *runCoordinator) Summarize(context.Context, string) error { + return nil +} +func (s *runCoordinator) Model() agent.Model { return agent.Model{} } +func (s *runCoordinator) UpdateModels(context.Context) error { return nil } +func (s *runCoordinator) GenerateTitle(context.Context, string, string) {} + +func (s *runCoordinator) capturedCtx() context.Context { + s.mu.Lock() + defer s.mu.Unlock() + return s.gotCtx +} + +// buildAgentWorkspace returns a controller wired to a backend whose +// single workspace exposes the given coordinator. The workspace +// shutdown hook is overridden to avoid driving a real [app.App] +// through teardown when the test exits. +func buildAgentWorkspace(t *testing.T, coord agent.Coordinator) (*controllerV1, string) { + t.Helper() + b := backend.New(context.Background(), nil, nil) + a := &app.App{AgentCoordinator: coord} + + ws := &backend.Workspace{ + ID: uuid.New().String(), + Path: t.TempDir(), + App: a, + } + backend.InsertWorkspaceForTest(b, ws) + backend.SetWorkspaceShutdownFnForTest(ws, func() {}) + + s := &Server{backend: b} + return &controllerV1{backend: b, server: s}, ws.ID +} + +func postAgent(t *testing.T, c *controllerV1, ctx context.Context, wsID, sessionID string) *httptest.ResponseRecorder { + t.Helper() + body, err := json.Marshal(proto.AgentMessage{SessionID: sessionID, Prompt: "hi"}) + require.NoError(t, err) + req := httptest.NewRequestWithContext(ctx, http.MethodPost, "/v1/workspaces/"+wsID+"/agent", bytes.NewReader(body)) + req.SetPathValue("id", wsID) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + c.handlePostWorkspaceAgent(rec, req) + return rec +} + +// TestPostAgent_ReturnsOKOnContextCanceled verifies that when another +// client cancels the session mid-turn, the prompting client's POST is +// unaffected: SendMessage is fire-and-forget, so the handler returns +// 200 immediately without waiting for the turn. A run that later +// returns context.Canceled never surfaces as a 500 to the prompter; +// the FinishReasonCanceled marker reaches SSE subscribers via the +// assistant message instead. +func TestPostAgent_ReturnsOKOnContextCanceled(t *testing.T) { + t.Parallel() + + coord := newRunCoordinator(func(context.Context) error { + return context.Canceled + }) + c, wsID := buildAgentWorkspace(t, coord) + + // The handler returns immediately, before the dispatched run is + // released, because the run no longer owns the HTTP response. + rec := postAgent(t, c, t.Context(), wsID, "S1") + require.Equal(t, http.StatusAccepted, rec.Code, "fire-and-forget SendMessage must return 202 without waiting for the run") + + // The run is dispatched on a goroutine; let it return + // context.Canceled. Nothing from that path reaches the (already + // returned) handler. + select { + case <-coord.entered: + case <-time.After(2 * time.Second): + t.Fatal("dispatched run was never entered") + } + close(coord.release) + + // Wait for the dispatched run to fully return. Backend.runAgent + // swallows context.Canceled, so it must not publish a + // notify.TypeAgentError. Publishing would dereference the synthetic + // workspace's nil notification broker and crash this goroutine, + // which is the explicit guard that a cancel produces no top-level + // error event. + require.Eventually(t, func() bool { + return coord.ranCount.Load() == 1 + }, 2*time.Second, 10*time.Millisecond) +} + +// TestHandleError_ContextCanceledFallsThroughTo500 documents the step 8 +// cleanup: the old context.Canceled special case in handleError was +// removed because runtime cancellation of an agent run can no longer +// reach handleError. The agent-prompt handler returns 202 before the run +// starts (fire-and-forget SendMessage) and Backend.runAgent swallows +// context.Canceled. Any context.Canceled that still reaches handleError +// is therefore an unexpected synchronous error and falls through to the +// default 500 like any other. +func TestHandleError_ContextCanceledFallsThroughTo500(t *testing.T) { + t.Parallel() + + c := &controllerV1{server: &Server{}} + rec := httptest.NewRecorder() + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/", nil) + + c.handleError(rec, req, context.Canceled) + + require.Equal(t, http.StatusInternalServerError, rec.Code) +} + +// TestPostAgent_DetachesRequestContext verifies that the dispatched run +// is bound to the workspace context, not the prompting client's HTTP +// request context. Canceling the request context must neither cancel +// the run nor be observed by the coordinator. +func TestPostAgent_DetachesRequestContext(t *testing.T) { + t.Parallel() + + coord := newRunCoordinator(func(context.Context) error { + return nil + }) + c, wsID := buildAgentWorkspace(t, coord) + + reqCtx, cancelReq := context.WithCancel(context.Background()) + + // The handler returns immediately; the run keeps executing on its + // own goroutine bound to the workspace context. + rec := postAgent(t, c, reqCtx, wsID, "S1") + require.Equal(t, http.StatusAccepted, rec.Code) + + select { + case <-coord.entered: + case <-time.After(2 * time.Second): + t.Fatal("dispatched run was never entered") + } + + // Drop the prompting client. This must not reach the run. + cancelReq() + + got := coord.capturedCtx() + require.NotNil(t, got) + // Compare by identity (pointer), not reflect.DeepEqual: deep + // comparison would traverse context internals that the runtime + // mutates concurrently. + require.False(t, got == reqCtx, "run ctx must not be the request ctx") + require.NoError(t, got.Err(), "run ctx must not inherit cancellation from the dropped request") + + // Release the run so it returns cleanly. + close(coord.release) + require.Eventually(t, func() bool { + return coord.ranCount.Load() == 1 + }, 2*time.Second, 10*time.Millisecond) +} diff --git a/internal/server/config.go b/internal/server/config.go new file mode 100644 index 0000000000000000000000000000000000000000..750a9ee89849a0849ed13cc10247f028efa23f35 --- /dev/null +++ b/internal/server/config.go @@ -0,0 +1,386 @@ +package server + +import ( + "encoding/json" + "net/http" + + "github.com/charmbracelet/crush/internal/proto" +) + +// handlePostWorkspaceConfigSet sets a configuration field. +func (c *controllerV1) handlePostWorkspaceConfigSet(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.ConfigSetRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + if err := c.backend.SetConfigField(id, req.Scope, req.Key, req.Value); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handlePostWorkspaceConfigRemove removes a configuration field. +func (c *controllerV1) handlePostWorkspaceConfigRemove(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.ConfigRemoveRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + if err := c.backend.RemoveConfigField(id, req.Scope, req.Key); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handlePostWorkspaceConfigModel updates the preferred model. +func (c *controllerV1) handlePostWorkspaceConfigModel(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.ConfigModelRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + if err := c.backend.UpdatePreferredModel(id, req.Scope, req.ModelType, req.Model); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handlePostWorkspaceConfigCompact sets compact mode. +func (c *controllerV1) handlePostWorkspaceConfigCompact(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.ConfigCompactRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + if err := c.backend.SetCompactMode(id, req.Scope, req.Enabled); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handlePostWorkspaceConfigProviderKey sets a provider API key. +func (c *controllerV1) handlePostWorkspaceConfigProviderKey(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.ConfigProviderKeyRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + apiKey, err := req.DecodeAPIKey() + if err != nil { + c.server.logError(r, "Failed to decode api key", "error", err, "kind", req.Kind) + jsonError(w, http.StatusBadRequest, err.Error()) + return + } + + if err := c.backend.SetProviderAPIKey(id, req.Scope, req.ProviderID, apiKey); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handlePostWorkspaceConfigImportCopilot imports Copilot credentials. +func (c *controllerV1) handlePostWorkspaceConfigImportCopilot(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + token, ok, err := c.backend.ImportCopilot(id) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, proto.ImportCopilotResponse{Token: token, Success: ok}) +} + +// handlePostWorkspaceConfigRefreshOAuth refreshes an OAuth token for a provider. +func (c *controllerV1) handlePostWorkspaceConfigRefreshOAuth(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.ConfigRefreshOAuthRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + if err := c.backend.RefreshOAuthToken(r.Context(), id, req.Scope, req.ProviderID); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handleGetWorkspaceProjectNeedsInit reports whether a project needs initialization. +func (c *controllerV1) handleGetWorkspaceProjectNeedsInit(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + needs, err := c.backend.ProjectNeedsInitialization(id) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, proto.ProjectNeedsInitResponse{NeedsInit: needs}) +} + +// handlePostWorkspaceProjectInit marks the project as initialized. +func (c *controllerV1) handlePostWorkspaceProjectInit(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + if err := c.backend.MarkProjectInitialized(id); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handleGetWorkspaceProjectInitPrompt returns the project initialization prompt. +func (c *controllerV1) handleGetWorkspaceProjectInitPrompt(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + prompt, err := c.backend.InitializePrompt(id) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, proto.ProjectInitPromptResponse{Prompt: prompt}) +} + +// handleGetWorkspaceSkills returns the effective visible skills for a workspace. +func (c *controllerV1) handleGetWorkspaceSkills(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + skills, err := c.backend.ListSkills(id) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, skills) +} + +// handlePostWorkspaceSkillRead reads a skill's content by ID. +func (c *controllerV1) handlePostWorkspaceSkillRead(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.ReadSkillRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + content, result, err := c.backend.ReadSkill(r.Context(), id, req.SkillID) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, proto.ReadSkillResponse{Content: content, Result: result}) +} + +// handlePostWorkspaceMCPEnableDocker enables the Docker MCP server. +func (c *controllerV1) handlePostWorkspaceMCPEnableDocker(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + if err := c.backend.EnableDockerMCP(r.Context(), id); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handlePostWorkspaceMCPDisableDocker disables the Docker MCP server. +func (c *controllerV1) handlePostWorkspaceMCPDisableDocker(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + if err := c.backend.DisableDockerMCP(id); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handlePostWorkspaceMCPRefreshTools refreshes tools for a named MCP server. +func (c *controllerV1) handlePostWorkspaceMCPRefreshTools(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.MCPNameRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + if err := c.backend.RefreshMCPTools(r.Context(), id, req.Name); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handlePostWorkspaceMCPReadResource reads a resource from an MCP server. +func (c *controllerV1) handlePostWorkspaceMCPReadResource(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.MCPReadResourceRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + contents, err := c.backend.ReadMCPResource(r.Context(), id, req.Name, req.URI) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, contents) +} + +// handleGetWorkspaceMCPPrompts returns the available MCP prompts for a workspace. +func (c *controllerV1) handleGetWorkspaceMCPPrompts(w http.ResponseWriter, r *http.Request) { + prompts, err := c.backend.ListMCPPrompts(r.PathValue("id")) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, prompts) +} + +// handlePostWorkspaceMCPGetPrompt retrieves a prompt from an MCP server. +func (c *controllerV1) handlePostWorkspaceMCPGetPrompt(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.MCPGetPromptRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + prompt, err := c.backend.GetMCPPrompt(id, req.ClientID, req.PromptID, req.Args) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, proto.MCPGetPromptResponse{Prompt: prompt}) +} + +// handleGetWorkspaceMCPStates returns the state of all MCP clients. +func (c *controllerV1) handleGetWorkspaceMCPStates(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + states := c.backend.MCPGetStates(id) + result := make(map[string]proto.MCPClientInfo, len(states)) + for k, v := range states { + result[k] = proto.MCPClientInfo{ + Name: v.Name, + State: proto.MCPState(v.State), + Error: v.Error, + ToolCount: v.Counts.Tools, + PromptCount: v.Counts.Prompts, + ResourceCount: v.Counts.Resources, + ConnectedAt: v.ConnectedAt, + } + } + jsonEncode(w, result) +} + +// handleGetWorkspaceMCPPendingAuth returns the MCP servers awaiting OAuth +// authentication for a workspace. +func (c *controllerV1) handleGetWorkspaceMCPPendingAuth(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + pending, err := c.backend.MCPPendingAuth(id) + if err != nil { + c.handleError(w, r, err) + return + } + result := make([]proto.MCPPendingAuthServer, len(pending)) + for i, p := range pending { + result[i] = proto.MCPPendingAuthServer{Name: p.Name, URL: p.URL} + } + jsonEncode(w, result) +} + +// handleGetWorkspaceMCPAuthURL returns the current OAuth authorization URL +// for a named MCP server, if a flow is in progress. +func (c *controllerV1) handleGetWorkspaceMCPAuthURL(w http.ResponseWriter, r *http.Request) { + name := r.URL.Query().Get("name") + if name == "" { + jsonError(w, http.StatusBadRequest, "name is required") + return + } + jsonEncode(w, proto.MCPAuthResponse{AuthURL: c.backend.MCPAuthURL(name)}) +} + +// handlePostWorkspaceMCPAuth runs the OAuth flow for a named MCP server. +// The local browser is suppressed on the server; the client polls +// pending-auth / auth-url to surface the authorization URL on the user's +// machine. The call blocks until the flow completes or the request context +// is cancelled. +func (c *controllerV1) handlePostWorkspaceMCPAuth(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.MCPNameRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + if err := c.backend.MCPAuthenticate(r.Context(), id, req.Name); err != nil { + // If the client went away the request context was cancelled; + // the error is still surfaced for logging but no response can + // be written. + c.handleError(w, r, err) + return + } + // The flow has finished by the time this returns, so there is no + // in-progress authorization URL to report; the client polls + // /mcp/auth-url for that while the flow runs. + jsonEncode(w, proto.MCPAuthResponse{}) +} + +// handlePostWorkspaceMCPRefreshPrompts refreshes prompts for a named MCP server. +func (c *controllerV1) handlePostWorkspaceMCPRefreshPrompts(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.MCPNameRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + c.backend.MCPRefreshPrompts(r.Context(), id, req.Name) + w.WriteHeader(http.StatusOK) +} + +// handlePostWorkspaceMCPRefreshResources refreshes resources for a named MCP server. +func (c *controllerV1) handlePostWorkspaceMCPRefreshResources(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.MCPNameRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + c.backend.MCPRefreshResources(r.Context(), id, req.Name) + w.WriteHeader(http.StatusOK) +} diff --git a/internal/server/e2e_test.go b/internal/server/e2e_test.go new file mode 100644 index 0000000000000000000000000000000000000000..f538632217d41a64d6abed70c44da58ea790e8c5 --- /dev/null +++ b/internal/server/e2e_test.go @@ -0,0 +1,632 @@ +package server + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "net/url" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/charmbracelet/crush/internal/app" + "github.com/charmbracelet/crush/internal/backend" + "github.com/charmbracelet/crush/internal/db" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/permission" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +// e2eHarness wires a Server, its Backend (with a custom shutdownFn we +// can observe), an httptest.NewServer, and a synthetic Workspace whose +// embedded App has a live event broker. It is the minimum scaffolding +// the multi-client end-to-end scenarios in PLAN item 6 need. +type e2eHarness struct { + httpSrv *httptest.Server + srv *Server + backend *backend.Backend + workspace *backend.Workspace + app *app.App + shutdownHit atomic.Bool + + // sseWG tracks every SSE reader goroutine spawned by + // [e2eHarness.subscribeSSE]. The harness's cleanup hook waits on + // it after the httptest server has been closed so that the test + // cannot leave behind background readers (and therefore unclosed + // response bodies) after returning. + sseWG sync.WaitGroup +} + +// installServer attaches a fresh Server (with a custom shutdown +// callback that flips [e2eHarness.shutdownHit]) wrapped in an +// [httptest.Server] onto h. It registers the cleanup hooks for the +// httptest server and the SSE reader WaitGroup in the order required +// by the LIFO contract documented on [newE2EHarness]. +// +// Callers that want a fully synthetic workspace use [newE2EHarness]; +// callers that want to drive the real CreateWorkspace HTTP path use +// [newRealCreateHarness] and then [e2eHarness.postWorkspace]. +func (h *e2eHarness) installServer(t *testing.T) { + t.Helper() + srv := &Server{} + srv.backend = backend.New(context.Background(), nil, func() { + h.shutdownHit.Store(true) + }) + srv.installHandler() + + hs := httptest.NewServer(srv.Handler()) + // Order matters: t.Cleanup is LIFO and the test's own per- + // stream cancels (cancelA/cancelB) run first. After those, we + // want hs.Close to fire first (so any handler still parked in + // its `select` returns), THEN sseWG.Wait so every reader + // goroutine exits and closes its response body. Any caller- + // owned cleanups registered *before* installServer (e.g. App + // teardown for the synthetic harness) therefore run LAST, + // after the readers have drained. + t.Cleanup(h.sseWG.Wait) + t.Cleanup(hs.Close) + + h.httpSrv = hs + h.srv = srv + h.backend = srv.backend +} + +// newE2EHarness builds an in-process server + a synthetic Workspace +// whose embedded App is a real [app.App] constructed via +// [app.NewForTest], so its event broker delivers everything the SSE +// pipeline expects. Used by the scenarios that do not need to +// exercise the path-dedupe behavior of [backend.CreateWorkspace]. +// +// Cleanup tears down the App's broker only after sseWG.Wait and +// hs.Close have run, so SSE readers cannot observe a dead broker. +func newE2EHarness(t *testing.T) *e2eHarness { + t.Helper() + + h := &e2eHarness{} + + // Register the App teardown FIRST so LIFO order puts it AFTER + // the cleanups that installServer registers below (hs.Close + + // sseWG.Wait). + appCtx, cancel := context.WithCancel(context.Background()) + a := app.NewForTest(appCtx) + t.Cleanup(func() { + cancel() + a.ShutdownForTest() + }) + + h.installServer(t) + + ws := &backend.Workspace{ + ID: uuid.New().String(), + Path: t.TempDir(), + App: a, + } + // Synthetic workspaces have an incomplete App; bypass the + // default teardown so the "last workspace removed" path can run + // without panicking inside [app.App.Shutdown]. + backend.SetWorkspaceShutdownFnForTest(ws, func() {}) + backend.InsertWorkspaceForTest(h.backend, ws) + + h.workspace = ws + h.app = a + return h +} + +// newRealCreateHarness builds an in-process server WITHOUT any +// pre-inserted workspace, intended for tests that drive the real +// [backend.CreateWorkspace] HTTP path (path-dedupe scenario). It +// isolates HOME/XDG_* via [t.Setenv] so [config.Init] doesn't read +// the host machine's config, which means callers MUST NOT mark the +// test as parallel. +func newRealCreateHarness(t *testing.T) *e2eHarness { + t.Helper() + t.Setenv("HOME", t.TempDir()) + t.Setenv("XDG_CACHE_HOME", t.TempDir()) + t.Setenv("XDG_CONFIG_HOME", t.TempDir()) + t.Setenv("XDG_DATA_HOME", t.TempDir()) + + h := &e2eHarness{} + h.installServer(t) + return h +} + +// postWorkspace drives the real POST /v1/workspaces handler and +// returns the resolved workspace proto. This is how scenario 1 +// exercises the path-dedupe behavior from PLAN item 1: two calls +// with the same Path and distinct ClientIDs must return the same +// workspace ID. +func (h *e2eHarness) postWorkspace(t *testing.T, args proto.Workspace) proto.Workspace { + t.Helper() + body, err := json.Marshal(args) + require.NoError(t, err) + req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, + h.httpSrv.URL+"/v1/workspaces", bytes.NewReader(body)) + require.NoError(t, err) + req.Header.Set("Content-Type", "application/json") + resp, err := h.httpSrv.Client().Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode, "POST /v1/workspaces must succeed") + var out proto.Workspace + require.NoError(t, json.NewDecoder(resp.Body).Decode(&out)) + require.NotEmpty(t, out.ID, "server must return a workspace id") + return out +} + +// subscribeSSE opens an SSE stream against the test server for the +// given workspace and client ID. It returns a channel of decoded +// envelopes plus a cancel function that closes the stream. The +// returned channel is closed when the stream ends. +func (h *e2eHarness) subscribeSSE(t *testing.T, ctx context.Context, workspaceID, clientID string) (<-chan any, context.CancelFunc) { + t.Helper() + streamCtx, cancel := context.WithCancel(ctx) + + q := url.Values{"client_id": []string{clientID}} + reqURL := h.httpSrv.URL + "/v1/workspaces/" + workspaceID + "/events?" + q.Encode() + req, err := http.NewRequestWithContext(streamCtx, http.MethodGet, reqURL, nil) + require.NoError(t, err) + req.Header.Set("Accept", "text/event-stream") + + resp, err := h.httpSrv.Client().Do(req) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode, "SSE subscribe should return 200") + + out := make(chan any, 64) + h.sseWG.Go(func() { + defer resp.Body.Close() + defer close(out) + reader := bufio.NewReader(resp.Body) + for { + line, err := reader.ReadBytes('\n') + if err != nil { + return + } + line = bytes.TrimSpace(line) + if len(line) == 0 { + continue + } + data, ok := bytes.CutPrefix(line, []byte("data:")) + if !ok { + continue + } + data = bytes.TrimSpace(data) + var p pubsub.Payload + if err := json.Unmarshal(data, &p); err != nil { + continue + } + ev, decoded := decodeSSEEnvelope(p) + if !decoded { + continue + } + select { + case out <- ev: + case <-streamCtx.Done(): + return + } + } + }) + return out, cancel +} + +// decodeSSEEnvelope decodes the discriminated SSE envelope into the +// concrete pubsub.Event[proto.X] payload the e2e tests care about. +// Unknown payload types are skipped so tests can match on type +// assertions without worrying about envelope noise. +func decodeSSEEnvelope(p pubsub.Payload) (any, bool) { + switch p.Type { + case pubsub.PayloadTypePermissionRequest: + var e pubsub.Event[proto.PermissionRequest] + if err := json.Unmarshal(p.Payload, &e); err != nil { + return nil, false + } + return e, true + case pubsub.PayloadTypePermissionNotification: + var e pubsub.Event[proto.PermissionNotification] + if err := json.Unmarshal(p.Payload, &e); err != nil { + return nil, false + } + return e, true + case pubsub.PayloadTypeQuestionNotification: + var e pubsub.Event[proto.QuestionNotification] + if err := json.Unmarshal(p.Payload, &e); err != nil { + return nil, false + } + return e, true + case pubsub.PayloadTypeMessage: + var e pubsub.Event[proto.Message] + if err := json.Unmarshal(p.Payload, &e); err != nil { + return nil, false + } + return e, true + case pubsub.PayloadTypeAgentEvent: + var e pubsub.Event[proto.AgentEvent] + if err := json.Unmarshal(p.Payload, &e); err != nil { + return nil, false + } + return e, true + case pubsub.PayloadTypeRunComplete: + var e pubsub.Event[proto.RunComplete] + if err := json.Unmarshal(p.Payload, &e); err != nil { + return nil, false + } + return e, true + } + return nil, false +} + +// grantPermission posts a permission grant via the HTTP surface and +// returns the server's "resolved" verdict. Mirrors the client-side +// GrantPermission flow without importing internal/client (which +// would create an import cycle from this in-package test). +func (h *e2eHarness) grantPermission(t *testing.T, ctx context.Context, workspaceID string, req proto.PermissionGrant) bool { + t.Helper() + body, err := json.Marshal(req) + require.NoError(t, err) + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, + h.httpSrv.URL+"/v1/workspaces/"+workspaceID+"/permissions/grant", + bytes.NewReader(body)) + require.NoError(t, err) + httpReq.Header.Set("Content-Type", "application/json") + resp, err := h.httpSrv.Client().Do(httpReq) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + var out proto.PermissionGrantResponse + require.NoError(t, json.NewDecoder(resp.Body).Decode(&out)) + return out.Resolved +} + +// waitForAttached spins until the workspace's clients map reports at +// least n entries with streams > 0. Catches the race where a test +// publishes events before the server-side AttachClient has completed. +func (h *e2eHarness) waitForAttached(t *testing.T, n int) { + t.Helper() + h.waitForAttachedOn(t, h.workspace, n) +} + +// waitForAttachedOn is the workspace-explicit form of waitForAttached. +// Tests that drive a workspace whose pointer is not stored on the +// harness (e.g. the real CreateWorkspace path) pass the workspace in. +func (h *e2eHarness) waitForAttachedOn(t *testing.T, ws *backend.Workspace, n int) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if backend.WorkspaceLiveStreamCountForTest(ws) >= n { + return + } + time.Sleep(5 * time.Millisecond) + } + t.Fatalf("expected %d attached streams, have %d", n, + backend.WorkspaceLiveStreamCountForTest(ws)) +} + +// drainUntil reads from evc until it sees an event of type T that +// satisfies match, or ctx expires. Returns the matching event and +// ok=true, or the zero value and ok=false on timeout. +func drainUntil[T any](ctx context.Context, evc <-chan any, match func(T) bool) (T, bool) { + var zero T + for { + select { + case <-ctx.Done(): + return zero, false + case ev, ok := <-evc: + if !ok { + return zero, false + } + typed, isT := ev.(T) + if !isT { + continue + } + if match == nil || match(typed) { + return typed, true + } + } + } +} + +// TestE2E_TwoClientsReceiveSameMessage covers PLAN item 6 scenario 1: +// two clients POST /v1/workspaces with the same Path and observe +// that the server returns a single workspace (path-dedupe from PLAN +// item 1) and that an event published on that workspace fans out to +// both SSE streams. +// +// Cannot run in parallel: it isolates HOME/XDG_* via t.Setenv so +// config.Init does not read the host machine's real config. +func TestE2E_TwoClientsReceiveSameMessage(t *testing.T) { + h := newRealCreateHarness(t) + // Shorten the lifecycle windows so the workspace's client holds + // release quickly during test cleanup once both SSE streams have + // been detached: the create grace covers the window before the + // streams attach, the detach grace the window after they drop. + // The pooled DB connection is only released once the last hold + // expires, and Windows cannot remove the temp data directory + // while that connection is still open. + h.backend.SetCreateGrace(200 * time.Millisecond) + h.backend.SetDetachGrace(200 * time.Millisecond) + + ctx, cancel := context.WithCancel(t.Context()) + t.Cleanup(cancel) + + cidA := uuid.New().String() + cidB := uuid.New().String() + + // Shared workspace path. Two POSTs with this path must + // deduplicate at the backend's pathIndex and return the same + // workspace id. + wsPath := t.TempDir() + dataDir := t.TempDir() + args := proto.Workspace{Path: wsPath, DataDir: dataDir} + + argsA := args + argsA.ClientID = cidA + wsRespA := h.postWorkspace(t, argsA) + + argsB := args + argsB.ClientID = cidB + wsRespB := h.postWorkspace(t, argsB) + + require.Equal(t, wsRespA.ID, wsRespB.ID, + "POST /v1/workspaces with the same Path must return the same workspace id") + + // Look up the resulting workspace on the backend so the test + // can publish events through its real [app.App] event broker. + ws, err := h.backend.GetWorkspace(wsRespA.ID) + require.NoError(t, err) + // Override the shutdown callback so test cleanup doesn't run + // the full app.Shutdown path (which would tear down LSP/MCP + // resources the test doesn't need to exercise), but still + // release the pooled DB connection so Windows can clean up + // the temp data directory. + wsDataDir := ws.Cfg.Config().Options.DataDirectory + backend.SetWorkspaceShutdownFnForTest(ws, func() { + _ = db.Release(wsDataDir) + }) + + evcA, cancelA := h.subscribeSSE(t, ctx, ws.ID, cidA) + t.Cleanup(cancelA) + evcB, cancelB := h.subscribeSSE(t, ctx, ws.ID, cidB) + t.Cleanup(cancelB) + + h.waitForAttachedOn(t, ws, 2) + + const sessionID = "s-e2e-1" + msg := message.Message{ + ID: "m-1", + SessionID: sessionID, + Role: message.Assistant, + Parts: []message.ContentPart{message.TextContent{Text: "hello multi-client"}}, + } + ws.SendEvent(pubsub.Event[message.Message]{ + Type: pubsub.CreatedEvent, + Payload: msg, + }) + + pickCtx, pickCancel := context.WithTimeout(ctx, 3*time.Second) + defer pickCancel() + gotA, okA := drainUntil(pickCtx, evcA, func(e pubsub.Event[proto.Message]) bool { + return e.Payload.ID == "m-1" + }) + require.True(t, okA, "client A must receive the MessageEvent") + require.Equal(t, sessionID, gotA.Payload.SessionID) + + gotB, okB := drainUntil(pickCtx, evcB, func(e pubsub.Event[proto.Message]) bool { + return e.Payload.ID == "m-1" + }) + require.True(t, okB, "client B must receive the same MessageEvent") + require.Equal(t, sessionID, gotB.Payload.SessionID) +} + +// TestE2E_PermissionFlowCrossClient covers PLAN item 6 scenario 2: +// a tool-driven permission request is granted by client A; client B +// observes a PermissionNotification; a redundant grant from B +// returns the "already resolved" indicator (resolved=false from the +// bool plumbing landed in item 3). +func TestE2E_PermissionFlowCrossClient(t *testing.T) { + t.Parallel() + h := newE2EHarness(t) + ctx, cancel := context.WithCancel(t.Context()) + t.Cleanup(cancel) + + cidA := uuid.New().String() + cidB := uuid.New().String() + + evcA, cancelA := h.subscribeSSE(t, ctx, h.workspace.ID, cidA) + t.Cleanup(cancelA) + evcB, cancelB := h.subscribeSSE(t, ctx, h.workspace.ID, cidB) + t.Cleanup(cancelB) + + h.waitForAttached(t, 2) + + // Drive the permission request from a goroutine simulating the + // tool path. Request blocks until resolved; capture the outcome. + const sessionID = "s-perm" + const toolCallID = "tc-1" + type result struct { + granted bool + err error + } + done := make(chan result, 1) + go func() { + granted, err := h.app.Permissions.Request(ctx, permission.CreatePermissionRequest{ + SessionID: sessionID, + ToolCallID: toolCallID, + ToolName: "view", + Description: "read a file", + Action: "read", + Path: h.workspace.Path, + }) + done <- result{granted: granted, err: err} + }() + + // Wait for the PermissionRequest to arrive on client A's SSE + // stream. We need its ID to drive the grant. + // 10s instead of the 3s used elsewhere because this context + // bounds three sequential waits, not one. + pickCtx, pickCancel := context.WithTimeout(ctx, 10*time.Second) + defer pickCancel() + reqEv, ok := drainUntil(pickCtx, evcA, func(e pubsub.Event[proto.PermissionRequest]) bool { + return e.Payload.ToolCallID == toolCallID + }) + require.True(t, ok, "client A must receive the PermissionRequest") + + // Client A grants — first grant must report resolved=true. + resolvedA := h.grantPermission(t, ctx, h.workspace.ID, proto.PermissionGrant{ + Permission: reqEv.Payload, + Action: proto.PermissionAllow, + }) + require.True(t, resolvedA, "client A's grant must resolve the pending request") + + // The blocked Request call must now return granted=true. + select { + case r := <-done: + require.NoError(t, r.err) + require.True(t, r.granted) + case <-pickCtx.Done(): + t.Fatal("permission Request did not return after grant") + } + + // Client B must receive a PermissionNotification with + // Granted=true for the same ToolCallID. The initial neither- + // granted-nor-denied notification published at the start of + // Request also lands on B's stream — match on the granted one. + notif, ok := drainUntil(pickCtx, evcB, func(e pubsub.Event[proto.PermissionNotification]) bool { + return e.Payload.ToolCallID == toolCallID && e.Payload.Granted + }) + require.True(t, ok, "client B must receive a granting PermissionNotification") + require.True(t, notif.Payload.Granted) + require.False(t, notif.Payload.Denied) + + // A follow-up grant from client B must report resolved=false + // (the request was already resolved by A). + resolvedB := h.grantPermission(t, ctx, h.workspace.ID, proto.PermissionGrant{ + Permission: reqEv.Payload, + Action: proto.PermissionAllow, + }) + require.False(t, resolvedB, "client B's follow-up grant must report already resolved") +} + +// TestE2E_KillingClientASSEDoesNotBreakClientB covers PLAN item 6 +// scenario 3: terminating client A's SSE stream does not affect +// client B's stream; client B continues to receive events. +func TestE2E_KillingClientASSEDoesNotBreakClientB(t *testing.T) { + t.Parallel() + h := newE2EHarness(t) + ctxB, cancelB := context.WithCancel(t.Context()) + t.Cleanup(cancelB) + ctxA, cancelA := context.WithCancel(t.Context()) + + cidA := uuid.New().String() + cidB := uuid.New().String() + + _, killA := h.subscribeSSE(t, ctxA, h.workspace.ID, cidA) + t.Cleanup(killA) + evcB, killB := h.subscribeSSE(t, ctxB, h.workspace.ID, cidB) + t.Cleanup(killB) + + h.waitForAttached(t, 2) + + // Kill A's stream. The server's deferred DetachClient should + // drop A's claim, leaving B as the sole attached client. + cancelA() + killA() + + require.Eventually(t, func() bool { + return backend.WorkspaceLiveStreamCountForTest(h.workspace) == 1 + }, 3*time.Second, 10*time.Millisecond, + "expected client A's stream to drop the attached count to 1") + + // Workspace must still exist (B is holding it open) and + // shutdown callback must not have fired yet. + _, err := h.backend.GetWorkspace(h.workspace.ID) + require.NoError(t, err, "workspace must still exist while B is attached") + require.False(t, h.shutdownHit.Load(), + "shutdown callback must not fire while B is still attached") + + // Publish a fresh event; B must still receive it. + const sessionID = "s-after-a-died" + msg := message.Message{ + ID: "m-after", + SessionID: sessionID, + Role: message.Assistant, + Parts: []message.ContentPart{message.TextContent{Text: "still alive"}}, + } + h.app.SendEvent(pubsub.Event[message.Message]{ + Type: pubsub.CreatedEvent, + Payload: msg, + }) + + pickCtx, pickCancel := context.WithTimeout(ctxB, 3*time.Second) + defer pickCancel() + got, ok := drainUntil(pickCtx, evcB, func(e pubsub.Event[proto.Message]) bool { + return e.Payload.ID == "m-after" + }) + require.True(t, ok, "client B must still receive events after A's stream is killed") + require.Equal(t, sessionID, got.Payload.SessionID) +} + +// TestE2E_ShutdownCallbackFiresWhenLastClientLeaves covers PLAN +// item 6 scenario 4: once both clients disconnect, the backend +// runs its "last workspace removed -> server shutdown" path — now +// after the idle-linger window rather than immediately. +func TestE2E_ShutdownCallbackFiresWhenLastClientLeaves(t *testing.T) { + t.Parallel() + h := newE2EHarness(t) + // Shorten both lifecycle windows so the test doesn't wait out the + // production values. Both stay non-zero, so teardown still travels the + // real path: detach grace, then idle linger. + h.backend.SetDetachGrace(100 * time.Millisecond) + h.backend.SetIdleShutdownDelay(200 * time.Millisecond) + + ctxA, cancelA := context.WithCancel(t.Context()) + ctxB, cancelB := context.WithCancel(t.Context()) + t.Cleanup(cancelA) + t.Cleanup(cancelB) + + cidA := uuid.New().String() + cidB := uuid.New().String() + _, killA := h.subscribeSSE(t, ctxA, h.workspace.ID, cidA) + t.Cleanup(killA) + _, killB := h.subscribeSSE(t, ctxB, h.workspace.ID, cidB) + t.Cleanup(killB) + + h.waitForAttached(t, 2) + require.False(t, h.shutdownHit.Load(), "shutdown must not fire while clients are attached") + + cancelA() + killA() + require.Eventually(t, func() bool { + return backend.WorkspaceLiveStreamCountForTest(h.workspace) == 1 + }, 3*time.Second, 10*time.Millisecond) + require.False(t, h.shutdownHit.Load(), + "shutdown must not fire after only one client disconnects") + + cancelB() + killB() + require.Eventually(t, h.shutdownHit.Load, + 3*time.Second, 10*time.Millisecond, + "shutdown callback must fire once the last client disconnects "+ + "(after the detach grace and the idle linger)") + + // Workspace must be gone from the index. + _, err := h.backend.GetWorkspace(h.workspace.ID) + require.ErrorIs(t, err, backend.ErrWorkspaceNotFound) + + // Subsequent GETs against the now-defunct workspace return + // 404, confirming the http surface still reflects the teardown. + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, + h.httpSrv.URL+"/v1/workspaces/"+h.workspace.ID, nil) + require.NoError(t, err) + r, err := h.httpSrv.Client().Do(req) + require.NoError(t, err) + _, _ = io.Copy(io.Discard, r.Body) + r.Body.Close() + require.Equal(t, http.StatusNotFound, r.StatusCode) +} diff --git a/internal/server/logging.go b/internal/server/logging.go new file mode 100644 index 0000000000000000000000000000000000000000..4a14fb60e86800d9c79ed9d4851f00f3fc21f730 --- /dev/null +++ b/internal/server/logging.go @@ -0,0 +1,53 @@ +package server + +import ( + "log/slog" + "net/http" + "time" +) + +func (s *Server) loggingHandler(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if s.logger == nil { + next.ServeHTTP(w, r) + return + } + + start := time.Now() + lrw := &loggingResponseWriter{ResponseWriter: w, statusCode: http.StatusOK} + s.logger.Debug( + "HTTP request", + slog.String("method", r.Method), + slog.String("path", r.URL.Path), + slog.String("remote_addr", r.RemoteAddr), + slog.String("user_agent", r.UserAgent()), + ) + + next.ServeHTTP(lrw, r) + duration := time.Since(start) + + s.logger.Debug( + "HTTP response", + slog.String("method", r.Method), + slog.String("path", r.URL.Path), + slog.Int("status", lrw.statusCode), + slog.Duration("duration", duration), + slog.String("remote_addr", r.RemoteAddr), + slog.String("user_agent", r.UserAgent()), + ) + }) +} + +type loggingResponseWriter struct { + http.ResponseWriter + statusCode int +} + +func (lrw *loggingResponseWriter) WriteHeader(code int) { + lrw.statusCode = code + lrw.ResponseWriter.WriteHeader(code) +} + +func (lrw *loggingResponseWriter) Unwrap() http.ResponseWriter { + return lrw.ResponseWriter +} diff --git a/internal/server/recover_test.go b/internal/server/recover_test.go new file mode 100644 index 0000000000000000000000000000000000000000..45b616677be086890b47c6176135221259c1c9b6 --- /dev/null +++ b/internal/server/recover_test.go @@ -0,0 +1,94 @@ +package server + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/charmbracelet/crush/internal/proto" + "github.com/stretchr/testify/require" +) + +// TestRecoverHandler_PanicReturns500 verifies that a panicking handler +// surfaces as a structured 500 to the client, rather than closing the +// connection silently and producing an opaque EOF. +func TestRecoverHandler_PanicReturns500(t *testing.T) { + t.Parallel() + + s := &Server{} + h := s.recoverHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + panic("kaboom") + })) + + rec := httptest.NewRecorder() + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/test", nil) + h.ServeHTTP(rec, req) + + require.Equal(t, http.StatusInternalServerError, rec.Code) + body, err := io.ReadAll(rec.Body) + require.NoError(t, err) + require.NotEmpty(t, body) + + var perr proto.Error + require.NoError(t, json.Unmarshal(body, &perr)) + require.NotEmpty(t, perr.Message) +} + +// TestRecoverHandler_NoPanicPassthrough verifies that the middleware +// does not interfere with successful responses. +func TestRecoverHandler_NoPanicPassthrough(t *testing.T) { + t.Parallel() + + s := &Server{} + h := s.recoverHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusTeapot) + _, _ = w.Write([]byte("ok")) + })) + + rec := httptest.NewRecorder() + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/test", nil) + h.ServeHTTP(rec, req) + + require.Equal(t, http.StatusTeapot, rec.Code) + require.Equal(t, "ok", rec.Body.String()) +} + +// TestRecoverHandler_PanicAfterWriteHeader verifies that if a handler +// panics after it has already started writing the response, the +// middleware does not attempt to overwrite the status (which would +// trigger a superfluous WriteHeader warning) but still logs and +// recovers. +func TestRecoverHandler_PanicAfterWriteHeader(t *testing.T) { + t.Parallel() + + s := &Server{} + h := s.recoverHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("partial")) + panic("late panic") + })) + + rec := httptest.NewRecorder() + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/test", nil) + require.NotPanics(t, func() { h.ServeHTTP(rec, req) }) + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, "partial", rec.Body.String()) +} + +// TestRecoverHandler_AbortHandlerPropagates verifies that the documented +// http.ErrAbortHandler sentinel is re-panicked so the net/http server +// can handle it normally (suppress logging, close connection). +func TestRecoverHandler_AbortHandlerPropagates(t *testing.T) { + t.Parallel() + + s := &Server{} + h := s.recoverHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + panic(http.ErrAbortHandler) + })) + + rec := httptest.NewRecorder() + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/test", nil) + require.PanicsWithValue(t, http.ErrAbortHandler, func() { h.ServeHTTP(rec, req) }) +} diff --git a/internal/server/server.go b/internal/server/server.go new file mode 100644 index 0000000000000000000000000000000000000000..cb1b54fb17b83b1f743c9d11d50e1a4c2f4fcc57 --- /dev/null +++ b/internal/server/server.go @@ -0,0 +1,232 @@ +package server + +import ( + "context" + "fmt" + "log/slog" + "net" + "net/http" + "net/url" + "os" + "os/user" + "path/filepath" + "runtime" + "strings" + + "github.com/charmbracelet/crush/internal/backend" + "github.com/charmbracelet/crush/internal/config" +) + +// maxUnixSocketPathLen is the maximum length of a Unix domain socket +// path. The macOS sun_path field is 104 bytes; Linux allows 108. We +// use 104 so the resulting path is portable across both platforms. +const maxUnixSocketPathLen = 104 + +// socketDir returns the directory used for the Crush Unix socket. +// It prefers $XDG_RUNTIME_DIR when set (systemd's per-user runtime +// directory on Linux), and otherwise falls back to [os.TempDir], +// which resolves to the per-user private $TMPDIR on macOS and to +// /tmp on Linux. +func socketDir() string { + if dir := os.Getenv("XDG_RUNTIME_DIR"); dir != "" { + return dir + } + return os.TempDir() +} + +// ErrServerClosed is returned when the server is closed. +var ErrServerClosed = http.ErrServerClosed + +// ParseHostURL parses a host URL into a [url.URL]. +func ParseHostURL(host string) (*url.URL, error) { + proto, addr, ok := strings.Cut(host, "://") + if !ok { + return nil, fmt.Errorf("invalid host format: %s", host) + } + + var basePath string + if proto == "tcp" { + parsed, err := url.Parse("tcp://" + addr) + if err != nil { + return nil, fmt.Errorf("invalid tcp address: %v", err) + } + addr = parsed.Host + basePath = parsed.Path + } + return &url.URL{ + Scheme: proto, + Host: addr, + Path: basePath, + }, nil +} + +// DefaultHost returns the default server host. +// +// On Windows the address is a named pipe under \\.\pipe\. On Unix +// platforms the socket lives in the per-user runtime directory +// returned by [socketDir] and is named crush-.sock, falling +// back to crush.sock when the current uid cannot be determined. If +// the composed path would exceed [maxUnixSocketPathLen] bytes (the +// macOS sun_path limit), we fall back to /tmp/crush-.sock so +// the socket remains bindable. +func DefaultHost() string { + sock := "crush.sock" + usr, err := user.Current() + if err == nil && usr.Uid != "" { + sock = fmt.Sprintf("crush-%s.sock", usr.Uid) + } + if runtime.GOOS == "windows" { + return fmt.Sprintf("npipe:////./pipe/%s", sock) + } + path := filepath.Join(socketDir(), sock) + if len(path) > maxUnixSocketPathLen { + path = filepath.Join("/tmp", sock) + } + return "unix://" + path +} + +// Server represents a Crush server bound to a specific address. +type Server struct { + // Addr can be a TCP address, a Unix socket path, or a Windows named pipe. + Addr string + network string + + h *http.Server + ln net.Listener + + backend *backend.Backend + logger *slog.Logger +} + +// SetLogger sets the logger for the server. +func (s *Server) SetLogger(logger *slog.Logger) { + s.logger = logger +} + +// Backend returns the server's backend. Intended for integration tests +// that drive lifecycle transitions (detach, grace tuning) against a live +// HTTP surface. +func (s *Server) Backend() *backend.Backend { + return s.backend +} + +// DefaultServer returns a new [Server] with the default address. +func DefaultServer(cfg *config.ConfigStore) *Server { + hostURL, err := ParseHostURL(DefaultHost()) + if err != nil { + panic("invalid default host") + } + return NewServer(cfg, hostURL.Scheme, hostURL.Host) +} + +// NewServer creates a new [Server] with the given network and address. +func NewServer(cfg *config.ConfigStore, network, address string) *Server { + s := new(Server) + s.Addr = address + s.network = network + + // The backend is created with a shutdown callback that triggers + // a graceful server shutdown (e.g. when the last workspace is + // removed). + s.backend = backend.New(context.Background(), cfg, func() { + go func() { + slog.Info("Shutting down server...") + if err := s.Shutdown(context.Background()); err != nil { + slog.Error("Failed to shutdown server", "error", err) + } + }() + }) + s.installHandler() + if network == "tcp" { + s.h.Addr = address + } + return s +} + +// installHandler builds the protocol/router around s.backend and +// assigns the resulting http.Server to s.h. Extracted from +// [NewServer] so test harnesses can wire a Server around a +// pre-constructed backend. +func (s *Server) installHandler() { + var p http.Protocols + p.SetHTTP1(true) + p.SetUnencryptedHTTP2(true) + c := &controllerV1{backend: s.backend, server: s} + mux := http.NewServeMux() + for _, e := range c.endpoints() { + mux.HandleFunc(e.Method()+" "+e.Path(), e.Handler()) + } + mux.HandleFunc("GET /v1/docs/", c.handleDocsIndex) + mux.HandleFunc("GET /v1/docs/openapi.json", c.handleDocsSpec) + s.h = &http.Server{ + Protocols: &p, + Handler: s.recoverHandler(s.loggingHandler(mux)), + } +} + +// Handler returns the server's HTTP handler. Exposed so test harnesses +// can wrap it in an httptest.Server without going through the +// production listener setup. +func (s *Server) Handler() http.Handler { + return s.h.Handler +} + +// Serve accepts incoming connections on the listener. +func (s *Server) Serve(ln net.Listener) error { + return s.h.Serve(ln) +} + +// ListenAndServe starts the server and begins accepting connections. +func (s *Server) ListenAndServe() error { + if s.ln != nil { + return fmt.Errorf("server already started") + } + ln, removedStale, err := listen(s.network, s.Addr) + if err != nil { + return fmt.Errorf("failed to listen on %s: %w", s.Addr, err) + } + if removedStale && s.logger != nil { + s.logger.Warn("Removed stale socket before binding", "address", s.Addr) + } + return s.Serve(ln) +} + +func (s *Server) closeListener() { + if s.ln != nil { + s.ln.Close() + s.ln = nil + } +} + +// Close force closes all listeners and connections. +func (s *Server) Close() error { + defer func() { s.closeListener() }() + return s.h.Close() +} + +// Shutdown gracefully shuts down the server without interrupting active +// connections. +func (s *Server) Shutdown(ctx context.Context) error { + defer func() { s.closeListener() }() + return s.h.Shutdown(ctx) +} + +func (s *Server) logDebug(r *http.Request, msg string, args ...any) { + if s.logger != nil { + s.logger.With( + slog.String("method", r.Method), + slog.String("url", r.URL.String()), + slog.String("remote_addr", r.RemoteAddr), + ).Debug(msg, args...) + } +} + +func (s *Server) logError(r *http.Request, msg string, args ...any) { + if s.logger != nil { + s.logger.With( + slog.String("method", r.Method), + slog.String("url", r.URL.String()), + slog.String("remote_addr", r.RemoteAddr), + ).Error(msg, args...) + } +} diff --git a/internal/server/socket_classify.go b/internal/server/socket_classify.go new file mode 100644 index 0000000000000000000000000000000000000000..7edd1fdbca2b212d3c065396ee5fd7a37eeb9d2a --- /dev/null +++ b/internal/server/socket_classify.go @@ -0,0 +1,27 @@ +package server + +import ( + "errors" + "io/fs" + "net" + "syscall" +) + +// IsStaleSocketErr reports whether err indicates that a Unix-domain +// socket file exists on disk but no process is listening on it (a stale +// or orphaned socket). It returns false for nil and for timeout errors. +// +// The classification is cross-platform: ECONNREFUSED and fs.ErrNotExist +// are defined on every supported OS, so callers on Windows can use the +// same helper even though stale-socket recovery only applies to Unix +// sockets in practice. +func IsStaleSocketErr(err error) bool { + if err == nil { + return false + } + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + return false + } + return errors.Is(err, syscall.ECONNREFUSED) || errors.Is(err, fs.ErrNotExist) +} diff --git a/internal/shell/background.go b/internal/shell/background.go new file mode 100644 index 0000000000000000000000000000000000000000..55f85b6ce2781faae78dfa8df9e6c34e77e25eef --- /dev/null +++ b/internal/shell/background.go @@ -0,0 +1,245 @@ +package shell + +import ( + "bytes" + "context" + "fmt" + "slices" + "sync" + "sync/atomic" + "time" + + "github.com/charmbracelet/crush/internal/csync" +) + +const ( + // MaxBackgroundJobs is the maximum number of concurrent background jobs allowed + MaxBackgroundJobs = 50 + // CompletedJobRetentionMinutes is how long to keep completed jobs before auto-cleanup (8 hours) + CompletedJobRetentionMinutes = 8 * 60 +) + +// syncBuffer is a thread-safe wrapper around bytes.Buffer. +type syncBuffer struct { + buf bytes.Buffer + mu sync.RWMutex +} + +func (sb *syncBuffer) Write(p []byte) (n int, err error) { + sb.mu.Lock() + defer sb.mu.Unlock() + return sb.buf.Write(p) +} + +func (sb *syncBuffer) WriteString(s string) (n int, err error) { + sb.mu.Lock() + defer sb.mu.Unlock() + return sb.buf.WriteString(s) +} + +func (sb *syncBuffer) String() string { + sb.mu.RLock() + defer sb.mu.RUnlock() + return sb.buf.String() +} + +// BackgroundShell represents a shell running in the background. +type BackgroundShell struct { + ID string + Command string + Description string + Shell *Shell + WorkingDir string + ctx context.Context + cancel context.CancelFunc + stdout *syncBuffer + stderr *syncBuffer + done chan struct{} + exitErr error + completedAt atomic.Int64 // Unix timestamp when job completed (0 if still running) +} + +// BackgroundShellManager manages background shell instances. +type BackgroundShellManager struct { + shells *csync.Map[string, *BackgroundShell] +} + +var ( + backgroundManager *BackgroundShellManager + backgroundManagerOnce sync.Once + idCounter atomic.Uint64 +) + +// newBackgroundShellManager creates a new BackgroundShellManager instance. +func newBackgroundShellManager() *BackgroundShellManager { + return &BackgroundShellManager{ + shells: csync.NewMap[string, *BackgroundShell](), + } +} + +// GetBackgroundShellManager returns the singleton background shell manager. +func GetBackgroundShellManager() *BackgroundShellManager { + backgroundManagerOnce.Do(func() { + backgroundManager = newBackgroundShellManager() + }) + return backgroundManager +} + +// Start creates and starts a new background shell with the given command. +func (m *BackgroundShellManager) Start(ctx context.Context, workingDir string, blockFuncs []BlockFunc, command string, description string) (*BackgroundShell, error) { + // Check job limit + if m.shells.Len() >= MaxBackgroundJobs { + return nil, fmt.Errorf("maximum number of background jobs (%d) reached. Please terminate or wait for some jobs to complete", MaxBackgroundJobs) + } + + id := fmt.Sprintf("%03X", idCounter.Add(1)) + + shell := NewShell(&Options{ + WorkingDir: workingDir, + BlockFuncs: blockFuncs, + }) + + shellCtx, cancel := context.WithCancel(ctx) + + bgShell := &BackgroundShell{ + ID: id, + Command: command, + Description: description, + WorkingDir: workingDir, + Shell: shell, + ctx: shellCtx, + cancel: cancel, + stdout: &syncBuffer{}, + stderr: &syncBuffer{}, + done: make(chan struct{}), + } + + m.shells.Set(id, bgShell) + + go func() { + defer close(bgShell.done) + + err := shell.ExecStream(shellCtx, command, bgShell.stdout, bgShell.stderr) + + bgShell.exitErr = err + bgShell.completedAt.Store(time.Now().Unix()) + }() + + return bgShell, nil +} + +// Get retrieves a background shell by ID. +func (m *BackgroundShellManager) Get(id string) (*BackgroundShell, bool) { + return m.shells.Get(id) +} + +// Remove removes a background shell from the manager without terminating it. +// This is useful when a shell has already completed and you just want to clean up tracking. +func (m *BackgroundShellManager) Remove(id string) error { + _, ok := m.shells.Take(id) + if !ok { + return fmt.Errorf("background shell not found: %s", id) + } + return nil +} + +// Kill terminates a background shell by ID. +func (m *BackgroundShellManager) Kill(id string) error { + shell, ok := m.shells.Take(id) + if !ok { + return fmt.Errorf("background shell not found: %s", id) + } + + shell.cancel() + <-shell.done + return nil +} + +// BackgroundShellInfo contains information about a background shell. +type BackgroundShellInfo struct { + ID string + Command string + Description string +} + +// List returns all background shell IDs. +func (m *BackgroundShellManager) List() []string { + ids := make([]string, 0, m.shells.Len()) + for id := range m.shells.Seq2() { + ids = append(ids, id) + } + return ids +} + +// Cleanup removes completed jobs that have been finished for more than the retention period +func (m *BackgroundShellManager) Cleanup() int { + now := time.Now().Unix() + retentionSeconds := int64(CompletedJobRetentionMinutes * 60) + + var toRemove []string + for shell := range m.shells.Seq() { + completedAt := shell.completedAt.Load() + if completedAt > 0 && now-completedAt > retentionSeconds { + toRemove = append(toRemove, shell.ID) + } + } + + for _, id := range toRemove { + m.Remove(id) + } + + return len(toRemove) +} + +// KillAll terminates all background shells. The provided context bounds how +// long the function waits for each shell to exit. +func (m *BackgroundShellManager) KillAll(ctx context.Context) { + shells := slices.Collect(m.shells.Seq()) + m.shells.Reset(map[string]*BackgroundShell{}) + + var wg sync.WaitGroup + for _, shell := range shells { + wg.Go(func() { + shell.cancel() + select { + case <-shell.done: + case <-ctx.Done(): + } + }) + } + wg.Wait() +} + +// GetOutput returns the current output of a background shell. +func (bs *BackgroundShell) GetOutput() (stdout string, stderr string, done bool, err error) { + select { + case <-bs.done: + return bs.stdout.String(), bs.stderr.String(), true, bs.exitErr + default: + return bs.stdout.String(), bs.stderr.String(), false, nil + } +} + +// IsDone checks if the background shell has finished execution. +func (bs *BackgroundShell) IsDone() bool { + select { + case <-bs.done: + return true + default: + return false + } +} + +// Wait blocks until the background shell completes. +func (bs *BackgroundShell) Wait() { + <-bs.done +} + +func (bs *BackgroundShell) WaitContext(ctx context.Context) bool { + select { + case <-bs.done: + return true + case <-ctx.Done(): + return false + } +} diff --git a/internal/shell/background_test.go b/internal/shell/background_test.go new file mode 100644 index 0000000000000000000000000000000000000000..063f5c31b3daaf3f75fc7a99604561cfcad07f9c --- /dev/null +++ b/internal/shell/background_test.go @@ -0,0 +1,330 @@ +package shell + +import ( + "context" + "runtime" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestBackgroundShellManager_Start(t *testing.T) { + t.Skip("Skipping this until I figure out why its flaky") + t.Parallel() + + ctx := t.Context() + workingDir := t.TempDir() + manager := newBackgroundShellManager() + + bgShell, err := manager.Start(ctx, workingDir, nil, "echo 'hello world'", "") + if err != nil { + t.Fatalf("failed to start background shell: %v", err) + } + + if bgShell.ID == "" { + t.Error("expected shell ID to be non-empty") + } + + // Wait for the command to complete + bgShell.Wait() + + stdout, stderr, done, err := bgShell.GetOutput() + if !done { + t.Error("expected shell to be done") + } + + if err != nil { + t.Errorf("expected no error, got: %v", err) + } + + if !strings.Contains(stdout, "hello world") { + t.Errorf("expected stdout to contain 'hello world', got: %s", stdout) + } + + if stderr != "" { + t.Errorf("expected empty stderr, got: %s", stderr) + } +} + +func TestBackgroundShellManager_Get(t *testing.T) { + t.Parallel() + + ctx := t.Context() + workingDir := t.TempDir() + manager := newBackgroundShellManager() + + bgShell, err := manager.Start(ctx, workingDir, nil, "echo 'test'", "") + if err != nil { + t.Fatalf("failed to start background shell: %v", err) + } + + // Retrieve the shell + retrieved, ok := manager.Get(bgShell.ID) + if !ok { + t.Error("expected to find the background shell") + } + + if retrieved.ID != bgShell.ID { + t.Errorf("expected shell ID %s, got %s", bgShell.ID, retrieved.ID) + } + + // Clean up + manager.Kill(bgShell.ID) +} + +func TestBackgroundShellManager_Kill(t *testing.T) { + t.Parallel() + + ctx := t.Context() + workingDir := t.TempDir() + manager := newBackgroundShellManager() + + // Start a long-running command + bgShell, err := manager.Start(ctx, workingDir, nil, "sleep 10", "") + if err != nil { + t.Fatalf("failed to start background shell: %v", err) + } + + // Kill it + err = manager.Kill(bgShell.ID) + if err != nil { + t.Errorf("failed to kill background shell: %v", err) + } + + // Verify it's no longer in the manager + _, ok := manager.Get(bgShell.ID) + if ok { + t.Error("expected shell to be removed after kill") + } + + // Verify the shell is done + if !bgShell.IsDone() { + t.Error("expected shell to be done after kill") + } +} + +func TestBackgroundShellManager_KillNonExistent(t *testing.T) { + t.Parallel() + + manager := newBackgroundShellManager() + + err := manager.Kill("non-existent-id") + if err == nil { + t.Error("expected error when killing non-existent shell") + } +} + +func TestBackgroundShell_IsDone(t *testing.T) { + t.Parallel() + + ctx := t.Context() + workingDir := t.TempDir() + manager := newBackgroundShellManager() + + bgShell, err := manager.Start(ctx, workingDir, nil, "echo 'quick'", "") + if err != nil { + t.Fatalf("failed to start background shell: %v", err) + } + + // Wait for the command to complete (Windows is slower to spin up). + require.Eventually(t, bgShell.IsDone, 5*time.Second, 50*time.Millisecond, "expected shell to be done") + + // Clean up + manager.Kill(bgShell.ID) +} + +func TestBackgroundShell_WithBlockFuncs(t *testing.T) { + t.Parallel() + + ctx := t.Context() + workingDir := t.TempDir() + manager := newBackgroundShellManager() + + blockFuncs := []BlockFunc{ + CommandsBlocker([]string{"curl", "wget"}), + } + + bgShell, err := manager.Start(ctx, workingDir, blockFuncs, "curl example.com", "") + if err != nil { + t.Fatalf("failed to start background shell: %v", err) + } + + // Wait for the command to complete + bgShell.Wait() + + stdout, stderr, done, execErr := bgShell.GetOutput() + if !done { + t.Error("expected shell to be done") + } + + // The command should have been blocked + output := stdout + stderr + if !strings.Contains(output, "not allowed") && execErr == nil { + t.Errorf("expected command to be blocked, got stdout: %s, stderr: %s, err: %v", stdout, stderr, execErr) + } + + // Clean up + manager.Kill(bgShell.ID) +} + +func TestBackgroundShellManager_List(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("skipping flacky test on windows") + } + + t.Parallel() + + ctx := t.Context() + workingDir := t.TempDir() + manager := newBackgroundShellManager() + + // Start two shells + bgShell1, err := manager.Start(ctx, workingDir, nil, "sleep 1", "") + if err != nil { + t.Fatalf("failed to start first background shell: %v", err) + } + + bgShell2, err := manager.Start(ctx, workingDir, nil, "sleep 1", "") + if err != nil { + t.Fatalf("failed to start second background shell: %v", err) + } + + ids := manager.List() + + // Check that both shells are in the list + found1 := false + found2 := false + for _, id := range ids { + if id == bgShell1.ID { + found1 = true + } + if id == bgShell2.ID { + found2 = true + } + } + + if !found1 { + t.Errorf("expected to find shell %s in list", bgShell1.ID) + } + if !found2 { + t.Errorf("expected to find shell %s in list", bgShell2.ID) + } + + // Clean up + manager.Kill(bgShell1.ID) + manager.Kill(bgShell2.ID) +} + +func TestBackgroundShellManager_KillAll(t *testing.T) { + t.Parallel() + + ctx := t.Context() + workingDir := t.TempDir() + manager := newBackgroundShellManager() + + // Start multiple long-running shells + shell1, err := manager.Start(ctx, workingDir, nil, "sleep 10", "") + if err != nil { + t.Fatalf("failed to start shell 1: %v", err) + } + + shell2, err := manager.Start(ctx, workingDir, nil, "sleep 10", "") + if err != nil { + t.Fatalf("failed to start shell 2: %v", err) + } + + shell3, err := manager.Start(ctx, workingDir, nil, "sleep 10", "") + if err != nil { + t.Fatalf("failed to start shell 3: %v", err) + } + + // Verify shells are running + if shell1.IsDone() || shell2.IsDone() || shell3.IsDone() { + t.Error("shells should not be done yet") + } + + // Kill all shells + manager.KillAll(t.Context()) + + // Verify all shells are done + if !shell1.IsDone() { + t.Error("shell1 should be done after KillAll") + } + if !shell2.IsDone() { + t.Error("shell2 should be done after KillAll") + } + if !shell3.IsDone() { + t.Error("shell3 should be done after KillAll") + } + + // Verify they're removed from the manager + if _, ok := manager.Get(shell1.ID); ok { + t.Error("shell1 should be removed from manager") + } + if _, ok := manager.Get(shell2.ID); ok { + t.Error("shell2 should be removed from manager") + } + if _, ok := manager.Get(shell3.ID); ok { + t.Error("shell3 should be removed from manager") + } + + // Verify list is empty (or doesn't contain our shells) + ids := manager.List() + for _, id := range ids { + if id == shell1.ID || id == shell2.ID || id == shell3.ID { + t.Errorf("shell %s should not be in list after KillAll", id) + } + } +} + +func TestBackgroundShellManager_KillAll_Timeout(t *testing.T) { + t.Parallel() + + // XXX: can't use synctest here - causes --race to trip. + + workingDir := t.TempDir() + manager := newBackgroundShellManager() + + // Start a shell that traps signals and ignores cancellation. + _, err := manager.Start(t.Context(), workingDir, nil, "trap '' TERM INT; sleep 60", "") + require.NoError(t, err) + + // Short timeout to test the timeout path. + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + t.Cleanup(cancel) + + start := time.Now() + manager.KillAll(ctx) + + elapsed := time.Since(start) + + // Must return promptly after timeout, not hang for 60 seconds. + require.Less(t, elapsed, 2*time.Second) +} + +func TestBackgroundShell_WaitContext_Completed(t *testing.T) { + t.Parallel() + + done := make(chan struct{}) + close(done) + + bgShell := &BackgroundShell{done: done} + + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + t.Cleanup(cancel) + + require.True(t, bgShell.WaitContext(ctx)) +} + +func TestBackgroundShell_WaitContext_Canceled(t *testing.T) { + t.Parallel() + + bgShell := &BackgroundShell{done: make(chan struct{})} + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + require.False(t, bgShell.WaitContext(ctx)) +} diff --git a/internal/shell/builtins_registry.go b/internal/shell/builtins_registry.go new file mode 100644 index 0000000000000000000000000000000000000000..509a498d139b2bdafb0210e8df19556f5a124156 --- /dev/null +++ b/internal/shell/builtins_registry.go @@ -0,0 +1,28 @@ +package shell + +import ( + "context" + "io" +) + +// BuiltinHandler is a function that handles a shell builtin command. +// It receives the full args slice (including the command name as args[0]), +// the context (which may carry a ConfigBuilder), and I/O streams. +type BuiltinHandler func(ctx context.Context, args []string, stdin io.Reader, stdout, stderr io.Writer) error + +// builtins holds all in-process builtin handlers, including jq (registered +// in this package) and config builtins (registered by shellconfig via +// RegisterBuiltin to avoid an import cycle). Dispatched by builtinHandler in +// run.go. +var builtins = make(map[string]BuiltinHandler) + +func init() { + RegisterBuiltin("jq", handleJQ) +} + +// RegisterBuiltin registers a builtin command handler. Must be called during +// init(), before any shell execution. If a handler is already registered for +// the same name, the new one replaces it. +func RegisterBuiltin(name string, handler BuiltinHandler) { + builtins[name] = handler +} diff --git a/internal/shell/builtins_registry_test.go b/internal/shell/builtins_registry_test.go new file mode 100644 index 0000000000000000000000000000000000000000..54721d5bdd9082deac0c5ecc37d85f07afd821dd --- /dev/null +++ b/internal/shell/builtins_registry_test.go @@ -0,0 +1,97 @@ +package shell + +import ( + "bytes" + "context" + "io" + "os" + "path/filepath" + "testing" +) + +// TestRegisterBuiltin_DispatchedThroughRun verifies that a handler registered +// via RegisterBuiltin is dispatched through the stateless Run path, receiving +// the full args slice and the I/O streams from the interpreter context. +func TestRegisterBuiltin_DispatchedThroughRun(t *testing.T) { + RegisterBuiltin("crushtest-echo", func(ctx context.Context, args []string, stdin io.Reader, stdout, stderr io.Writer) error { + io.WriteString(stdout, "builtin:"+args[1]) + return nil + }) + t.Cleanup(func() { delete(builtins, "crushtest-echo") }) + + var stdout bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: "crushtest-echo hello", + Cwd: t.TempDir(), + Stdout: &stdout, + }) + if err != nil { + t.Fatalf("Run returned error: %v", err) + } + if got := stdout.String(); got != "builtin:hello" { + t.Fatalf("stdout = %q, want %q", got, "builtin:hello") + } +} + +// TestRegisterBuiltin_OverridesPATHBinary verifies that a registered builtin +// takes precedence over an executable of the same name found on PATH. This is +// the stated design intent of builtinHandler sitting ahead of the exec layer. +func TestRegisterBuiltin_OverridesPATHBinary(t *testing.T) { + dir := t.TempDir() + // Put a fake "crushtest-shadowed" binary on PATH that would print + // "from-path" if it ran. + binDir := filepath.Join(dir, "bin") + if err := os.MkdirAll(binDir, 0o755); err != nil { + t.Fatal(err) + } + script := filepath.Join(binDir, "crushtest-shadowed") + if err := os.WriteFile(script, []byte("#!/bin/sh\necho from-path\n"), 0o755); err != nil { + t.Fatal(err) + } + t.Setenv("PATH", binDir+string(os.PathListSeparator)+os.Getenv("PATH")) + + RegisterBuiltin("crushtest-shadowed", func(ctx context.Context, args []string, stdin io.Reader, stdout, stderr io.Writer) error { + io.WriteString(stdout, "from-builtin") + return nil + }) + t.Cleanup(func() { delete(builtins, "crushtest-shadowed") }) + + var stdout bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: "crushtest-shadowed", + Cwd: dir, + Env: os.Environ(), + Stdout: &stdout, + }) + if err != nil { + t.Fatalf("Run returned error: %v", err) + } + if got := stdout.String(); got != "from-builtin" { + t.Fatalf("builtin did not shadow PATH binary: stdout = %q, want %q", got, "from-builtin") + } +} + +// TestBuiltinHandler_UnknownFallsThrough verifies that an unregistered command +// name is passed to the underlying exec layer rather than being swallowed. +func TestBuiltinHandler_UnknownFallsThrough(t *testing.T) { + var stdout bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: "echo fallthrough", + Cwd: t.TempDir(), + Stdout: &stdout, + }) + if err != nil { + t.Fatalf("Run returned error: %v", err) + } + if got := stdout.String(); got != "fallthrough\n" { + t.Fatalf("stdout = %q, want %q", got, "fallthrough\n") + } +} + +// TestRegisterBuiltin_JQRegistered verifies that jq is registered through the +// same registry as every other builtin (no special-cased dispatch path). +func TestRegisterBuiltin_JQRegistered(t *testing.T) { + if _, ok := builtins["jq"]; !ok { + t.Fatal("jq is not registered in the builtins map") + } +} diff --git a/internal/shell/command_block_test.go b/internal/shell/command_block_test.go new file mode 100644 index 0000000000000000000000000000000000000000..7cfdf6afe5f8065afb90d866dfed41ba51b95cb5 --- /dev/null +++ b/internal/shell/command_block_test.go @@ -0,0 +1,376 @@ +package shell + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestCommandBlocking(t *testing.T) { + tests := []struct { + name string + blockFuncs []BlockFunc + command string + shouldBlock bool + }{ + { + name: "block simple command", + blockFuncs: []BlockFunc{ + func(args []string) bool { + return len(args) > 0 && args[0] == "curl" + }, + }, + command: "curl https://example.com", + shouldBlock: true, + }, + { + name: "allow non-blocked command", + blockFuncs: []BlockFunc{ + func(args []string) bool { + return len(args) > 0 && args[0] == "curl" + }, + }, + command: "echo hello", + shouldBlock: false, + }, + { + name: "block subcommand", + blockFuncs: []BlockFunc{ + func(args []string) bool { + return len(args) >= 2 && args[0] == "brew" && args[1] == "install" + }, + }, + command: "brew install wget", + shouldBlock: true, + }, + { + name: "allow different subcommand", + blockFuncs: []BlockFunc{ + func(args []string) bool { + return len(args) >= 2 && args[0] == "brew" && args[1] == "install" + }, + }, + command: "brew list", + shouldBlock: false, + }, + { + name: "block npm global install with -g", + blockFuncs: []BlockFunc{ + ArgumentsBlocker("npm", []string{"install"}, []string{"-g"}), + }, + command: "npm install -g typescript", + shouldBlock: true, + }, + { + name: "block npm global install with --global", + blockFuncs: []BlockFunc{ + ArgumentsBlocker("npm", []string{"install"}, []string{"--global"}), + }, + command: "npm install --global typescript", + shouldBlock: true, + }, + { + name: "allow npm local install", + blockFuncs: []BlockFunc{ + ArgumentsBlocker("npm", []string{"install"}, []string{"-g"}), + ArgumentsBlocker("npm", []string{"install"}, []string{"--global"}), + }, + command: "npm install typescript", + shouldBlock: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create a temporary directory for each test + tmpDir := t.TempDir() + + shell := NewShell(&Options{ + WorkingDir: tmpDir, + BlockFuncs: tt.blockFuncs, + }) + + _, _, err := shell.Exec(t.Context(), tt.command) + + if tt.shouldBlock { + if err == nil { + t.Errorf("Expected command to be blocked, but it was allowed") + } else if !strings.Contains(err.Error(), "not allowed for security reasons") { + t.Errorf("Expected security error, got: %v", err) + } + } else { + // For non-blocked commands, we might get other errors (like command not found) + // but we shouldn't get the security error + if err != nil && strings.Contains(err.Error(), "not allowed for security reasons") { + t.Errorf("Command was unexpectedly blocked: %v", err) + } + } + }) + } +} + +func TestArgumentsBlocker(t *testing.T) { + tests := []struct { + name string + cmd string + args []string + flags []string + input []string + shouldBlock bool + }{ + // Basic command blocking + { + name: "block exact command match", + cmd: "npm", + args: []string{"install"}, + flags: nil, + input: []string{"npm", "install", "package"}, + shouldBlock: true, + }, + { + name: "allow different command", + cmd: "npm", + args: []string{"install"}, + flags: nil, + input: []string{"yarn", "install", "package"}, + shouldBlock: false, + }, + { + name: "allow different subcommand", + cmd: "npm", + args: []string{"install"}, + flags: nil, + input: []string{"npm", "list"}, + shouldBlock: false, + }, + + // Flag-based blocking + { + name: "block with single flag", + cmd: "npm", + args: []string{"install"}, + flags: []string{"-g"}, + input: []string{"npm", "install", "-g", "typescript"}, + shouldBlock: true, + }, + { + name: "block with flag in different position", + cmd: "npm", + args: []string{"install"}, + flags: []string{"-g"}, + input: []string{"npm", "install", "typescript", "-g"}, + shouldBlock: true, + }, + { + name: "allow without required flag", + cmd: "npm", + args: []string{"install"}, + flags: []string{"-g"}, + input: []string{"npm", "install", "typescript"}, + shouldBlock: false, + }, + { + name: "block with multiple flags", + cmd: "pip", + args: []string{"install"}, + flags: []string{"--user"}, + input: []string{"pip", "install", "--user", "--upgrade", "package"}, + shouldBlock: true, + }, + + // Complex argument patterns + { + name: "block multi-arg subcommand", + cmd: "yarn", + args: []string{"global", "add"}, + flags: nil, + input: []string{"yarn", "global", "add", "typescript"}, + shouldBlock: true, + }, + { + name: "allow partial multi-arg match", + cmd: "yarn", + args: []string{"global", "add"}, + flags: nil, + input: []string{"yarn", "global", "list"}, + shouldBlock: false, + }, + + // Edge cases + { + name: "handle empty input", + cmd: "npm", + args: []string{"install"}, + flags: nil, + input: []string{}, + shouldBlock: false, + }, + { + name: "handle command only", + cmd: "npm", + args: []string{"install"}, + flags: nil, + input: []string{"npm"}, + shouldBlock: false, + }, + { + name: "block pacman with -S flag", + cmd: "pacman", + args: nil, + flags: []string{"-S"}, + input: []string{"pacman", "-S", "package"}, + shouldBlock: true, + }, + { + name: "allow pacman without -S flag", + cmd: "pacman", + args: nil, + flags: []string{"-S"}, + input: []string{"pacman", "-Q", "package"}, + shouldBlock: false, + }, + + // `go test -exec` + { + name: "go test exec", + cmd: "go", + args: []string{"test"}, + flags: []string{"-exec"}, + input: []string{"go", "test", "-exec", "bash -c 'echo hello'"}, + shouldBlock: true, + }, + { + name: "go test exec", + cmd: "go", + args: []string{"test"}, + flags: []string{"-exec"}, + input: []string{"go", "test", `-exec="bash -c 'echo hello'"`}, + shouldBlock: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + blocker := ArgumentsBlocker(tt.cmd, tt.args, tt.flags) + result := blocker(tt.input) + require.Equal(t, tt.shouldBlock, result, + "Expected block=%v for input %v", tt.shouldBlock, tt.input) + }) + } +} + +func TestCommandsBlocker(t *testing.T) { + tests := []struct { + name string + banned []string + input []string + shouldBlock bool + }{ + { + name: "block single banned command", + banned: []string{"curl"}, + input: []string{"curl", "https://example.com"}, + shouldBlock: true, + }, + { + name: "allow non-banned command", + banned: []string{"curl", "wget"}, + input: []string{"echo", "hello"}, + shouldBlock: false, + }, + { + name: "block from multiple banned", + banned: []string{"curl", "wget", "nc"}, + input: []string{"wget", "https://example.com"}, + shouldBlock: true, + }, + { + name: "handle empty input", + banned: []string{"curl"}, + input: []string{}, + shouldBlock: false, + }, + { + name: "case sensitive matching", + banned: []string{"curl"}, + input: []string{"CURL", "https://example.com"}, + shouldBlock: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + blocker := CommandsBlocker(tt.banned) + result := blocker(tt.input) + require.Equal(t, tt.shouldBlock, result, + "Expected block=%v for input %v", tt.shouldBlock, tt.input) + }) + } +} + +func TestSplitArgsFlags(t *testing.T) { + tests := []struct { + name string + input []string + wantArgs []string + wantFlags []string + }{ + { + name: "only args", + input: []string{"install", "package", "another"}, + wantArgs: []string{"install", "package", "another"}, + wantFlags: []string{}, + }, + { + name: "only flags", + input: []string{"-g", "--verbose", "-f"}, + wantArgs: []string{}, + wantFlags: []string{"-g", "--verbose", "-f"}, + }, + { + name: "mixed args and flags", + input: []string{"install", "-g", "package", "--verbose"}, + wantArgs: []string{"install", "package"}, + wantFlags: []string{"-g", "--verbose"}, + }, + { + name: "empty input", + input: []string{}, + wantArgs: []string{}, + wantFlags: []string{}, + }, + { + name: "single dash flag", + input: []string{"-S", "package"}, + wantArgs: []string{"package"}, + wantFlags: []string{"-S"}, + }, + { + name: "flag with equals sign", + input: []string{"-exec=bash", "package"}, + wantArgs: []string{"package"}, + wantFlags: []string{"-exec"}, + }, + { + name: "long flag with equals sign", + input: []string{"--config=/path/to/config", "run"}, + wantArgs: []string{"run"}, + wantFlags: []string{"--config"}, + }, + { + name: "flag with complex value", + input: []string{`-exec="bash -c 'echo hello'"`, "test"}, + wantArgs: []string{"test"}, + wantFlags: []string{"-exec"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + args, flags := splitArgsFlags(tt.input) + require.Equal(t, tt.wantArgs, args, "args mismatch") + require.Equal(t, tt.wantFlags, flags, "flags mismatch") + }) + } +} diff --git a/internal/shell/comparison_test.go b/internal/shell/comparison_test.go new file mode 100644 index 0000000000000000000000000000000000000000..84d939aa247bbff6f4de9ffbc423d1cbe4aaf7f6 --- /dev/null +++ b/internal/shell/comparison_test.go @@ -0,0 +1,41 @@ +package shell + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestShellPerformanceComparison(t *testing.T) { + shell := NewShell(&Options{WorkingDir: t.TempDir()}) + + // Test quick command + start := time.Now() + stdout, stderr, err := shell.Exec(t.Context(), "echo 'hello'") + exitCode := ExitCode(err) + duration := time.Since(start) + + require.NoError(t, err) + require.Equal(t, 0, exitCode) + require.Contains(t, stdout, "hello") + require.Empty(t, stderr) + + t.Logf("Quick command took: %v", duration) +} + +// Benchmark CPU usage during polling +func BenchmarkShellPolling(b *testing.B) { + shell := NewShell(&Options{WorkingDir: b.TempDir()}) + + b.ReportAllocs() + + for b.Loop() { + // Use a short sleep to measure polling overhead + _, _, err := shell.Exec(b.Context(), "sleep 0.02") + exitCode := ExitCode(err) + if err != nil || exitCode != 0 { + b.Fatalf("Command failed: %v, exit code: %d", err, exitCode) + } + } +} diff --git a/internal/shell/coreutils.go b/internal/shell/coreutils.go new file mode 100644 index 0000000000000000000000000000000000000000..ad5401d5f98b2f51dcc8b38da357a70e085703c6 --- /dev/null +++ b/internal/shell/coreutils.go @@ -0,0 +1,19 @@ +package shell + +import ( + "os" + "runtime" + "strconv" +) + +var useGoCoreUtils bool + +func init() { + // If CRUSH_CORE_UTILS is set to either true or false, respect that. + // By default, enable on Windows only. + if v, err := strconv.ParseBool(os.Getenv("CRUSH_CORE_UTILS")); err == nil { + useGoCoreUtils = v + } else { + useGoCoreUtils = runtime.GOOS == "windows" + } +} diff --git a/internal/shell/coreutils_exec.go b/internal/shell/coreutils_exec.go new file mode 100644 index 0000000000000000000000000000000000000000..699904a5cf6eb567910a3297ca3b120c7b6109a2 --- /dev/null +++ b/internal/shell/coreutils_exec.go @@ -0,0 +1,12 @@ +//go:build !solaris && !illumos + +package shell + +import ( + "mvdan.cc/sh/x/coreutils" +) + +// coreUtilsExecHandler installs mvdan.cc/sh's Go coreutils. It is only wired +// in when useGoCoreUtils is set. On illumos and solaris, coreutils_exec_stub.go +// supplies a nil handler because the coreutils package does not build there. +var coreUtilsExecHandler execMiddleware = coreutils.ExecHandler diff --git a/internal/shell/coreutils_exec_stub.go b/internal/shell/coreutils_exec_stub.go new file mode 100644 index 0000000000000000000000000000000000000000..a35e8e43cc17fe172f87449ffccc0af602f1c573 --- /dev/null +++ b/internal/shell/coreutils_exec_stub.go @@ -0,0 +1,8 @@ +//go:build solaris || illumos + +package shell + +// coreUtilsExecHandler is nil on illumos and solaris: mvdan.cc/sh's coreutils +// package does not compile there, so the Go coreutils middleware is simply +// unavailable. standardHandlers skips it when the handler is nil. +var coreUtilsExecHandler execMiddleware diff --git a/internal/shell/dispatch.go b/internal/shell/dispatch.go new file mode 100644 index 0000000000000000000000000000000000000000..9f75726ccd65b9349500729d85fab0a4e5abaa89 --- /dev/null +++ b/internal/shell/dispatch.go @@ -0,0 +1,425 @@ +package shell + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "io/fs" + "log/slog" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + + "github.com/charmbracelet/crush/internal/filepathext" + "mvdan.cc/sh/v3/expand" + "mvdan.cc/sh/v3/interp" + "mvdan.cc/sh/v3/syntax" +) + +// probeWindow is how many bytes we read from the head of a file to decide +// how to dispatch it. 128 is plenty for a shebang line and for magic-byte +// inspection, while small enough to make the probe cheap for users whose +// hooks invoke many scripts. +const probeWindow = 128 + +// scriptDispatchHandler returns middleware that intercepts exec of a +// path-prefixed argv[0] (e.g. ./foo.sh, /opt/bin/tool, C:\foo\bar.exe) and +// dispatches based on the file's contents: +// +// 1. Shebang line (#!...) → exec the named interpreter via os/exec. The +// interpreter is resolved literally first, then via PATH on the +// basename as a permissive fallback (so #!/bin/bash works on Windows +// boxes where Git for Windows puts bash.exe on PATH). +// 2. Known binary magic (MZ, ELF, Mach-O) or a NUL byte in the probe +// window → pass through to the next handler (mvdan's default exec). +// 3. Otherwise → treat the file as shell source and run it in-process via +// a nested interp.Runner that reuses the same handler stack. +// +// Non-path-prefixed argv[0] and empty args are passed straight through; this +// handler is a no-op for ordinary commands like `echo` or `jq`. +// +// blockFuncs is the block list used when building the nested runner for the +// shell-source case, so deny rules apply recursively to commands invoked +// from in-process scripts. +func scriptDispatchHandler(blockFuncs []BlockFunc) execMiddleware { + return func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc { + return func(ctx context.Context, args []string) error { + if len(args) == 0 || !isPathPrefixed(args[0]) { + return next(ctx, args) + } + + // Resolve relative paths against the interpreter's cwd, not + // the process cwd — hook commands are authored with the hook + // Runner's cwd in mind and sub-shells can cd before an exec. + scriptPath := filepathext.SmartJoin(interp.HandlerCtx(ctx).Dir, args[0]) + probe, err := probeFile(scriptPath) + if err != nil { + return err + } + + switch { + case hasShebang(probe): + return dispatchShebang(ctx, scriptPath, probe, args) + case isBinary(probe): + return next(ctx, args) + default: + return runShellSource(ctx, scriptPath, args, blockFuncs) + } + } + } +} + +// isPathPrefixed reports whether argv[0] is a file reference (as opposed +// to a bare command to be resolved via PATH). A path reference starts with +// `./`, `../`, `/`, or — on Windows — a drive-letter prefix. +// +// Note: mvdan already performs tilde expansion during word expansion, so +// `~/script.sh` arrives here as an absolute path. We still call the helper +// on the raw string to stay robust if a future change ever bypasses that +// expansion; cover that path with a regression test. +func isPathPrefixed(arg string) bool { + switch { + case strings.HasPrefix(arg, "./"), + strings.HasPrefix(arg, "../"), + strings.HasPrefix(arg, "/"): + return true + } + if runtime.GOOS == "windows" { + // Drive-letter paths: C:\foo or C:/foo (length check avoids + // accidentally matching a single letter followed by a colon). + if len(arg) >= 3 && isDriveLetter(arg[0]) && arg[1] == ':' && + (arg[2] == '\\' || arg[2] == '/') { + return true + } + // Also treat backslash-prefixed UNC-like paths as path-prefixed. + if strings.HasPrefix(arg, "\\") { + return true + } + } + return false +} + +func isDriveLetter(b byte) bool { + return (b >= 'A' && b <= 'Z') || (b >= 'a' && b <= 'z') +} + +// probeFile reads the first probeWindow bytes of the target path. It +// deliberately does not slurp the whole file: callers that need the full +// contents (only the shell-source branch) re-open via os.ReadFile. This +// keeps memory bounded when argv[0] turns out to be a large binary. +// +// Returns errors surfaced by os.Open/os.Stat directly so callers see the +// real reason: ENOENT, EACCES, EISDIR, ELOOP, etc. +func probeFile(path string) ([]byte, error) { + f, err := os.Open(path) + if err != nil { + return nil, err + } + defer f.Close() + fi, err := f.Stat() + if err != nil { + return nil, err + } + if fi.IsDir() { + return nil, fmt.Errorf("%s: is a directory", path) + } + probe := make([]byte, probeWindow) + n, err := io.ReadFull(f, probe) + if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF { + return nil, err + } + return probe[:n], nil +} + +// hasShebang reports whether probe starts with the `#!` marker. A +// one-byte file that happens to be `#` is not a shebang. +func hasShebang(probe []byte) bool { + return len(probe) >= 2 && probe[0] == '#' && probe[1] == '!' +} + +// isBinary heuristically classifies probe as an executable or otherwise +// non-text file. A NUL byte in the first probeWindow bytes is the classic +// Unix-y text-vs-binary signal; we additionally recognize known magic +// numbers so we can fast-path well-formed binaries that happen to have no +// NUL in the first 128 bytes (rare but possible for small binaries). +func isBinary(probe []byte) bool { + if bytes.IndexByte(probe, 0) >= 0 { + return true + } + magics := [][]byte{ + {'M', 'Z'}, // Windows PE / DOS MZ. + {0x7F, 'E', 'L', 'F'}, // ELF. + {0xFE, 0xED, 0xFA, 0xCE}, // Mach-O 32-bit BE. + {0xFE, 0xED, 0xFA, 0xCF}, // Mach-O 64-bit BE. + {0xCF, 0xFA, 0xED, 0xFE}, // Mach-O 64-bit LE. + {0xCE, 0xFA, 0xED, 0xFE}, // Mach-O 32-bit LE. + {0xCA, 0xFE, 0xBA, 0xBE}, // Mach-O fat binary. + } + for _, m := range magics { + if bytes.HasPrefix(probe, m) { + return true + } + } + return false +} + +// dispatchShebang parses probe's shebang line and execs the resolved +// interpreter via os/exec, inheriting the parent runner's cwd, env, and +// stdio. Returns interp.ExitStatus on non-zero interpreter exit so the +// parent interpreter sees it as a normal non-zero status. +func dispatchShebang(ctx context.Context, scriptPath string, probe []byte, args []string) error { + sb, err := parseShebang(probe) + if err != nil { + hc := interp.HandlerCtx(ctx) + fmt.Fprintf(hc.Stderr, "crush: %s: %s\n", scriptPath, err) + return interp.ExitStatus(126) + } + + interpreter, err := resolveInterpreter(sb.interpreter) + if err != nil { + hc := interp.HandlerCtx(ctx) + fmt.Fprintf(hc.Stderr, "crush: %s: %s\n", scriptPath, err) + return interp.ExitStatus(127) + } + + cmdArgs := append([]string{}, sb.args...) + cmdArgs = append(cmdArgs, scriptPath) + cmdArgs = append(cmdArgs, args[1:]...) + + cmd := exec.CommandContext(ctx, interpreter, cmdArgs...) + hc := interp.HandlerCtx(ctx) + cmd.Dir = hc.Dir + cmd.Env = execEnvList(hc.Env) + cmd.Stdin = hc.Stdin + cmd.Stdout = hc.Stdout + cmd.Stderr = hc.Stderr + isolateProcess(cmd) + + if err := cmd.Run(); err != nil { + var exitErr *exec.ExitError + if errors.As(err, &exitErr) { + code := exitErr.ExitCode() + if code < 0 { + code = 1 + } + return interp.ExitStatus(uint8(code)) + } + return err + } + return nil +} + +// resolveInterpreter tries the literal shebang path first, then falls back +// to PATH-lookup on its basename — but only when the literal path is +// genuinely missing. A file that exists but fails stat for another reason +// (EACCES, ELOOP, etc.) surfaces the real error: silently resolving a +// different binary off PATH in that case would hide a real problem and +// produce surprising behavior for the user. +// +// The permissive fallback is what makes #!/bin/bash portable to Windows +// boxes where Git for Windows puts bash.exe on PATH but there is no +// /bin/bash on disk. +func resolveInterpreter(path string) (string, error) { + _, statErr := os.Stat(path) + if statErr == nil { + return path, nil + } + if !errors.Is(statErr, fs.ErrNotExist) { + return "", statErr + } + + base := filepath.Base(path) + if base == "" || base == path && !strings.ContainsAny(path, `/\`) { + // Already a bare name — just do a PATH lookup. + resolved, err := exec.LookPath(path) + if err != nil { + return "", fmt.Errorf("interpreter %q not found in PATH", path) + } + return resolved, nil + } + resolved, err := exec.LookPath(base) + if err != nil { + return "", fmt.Errorf("interpreter %q not found and %q not in PATH", path, base) + } + slog.Debug("Shebang interpreter not found; falling back to PATH", + "requested", path, "resolved", resolved) + return resolved, nil +} + +// shebang captures the parsed `#!` line. interpreter is the program to +// invoke; args is the list of extra arguments to pass before the script +// path. The kernel's single-arg semantics (for literal paths and for env +// without `-S`) is encoded by returning a single-element args slice +// containing the un-tokenized remainder. +type shebang struct { + interpreter string + args []string +} + +// parseShebang extracts the interpreter invocation from probe. It tolerates +// CRLF line endings and a single leading space between `#!` and the path. +// env special-cases: `/usr/bin/env NAME [args...]` unwraps to NAME with +// kernel single-arg semantics; `-S` enables tokenized argument splitting. +func parseShebang(probe []byte) (*shebang, error) { + if !hasShebang(probe) { + return nil, errors.New("not a shebang") + } + line := probe[2:] + // Take up to the first newline. + if idx := bytes.IndexByte(line, '\n'); idx >= 0 { + line = line[:idx] + } + // Strip trailing CR (CRLF-authored scripts). + line = bytes.TrimRight(line, "\r") + // Strip leading whitespace ("#! /usr/bin/env bash" is legal). + line = bytes.TrimLeft(line, " \t") + if len(line) == 0 { + return nil, errors.New("empty shebang") + } + + var pathStr, rest string + if idx := bytes.IndexAny(line, " \t"); idx >= 0 { + pathStr = string(line[:idx]) + rest = strings.TrimLeft(string(line[idx+1:]), " \t") + } else { + pathStr = string(line) + } + + if isEnvShebang(pathStr) { + return parseEnvShebang(rest) + } + + // Literal-path shebang: kernel semantics pass the remainder as a + // single argv[1], not tokenized. + sb := &shebang{interpreter: pathStr} + if rest != "" { + sb.args = []string{rest} + } + return sb, nil +} + +// isEnvShebang reports whether the shebang path targets `env`. We accept +// both common absolute paths and a bare `env` so that unusual setups +// (NixOS, BSDs) still work. +func isEnvShebang(p string) bool { + if p == "/usr/bin/env" || p == "/bin/env" { + return true + } + return filepath.Base(p) == "env" +} + +// parseEnvShebang handles `/usr/bin/env` rewriting. Without `-S`, the +// remainder after the program name is a single argv[1] (kernel +// single-arg semantics via env, even though real env would fail to find a +// program named "bash -x"). With `-S`, the remainder is tokenized on +// whitespace. Any other `env` flag is rejected — forwarding unknown flags +// to a /usr/bin/env on disk is a subtle portability footgun we don't want. +func parseEnvShebang(rest string) (*shebang, error) { + if rest == "" { + return nil, errors.New("env: missing program name") + } + + useSplit := false + if strings.HasPrefix(rest, "-") { + var flag, after string + if idx := strings.IndexAny(rest, " \t"); idx >= 0 { + flag = rest[:idx] + after = strings.TrimLeft(rest[idx+1:], " \t") + } else { + flag = rest + after = "" + } + if flag != "-S" { + return nil, fmt.Errorf("unsupported env flag: %s", flag) + } + useSplit = true + rest = after + if rest == "" { + return nil, errors.New("env -S requires a program") + } + } + + if rest == "" { + return nil, errors.New("env: missing program name") + } + + var prog, remainder string + if idx := strings.IndexAny(rest, " \t"); idx >= 0 { + prog = rest[:idx] + remainder = strings.TrimLeft(rest[idx+1:], " \t") + } else { + prog = rest + } + + sb := &shebang{interpreter: prog} + if remainder != "" { + if useSplit { + sb.args = strings.Fields(remainder) + } else { + sb.args = []string{remainder} + } + } + return sb, nil +} + +// runShellSource parses path's contents as POSIX shell and runs it +// in-process via a nested interp.Runner. It reuses the parent runner's cwd, +// env, and stdio, and rebuilds the Crush handler stack so builtins and the +// dispatch handler itself remain available to anything the script invokes. +// Positional parameters ($1, $2, …) come from args[1:]. +// +// This is the only branch that reads the full file; probeFile keeps its +// read to probeWindow bytes so the binary/shebang paths never touch more +// than 128 bytes of I/O. +func runShellSource(ctx context.Context, path string, args []string, blockFuncs []BlockFunc) error { + data, err := os.ReadFile(path) + if err != nil { + return err + } + + file, err := syntax.NewParser().Parse(bytes.NewReader(data), path) + if err != nil { + return fmt.Errorf("could not parse %s: %w", path, err) + } + + hc := interp.HandlerCtx(ctx) + + opts := []interp.RunnerOption{ + interp.StdIO(hc.Stdin, hc.Stdout, hc.Stderr), + interp.Interactive(false), + interp.Env(hc.Env), + interp.Dir(hc.Dir), + execHandlerOption(blockFuncs), + } + if len(args) > 1 { + // Params with a leading "--" avoids any of args[1:] being + // misinterpreted as set-options (e.g. a user passing "-e" as + // a positional arg to their script). + params := append([]string{"--"}, args[1:]...) + opts = append(opts, interp.Params(params...)) + } + + runner, err := interp.New(opts...) + if err != nil { + return fmt.Errorf("could not build runner for %s: %w", path, err) + } + return runner.Run(ctx, file) +} + +// execEnvList converts an expand.Environ to the []string form that +// os/exec.Cmd.Env expects. Only exported string variables are included, +// matching what a real shell would pass to a child process. +func execEnvList(env expand.Environ) []string { + var out []string + env.Each(func(name string, vr expand.Variable) bool { + if vr.Exported && vr.Kind == expand.String { + out = append(out, name+"="+vr.Str) + } + return true + }) + return out +} diff --git a/internal/shell/dispatch_test.go b/internal/shell/dispatch_test.go new file mode 100644 index 0000000000000000000000000000000000000000..b1fb4dc650e461d3ba425ce865ad40e3ff16e82a --- /dev/null +++ b/internal/shell/dispatch_test.go @@ -0,0 +1,598 @@ +package shell + +import ( + "bytes" + "crypto/rand" + "encoding/hex" + "os" + "os/exec" + "path/filepath" + "reflect" + "runtime" + "strings" + "testing" +) + +// writeScript is a small helper that drops a file with the given contents +// and executable mode into dir. Tests that need exec semantics rely on the +// 0o755 mode on Unix; Windows ignores file modes but doesn't need them +// because dispatch decides what to do from file contents, not permissions. +func writeScript(t *testing.T, dir, name, contents string) string { + t.Helper() + path := filepath.Join(dir, name) + if err := os.WriteFile(path, []byte(contents), 0o755); err != nil { + t.Fatalf("write %s: %v", name, err) + } + return filepath.ToSlash(path) +} + +// randSuffix returns a short random hex string, used to build +// intentionally-unique paths that won't collide with anything on disk. +func randSuffix() string { + var b [4]byte + _, _ = rand.Read(b[:]) + return hex.EncodeToString(b[:]) +} + +// TestIsPathPrefixed covers the classification rules used by the dispatch +// handler to decide whether argv[0] is a file reference. +func TestIsPathPrefixed(t *testing.T) { + cases := []struct { + in string + want bool + }{ + {"./foo.sh", true}, + {"../foo.sh", true}, + {"/usr/bin/foo", true}, + {"foo", false}, + {"foo.sh", false}, + {"jq", false}, + {"", false}, + } + for _, c := range cases { + if got := isPathPrefixed(c.in); got != c.want { + t.Errorf("isPathPrefixed(%q) = %v, want %v", c.in, got, c.want) + } + } + + if runtime.GOOS == "windows" { + winCases := []struct { + in string + want bool + }{ + {`C:\foo\bar.exe`, true}, + {`C:/foo/bar.exe`, true}, + {`c:\foo`, true}, + {`Z:/x`, true}, + {`C:`, false}, // just a drive, no path. + {`\\server\share`, true}, + } + for _, c := range winCases { + if got := isPathPrefixed(c.in); got != c.want { + t.Errorf("isPathPrefixed(%q) = %v, want %v", c.in, got, c.want) + } + } + } +} + +// TestParseShebang covers the shebang grammar: literal paths, env, +// env -S, kernel single-arg semantics, CRLF tolerance, and every +// enumerated error case. +func TestParseShebang(t *testing.T) { + type want struct { + interp string + args []string + errSub string // substring expected in error message (empty → no error) + } + cases := []struct { + name string + in string + want want + }{ + { + name: "literal-no-args", + in: "#!/bin/bash\necho body\n", + want: want{interp: "/bin/bash"}, + }, + { + name: "literal-kernel-single-arg", + in: "#!/bin/bash -x -y\n", + want: want{interp: "/bin/bash", args: []string{"-x -y"}}, + }, + { + name: "env-basic", + in: "#!/usr/bin/env bash\n", + want: want{interp: "bash"}, + }, + { + name: "env-kernel-single-arg", + in: "#!/usr/bin/env bash -x\n", + want: want{interp: "bash", args: []string{"-x"}}, + }, + { + name: "env-dash-S-splits", + in: "#!/usr/bin/env -S bash -x\n", + want: want{interp: "bash", args: []string{"-x"}}, + }, + { + name: "env-dash-S-multi-args", + in: "#!/usr/bin/env -S bash -x --noprofile\n", + want: want{interp: "bash", args: []string{"-x", "--noprofile"}}, + }, + { + name: "leading-space", + in: "#! /usr/bin/env bash\n", + want: want{interp: "bash"}, + }, + { + name: "crlf", + in: "#!/bin/bash\r\n", + want: want{interp: "/bin/bash"}, + }, + { + name: "bare-env-name", + in: "#!env bash\n", + want: want{interp: "bash"}, + }, + { + name: "empty-after-hashbang", + in: "#!\n", + want: want{errSub: "empty shebang"}, + }, + { + name: "env-alone", + in: "#!/usr/bin/env\n", + want: want{errSub: "missing program name"}, + }, + { + name: "env-dash-S-alone", + in: "#!/usr/bin/env -S\n", + want: want{errSub: "env -S requires a program"}, + }, + { + name: "env-unknown-flag", + in: "#!/usr/bin/env -x bash\n", + want: want{errSub: "unsupported env flag"}, + }, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + sb, err := parseShebang([]byte(c.in)) + if c.want.errSub != "" { + if err == nil || !strings.Contains(err.Error(), c.want.errSub) { + t.Fatalf("expected error containing %q, got: %v", c.want.errSub, err) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if sb.interpreter != c.want.interp { + t.Errorf("interpreter = %q, want %q", sb.interpreter, c.want.interp) + } + if !equalStringSlice(sb.args, c.want.args) { + t.Errorf("args = %v, want %v", sb.args, c.want.args) + } + }) + } +} + +func equalStringSlice(a, b []string) bool { + if len(a) == 0 && len(b) == 0 { + return true + } + return reflect.DeepEqual(a, b) +} + +// TestIsBinary covers the NUL-byte and magic-byte classification used to +// keep compiled executables off the in-process shell-source path. +func TestIsBinary(t *testing.T) { + cases := []struct { + name string + in []byte + want bool + }{ + {"shell", []byte("echo hi\n"), false}, + {"nul", []byte("hello\x00world"), true}, + {"elf", []byte{0x7F, 'E', 'L', 'F', 0x02, 0x01}, true}, + {"mz", []byte("MZ\x90\x00"), true}, + {"macho-64-le", []byte{0xCF, 0xFA, 0xED, 0xFE}, true}, + {"short-non-binary", []byte("a"), false}, + } + for _, c := range cases { + if got := isBinary(c.in); got != c.want { + t.Errorf("%s: isBinary = %v, want %v", c.name, got, c.want) + } + } +} + +// TestDispatch_ShellSourceNoShebang exercises the in-process shell-source +// branch: a file without a shebang runs via a nested runner and sees +// positional params from argv[1:]. +func TestDispatch_ShellSourceNoShebang(t *testing.T) { + dir := t.TempDir() + script := writeScript(t, dir, "args.sh", `echo "$1 $2"`) + + var stdout bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: script + " alpha beta", + Cwd: dir, + Stdout: &stdout, + }) + if err != nil { + t.Fatalf("Run returned error: %v", err) + } + if got := stdout.String(); got != "alpha beta\n" { + t.Fatalf("stdout = %q, want %q", got, "alpha beta\n") + } +} + +// TestDispatch_EmptyFile confirms a zero-byte script runs as empty shell +// source (exit 0, no output). +func TestDispatch_EmptyFile(t *testing.T) { + dir := t.TempDir() + script := writeScript(t, dir, "empty.sh", "") + + var stdout, stderr bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: script, + Cwd: dir, + Stdout: &stdout, + Stderr: &stderr, + }) + if err != nil { + t.Fatalf("Run returned error: %v (stderr=%q)", err, stderr.String()) + } + if stdout.Len() != 0 || stderr.Len() != 0 { + t.Fatalf("expected empty output, got stdout=%q stderr=%q", stdout.String(), stderr.String()) + } +} + +// TestDispatch_ShellSourceComposesWithPipe confirms the dispatch handler +// plays nicely with mvdan's pipeline logic: a shell-source script on the +// left feeds the jq builtin on the right. +func TestDispatch_ShellSourceComposesWithPipe(t *testing.T) { + dir := t.TempDir() + script := writeScript(t, dir, "emit.sh", `printf '"value"'`) + + var stdout bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: script + ` | jq -r .`, + Cwd: dir, + Stdout: &stdout, + }) + if err != nil { + t.Fatalf("Run returned error: %v", err) + } + if got := stdout.String(); got != "value\n" { + t.Fatalf("stdout = %q, want %q", got, "value\n") + } +} + +// TestDispatch_MissingFile returns a clean error for a non-existent path. +func TestDispatch_MissingFile(t *testing.T) { + dir := t.TempDir() + missing := filepath.Join(dir, "nope.sh") + err := Run(t.Context(), RunOptions{ + Command: missing, + Cwd: dir, + }) + if err == nil { + t.Fatal("expected error for missing script, got nil") + } +} + +// TestDispatch_DirectoryNotFile surfaces a distinct error when the path +// resolves to a directory. +func TestDispatch_DirectoryNotFile(t *testing.T) { + dir := t.TempDir() + subDir := filepath.Join(dir, "adir") + if err := os.MkdirAll(subDir, 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + + var stderr bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: "./adir", + Cwd: dir, + Stderr: &stderr, + }) + if err == nil { + t.Fatal("expected error when invoking a directory, got nil") + } + if !strings.Contains(err.Error(), "is a directory") { + t.Fatalf("expected 'is a directory' in error, got: %v", err) + } +} + +// TestDispatch_BashShebang runs a #!/bin/bash script via os/exec. Skipped +// if bash isn't available (rare in CI, but keep the test robust). +func TestDispatch_BashShebang(t *testing.T) { + bash, err := exec.LookPath("bash") + if err != nil { + t.Skipf("bash not in PATH: %v", err) + } + _ = bash + + dir := t.TempDir() + script := writeScript(t, dir, "bash-echo.sh", "#!/usr/bin/env bash\necho bashout\n") + + var stdout, stderr bytes.Buffer + err = Run(t.Context(), RunOptions{ + Command: script, + Cwd: dir, + Stdout: &stdout, + Stderr: &stderr, + }) + if err != nil { + t.Fatalf("Run returned error: %v (stderr=%q)", err, stderr.String()) + } + if got := stdout.String(); got != "bashout\n" { + t.Fatalf("stdout = %q, want %q", got, "bashout\n") + } +} + +// TestDispatch_ShebangPassesExitCode maps interpreter exit codes through to +// interp.ExitStatus so the caller can inspect them with ExitCode. +func TestDispatch_ShebangPassesExitCode(t *testing.T) { + if _, err := exec.LookPath("bash"); err != nil { + t.Skipf("bash not in PATH: %v", err) + } + dir := t.TempDir() + script := writeScript(t, dir, "fail.sh", "#!/usr/bin/env bash\nexit 5\n") + + err := Run(t.Context(), RunOptions{ + Command: script, + Cwd: dir, + }) + if err == nil { + t.Fatal("expected non-nil error from exit 5") + } + if code := ExitCode(err); code != 5 { + t.Fatalf("ExitCode = %d, want 5", code) + } +} + +// TestDispatch_MissingInterpreter surfaces a clear error (and non-zero +// exit) when the shebang points to a binary that doesn't exist and has +// no PATH fallback. +func TestDispatch_MissingInterpreter(t *testing.T) { + dir := t.TempDir() + script := writeScript(t, dir, "bad.sh", "#!/no/such/interpreter-"+randSuffix()+"\n:\n") + + var stderr bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: script, + Cwd: dir, + Stderr: &stderr, + }) + if err == nil { + t.Fatal("expected error for missing interpreter, got nil") + } + if ExitCode(err) == 0 { + t.Fatalf("expected non-zero exit code, got 0") + } + if !strings.Contains(stderr.String(), "not found") { + t.Fatalf("expected 'not found' in stderr, got: %q", stderr.String()) + } +} + +// TestDispatch_BarePathNotHandled confirms the handler ignores +// non-path-prefixed argv[0] entirely: a benign bare `true` command must +// not try to open a file in cwd. If dispatch were (incorrectly) firing +// on bare commands, this test would see probeFile's ENOENT. +func TestDispatch_BarePathNotHandled(t *testing.T) { + dir := t.TempDir() + err := Run(t.Context(), RunOptions{ + Command: "true", + Cwd: dir, + }) + if err != nil { + t.Fatalf("bare `true` should not trigger dispatch: %v", err) + } +} + +// TestDispatch_ProbeWindowClassifiesByHead confirms that classification is +// done on the first probeWindow bytes even when the file is much larger; +// a file whose head is shell source but whose tail contains NUL bytes is +// classified as shell source, not binary. +func TestDispatch_ProbeWindowClassifiesByHead(t *testing.T) { + dir := t.TempDir() + head := "echo prefix\n" + // Pad past probeWindow, then append some NULs. + padding := strings.Repeat(" ", probeWindow) + contents := head + padding + "\x00\x00\x00" + script := writeScript(t, dir, "long.sh", contents) + + var stdout bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: script, + Cwd: dir, + Stdout: &stdout, + }) + if err != nil { + t.Fatalf("Run returned error: %v", err) + } + if got := stdout.String(); !strings.HasPrefix(got, "prefix\n") { + t.Fatalf("stdout = %q, want prefix %q", got, "prefix\n") + } +} + +// TestDispatch_BinaryPassthroughExecutes copies a real binary from PATH +// into a tempdir, invokes it via a path-prefixed argv[0], and verifies it +// ran — i.e. the binary branch correctly returns through `next` to the +// default exec handler. It skips on Windows, where the stock binaries +// don't share names and the Go test binary approach is heavier than this +// test deserves. +// +// The copy keeps the source's filename. Some coreutils builds (Nix's, for +// one) ship every utility as a single multi-call binary that picks its +// behaviour from argv[0], so a renamed copy of `true` exits non-zero with +// "unknown program" instead of succeeding. +func TestDispatch_BinaryPassthroughExecutes(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("relies on a Unix-style PATH binary") + } + src, err := exec.LookPath("true") + if err != nil { + t.Skipf("no `true` binary on PATH: %v", err) + } + data, err := os.ReadFile(src) + if err != nil { + t.Fatalf("read %s: %v", src, err) + } + dir := t.TempDir() + dst := filepath.Join(dir, filepath.Base(src)) + if err := os.WriteFile(dst, data, 0o755); err != nil { + t.Fatalf("write %s: %v", dst, err) + } + + runErr := Run(t.Context(), RunOptions{ + Command: dst, + Cwd: dir, + // Default handler needs PATH to resolve dynamic linker / loader + // helpers on some systems; inherit the process env so the copy + // can actually start. + Env: os.Environ(), + }) + if runErr != nil { + t.Fatalf("expected copy of %s to exit 0, got: %v", src, runErr) + } +} + +// TestDispatch_UnreadableFile confirms an EACCES on the script surfaces +// as a clean error rather than a silent fallback or a mis-classified +// shell-source attempt. POSIX-only: Windows doesn't have the same +// permission model and running as root would bypass the check anyway. +func TestDispatch_UnreadableFile(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("POSIX permission model") + } + if os.Geteuid() == 0 { + t.Skip("root bypasses file mode permission checks") + } + dir := t.TempDir() + script := writeScript(t, dir, "unreadable.sh", "echo nope\n") + if err := os.Chmod(script, 0o000); err != nil { + t.Fatalf("chmod: %v", err) + } + t.Cleanup(func() { _ = os.Chmod(script, 0o644) }) + + err := Run(t.Context(), RunOptions{ + Command: script, + Cwd: dir, + }) + if err == nil { + t.Fatal("expected permission error, got nil") + } + if !strings.Contains(err.Error(), "permission") { + t.Fatalf("expected 'permission' in error, got: %v", err) + } +} + +// TestDispatch_SymlinkLoop confirms that an ELOOP-returning path surfaces +// cleanly. POSIX-only: creating symlinks reliably on Windows requires +// elevated privileges or developer mode, and neither is guaranteed in CI. +func TestDispatch_SymlinkLoop(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("symlink creation requires special privileges on Windows") + } + dir := t.TempDir() + a := filepath.Join(dir, "a") + b := filepath.Join(dir, "b") + if err := os.Symlink(b, a); err != nil { + t.Fatalf("symlink a→b: %v", err) + } + if err := os.Symlink(a, b); err != nil { + t.Fatalf("symlink b→a: %v", err) + } + + err := Run(t.Context(), RunOptions{ + Command: a, + Cwd: dir, + }) + if err == nil { + t.Fatal("expected loop error, got nil") + } + // The exact error varies by OS; any of these message fragments is + // acceptable evidence that the loop was detected. + msg := err.Error() + if !strings.Contains(msg, "too many") && + !strings.Contains(msg, "loop") && + !strings.Contains(msg, "level") { + t.Fatalf("expected symlink-loop-ish error, got: %v", err) + } +} + +// TestResolveInterpreter_PermissiveFallback confirms the key portability +// behavior: a literal shebang path that doesn't exist falls back to a +// PATH-lookup on its basename. This is what makes #!/bin/bash work on a +// Windows box where bash.exe lives somewhere else on PATH. We construct a +// fake PATH in a tempdir rather than depending on what the host has +// installed so the test is deterministic everywhere. +func TestResolveInterpreter_PermissiveFallback(t *testing.T) { + if runtime.GOOS == "windows" { + // exec.LookPath on Windows requires a recognized extension + // (.exe/.bat/.cmd). Producing one of those without a compiler + // run is more ceremony than this smoke test deserves; the + // logic under test is exercised by the Unix run. + t.Skip("Windows PATH lookup requires an extension-matched binary") + } + dir := t.TempDir() + fake := filepath.Join(dir, "bash") + if err := os.WriteFile(fake, []byte("#!/bin/sh\nexit 0\n"), 0o755); err != nil { + t.Fatalf("write fake bash: %v", err) + } + t.Setenv("PATH", dir) + + // Basename must match the fake we planted on PATH; the directory + // prefix must not exist so the literal stat fails. + missingDir := filepath.Join(dir, "definitely-not-here-"+randSuffix()) + resolved, err := resolveInterpreter(filepath.Join(missingDir, "bash")) + if err != nil { + t.Fatalf("expected fallback to succeed, got: %v", err) + } + if resolved != fake { + t.Fatalf("resolved = %q, want %q", resolved, fake) + } +} + +// TestResolveInterpreter_NonENOENTErrorsSurface guards against silently +// falling back to PATH when stat fails for a reason other than the file +// being missing. With a directory at the shebang path, os.Stat succeeds +// (no fallback needed), but with an EACCES'd file it fails with a non- +// ENOENT error that must be surfaced — otherwise we'd silently resolve a +// different binary off PATH and hide the real problem. +func TestResolveInterpreter_NonENOENTErrorsSurface(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("POSIX permission model") + } + if os.Geteuid() == 0 { + t.Skip("root bypasses dir mode permission checks") + } + dir := t.TempDir() + // Put a candidate interpreter inside an unreadable/untraversable dir. + inner := filepath.Join(dir, "private") + if err := os.Mkdir(inner, 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + interp := filepath.Join(inner, "bash") + if err := os.WriteFile(interp, []byte("#!/bin/sh\nexit 0\n"), 0o755); err != nil { + t.Fatalf("write interpreter: %v", err) + } + // Drop search permission on inner so os.Stat(interp) returns EACCES. + if err := os.Chmod(inner, 0o000); err != nil { + t.Fatalf("chmod: %v", err) + } + t.Cleanup(func() { _ = os.Chmod(inner, 0o755) }) + + _, err := resolveInterpreter(interp) + if err == nil { + t.Fatal("expected error for unreadable interpreter, got nil") + } + // Must NOT have silently fallen back — the returned path shouldn't + // be a valid resolution; either way, the error has to surface. + if !strings.Contains(err.Error(), "permission") { + t.Fatalf("expected permission-denied error to surface, got: %v", err) + } +} diff --git a/internal/shell/dispatch_windows_test.go b/internal/shell/dispatch_windows_test.go new file mode 100644 index 0000000000000000000000000000000000000000..529945ad5a33eac5d943bdd2a562a46d8d4ff000 --- /dev/null +++ b/internal/shell/dispatch_windows_test.go @@ -0,0 +1,40 @@ +//go:build windows + +package shell + +import ( + "os" + "path/filepath" + "testing" +) + +// TestResolveInterpreter_PermissiveFallback_Windows is the Windows-native +// counterpart to the POSIX permissive-fallback test. It proves the one +// behavior that makes `#!/bin/bash` hooks work on a stock Windows box +// with Git for Windows installed: when the literal interpreter path does +// not exist, we fall back to a PATH-lookup on the basename and that +// lookup accepts any executable extension Windows honors (here, `.bat`). +// +// We plant a bash.bat in a tempdir rather than a .exe because producing +// a .exe would require a toolchain step; LookPath on Windows resolves +// PATHEXT extensions, so .bat is just as valid for the lookup codepath. +func TestResolveInterpreter_PermissiveFallback_Windows(t *testing.T) { + dir := t.TempDir() + fake := filepath.Join(dir, "bash.bat") + contents := "@echo off\r\nexit /b 0\r\n" + if err := os.WriteFile(fake, []byte(contents), 0o755); err != nil { + t.Fatalf("write fake bash.bat: %v", err) + } + t.Setenv("PATH", dir) + t.Setenv("PATHEXT", ".BAT;.CMD;.EXE") + + // Literal path must be absent so the stat fails with ENOENT. + missing := filepath.Join(dir, "definitely-not-here-"+randSuffix(), "bash") + resolved, err := resolveInterpreter(missing) + if err != nil { + t.Fatalf("expected fallback to succeed, got: %v", err) + } + if resolved != fake { + t.Fatalf("resolved = %q, want %q", resolved, fake) + } +} diff --git a/internal/shell/doc.go b/internal/shell/doc.go new file mode 100644 index 0000000000000000000000000000000000000000..feed86db011d04178a1643d757d618d313c1f579 --- /dev/null +++ b/internal/shell/doc.go @@ -0,0 +1,25 @@ +package shell + +// Example usage of the shell package: +// +// 1. For one-off commands: +// +// shell := shell.NewShell(nil) +// stdout, stderr, err := shell.Exec(context.Background(), "echo hello") +// +// 2. For maintaining state across commands: +// +// shell := shell.NewShell(&shell.Options{ +// WorkingDir: "/tmp", +// Logger: myLogger, +// }) +// shell.Exec(ctx, "export FOO=bar") +// shell.Exec(ctx, "echo $FOO") // Will print "bar" +// +// 3. Managing environment and working directory: +// +// shell := shell.NewShell(nil) +// shell.SetEnv("MY_VAR", "value") +// shell.SetWorkingDir("/tmp") +// cwd := shell.GetWorkingDir() +// env := shell.GetEnv() diff --git a/internal/shell/exec_unix.go b/internal/shell/exec_unix.go new file mode 100644 index 0000000000000000000000000000000000000000..b3488c24575d26ccb2a4c80a1674b7770d520341 --- /dev/null +++ b/internal/shell/exec_unix.go @@ -0,0 +1,108 @@ +//go:build !windows + +package shell + +import ( + "context" + "fmt" + "io" + "os/exec" + "syscall" + "time" + + "mvdan.cc/sh/v3/interp" +) + +// defaultKillTimeout matches mvdan's DefaultExecHandler default. Extracted +// so the coupling to upstream is explicit rather than a buried literal. +const defaultKillTimeout = 2 * time.Second + +// isolateProcess sets SysProcAttr so the child runs in its own session, +// fully detached from Crush's controlling terminal. This prevents shells +// like zsh from grabbing the TTY for job control, which would cause +// "suspended (tty input)" and corrupt Bubble Tea's terminal rendering. +// It also prevents child processes from sending signals to Crush's process +// group. +func isolateProcess(cmd *exec.Cmd) { + if cmd.SysProcAttr == nil { + cmd.SysProcAttr = &syscall.SysProcAttr{} + } + cmd.SysProcAttr.Setsid = true +} + +// processGroupExecHandler returns an ExecHandlerFunc that replaces +// interp.DefaultExecHandler with one that fully isolates child processes +// from Crush's session and controlling terminal. +// +// Without this, shells like zsh that enable job control when sourcing +// framework files will attempt to take over the TTY, causing SIGTTIN/SIGTTOU +// signals and corrupting the parent terminal state. +// +// The implementation mirrors interp.DefaultExecHandler with two additions: +// isolateProcess(&cmd) after construction, and negative-PID signal targeting +// in the cancellation callback so the entire child process group is killed. +func processGroupExecHandler(killTimeout time.Duration) interp.ExecHandlerFunc { + return func(ctx context.Context, args []string) error { + hc := interp.HandlerCtx(ctx) + path, err := interp.LookPathDir(hc.Dir, hc.Env, args[0]) + if err != nil { + fmt.Fprintln(hc.Stderr, err) + return interp.ExitStatus(127) + } + + cmd := exec.Cmd{ + Path: path, + Args: args, + Env: execEnvList(hc.Env), + Dir: hc.Dir, + Stdin: hc.Stdin, + Stdout: hc.Stdout, + Stderr: hc.Stderr, + } + isolateProcess(&cmd) + + err = cmd.Start() + if err == nil { + stopf := context.AfterFunc(ctx, func() { + if killTimeout <= 0 { + _ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) + return + } + // Signal the child's process group (negative PID) so + // grandchildren also receive it. + _ = syscall.Kill(-cmd.Process.Pid, syscall.SIGINT) + time.Sleep(killTimeout) + _ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) + }) + defer stopf() + + err = cmd.Wait() + } + + return exitStatusFromError(ctx, hc.Stderr, err) + } +} + +// exitStatusFromError translates an exec error into an interp exit status, +// matching the conventions of interp.DefaultExecHandler. Extracted so the +// platform-specific exec handler stays focused on isolation mechanics. +func exitStatusFromError(ctx context.Context, stderr io.Writer, err error) error { + if err == nil { + return nil + } + switch err := err.(type) { + case *exec.ExitError: + if status, ok := err.Sys().(syscall.WaitStatus); ok && status.Signaled() { + if ctx.Err() != nil { + return ctx.Err() + } + return interp.ExitStatus(128 + uint8(status.Signal())) + } + return interp.ExitStatus(uint8(err.ExitCode())) + case *exec.Error: + fmt.Fprintf(stderr, "%v\n", err) + return interp.ExitStatus(127) + default: + return err + } +} diff --git a/internal/shell/exec_windows.go b/internal/shell/exec_windows.go new file mode 100644 index 0000000000000000000000000000000000000000..3609c68e3b05556d2f6ad7ab60c940cc27811453 --- /dev/null +++ b/internal/shell/exec_windows.go @@ -0,0 +1,23 @@ +//go:build windows + +package shell + +import ( + "os/exec" + "time" + + "mvdan.cc/sh/v3/interp" +) + +// defaultKillTimeout matches mvdan's DefaultExecHandler default. +const defaultKillTimeout = 2 * time.Second + +// isolateProcess is a no-op on Windows. Session isolation via Setsid is a +// Unix-only concept; Windows uses CREATE_NEW_PROCESS_GROUP which mvdan's +// default handler already handles adequately. +func isolateProcess(_ *exec.Cmd) {} + +// processGroupExecHandler returns interp.DefaultExecHandler on Windows. +func processGroupExecHandler(killTimeout time.Duration) interp.ExecHandlerFunc { + return interp.DefaultExecHandler(killTimeout) +} diff --git a/internal/shell/expand.go b/internal/shell/expand.go new file mode 100644 index 0000000000000000000000000000000000000000..04bb8c2e337fba452b4a78c064f1fd50ef9d3585 --- /dev/null +++ b/internal/shell/expand.go @@ -0,0 +1,162 @@ +package shell + +import ( + "bytes" + "context" + "fmt" + "io" + "os" + "strings" + "sync/atomic" + + "mvdan.cc/sh/v3/expand" + "mvdan.cc/sh/v3/interp" + "mvdan.cc/sh/v3/syntax" +) + +// maxInnerStderrBytes bounds how much stderr from a failing $(...) is +// surfaced in the returned error, to avoid leaking a secret that happened +// to be embedded in a failing inner command. +const maxInnerStderrBytes = 512 + +// NoUnset controls whether ExpandValue treats unset variables as an +// error. Default false matches bash: $UNSET expands to "". Store true +// to re-enable strict mode globally. Not exposed in crush.json; this is +// an internal escape hatch in case the lenient default turns out to be +// the wrong call. +// +// Declared atomic because ExpandValue is invoked concurrently (multiple +// MCP / LSP / provider loads in flight at startup, hook execution, etc.) +// and an unsynchronised read/write pair is a data race under the Go +// memory model regardless of test-level happens-before reasoning. The +// atomic load on the hot path is negligible against the cost of parsing +// and running through mvdan. +var NoUnset atomic.Bool + +// ExpandValue expands shell-style substitutions in a single config value. +// +// Supported constructs match the bash tool: +// +// - $VAR and ${VAR}. +// - ${VAR:-default} / ${VAR:+alt} / ${VAR:?msg}. +// - $(command) with full quoting and nesting. +// - escaped and quoted strings ("...", '...'). +// +// Contract: +// +// - Returns exactly one string. No field splitting, no globbing, no +// pathname generation. Multi-word command output is preserved +// verbatim; it is never split into multiple values. +// - Nounset is off by default, matching bash: unset variables expand +// to "". Opt in to strict behaviour per-reference with +// ${VAR:?msg}, which errors loudly when VAR is unset regardless of +// the global toggle. Flip the global default via +// shell.NoUnset.Store(true) as an internal escape hatch. +// - Embedded whitespace and newlines in the input are preserved +// verbatim. Command substitution strips trailing newlines only +// (POSIX), never leading or internal whitespace. +// - Errors wrap the failing inner command's exit code and a bounded +// prefix of its stderr. Callers that surface the error to users +// should additionally scrub it for the original template text. +func ExpandValue(ctx context.Context, value string, env []string) (string, error) { + // Fast path: a value with no shell metacharacters expands to itself. + // Parsing and running it through the interpreter would produce the + // same string at significant cost. Most config values (literal API + // keys, fixed URLs) hit this path, which matters because ExpandValue + // runs over every provider/MCP/LSP value on each config reload. + if !strings.ContainsAny(value, "$`\\'\"") { + return value, nil + } + + // Parse the value as a here-doc style word: no word splitting, no + // globbing, but full support for $VAR, ${VAR...}, $(...), and + // quoted/escaped strings. + word, err := syntax.NewParser().Document(strings.NewReader(value)) + if err != nil { + return "", fmt.Errorf("parse: %w", err) + } + + // Build a minimal Shell value purely to reuse its handler chain + // (builtins, block funcs, optional Go coreutils) inside $(...). + // We deliberately skip NewShell so the passed-in env is used + // verbatim, with no CRUSH/AGENT/AI_AGENT injection: callers of + // ExpandValue control the env, and nounset must treat any name + // not in env as unset. + // + // The working directory is only needed for $(...) command + // substitution, so we resolve it lazily on first use. Values that + // only reference variables or use quoting (the common case) avoid + // the os.Getwd() syscall entirely. + s := &Shell{ + env: env, + logger: noopLogger{}, + } + cwdResolved := false + + strict := NoUnset.Load() + + var stderrBuf bytes.Buffer + cfg := &expand.Config{ + Env: expand.ListEnviron(env...), + NoUnset: strict, + CmdSubst: func(w io.Writer, cs *syntax.CmdSubst) error { + if !cwdResolved { + s.cwd, _ = os.Getwd() + cwdResolved = true + } + stderrBuf.Reset() + runnerOpts := []interp.RunnerOption{ + interp.StdIO(nil, w, &stderrBuf), + interp.Interactive(false), + interp.Env(expand.ListEnviron(env...)), + interp.Dir(s.cwd), + execHandlerOption(s.blockFuncs), + } + if strict { + // Match the outer NoUnset: an unset $VAR inside + // $(...) is also an error, not a silent empty. + runnerOpts = append(runnerOpts, interp.Params("-u")) + } + runner, rerr := interp.New(runnerOpts...) + if rerr != nil { + return rerr + } + if rerr := runner.Run(ctx, &syntax.File{Stmts: cs.Stmts}); rerr != nil { + return wrapCmdSubstErr(rerr, stderrBuf.Bytes()) + } + return nil + }, + // ReadDir / ReadDir2 left nil: globbing is disabled. + } + + return expand.Document(cfg, word) +} + +// wrapCmdSubstErr attaches a bounded prefix of the inner command's stderr +// to the original error, if any. +func wrapCmdSubstErr(err error, stderrBytes []byte) error { + msg := sanitizeStderr(stderrBytes) + if msg == "" { + return err + } + return fmt.Errorf("%w: %s", err, msg) +} + +// sanitizeStderr trims, bounds, and scrubs non-printable bytes from the +// stderr of a failing command so the result is safe to include in an +// error message shown to the user. +func sanitizeStderr(b []byte) string { + b = bytes.TrimRight(b, "\n") + if len(b) > maxInnerStderrBytes { + b = b[:maxInnerStderrBytes] + } + out := make([]byte, len(b)) + for i, c := range b { + if c == '\t' || c == '\n' || (c >= 0x20 && c < 0x7f) { + out[i] = c + } else { + out[i] = '?' + } + } + return string(out) +} diff --git a/internal/shell/expand_test.go b/internal/shell/expand_test.go new file mode 100644 index 0000000000000000000000000000000000000000..77d04454abf4447405e5b319b785dcff6e9fc760 --- /dev/null +++ b/internal/shell/expand_test.go @@ -0,0 +1,241 @@ +package shell + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestExpandValue_Success(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + value string + env []string + want string + }{ + { + name: "plain string round trip", + value: "hello world", + want: "hello world", + }, + { + name: "plain var from env", + value: "$FOO", + env: []string{"FOO=bar"}, + want: "bar", + }, + { + name: "braced var from env", + value: "pre-${FOO}-post", + env: []string{"FOO=bar"}, + want: "pre-bar-post", + }, + { + name: "default syntax on unset", + value: "${MISSING:-fallback}", + want: "fallback", + }, + { + name: "default syntax on set preserves value", + value: "${SET:-fallback}", + env: []string{"SET=used"}, + want: "used", + }, + { + name: "command substitution", + value: "$(echo hi)", + want: "hi", + }, + { + name: "command substitution preserves internal spaces", + value: `$(echo "a b")`, + want: "a b", + }, + { + name: "command substitution strips only trailing newline", + value: "$(printf 'a\\nb\\n')", + want: "a\nb", + }, + { + name: "literal spaces around cmdsubst are preserved", + value: " $(echo v) ", + want: " v ", + }, + { + name: "paren inside quoted arg to echo", + value: `$(echo ")")`, + want: ")", + }, + { + name: "nested command substitution", + value: "$(echo $(echo hi))", + want: "hi", + }, + { + name: "glob-like input round trips unchanged", + value: "*.go", + want: "*.go", + }, + { + name: "tilde round trips unchanged", + value: "~", + want: "~", + }, + { + name: "env var inside cmdsubst", + value: `$(printf '%s' "$FOO")`, + env: []string{"FOO=bar"}, + want: "bar", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got, err := ExpandValue(t.Context(), tc.value, tc.env) + require.NoError(t, err) + require.Equal(t, tc.want, got) + }) + } +} + +func TestExpandValue_Errors(t *testing.T) { + t.Parallel() + + t.Run("unset var expands to empty under lenient default", func(t *testing.T) { + t.Parallel() + got, err := ExpandValue(t.Context(), "$MISSING", nil) + require.NoError(t, err) + require.Equal(t, "", got) + }) + + t.Run("unset var inside braces expands to empty", func(t *testing.T) { + t.Parallel() + got, err := ExpandValue(t.Context(), "${MISSING}", nil) + require.NoError(t, err) + require.Equal(t, "", got) + }) + + t.Run("unset var inside cmdsubst expands to empty", func(t *testing.T) { + t.Parallel() + got, err := ExpandValue(t.Context(), `$(printf '%s' "$MISSING")`, nil) + require.NoError(t, err) + require.Equal(t, "", got) + }) + + t.Run("bad syntax returns error", func(t *testing.T) { + t.Parallel() + _, err := ExpandValue(t.Context(), "$(", nil) + require.Error(t, err) + }) + + t.Run("inner non-zero exit returns error with exit code", func(t *testing.T) { + t.Parallel() + _, err := ExpandValue(t.Context(), "$(false)", nil) + require.Error(t, err) + require.Contains(t, err.Error(), "exit status 1") + }) + + t.Run("inner explicit exit code is surfaced", func(t *testing.T) { + t.Parallel() + _, err := ExpandValue(t.Context(), "$(exit 7)", nil) + require.Error(t, err) + require.Contains(t, err.Error(), "exit status 7") + }) + + t.Run("inner stderr is surfaced", func(t *testing.T) { + t.Parallel() + _, err := ExpandValue( + t.Context(), + `$(printf 'boom\n' 1>&2; exit 1)`, + nil, + ) + require.Error(t, err) + require.Contains(t, err.Error(), "boom") + }) + + t.Run("inner stderr is truncated to byte budget", func(t *testing.T) { + t.Parallel() + // Emit more than maxInnerStderrBytes bytes of 'X' on stderr. + long := strings.Repeat("X", maxInnerStderrBytes*2) + _, err := ExpandValue( + t.Context(), + `$(printf '`+long+`' 1>&2; exit 1)`, + nil, + ) + require.Error(t, err) + require.NotContains( + t, + err.Error(), + strings.Repeat("X", maxInnerStderrBytes+1), + "stderr should be bounded", + ) + }) +} + +// TestExpandValue_StrictToggle pins the NoUnset escape hatch: when a +// caller flips strict mode on, bare $UNSET must error instead of +// expanding to the empty string. Must not run in parallel: it mutates +// the package-level NoUnset atomic, so a parallel peer observing the +// flipped value would break the lenient default other tests assume. +func TestExpandValue_StrictToggle(t *testing.T) { + NoUnset.Store(true) + t.Cleanup(func() { NoUnset.Store(false) }) + + _, err := ExpandValue(t.Context(), "$UNSET", nil) + require.Error(t, err) + + _, err = ExpandValue(t.Context(), "${UNSET}", nil) + require.Error(t, err) + + _, err = ExpandValue(t.Context(), `$(printf '%s' "$UNSET")`, nil) + require.Error(t, err) +} + +// TestExpandValue_RequiredOptIn pins the per-reference opt-in strict +// idiom ${VAR:?msg}: it must error whether or not the global NoUnset +// toggle is on, so config authors can mark individual credentials as +// required without flipping the global default. +func TestExpandValue_RequiredOptIn(t *testing.T) { + t.Parallel() + + _, err := ExpandValue(t.Context(), "${REQUIRED:?must be set}", nil) + require.Error(t, err) + require.Contains(t, err.Error(), "must be set") + + got, err := ExpandValue( + t.Context(), + "${REQUIRED:?must be set}", + []string{"REQUIRED=ok"}, + ) + require.NoError(t, err) + require.Equal(t, "ok", got) +} + +func TestSanitizeStderr(t *testing.T) { + t.Parallel() + + t.Run("trims trailing newlines", func(t *testing.T) { + t.Parallel() + require.Equal(t, "hi", sanitizeStderr([]byte("hi\n\n"))) + }) + + t.Run("preserves tabs and embedded newlines", func(t *testing.T) { + t.Parallel() + require.Equal(t, "a\tb\nc", sanitizeStderr([]byte("a\tb\nc"))) + }) + + t.Run("replaces control characters", func(t *testing.T) { + t.Parallel() + require.Equal(t, "a?b", sanitizeStderr([]byte{'a', 0x01, 'b'})) + }) + + t.Run("bounds output", func(t *testing.T) { + t.Parallel() + got := sanitizeStderr([]byte(strings.Repeat("x", maxInnerStderrBytes*2))) + require.Len(t, got, maxInnerStderrBytes) + }) +} diff --git a/internal/shell/isolation_unix_test.go b/internal/shell/isolation_unix_test.go new file mode 100644 index 0000000000000000000000000000000000000000..e8d6da7e357f5349bf01accdc215dbbe5e4dcbb7 --- /dev/null +++ b/internal/shell/isolation_unix_test.go @@ -0,0 +1,143 @@ +//go:build !windows + +package shell + +import ( + "bytes" + "context" + "os" + "os/exec" + "os/signal" + "strings" + "syscall" + "testing" + "time" +) + +// TestProcessIsolation_SignalDoesNotReachParent verifies that a child +// sending SIGINT to its own process group does not deliver the signal to +// the parent (test) process. Without Setsid isolation, the child shares +// the parent's process group and the signal would kill the test. +func TestProcessIsolation_SignalDoesNotReachParent(t *testing.T) { + t.Parallel() + + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, syscall.SIGINT) + defer signal.Stop(sigCh) + + var stderr bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: `sh -c 'kill -INT -$$'`, + Cwd: t.TempDir(), + Env: os.Environ(), + Stderr: &stderr, + }) + + if err == nil { + t.Log("child exited 0 after self-signal (unexpected but not a failure)") + } + + select { + case sig := <-sigCh: + t.Fatalf("SIGINT leaked to parent process: %v (stderr=%q)", sig, stderr.String()) + case <-time.After(200 * time.Millisecond): + } +} + +// TestProcessIsolation_TTYAccessDoesNotBlock verifies that a child trying +// to read from /dev/tty does not block or suspend the parent. With Setsid, +// the child has no controlling terminal, so /dev/tty access fails immediately +// with ENXIO. +func TestProcessIsolation_TTYAccessDoesNotBlock(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(t.Context(), 3*time.Second) + defer cancel() + + var stderr bytes.Buffer + done := make(chan error, 1) + go func() { + done <- Run(ctx, RunOptions{ + Command: `sh -c 'cat < /dev/tty'`, + Cwd: t.TempDir(), + Env: os.Environ(), + Stderr: &stderr, + }) + }() + + select { + case err := <-done: + if err == nil { + t.Fatal("expected error from /dev/tty access in detached session, got nil") + } + case <-time.After(3 * time.Second): + t.Fatalf("/dev/tty access blocked for >3s; session isolation may be broken (stderr=%q)", stderr.String()) + } +} + +// TestProcessIsolation_ZshJobControlDoesNotSuspend simulates the exact +// scenario that caused the original bug: running zsh with job control +// enabled. Without session isolation, zsh tries to become the foreground +// process group leader on the controlling terminal, gets SIGTTIN, and +// suspends. Skips if zsh is not installed. +func TestProcessIsolation_ZshJobControlDoesNotSuspend(t *testing.T) { + t.Parallel() + + if _, err := exec.LookPath("zsh"); err != nil { + t.Skip("zsh not installed") + } + + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + + var stdout, stderr bytes.Buffer + done := make(chan error, 1) + go func() { + done <- Run(ctx, RunOptions{ + Command: `zsh -i -c 'echo ok'`, + Cwd: t.TempDir(), + Env: os.Environ(), + Stdout: &stdout, + Stderr: &stderr, + }) + }() + + select { + case err := <-done: + if err != nil { + t.Logf("zsh exited with error (may be expected): %v", err) + } + out := strings.TrimSpace(stdout.String()) + if !strings.Contains(out, "ok") { + t.Errorf("expected 'ok' in stdout, got %q (stderr=%q)", out, stderr.String()) + } + case <-time.After(5 * time.Second): + t.Fatalf("zsh -i -c did not complete within 5s; likely suspended on TTY (stderr=%q)", stderr.String()) + } +} + +// TestProcessIsolation_ChildProcessGroupKill verifies that when context +// is cancelled, the signal reaches the entire child process group (not +// just the direct child). +func TestProcessIsolation_ChildProcessGroupKill(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(t.Context(), 500*time.Millisecond) + defer cancel() + + start := time.Now() + err := Run(ctx, RunOptions{ + Command: `sh -c 'sleep 60 & wait'`, + Cwd: t.TempDir(), + Env: os.Environ(), + }) + elapsed := time.Since(start) + + if err == nil { + t.Fatal("expected error from cancelled context") + } + // Allow up to 4s: 500ms context timeout + 2s kill timeout + margin. + if elapsed > 4*time.Second { + t.Fatalf("cancellation took %v; expected prompt process group kill", elapsed) + } +} diff --git a/internal/shell/jq.go b/internal/shell/jq.go new file mode 100644 index 0000000000000000000000000000000000000000..99030b223379195afffdae2b5f73b5de147fea05 --- /dev/null +++ b/internal/shell/jq.go @@ -0,0 +1,348 @@ +package shell + +import ( + "context" + "encoding/json" + "fmt" + "io" + "os" + "strings" + + "github.com/itchyny/gojq" + "mvdan.cc/sh/v3/interp" +) + +const jqUsage = `jq - Go implementation of jq (gojq 0.12.19 builtin) + +Synopsis: + %% echo '{"foo": 128}' | jq '.foo' + +Usage: + jq [OPTIONS] [FILTER] [FILE...] + +Options: + -r, --raw-output output raw strings + -j, --join-output implies -r with no newline delimiter + -c, --compact-output output without pretty-printing + -s, --slurp read all inputs into an array + -n, --null-input use null as input value + -e, --exit-status exit 1 when the last value is false or null + -R, --raw-input read input as raw strings + --arg name value set a string value to a variable + --argjson name value set a JSON value to a variable + -h, --help display this help +` + +// handleJQ implements the jq builtin using gojq. It supports a subset of jq +// flags: -r (raw output), -c (compact output), -s (slurp), -n (null input), +// -e (exit status), -R (raw input), and --arg name value. +// +// ctx is polled at each iteration of the output loop and at each reader in +// [readInputs] so that hook timeouts or other cancellations can interrupt +// long-running queries. A cancelled context surfaces as ctx.Err(), not an +// [interp.ExitStatus], so callers (e.g. the hook runner) can distinguish +// "filter exited non-zero" from "we ran out of time". +// +// Note that this is somewhat of a reimplmentation of the CLI of the glorious +// github.com/itchyny/gojq, and we'd ideally get the CLI exposed upstream to +// avoid this falling out of sync. +func handleJQ(ctx context.Context, args []string, stdin io.Reader, stdout, stderr io.Writer) error { + // Fast-fail when ctx is already cancelled so callers don't pay for + // flag parsing and gojq compilation on a doomed request. + if err := ctx.Err(); err != nil { + return err + } + + var ( + rawOutput bool + compact bool + slurp bool + nullInput bool + exitStatus bool + rawInput bool + joinOutput bool + argNames []string + argValues []any + ) + + // Parse flags and extract the query. + var queryStr string + var fileArgs []string + i := 1 // skip "jq" + for i < len(args) { + arg := args[i] + switch { + case arg == "-h" || arg == "--help": + fmt.Fprint(stdout, jqUsage) + return nil + case arg == "-r" || arg == "--raw-output": + rawOutput = true + case arg == "-j" || arg == "--join-output": + joinOutput = true + rawOutput = true + case arg == "-c" || arg == "--compact-output": + compact = true + case arg == "-s" || arg == "--slurp": + slurp = true + case arg == "-n" || arg == "--null-input": + nullInput = true + case arg == "-e" || arg == "--exit-status": + exitStatus = true + case arg == "-R" || arg == "--raw-input": + rawInput = true + case arg == "--arg": + if i+2 >= len(args) { + fmt.Fprintf(stderr, "jq: --arg requires name and value\n") + return interp.ExitStatus(2) + } + argNames = append(argNames, "$"+args[i+1]) + argValues = append(argValues, args[i+2]) + i += 2 + case arg == "--argjson": + if i+2 >= len(args) { + fmt.Fprintf(stderr, "jq: --argjson requires name and value\n") + return interp.ExitStatus(2) + } + var val any + if err := json.Unmarshal([]byte(args[i+2]), &val); err != nil { + fmt.Fprintf(stderr, "jq: invalid JSON for --argjson %s: %s\n", args[i+1], err) + return interp.ExitStatus(2) + } + argNames = append(argNames, "$"+args[i+1]) + argValues = append(argValues, val) + i += 2 + case arg == "--": + i++ + // Remaining args are file arguments. + for i < len(args) { + fileArgs = append(fileArgs, args[i]) + i++ + } + continue + case strings.HasPrefix(arg, "-") && queryStr != "": + fmt.Fprintf(stderr, "jq: unknown option: %s\n", arg) + return interp.ExitStatus(2) + default: + if queryStr == "" { + queryStr = arg + } else { + fileArgs = append(fileArgs, arg) + } + } + i++ + } + + if queryStr == "" { + queryStr = "." + } + + query, err := gojq.Parse(queryStr) + if err != nil { + fmt.Fprintf(stderr, "jq: %s\n", err) + return interp.ExitStatus(3) + } + + opts := []gojq.CompilerOption{ + gojq.WithEnvironLoader(os.Environ), + } + if len(argNames) > 0 { + opts = append(opts, gojq.WithVariables(argNames)) + } + + code, err := gojq.Compile(query, opts...) + if err != nil { + fmt.Fprintf(stderr, "jq: %s\n", err) + return interp.ExitStatus(3) + } + + // Build input values. + inputs, err := readInputs(ctx, stdin, fileArgs, nullInput, rawInput, slurp) + if err != nil { + // Prefer surfacing ctx cancellation verbatim so timeouts are + // distinguishable from user input errors. + if ctxErr := ctx.Err(); ctxErr != nil { + return ctxErr + } + fmt.Fprintf(stderr, "jq: %s\n", err) + return interp.ExitStatus(2) + } + + var lastFalsy bool + for _, input := range inputs { + iter := code.Run(input, argValues...) + for { + // Poll ctx on every value so a long-running filter (e.g. a + // generator over a slurped array) can be interrupted by hook + // timeouts without waiting for iter.Next to yield. + if err := ctx.Err(); err != nil { + return err + } + v, ok := iter.Next() + if !ok { + break + } + if err, ok := v.(error); ok { + fmt.Fprintf(stderr, "jq: %s\n", err) + return interp.ExitStatus(5) + } + if exitStatus { + lastFalsy = v == nil || v == false + } + if err := writeValue(stdout, v, rawOutput, compact, joinOutput); err != nil { + return err + } + } + } + + if exitStatus && lastFalsy { + return interp.ExitStatus(1) + } + return nil +} + +// readInputs reads JSON (or raw) input values from stdin or files. +// +// ctx is polled in three places so that a cancellation observed mid-read +// short-circuits promptly: +// - between readers (before opening the next file / consuming stdin); +// - on every io.Read call via ctxReader, so io.ReadAll on a large but +// non-blocking source (e.g. the bytes.NewReader payload the hook +// runner supplies) returns ctx.Err() on the next chunk boundary; +// - inside the post-read value accumulation loops (raw-input line +// split and JSON stream decode), which are otherwise unbounded in +// the size of the input. +// +// A reader that blocks forever in Read (e.g. an unterminated pipe) can +// still outlast ctx; the outer abandon-goroutine path in the hook +// runner (internal/hooks/runner.go) is the authoritative enforcer for +// that case. +func readInputs(ctx context.Context, stdin io.Reader, files []string, nullInput, rawInput, slurp bool) ([]any, error) { + if nullInput { + return []any{nil}, nil + } + + var readers []io.Reader + if len(files) > 0 { + for _, f := range files { + file, err := os.Open(f) + if err != nil { + return nil, err + } + defer file.Close() + readers = append(readers, file) + } + } else { + readers = []io.Reader{stdin} + } + + var vals []any + for _, r := range readers { + if err := ctx.Err(); err != nil { + return nil, err + } + data, err := io.ReadAll(ctxReader{ctx: ctx, r: r}) + if err != nil { + // ctxReader surfaces ctx.Err() verbatim; preserve it so the + // caller can distinguish cancellation from a parse error. + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, ctxErr + } + return nil, err + } + + if rawInput { + lines := strings.Split(string(data), "\n") + if slurp { + vals = append(vals, strings.Join(lines, "\n")) + } else { + for _, line := range lines { + if err := ctx.Err(); err != nil { + return nil, err + } + if line != "" || !slurp { + vals = append(vals, line) + } + } + } + continue + } + + // Decode potentially multiple JSON values from the stream. + dec := json.NewDecoder(strings.NewReader(string(data))) + var streamVals []any + for { + if err := ctx.Err(); err != nil { + return nil, err + } + var v any + if err := dec.Decode(&v); err != nil { + if err == io.EOF { + break + } + return nil, fmt.Errorf("parse error: %w", err) + } + streamVals = append(streamVals, v) + } + + if slurp { + vals = append(vals, streamVals) + } else { + vals = append(vals, streamVals...) + } + } + + if len(vals) == 0 { + return []any{nil}, nil + } + return vals, nil +} + +// ctxReader wraps an io.Reader so that each Read call checks ctx first. +// This makes io.ReadAll over a large but non-blocking source (e.g. a +// bytes.Reader of the hook stdin payload) cancellable on the next chunk +// boundary. A reader that itself blocks in Read will still outlast ctx — +// the hook runner's abandon-goroutine path is the enforcer of last resort +// for that case. +type ctxReader struct { + ctx context.Context + r io.Reader +} + +func (cr ctxReader) Read(p []byte) (int, error) { + if err := cr.ctx.Err(); err != nil { + return 0, err + } + return cr.r.Read(p) +} + +// writeValue writes a single jq output value. +func writeValue(w io.Writer, v any, raw, compact, join bool) error { + if raw { + if s, ok := v.(string); ok { + if _, err := fmt.Fprint(w, s); err != nil { + return err + } + if !join { + _, err := fmt.Fprint(w, "\n") + return err + } + return nil + } + } + + var bs []byte + var err error + if compact { + bs, err = gojq.Marshal(v) + } else { + bs, err = json.MarshalIndent(v, "", " ") + } + if err != nil { + return err + } + if _, writeErr := w.Write(bs); writeErr != nil { + return writeErr + } + _, err = fmt.Fprint(w, "\n") + return err +} diff --git a/internal/shell/jq_test.go b/internal/shell/jq_test.go new file mode 100644 index 0000000000000000000000000000000000000000..73a33f58c376f431464faa04035d1fd11d1a75b9 --- /dev/null +++ b/internal/shell/jq_test.go @@ -0,0 +1,212 @@ +package shell + +import ( + "bytes" + "context" + "errors" + "io" + "strings" + "testing" + "time" +) + +// TestJQ_CtxCancel verifies that handleJQ polls ctx during iteration and +// returns ctx.Err() (not an interp.ExitStatus) when the context is +// cancelled. This is what lets hook timeouts interrupt long-running jq +// filters rather than waiting for the iterator to terminate naturally. +func TestJQ_CtxCancel(t *testing.T) { + t.Parallel() + + // `range(N)` generates a large stream of values. With a slurped input + // the filter produces all N values in sequence; ctx cancellation + // between values should short-circuit the loop. + const filter = "range(10000000)" + stdin := strings.NewReader("null\n") + + ctx, cancel := context.WithCancel(t.Context()) + // Cancel almost immediately so we catch the next iteration check. + cancel() + + err := handleJQ(ctx, []string{"jq", filter}, stdin, io.Discard, io.Discard) + if err == nil { + t.Fatal("expected ctx cancel error, got nil") + } + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } +} + +// TestJQ_CtxCancel_DuringFilter verifies cancellation mid-stream: ctx is +// cancelled after jq has started producing output, and the loop must +// observe the cancel on the next iteration rather than running to +// completion. +func TestJQ_CtxCancel_DuringFilter(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) + defer cancel() + + // 100M values; without ctx polling this would take many seconds to + // fully emit. With ctx polling the loop exits shortly after the + // deadline. + stdin := strings.NewReader("null\n") + var stdout, stderr bytes.Buffer + + start := time.Now() + err := handleJQ(ctx, []string{"jq", "-c", "range(100000000)"}, stdin, &stdout, &stderr) + elapsed := time.Since(start) + + if err == nil { + t.Fatal("expected ctx timeout error, got nil") + } + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected context.DeadlineExceeded, got %v", err) + } + // Allow generous slack for slow CI; the important invariant is that we + // don't run all 100M iterations (which would take orders of magnitude + // longer than 1s). + if elapsed > time.Second { + t.Fatalf("handleJQ took %v after 50ms timeout; ctx polling is not tight enough", elapsed) + } +} + +// slowReader serves bytes in small chunks with a fixed delay between +// Read calls. It never blocks indefinitely — each Read returns after +// chunkDelay — so cancellation must be observed via ctxReader's ctx +// check, not by the underlying reader itself. That isolates the +// behavior we want to test: the wrapper polling ctx between chunks. +type slowReader struct { + remaining []byte + chunk int + chunkDelay time.Duration +} + +func (s *slowReader) Read(p []byte) (int, error) { + if len(s.remaining) == 0 { + return 0, io.EOF + } + time.Sleep(s.chunkDelay) + n := min(len(p), min(s.chunk, len(s.remaining))) + copy(p, s.remaining[:n]) + s.remaining = s.remaining[n:] + return n, nil +} + +// TestJQ_CtxCancel_MidReadAll verifies that ctx cancellation observed +// *during* io.ReadAll — after several chunks have already been consumed +// — short-circuits the read via ctxReader, rather than draining the +// whole source. This is the guarantee the hook runner relies on when +// it feeds a large bytes.Reader payload. +// +// The reader serves bytes in 512-byte chunks with a 5ms gap between +// reads. ctx is cancelled after ~50ms, so several chunks have already +// been read when ctxReader first observes the cancellation. The test +// asserts that (a) we got a context.Canceled error and (b) the call +// returned well before the reader would have been fully drained. +func TestJQ_CtxCancel_MidReadAll(t *testing.T) { + t.Parallel() + + const ( + size = 64 * 1024 * 1024 // 64 MiB + chunk = 512 + chunkDelay = 5 * time.Millisecond + ) + // At 512 bytes / 5ms, draining 64 MiB would take ~11 minutes. Any + // return within a second proves cancel was observed mid-stream, not + // after EOF. + reader := &slowReader{ + remaining: bytes.Repeat([]byte("a"), size), + chunk: chunk, + chunkDelay: chunkDelay, + } + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + // Cancel after enough time that several Read calls have completed + // and io.ReadAll is actively consuming the source. + go func() { + time.Sleep(50 * time.Millisecond) + cancel() + }() + + start := time.Now() + err := handleJQ(ctx, []string{"jq", "-R", "."}, reader, io.Discard, io.Discard) + elapsed := time.Since(start) + + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } + // Generous slack for slow CI; the invariant is orders-of-magnitude + // faster than draining the full source. + if elapsed > time.Second { + t.Fatalf("mid-ReadAll cancel took %v; ctxReader is not polling between chunks", elapsed) + } + // Sanity check: we should have been cancelled mid-stream, not + // before any reads happened. If remaining == size, cancel fired so + // early nothing was consumed — that's a fast-fail path, not the + // mid-read guarantee we want to verify. + consumed := size - len(reader.remaining) + if consumed == 0 { + t.Fatal("reader was never read from; test did not exercise mid-ReadAll cancel") + } + if consumed >= size { + t.Fatal("reader was fully drained; cancel was not observed mid-read") + } +} + +// failOnReadReader fails the test if Read is ever called. It proves that a +// code path short-circuited before touching the input, without relying on +// wall-clock timing. +type failOnReadReader struct { + t *testing.T +} + +func (r failOnReadReader) Read(p []byte) (int, error) { + r.t.Error("Read called; outer guard did not short-circuit before io.ReadAll") + return 0, io.EOF +} + +// TestJQ_CtxCancel_PreCancel verifies the fast-fail path: a ctx already +// cancelled before handleJQ is called returns context.Canceled +// immediately via the outer-loop guard, never entering io.ReadAll. +// Complements TestJQ_CtxCancel_MidReadAll. +// +// The invariant — not wall-clock timing — is what proves the guard fired: +// the input reader fails the test if it is ever read from. A previous +// version asserted a 100ms ceiling, which measured scheduler latency rather +// than the guard and flaked on contended Windows CI runners under -race. +func TestJQ_CtxCancel_PreCancel(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + err := handleJQ(ctx, []string{"jq", "-R", "."}, + failOnReadReader{t: t}, + io.Discard, io.Discard) + + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected context.Canceled, got %v", err) + } +} + +// TestJQ_Success confirms the ctx-aware refactor did not regress the +// success path. +func TestJQ_Success(t *testing.T) { + t.Parallel() + + var stdout bytes.Buffer + err := handleJQ( + t.Context(), + []string{"jq", "-c", ".a"}, + strings.NewReader(`{"a":1}`), + &stdout, io.Discard, + ) + if err != nil { + t.Fatalf("handleJQ returned error: %v", err) + } + if got := stdout.String(); got != "1\n" { + t.Fatalf("stdout = %q, want %q", got, "1\n") + } +} diff --git a/internal/shell/persist_message.go b/internal/shell/persist_message.go new file mode 100644 index 0000000000000000000000000000000000000000..bca6b71be6ea37b921d287414e2b0e35cebd3f1b --- /dev/null +++ b/internal/shell/persist_message.go @@ -0,0 +1,44 @@ +package shell + +import ( + "context" + "log/slog" + "strings" + + "github.com/charmbracelet/crush/internal/message" +) + +// PersistOutput stores a bang-mode shell command result as a user message. +// If the target session no longer exists (deleted before or during the +// command), persistence is skipped without surfacing an error. +func PersistOutput( + ctx context.Context, + messages message.Service, + sessionID, command, output string, + exitCode int, +) error { + if sessionID == "" { + return nil + } + + _, err := messages.Create(ctx, sessionID, message.CreateMessageParams{ + Role: message.User, + Parts: []message.ContentPart{message.ShellCommand{ + Command: command, + Output: output, + ExitCode: exitCode, + }}, + }) + // The messages table has a single foreign key (session_id), so an FK + // failure here can only mean the session is gone. We match on the error + // text because the codebase builds against two swappable SQLite drivers + // (modernc and ncruces) and this is the one signal stable across both. + if err != nil && strings.Contains(err.Error(), "FOREIGN KEY constraint failed") { + slog.Debug( + "Skipping shell command persistence: session no longer exists", + "session_id", sessionID, + ) + return nil + } + return err +} diff --git a/internal/shell/persist_message_test.go b/internal/shell/persist_message_test.go new file mode 100644 index 0000000000000000000000000000000000000000..b743533686551ff561bcc49b17bce7896bc7b21f --- /dev/null +++ b/internal/shell/persist_message_test.go @@ -0,0 +1,69 @@ +package shell + +import ( + "testing" + + "github.com/charmbracelet/crush/internal/db" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/session" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func TestPersistOutput_SkipsMissingSession(t *testing.T) { + t.Parallel() + + conn, err := db.Connect(t.Context(), t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { conn.Close() }) + + messages := message.NewService(db.New(conn)) + + missingID := uuid.New().String() + err = PersistOutput(t.Context(), messages, missingID, "cat file.txt", "hello", 0) + require.NoError(t, err) + + stored, err := messages.List(t.Context(), missingID) + require.NoError(t, err) + require.Empty(t, stored) +} + +func TestPersistOutput_NoOpForEmptySessionID(t *testing.T) { + t.Parallel() + + conn, err := db.Connect(t.Context(), t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { conn.Close() }) + + messages := message.NewService(db.New(conn)) + + require.NoError(t, PersistOutput(t.Context(), messages, "", "echo hi", "hi", 0)) +} + +func TestPersistOutput_PersistsForExistingSession(t *testing.T) { + t.Parallel() + + conn, err := db.Connect(t.Context(), t.TempDir()) + require.NoError(t, err) + t.Cleanup(func() { conn.Close() }) + + q := db.New(conn) + sessions := session.NewService(q, conn) + messages := message.NewService(q) + + sess, err := sessions.Create(t.Context(), "shell test") + require.NoError(t, err) + + err = PersistOutput(t.Context(), messages, sess.ID, "cat file.txt", "hello", 0) + require.NoError(t, err) + + stored, err := messages.List(t.Context(), sess.ID) + require.NoError(t, err) + require.Len(t, stored, 1) + require.Equal(t, message.User, stored[0].Role) + shellParts := stored[0].ShellCommands() + require.Len(t, shellParts, 1) + require.Equal(t, "cat file.txt", shellParts[0].Command) + require.Equal(t, "hello", shellParts[0].Output) + require.Zero(t, shellParts[0].ExitCode) +} diff --git a/internal/shell/run.go b/internal/shell/run.go new file mode 100644 index 0000000000000000000000000000000000000000..f611f5654880ef84fff4b4562b2bb53c4c190ef7 --- /dev/null +++ b/internal/shell/run.go @@ -0,0 +1,349 @@ +package shell + +import ( + "bytes" + "context" + "fmt" + "io" + "log/slog" + "os" + "slices" + "strings" + + "mvdan.cc/sh/v3/expand" + "mvdan.cc/sh/v3/interp" + "mvdan.cc/sh/v3/syntax" +) + +// RunOptions configures a single stateless shell execution via [Run]. +// +// The zero value is not useful; at minimum Command must be set. Stdin, +// Stdout, and Stderr may be nil (nil readers/writers are treated as +// empty/discard). BlockFuncs may be nil to disable block-list enforcement — +// hooks use this to run user-authored commands with the same trust level as +// a shell alias. +type RunOptions struct { + // Command is the shell source to parse and execute. + Command string + // Cwd is the working directory for the execution. Required: callers + // must supply a non-empty value. Run does not silently fall back to + // the Crush process cwd — hooks and the bash tool have different + // notions of "default" and each owns that decision. + Cwd string + // Env is the full environment visible to the command. The caller is + // responsible for inheriting from os.Environ() if that's desired. + Env []string + // Stdin is the command's standard input. nil is equivalent to an empty + // input stream. + Stdin io.Reader + // Stdout receives the command's standard output. nil discards output. + Stdout io.Writer + // Stderr receives the command's standard error. nil discards output. + Stderr io.Writer + // BlockFuncs is an optional list of deny-list matchers applied before + // each command reaches the exec layer. nil disables blocking entirely. + BlockFuncs []BlockFunc + // TermWidth is the terminal width in columns for PTY execution. + // Zero uses a default of 200. + TermWidth int +} + +// Run parses and executes a shell command using the same mvdan.cc/sh +// interpreter stack that the stateful [Shell] type uses (builtins, +// optional block list, optional Go coreutils). It is safe to call +// concurrently from multiple goroutines: each call builds its own +// [interp.Runner] and shares no state with other callers or with any +// [Shell] instance. +// +// Errors returned from the command itself (non-zero exit, context +// cancellation, parse failures) follow the same conventions as +// [Shell.Exec]: inspect with [IsInterrupt] and [ExitCode]. +func Run(ctx context.Context, opts RunOptions) (err error) { + defer func() { + if r := recover(); r != nil { + err = fmt.Errorf("command execution panic: %v", r) + } + }() + + if opts.Cwd == "" { + return fmt.Errorf("shell.Run: Cwd is required") + } + + stdout := opts.Stdout + if stdout == nil { + stdout = io.Discard + } + stderr := opts.Stderr + if stderr == nil { + stderr = io.Discard + } + + line, err := syntax.NewParser().Parse(strings.NewReader(opts.Command), "") + if err != nil { + return fmt.Errorf("could not parse command: %w", err) + } + + runner, err := newRunner(opts.Cwd, opts.Env, opts.Stdin, stdout, stderr, opts.BlockFuncs) + if err != nil { + return fmt.Errorf("could not run command: %w", err) + } + + return runner.Run(ctx, line) +} + +// CaptureResult holds the combined output and exit code from a +// captured shell execution. +type CaptureResult struct { + Output string + ExitCode int +} + +// PersistFunc is a callback that persists a shell command result. +// Used by RunAndPersist to decouple execution from storage. +type PersistFunc func(command, output string, exitCode int) error + +// RunAndPersist executes a shell command via PTY and optionally +// persists the result through the provided callback. This unifies +// the run-and-save pattern used by both AppWorkspace and Backend. +func RunAndPersist(ctx context.Context, opts RunOptions, persist PersistFunc) (CaptureResult, error) { + result, err := RunAndCapturePTY(ctx, opts) + if err != nil { + return CaptureResult{}, err + } + + if persist != nil { + if persistErr := persist(opts.Command, result.Output, result.ExitCode); persistErr != nil { + slog.Error("Failed to persist shell command output", "error", persistErr, "command", opts.Command) + } + } + + return result, nil +} + +// RunAndCapture executes a shell command and returns its combined +// stdout/stderr output along with the exit code. It inherits the +// current process environment when opts.Env is nil. +func RunAndCapture(ctx context.Context, opts RunOptions) (CaptureResult, error) { + if opts.Env == nil { + opts.Env = os.Environ() + } + + var stdout, stderr bytes.Buffer + opts.Stdout = &stdout + opts.Stderr = &stderr + + runErr := Run(ctx, opts) + + exitCode := 0 + if runErr != nil { + exitCode = ExitCode(runErr) + } + + output := stdout.String() + if stderr.Len() > 0 { + if output != "" { + output += "\n" + } + output += stderr.String() + } + + return CaptureResult{ + Output: output, + ExitCode: exitCode, + }, nil +} + +// ptyColorEnvVars force color output for programs running inside a +// PTY. These are only applied in RunAndCapturePTY, not in the plain +// Run/RunAndCapture paths where ANSI codes would be noise. +var ptyColorEnvVars = []string{ + "COLORTERM=truecolor", + "CLICOLOR_FORCE=1", + "FORCE_COLOR=1", +} + +// RunAndCapturePTY executes a shell command through the mvdan.cc/sh +// interpreter with color-forcing environment variables set. Programs +// that respect FORCE_COLOR or CLICOLOR_FORCE (git, cargo, npm, eza, +// bat, ripgrep, etc.) will emit ANSI color sequences even without a +// real PTY. This approach is fully cross-platform — no /bin/sh or +// unix PTY required. +// +// The name is preserved for API compatibility; the PTY path has been +// replaced by the portable interpreter + env-var approach. +func RunAndCapturePTY(ctx context.Context, opts RunOptions) (CaptureResult, error) { + if opts.Env == nil { + opts.Env = os.Environ() + } + opts.Env = append(opts.Env, ptyColorEnvVars...) + return RunAndCapture(ctx, opts) +} + +// newRunner constructs an [interp.Runner] configured with the standard +// Crush handler stack. Shared by the stateless [Run] entrypoint and the +// stateful [Shell] so the two surfaces cannot drift. +func newRunner(cwd string, env []string, stdin io.Reader, stdout, stderr io.Writer, blockFuncs []BlockFunc) (*interp.Runner, error) { + env = withNonInteractiveEnv(env) + return interp.New( + interp.StdIO(stdin, stdout, stderr), + interp.Interactive(false), + interp.Env(expand.ListEnviron(env...)), + interp.Dir(cwd), + execHandlerOption(blockFuncs), + ) +} + +// execHandlerOption returns an interp.RunnerOption that installs the +// standard Crush middleware chain (builtins, script dispatch, block list) +// on top of a process-group-isolated base exec handler. +// +// We use interp.ExecHandler (singular) with a manually-built chain rather +// than interp.ExecHandlers because the latter always appends +// interp.DefaultExecHandler as the final handler, which lacks process group +// isolation. Without isolation, shells like zsh that set up job control +// when sourcing framework files can send SIGINT/SIGTERM to Crush's process +// group and crash the parent. +func execHandlerOption(blockFuncs []BlockFunc) interp.RunnerOption { + base := processGroupExecHandler(defaultKillTimeout) + handler := base + for _, mw := range slices.Backward(standardHandlers(blockFuncs)) { + handler = mw(handler) + } + // ExecHandlers always appends DefaultExecHandler which lacks process + // group isolation, so we use the deprecated ExecHandler instead. + return interp.ExecHandler(handler) +} + +// nonInteractiveEnvVars are forced on every shell execution to prevent +// commands from hanging on a nonexistent TTY. These are always applied +// regardless of the caller's environment because Crush shells are never +// interactive — preserving user preferences like EDITOR=nvim only causes +// hangs, not useful behavior. +var nonInteractiveEnvVars = []string{ + "TERM=xterm-256color", + "GIT_EDITOR=false", + "EDITOR=false", + "VISUAL=false", + "JJ_EDITOR=false", + "JJ_PAGER=cat", + "GIT_PAGER=cat", + "PAGER=cat", +} + +// withNonInteractiveEnv returns env with nonInteractiveEnvVars forced in, +// replacing any existing values for those keys. The returned slice is a +// new allocation safe to use concurrently with the input. +func withNonInteractiveEnv(env []string) []string { + // Build a set of override keys for fast lookup. + overrideKeys := make(map[string]bool, len(nonInteractiveEnvVars)) + for _, kv := range nonInteractiveEnvVars { + if key, _, ok := strings.Cut(kv, "="); ok { + overrideKeys[key] = true + } + } + + // Copy env, filtering out any keys we will override. + result := make([]string, 0, len(env)+len(nonInteractiveEnvVars)) + for _, e := range env { + if key, _, ok := strings.Cut(e, "="); ok && overrideKeys[key] { + continue + } + result = append(result, e) + } + + return append(result, nonInteractiveEnvVars...) +} + +// herdrEnvVars are the environment variables herdr injects into panes +// so agents can report state over its Unix socket API. Subprocesses +// must not inherit these: a child process that calls herdr.Init() +// would attach to the parent's pane and, on exit, release its agent +// authority — making the status vanish. Stripping them here closes +// that gap for every command the bash tool runs. +var herdrEnvVars = []string{ + "HERDR_ENV", + "HERDR_SOCKET_PATH", + "HERDR_PANE_ID", +} + +// withoutHerdrEnv returns env with all HERDR_* variables removed. +// The returned slice is a new allocation safe to use concurrently +// with the input. +func withoutHerdrEnv(env []string) []string { + strip := make(map[string]bool, len(herdrEnvVars)) + for _, k := range herdrEnvVars { + strip[k] = true + } + result := make([]string, 0, len(env)) + for _, e := range env { + if key, _, ok := strings.Cut(e, "="); ok && strip[key] { + continue + } + result = append(result, e) + } + return result +} + +// execMiddleware wraps a base [interp.ExecHandlerFunc], composing like HTTP +// middleware: each layer either handles a command itself or delegates to the +// next handler in the chain. +type execMiddleware = func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc + +// standardHandlers returns the exec-handler middleware chain used by both +// [Run] and [Shell]. Order matters: +// 1. builtins first (so Crush's in-process jq wins over any PATH binary); +// 2. script dispatch (shebang / binary / shell-source for path-prefixed +// argv[0], no-op for bare commands) — runs before the block list so +// that deny rules see the already-resolved argv of anything the +// script exec's rather than the outer path-prefixed wrapper; +// 3. block list; +// 4. optional Go coreutils (only when useGoCoreUtils is on). +func standardHandlers(blockFuncs []BlockFunc) []execMiddleware { + handlers := []execMiddleware{ + builtinHandler(), + scriptDispatchHandler(blockFuncs), + blockHandler(blockFuncs), + } + if useGoCoreUtils && coreUtilsExecHandler != nil { + handlers = append(handlers, coreUtilsExecHandler) + } + return handlers +} + +// builtinHandler returns middleware that dispatches recognized Crush +// builtins to their in-process Go implementations. All builtins (jq plus the +// config builtins registered by shellconfig) live in the builtins map; config +// builtins are no-ops without a ConfigBuilder on the context. +func builtinHandler() execMiddleware { + return func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc { + return func(ctx context.Context, args []string) error { + if len(args) == 0 { + return next(ctx, args) + } + if h, ok := builtins[args[0]]; ok { + hc := interp.HandlerCtx(ctx) + return h(ctx, args, hc.Stdin, hc.Stdout, hc.Stderr) + } + return next(ctx, args) + } + } +} + +// blockHandler returns middleware that rejects commands matched by any of +// the provided [BlockFunc]s before they reach the underlying exec path. +// A nil or empty blockFuncs slice is a no-op. +func blockHandler(blockFuncs []BlockFunc) execMiddleware { + return func(next interp.ExecHandlerFunc) interp.ExecHandlerFunc { + return func(ctx context.Context, args []string) error { + if len(args) == 0 { + return next(ctx, args) + } + for _, blockFunc := range blockFuncs { + if blockFunc(args) { + return fmt.Errorf("command is not allowed for security reasons: %q", args[0]) + } + } + return next(ctx, args) + } + } +} diff --git a/internal/shell/run_test.go b/internal/shell/run_test.go new file mode 100644 index 0000000000000000000000000000000000000000..e80aaf9fe97f51e26799e7a37856f5256af87a97 --- /dev/null +++ b/internal/shell/run_test.go @@ -0,0 +1,364 @@ +package shell + +import ( + "bytes" + "context" + "errors" + "fmt" + "slices" + "strings" + "sync" + "testing" + "time" +) + +func TestRun_Echo(t *testing.T) { + var stdout, stderr bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: "echo hi", + Cwd: t.TempDir(), + Stdout: &stdout, + Stderr: &stderr, + }) + if err != nil { + t.Fatalf("Run returned error: %v (stderr=%q)", err, stderr.String()) + } + if got := stdout.String(); got != "hi\n" { + t.Fatalf("stdout = %q, want %q", got, "hi\n") + } +} + +func TestRun_ExitCode(t *testing.T) { + err := Run(t.Context(), RunOptions{ + Command: "exit 7", + Cwd: t.TempDir(), + }) + if err == nil { + t.Fatal("expected error for exit 7, got nil") + } + if code := ExitCode(err); code != 7 { + t.Fatalf("ExitCode = %d, want 7", code) + } +} + +func TestRun_Stdin(t *testing.T) { + // Use the `read` shell builtin so the test doesn't depend on any + // external binary being on PATH (we pass an empty Env here). + var stdout bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: "read line; echo got:$line", + Cwd: t.TempDir(), + Stdin: strings.NewReader("hello\n"), + Stdout: &stdout, + }) + if err != nil { + t.Fatalf("Run returned error: %v", err) + } + if got := stdout.String(); got != "got:hello\n" { + t.Fatalf("stdout = %q, want %q", got, "got:hello\n") + } +} + +func TestRun_Env(t *testing.T) { + var stdout bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: `echo "$FOO"`, + Cwd: t.TempDir(), + Env: []string{"FOO=bar"}, + Stdout: &stdout, + }) + if err != nil { + t.Fatalf("Run returned error: %v", err) + } + if got := stdout.String(); got != "bar\n" { + t.Fatalf("stdout = %q, want %q", got, "bar\n") + } +} + +func TestRun_Cwd(t *testing.T) { + dir := t.TempDir() + var stdout bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: "pwd", + Cwd: dir, + Stdout: &stdout, + }) + if err != nil { + t.Fatalf("Run returned error: %v", err) + } + // mvdan's pwd builtin resolves symlinks (e.g. /var -> /private/var on + // macOS). Compare against a suffix so we don't get bitten by that. + got := strings.TrimRight(stdout.String(), "\n") + if !strings.HasSuffix(got, dir) && !strings.HasSuffix(dir, got) { + t.Fatalf("pwd = %q, want it to match %q", got, dir) + } +} + +func TestRun_JqBuiltin(t *testing.T) { + var stdout bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: `echo '{"a":1}' | jq .a`, + Cwd: t.TempDir(), + Stdout: &stdout, + }) + if err != nil { + t.Fatalf("Run returned error: %v", err) + } + if got := stdout.String(); got != "1\n" { + t.Fatalf("stdout = %q, want %q", got, "1\n") + } +} + +func TestRun_ParallelIsolation(t *testing.T) { + const n = 10 + var wg sync.WaitGroup + wg.Add(n) + errs := make([]error, n) + outs := make([]string, n) + dirs := make([]string, n) + for i := range n { + dirs[i] = t.TempDir() + go func(i int) { + defer wg.Done() + var stdout bytes.Buffer + errs[i] = Run(t.Context(), RunOptions{ + Command: `echo "$MARKER"`, + Cwd: dirs[i], + Env: []string{fmt.Sprintf("MARKER=id-%d", i)}, + Stdout: &stdout, + }) + outs[i] = stdout.String() + }(i) + } + wg.Wait() + for i := range n { + if errs[i] != nil { + t.Errorf("goroutine %d: err = %v", i, errs[i]) + continue + } + want := fmt.Sprintf("id-%d\n", i) + if outs[i] != want { + t.Errorf("goroutine %d: stdout = %q, want %q", i, outs[i], want) + } + } +} + +// TestRun_CtxCancel_BusyLoop verifies that a pure-shell loop respects ctx +// cancellation. mvdan's interpreter checks ctx between statements, so this +// should return quickly even without any external command. The test bounds +// its own wait via a select so a regression can't hang CI. +func TestRun_CtxCancel_BusyLoop(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 500*time.Millisecond) + t.Cleanup(cancel) + + done := make(chan error, 1) + go func() { + done <- Run(ctx, RunOptions{ + Command: "while true; do :; done", + Cwd: t.TempDir(), + }) + }() + + select { + case err := <-done: + if !IsInterrupt(err) && !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected interrupt/deadline error, got: %v", err) + } + case <-time.After(1500 * time.Millisecond): + t.Fatal("Run did not return within 1.5s after ctx cancel") + } +} + +// TestRun_CtxCancel_ExternalSleep verifies ctx cancellation reaches an +// external process via mvdan's default exec. Uses sleep, which lives in +// coreutils on Windows and /bin on Unix. +func TestRun_CtxCancel_ExternalSleep(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), 200*time.Millisecond) + t.Cleanup(cancel) + + done := make(chan error, 1) + start := time.Now() + go func() { + done <- Run(ctx, RunOptions{ + Command: "sleep 30", + Cwd: t.TempDir(), + }) + }() + + select { + case err := <-done: + elapsed := time.Since(start) + if elapsed > time.Second { + t.Fatalf("sleep took too long to cancel: %v", elapsed) + } + if err == nil { + t.Fatal("expected non-nil error from cancelled sleep") + } + case <-time.After(time.Second): + t.Fatal("Run did not return within 1s after ctx cancel") + } +} + +func TestRun_ParseError(t *testing.T) { + err := Run(t.Context(), RunOptions{ + Command: "echo 'unterminated", + Cwd: t.TempDir(), + }) + if err == nil { + t.Fatal("expected parse error, got nil") + } + if !strings.Contains(err.Error(), "parse") { + t.Fatalf("error should mention parse: %v", err) + } +} + +func TestRun_BlockFuncs(t *testing.T) { + block := CommandsBlocker([]string{"forbidden"}) + var stderr bytes.Buffer + err := Run(t.Context(), RunOptions{ + Command: "forbidden", + Cwd: t.TempDir(), + Stderr: &stderr, + BlockFuncs: []BlockFunc{block}, + }) + if err == nil { + t.Fatal("expected error when running blocked command") + } + if !strings.Contains(err.Error(), "not allowed") { + t.Fatalf("expected 'not allowed' error, got: %v", err) + } +} + +func TestRun_RequiresCwd(t *testing.T) { + err := Run(t.Context(), RunOptions{ + Command: "echo hi", + }) + if err == nil { + t.Fatal("expected error when Cwd is empty, got nil") + } + if !strings.Contains(err.Error(), "Cwd is required") { + t.Fatalf("error should mention Cwd requirement: %v", err) + } +} + +func TestWithNonInteractiveEnv_Empty(t *testing.T) { + t.Parallel() + result := withNonInteractiveEnv(nil) + // All defaults must be present. + for _, want := range nonInteractiveEnvVars { + if !slices.Contains(result, want) { + t.Errorf("missing default %q in result", want) + } + } +} + +func TestWithNonInteractiveEnv_OverridesExisting(t *testing.T) { + t.Parallel() + env := []string{"EDITOR=nvim", "PAGER=less", "FOO=bar"} + result := withNonInteractiveEnv(env) + + // EDITOR and PAGER must be overridden, not preserved. + for _, e := range result { + if e == "EDITOR=nvim" { + t.Error("EDITOR=nvim should have been overridden") + } + if e == "PAGER=less" { + t.Error("PAGER=less should have been overridden") + } + } + // FOO must survive. + if !slices.Contains(result, "FOO=bar") { + t.Error("FOO=bar should be preserved") + } +} + +func TestWithNonInteractiveEnv_NoPrefixCollision(t *testing.T) { + t.Parallel() + // EDITORIAL should NOT match EDITOR. + env := []string{"EDITORIAL=yes", "GITHUB_TOKEN=secret"} + result := withNonInteractiveEnv(env) + + foundEditorial := false + foundGithub := false + for _, e := range result { + if e == "EDITORIAL=yes" { + foundEditorial = true + } + if e == "GITHUB_TOKEN=secret" { + foundGithub = true + } + } + if !foundEditorial { + t.Error("EDITORIAL=yes should not be removed by EDITOR override") + } + if !foundGithub { + t.Error("GITHUB_TOKEN=secret should not be removed") + } +} + +func TestWithNonInteractiveEnv_SliceIndependence(t *testing.T) { + t.Parallel() + env := []string{"FOO=bar"} + result := withNonInteractiveEnv(env) + // Mutating the input must not affect the result. + env[0] = "FOO=baz" + for _, e := range result { + if e == "FOO=baz" { + t.Error("result shares backing array with input") + } + } +} + +func TestWithoutHerdrEnv_StripsAllVars(t *testing.T) { + t.Parallel() + env := []string{ + "HERDR_ENV=1", + "HERDR_SOCKET_PATH=/tmp/herdr.sock", + "HERDR_PANE_ID=wA:p1", + "PATH=/usr/bin", + "HOME=/home/user", + } + result := withoutHerdrEnv(env) + for _, e := range result { + if strings.HasPrefix(e, "HERDR_") { + t.Errorf("herdr var not stripped: %s", e) + } + } + if !slices.Contains(result, "PATH=/usr/bin") { + t.Error("non-herdr var PATH was incorrectly removed") + } + if !slices.Contains(result, "HOME=/home/user") { + t.Error("non-herdr var HOME was incorrectly removed") + } +} + +func TestWithoutHerdrEnv_EmptyInput(t *testing.T) { + t.Parallel() + result := withoutHerdrEnv(nil) + if len(result) != 0 { + t.Errorf("expected empty result for nil input, got %v", result) + } +} + +func TestWithoutHerdrEnv_SliceIndependence(t *testing.T) { + t.Parallel() + env := []string{"HERDR_ENV=1", "FOO=bar"} + result := withoutHerdrEnv(env) + env[1] = "FOO=baz" + for _, e := range result { + if e == "FOO=baz" { + t.Error("result shares backing array with input") + } + } +} + +func TestRun_DiscardsNilWriters(t *testing.T) { + // No panic when Stdout/Stderr are nil. + err := Run(t.Context(), RunOptions{ + Command: "echo hi; echo err >&2", + Cwd: t.TempDir(), + }) + if err != nil { + t.Fatalf("Run returned error: %v", err) + } +} diff --git a/internal/shell/shell.go b/internal/shell/shell.go new file mode 100644 index 0000000000000000000000000000000000000000..eef3f424e21577e31ee707a5f9f651446b6ac4d5 --- /dev/null +++ b/internal/shell/shell.go @@ -0,0 +1,312 @@ +// Package shell provides cross-platform shell execution capabilities. +// +// This package provides Shell instances for executing commands with their own +// working directory and environment. Each shell execution is independent. +// +// WINDOWS COMPATIBILITY: +// This implementation provides POSIX shell emulation (mvdan.cc/sh/v3) even on +// Windows. Commands should use forward slashes (/) as path separators to work +// correctly on all platforms. +package shell + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "os" + "slices" + "strings" + "sync" + + "github.com/charmbracelet/x/exp/slice" + "mvdan.cc/sh/v3/interp" + "mvdan.cc/sh/v3/syntax" +) + +// ShellType represents the type of shell to use +type ShellType int + +const ( + ShellTypePOSIX ShellType = iota + ShellTypeCmd + ShellTypePowerShell +) + +// CrushEnvMarkers returns a fresh slice of the environment variables that +// Crush unconditionally sets on every shell it spawns — both the interactive +// bash tool's [Shell] and the hook runner's [Run] calls. Tools that want to +// detect "am I being invoked by an AI agent?" can check any of these. +// Keeping them in one place guarantees the two shell surfaces cannot drift. +// A fresh slice is returned on every call so callers may append freely. +func CrushEnvMarkers() []string { + return []string{ + "CRUSH=1", + "AGENT=crush", + "AI_AGENT=crush", + } +} + +// Logger interface for optional logging +type Logger interface { + InfoPersist(msg string, keysAndValues ...any) +} + +// noopLogger is a logger that does nothing +type noopLogger struct{} + +func (noopLogger) InfoPersist(msg string, keysAndValues ...any) {} + +// BlockFunc is a function that determines if a command should be blocked +type BlockFunc func(args []string) bool + +// Shell provides cross-platform shell execution with optional state persistence +type Shell struct { + env []string + cwd string + mu sync.Mutex + logger Logger + blockFuncs []BlockFunc +} + +// Options for creating a new shell +type Options struct { + WorkingDir string + Env []string + Logger Logger + BlockFuncs []BlockFunc +} + +// NewShell creates a new shell instance with the given options +func NewShell(opts *Options) *Shell { + if opts == nil { + opts = &Options{} + } + + cwd := opts.WorkingDir + if cwd == "" { + cwd, _ = os.Getwd() + } + + env := opts.Env + if env == nil { + env = os.Environ() + } + + // Strip herdr pane-ownership vars so subprocesses (including test + // binaries and nested crush instances) can't attach to or release + // the parent pane's agent authority. + env = withoutHerdrEnv(env) + + // Allow tools to detect execution by Crush. + env = append(env, CrushEnvMarkers()...) + + logger := opts.Logger + if logger == nil { + logger = noopLogger{} + } + + return &Shell{ + cwd: cwd, + env: env, + logger: logger, + blockFuncs: opts.BlockFuncs, + } +} + +// Exec executes a command in the shell +func (s *Shell) Exec(ctx context.Context, command string) (string, string, error) { + s.mu.Lock() + defer s.mu.Unlock() + + return s.exec(ctx, command) +} + +// ExecStream executes a command in the shell with streaming output to provided writers +func (s *Shell) ExecStream(ctx context.Context, command string, stdout, stderr io.Writer) error { + s.mu.Lock() + defer s.mu.Unlock() + + return s.execStream(ctx, command, stdout, stderr) +} + +// GetWorkingDir returns the current working directory +func (s *Shell) GetWorkingDir() string { + s.mu.Lock() + defer s.mu.Unlock() + return s.cwd +} + +// SetWorkingDir sets the working directory +func (s *Shell) SetWorkingDir(dir string) error { + s.mu.Lock() + defer s.mu.Unlock() + + // Verify the directory exists + if _, err := os.Stat(dir); err != nil { + return fmt.Errorf("directory does not exist: %w", err) + } + + s.cwd = dir + return nil +} + +// GetEnv returns a copy of the environment variables +func (s *Shell) GetEnv() []string { + s.mu.Lock() + defer s.mu.Unlock() + + env := make([]string, len(s.env)) + copy(env, s.env) + return env +} + +// SetEnv sets an environment variable +func (s *Shell) SetEnv(key, value string) { + s.mu.Lock() + defer s.mu.Unlock() + + // Update or add the environment variable + keyPrefix := key + "=" + for i, env := range s.env { + if strings.HasPrefix(env, keyPrefix) { + s.env[i] = keyPrefix + value + return + } + } + s.env = append(s.env, keyPrefix+value) +} + +// SetBlockFuncs sets the command block functions for the shell +func (s *Shell) SetBlockFuncs(blockFuncs []BlockFunc) { + s.mu.Lock() + defer s.mu.Unlock() + s.blockFuncs = blockFuncs +} + +// CommandsBlocker creates a BlockFunc that blocks exact command matches +func CommandsBlocker(cmds []string) BlockFunc { + bannedSet := make(map[string]struct{}) + for _, cmd := range cmds { + bannedSet[cmd] = struct{}{} + } + + return func(args []string) bool { + if len(args) == 0 { + return false + } + _, ok := bannedSet[args[0]] + return ok + } +} + +// ArgumentsBlocker creates a BlockFunc that blocks specific subcommand +func ArgumentsBlocker(cmd string, args []string, flags []string) BlockFunc { + return func(parts []string) bool { + if len(parts) == 0 || parts[0] != cmd { + return false + } + + argParts, flagParts := splitArgsFlags(parts[1:]) + if len(argParts) < len(args) || len(flagParts) < len(flags) { + return false + } + + argsMatch := slices.Equal(argParts[:len(args)], args) + flagsMatch := slice.IsSubset(flags, flagParts) + + return argsMatch && flagsMatch + } +} + +func splitArgsFlags(parts []string) (args []string, flags []string) { + args = make([]string, 0, len(parts)) + flags = make([]string, 0, len(parts)) + for _, part := range parts { + if strings.HasPrefix(part, "-") { + // Extract flag name before '=' if present + flag := part + if before, _, ok := strings.Cut(part, "="); ok { + flag = before + } + flags = append(flags, flag) + } else { + args = append(args, part) + } + } + return args, flags +} + +// newInterp creates a new interpreter with the current shell state. A nil +// stdin is equivalent to an empty input stream. +func (s *Shell) newInterp(stdin io.Reader, stdout, stderr io.Writer) (*interp.Runner, error) { + return newRunner(s.cwd, s.env, stdin, stdout, stderr, s.blockFuncs) +} + +// updateShellFromRunner updates the shell from the interpreter after execution. +func (s *Shell) updateShellFromRunner(runner *interp.Runner) { + s.cwd = runner.Dir + s.env = s.env[:0] + for name, vr := range runner.Vars { + if vr.Exported { + s.env = append(s.env, name+"="+vr.Str) + } + } +} + +// execCommon is the shared implementation for executing commands +func (s *Shell) execCommon(ctx context.Context, command string, stdout, stderr io.Writer) (err error) { + var runner *interp.Runner + defer func() { + if r := recover(); r != nil { + err = fmt.Errorf("command execution panic: %v", r) + } + if runner != nil { + s.updateShellFromRunner(runner) + } + s.logger.InfoPersist("command finished", "command", command, "err", err) + }() + + line, err := syntax.NewParser().Parse(strings.NewReader(command), "") + if err != nil { + return fmt.Errorf("could not parse command: %w", err) + } + + runner, err = s.newInterp(nil, stdout, stderr) + if err != nil { + return fmt.Errorf("could not run command: %w", err) + } + + err = runner.Run(ctx, line) + return err +} + +// exec executes commands using a cross-platform shell interpreter. +func (s *Shell) exec(ctx context.Context, command string) (string, string, error) { + var stdout, stderr bytes.Buffer + err := s.execCommon(ctx, command, &stdout, &stderr) + return stdout.String(), stderr.String(), err +} + +// execStream executes commands using POSIX shell emulation with streaming output +func (s *Shell) execStream(ctx context.Context, command string, stdout, stderr io.Writer) error { + return s.execCommon(ctx, command, stdout, stderr) +} + +// IsInterrupt checks if an error is due to interruption +func IsInterrupt(err error) bool { + return errors.Is(err, context.Canceled) || + errors.Is(err, context.DeadlineExceeded) +} + +// ExitCode extracts the exit code from an error +func ExitCode(err error) int { + if err == nil { + return 0 + } + if exitErr, ok := errors.AsType[interp.ExitStatus](err); ok { + return int(exitErr) + } + return 1 +} diff --git a/internal/shell/shell_test.go b/internal/shell/shell_test.go new file mode 100644 index 0000000000000000000000000000000000000000..ae53fb1c3ecc2c1fe3760566b122337eb0a2782f --- /dev/null +++ b/internal/shell/shell_test.go @@ -0,0 +1,120 @@ +package shell + +import ( + "context" + "path/filepath" + "runtime" + "strings" + "testing" + "time" +) + +// Benchmark to measure CPU efficiency +func BenchmarkShellQuickCommands(b *testing.B) { + shell := NewShell(&Options{WorkingDir: b.TempDir()}) + + b.ReportAllocs() + + for b.Loop() { + _, _, err := shell.Exec(b.Context(), "echo test") + exitCode := ExitCode(err) + if err != nil || exitCode != 0 { + b.Fatalf("Command failed: %v, exit code: %d", err, exitCode) + } + } +} + +func TestTestTimeout(t *testing.T) { + // XXX(@andreynering): This fails on Windows. Address once possible. + if runtime.GOOS == "windows" { + t.Skip("Skipping test on Windows") + } + + ctx, cancel := context.WithTimeout(t.Context(), time.Millisecond) + t.Cleanup(cancel) + + shell := NewShell(&Options{WorkingDir: t.TempDir()}) + _, _, err := shell.Exec(ctx, "sleep 10") + if status := ExitCode(err); status == 0 { + t.Fatalf("Expected non-zero exit status, got %d", status) + } + if !IsInterrupt(err) { + t.Fatalf("Expected command to be interrupted, but it was not") + } + if err == nil { + t.Fatalf("Expected an error due to timeout, but got none") + } +} + +func TestTestCancel(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + cancel() // immediately cancel the context + + shell := NewShell(&Options{WorkingDir: t.TempDir()}) + _, _, err := shell.Exec(ctx, "sleep 10") + if status := ExitCode(err); status == 0 { + t.Fatalf("Expected non-zero exit status, got %d", status) + } + if !IsInterrupt(err) { + t.Fatalf("Expected command to be interrupted, but it was not") + } + if err == nil { + t.Fatalf("Expected an error due to cancel, but got none") + } +} + +func TestRunCommandError(t *testing.T) { + shell := NewShell(&Options{WorkingDir: t.TempDir()}) + _, _, err := shell.Exec(t.Context(), "nopenopenope") + if status := ExitCode(err); status == 0 { + t.Fatalf("Expected non-zero exit status, got %d", status) + } + if IsInterrupt(err) { + t.Fatalf("Expected command to not be interrupted, but it was") + } + if err == nil { + t.Fatalf("Expected an error, got nil") + } +} + +func TestRunContinuity(t *testing.T) { + tempDir1 := t.TempDir() + tempDir2 := t.TempDir() + + shell := NewShell(&Options{WorkingDir: tempDir1}) + if _, _, err := shell.Exec(t.Context(), "export FOO=bar"); err != nil { + t.Fatalf("failed to set env: %v", err) + } + if _, _, err := shell.Exec(t.Context(), "cd "+filepath.ToSlash(tempDir2)); err != nil { + t.Fatalf("failed to change directory: %v", err) + } + out, _, err := shell.Exec(t.Context(), "echo $FOO ; pwd") + if err != nil { + t.Fatalf("failed to echo: %v", err) + } + expect := "bar\n" + tempDir2 + "\n" + if out != expect { + t.Fatalf("expected output %q, got %q", expect, out) + } +} + +func TestCrossPlatformExecution(t *testing.T) { + shell := NewShell(&Options{WorkingDir: "."}) + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + + // Test a simple command that should work on all platforms + stdout, stderr, err := shell.Exec(ctx, "echo hello") + if err != nil { + t.Fatalf("Echo command failed: %v, stderr: %s", err, stderr) + } + + if stdout == "" { + t.Error("Echo command produced no output") + } + + // The output should contain "hello" regardless of platform + if !strings.Contains(strings.ToLower(stdout), "hello") { + t.Errorf("Echo output should contain 'hello', got: %q", stdout) + } +} diff --git a/internal/shell/stream.go b/internal/shell/stream.go new file mode 100644 index 0000000000000000000000000000000000000000..1fbc462c07f2ba36aa7580b0f396f53360f0ee4a --- /dev/null +++ b/internal/shell/stream.go @@ -0,0 +1,52 @@ +package shell + +import ( + "bytes" + "context" + "os" + "sync" +) + +// progressWriter wraps an io.Writer and calls onProgress with each write. +// It is safe for concurrent use from multiple goroutines (e.g. stdout and +// stderr writing simultaneously). +type progressWriter struct { + mu sync.Mutex + buf bytes.Buffer + onProgress func(string) +} + +func (w *progressWriter) Write(p []byte) (int, error) { + w.mu.Lock() + defer w.mu.Unlock() + n, err := w.buf.Write(p) + if n > 0 && w.onProgress != nil { + w.onProgress(string(p[:n])) + } + return n, err +} + +// RunAndCaptureStream executes a shell command and streams output chunks +// to onProgress as they arrive. Returns the complete output and exit code. +func RunAndCaptureStream(ctx context.Context, opts RunOptions, onProgress func(string)) (CaptureResult, error) { + if opts.Env == nil { + opts.Env = os.Environ() + } + opts.Env = append(opts.Env, ptyColorEnvVars...) + + buf := &progressWriter{onProgress: onProgress} + opts.Stdout = buf + opts.Stderr = buf + + runErr := Run(ctx, opts) + + exitCode := 0 + if runErr != nil { + exitCode = ExitCode(runErr) + } + + return CaptureResult{ + Output: buf.buf.String(), + ExitCode: exitCode, + }, nil +} diff --git a/internal/workspace/app_workspace.go b/internal/workspace/app_workspace.go new file mode 100644 index 0000000000000000000000000000000000000000..c2a36fff82bb80acf651d0057bedb1f173de5f99 --- /dev/null +++ b/internal/workspace/app_workspace.go @@ -0,0 +1,506 @@ +package workspace + +import ( + "context" + "errors" + "fmt" + "log/slog" + "time" + + tea "charm.land/bubbletea/v2" + "github.com/charmbracelet/crush/internal/agent" + mcptools "github.com/charmbracelet/crush/internal/agent/tools/mcp" + "github.com/charmbracelet/crush/internal/app" + "github.com/charmbracelet/crush/internal/commands" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/history" + "github.com/charmbracelet/crush/internal/lsp" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/oauth" + "github.com/charmbracelet/crush/internal/permission" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/question" + "github.com/charmbracelet/crush/internal/session" + "github.com/charmbracelet/crush/internal/shell" + "github.com/charmbracelet/crush/internal/skills" +) + +// AppWorkspace implements the Workspace interface by delegating +// directly to an in-process [app.App] instance. This is the default +// mode when the client/server architecture is not enabled. +type AppWorkspace struct { + app *app.App + store *config.ConfigStore +} + +// NewAppWorkspace creates a new AppWorkspace wrapping the given app +// and config store. +func NewAppWorkspace(a *app.App, store *config.ConfigStore) *AppWorkspace { + return &AppWorkspace{ + app: a, + store: store, + } +} + +// -- Sessions -- + +func (w *AppWorkspace) CreateSession(ctx context.Context, title string) (session.Session, error) { + return w.app.Sessions.Create(ctx, title) +} + +func (w *AppWorkspace) GetSession(ctx context.Context, sessionID string) (session.Session, error) { + return w.app.Sessions.Get(ctx, sessionID) +} + +func (w *AppWorkspace) ListSessions(ctx context.Context) ([]session.Session, error) { + return w.app.Sessions.List(ctx) +} + +func (w *AppWorkspace) SaveSession(ctx context.Context, sess session.Session) (session.Session, error) { + return w.app.Sessions.Save(ctx, sess) +} + +func (w *AppWorkspace) DeleteSession(ctx context.Context, sessionID string) error { + return w.app.Sessions.Delete(ctx, sessionID) +} + +func (w *AppWorkspace) CreateAgentToolSessionID(messageID, toolCallID string) string { + return w.app.Sessions.CreateAgentToolSessionID(messageID, toolCallID) +} + +func (w *AppWorkspace) ParseAgentToolSessionID(sessionID string) (string, string, bool) { + return w.app.Sessions.ParseAgentToolSessionID(sessionID) +} + +// SetCurrentSession reports the active session to herdr so the pane +// can persist a resumable reference. Multi-client presence tracking +// is irrelevant in single-client local mode, but herdr still needs +// to know which session is live to support agent resume. +func (w *AppWorkspace) SetCurrentSession(ctx context.Context, sessionID string) error { + w.app.ReportCurrentSession(sessionID) + return nil +} + +// -- Messages -- + +func (w *AppWorkspace) ListMessages(ctx context.Context, sessionID string) ([]message.Message, error) { + // Drain any debounced updates so the caller observes the latest + // in-memory state. message.Service buffers streaming deltas and a + // cold List would otherwise miss them at session-switch time. + if err := w.app.Messages.FlushAll(ctx); err != nil { + return nil, err + } + return w.app.Messages.List(ctx, sessionID) +} + +func (w *AppWorkspace) ListUserMessages(ctx context.Context, sessionID string) ([]message.Message, error) { + return w.app.Messages.ListUserMessages(ctx, sessionID) +} + +func (w *AppWorkspace) ListAllUserMessages(ctx context.Context) ([]message.Message, error) { + return w.app.Messages.ListAllUserMessages(ctx) +} + +// -- Agent -- + +func (w *AppWorkspace) AgentRun(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) error { + if w.app.AgentCoordinator == nil { + return errors.New("agent coordinator not initialized") + } + _, err := w.app.AgentCoordinator.Run(ctx, sessionID, prompt, attachments...) + return err +} + +func (w *AppWorkspace) AgentRunShellCommand(ctx context.Context, sessionID, command string, termWidth int, onProgress func(string), isFirstMessage bool) (proto.ShellCommandResponse, error) { + var persist shell.PersistFunc + if sessionID != "" { + persist = func(cmd, output string, exitCode int) error { + return shell.PersistOutput(ctx, w.app.Messages, sessionID, cmd, output, exitCode) + } + } + + opts := shell.RunOptions{ + Command: command, + Cwd: w.store.WorkingDir(), + TermWidth: termWidth, + } + + var result shell.CaptureResult + var err error + + if onProgress != nil { + result, err = shell.RunAndCaptureStream(ctx, opts, onProgress) + } else { + result, err = shell.RunAndPersist(ctx, opts, persist) + } + + if err != nil && onProgress == nil { + return proto.ShellCommandResponse{}, err + } + + // Persist if we used the streaming path (persist wasn't called by RunAndPersist). + if onProgress != nil && persist != nil { + if persistErr := persist(command, result.Output, result.ExitCode); persistErr != nil { + slog.Error("Failed to persist shell command output", "error", persistErr, "command", command) + } + } + + // Generate a title from the shell command if it was the first message. + if isFirstMessage && w.app.AgentCoordinator != nil { + titleCtx := context.WithoutCancel(ctx) + w.app.AgentCoordinator.GenerateTitle(titleCtx, sessionID, "$ "+command) + } + + return proto.ShellCommandResponse{ + Output: result.Output, + ExitCode: result.ExitCode, + }, nil +} + +func (w *AppWorkspace) AgentCancel(sessionID string) { + if w.app.AgentCoordinator != nil { + w.app.AgentCoordinator.Cancel(sessionID) + } +} + +func (w *AppWorkspace) AgentIsBusy() bool { + if w.app.AgentCoordinator == nil { + return false + } + return w.app.AgentCoordinator.IsBusy() +} + +func (w *AppWorkspace) AgentIsSessionBusy(sessionID string) bool { + if w.app.AgentCoordinator == nil { + return false + } + return w.app.AgentCoordinator.IsSessionBusy(sessionID) +} + +func (w *AppWorkspace) AgentModel() AgentModel { + if w.app.AgentCoordinator == nil { + return AgentModel{} + } + m := w.app.AgentCoordinator.Model() + return AgentModel{ + CatwalkCfg: m.CatwalkCfg, + ModelCfg: m.ModelCfg, + } +} + +func (w *AppWorkspace) AgentIsReady() bool { + return w.app.AgentCoordinator != nil +} + +func (w *AppWorkspace) AgentReadyErr() error { + if w.app.AgentCoordinator == nil { + return ErrAgentNotInitialized + } + return nil +} + +func (w *AppWorkspace) AgentQueuedPrompts(sessionID string) int { + if w.app.AgentCoordinator == nil { + return 0 + } + return w.app.AgentCoordinator.QueuedPrompts(sessionID) +} + +func (w *AppWorkspace) AgentQueuedPromptsList(sessionID string) []string { + if w.app.AgentCoordinator == nil { + return nil + } + return w.app.AgentCoordinator.QueuedPromptsList(sessionID) +} + +func (w *AppWorkspace) AgentClearQueue(sessionID string) { + if w.app.AgentCoordinator != nil { + w.app.AgentCoordinator.ClearQueue(sessionID) + } +} + +func (w *AppWorkspace) AgentSummarize(ctx context.Context, sessionID string) error { + if w.app.AgentCoordinator == nil { + return errors.New("agent coordinator not initialized") + } + return w.app.AgentCoordinator.Summarize(ctx, sessionID) +} + +func (w *AppWorkspace) UpdateAgentModel(ctx context.Context) error { + return w.app.UpdateAgentModel(ctx) +} + +func (w *AppWorkspace) InitCoderAgent(ctx context.Context) error { + return w.app.InitCoderAgent(ctx) +} + +func (w *AppWorkspace) InitCoderAgentNonInteractive(ctx context.Context) error { + return w.app.InitCoderAgentNonInteractive(ctx) +} + +func (w *AppWorkspace) GetDefaultSmallModel(providerID string) config.SelectedModel { + return w.app.GetDefaultSmallModel(providerID) +} + +// -- Permissions -- + +func (w *AppWorkspace) PermissionGrant(perm permission.PermissionRequest) bool { + return w.app.Permissions.Grant(perm) +} + +func (w *AppWorkspace) PermissionGrantPersistent(perm permission.PermissionRequest) bool { + return w.app.Permissions.GrantPersistent(perm) +} + +func (w *AppWorkspace) PermissionDeny(perm permission.PermissionRequest) bool { + return w.app.Permissions.Deny(perm) +} + +func (w *AppWorkspace) PermissionSkipRequests() bool { + return w.app.Permissions.SkipRequests() +} + +func (w *AppWorkspace) PermissionSetSkipRequests(skip bool) { + w.app.Permissions.SetSkipRequests(skip) +} + +// -- Questions -- + +func (w *AppWorkspace) QuestionAnswer(responses []question.Answer) bool { + return w.app.Questions.Answer(responses) +} + +func (w *AppWorkspace) QuestionCancel() bool { + return w.app.Questions.Cancel() +} + +// -- FileTracker -- + +func (w *AppWorkspace) FileTrackerRecordRead(ctx context.Context, sessionID, path string) { + w.app.FileTracker.RecordRead(ctx, sessionID, path) +} + +func (w *AppWorkspace) FileTrackerLastReadTime(ctx context.Context, sessionID, path string) time.Time { + return w.app.FileTracker.LastReadTime(ctx, sessionID, path) +} + +func (w *AppWorkspace) FileTrackerListReadFiles(ctx context.Context, sessionID string) ([]string, error) { + return w.app.FileTracker.ListReadFiles(ctx, sessionID) +} + +// -- History -- + +func (w *AppWorkspace) ListSessionHistory(ctx context.Context, sessionID string) ([]history.File, error) { + return w.app.History.ListBySession(ctx, sessionID) +} + +// -- LSP -- + +func (w *AppWorkspace) LSPStart(ctx context.Context, path string) { + w.app.LSPManager.Start(ctx, path) +} + +func (w *AppWorkspace) LSPStopAll(ctx context.Context) { + w.app.LSPManager.StopAll(ctx) +} + +func (w *AppWorkspace) LSPGetStates() map[string]LSPClientInfo { + states := app.GetLSPStates() + result := make(map[string]LSPClientInfo, len(states)) + for k, v := range states { + result[k] = LSPClientInfo{ + Name: v.Name, + State: v.State, + Error: v.Error, + DiagnosticCount: v.DiagnosticCount, + ConnectedAt: v.ConnectedAt, + } + } + return result +} + +func (w *AppWorkspace) LSPGetDiagnosticCounts(name string) lsp.DiagnosticCounts { + state, ok := app.GetLSPState(name) + if !ok || state.Client == nil { + return lsp.DiagnosticCounts{} + } + return state.Client.GetDiagnosticCounts() +} + +// -- Config (read-only) -- + +func (w *AppWorkspace) Config() *config.Config { + return w.store.Config() +} + +func (w *AppWorkspace) WorkingDir() string { + return w.store.WorkingDir() +} + +func (w *AppWorkspace) Resolver() config.VariableResolver { + return w.store.Resolver() +} + +// -- Config mutations -- + +func (w *AppWorkspace) UpdatePreferredModel(scope config.Scope, modelType config.SelectedModelType, model config.SelectedModel) error { + return w.store.UpdatePreferredModel(scope, modelType, model) +} + +func (w *AppWorkspace) SetCompactMode(scope config.Scope, enabled bool) error { + return w.store.SetCompactMode(scope, enabled) +} + +func (w *AppWorkspace) SetProviderAPIKey(scope config.Scope, providerID string, apiKey any) error { + if err := w.store.SetProviderAPIKey(scope, providerID, apiKey); err != nil { + return err + } + w.store.SignalAuthComplete(providerID) + return nil +} + +func (w *AppWorkspace) SetConfigField(scope config.Scope, key string, value any) error { + return w.store.SetConfigField(scope, key, value) +} + +func (w *AppWorkspace) RemoveConfigField(scope config.Scope, key string) error { + return w.store.RemoveConfigField(scope, key) +} + +func (w *AppWorkspace) ImportCopilot() (*oauth.Token, bool) { + return w.store.ImportCopilot() +} + +func (w *AppWorkspace) RefreshOAuthToken(ctx context.Context, scope config.Scope, providerID string) error { + return w.store.RefreshOAuthToken(ctx, scope, providerID) +} + +// -- Project lifecycle -- + +func (w *AppWorkspace) ProjectNeedsInitialization() (bool, error) { + return config.ProjectNeedsInitialization(w.store) +} + +func (w *AppWorkspace) MarkProjectInitialized() error { + return config.MarkProjectInitialized(w.store) +} + +func (w *AppWorkspace) InitializePrompt() (string, error) { + return agent.InitializePrompt(w.store) +} + +func (w *AppWorkspace) ListSkills(_ context.Context) ([]skills.CatalogEntry, error) { + mgr := w.app.Skills + return skills.Catalog(mgr.ActiveSkills(), mgr.ResolvedPaths(), mgr.WorkingDir()), nil +} + +func (w *AppWorkspace) ReadSkill(_ context.Context, skillID string) ([]byte, skills.SkillReadResult, error) { + mgr := w.app.Skills + return skills.ReadContent(mgr.ActiveSkills(), mgr.ResolvedPaths(), mgr.WorkingDir(), skillID) +} + +// -- MCP operations -- + +func (w *AppWorkspace) MCPGetStates() map[string]mcptools.ClientInfo { + return mcptools.GetStates() +} + +func (w *AppWorkspace) MCPRefreshPrompts(ctx context.Context, name string) { + mcptools.RefreshPrompts(ctx, name) +} + +func (w *AppWorkspace) MCPRefreshResources(ctx context.Context, name string) { + mcptools.RefreshResources(ctx, name) +} + +func (w *AppWorkspace) RefreshMCPTools(ctx context.Context, name string) { + mcptools.RefreshTools(ctx, w.store, name) +} + +func (w *AppWorkspace) ReadMCPResource(ctx context.Context, name, uri string) ([]MCPResourceContents, error) { + contents, err := mcptools.ReadResource(ctx, w.store, name, uri) + if err != nil { + return nil, err + } + result := make([]MCPResourceContents, len(contents)) + for i, c := range contents { + result[i] = MCPResourceContents{ + URI: c.URI, + MIMEType: c.MIMEType, + Text: c.Text, + Blob: c.Blob, + } + } + return result, nil +} + +func (w *AppWorkspace) ListMCPPrompts(context.Context) ([]commands.MCPPrompt, error) { + return commands.LoadMCPPrompts() +} + +func (w *AppWorkspace) GetMCPPrompt(clientID, promptID string, args map[string]string) (string, error) { + return commands.GetMCPPrompt(w.store, clientID, promptID, args) +} + +func (w *AppWorkspace) EnableDockerMCP(ctx context.Context) error { + mcpConfig, err := w.store.PrepareDockerMCPConfig() + if err != nil { + return err + } + + if err := mcptools.InitializeSingle(ctx, config.DockerMCPName, w.store); err != nil { + disableErr := mcptools.DisableSingle(w.store, config.DockerMCPName) + w.store.RemoveDockerMCPInMemory() + return fmt.Errorf("failed to start docker MCP: %w", errors.Join(err, disableErr)) + } + + if err := w.store.PersistDockerMCPConfig(mcpConfig); err != nil { + disableErr := mcptools.DisableSingle(w.store, config.DockerMCPName) + w.store.RemoveDockerMCPInMemory() + return fmt.Errorf("docker MCP started but failed to persist configuration: %w", errors.Join(err, disableErr)) + } + + return nil +} + +func (w *AppWorkspace) DisableDockerMCP() error { + if err := mcptools.DisableSingle(w.store, config.DockerMCPName); err != nil { + return fmt.Errorf("failed to disable docker MCP: %w", err) + } + return w.store.DisableDockerMCP() +} + +func (w *AppWorkspace) MCPAuthenticate(ctx context.Context, name string) error { + return mcptools.AuthenticateMCP(ctx, w.store, name) +} + +func (w *AppWorkspace) MCPPendingAuth() []mcptools.PendingAuthServer { + return mcptools.PendingAuthMCPs(w.store) +} + +func (w *AppWorkspace) MCPAuthURL(name string) string { + return mcptools.MCPAuthURL(name) +} + +// -- Lifecycle -- + +func (w *AppWorkspace) Subscribe(program *tea.Program) { + w.app.Subscribe(program) +} + +func (w *AppWorkspace) Shutdown() { + w.app.Shutdown() +} + +// App returns the underlying app.App instance. +func (w *AppWorkspace) App() *app.App { + return w.app +} + +// Store returns the underlying config store. +func (w *AppWorkspace) Store() *config.ConfigStore { + return w.store +} + +// Compile-time check that AppWorkspace implements Workspace. +var _ Workspace = (*AppWorkspace)(nil) diff --git a/internal/workspace/client_workspace.go b/internal/workspace/client_workspace.go new file mode 100644 index 0000000000000000000000000000000000000000..df65ce7f57cea989126dd92eab5a46df305cf7e6 --- /dev/null +++ b/internal/workspace/client_workspace.go @@ -0,0 +1,1439 @@ +package workspace + +import ( + "context" + "errors" + "fmt" + "log/slog" + "strings" + "sync" + "sync/atomic" + "time" + + tea "charm.land/bubbletea/v2" + "github.com/charmbracelet/crush/internal/agent/notify" + "github.com/charmbracelet/crush/internal/agent/tools/mcp" + "github.com/charmbracelet/crush/internal/app" + "github.com/charmbracelet/crush/internal/client" + "github.com/charmbracelet/crush/internal/commands" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/herdr" + "github.com/charmbracelet/crush/internal/history" + "github.com/charmbracelet/crush/internal/log" + "github.com/charmbracelet/crush/internal/lsp" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/oauth" + "github.com/charmbracelet/crush/internal/permission" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/charmbracelet/crush/internal/question" + "github.com/charmbracelet/crush/internal/session" + "github.com/charmbracelet/crush/internal/skills" + "github.com/charmbracelet/crush/internal/version" + "github.com/charmbracelet/x/powernap/pkg/lsp/protocol" + "github.com/pkg/browser" +) + +// ClientWorkspace implements the Workspace interface by delegating all +// operations to a remote server via the client SDK. It caches the +// proto.Workspace returned at creation time and refreshes it after +// config-mutating operations. +type ClientWorkspace struct { + client *client.Client + + mu sync.RWMutex + ws proto.Workspace + skills *skills.Manager + // lastSession is the most recent session ID reported via + // SetCurrentSession. The subscription loop re-asserts it after a + // reconnect, because the server's per-client presence entry (or the + // whole workspace) may have been re-created in the meantime. + lastSession string + + // subCtx bounds the lifetime of the event subscription (and its + // reconnect loop). Shutdown cancels it so Subscribe stops + // reconnecting instead of racing the teardown. + subCtx context.Context + subCancel context.CancelFunc + // subStarted reports whether the subscription loop ever ran, and + // subDone is closed when it returns. Shutdown uses them to let an + // in-flight workspace recovery finish before it says goodbye to the + // server, so the workspace it releases is the one recovery just + // minted. + subStarted atomic.Bool + subDone chan struct{} + + // herdrClient reports agent state to herdr when running inside + // a herdr-managed pane. Nil when not in a herdr environment. + herdrClient *herdr.Client +} + +// SSE reconnect backoff bounds for the workspace event stream. Declared +// as vars (not consts) so tests can shrink the delays. +var ( + sseReconnectInitialBackoff = 250 * time.Millisecond + sseReconnectMaxBackoff = 10 * time.Second +) + +// NewClientWorkspace creates a new ClientWorkspace that proxies all +// operations through the given client SDK. The ws parameter is the +// proto.Workspace snapshot returned by the server at creation time. The +// snapshot's Skills field seeds a process-local skills.Manager so the +// TUI sees discovery state before the first SSE event arrives. The +// manager is constructed with WithGlobalMirror because the client +// process represents exactly one workspace and the TUI reads +// skills.GetLatestStates directly at construction time. +func NewClientWorkspace(c *client.Client, ws proto.Workspace) *ClientWorkspace { + if ws.Config != nil { + ws.Config.SetupAgents() + ws.Config.NormalizeOptions() + } + states := protoToSkillStates(ws.Skills) + mgr := skills.NewManager(nil, nil, states, skills.WithGlobalMirror()) + subCtx, subCancel := context.WithCancel(context.Background()) + return &ClientWorkspace{ + client: c, + ws: ws, + skills: mgr, + subCtx: subCtx, + subCancel: subCancel, + subDone: make(chan struct{}), + herdrClient: herdr.Init(), + } +} + +// refreshWorkspace re-fetches the workspace from the server, updating +// the cached snapshot. Called after config-mutating operations. +func (w *ClientWorkspace) refreshWorkspace() { + updated, err := w.client.GetWorkspace(context.Background(), w.workspaceID()) + if err != nil { + slog.Error("Failed to refresh workspace", "error", err) + return + } + if updated.Config != nil { + updated.Config.SetupAgents() + updated.Config.NormalizeOptions() + } + w.mu.Lock() + w.ws = *updated + w.mu.Unlock() +} + +// cached returns a snapshot of the cached workspace. +func (w *ClientWorkspace) cached() proto.Workspace { + w.mu.RLock() + defer w.mu.RUnlock() + return w.ws +} + +// workspaceID returns the cached workspace ID. +func (w *ClientWorkspace) workspaceID() string { + return w.cached().ID +} + +// -- Sessions -- + +func (w *ClientWorkspace) CreateSession(ctx context.Context, title string) (session.Session, error) { + sess, err := w.client.CreateSession(ctx, w.workspaceID(), title) + if err != nil { + return session.Session{}, err + } + return protoToSession(*sess), nil +} + +func (w *ClientWorkspace) GetSession(ctx context.Context, sessionID string) (session.Session, error) { + sess, err := w.client.GetSession(ctx, w.workspaceID(), sessionID) + if err != nil { + return session.Session{}, err + } + return protoToSession(*sess), nil +} + +func (w *ClientWorkspace) ListSessions(ctx context.Context) ([]session.Session, error) { + protoSessions, err := w.client.ListSessions(ctx, w.workspaceID()) + if err != nil { + return nil, err + } + sessions := make([]session.Session, len(protoSessions)) + for i, s := range protoSessions { + sessions[i] = protoToSession(s) + } + return sessions, nil +} + +func (w *ClientWorkspace) SaveSession(ctx context.Context, sess session.Session) (session.Session, error) { + saved, err := w.client.SaveSession(ctx, w.workspaceID(), sessionToProto(sess)) + if err != nil { + return session.Session{}, err + } + return protoToSession(*saved), nil +} + +func (w *ClientWorkspace) DeleteSession(ctx context.Context, sessionID string) error { + return w.client.DeleteSession(ctx, w.workspaceID(), sessionID) +} + +func (w *ClientWorkspace) CreateAgentToolSessionID(messageID, toolCallID string) string { + return fmt.Sprintf("%s$$%s", messageID, toolCallID) +} + +func (w *ClientWorkspace) ParseAgentToolSessionID(sessionID string) (string, string, bool) { + parts := strings.Split(sessionID, "$$") + if len(parts) != 2 { + return "", "", false + } + return parts[0], parts[1], true +} + +// SetCurrentSession reports the session this client is currently +// viewing to the server. Empty sessionID clears the entry. Errors +// are propagated to the caller; the TUI logs and ignores them since +// the presence record is a hint, not correctness-critical state. +func (w *ClientWorkspace) SetCurrentSession(ctx context.Context, sessionID string) error { + w.herdrClient.SetSessionID(sessionID) + w.mu.Lock() + w.lastSession = sessionID + w.mu.Unlock() + return w.client.SetCurrentSession(ctx, w.workspaceID(), sessionID) +} + +// -- Messages -- + +func (w *ClientWorkspace) ListMessages(ctx context.Context, sessionID string) ([]message.Message, error) { + msgs, err := w.client.ListMessages(ctx, w.workspaceID(), sessionID) + if err != nil { + return nil, err + } + return protoToMessages(msgs), nil +} + +func (w *ClientWorkspace) ListUserMessages(ctx context.Context, sessionID string) ([]message.Message, error) { + msgs, err := w.client.ListUserMessages(ctx, w.workspaceID(), sessionID) + if err != nil { + return nil, err + } + return protoToMessages(msgs), nil +} + +func (w *ClientWorkspace) ListAllUserMessages(ctx context.Context) ([]message.Message, error) { + msgs, err := w.client.ListAllUserMessages(ctx, w.workspaceID()) + if err != nil { + return nil, err + } + return protoToMessages(msgs), nil +} + +// -- Agent -- + +func (w *ClientWorkspace) AgentRun(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) error { + // The interactive TUI does not consume notify.RunComplete for + // completion detection (it observes message events directly), + // so passing an empty RunID is correct here: it skips the + // correlator stamping path without functional consequences. + return w.client.SendMessage(ctx, w.workspaceID(), sessionID, "", prompt, attachments...) +} + +func (w *ClientWorkspace) AgentRunShellCommand(ctx context.Context, sessionID, command string, termWidth int, _ func(string), _ bool) (proto.ShellCommandResponse, error) { + return w.client.RunShellCommand(ctx, w.workspaceID(), sessionID, command, termWidth) +} + +func (w *ClientWorkspace) AgentCancel(sessionID string) { + _ = w.client.CancelAgentSession(context.Background(), w.workspaceID(), sessionID) +} + +func (w *ClientWorkspace) AgentIsBusy() bool { + info, err := w.client.GetAgentInfo(context.Background(), w.workspaceID()) + if err != nil { + return false + } + return info.IsBusy +} + +func (w *ClientWorkspace) AgentIsSessionBusy(sessionID string) bool { + info, err := w.client.GetAgentSessionInfo(context.Background(), w.workspaceID(), sessionID) + if err != nil { + return false + } + return info.IsBusy +} + +func (w *ClientWorkspace) AgentModel() AgentModel { + info, err := w.client.GetAgentInfo(context.Background(), w.workspaceID()) + if err != nil { + return AgentModel{} + } + return AgentModel{ + CatwalkCfg: info.Model, + ModelCfg: info.ModelCfg, + } +} + +func (w *ClientWorkspace) AgentIsReady() bool { + return w.AgentReadyErr() == nil +} + +func (w *ClientWorkspace) AgentReadyErr() error { + info, err := w.client.GetAgentInfo(context.Background(), w.workspaceID()) + if err != nil { + if errors.Is(err, client.ErrNotFound) { + // The server answered, it just does not know this workspace + // any more. The subscription loop is already re-registering; + // saying "lost connection" here would be plainly wrong. + return ErrWorkspaceGone + } + // The workspace/server could not be reached. This is distinct + // from an initialized-but-not-ready agent: the server may have + // torn the workspace down or restarted underneath us. + return fmt.Errorf("%w: %v", ErrServerUnreachable, err) + } + if !info.IsReady { + return ErrAgentNotInitialized + } + return nil +} + +func (w *ClientWorkspace) AgentQueuedPrompts(sessionID string) int { + count, err := w.client.GetAgentSessionQueuedPrompts(context.Background(), w.workspaceID(), sessionID) + if err != nil { + return 0 + } + return count +} + +func (w *ClientWorkspace) AgentQueuedPromptsList(sessionID string) []string { + prompts, err := w.client.GetAgentSessionQueuedPromptsList(context.Background(), w.workspaceID(), sessionID) + if err != nil { + return nil + } + return prompts +} + +func (w *ClientWorkspace) AgentClearQueue(sessionID string) { + _ = w.client.ClearAgentSessionQueuedPrompts(context.Background(), w.workspaceID(), sessionID) +} + +func (w *ClientWorkspace) AgentSummarize(ctx context.Context, sessionID string) error { + return w.client.AgentSummarizeSession(ctx, w.workspaceID(), sessionID) +} + +func (w *ClientWorkspace) UpdateAgentModel(ctx context.Context) error { + return w.client.UpdateAgent(ctx, w.workspaceID()) +} + +func (w *ClientWorkspace) InitCoderAgent(ctx context.Context) error { + return w.client.InitiateAgentProcessing(ctx, w.workspaceID(), true) +} + +func (w *ClientWorkspace) InitCoderAgentNonInteractive(ctx context.Context) error { + return w.client.InitiateAgentProcessing(ctx, w.workspaceID(), false) +} + +func (w *ClientWorkspace) GetDefaultSmallModel(providerID string) config.SelectedModel { + model, err := w.client.GetDefaultSmallModel(context.Background(), w.workspaceID(), providerID) + if err != nil { + return config.SelectedModel{} + } + return *model +} + +// -- Permissions -- + +func (w *ClientWorkspace) PermissionGrant(perm permission.PermissionRequest) bool { + resolved, _ := w.client.GrantPermission(context.Background(), w.workspaceID(), proto.PermissionGrant{ + Permission: proto.PermissionRequest{ + ID: perm.ID, + SessionID: perm.SessionID, + ToolCallID: perm.ToolCallID, + ToolName: perm.ToolName, + Description: perm.Description, + Action: perm.Action, + Path: perm.Path, + Params: perm.Params, + }, + Action: proto.PermissionAllow, + }) + return resolved +} + +func (w *ClientWorkspace) PermissionGrantPersistent(perm permission.PermissionRequest) bool { + resolved, _ := w.client.GrantPermission(context.Background(), w.workspaceID(), proto.PermissionGrant{ + Permission: proto.PermissionRequest{ + ID: perm.ID, + SessionID: perm.SessionID, + ToolCallID: perm.ToolCallID, + ToolName: perm.ToolName, + Description: perm.Description, + Action: perm.Action, + Path: perm.Path, + Params: perm.Params, + }, + Action: proto.PermissionAllowForSession, + }) + return resolved +} + +func (w *ClientWorkspace) PermissionDeny(perm permission.PermissionRequest) bool { + resolved, _ := w.client.GrantPermission(context.Background(), w.workspaceID(), proto.PermissionGrant{ + Permission: proto.PermissionRequest{ + ID: perm.ID, + SessionID: perm.SessionID, + ToolCallID: perm.ToolCallID, + ToolName: perm.ToolName, + Description: perm.Description, + Action: perm.Action, + Path: perm.Path, + Params: perm.Params, + }, + Action: proto.PermissionDeny, + }) + return resolved +} + +func (w *ClientWorkspace) PermissionSkipRequests() bool { + skip, err := w.client.GetPermissionsSkipRequests(context.Background(), w.workspaceID()) + if err != nil { + return false + } + return skip +} + +func (w *ClientWorkspace) PermissionSetSkipRequests(skip bool) { + _ = w.client.SetPermissionsSkipRequests(context.Background(), w.workspaceID(), skip) +} + +// -- Questions -- + +// QuestionAnswer submits answers for a question via the client SDK. +func (w *ClientWorkspace) QuestionAnswer(responses []question.Answer) bool { + protoResp := proto.QuestionAnswer{ + Responses: make([]proto.QuestionResponse, len(responses)), + } + for i, r := range responses { + protoResp.Responses[i] = proto.QuestionResponse{ + QuestionID: r.QuestionID, + SelectedIDs: r.SelectedIDs, + FillInText: r.FillInText, + Yes: r.Yes, + Notes: r.Notes, + } + } + resolved, err := w.client.AnswerQuestionBatch(context.Background(), w.workspaceID(), protoResp) + if err != nil { + slog.Error("Failed to answer question", "error", err) + return false + } + return resolved +} + +// QuestionCancel cancels the pending question via the client SDK. +func (w *ClientWorkspace) QuestionCancel() bool { + cancelled, err := w.client.CancelQuestionBatch(context.Background(), w.workspaceID()) + if err != nil { + slog.Error("Failed to cancel question", "error", err) + return false + } + return cancelled +} + +// -- FileTracker -- + +func (w *ClientWorkspace) FileTrackerRecordRead(ctx context.Context, sessionID, path string) { + _ = w.client.FileTrackerRecordRead(ctx, w.workspaceID(), sessionID, path) +} + +func (w *ClientWorkspace) FileTrackerLastReadTime(ctx context.Context, sessionID, path string) time.Time { + t, err := w.client.FileTrackerLastReadTime(ctx, w.workspaceID(), sessionID, path) + if err != nil { + return time.Time{} + } + return t +} + +func (w *ClientWorkspace) FileTrackerListReadFiles(ctx context.Context, sessionID string) ([]string, error) { + return w.client.FileTrackerListReadFiles(ctx, w.workspaceID(), sessionID) +} + +// -- History -- + +func (w *ClientWorkspace) ListSessionHistory(ctx context.Context, sessionID string) ([]history.File, error) { + files, err := w.client.ListSessionHistoryFiles(ctx, w.workspaceID(), sessionID) + if err != nil { + return nil, err + } + return protoToFiles(files), nil +} + +// -- LSP -- + +func (w *ClientWorkspace) LSPStart(ctx context.Context, path string) { + _ = w.client.LSPStart(ctx, w.workspaceID(), path) +} + +func (w *ClientWorkspace) LSPStopAll(ctx context.Context) { + _ = w.client.LSPStopAll(ctx, w.workspaceID()) +} + +func (w *ClientWorkspace) LSPGetStates() map[string]LSPClientInfo { + states, err := w.client.GetLSPs(context.Background(), w.workspaceID()) + if err != nil { + return nil + } + result := make(map[string]LSPClientInfo, len(states)) + for k, v := range states { + result[k] = LSPClientInfo{ + Name: v.Name, + State: v.State, + Error: v.Error, + DiagnosticCount: v.DiagnosticCount, + ConnectedAt: v.ConnectedAt, + } + } + return result +} + +func (w *ClientWorkspace) LSPGetDiagnosticCounts(name string) lsp.DiagnosticCounts { + diags, err := w.client.GetLSPDiagnostics(context.Background(), w.workspaceID(), name) + if err != nil { + return lsp.DiagnosticCounts{} + } + var counts lsp.DiagnosticCounts + for _, fileDiags := range diags { + for _, d := range fileDiags { + switch d.Severity { + case protocol.SeverityError: + counts.Error++ + case protocol.SeverityWarning: + counts.Warning++ + case protocol.SeverityInformation: + counts.Information++ + case protocol.SeverityHint: + counts.Hint++ + } + } + } + return counts +} + +// -- Config (read-only) -- + +func (w *ClientWorkspace) Config() *config.Config { + return w.cached().Config +} + +func (w *ClientWorkspace) WorkingDir() string { + return w.cached().Path +} + +func (w *ClientWorkspace) Resolver() config.VariableResolver { + return config.IdentityResolver() +} + +// -- Config mutations -- + +func (w *ClientWorkspace) UpdatePreferredModel(scope config.Scope, modelType config.SelectedModelType, model config.SelectedModel) error { + err := w.client.UpdatePreferredModel(context.Background(), w.workspaceID(), scope, modelType, model) + if err == nil { + w.refreshWorkspace() + } + return err +} + +func (w *ClientWorkspace) SetCompactMode(scope config.Scope, enabled bool) error { + err := w.client.SetCompactMode(context.Background(), w.workspaceID(), scope, enabled) + if err == nil { + w.refreshWorkspace() + } + return err +} + +func (w *ClientWorkspace) SetProviderAPIKey(scope config.Scope, providerID string, apiKey any) error { + err := w.client.SetProviderAPIKey(context.Background(), w.workspaceID(), scope, providerID, apiKey) + if err == nil { + w.refreshWorkspace() + } + return err +} + +func (w *ClientWorkspace) SetConfigField(scope config.Scope, key string, value any) error { + err := w.client.SetConfigField(context.Background(), w.workspaceID(), scope, key, value) + if err == nil { + w.refreshWorkspace() + } + return err +} + +func (w *ClientWorkspace) RemoveConfigField(scope config.Scope, key string) error { + err := w.client.RemoveConfigField(context.Background(), w.workspaceID(), scope, key) + if err == nil { + w.refreshWorkspace() + } + return err +} + +func (w *ClientWorkspace) ImportCopilot() (*oauth.Token, bool) { + token, ok, err := w.client.ImportCopilot(context.Background(), w.workspaceID()) + if err != nil { + return nil, false + } + if ok { + w.refreshWorkspace() + } + return token, ok +} + +func (w *ClientWorkspace) RefreshOAuthToken(ctx context.Context, scope config.Scope, providerID string) error { + err := w.client.RefreshOAuthToken(ctx, w.workspaceID(), scope, providerID) + if err == nil { + w.refreshWorkspace() + } + return err +} + +// -- Project lifecycle -- + +func (w *ClientWorkspace) ProjectNeedsInitialization() (bool, error) { + return w.client.ProjectNeedsInitialization(context.Background(), w.workspaceID()) +} + +func (w *ClientWorkspace) MarkProjectInitialized() error { + return w.client.MarkProjectInitialized(context.Background(), w.workspaceID()) +} + +func (w *ClientWorkspace) InitializePrompt() (string, error) { + return w.client.GetInitializePrompt(context.Background(), w.workspaceID()) +} + +func (w *ClientWorkspace) ListSkills(ctx context.Context) ([]skills.CatalogEntry, error) { + entries, err := w.client.ListSkills(ctx, w.workspaceID()) + if err != nil { + return nil, err + } + result := make([]skills.CatalogEntry, len(entries)) + for i, entry := range entries { + result[i] = skills.CatalogEntry{ + ID: entry.ID, + Name: entry.Name, + Description: entry.Description, + Label: entry.Label, + Source: skills.SourceType(entry.Source), + UserInvocable: entry.UserInvocable, + } + } + return result, nil +} + +func (w *ClientWorkspace) ReadSkill(ctx context.Context, skillID string) ([]byte, skills.SkillReadResult, error) { + resp, err := w.client.ReadSkill(ctx, w.workspaceID(), skillID) + if err != nil { + return nil, skills.SkillReadResult{}, err + } + return resp.Content, skills.SkillReadResult{ + Name: resp.Result.Name, + Description: resp.Result.Description, + Source: skills.SourceType(resp.Result.Source), + Builtin: resp.Result.Builtin, + }, nil +} + +// -- MCP operations -- + +func (w *ClientWorkspace) MCPGetStates() map[string]mcp.ClientInfo { + states, err := w.client.MCPGetStates(context.Background(), w.workspaceID()) + if err != nil { + return nil + } + result := make(map[string]mcp.ClientInfo, len(states)) + for k, v := range states { + result[k] = mcp.ClientInfo{ + Name: v.Name, + State: mcp.State(v.State), + Error: v.Error, + Counts: mcp.Counts{ + Tools: v.ToolCount, + Prompts: v.PromptCount, + Resources: v.ResourceCount, + }, + ConnectedAt: v.ConnectedAt, + } + } + return result +} + +func (w *ClientWorkspace) MCPRefreshPrompts(ctx context.Context, name string) { + _ = w.client.MCPRefreshPrompts(ctx, w.workspaceID(), name) +} + +func (w *ClientWorkspace) MCPRefreshResources(ctx context.Context, name string) { + _ = w.client.MCPRefreshResources(ctx, w.workspaceID(), name) +} + +func (w *ClientWorkspace) RefreshMCPTools(ctx context.Context, name string) { + _ = w.client.RefreshMCPTools(ctx, w.workspaceID(), name) +} + +func (w *ClientWorkspace) ReadMCPResource(ctx context.Context, name, uri string) ([]MCPResourceContents, error) { + contents, err := w.client.ReadMCPResource(ctx, w.workspaceID(), name, uri) + if err != nil { + return nil, err + } + result := make([]MCPResourceContents, len(contents)) + for i, c := range contents { + result[i] = MCPResourceContents{ + URI: c.URI, + MIMEType: c.MIMEType, + Text: c.Text, + Blob: c.Blob, + } + } + return result, nil +} + +func (w *ClientWorkspace) ListMCPPrompts(ctx context.Context) ([]commands.MCPPrompt, error) { + prompts, err := w.client.ListMCPPrompts(ctx, w.workspaceID()) + if err != nil { + return nil, err + } + result := make([]commands.MCPPrompt, len(prompts)) + for i, prompt := range prompts { + arguments := make([]commands.Argument, len(prompt.Arguments)) + for j, argument := range prompt.Arguments { + arguments[j] = commands.Argument{ + ID: argument.ID, + Title: argument.Title, + Description: argument.Description, + Required: argument.Required, + } + } + result[i] = commands.MCPPrompt{ + ID: prompt.ID, + Title: prompt.Title, + Description: prompt.Description, + PromptID: prompt.PromptID, + ClientID: prompt.ClientID, + Arguments: arguments, + } + } + return result, nil +} + +func (w *ClientWorkspace) GetMCPPrompt(clientID, promptID string, args map[string]string) (string, error) { + return w.client.GetMCPPrompt(context.Background(), w.workspaceID(), clientID, promptID, args) +} + +func (w *ClientWorkspace) EnableDockerMCP(ctx context.Context) error { + return w.client.EnableDockerMCP(ctx, w.workspaceID()) +} + +func (w *ClientWorkspace) DisableDockerMCP() error { + return w.client.DisableDockerMCP(context.Background(), w.workspaceID()) +} + +func (w *ClientWorkspace) MCPAuthenticate(ctx context.Context, name string) error { + // The server suppresses its own browser open for this flow; the client + // polls the auth URL and opens it locally so the user authorizes on + // their own machine. The OAuth callback listener runs on the server + // (localhost ports shared when server and client are co-located). + authErr := make(chan error, 1) + go func() { + authErr <- w.client.MCPAuthenticate(ctx, w.workspaceID(), name) + }() + + // Poll for the authorization URL so we can open it in the local + // browser as soon as the flow generates one. + var opened bool + ticker := time.NewTicker(200 * time.Millisecond) + defer ticker.Stop() + + for { + select { + case err := <-authErr: + return err + case <-ctx.Done(): + return ctx.Err() + case <-ticker.C: + if opened { + continue + } + if u := w.MCPAuthURL(name); u != "" { + if err := browser.OpenURL(u); err != nil { + slog.Warn("Failed to open MCP OAuth URL in browser", "error", err) + } + opened = true + } + } + } +} + +func (w *ClientWorkspace) MCPPendingAuth() []mcp.PendingAuthServer { + pending, err := w.client.MCPPendingAuth(context.Background(), w.workspaceID()) + if err != nil { + slog.Warn("Failed to fetch MCP pending auth", "error", err) + return nil + } + result := make([]mcp.PendingAuthServer, len(pending)) + for i, p := range pending { + result[i] = mcp.PendingAuthServer{Name: p.Name, URL: p.URL} + } + return result +} + +func (w *ClientWorkspace) MCPAuthURL(name string) string { + // The server's in-progress authorization URL is exposed through the + // pending-auth list while the flow runs; a server in StateNeedsAuth + // paired with an active flow reports its URL here. Poll the server + // for the in-flight URL. + u, err := w.client.MCPAuthURL(context.Background(), w.workspaceID(), name) + if err != nil { + return "" + } + return u +} + +// -- Lifecycle -- + +func (w *ClientWorkspace) Subscribe(program *tea.Program) { + defer log.RecoverPanic("ClientWorkspace.Subscribe", func() { + slog.Info("TUI subscription panic: attempting graceful shutdown") + program.Quit() + }) + + w.runSubscription(program.Send) +} + +// maxRecoveryEscalate is the number of consecutive failed workspace +// recovery attempts after which the loop tells the UI the connection +// looks unrecoverable. It keeps retrying regardless: a hard stop would +// strand a user whose server comes back a minute later, and Shutdown can +// always cancel it. +const maxRecoveryEscalate = 20 + +// recoveryCreateTimeout bounds a single re-registration attempt. It is +// generous because workspace startup is slow (config, database, LSP, MCP); +// it exists only so an unresponsive server cannot pin the subscription +// goroutine indefinitely, and the loop simply retries when it trips. +// A var, not a const, so tests can shrink it. +var recoveryCreateTimeout = 30 * time.Second + +// runSubscription subscribes to the workspace event stream and forwards +// translated events to send, reconnecting with capped exponential +// backoff whenever the stream drops. It returns only when the +// subscription context is cancelled (via Shutdown). Split out from +// Subscribe so it can be tested without a real *tea.Program. +// +// Two failures need more than a retry. A 404 means the server no longer +// knows this workspace, so resubscribing with the same ID can never +// succeed and the loop re-registers instead. And any stream that closes +// loses whatever was published while the client was away, so every +// re-established stream — even one that reconnects on the first try — +// re-asserts the client's session and asks the UI to resync. +func (w *ClientWorkspace) runSubscription(send func(tea.Msg)) { + w.subStarted.Store(true) + defer close(w.subDone) + + backoff := sseReconnectInitialBackoff + degraded := false + recoveryFailures := 0 + markDegraded := func(err error, stuck bool) { + if degraded && !stuck { + return + } + degraded = true + send(ConnectionEvent{State: ConnectionDegraded, Err: err, Stuck: stuck}) + } + + for { + if w.subCtx.Err() != nil { + return + } + + evc, err := w.client.SubscribeEvents(w.subCtx, w.workspaceID()) + if err != nil { + if w.subCtx.Err() != nil { + return + } + markDegraded(err, false) + if !errors.Is(err, client.ErrNotFound) { + slog.Error("Failed to subscribe to workspace events; retrying", + "error", err, "retry_in", backoff) + } else if w.recoverWorkspace() == nil { + // Re-registered: resubscribe immediately under the fresh + // workspace ID. + backoff = sseReconnectInitialBackoff + continue + } else if w.subCtx.Err() == nil { + recoveryFailures++ + if recoveryFailures == maxRecoveryEscalate { + markDegraded(ErrWorkspaceGone, true) + } + } + if !w.sleepOrDone(backoff) { + return + } + backoff = min(backoff*2, sseReconnectMaxBackoff) + continue + } + + if degraded { + degraded = false + recoveryFailures = 0 + w.afterReconnect(send) + } + backoff = sseReconnectInitialBackoff + w.consumeEvents(evc, send) + + // The event channel closed: the server restarted, the stream was + // interrupted, or the workspace briefly went away. Reconnect + // after a short delay instead of leaving the TUI permanently + // orphaned, which is what surfaced as a stuck "coder agent is + // offline". + if w.subCtx.Err() != nil { + return + } + markDegraded(ErrStreamClosed, false) + slog.Warn("Workspace event stream closed; reconnecting", "retry_in", backoff) + if !w.sleepOrDone(backoff) { + return + } + backoff = min(backoff*2, sseReconnectMaxBackoff) + } +} + +// recoverWorkspace re-registers the workspace after the server reported it +// gone: it re-creates it from the cached snapshot (the server's own view of +// path, data dir, flags and env), adopts the new ID, and re-initializes the +// coder agent when the config is ready, mirroring the startup handshake. The +// server's path dedupe means this either rejoins a live sibling workspace or +// mints a fresh one. It must only be called from the subscription goroutine, +// the only writer of the cached ID. +// +// The create deliberately runs detached from the subscription context: the +// server does not abandon a create when the requesting connection goes away, +// so cancelling would hide the outcome while the workspace got registered +// anyway. Riding it out means the ID is known by the time Shutdown looks, +// and retirement covers a lost response regardless. Its own timeout keeps a +// wedged server from pinning the subscription goroutine forever; the client +// SDK sets no request timeout of its own. +func (w *ClientWorkspace) recoverWorkspace() error { + ctx, cancel := context.WithTimeout( + context.WithoutCancel(w.subCtx), recoveryCreateTimeout, + ) + defer cancel() + created, err := w.client.CreateWorkspace(ctx, w.recreateArgs()) + if err != nil { + slog.Error("Failed to re-register workspace; retrying", "error", err) + return err + } + if created.Config != nil { + created.Config.SetupAgents() + created.Config.NormalizeOptions() + } + w.mu.Lock() + oldID := w.ws.ID + w.ws = *created + w.mu.Unlock() + slog.Info("Re-registered workspace after server-side loss", + "old_id", oldID, "new_id", created.ID) + + if created.Config != nil && created.Config.IsConfigured() { + if err := w.InitCoderAgent(w.subCtx); err != nil { + // Matches the startup handshake: agent init failure is + // logged, not fatal, since the user can still pick a model. + slog.Error("Failed to initialize coder agent after workspace recovery", "error", err) + } + } + return nil +} + +// recreateArgs derives the CreateWorkspace request used for recovery from +// the cached snapshot. The ID is dropped so the server can dedupe by +// path or mint a fresh workspace, and Version carries this client's +// version, matching the startup handshake. +func (w *ClientWorkspace) recreateArgs() proto.Workspace { + ws := w.cached() + return proto.Workspace{ + Path: ws.Path, + DataDir: ws.DataDir, + Debug: ws.Debug, + YOLO: ws.YOLO, + Channels: ws.Channels, + Env: ws.Env, + Version: version.Version, + } +} + +// afterReconnect runs once a degraded subscription is re-established. It +// re-asserts the client's current-session selection, since the server's +// presence entry (or the whole workspace) may have been re-created while we +// were away, and tells the UI to resync state published while detached. The +// SSE handler attaches the client before writing its 200, so the presence +// call cannot be rejected as not-attached here. +func (w *ClientWorkspace) afterReconnect(send func(tea.Msg)) { + w.mu.RLock() + sid := w.lastSession + w.mu.RUnlock() + if sid != "" { + if err := w.SetCurrentSession(w.subCtx, sid); err != nil { + slog.Warn("Failed to re-assert current session after reconnect", "error", err) + } + } + send(ConnectionEvent{State: ConnectionRecovered}) +} + +// sleepOrDone waits for d or until the subscription context is +// cancelled. It reports false when the context was cancelled, signalling +// the caller to stop reconnecting. +func (w *ClientWorkspace) sleepOrDone(d time.Duration) bool { + t := time.NewTimer(d) + defer t.Stop() + select { + case <-t.C: + return true + case <-w.subCtx.Done(): + return false + } +} + +// consumeEvents drives the workspace event loop. It is split out from +// Subscribe so tests can drive it without a real *tea.Program. +// ConfigChanged events trigger a workspace refresh; all other events +// are translated into domain types and forwarded to send. +func (w *ClientWorkspace) consumeEvents(evc <-chan any, send func(tea.Msg)) { + for ev := range evc { + // Forward events to herdr if running inside a herdr pane. + if hev := herdr.Translate(ev); hev != nil { + w.herdrClient.HandleEvent(hev) + } + + if _, ok := ev.(pubsub.Event[proto.ConfigChanged]); ok { + w.refreshWorkspace() + continue + } + translated := w.translateEvent(ev) + if translated != nil && send != nil { + send(translated) + } + } +} + +// shutdownDrainTimeout bounds how long Shutdown waits for the subscription +// loop to stop. Exceeding it is not a correctness problem — retiring the +// client releases whatever a late recovery registers — it only makes the +// goodbye less tidy. +const shutdownDrainTimeout = 5 * time.Second + +func (w *ClientWorkspace) Shutdown() { + // Stop the reconnect/recovery loop first, then wait for it: cancelling + // alone does not unwind a workspace recovery that is already in + // flight, and we want to release the workspace that recovery ended up + // with rather than one it is about to replace. + if w.subCancel != nil { + w.subCancel() + } + w.awaitSubscription() + w.herdrClient.Close() + + // Retiring the client releases every claim it holds, on every workspace, + // and blocks any further create from this client ID. That is what makes + // teardown exact even when a recovery create's response was lost: the + // create either landed before this call, and its claim is released here, + // or it arrives afterwards and registers nothing. + err := w.client.RetireClient(context.Background()) + if err == nil { + return + } + if !errors.Is(err, client.ErrUnsupported) { + slog.Warn("Failed to retire client on the server", "error", err) + return + } + // The server predates client retirement, so fall back to releasing + // the workspace we know about. Nothing better is possible against an + // older server. + _ = w.client.DeleteWorkspace(context.Background(), w.workspaceID()) +} + +// awaitSubscription waits for the subscription loop to return. It returns +// immediately when the loop never started, which is the case for +// workspaces shut down before Subscribe runs. +func (w *ClientWorkspace) awaitSubscription() { + if !w.subStarted.Load() || w.subDone == nil { + return + } + t := time.NewTimer(shutdownDrainTimeout) + defer t.Stop() + select { + case <-w.subDone: + case <-t.C: + slog.Warn("Timed out waiting for the workspace subscription to stop") + } +} + +// translateEvent converts proto-typed SSE events into the domain types +// that the TUI's Update() method expects. Skills events also update the +// process-local skills.Manager so callers reading +// skills.GetLatestStates see fresh data. +func (w *ClientWorkspace) translateEvent(ev any) tea.Msg { + switch e := ev.(type) { + case pubsub.Event[proto.LSPEvent]: + return pubsub.Event[LSPEvent]{ + Type: e.Type, + Payload: LSPEvent{ + Type: LSPEventType(e.Payload.Type), + Name: e.Payload.Name, + State: e.Payload.State, + Error: e.Payload.Error, + DiagnosticCount: e.Payload.DiagnosticCount, + }, + } + case pubsub.Event[proto.MCPEvent]: + return pubsub.Event[mcp.Event]{ + Type: e.Type, + Payload: mcp.Event{ + Type: protoToMCPEventType(e.Payload.Type), + Name: e.Payload.Name, + State: mcp.State(e.Payload.State), + Error: e.Payload.Error, + Counts: mcp.Counts{ + Tools: e.Payload.ToolCount, + Prompts: e.Payload.PromptCount, + Resources: e.Payload.ResourceCount, + }, + }, + } + case pubsub.Event[proto.PermissionRequest]: + return pubsub.Event[permission.PermissionRequest]{ + Type: e.Type, + Payload: permission.PermissionRequest{ + ID: e.Payload.ID, + SessionID: e.Payload.SessionID, + ToolCallID: e.Payload.ToolCallID, + ToolName: e.Payload.ToolName, + Description: e.Payload.Description, + Action: e.Payload.Action, + Path: e.Payload.Path, + Params: e.Payload.Params, + }, + } + case pubsub.Event[proto.PermissionNotification]: + return pubsub.Event[permission.PermissionNotification]{ + Type: e.Type, + Payload: permission.PermissionNotification{ + ToolCallID: e.Payload.ToolCallID, + Granted: e.Payload.Granted, + Denied: e.Payload.Denied, + }, + } + case pubsub.Event[proto.QuestionRequest]: + return pubsub.Event[question.Request]{ + Type: e.Type, + Payload: question.Request{ + ID: e.Payload.ID, + SessionID: e.Payload.SessionID, + ToolCallID: e.Payload.ToolCallID, + Questions: protoQuestionsToDomain(e.Payload.Questions), + ConfirmTitle: e.Payload.ConfirmTitle, + ConfirmDescription: e.Payload.ConfirmDescription, + }, + } + case pubsub.Event[proto.QuestionNotification]: + return pubsub.Event[question.Notification]{ + Type: e.Type, + Payload: question.Notification{ + BatchID: e.Payload.BatchID, + }, + } + case pubsub.Event[proto.Message]: + return pubsub.Event[message.Message]{ + Type: e.Type, + Payload: protoToMessage(e.Payload), + } + case pubsub.Event[proto.Session]: + return pubsub.Event[session.Session]{ + Type: e.Type, + Payload: protoToSession(e.Payload), + } + case pubsub.Event[proto.File]: + return pubsub.Event[history.File]{ + Type: e.Type, + Payload: protoToFile(e.Payload), + } + case pubsub.Event[proto.AgentEvent]: + n := notify.Notification{ + SessionID: e.Payload.SessionID, + SessionTitle: e.Payload.SessionTitle, + RunID: e.Payload.RunID, + Type: notify.Type(e.Payload.Type), + AWSSOCommand: e.Payload.AWSSOCommand, + AWSSOURL: e.Payload.AWSSOURL, + } + if e.Payload.Error != nil { + n.Message = e.Payload.Error.Error() + } + return pubsub.Event[notify.Notification]{ + Type: e.Type, + Payload: n, + } + case pubsub.Event[proto.RunComplete]: + // Translate the wire-level proto.RunComplete back into the + // agent's domain notify.RunComplete. Without this case the + // default branch below warns on every run completion in the + // server-mode TUI, even though the TUI itself doesn't act + // on RunComplete — converting silently keeps the workspace + // event bridge symmetric with the server-side wrapEvent. + return pubsub.Event[notify.RunComplete]{ + Type: e.Type, + Payload: notify.RunComplete{ + SessionID: e.Payload.SessionID, + RunID: e.Payload.RunID, + MessageID: e.Payload.MessageID, + Text: e.Payload.Text, + Error: e.Payload.Error, + Cancelled: e.Payload.Cancelled, + }, + } + case pubsub.Event[proto.SkillsEvent]: + states := protoToSkillStates(e.Payload.States) + if w.skills != nil { + w.skills.SetLatestStates(states) + } + return pubsub.Event[skills.Event]{ + Type: e.Type, + Payload: skills.Event{States: states}, + } + case pubsub.Event[proto.UpdateAvailable]: + return app.UpdateAvailableMsg{ + CurrentVersion: e.Payload.CurrentVersion, + LatestVersion: e.Payload.LatestVersion, + IsDevelopment: e.Payload.IsDevelopment, + } + default: + slog.Warn("Unknown event type in translateEvent", "type", fmt.Sprintf("%T", ev)) + return nil + } +} + +func protoToMCPEventType(t proto.MCPEventType) mcp.EventType { + switch t { + case proto.MCPEventStateChanged: + return mcp.EventStateChanged + case proto.MCPEventToolsListChanged: + return mcp.EventToolsListChanged + case proto.MCPEventPromptsListChanged: + return mcp.EventPromptsListChanged + case proto.MCPEventResourcesListChanged: + return mcp.EventResourcesListChanged + default: + return mcp.EventStateChanged + } +} + +// protoToSession converts a wire-level proto.Session into the domain +// session.Session. Fields that exist only on the wire (computed-on-read +// signals like IsBusy, and any future presence counters) are +// intentionally dropped here: session.Session models persisted state, +// not transient runtime signals. UI features that need those signals +// should either extend session.Session or read them from the proto +// payload directly before this conversion runs. +func protoToSession(s proto.Session) session.Session { + return session.Session{ + ID: s.ID, + ParentSessionID: s.ParentSessionID, + Title: s.Title, + SummaryMessageID: s.SummaryMessageID, + MessageCount: s.MessageCount, + PromptTokens: s.PromptTokens, + CompletionTokens: s.CompletionTokens, + Cost: s.Cost, + Todos: protoToTodos(s.Todos), + CreatedAt: s.CreatedAt, + UpdatedAt: s.UpdatedAt, + } +} + +func protoToTodos(todos []proto.Todo) []session.Todo { + if len(todos) == 0 { + return nil + } + out := make([]session.Todo, len(todos)) + for i, t := range todos { + out[i] = session.Todo{ + Content: t.Content, + Status: session.TodoStatus(t.Status), + ActiveForm: t.ActiveForm, + } + } + return out +} + +func protoToFile(f proto.File) history.File { + return history.File{ + ID: f.ID, + SessionID: f.SessionID, + Path: f.Path, + Content: f.Content, + Version: f.Version, + CreatedAt: f.CreatedAt, + UpdatedAt: f.UpdatedAt, + } +} + +func protoToMessage(m proto.Message) message.Message { + msg := message.Message{ + ID: m.ID, + SessionID: m.SessionID, + Role: message.MessageRole(m.Role), + Model: m.Model, + Provider: m.Provider, + PrismModelID: m.PrismModelID, + PrismModelName: m.PrismModelName, + PrismHypercreditSavings: m.PrismHypercreditSavings, + PrismDollarSavings: m.PrismDollarSavings, + CreatedAt: m.CreatedAt, + UpdatedAt: m.UpdatedAt, + IsSummaryMessage: m.IsSummaryMessage, + } + + for _, p := range m.Parts { + switch v := p.(type) { + case proto.TextContent: + msg.Parts = append(msg.Parts, message.TextContent{Text: v.Text}) + case proto.ReasoningContent: + msg.Parts = append(msg.Parts, message.ReasoningContent{ + Thinking: v.Thinking, + Signature: v.Signature, + StartedAt: v.StartedAt, + FinishedAt: v.FinishedAt, + }) + case proto.ToolCall: + msg.Parts = append(msg.Parts, message.ToolCall{ + ID: v.ID, + Name: v.Name, + Input: v.Input, + Finished: v.Finished, + }) + case proto.ToolResult: + msg.Parts = append(msg.Parts, message.ToolResult{ + ToolCallID: v.ToolCallID, + Name: v.Name, + Content: v.Content, + Data: v.Data, + MIMEType: v.MIMEType, + Metadata: v.Metadata, + IsError: v.IsError, + }) + case proto.Finish: + msg.Parts = append(msg.Parts, message.Finish{ + Reason: message.FinishReason(v.Reason), + Time: v.Time, + Message: v.Message, + Details: v.Details, + }) + case proto.ImageURLContent: + msg.Parts = append(msg.Parts, message.ImageURLContent{URL: v.URL, Detail: v.Detail}) + case proto.BinaryContent: + msg.Parts = append(msg.Parts, message.BinaryContent{Path: v.Path, MIMEType: v.MIMEType, Data: v.Data}) + case proto.ShellCommand: + msg.Parts = append(msg.Parts, message.ShellCommand{ + Command: v.Command, + Output: v.Output, + ExitCode: v.ExitCode, + }) + } + } + + return msg +} + +func protoToMessages(msgs []proto.Message) []message.Message { + out := make([]message.Message, len(msgs)) + for i, m := range msgs { + out[i] = protoToMessage(m) + } + return out +} + +func protoToFiles(files []proto.File) []history.File { + out := make([]history.File, len(files)) + for i, f := range files { + out[i] = protoToFile(f) + } + return out +} + +func sessionToProto(s session.Session) proto.Session { + return proto.Session{ + ID: s.ID, + ParentSessionID: s.ParentSessionID, + Title: s.Title, + SummaryMessageID: s.SummaryMessageID, + MessageCount: s.MessageCount, + PromptTokens: s.PromptTokens, + CompletionTokens: s.CompletionTokens, + Cost: s.Cost, + Todos: todosToProto(s.Todos), + CreatedAt: s.CreatedAt, + UpdatedAt: s.UpdatedAt, + } +} + +// protoToSkillStates reconstructs internal skill state slices from +// their wire representation. Non-empty Error strings are turned into +// synthetic error values; the TUI never type-asserts on Err. +func protoToSkillStates(in []proto.SkillState) []*skills.SkillState { + if len(in) == 0 { + return nil + } + out := make([]*skills.SkillState, len(in)) + for i, s := range in { + state := &skills.SkillState{ + Name: s.Name, + Path: s.Path, + State: skills.DiscoveryState(s.State), + } + if s.Error != "" { + state.Err = errors.New(s.Error) + } + out[i] = state + } + return out +} + +func todosToProto(todos []session.Todo) []proto.Todo { + if len(todos) == 0 { + return nil + } + out := make([]proto.Todo, len(todos)) + for i, t := range todos { + out[i] = proto.Todo{ + Content: t.Content, + Status: string(t.Status), + ActiveForm: t.ActiveForm, + } + } + return out +} + +func protoQuestionsToDomain(qs []proto.QuestionItem) []question.Question { + if len(qs) == 0 { + return nil + } + out := make([]question.Question, len(qs)) + for i, q := range qs { + choices := make([]question.Choice, len(q.Choices)) + for j, c := range q.Choices { + choices[j] = question.Choice{ + ID: c.ID, + Label: c.Label, + Description: c.Description, + } + } + out[i] = question.Question{ + ID: q.ID, + Type: question.Type(q.Type), + Label: q.Label, + Text: q.Question, + Description: q.Description, + Choices: choices, + } + } + return out +} diff --git a/internal/workspace/client_workspace_test.go b/internal/workspace/client_workspace_test.go new file mode 100644 index 0000000000000000000000000000000000000000..62e7537e55ae1f0f6ad378e9be17714a45a457ca --- /dev/null +++ b/internal/workspace/client_workspace_test.go @@ -0,0 +1,976 @@ +package workspace + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + tea "charm.land/bubbletea/v2" + "github.com/charmbracelet/crush/internal/app" + "github.com/charmbracelet/crush/internal/client" + "github.com/charmbracelet/crush/internal/commands" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/permission" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/charmbracelet/crush/internal/skills" + "github.com/stretchr/testify/require" +) + +// TestProtoToMessageToolResult ensures that ToolResult metadata, +// data, and MIME type survive the conversion from proto on the +// client. Without these fields the TUI cannot render rich tool +// output (e.g. syntax-highlighted code from view, diffs from edit, +// images, etc.) and falls back to the raw LLM-facing string. +func TestProtoToMessageToolResult(t *testing.T) { + t.Parallel() + + src := proto.Message{ + ID: "m1", + Role: proto.Tool, + Parts: []proto.ContentPart{ + proto.ToolResult{ + ToolCallID: "call-1", + Name: "view", + Content: "\n 1| hi\n", + Data: "base64data", + MIMEType: "image/png", + Metadata: `{"file_path":"/tmp/x","content":"hi"}`, + IsError: false, + }, + }, + } + + got := protoToMessage(src) + require.Len(t, got.Parts, 1) + tr, ok := got.Parts[0].(message.ToolResult) + require.True(t, ok, "expected message.ToolResult, got %T", got.Parts[0]) + require.Equal(t, "call-1", tr.ToolCallID) + require.Equal(t, "view", tr.Name) + require.Equal(t, "\n 1| hi\n", tr.Content) + require.Equal(t, "base64data", tr.Data) + require.Equal(t, "image/png", tr.MIMEType) + require.Equal(t, `{"file_path":"/tmp/x","content":"hi"}`, tr.Metadata) + require.False(t, tr.IsError) +} + +// TestProtoToMessagePrismModel ensures the Prism-routed model fields survive +// the conversion from proto. Without them the client TUI cannot show which +// model actually served each turn on Hyper's model router. +func TestProtoToMessagePrismModel(t *testing.T) { + t.Parallel() + + src := proto.Message{ + ID: "m1", + Role: proto.Assistant, + Model: "prism-model", + Provider: "hyper", + PrismModelID: "prism-42", + PrismModelName: "GLM 5.3", + + PrismHypercreditSavings: protoPtrFloat(1.5), + PrismDollarSavings: protoPtrFloat(0.002), + } + + got := protoToMessage(src) + require.Equal(t, "prism-42", got.PrismModelID) + require.Equal(t, "GLM 5.3", got.PrismModelName) + require.NotNil(t, got.PrismHypercreditSavings) + require.Equal(t, 1.5, *got.PrismHypercreditSavings) + require.NotNil(t, got.PrismDollarSavings) + require.Equal(t, 0.002, *got.PrismDollarSavings) +} + +func protoPtrFloat(v float64) *float64 { return &v } + +// TestClientWorkspace_PermissionGrantMapping verifies that +// PermissionGrant on the ClientWorkspace serializes a one-time grant +// (proto.PermissionAllow) and PermissionGrantPersistent serializes a +// persistent grant (proto.PermissionAllowForSession). A swap between +// these two would silently flip "allow once" into "remember for the +// session", and vice versa, so we pin the wire mapping here. +func TestClientWorkspace_PermissionGrantMapping(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + call func(*ClientWorkspace, permission.PermissionRequest) + want proto.PermissionAction + }{ + { + name: "Grant -> PermissionAllow", + call: func(w *ClientWorkspace, p permission.PermissionRequest) { + w.PermissionGrant(p) + }, + want: proto.PermissionAllow, + }, + { + name: "GrantPersistent -> PermissionAllowForSession", + call: func(w *ClientWorkspace, p permission.PermissionRequest) { + w.PermissionGrantPersistent(p) + }, + want: proto.PermissionAllowForSession, + }, + { + name: "Deny -> PermissionDeny", + call: func(w *ClientWorkspace, p permission.PermissionRequest) { + w.PermissionDeny(p) + }, + want: proto.PermissionDeny, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + var got proto.PermissionGrant + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, http.MethodPost, r.Method) + require.Equal(t, "/v1/workspaces/ws-1/permissions/grant", r.URL.Path) + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + require.NoError(t, json.Unmarshal(body, &got)) + require.NoError(t, json.NewEncoder(w).Encode(proto.PermissionGrantResponse{Resolved: true})) + })) + defer srv.Close() + + u, err := url.Parse(srv.URL) + require.NoError(t, err) + c, err := client.NewClient(t.TempDir(), "tcp", u.Host) + require.NoError(t, err) + + ws := NewClientWorkspace(c, proto.Workspace{ID: "ws-1"}) + + perm := permission.PermissionRequest{ + ID: "req-1", + SessionID: "sess-1", + ToolCallID: "tc-1", + ToolName: "tool", + Description: "do thing", + Action: "act", + Path: "/tmp/p", + } + tc.call(ws, perm) + + require.Equal(t, tc.want, got.Action) + require.Equal(t, "req-1", got.Permission.ID) + require.Equal(t, "sess-1", got.Permission.SessionID) + require.Equal(t, "tc-1", got.Permission.ToolCallID) + require.Equal(t, "tool", got.Permission.ToolName) + require.Equal(t, "act", got.Permission.Action) + require.Equal(t, "/tmp/p", got.Permission.Path) + }) + } +} + +// TestProtoToSkillStates verifies that the wire representation of skill +// discovery states reconstructs identical values on the client, +// including synthetic errors derived from Error strings. +func TestProtoToSkillStates(t *testing.T) { + t.Parallel() + + in := []proto.SkillState{ + {Name: "ok", Path: "/p/ok", State: proto.SkillStateNormal}, + {Name: "broken", Path: "/p/broken", State: proto.SkillStateError, Error: "bad frontmatter"}, + } + + got := protoToSkillStates(in) + require.Len(t, got, 2) + require.Equal(t, "ok", got[0].Name) + require.Equal(t, skills.StateNormal, got[0].State) + require.NoError(t, got[0].Err) + require.Equal(t, "broken", got[1].Name) + require.Equal(t, skills.StateError, got[1].State) + require.EqualError(t, got[1].Err, "bad frontmatter") +} + +// TestTranslateEvent_Skills verifies that an incoming proto.SkillsEvent +// is converted into pubsub.Event[skills.Event] and that the +// client-process skill cache is updated as a side effect, so callers +// reading skills.GetLatestStates see fresh data after each delta. +func TestTranslateEvent_Skills(t *testing.T) { + // Not parallel - touches the package-level skills cache via the + // manager constructed with WithGlobalMirror. + prev := skills.GetLatestStates() + t.Cleanup(func() { skills.SetLatestStates(prev) }) + + skills.SetLatestStates(nil) + + w := NewClientWorkspace(nil, proto.Workspace{}) + ev := pubsub.Event[proto.SkillsEvent]{ + Type: pubsub.UpdatedEvent, + Payload: proto.SkillsEvent{ + States: []proto.SkillState{ + {Name: "from-server", Path: "/p", State: proto.SkillStateNormal}, + }, + }, + } + + out := w.translateEvent(ev) + got, ok := out.(pubsub.Event[skills.Event]) + require.True(t, ok, "expected pubsub.Event[skills.Event], got %T", out) + require.Len(t, got.Payload.States, 1) + require.Equal(t, "from-server", got.Payload.States[0].Name) + + // Manager (with WithGlobalMirror) propagated to the package cache. + cached := skills.GetLatestStates() + require.Len(t, cached, 1) + require.Equal(t, "from-server", cached[0].Name) +} + +// TestNewClientWorkspace_SeedsSkillsCache verifies that the snapshot in +// proto.Workspace.Skills populates the package-level cache the TUI +// reads at construction time, eliminating the race between TUI startup +// and the first SSE event. +func TestNewClientWorkspace_SeedsSkillsCache(t *testing.T) { + // Not parallel - touches the package-level skills cache. + prev := skills.GetLatestStates() + t.Cleanup(func() { skills.SetLatestStates(prev) }) + + skills.SetLatestStates(nil) + + _ = NewClientWorkspace(nil, proto.Workspace{ + Skills: []proto.SkillState{ + {Name: "seeded", Path: "/p", State: proto.SkillStateNormal}, + }, + }) + + got := skills.GetLatestStates() + require.Len(t, got, 1) + require.Equal(t, "seeded", got[0].Name) +} + +// TestTranslateEvent_UpdateAvailable verifies that an incoming +// proto.UpdateAvailable event is converted back into the +// app.UpdateAvailableMsg that the TUI expects, so client/server mode +// shows the same update notification as local mode. +func TestTranslateEvent_UpdateAvailable(t *testing.T) { + t.Parallel() + + w := NewClientWorkspace(nil, proto.Workspace{}) + ev := pubsub.Event[proto.UpdateAvailable]{ + Type: pubsub.UpdatedEvent, + Payload: proto.UpdateAvailable{ + CurrentVersion: "1.0.0", + LatestVersion: "1.1.0", + IsDevelopment: true, + }, + } + + out := w.translateEvent(ev) + got, ok := out.(app.UpdateAvailableMsg) + require.True(t, ok, "expected app.UpdateAvailableMsg, got %T", out) + require.Equal(t, "1.0.0", got.CurrentVersion) + require.Equal(t, "1.1.0", got.LatestVersion) + require.True(t, got.IsDevelopment) +} + +func TestClientWorkspaceListMCPPrompts(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/v1/workspaces/ws-1/mcp/prompts", r.URL.Path) + require.NoError(t, json.NewEncoder(w).Encode([]proto.MCPPrompt{ + { + ID: "server:review", + PromptID: "review", + ClientID: "server", + Arguments: []proto.MCPPromptArgument{ + {ID: "focus", Title: "Focus", Required: true}, + }, + }, + })) + })) + defer srv.Close() + + u, err := url.Parse(srv.URL) + require.NoError(t, err) + c, err := client.NewClient(t.TempDir(), "tcp", u.Host) + require.NoError(t, err) + workspace := NewClientWorkspace(c, proto.Workspace{ID: "ws-1"}) + + got, err := workspace.ListMCPPrompts(t.Context()) + require.NoError(t, err) + require.Equal(t, []commands.MCPPrompt{ + { + ID: "server:review", + PromptID: "review", + ClientID: "server", + Arguments: []commands.Argument{ + {ID: "focus", Title: "Focus", Required: true}, + }, + }, + }, got) +} + +func TestClientWorkspaceListMCPPromptsServerError(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + defer srv.Close() + + u, err := url.Parse(srv.URL) + require.NoError(t, err) + c, err := client.NewClient(t.TempDir(), "tcp", u.Host) + require.NoError(t, err) + workspace := NewClientWorkspace(c, proto.Workspace{ID: "ws-1"}) + + _, err = workspace.ListMCPPrompts(t.Context()) + require.Error(t, err) + require.Contains(t, err.Error(), "status code 500") +} + +// TestClientWorkspace_ReconnectsOnStreamDrop verifies that the event +// subscription loop reconnects after the SSE stream drops instead of +// leaving the TUI permanently orphaned (which surfaced as a stuck +// "coder agent is offline"), and that Shutdown stops the loop. +func TestClientWorkspace_ReconnectsOnStreamDrop(t *testing.T) { + // Shrink the backoff so several reconnects happen quickly. + origInitial, origMax := sseReconnectInitialBackoff, sseReconnectMaxBackoff + sseReconnectInitialBackoff = 5 * time.Millisecond + sseReconnectMaxBackoff = 20 * time.Millisecond + t.Cleanup(func() { + sseReconnectInitialBackoff = origInitial + sseReconnectMaxBackoff = origMax + }) + + var subscribes atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.HasSuffix(r.URL.Path, "/events") { + // Any other bookkeeping call (e.g. GetWorkspace) just + // gets an empty OK; the test only cares about the stream. + w.WriteHeader(http.StatusOK) + return + } + subscribes.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + // Drop the stream immediately so the client must reconnect. + })) + defer srv.Close() + + u, err := url.Parse(srv.URL) + require.NoError(t, err) + c, err := client.NewClient(t.TempDir(), "tcp", u.Host) + require.NoError(t, err) + + ws := NewClientWorkspace(c, proto.Workspace{ID: "ws-1"}) + + done := make(chan struct{}) + go func() { + ws.runSubscription(func(tea.Msg) {}) + close(done) + }() + + // The loop must reconnect several times as the server keeps + // dropping the stream. + require.Eventually(t, func() bool { return subscribes.Load() >= 3 }, + 2*time.Second, 5*time.Millisecond, + "subscription loop should reconnect after the stream drops") + + // Shutdown cancels the subscription context; the loop must return. + ws.Shutdown() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("runSubscription did not return after Shutdown") + } +} + +// TestClientWorkspace_SubscriptionStopsWhenServerDown verifies the +// reconnect loop does not spin forever after Shutdown even when it can +// never connect (server unreachable). +func TestClientWorkspace_SubscriptionStopsWhenServerDown(t *testing.T) { + origInitial, origMax := sseReconnectInitialBackoff, sseReconnectMaxBackoff + sseReconnectInitialBackoff = 5 * time.Millisecond + sseReconnectMaxBackoff = 20 * time.Millisecond + t.Cleanup(func() { + sseReconnectInitialBackoff = origInitial + sseReconnectMaxBackoff = origMax + }) + + // Port 1 is not listening: SubscribeEvents fails on every attempt. + c, err := client.NewClient(t.TempDir(), "tcp", "127.0.0.1:1") + require.NoError(t, err) + ws := NewClientWorkspace(c, proto.Workspace{ID: "ws-1"}) + + done := make(chan struct{}) + go func() { + ws.runSubscription(func(tea.Msg) {}) + close(done) + }() + + // Let it retry a few times, then shut down. + time.Sleep(30 * time.Millisecond) + ws.Shutdown() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("runSubscription did not return after Shutdown while server was down") + } +} + +// TestClientWorkspace_AgentReadyErr distinguishes a server that reports +// an uninitialized agent from a server that cannot be reached, so the UI +// can show an actionable message instead of a blanket "agent offline". +func TestClientWorkspace_AgentReadyErr(t *testing.T) { + t.Parallel() + + t.Run("ready", func(t *testing.T) { + t.Parallel() + ws := agentInfoWorkspace(t, proto.AgentInfo{IsReady: true}) + require.NoError(t, ws.AgentReadyErr()) + require.True(t, ws.AgentIsReady()) + }) + + t.Run("not initialized", func(t *testing.T) { + t.Parallel() + ws := agentInfoWorkspace(t, proto.AgentInfo{IsReady: false}) + err := ws.AgentReadyErr() + require.ErrorIs(t, err, ErrAgentNotInitialized) + require.NotErrorIs(t, err, ErrServerUnreachable) + require.False(t, ws.AgentIsReady()) + }) + + t.Run("server unreachable", func(t *testing.T) { + t.Parallel() + c, err := client.NewClient(t.TempDir(), "tcp", "127.0.0.1:1") + require.NoError(t, err) + ws := NewClientWorkspace(c, proto.Workspace{ID: "ws-1"}) + readyErr := ws.AgentReadyErr() + require.ErrorIs(t, readyErr, ErrServerUnreachable) + require.NotErrorIs(t, readyErr, ErrAgentNotInitialized) + require.False(t, ws.AgentIsReady()) + }) +} + +// agentInfoWorkspace returns a ClientWorkspace whose server answers the +// agent-info endpoint with the given info. +func agentInfoWorkspace(t *testing.T, info proto.AgentInfo) *ClientWorkspace { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/v1/workspaces/ws-1/agent", r.URL.Path) + require.NoError(t, json.NewEncoder(w).Encode(info)) + })) + t.Cleanup(srv.Close) + + u, err := url.Parse(srv.URL) + require.NoError(t, err) + c, err := client.NewClient(t.TempDir(), "tcp", u.Host) + require.NoError(t, err) + return NewClientWorkspace(c, proto.Workspace{ID: "ws-1"}) +} + +// -- Workspace recovery -- +// +// The bug these cover: the reconnect loop retried SubscribeEvents with +// the workspace ID captured at startup, forever. Once the server no +// longer knew that ID — because the workspace was torn down under the +// client, or the server was replaced — every retry and every other +// request answered 404 permanently, and a sibling session kept the +// server (and its 404s) alive. + +// recoveryServer is a scripted server for the recovery tests. It 404s +// event subscriptions for stale workspace IDs, mints a new ID on create, +// and records what the client did. +type recoveryServer struct { + mu sync.Mutex + // liveID is the only workspace ID whose event stream is served. + liveID string + // nextID names the workspace the next create hands back. + nextID string + // createErr, when set, fails creates instead of serving them. + createErr func() int + // hangUpOnCreate drops the connection mid-response so the client + // never learns the ID of a workspace the server did register. + hangUpOnCreate bool + + creates int + streams int + sessionPosts []string + retired []string + deleted []string +} + +func (s *recoveryServer) start(t *testing.T) *client.Client { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + s.mu.Lock() + defer s.mu.Unlock() + switch { + case r.Method == http.MethodPost && r.URL.Path == "/v1/workspaces": + s.creates++ + if s.createErr != nil { + http.Error(w, "no", s.createErr()) + return + } + if s.hangUpOnCreate { + // Register the workspace, then hang up before the client + // can read the ID. This is the case no amount of client + // bookkeeping can name. + s.liveID = s.nextID + if hj, ok := w.(http.Hijacker); ok { + conn, _, err := hj.Hijack() + require.NoError(t, err) + _ = conn.Close() + } + return + } + s.liveID = s.nextID + require.NoError(t, json.NewEncoder(w).Encode(proto.Workspace{ + ID: s.liveID, Path: "/tmp/recover", + })) + case r.Method == http.MethodDelete && strings.HasPrefix(r.URL.Path, "/v1/clients/"): + s.retired = append(s.retired, strings.TrimPrefix(r.URL.Path, "/v1/clients/")) + s.liveID = "" + case r.Method == http.MethodDelete && strings.HasPrefix(r.URL.Path, "/v1/workspaces/"): + s.deleted = append(s.deleted, strings.TrimPrefix(r.URL.Path, "/v1/workspaces/")) + case strings.HasSuffix(r.URL.Path, "/current-session"): + id := strings.TrimSuffix(strings.TrimPrefix(r.URL.Path, "/v1/workspaces/"), "/current-session") + if id != s.liveID { + http.Error(w, "workspace not found", http.StatusNotFound) + return + } + var req proto.CurrentSession + require.NoError(t, json.NewDecoder(r.Body).Decode(&req)) + s.sessionPosts = append(s.sessionPosts, req.SessionID) + case strings.HasSuffix(r.URL.Path, "/events"): + id := strings.TrimSuffix(strings.TrimPrefix(r.URL.Path, "/v1/workspaces/"), "/events") + if id != s.liveID { + http.Error(w, "workspace not found", http.StatusNotFound) + return + } + s.streams++ + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + // Drop the stream so the loop keeps cycling. + default: + w.WriteHeader(http.StatusOK) + } + })) + t.Cleanup(srv.Close) + + u, err := url.Parse(srv.URL) + require.NoError(t, err) + c, err := client.NewClient(t.TempDir(), "tcp", u.Host) + require.NoError(t, err) + return c +} + +func (s *recoveryServer) snapshot(f func(*recoveryServer)) { + s.mu.Lock() + defer s.mu.Unlock() + f(s) +} + +// connectionRecorder collects the ConnectionEvents the subscription loop +// reports to the UI. +type connectionRecorder struct { + mu sync.Mutex + events []ConnectionEvent +} + +func (r *connectionRecorder) send(msg tea.Msg) { + ev, ok := msg.(ConnectionEvent) + if !ok { + return + } + r.mu.Lock() + defer r.mu.Unlock() + r.events = append(r.events, ev) +} + +func (r *connectionRecorder) states() []ConnectionState { + r.mu.Lock() + defer r.mu.Unlock() + out := make([]ConnectionState, len(r.events)) + for i, ev := range r.events { + out[i] = ev.State + } + return out +} + +func (r *connectionRecorder) sawStuck() bool { + r.mu.Lock() + defer r.mu.Unlock() + for _, ev := range r.events { + if ev.Stuck { + return true + } + } + return false +} + +// TestClientWorkspace_RecoversFromWorkspaceGone is the headline client +// regression: a 404 for the cached workspace ID must make the client +// re-register, adopt the new ID, re-assert the session it was viewing, +// and tell the UI to resync — not retry a dead ID forever. +func TestClientWorkspace_RecoversFromWorkspaceGone(t *testing.T) { + t.Cleanup(SetSSEBackoffForTest(time.Millisecond, 5*time.Millisecond)) + + srv := &recoveryServer{liveID: "", nextID: "ws-2"} + c := srv.start(t) + + ws := NewClientWorkspace(c, proto.Workspace{ID: "ws-1", Path: "/tmp/recover"}) + // The client was viewing a session, so recovery has to restore that + // selection: the server's presence entry died with the workspace. + ws.mu.Lock() + ws.lastSession = "sess-1" + ws.mu.Unlock() + + rec := &connectionRecorder{} + done := make(chan struct{}) + go func() { + ws.runSubscription(rec.send) + close(done) + }() + + require.Eventually(t, func() bool { return ws.workspaceID() == "ws-2" }, + 3*time.Second, 5*time.Millisecond, + "a 404 must make the client re-register instead of retrying a dead ID") + require.Eventually(t, func() bool { + var posts []string + srv.snapshot(func(s *recoveryServer) { posts = append(posts, s.sessionPosts...) }) + return len(posts) > 0 && posts[0] == "sess-1" + }, 3*time.Second, 5*time.Millisecond, + "the recovered workspace must be told which session the client is on") + + require.Eventually(t, func() bool { + states := rec.states() + var degraded, recovered bool + for _, s := range states { + degraded = degraded || s == ConnectionDegraded + recovered = recovered || s == ConnectionRecovered + } + return degraded && recovered + }, 3*time.Second, 5*time.Millisecond, + "the UI must be told to resync after the workspace was re-created") + + ws.Shutdown() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("runSubscription did not return after Shutdown") + } +} + +// TestClientWorkspace_ResyncsAfterPlainStreamDrop covers the quieter +// failure: the stream closes and the very next subscribe succeeds. That +// used to be treated as a non-event, so the TUI silently showed state +// frozen at the moment of the drop — everything published while the +// client was away is gone, and the server had dropped its presence entry. +func TestClientWorkspace_ResyncsAfterPlainStreamDrop(t *testing.T) { + t.Cleanup(SetSSEBackoffForTest(time.Millisecond, 5*time.Millisecond)) + + srv := &recoveryServer{liveID: "ws-1"} + c := srv.start(t) + + ws := NewClientWorkspace(c, proto.Workspace{ID: "ws-1", Path: "/tmp/blip"}) + ws.mu.Lock() + ws.lastSession = "sess-9" + ws.mu.Unlock() + + rec := &connectionRecorder{} + done := make(chan struct{}) + go func() { + ws.runSubscription(rec.send) + close(done) + }() + + require.Eventually(t, func() bool { + states := rec.states() + var degraded, recovered bool + for _, s := range states { + degraded = degraded || s == ConnectionDegraded + recovered = recovered || s == ConnectionRecovered + } + return degraded && recovered + }, 3*time.Second, 5*time.Millisecond, + "a reconnect that succeeds first try must still resync") + + var posts int + var creates int + srv.snapshot(func(s *recoveryServer) { posts, creates = len(s.sessionPosts), s.creates }) + require.Positive(t, posts, "the session selection must be re-asserted after the drop") + require.Zero(t, creates, "a live workspace must not be re-created") + require.Equal(t, "ws-1", ws.workspaceID()) + + ws.Shutdown() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("runSubscription did not return after Shutdown") + } +} + +// TestClientWorkspace_EscalatesUnrecoverableConnection checks the loop +// keeps trying when re-registration itself keeps failing, and escalates +// the UI notice rather than stopping: a hard stop would strand a user +// whose server comes back a minute later. +func TestClientWorkspace_EscalatesUnrecoverableConnection(t *testing.T) { + t.Cleanup(SetSSEBackoffForTest(time.Millisecond, 2*time.Millisecond)) + + srv := &recoveryServer{createErr: func() int { return http.StatusInternalServerError }} + c := srv.start(t) + + ws := NewClientWorkspace(c, proto.Workspace{ID: "ws-1", Path: "/tmp/hopeless"}) + rec := &connectionRecorder{} + done := make(chan struct{}) + go func() { + ws.runSubscription(rec.send) + close(done) + }() + + require.Eventually(t, rec.sawStuck, 5*time.Second, 5*time.Millisecond, + "repeated recovery failures must escalate to a persistent notice") + var creates int + srv.snapshot(func(s *recoveryServer) { creates = s.creates }) + require.GreaterOrEqual(t, creates, maxRecoveryEscalate, + "the loop must keep retrying rather than giving up") + + ws.Shutdown() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("runSubscription did not return after Shutdown") + } +} + +// TestClientWorkspace_ShutdownRetiresAfterLostCreateResponse is the +// quit-during-recovery case, in its nastiest form: the client asks for a +// replacement workspace and the connection dies before it can read the +// ID. The server registered the workspace anyway, so the client is +// holding a claim it cannot name. Retiring the client is what makes +// teardown exact here — there is nothing to guess and nothing to scan. +func TestClientWorkspace_ShutdownRetiresAfterLostCreateResponse(t *testing.T) { + t.Cleanup(SetSSEBackoffForTest(time.Millisecond, 5*time.Millisecond)) + + srv := &recoveryServer{nextID: "ws-orphan", hangUpOnCreate: true} + c := srv.start(t) + + ws := NewClientWorkspace(c, proto.Workspace{ID: "ws-1", Path: "/tmp/recover"}) + done := make(chan struct{}) + go func() { + ws.runSubscription(func(tea.Msg) {}) + close(done) + }() + + // Wait until the server has registered a workspace the client never + // learned about: that is the state Shutdown has to clean up. + require.Eventually(t, func() bool { + var live string + srv.snapshot(func(s *recoveryServer) { live = s.liveID }) + return live == "ws-orphan" + }, 3*time.Second, 5*time.Millisecond) + require.NotEqual(t, "ws-orphan", ws.workspaceID(), + "the client must genuinely not know the workspace's ID") + + ws.Shutdown() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("runSubscription did not return after Shutdown") + } + + srv.snapshot(func(s *recoveryServer) { + require.Equal(t, []string{c.ClientID()}, s.retired, + "Shutdown must retire the client, which releases claims it cannot name") + require.Empty(t, s.liveID, "no workspace may be left holding the server open") + }) +} + +// TestClientWorkspace_ShutdownFallsBackForLegacyServer keeps quits working +// against a server that predates client retirement: a 404 for the retire +// endpoint must fall back to releasing the workspace by ID, not be +// reported as a failed teardown. +func TestClientWorkspace_ShutdownFallsBackForLegacyServer(t *testing.T) { + t.Parallel() + + var mu sync.Mutex + var deleted []string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + defer mu.Unlock() + switch { + case strings.HasPrefix(r.URL.Path, "/v1/clients/"): + http.NotFound(w, r) + case r.Method == http.MethodDelete: + deleted = append(deleted, strings.TrimPrefix(r.URL.Path, "/v1/workspaces/")) + default: + w.WriteHeader(http.StatusOK) + } + })) + t.Cleanup(srv.Close) + + u, err := url.Parse(srv.URL) + require.NoError(t, err) + c, err := client.NewClient(t.TempDir(), "tcp", u.Host) + require.NoError(t, err) + + NewClientWorkspace(c, proto.Workspace{ID: "ws-1"}).Shutdown() + + mu.Lock() + defer mu.Unlock() + require.Equal(t, []string{"ws-1"}, deleted) +} + +// TestClientWorkspace_AgentReadyErr_WorkspaceGone checks the status the +// UI is given while recovery runs. A 404 from a live server used to print +// "lost connection to the crush server: ... status code 404", which is +// both wrong and unactionable. +func TestClientWorkspace_AgentReadyErr_WorkspaceGone(t *testing.T) { + t.Parallel() + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.NotFound(w, r) + })) + t.Cleanup(srv.Close) + + u, err := url.Parse(srv.URL) + require.NoError(t, err) + c, err := client.NewClient(t.TempDir(), "tcp", u.Host) + require.NoError(t, err) + + ws := NewClientWorkspace(c, proto.Workspace{ID: "ws-1"}) + readyErr := ws.AgentReadyErr() + require.ErrorIs(t, readyErr, ErrWorkspaceGone) + require.NotErrorIs(t, readyErr, ErrServerUnreachable) + require.NotErrorIs(t, readyErr, ErrAgentNotInitialized) +} + +// TestClientWorkspace_ShutdownWaitsForInFlightRecovery pins the ordering +// requirement: quitting must stop recovery before saying goodbye to the +// server. Cancelling alone does not unwind a create that is already in +// flight, so a goodbye sent first would be followed by the workspace the +// create went on to register — an orphan keeping the server alive and +// blocking the next upgrade. +func TestClientWorkspace_ShutdownWaitsForInFlightRecovery(t *testing.T) { + t.Cleanup(SetSSEBackoffForTest(time.Millisecond, 5*time.Millisecond)) + + var mu sync.Mutex + var ops []string + creating := make(chan struct{}) + var once sync.Once + live := "" + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodPost && r.URL.Path == "/v1/workspaces": + once.Do(func() { close(creating) }) + // Slow enough that Shutdown lands mid-create. + //nolint:forbidigo // The overlap is the point of the test. + time.Sleep(300 * time.Millisecond) + mu.Lock() + ops = append(ops, "create") + live = "ws-2" + mu.Unlock() + require.NoError(t, json.NewEncoder(w).Encode(proto.Workspace{ID: "ws-2"})) + case strings.HasPrefix(r.URL.Path, "/v1/clients/"): + mu.Lock() + ops = append(ops, "retire") + live = "" + mu.Unlock() + case strings.HasSuffix(r.URL.Path, "/events"): + mu.Lock() + known := live == strings.TrimSuffix(strings.TrimPrefix(r.URL.Path, "/v1/workspaces/"), "/events") + mu.Unlock() + if !known { + http.Error(w, "workspace not found", http.StatusNotFound) + return + } + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + default: + w.WriteHeader(http.StatusOK) + } + })) + t.Cleanup(srv.Close) + + u, err := url.Parse(srv.URL) + require.NoError(t, err) + c, err := client.NewClient(t.TempDir(), "tcp", u.Host) + require.NoError(t, err) + + ws := NewClientWorkspace(c, proto.Workspace{ID: "ws-1", Path: "/tmp/quit-mid-recovery"}) + done := make(chan struct{}) + go func() { + ws.runSubscription(func(tea.Msg) {}) + close(done) + }() + + <-creating + ws.Shutdown() + select { + case <-done: + case <-time.After(3 * time.Second): + t.Fatal("runSubscription did not return after Shutdown") + } + + mu.Lock() + defer mu.Unlock() + require.Equal(t, []string{"create", "retire"}, ops, + "the goodbye must follow the create it was racing, not precede it") + require.Empty(t, live, "the recovered workspace must not outlive the client") + require.Equal(t, "ws-2", ws.workspaceID(), + "the client must have adopted the workspace it created before releasing it") +} + +// TestClientWorkspace_RecoveryCreateIsBounded checks a wedged server cannot +// pin the recovery attempt forever. The client SDK sets no request timeout, +// so without a bound here the subscription goroutine would block for good +// on a server that accepts the create and never answers. +func TestClientWorkspace_RecoveryCreateIsBounded(t *testing.T) { + t.Parallel() + + release := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost && r.URL.Path == "/v1/workspaces" { + <-release // Accept the create and never answer it. + return + } + http.NotFound(w, r) + })) + // Unwedge the handler before closing the server, which waits on it. + t.Cleanup(func() { + close(release) + srv.Close() + }) + + u, err := url.Parse(srv.URL) + require.NoError(t, err) + c, err := client.NewClient(t.TempDir(), "tcp", u.Host) + require.NoError(t, err) + + ws := NewClientWorkspace(c, proto.Workspace{ID: "ws-1", Path: "/tmp/wedged"}) + // Shrink the bound so the test does not wait out the production window. + // The recovery must give up on its own, without anything cancelling it: + // the create is deliberately detached from the subscription context. + orig := recoveryCreateTimeoutForTest(50 * time.Millisecond) + t.Cleanup(orig) + + done := make(chan error, 1) + go func() { done <- ws.recoverWorkspace() }() + + select { + case err := <-done: + require.Error(t, err, "a create that never answers must not report success") + case <-time.After(5 * time.Second): + t.Fatal("recoverWorkspace blocked on an unresponsive server") + } +} diff --git a/internal/workspace/export_test.go b/internal/workspace/export_test.go new file mode 100644 index 0000000000000000000000000000000000000000..d45fa54799bf1e0a4c1dade834415ff44090d02d --- /dev/null +++ b/internal/workspace/export_test.go @@ -0,0 +1,47 @@ +package workspace + +import ( + "time" + + tea "charm.land/bubbletea/v2" +) + +// ConsumeEventsForTest runs the event-handling loop on the given +// channel, invoking send for translated domain messages and refreshing +// the cached workspace snapshot on ConfigChanged. Exposed for +// cross-package integration tests that cannot rely on a real +// *tea.Program. It returns when evc is closed. +func (w *ClientWorkspace) ConsumeEventsForTest(evc <-chan any, send func(tea.Msg)) { + w.consumeEvents(evc, send) +} + +// RunSubscriptionForTest runs the full subscribe/reconnect/recovery loop. +// Exposed for cross-package integration tests. It returns once the +// workspace is shut down. +func (w *ClientWorkspace) RunSubscriptionForTest(send func(tea.Msg)) { + w.runSubscription(send) +} + +// WorkspaceIDForTest returns the currently cached workspace ID, which +// recovery may have re-minted. +func (w *ClientWorkspace) WorkspaceIDForTest() string { + return w.workspaceID() +} + +// SetSSEBackoffForTest shrinks the subscription reconnect backoff and +// returns a restore function for t.Cleanup. +func SetSSEBackoffForTest(initial, maxBackoff time.Duration) (restore func()) { + origInitial, origMax := sseReconnectInitialBackoff, sseReconnectMaxBackoff + sseReconnectInitialBackoff, sseReconnectMaxBackoff = initial, maxBackoff + return func() { + sseReconnectInitialBackoff, sseReconnectMaxBackoff = origInitial, origMax + } +} + +// recoveryCreateTimeoutForTest shrinks the bound on a single workspace +// re-registration attempt and returns a restore function for t.Cleanup. +func recoveryCreateTimeoutForTest(d time.Duration) (restore func()) { + orig := recoveryCreateTimeout + recoveryCreateTimeout = d + return func() { recoveryCreateTimeout = orig } +} diff --git a/internal/workspace/multiclient_integration_test.go b/internal/workspace/multiclient_integration_test.go new file mode 100644 index 0000000000000000000000000000000000000000..c9b1076c384bf7c63bd7dc6f3990d3c0b4e9837e --- /dev/null +++ b/internal/workspace/multiclient_integration_test.go @@ -0,0 +1,362 @@ +package workspace_test + +import ( + "context" + "net/http/httptest" + "net/url" + "testing" + "time" + + tea "charm.land/bubbletea/v2" + "github.com/charmbracelet/crush/internal/client" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/charmbracelet/crush/internal/server" + "github.com/charmbracelet/crush/internal/workspace" + "github.com/stretchr/testify/require" +) + +// xdgIsolate redirects HOME and XDG_* to fresh temp dirs so config +// loading does not touch the host's real config. +func xdgIsolate(t *testing.T) { + t.Helper() + t.Setenv("HOME", t.TempDir()) + t.Setenv("XDG_CACHE_HOME", t.TempDir()) + t.Setenv("XDG_CONFIG_HOME", t.TempDir()) + t.Setenv("XDG_DATA_HOME", t.TempDir()) +} + +// runtimeServer wires the production server handler around an +// httptest.NewServer for integration testing. +type runtimeServer struct { + srv *server.Server + httpSrv *httptest.Server + host string +} + +func newRuntimeServer(t *testing.T) *runtimeServer { + t.Helper() + s := server.NewServer(nil, "tcp", "127.0.0.1:0") + hs := httptest.NewServer(s.Handler()) + t.Cleanup(hs.Close) + + u, err := url.Parse(hs.URL) + require.NoError(t, err) + return &runtimeServer{srv: s, httpSrv: hs, host: u.Host} +} + +func (r *runtimeServer) newClient(t *testing.T, path string) *client.Client { + t.Helper() + c, err := client.NewClient(path, "tcp", r.host) + require.NoError(t, err) + // Retire the client during cleanup so the server releases every + // claim it holds and tears the workspace down at once, closing the + // pooled DB connection. Without this a test that leaves clients + // attached (the SSE cache tests never shut theirs down) keeps the + // workspace, and its open crush.db, alive past t.TempDir cleanup, + // which Windows cannot remove while the file is locked. + t.Cleanup(func() { _ = c.RetireClient(context.Background()) }) + return c +} + +// TestClientWorkspace_ConfigChangedRefreshesSiblingCache is the +// cross-client refresh end-to-end test required by PLAN item 4. Two +// ClientWorkspace instances pointed at the same backend workspace +// subscribe to events; when one mutates configuration via the server, +// the other's cached Config snapshot reflects the new value without +// a manual refresh. +func TestClientWorkspace_ConfigChangedRefreshesSiblingCache(t *testing.T) { + xdgIsolate(t) + rt := newRuntimeServer(t) + + cwd := t.TempDir() + dataDir := t.TempDir() + + cA := rt.newClient(t, cwd) + cB := rt.newClient(t, cwd) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + wsProto, err := cA.CreateWorkspace(ctx, proto.Workspace{Path: cwd, DataDir: dataDir}) + require.NoError(t, err) + // Client B joins the same workspace by path; the server + // deduplicates and returns the existing workspace. + wsProtoB, err := cB.CreateWorkspace(ctx, proto.Workspace{Path: cwd, DataDir: dataDir}) + require.NoError(t, err) + require.Equal(t, wsProto.ID, wsProtoB.ID) + + wsA := workspace.NewClientWorkspace(cA, *wsProto) + wsB := workspace.NewClientWorkspace(cB, *wsProtoB) + + // Both clients attach event streams. They run for the + // lifetime of the test; cancelling via context tears them + // down. consumeEvents is exercised by Subscribe in production; + // here we run it inline so we don't need a real *tea.Program. + evcA, err := cA.SubscribeEvents(ctx, wsProto.ID) + require.NoError(t, err) + evcB, err := cB.SubscribeEvents(ctx, wsProto.ID) + require.NoError(t, err) + + go wsA.ConsumeEventsForTest(evcA, func(tea.Msg) {}) + go wsB.ConsumeEventsForTest(evcB, func(tea.Msg) {}) + + // Pre-condition: neither cache has compact mode enabled yet. + require.NotNil(t, wsA.Config()) + require.NotNil(t, wsB.Config()) + require.False(t, compactMode(wsA.Config()), "compact mode must start disabled on client A") + require.False(t, compactMode(wsB.Config()), "compact mode must start disabled on client B") + + // Client A flips a real config-mutating workspace operation + // (SetCompactMode) via the server. PLAN item 4 acceptance: + // B's cached ws.Config must reflect this change without restart. + // SetCompactMode is used over UpdatePreferredModel because the + // latter's autoReload reverts unknown-provider models back to + // defaults during configureSelectedModels, which would make the + // assertion test infrastructure rather than the cache wiring. + require.NoError(t, wsA.SetCompactMode(config.ScopeGlobal, true)) + + // Client A writes and refreshes synchronously inside + // SetCompactMode, so its cache must already reflect the change. + // Eventually here absorbs any background work but should pass + // immediately. + require.Eventually(t, func() bool { return compactMode(wsA.Config()) }, + 3*time.Second, 25*time.Millisecond, + "client A cache must reflect its own compact-mode mutation") + + // Client B must see the same change via the ConfigChanged SSE + // event triggering its own cached refresh. + require.Eventually(t, func() bool { return compactMode(wsB.Config()) }, + 3*time.Second, 25*time.Millisecond, + "client B cache must reflect A's compact-mode mutation via SSE") +} + +// compactMode is a tiny accessor that survives nil intermediates so +// the Eventually polling loop can call it on a transient cache state. +func compactMode(cfg *config.Config) bool { + if cfg == nil || cfg.Options == nil { + return false + } + return cfg.Options.TUI.CompactMode +} + +// TestClientWorkspace_ConfigChangedSignalArrives is a smaller test +// that asserts the SSE wiring delivers a ConfigChanged event to the +// raw client subscription. It catches breakage in the +// wrapEvent/decoder bridge independent of the workspace cache. +func TestClientWorkspace_ConfigChangedSignalArrives(t *testing.T) { + xdgIsolate(t) + rt := newRuntimeServer(t) + + cwd := t.TempDir() + dataDir := t.TempDir() + + c := rt.newClient(t, cwd) + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + wsProto, err := c.CreateWorkspace(ctx, proto.Workspace{Path: cwd, DataDir: dataDir}) + require.NoError(t, err) + + evc, err := c.SubscribeEvents(ctx, wsProto.ID) + require.NoError(t, err) + + require.NoError(t, c.SetConfigField(ctx, wsProto.ID, config.ScopeGlobal, "options.debug", true)) + + gotConfigChanged := false + deadline := time.After(3 * time.Second) +loop: + for !gotConfigChanged { + select { + case ev, ok := <-evc: + if !ok { + break loop + } + if cc, isCC := ev.(pubsub.Event[proto.ConfigChanged]); isCC { + require.Equal(t, wsProto.ID, cc.Payload.WorkspaceID) + gotConfigChanged = true + } + case <-deadline: + break loop + } + } + require.True(t, gotConfigChanged, "expected ConfigChanged event over SSE") +} + +// -- Concurrent session lifecycle -- +// +// These drive the real HTTP surface, because the bug they cover only +// appears when two sessions share a server: the first client's workspace +// died, and the sibling kept the server (and its 404s) alive so the first +// client never got a fresh start. + +// TestServer_ConcurrentSessionsSameDirShareOneWorkspace checks the +// baseline for two sessions in one directory: they share a workspace, and +// one of them leaving must not disturb the other. +func TestServer_ConcurrentSessionsSameDirShareOneWorkspace(t *testing.T) { + xdgIsolate(t) + rt := newRuntimeServer(t) + rt.srv.Backend().SetDetachGrace(0) + + cwd, dataDir := t.TempDir(), t.TempDir() + cA, cB := rt.newClient(t, cwd), rt.newClient(t, cwd) + ctx := t.Context() + + wsA, err := cA.CreateWorkspace(ctx, proto.Workspace{Path: cwd, DataDir: dataDir}) + require.NoError(t, err) + wsB, err := cB.CreateWorkspace(ctx, proto.Workspace{Path: cwd, DataDir: dataDir}) + require.NoError(t, err) + require.Equal(t, wsA.ID, wsB.ID, "same directory must share one workspace") + + // A leaves. B's claim must keep the workspace addressable. + require.NoError(t, cA.RetireClient(ctx)) + got, err := cB.GetWorkspace(ctx, wsB.ID) + require.NoError(t, err, "one client leaving must not destroy the shared workspace") + require.Equal(t, wsB.ID, got.ID) +} + +// TestServer_ConcurrentSessionsDifferentDirsAreIndependent is the other +// half: two directories get two workspaces on one server, and retiring +// one client must leave the other's workspace completely alone. +func TestServer_ConcurrentSessionsDifferentDirsAreIndependent(t *testing.T) { + xdgIsolate(t) + rt := newRuntimeServer(t) + rt.srv.Backend().SetDetachGrace(0) + + cwdA, cwdB := t.TempDir(), t.TempDir() + cA, cB := rt.newClient(t, cwdA), rt.newClient(t, cwdB) + ctx := t.Context() + + wsA, err := cA.CreateWorkspace(ctx, proto.Workspace{Path: cwdA, DataDir: t.TempDir()}) + require.NoError(t, err) + wsB, err := cB.CreateWorkspace(ctx, proto.Workspace{Path: cwdB, DataDir: t.TempDir()}) + require.NoError(t, err) + require.NotEqual(t, wsA.ID, wsB.ID) + + require.NoError(t, cA.RetireClient(ctx)) + + _, err = cB.GetWorkspace(ctx, wsB.ID) + require.NoError(t, err, "a sibling session's workspace must survive") + _, err = cB.GetWorkspace(ctx, wsA.ID) + require.ErrorIs(t, err, client.ErrNotFound, "the retired client's workspace must be gone") +} + +// TestServer_RefusesShutdownWhileWorkspaceLive is the invariant that keeps +// an upgrading client from killing a running session. The refusal has to +// come from the server: a client checking idleness itself needs a second +// round trip that a new session can slip into. +func TestServer_RefusesShutdownWhileWorkspaceLive(t *testing.T) { + xdgIsolate(t) + rt := newRuntimeServer(t) + + cwd := t.TempDir() + cLive := rt.newClient(t, cwd) + ctx := t.Context() + + ws, err := cLive.CreateWorkspace(ctx, proto.Workspace{Path: cwd, DataDir: t.TempDir()}) + require.NoError(t, err) + + // A second client, standing in for one that found a version mismatch. + err = rt.newClient(t, t.TempDir()).ShutdownServerIfIdle(ctx) + require.ErrorIs(t, err, client.ErrServerBusy, + "a server hosting a workspace must refuse to stand down") + + _, err = cLive.GetWorkspace(ctx, ws.ID) + require.NoError(t, err, "the live session's workspace must be untouched") +} + +// TestServer_DetachGraceSurvivesStreamBlip is the server-side half of the +// SSE regression. Cutting the stream used to destroy the workspace +// instantly, so the client's reconnect — 250ms later — came back to an ID +// the server no longer knew, and 404'd from then on. +func TestServer_DetachGraceSurvivesStreamBlip(t *testing.T) { + xdgIsolate(t) + rt := newRuntimeServer(t) + rt.srv.Backend().SetDetachGrace(30 * time.Second) + + cwd := t.TempDir() + c := rt.newClient(t, cwd) + + ws, err := c.CreateWorkspace(t.Context(), proto.Workspace{Path: cwd, DataDir: t.TempDir()}) + require.NoError(t, err) + + streamCtx, killStream := context.WithCancel(t.Context()) + evc, err := c.SubscribeEvents(streamCtx, ws.ID) + require.NoError(t, err) + + // Cut the stream the way a network blip does: no release first. + killStream() + for range evc { //nolint:revive // Drain until the server closes it. + } + + require.Eventually(t, func() bool { + _, err := c.GetWorkspace(t.Context(), ws.ID) + return err == nil + }, 3*time.Second, 25*time.Millisecond, + "the workspace must survive the drop so the reconnect can re-attach") + + // The reconnect lands and the workspace is still the same one. + evc2, err := c.SubscribeEvents(t.Context(), ws.ID) + require.NoError(t, err, "the client must be able to re-attach to its workspace") + require.NotNil(t, evc2) +} + +// TestClientWorkspace_RecoversAfterServerSideTeardown is the whole bug, +// end to end: the server drops the client's workspace while the server +// itself stays up (a sibling session keeps it alive). The client must +// notice the 404, re-register, and end up on a workspace it can use again +// rather than 404ing forever. +func TestClientWorkspace_RecoversAfterServerSideTeardown(t *testing.T) { + xdgIsolate(t) + t.Cleanup(workspace.SetSSEBackoffForTest(5*time.Millisecond, 25*time.Millisecond)) + + rt := newRuntimeServer(t) + rt.srv.Backend().SetDetachGrace(0) + + cwd, dataDir := t.TempDir(), t.TempDir() + // A sibling session in another directory keeps the server alive, which + // is what made the 404s permanent instead of self-healing. + sibling := rt.newClient(t, t.TempDir()) + siblingCwd := t.TempDir() + _, err := sibling.CreateWorkspace(t.Context(), proto.Workspace{ + Path: siblingCwd, DataDir: t.TempDir(), + }) + require.NoError(t, err) + + c := rt.newClient(t, cwd) + wsProto, err := c.CreateWorkspace(t.Context(), proto.Workspace{Path: cwd, DataDir: dataDir}) + require.NoError(t, err) + originalID := wsProto.ID + + // The server drops the workspace while the client still holds the + // snapshot naming it. That is what an upgrade, or a teardown racing a + // stream drop, leaves behind: a live server that answers 404 for the + // only workspace ID this client knows. + require.NoError(t, c.DeleteWorkspace(t.Context(), originalID)) + _, err = c.GetWorkspace(t.Context(), originalID) + require.ErrorIs(t, err, client.ErrNotFound, "the workspace must really be gone") + + ws := workspace.NewClientWorkspace(c, *wsProto) + done := make(chan struct{}) + go func() { + ws.RunSubscriptionForTest(func(tea.Msg) {}) + close(done) + }() + + require.Eventually(t, func() bool { + id := ws.WorkspaceIDForTest() + return id != "" && id != originalID + }, 10*time.Second, 25*time.Millisecond, + "the client must re-register instead of retrying a workspace the server forgot") + + // The recovered workspace is usable, which is what 404s prevented. + _, err = ws.ListSessions(t.Context()) + require.NoError(t, err, "the recovered workspace must serve requests again") + + ws.Shutdown() + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("the subscription loop did not stop after Shutdown") + } +} diff --git a/internal/workspace/workspace.go b/internal/workspace/workspace.go new file mode 100644 index 0000000000000000000000000000000000000000..2e4fdfb8142a43763a4141a1cd772be54d5ec3cb --- /dev/null +++ b/internal/workspace/workspace.go @@ -0,0 +1,245 @@ +// Package workspace defines the Workspace interface used by all +// frontends (TUI, CLI) to interact with a running workspace. Two +// implementations exist: one wrapping a local app.App instance and one +// wrapping the HTTP client SDK. +package workspace + +import ( + "context" + "errors" + "time" + + tea "charm.land/bubbletea/v2" + "charm.land/catwalk/pkg/catwalk" + mcptools "github.com/charmbracelet/crush/internal/agent/tools/mcp" + "github.com/charmbracelet/crush/internal/commands" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/history" + "github.com/charmbracelet/crush/internal/lsp" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/oauth" + "github.com/charmbracelet/crush/internal/permission" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/question" + "github.com/charmbracelet/crush/internal/session" + "github.com/charmbracelet/crush/internal/skills" +) + +// Reasons the coder agent may be unavailable, returned by +// Workspace.AgentReadyErr so callers can tell a genuinely +// uninitialized agent apart from a lost server connection. +var ( + // ErrAgentNotInitialized means the workspace exists but its coder + // agent has not been configured/initialized (e.g. no model set). + ErrAgentNotInitialized = errors.New("coder agent is not initialized") + // ErrServerUnreachable means the client could not reach the server + // to determine the agent's status (server down, or the workspace was + // torn down out from under the client). + ErrServerUnreachable = errors.New("lost connection to the crush server") + // ErrWorkspaceGone means the server is reachable but no longer knows + // this client's workspace: it was torn down, or the server was + // replaced underneath the client. The subscription loop re-registers + // the workspace in the background when it sees this. + ErrWorkspaceGone = errors.New("the server reset this workspace; reconnecting") + // ErrStreamClosed means an established event stream ended. + // Resubscribing usually succeeds immediately, but events published in + // the meantime are lost for good, so the client treats it as a + // degraded link that requires a resync. + ErrStreamClosed = errors.New("the event stream closed; reconnecting") +) + +// ConnectionState describes the health of the client-server link as +// reported by the [ClientWorkspace] subscription loop. +type ConnectionState int + +const ( + // ConnectionDegraded means the event stream is down (or the workspace + // was lost server-side) and the client is retrying or re-registering + // in the background. + ConnectionDegraded ConnectionState = iota + // ConnectionRecovered means the event stream was re-established, + // possibly against a re-created workspace. + ConnectionRecovered +) + +// ConnectionEvent is delivered to the TUI as a tea.Msg on degraded and +// recovered transitions of the client-server link. Local (in-process) +// workspaces never emit it. +type ConnectionEvent struct { + State ConnectionState + // Err is the most recent failure, set when State is + // ConnectionDegraded. + Err error + // Stuck marks a degraded connection that has resisted repeated + // recovery attempts. The loop keeps retrying regardless; the UI + // should escalate from a transient notice to a persistent error. + Stuck bool +} + +// LSPClientInfo holds information about an LSP client's state. This is +// the frontend-facing type; implementations translate from the +// underlying app or proto representation. +type LSPClientInfo struct { + Name string + State lsp.ServerState + Error error + DiagnosticCount int + ConnectedAt time.Time +} + +// LSPEventType represents the type of LSP event. +type LSPEventType string + +const ( + LSPEventStateChanged LSPEventType = "state_changed" + LSPEventDiagnosticsChanged LSPEventType = "diagnostics_changed" +) + +// LSPEvent represents an LSP event forwarded to the TUI. +type LSPEvent struct { + Type LSPEventType + Name string + State lsp.ServerState + Error error + DiagnosticCount int +} + +// AgentModel holds the model information exposed to the UI. +type AgentModel struct { + CatwalkCfg catwalk.Model + ModelCfg config.SelectedModel +} + +// Workspace is the main abstraction consumed by the TUI and CLI. It +// groups every operation a frontend needs to perform against a running +// workspace, regardless of whether the workspace is in-process or +// remote. +type Workspace interface { + // Sessions + CreateSession(ctx context.Context, title string) (session.Session, error) + GetSession(ctx context.Context, sessionID string) (session.Session, error) + ListSessions(ctx context.Context) ([]session.Session, error) + SaveSession(ctx context.Context, sess session.Session) (session.Session, error) + DeleteSession(ctx context.Context, sessionID string) error + CreateAgentToolSessionID(messageID, toolCallID string) string + ParseAgentToolSessionID(sessionID string) (messageID string, toolCallID string, ok bool) + // SetCurrentSession reports the session this client is currently + // viewing. Empty sessionID clears the entry (e.g. landing screen). + // In single-client local mode this is a no-op. In client/server + // mode it informs the server's per-client presence map so other + // observers can compute attached-client counts per session. + SetCurrentSession(ctx context.Context, sessionID string) error + + // Messages + ListMessages(ctx context.Context, sessionID string) ([]message.Message, error) + ListUserMessages(ctx context.Context, sessionID string) ([]message.Message, error) + ListAllUserMessages(ctx context.Context) ([]message.Message, error) + + // Agent + AgentRun(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) error + AgentRunShellCommand(ctx context.Context, sessionID, command string, termWidth int, onProgress func(string), isFirstMessage bool) (proto.ShellCommandResponse, error) + AgentCancel(sessionID string) + AgentIsBusy() bool + AgentIsSessionBusy(sessionID string) bool + AgentModel() AgentModel + AgentIsReady() bool + // AgentReadyErr reports nil when the coder agent is ready to accept + // work, or a descriptive error otherwise: ErrAgentNotInitialized + // when the agent simply isn't set up, or ErrServerUnreachable + // (wrapped) when the client could not reach the server to find out. + // It lets the UI show an actionable message instead of collapsing + // both cases into "agent offline". + AgentReadyErr() error + AgentQueuedPrompts(sessionID string) int + AgentQueuedPromptsList(sessionID string) []string + AgentClearQueue(sessionID string) + AgentSummarize(ctx context.Context, sessionID string) error + UpdateAgentModel(ctx context.Context) error + InitCoderAgent(ctx context.Context) error + InitCoderAgentNonInteractive(ctx context.Context) error + GetDefaultSmallModel(providerID string) config.SelectedModel + + // Permissions + // + // PermissionGrant, PermissionGrantPersistent, and PermissionDeny + // return true if the call resolved the pending request and false if + // it had already been resolved by another subscriber (or is no + // longer pending). A false return is not an error; the modal can + // still close locally because the resolution will arrive via the + // PermissionNotification event stream regardless of which client + // won the race. + PermissionGrant(perm permission.PermissionRequest) bool + PermissionGrantPersistent(perm permission.PermissionRequest) bool + PermissionDeny(perm permission.PermissionRequest) bool + PermissionSkipRequests() bool + PermissionSetSkipRequests(skip bool) + + // Questions + // + // QuestionAnswer resolves the pending question with responses. + QuestionAnswer(responses []question.Answer) bool + + // QuestionCancel cancels the pending question. + QuestionCancel() bool + + // FileTracker + FileTrackerRecordRead(ctx context.Context, sessionID, path string) + FileTrackerLastReadTime(ctx context.Context, sessionID, path string) time.Time + FileTrackerListReadFiles(ctx context.Context, sessionID string) ([]string, error) + + // History + ListSessionHistory(ctx context.Context, sessionID string) ([]history.File, error) + + // LSP + LSPStart(ctx context.Context, path string) + LSPStopAll(ctx context.Context) + LSPGetStates() map[string]LSPClientInfo + LSPGetDiagnosticCounts(name string) lsp.DiagnosticCounts + + // Config (read-only data) + Config() *config.Config + WorkingDir() string + Resolver() config.VariableResolver + + // Config mutations (proxied to server in client mode) + UpdatePreferredModel(scope config.Scope, modelType config.SelectedModelType, model config.SelectedModel) error + SetCompactMode(scope config.Scope, enabled bool) error + SetProviderAPIKey(scope config.Scope, providerID string, apiKey any) error + SetConfigField(scope config.Scope, key string, value any) error + RemoveConfigField(scope config.Scope, key string) error + ImportCopilot() (*oauth.Token, bool) + RefreshOAuthToken(ctx context.Context, scope config.Scope, providerID string) error + + // Project lifecycle + ProjectNeedsInitialization() (bool, error) + MarkProjectInitialized() error + InitializePrompt() (string, error) + ListSkills(ctx context.Context) ([]skills.CatalogEntry, error) + ReadSkill(ctx context.Context, skillID string) ([]byte, skills.SkillReadResult, error) + + // MCP operations (server-side in client mode) + MCPGetStates() map[string]mcptools.ClientInfo + MCPRefreshPrompts(ctx context.Context, name string) + MCPRefreshResources(ctx context.Context, name string) + RefreshMCPTools(ctx context.Context, name string) + ReadMCPResource(ctx context.Context, name, uri string) ([]MCPResourceContents, error) + ListMCPPrompts(ctx context.Context) ([]commands.MCPPrompt, error) + GetMCPPrompt(clientID, promptID string, args map[string]string) (string, error) + EnableDockerMCP(ctx context.Context) error + DisableDockerMCP() error + MCPAuthenticate(ctx context.Context, name string) error + MCPPendingAuth() []mcptools.PendingAuthServer + MCPAuthURL(name string) string + + // Events + Subscribe(program *tea.Program) + Shutdown() +} + +// MCPResourceContents holds the contents of an MCP resource. +type MCPResourceContents struct { + URI string `json:"uri"` + MIMEType string `json:"mime_type,omitempty"` + Text string `json:"text,omitempty"` + Blob []byte `json:"blob,omitempty"` +}