diff --git a/internal/agent/.env.sample b/internal/agent/.env.sample new file mode 100644 index 0000000000000000000000000000000000000000..f32266e36f443592433a076b1910d588b5f335d3 --- /dev/null +++ b/internal/agent/.env.sample @@ -0,0 +1,5 @@ +CRUSH_ANTHROPIC_API_KEY= +CRUSH_HYPER_API_KEY= +CRUSH_OPENAI_API_KEY= +CRUSH_OPENROUTER_API_KEY= +CRUSH_ZAI_API_KEY= diff --git a/internal/agent/accepted_run_test.go b/internal/agent/accepted_run_test.go new file mode 100644 index 0000000000000000000000000000000000000000..14aec44265d44fa0f5b055ef14af3086e13a0cf3 --- /dev/null +++ b/internal/agent/accepted_run_test.go @@ -0,0 +1,231 @@ +package agent + +import ( + "context" + "testing" + + "github.com/charmbracelet/crush/internal/message" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newCancelTestAgent builds a DB-backed sessionAgent with no model. The +// tests here exercise the dispatch/cancel/persist paths, none of which +// reach agent.Stream, so a model is unnecessary. +func newCancelTestAgent(t *testing.T) (*sessionAgent, fakeEnv) { + t.Helper() + env := testEnv(t) + sa := NewSessionAgent(SessionAgentOptions{ + Sessions: env.sessions, + Messages: env.messages, + }).(*sessionAgent) + return sa, env +} + +func (a *sessionAgent) acceptedCount(sessionID string) int { + c, _ := a.acceptedRuns.Get(sessionID) + return c +} + +func (a *sessionAgent) hasPendingCancel(sessionID string) bool { + mark, ok := a.cancelMark.Get(sessionID) + return ok && mark > 0 +} + +func (a *sessionAgent) pendingCancelMark(sessionID string) uint64 { + mark, _ := a.cancelMark.Get(sessionID) + return mark +} + +func TestAcceptedRun_CloseIsIdempotent(t *testing.T) { + t.Parallel() + sa, _ := newCancelTestAgent(t) + + accept := sa.BeginAccepted("sid") + require.Equal(t, "sid", accept.SessionID()) + require.Equal(t, 1, sa.acceptedCount("sid")) + + accept.Close() + require.Equal(t, 0, sa.acceptedCount("sid")) + + // Repeated Close must not underflow the counter. + accept.Close() + accept.Close() + require.Equal(t, 0, sa.acceptedCount("sid")) +} + +func TestAcceptedRun_MultipleReservations(t *testing.T) { + t.Parallel() + sa, _ := newCancelTestAgent(t) + + a1 := sa.BeginAccepted("sid") + a2 := sa.BeginAccepted("sid") + require.Equal(t, 2, sa.acceptedCount("sid")) + + a1.Close() + require.Equal(t, 1, sa.acceptedCount("sid")) + + a2.Close() + require.Equal(t, 0, sa.acceptedCount("sid")) +} + +func TestAcceptedRun_NilSafe(t *testing.T) { + t.Parallel() + var accept *AcceptedRun + require.Equal(t, "", accept.SessionID()) + // Must not panic. + accept.Close() +} + +func TestCancel_IdleDoesNotRecordPendingCancel(t *testing.T) { + t.Parallel() + sa, _ := newCancelTestAgent(t) + + // No accepted run, no active request: a true no-op. + sa.Cancel("sid") + require.False(t, sa.hasPendingCancel("sid")) +} + +func TestCancel_AcceptedRecordsPendingCancel(t *testing.T) { + t.Parallel() + sa, _ := newCancelTestAgent(t) + + accept := sa.BeginAccepted("sid") + defer accept.Close() + + sa.Cancel("sid") + require.True(t, sa.hasPendingCancel("sid")) +} + +func TestCancel_SecondCancelWhilePendingIsNoOp(t *testing.T) { + t.Parallel() + sa, _ := newCancelTestAgent(t) + + accept := sa.BeginAccepted("sid") + defer accept.Close() + + sa.Cancel("sid") + require.True(t, sa.hasPendingCancel("sid")) + + // A second cancel while a pending cancel is already recorded must + // remain a single pending cancel; one Run consumes exactly one. + sa.Cancel("sid") + require.True(t, sa.hasPendingCancel("sid")) +} + +func TestRun_CancelOnEntryPersistsCanceledTurn(t *testing.T) { + t.Parallel() + sa, env := newCancelTestAgent(t) + + sess, err := env.sessions.Create(t.Context(), "session") + require.NoError(t, err) + + accept := sa.BeginAccepted(sess.ID) + // A cancel arrives in the accepted-but-not-yet-active window. + sa.Cancel(sess.ID) + require.True(t, sa.hasPendingCancel(sess.ID)) + + result, err := sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "hello", + Accepted: accept, + }) + require.NoError(t, err) + require.Nil(t, result) + + // The pending cancel was consumed and the accept released. + require.False(t, sa.hasPendingCancel(sess.ID)) + require.Equal(t, 0, sa.acceptedCount(sess.ID)) + + msgs, err := env.messages.List(t.Context(), sess.ID) + require.NoError(t, err) + require.Len(t, msgs, 2) + assert.Equal(t, message.User, msgs[0].Role) + assert.Equal(t, message.Assistant, msgs[1].Role) + assert.Equal(t, message.FinishReasonCanceled, msgs[1].FinishReason()) +} + +func TestPersistCanceledTurn_WritesBothWhenUserMissing(t *testing.T) { + t.Parallel() + sa, env := newCancelTestAgent(t) + + sess, err := env.sessions.Create(t.Context(), "session") + require.NoError(t, err) + + err = sa.persistCanceledTurn(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "hello", + }, false) + require.NoError(t, err) + + msgs, err := env.messages.List(t.Context(), sess.ID) + require.NoError(t, err) + require.Len(t, msgs, 2) + assert.Equal(t, message.User, msgs[0].Role) + assert.Equal(t, message.Assistant, msgs[1].Role) + assert.Equal(t, message.FinishReasonCanceled, msgs[1].FinishReason()) +} + +func TestPersistCanceledTurn_WritesAssistantOnlyWhenUserCreated(t *testing.T) { + t.Parallel() + sa, env := newCancelTestAgent(t) + + sess, err := env.sessions.Create(t.Context(), "session") + require.NoError(t, err) + + // Simulate PrepareStep having already created the user message. + _, err = sa.createUserMessage(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "hello", + }) + require.NoError(t, err) + + err = sa.persistCanceledTurn(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "hello", + }, true) + require.NoError(t, err) + + msgs, err := env.messages.List(t.Context(), sess.ID) + require.NoError(t, err) + require.Len(t, msgs, 2) + assert.Equal(t, message.User, msgs[0].Role) + assert.Equal(t, message.Assistant, msgs[1].Role) + assert.Equal(t, message.FinishReasonCanceled, msgs[1].FinishReason()) +} + +func TestPersistCanceledTurn_SucceedsWithCanceledContext(t *testing.T) { + t.Parallel() + sa, env := newCancelTestAgent(t) + + sess, err := env.sessions.Create(t.Context(), "session") + require.NoError(t, err) + + // Simulate workspace shutdown having already canceled the run + // context. WithoutCancel must let the writes through. + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + err = sa.persistCanceledTurn(ctx, SessionAgentCall{ + SessionID: sess.ID, + Prompt: "hello", + }, false) + require.NoError(t, err) + + msgs, err := env.messages.List(t.Context(), sess.ID) + require.NoError(t, err) + require.Len(t, msgs, 2) +} + +func TestClearPendingCancel(t *testing.T) { + t.Parallel() + sa, _ := newCancelTestAgent(t) + + accept := sa.BeginAccepted("sid") + defer accept.Close() + sa.Cancel("sid") + require.True(t, sa.hasPendingCancel("sid")) + + sa.clearPendingCancel("sid") + require.False(t, sa.hasPendingCancel("sid")) +} diff --git a/internal/agent/agent.go b/internal/agent/agent.go new file mode 100644 index 0000000000000000000000000000000000000000..92d289bd6fc3ea3f07959fbe7ec1b1fae6ed5872 --- /dev/null +++ b/internal/agent/agent.go @@ -0,0 +1,2331 @@ +// Package agent is the core orchestration layer for Crush AI agents. +// +// It provides session-based AI agent functionality for managing +// conversations, tool execution, and message handling. It coordinates +// interactions between language models, messages, sessions, and tools while +// handling features like automatic summarization, queuing, and token +// management. +package agent + +import ( + "cmp" + "context" + _ "embed" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "log/slog" + "math" + "net/http" + "os" + "regexp" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" + + "charm.land/catwalk/pkg/catwalk" + "charm.land/fantasy" + "charm.land/fantasy/providers/anthropic" + "charm.land/fantasy/providers/bedrock" + "charm.land/fantasy/providers/google" + "charm.land/fantasy/providers/openai" + "charm.land/fantasy/providers/openrouter" + "charm.land/fantasy/providers/vercel" + "charm.land/lipgloss/v2" + "github.com/charmbracelet/crush/internal/agent/hyper" + "github.com/charmbracelet/crush/internal/agent/notify" + "github.com/charmbracelet/crush/internal/agent/tools" + "github.com/charmbracelet/crush/internal/agent/tools/mcp" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/csync" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/charmbracelet/crush/internal/session" + "github.com/charmbracelet/crush/internal/stringext" + "github.com/charmbracelet/crush/internal/version" + "github.com/charmbracelet/x/ansi" + "github.com/charmbracelet/x/exp/charmtone" +) + +const ( + DefaultSessionName = "Untitled Session" + + // Constants for auto-summarization thresholds + largeContextWindowThreshold = 200_000 + largeContextWindowBuffer = 20_000 + smallContextWindowRatio = 0.2 +) + +var userAgent = fmt.Sprintf("Charm-Crush/%s (https://charm.land/crush)", version.Version) + +//go:embed templates/title.md +var titlePrompt []byte + +//go:embed templates/summary.md +var summaryPrompt []byte + +// Used to remove tags from generated titles. +var ( + thinkTagRegex = regexp.MustCompile(`(?s).*?`) + orphanThinkTagRegex = regexp.MustCompile(``) +) + +type SessionAgentCall struct { + SessionID string + // RunID, when non-empty, is the caller-supplied correlator that + // gets echoed back on the notify.RunComplete event emitted for + // this turn. It is preserved when the call is enqueued behind a + // busy session so the queued turn's terminal event is still + // recognisable to the original caller. Callers that need a + // reliable completion contract (e.g. `crush run` against a + // session that may be busy) MUST set it; SessionID alone is + // ambiguous when concurrent turns share the same session. + RunID string + Prompt string + ProviderOptions fantasy.ProviderOptions + Attachments []message.Attachment + MaxOutputTokens int64 + Temperature *float64 + TopP *float64 + TopK *int64 + FrequencyPenalty *float64 + PresencePenalty *float64 + NonInteractive bool + // OnComplete, when non-nil, replaces the default RunComplete + // publish path: the inner Run hands the terminal payload to this + // callback instead of emitting it on the RunComplete broker. The + // coordinator uses this hook to coalesce the unauthorized → + // re-auth → retry chain into a single user-visible terminal + // event, so non-interactive clients (e.g. `crush run`) don't + // exit on a stale failed-attempt RunComplete before the + // successful retry. It is intentionally stripped when queueing + // a busy-session call (see Run): the originating + // coordinator.Run has long returned by the time the queued + // recursion drains, so falling back to the default broker + // publish keeps the event visible to subscribers. + OnComplete func(notify.RunComplete) + // Accepted, when non-nil, is the accept reservation taken by + // BeginAccepted before the call was dispatched onto a goroutine + // (the client/server fire-and-forget path). Run consumes it under + // dispatchMu[SessionID] once the accepted -> (cancel-on-entry | + // queued | active) transition has been chosen. When nil + // (in-process / local callers like AppWorkspace), behavior is + // unchanged and no accept tracking applies. + Accepted *AcceptedRun + // acceptSeq carries the accept sequence of the handle that produced + // this call after it has been enqueued and its Accepted handle + // stripped. The queue-drain paths compare it against a session's + // cancel mark so a follow-up queued before a cancel is dropped while + // one queued after the cancel survives. 0 means untracked (an + // in-process enqueue with no accept reservation), which the drain + // paths treat as covered by any present mark, preserving the + // pre-sequence behavior. + acceptSeq uint64 + // OnAuthRefresh, when non-nil, is called by fantasy when a stream + // fails with an authentication error (HTTP 401). The callback should + // refresh credentials and return nil on success, in which case + // fantasy retries the stream transparently. Returning an error + // surfaces the original auth error without retry. + OnAuthRefresh func(ctx context.Context, err *fantasy.ProviderError) error +} + +type SessionAgent interface { + Run(context.Context, SessionAgentCall) (*fantasy.AgentResult, error) + BeginAccepted(sessionID string) *AcceptedRun + SetModels(large Model, small Model) + SetTools(tools []fantasy.AgentTool) + SetSystemPrompt(systemPrompt string) + Cancel(sessionID string) + CancelAll() + IsSessionBusy(sessionID string) bool + IsBusy() bool + QueuedPrompts(sessionID string) int + QueuedPromptsList(sessionID string) []string + ClearQueue(sessionID string) + Summarize(context.Context, string, fantasy.ProviderOptions, func(context.Context, *fantasy.ProviderError) error) error + Model() Model + GenerateTitle(ctx context.Context, sessionID, userPrompt string) +} + +type Model struct { + Model fantasy.LanguageModel + CatwalkCfg catwalk.Model + ModelCfg config.SelectedModel + FlatRate bool +} + +// activeCancel wraps a context.CancelFunc with a unique pointer identity. +// The pointer is used for compare-and-delete in the dispatch completion path: +// when a finishing run's deferred cleanup fires, it must only remove its own +// entry — not a newer run's entry that was installed in the window between +// the explicit Del and the function return. +type activeCancel struct { + cancel context.CancelFunc +} + +type sessionAgent struct { + largeModel *csync.Value[Model] + smallModel *csync.Value[Model] + systemPromptPrefix *csync.Value[string] + systemPrompt *csync.Value[string] + tools *csync.Slice[fantasy.AgentTool] + + isSubAgent bool + sessions session.Service + messages message.Service + disableAutoSummarize bool + isYolo bool + notify pubsub.Publisher[notify.Notification] + runComplete pubsub.Publisher[notify.RunComplete] + + messageQueue *csync.Map[string, []SessionAgentCall] + activeRequests *csync.Map[string, *activeCancel] + + // dispatchMu holds a per-session mutex that serializes the + // accepted -> (cancel-on-entry | queued | active) transition in + // Run against a concurrent Cancel. The lock is held only during + // the brief handoff (no DB or LLM I/O under the lock). + dispatchMu *csync.Map[string, *sync.Mutex] + // acceptedRuns counts dispatched-but-not-yet-active runs per + // session. A counter > 0 means a dispatched prompt is in flight + // and has not yet completed the dispatch handoff in Run. Only + // BeginAccepted increments it; only AcceptedRun.Close decrements + // it. + acceptedRuns *csync.Map[string, int] + // cancelMark records, per session, a high-water accept sequence: an + // accepted handle is canceled by it iff the handle's sequence is at + // or below the mark. Cancel raises the mark to the latest sequence + // assigned at cancel time, so a single Cancel covers every prompt + // accepted-but-not-yet-active then, while a prompt accepted later + // (higher sequence) is never poisoned. Absent or 0 means no pending + // cancel. It is only raised by Cancel when acceptedRuns > 0, so an + // idle Escape never records a mark. + cancelMark *csync.Map[string, uint64] + // dispatchMuCreate guards lazy creation of per-session entries in + // dispatchMu so two goroutines can't race to lock different mutex + // instances for the same session. + dispatchMuCreate sync.Mutex + // acceptedMu serializes increments/decrements of acceptedRuns and + // the assignment of accept sequence numbers from acceptSeqGen. It + // is separate from dispatchMu so AcceptedRun.Close (which may run + // while Run holds dispatchMu for the same session) does not + // deadlock by re-entering the dispatch lock. + acceptedMu sync.Mutex + // acceptSeqGen is the monotonic source of accept sequence numbers. + // Each BeginAccepted increments it under acceptedMu and stamps the + // returned handle, so sequences strictly increase in accept order + // across the agent. Cancel uses its current value as the per-session + // high-water mark. + acceptSeqGen uint64 +} + +type SessionAgentOptions struct { + LargeModel Model + SmallModel Model + SystemPromptPrefix string + SystemPrompt string + IsSubAgent bool + DisableAutoSummarize bool + IsYolo bool + Sessions session.Service + Messages message.Service + Tools []fantasy.AgentTool + Notify pubsub.Publisher[notify.Notification] + RunComplete pubsub.Publisher[notify.RunComplete] +} + +func NewSessionAgent( + opts SessionAgentOptions, +) SessionAgent { + return &sessionAgent{ + largeModel: csync.NewValue(opts.LargeModel), + smallModel: csync.NewValue(opts.SmallModel), + systemPromptPrefix: csync.NewValue(opts.SystemPromptPrefix), + systemPrompt: csync.NewValue(opts.SystemPrompt), + isSubAgent: opts.IsSubAgent, + sessions: opts.Sessions, + messages: opts.Messages, + disableAutoSummarize: opts.DisableAutoSummarize, + tools: csync.NewSliceFrom(opts.Tools), + isYolo: opts.IsYolo, + notify: opts.Notify, + runComplete: opts.RunComplete, + messageQueue: csync.NewMap[string, []SessionAgentCall](), + activeRequests: csync.NewMap[string, *activeCancel](), + dispatchMu: csync.NewMap[string, *sync.Mutex](), + acceptedRuns: csync.NewMap[string, int](), + cancelMark: csync.NewMap[string, uint64](), + } +} + +// AcceptedRun owns exactly one accept reservation taken by +// BeginAccepted. It is the only carrier of accept-state across the +// backend.runAgent / Coordinator.Run / sessionAgent.Run layers: a +// counter > 0 means a dispatched prompt is in flight and has not yet +// completed the dispatch handoff in Run. Close is the only way to +// release the reservation and is idempotent. +type AcceptedRun struct { + agent *sessionAgent + sessionID string + // seq is the monotonic accept sequence stamped by BeginAccepted. A + // cancel covers this handle iff seq is at or below the session's + // cancel mark, so a handle accepted after a cancel (higher seq) is + // never poisoned by it. + seq uint64 + done atomic.Bool +} + +// Close decrements the accept counter for this reservation. It is safe +// to call multiple times; only the first call has effect. +func (r *AcceptedRun) Close() { + if r == nil { + return + } + if !r.done.CompareAndSwap(false, true) { + return + } + r.agent.endAccepted(r.sessionID) +} + +// SessionID exposes the session this reservation is for so the run path +// can use it without an extra parameter. +func (r *AcceptedRun) SessionID() string { + if r == nil { + return "" + } + return r.sessionID +} + +// BeginAccepted increments the accept counter for sessionID and returns +// a handle whose Close is the only way to decrement it. It is the only +// entry point that mutates acceptedRuns. +func (a *sessionAgent) BeginAccepted(sessionID string) *AcceptedRun { + a.acceptedMu.Lock() + defer a.acceptedMu.Unlock() + count, _ := a.acceptedRuns.Get(sessionID) + a.acceptedRuns.Set(sessionID, count+1) + a.acceptSeqGen++ + return &AcceptedRun{agent: a, sessionID: sessionID, seq: a.acceptSeqGen} +} + +// endAccepted decrements the accept counter for sessionID. It is only +// called via AcceptedRun.Close. It uses a dedicated lock (not the +// per-session dispatch mutex) so it can run while Run holds dispatchMu +// for the same session without deadlocking. +// +// When the count reaches zero the session's cancel mark is dropped: no +// accepted handle remains for it to cover, and any handle accepted later +// gets a strictly higher sequence that the mark would not match anyway. +// Handles canceled on entry never reach RunComplete, so this is the only +// place that clears the mark for an all-canceled batch. Sibling handles +// covered by the same mark are serialized on the per-session dispatch +// mutex and read the mark before they Close, so this never clears it out +// from under a covered handle still waiting to enter Run. +func (a *sessionAgent) endAccepted(sessionID string) { + a.acceptedMu.Lock() + defer a.acceptedMu.Unlock() + count, ok := a.acceptedRuns.Get(sessionID) + if !ok || count <= 1 { + a.acceptedRuns.Del(sessionID) + a.cancelMark.Del(sessionID) + return + } + a.acceptedRuns.Set(sessionID, count-1) +} + +// sessionMu returns the per-session dispatch mutex, creating it on first +// use. Creation is guarded so concurrent callers always observe the same +// mutex instance for a given session. +func (a *sessionAgent) sessionMu(sessionID string) *sync.Mutex { + if mu, ok := a.dispatchMu.Get(sessionID); ok { + return mu + } + a.dispatchMuCreate.Lock() + defer a.dispatchMuCreate.Unlock() + if mu, ok := a.dispatchMu.Get(sessionID); ok { + return mu + } + mu := &sync.Mutex{} + a.dispatchMu.Set(sessionID, mu) + return mu +} + +// enqueueCall appends call to the session's message queue. The +// OnComplete hook is stripped: the caller that supplied it (typically +// coordinator.Run) has its own retry/coalesce scope that ends when it +// returns, so by the time the queue drains nobody is left to consume the +// buffered terminal event. The recursive Run falls back to the default +// broker publish, which is what existing subscribers expect for queued +// turns. +func (a *sessionAgent) enqueueCall(call SessionAgentCall) { + existing, ok := a.messageQueue.Get(call.SessionID) + if !ok { + existing = []SessionAgentCall{} + } + queued := call + if call.Accepted != nil { + // Preserve the accept sequence after the handle is stripped so + // the queue-drain paths can tell a follow-up queued before a + // cancel (covered by the mark) from one queued after it. + queued.acceptSeq = call.Accepted.seq + } + queued.OnComplete = nil + queued.Accepted = nil + existing = append(existing, queued) + a.messageQueue.Set(call.SessionID, existing) +} + +// drainQueueForStep partitions the session's queued calls for the current +// streaming step under the per-session dispatch mutex so the filtering is +// atomic against a concurrent Cancel: canceledBySeq requires the caller to +// hold that mutex, and evaluating it here (rather than after unlocking) +// prevents a cancel recorded between the drain and the check from being +// observed inconsistently. +// +// Calls covered by a pending cancel are dropped; the dropped ones that +// carry a RunID are returned in canceledWithRunID so the caller can +// publish their terminal cancelled RunComplete (a caller waiting on that +// RunID, e.g. `crush run`, would otherwise hang). Uncanceled calls without +// a RunID are returned in fold to be folded into the active turn, +// preserving the existing follow-up behavior. Uncanceled calls that carry +// a RunID are left in the queue so each runs as its own turn via the +// recursive run path and publishes its own RunComplete, giving every +// RunID-bearing prompt an explicit lifecycle instead of being silently +// absorbed into another turn. fold is processed by the caller without the +// lock held. +func (a *sessionAgent) drainQueueForStep(sessionID string) (fold, canceledWithRunID []SessionAgentCall) { + dispatchLock := a.sessionMu(sessionID) + dispatchLock.Lock() + defer dispatchLock.Unlock() + queuedCalls, _ := a.messageQueue.Get(sessionID) + var keep []SessionAgentCall + for _, queued := range queuedCalls { + if a.canceledBySeq(sessionID, queued.acceptSeq) { + if queued.RunID != "" { + canceledWithRunID = append(canceledWithRunID, queued) + } + continue + } + if queued.RunID != "" { + keep = append(keep, queued) + continue + } + fold = append(fold, queued) + } + if len(keep) == 0 { + a.messageQueue.Del(sessionID) + } else { + a.messageQueue.Set(sessionID, keep) + } + return fold, canceledWithRunID +} + +// publishCanceledQueueDrops emits a terminal cancelled RunComplete for +// every dropped queued call that carries a RunID. A queued prompt removed +// from the queue without ever running — covered by a pending cancel, or +// cleared by Cancel/ClearQueue — would otherwise leave a caller blocked on +// that RunID: `crush run` ignores live message events and exits only on a +// RunComplete whose RunID matches. Calls without a RunID had no such waiter +// and are dropped silently as before. A detached, bounded context keeps the +// must-deliver publish alive even when the run context that triggered the +// drop is already canceled. +func (a *sessionAgent) publishCanceledQueueDrops(drops []SessionAgentCall) { + var hasRunID bool + for _, d := range drops { + if d.RunID != "" { + hasRunID = true + break + } + } + if !hasRunID { + return + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + for _, d := range drops { + if d.RunID == "" { + continue + } + a.publishRunComplete(ctx, d, notify.RunComplete{ + SessionID: d.SessionID, + RunID: d.RunID, + Cancelled: true, + }) + } +} + +// clearQueueAndNotify removes all queued prompts for the session and +// publishes a terminal cancelled RunComplete for any that carried a RunID, +// so callers waiting on those RunIDs (e.g. `crush run`) are not left +// hanging when their queued prompt is discarded without running. +func (a *sessionAgent) clearQueueAndNotify(sessionID string) { + queued, ok := a.messageQueue.Get(sessionID) + a.messageQueue.Del(sessionID) + if !ok { + return + } + a.publishCanceledQueueDrops(queued) +} + +// clearPendingCancel removes any pending-cancel mark for sessionID. It +// takes the per-session dispatch lock so it is ordered against Cancel +// and the dispatch handoff. +func (a *sessionAgent) clearPendingCancel(sessionID string) { + mu := a.sessionMu(sessionID) + mu.Lock() + defer mu.Unlock() + a.cancelMark.Del(sessionID) +} + +// canceledBySeq reports whether an accepted handle or queued call with +// the given accept sequence is covered by a pending cancel for the +// session. Callers must hold the session's dispatch mutex. A tracked +// sequence (seq > 0) is covered only when it is at or below the cancel +// high-water mark, so a prompt accepted after the cancel (higher seq) is +// never poisoned. An untracked sequence (seq == 0, an in-process enqueue +// with no accept reservation) is covered whenever any mark is present, +// preserving the pre-sequence behavior. The mark is not consumed: it +// stays so every sibling handle it covers observes the same cancel, and +// a later handle (higher seq) ignores it regardless. +func (a *sessionAgent) canceledBySeq(sessionID string, seq uint64) bool { + mark, ok := a.cancelMark.Get(sessionID) + if !ok || mark == 0 { + return false + } + return seq == 0 || seq <= mark +} + +// persistCanceledTurn writes the user/assistant records for a turn that +// was canceled before (or just as) streaming would have produced them. +// It creates the user message only when it was not already created by an +// earlier createUserMessage call (userMsgCreated), then writes an +// assistant message with FinishReasonCanceled. Both writes use +// context.WithoutCancel(ctx) so workspace shutdown (which cancels the run +// context) can't drop them. +func (a *sessionAgent) persistCanceledTurn(ctx context.Context, call SessionAgentCall, userMsgCreated bool) error { + writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + if !userMsgCreated { + if _, err := a.createUserMessage(writeCtx, call); err != nil { + return err + } + } + largeModel := a.largeModel.Get() + assistant, err := a.messages.Create(writeCtx, call.SessionID, message.CreateMessageParams{ + Role: message.Assistant, + Parts: []message.ContentPart{}, + Model: largeModel.ModelCfg.Model, + Provider: largeModel.ModelCfg.Provider, + }) + if err != nil { + return err + } + assistant.AddFinish(message.FinishReasonCanceled, "User canceled request", "") + return a.messages.Update(writeCtx, assistant) +} + +// publishRunComplete emits the authoritative terminal event for a turn. +// It honors the per-call OnComplete hook when set (so the coordinator can +// coalesce retries) and otherwise falls back to the RunComplete broker. +// ctx is used only for the bounded-blocking must-deliver publish; the +// terminal payload is supplied by the caller. This is the single emit path +// shared by the streaming defer and the cancel-on-entry early return so a +// caller waiting on RunComplete (e.g. `crush run` with a RunID) always +// observes exactly one terminal event regardless of which Run branch ends +// the turn. +func (a *sessionAgent) publishRunComplete(ctx context.Context, call SessionAgentCall, complete notify.RunComplete) { + if call.OnComplete != nil { + call.OnComplete(complete) + return + } + if a.runComplete == nil { + return + } + a.runComplete.PublishMustDeliver(ctx, pubsub.UpdatedEvent, complete) +} + +// ValidateCall performs the cheap structural validation that +// sessionAgent.Run requires before a call can be dispatched: a call must +// carry either a non-empty prompt or a text attachment, and it must name a +// session. It is exported so callers that accept a run before dispatching it +// (e.g. backend.SendMessage) can apply the same checks and keep the error +// contract consistent. +func ValidateCall(call SessionAgentCall) error { + if call.Prompt == "" && !message.ContainsTextAttachment(call.Attachments) { + return ErrEmptyPrompt + } + if call.SessionID == "" { + return ErrSessionMissing + } + return nil +} + +func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (result *fantasy.AgentResult, retErr error) { + if err := ValidateCall(call); err != nil { + return nil, err + } + + // genCtx/cancel are the run context and its cancel func, created under + // the per-session dispatch mutex below so a concurrent Cancel can observe + // the activeRequests entry before the assistant message exists. + var ( + genCtx context.Context + cancel context.CancelFunc + userMsgCreated bool + ) + + // Serialize the dispatch decision (cancel-on-entry | queued | active) + // against a concurrent Cancel. Cancel takes the same per-session lock, so + // every cancel observes at least one of: a cancel mark, an activeRequests + // entry, or a messageQueue entry it then clears. Holding the lock across + // the busy check and the active registration also makes them atomic, so + // two concurrent in-process callers — a burst of channel events, or a + // channel event racing a typed prompt — cannot both pass the busy check + // and start two runs on the same session. + sessMu := a.sessionMu(call.SessionID) + sessMu.Lock() + + if call.Accepted != nil && a.canceledBySeq(call.SessionID, call.Accepted.seq) { + // Cancel-on-entry: a cancel arrived while this accepted run was + // dispatched but not yet active, and this handle's accept sequence + // is at or below the session's cancel mark. The mark is left in + // place so sibling handles it also covers observe the same cancel; + // release the accept reservation, drop the lock, and persist a + // canceled turn without entering Stream. + // + // This path returns before the streaming defer that publishes + // RunComplete is installed, so emit the terminal event explicitly. + // Without it, a caller waiting on RunComplete for this RunID (e.g. + // `crush run`, which ignores message events and blocks on + // RunComplete) would hang on an immediately-canceled accepted run. + call.Accepted.Close() + sessMu.Unlock() + complete := notify.RunComplete{ + SessionID: call.SessionID, + RunID: call.RunID, + Cancelled: true, + } + if err := a.persistCanceledTurn(ctx, call, false); err != nil { + complete.Error = err.Error() + a.publishRunComplete(ctx, call, complete) + return nil, err + } + a.publishRunComplete(ctx, call, complete) + return nil, nil + } + + if a.IsSessionBusy(call.SessionID) { + // Busy: an earlier prompt is active. Queue this call so it is + // folded into (or sequenced after) the active turn, and release any + // accept reservation. A Cancel arriving after this point sees the + // active entry and clears the queue. + // + // enqueueCall strips OnComplete: the caller that supplied the hook + // (typically coordinator.Run) has its own retry/coalesce scope that + // ends when it returns, so by the time the queue drains nobody is + // left to consume the buffered terminal event. The queued turn falls + // back to the default broker publish, which is what existing + // subscribers expect. + a.enqueueCall(call) + if call.Accepted != nil { + call.Accepted.Close() + } + sessMu.Unlock() + return nil, nil + } + + // Idle: become the active run. Register the cancel func before dropping + // the lock so a Cancel that arrives between here and assistant creation + // is not lost. + runCtx := context.WithValue(ctx, tools.SessionIDContextKey, call.SessionID) + genCtx, cancel = context.WithCancel(runCtx) + ac := &activeCancel{cancel: cancel} + a.activeRequests.Set(call.SessionID, ac) + if call.Accepted != nil { + call.Accepted.Close() + } + sessMu.Unlock() + + defer cancel() + // Conditional cleanup: only remove our entry if it hasn't been replaced + // by a newer run. Without this guard, the deferred Del fires after a + // concurrent run registers in the completion window, silently wiping + // the new run's cancel and breaking cancellation. + defer a.activeRequests.CompareAndDelete(call.SessionID, ac) + + // Copy mutable fields under lock to avoid races with SetTools/SetModels. + agentTools := a.tools.Copy() + largeModel := a.largeModel.Get() + systemPrompt := a.systemPrompt.Get() + promptPrefix := a.systemPromptPrefix.Get() + var instructions strings.Builder + + for _, server := range mcp.GetStates() { + if server.State != mcp.StateConnected { + continue + } + if s := server.Client.InitializeResult().Instructions; s != "" { + instructions.WriteString(s) + instructions.WriteString("\n\n") + } + } + + if s := instructions.String(); s != "" { + systemPrompt += "\n\n\n" + s + "\n" + } + + if len(agentTools) > 0 { + // Add Anthropic caching to the last tool. + agentTools[len(agentTools)-1].SetProviderOptions(a.getCacheControlOptions()) + } + + agent := fantasy.NewAgent( + largeModel.Model, + fantasy.WithSystemPrompt(systemPrompt), + fantasy.WithTools(agentTools...), + fantasy.WithUserAgent(userAgent), + ) + + sessionLock := sync.Mutex{} + currentSession, err := a.sessions.Get(ctx, call.SessionID) + if err != nil { + return nil, fmt.Errorf("failed to get session: %w", err) + } + + msgs, err := a.getSessionMessages(ctx, currentSession) + if err != nil { + return nil, fmt.Errorf("failed to get session messages: %w", err) + } + + // Generate title from the first real (non-shell) user prompt. + // can take tens of seconds. Blocking Run on it delays the + // response to the caller. Use a detached context so the title + // goroutine survives Run's cancel. + if !hasUserTextMessage(msgs) { + titleCtx := context.WithoutCancel(ctx) + go a.GenerateTitle(titleCtx, call.SessionID, call.Prompt) + } + + // Add the user message to the session. + _, err = a.createUserMessage(ctx, call) + if err != nil { + return nil, err + } + userMsgCreated = true + + // Add the session to the context. The run context (genCtx) and its + // cancel func were already created and registered under the dispatch + // mutex above for both the accepted and in-process paths. + ctx = context.WithValue(ctx, tools.SessionIDContextKey, call.SessionID) + // skipRunComplete is set just before the queued-recursion path so + // the outer Run doesn't publish a RunComplete that would race + // with — and be superseded by — the recursive call's own + // RunComplete (each queued user prompt is its own turn and + // publishes exactly one terminal event). + var skipRunComplete bool + // currentAssistant is declared here so the deferred RunComplete + // publish below can capture the pointer that PrepareStep will + // later (re)assign for each streaming step. The final assistant + // message of the turn is the value reachable through this + // pointer when the defer runs. + var currentAssistant *message.Message + // Drain any debounced message updates before returning. message.Service + // already flushes synchronously on terminal updates, but a defer here + // guarantees the contract at every Run exit (success, error, panic + // recovery upstream) without callers needing to know. + // + // After the flush completes — meaning all per-message + // Publish(UpdatedEvent) calls have fired and been buffered into + // every subscriber's channel — publish the authoritative + // RunComplete event for this turn. The flush-then-publish order + // gives well-behaved clients the best chance of seeing the final + // message event before RunComplete; the embedded Text field + // reconciles for clients that observe the events out of order + // (the pubsub broker fan-in does not serialize publishes from + // different upstream brokers). + defer func() { + // Use a context detached from the run context: workspace + // shutdown cancels ctx before this goroutine returns, but the + // buffered streaming deltas must still land before the DB is + // closed. A short timeout bounds the flush. + flushCtx, flushCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer flushCancel() + if flushErr := a.messages.FlushAll(flushCtx); flushErr != nil { + slog.Error("Failed to flush pending message updates after run", "error", flushErr) + } + if skipRunComplete { + return + } + complete := notify.RunComplete{SessionID: call.SessionID, RunID: call.RunID} + if currentAssistant != nil { + complete.MessageID = currentAssistant.ID + complete.Text = currentAssistant.Content().String() + } + if retErr != nil { + complete.Error = retErr.Error() + complete.Cancelled = errors.Is(retErr, context.Canceled) + } else if ctx.Err() != nil { + complete.Cancelled = true + } + // Prefer the per-call hook when supplied so the coordinator + // can coalesce retries (e.g. unauthorized → re-auth → retry) + // into a single user-visible terminal event. The fallback + // must-deliver publish applies bounded-blocking semantics to + // the authoritative terminal event so a momentarily-full + // subscriber channel can't silently drop it and hang + // non-interactive clients waiting on RunComplete. + a.publishRunComplete(ctx, call, complete) + }() + + history, files := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages, call.Attachments...) + + startTime := time.Now() + a.eventPromptSent(call.SessionID) + + var stepMessages []fantasy.Message + var shouldSummarize bool + sanitizedToolCalls := make(map[string]bool) + // Don't send MaxOutputTokens if 0 — some providers (e.g. LM Studio) reject it + var maxOutputTokens *int64 + if call.MaxOutputTokens > 0 { + maxOutputTokens = &call.MaxOutputTokens + } + result, err = agent.Stream(genCtx, fantasy.AgentStreamCall{ + Prompt: message.PromptWithTextAttachments(call.Prompt, call.Attachments), + Files: files, + Messages: history, + Headers: sessionHeaders(call.SessionID), + ProviderOptions: call.ProviderOptions, + MaxOutputTokens: maxOutputTokens, + TopP: call.TopP, + Temperature: call.Temperature, + PresencePenalty: call.PresencePenalty, + TopK: call.TopK, + FrequencyPenalty: call.FrequencyPenalty, + PrepareStep: func(callContext context.Context, options fantasy.PrepareStepFunctionOptions) (_ context.Context, prepared fantasy.PrepareStepResult, err error) { + prepared.Messages = options.Messages + for i := range prepared.Messages { + prepared.Messages[i].ProviderOptions = nil + } + + // Use latest tools (updated by SetTools when MCP tools change). + prepared.Tools = a.tools.Copy() + + // Drain queued follow-up prompts for this step. Calls covered + // by a cancel recorded while they sat in the queue are dropped: + // a cancel that arrived after a prompt was queued must not let + // it run as part of this step. Coverage is per-call by accept + // sequence so a follow-up queued after the cancel (higher seq) + // is not dropped. A dropped prompt carrying a RunID still gets + // its terminal cancelled RunComplete so a caller waiting on it + // does not hang. Uncanceled prompts without a RunID are folded + // into this turn; uncanceled prompts with a RunID are left + // queued so each runs as its own turn (with its own + // RunComplete) via the recursive run path below. + fold, canceledRunIDs := a.drainQueueForStep(call.SessionID) + a.publishCanceledQueueDrops(canceledRunIDs) + for _, queued := range fold { + userMessage, createErr := a.createUserMessage(callContext, queued) + if createErr != nil { + return callContext, prepared, createErr + } + prepared.Messages = append(prepared.Messages, userMessage.ToAIMessage()...) + } + + prepared.Messages = a.workaroundProviderMediaLimitations(prepared.Messages, largeModel) + + lastSystemRoleInx := 0 + systemMessageUpdated := false + for i, msg := range prepared.Messages { + // Only add cache control to the last message. + if msg.Role == fantasy.MessageRoleSystem { + lastSystemRoleInx = i + } else if !systemMessageUpdated { + prepared.Messages[lastSystemRoleInx].ProviderOptions = a.getCacheControlOptions() + systemMessageUpdated = true + } + // Than add cache control to the last 2 messages. + if i > len(prepared.Messages)-3 { + prepared.Messages[i].ProviderOptions = a.getCacheControlOptions() + } + } + + if promptPrefix != "" { + prepared.Messages = append([]fantasy.Message{fantasy.NewSystemMessage(promptPrefix)}, prepared.Messages...) + } + + sessionLock.Lock() + stepMessages = cloneFantasyMessages(prepared.Messages) + sessionLock.Unlock() + + var assistantMsg message.Message + assistantMsg, err = a.messages.Create(callContext, call.SessionID, message.CreateMessageParams{ + Role: message.Assistant, + Parts: []message.ContentPart{}, + Model: largeModel.ModelCfg.Model, + Provider: largeModel.ModelCfg.Provider, + }) + if err != nil { + return callContext, prepared, err + } + callContext = context.WithValue(callContext, tools.MessageIDContextKey, assistantMsg.ID) + callContext = context.WithValue(callContext, tools.SupportsImagesContextKey, largeModel.CatwalkCfg.SupportsImages) + callContext = context.WithValue(callContext, tools.ModelNameContextKey, largeModel.CatwalkCfg.Name) + currentAssistant = &assistantMsg + return callContext, prepared, err + }, + OnReasoningStart: func(id string, reasoning fantasy.ReasoningContent) error { + currentAssistant.AppendReasoningContent(reasoning.Text) + return a.messages.Update(genCtx, *currentAssistant) + }, + OnReasoningDelta: func(id string, text string) error { + currentAssistant.AppendReasoningContent(text) + return a.messages.Update(genCtx, *currentAssistant) + }, + OnReasoningEnd: func(id string, reasoning fantasy.ReasoningContent) error { + // handle anthropic signature + if anthropicData, ok := reasoning.ProviderMetadata[anthropic.Name]; ok { + if reasoning, ok := anthropicData.(*anthropic.ReasoningOptionMetadata); ok { + currentAssistant.AppendReasoningSignature(reasoning.Signature) + } + } + if googleData, ok := reasoning.ProviderMetadata[google.Name]; ok { + if reasoning, ok := googleData.(*google.ReasoningMetadata); ok { + currentAssistant.AppendThoughtSignature(reasoning.Signature, reasoning.ToolID) + } + } + if openaiData, ok := reasoning.ProviderMetadata[openai.Name]; ok { + if reasoning, ok := openaiData.(*openai.ResponsesReasoningMetadata); ok { + currentAssistant.SetReasoningResponsesData(reasoning) + } + } + currentAssistant.FinishThinking() + return a.messages.Update(genCtx, *currentAssistant) + }, + OnTextDelta: func(id string, text string) error { + // Strip leading newline from initial text content. This is is + // particularly important in non-interactive mode where leading + // newlines are very visible. + if len(currentAssistant.Parts) == 0 { + text = strings.TrimPrefix(text, "\n") + } + + currentAssistant.AppendContent(text) + return a.messages.Update(genCtx, *currentAssistant) + }, + OnToolInputStart: func(id string, toolName string) error { + toolCall := message.ToolCall{ + ID: id, + Name: toolName, + ProviderExecuted: false, + Finished: false, + } + currentAssistant.AddToolCall(toolCall) + // Use parent ctx instead of genCtx to ensure the update succeeds + // even if the request is canceled mid-stream + return a.messages.Update(ctx, *currentAssistant) + }, + OnRetry: func(err *fantasy.ProviderError, delay time.Duration) { + slog.Warn("Provider request failed, retrying", providerRetryLogFields(err, delay)...) + // Reset streamed content so the retried response doesn't + // concatenate with partial content from the failed attempt. + // On the final attempt (no more retries), any partial content + // stays in the message as useful context beneath the error. + currentAssistant.ResetStreamedContent() + if updateErr := a.messages.Update(genCtx, *currentAssistant); updateErr != nil { + slog.Error("Failed to reset message on retry", "error", updateErr) + } + }, + OnAuthRefresh: call.OnAuthRefresh, + ModelProvider: func() fantasy.LanguageModel { + m := a.largeModel.Get() + slog.Info("ModelProvider called", + "provider", m.ModelCfg.Provider, + "model", m.ModelCfg.Model) + return m.Model + }, + OnToolCall: func(tc fantasy.ToolCallContent) error { + input, wasSanitized := sanitizeToolInput(tc.ToolName, tc.ToolCallID, tc.Input) + if wasSanitized { + sanitizedToolCalls[tc.ToolCallID] = true + } + toolCall := message.ToolCall{ + ID: tc.ToolCallID, + Name: tc.ToolName, + Input: input, + ProviderExecuted: false, + Finished: true, + } + currentAssistant.AddToolCall(toolCall) + // Use parent ctx instead of genCtx to ensure the update succeeds + // even if the request is canceled mid-stream + return a.messages.Update(ctx, *currentAssistant) + }, + OnToolResult: func(result fantasy.ToolResultContent) error { + toolResult := a.convertToToolResult(result) + if sanitizedToolCalls[result.ToolCallID] { + toolResult.Content = "Tool call failed: arguments were not valid JSON. Please check your tool call format and try again." + toolResult.IsError = true + } + // Use parent ctx instead of genCtx to ensure the message is created + // even if the request is canceled mid-stream + _, createMsgErr := a.messages.Create(ctx, currentAssistant.SessionID, message.CreateMessageParams{ + Role: message.Tool, + Parts: []message.ContentPart{ + toolResult, + }, + }) + return createMsgErr + }, + OnStepFinish: func(stepResult fantasy.StepResult) error { + for _, w := range stepResult.Warnings { + slog.Warn("Provider warning", "type", w.Type, "message", w.Message) + } + finishReason := message.FinishReasonUnknown + switch stepResult.FinishReason { + case fantasy.FinishReasonLength: + finishReason = message.FinishReasonMaxTokens + case fantasy.FinishReasonStop: + finishReason = message.FinishReasonEndTurn + case fantasy.FinishReasonToolCalls: + finishReason = message.FinishReasonToolUse + case fantasy.FinishReasonContentFilter: + // Provider safety classifier stopped the model + // (Anthropic stop_reason=refusal, OpenAI content_filter). + // The TUI owns the display copy; we only persist the + // reason so the UI can show a REFUSED banner. + finishReason = message.FinishReasonContentFilter + slog.Warn( + "Provider content filter stopped the model", + "session_id", call.SessionID, + "finish_reason", string(stepResult.FinishReason), + ) + } + // If a tool result halted the turn (e.g. a hook halt or a + // permission denial), the step ends on FinishReasonToolCalls but + // the model will not be called again. Treat it as the end of the + // turn so the UI can render the assistant footer. + if finishReason == message.FinishReasonToolUse { + for _, tr := range stepResult.Content.ToolResults() { + if tr.StopTurn { + finishReason = message.FinishReasonEndTurn + break + } + } + } + currentAssistant.AddFinish(finishReason, "", "") + currentAssistant.PrismModelID, currentAssistant.PrismModelName = extractPrismModel(stepResult.ProviderMetadata) + currentAssistant.PrismHypercreditSavings, currentAssistant.PrismDollarSavings = extractPrismSavings(stepResult.ProviderMetadata) + sessionLock.Lock() + defer sessionLock.Unlock() + + updatedSession, getSessionErr := a.sessions.Get(ctx, call.SessionID) + if getSessionErr != nil { + return getSessionErr + } + usage, estimated := fallbackStepUsage(stepMessages, stepResult) + a.updateSessionUsage(largeModel, &updatedSession, usage, a.openrouterCost(stepResult.ProviderMetadata), estimated) + extractHyperCredits(stepResult.ProviderMetadata) + _, sessionErr := a.sessions.Save(ctx, updatedSession) + if sessionErr != nil { + return sessionErr + } + currentSession = updatedSession + return a.messages.Update(genCtx, *currentAssistant) + }, + StopWhen: []fantasy.StopCondition{ + func(_ []fantasy.StepResult) bool { + cw := int64(largeModel.CatwalkCfg.ContextWindow) + // If context window is unknown (0), skip auto-summarize + // to avoid immediately truncating custom/local models. + if cw == 0 { + return false + } + tokens := currentSession.CompletionTokens + currentSession.PromptTokens + remaining := cw - tokens + var threshold int64 + if cw > largeContextWindowThreshold { + threshold = largeContextWindowBuffer + } else { + threshold = int64(float64(cw) * smallContextWindowRatio) + } + if (remaining <= threshold) && !a.disableAutoSummarize { + shouldSummarize = true + return true + } + return false + }, + func(steps []fantasy.StepResult) bool { + return hasRepeatedToolCalls(steps, loopDetectionWindowSize, loopDetectionMaxRepeats) + }, + }, + }) + + a.eventPromptResponded(call.SessionID, time.Since(startTime).Truncate(time.Second)) + + if err != nil { + isHyper := largeModel.ModelCfg.Provider == hyper.Name + isCancelErr := errors.Is(err, context.Canceled) + slog.Info("Agent stream returned error", + "error", err.Error(), + "error_type", fmt.Sprintf("%T", err), + "is_hyper", isHyper, + "is_cancel", isCancelErr) + if currentAssistant == nil { + // Cancel-before-assistant-creation window: the run was + // canceled after activeRequests.Set but before PrepareStep + // created the assistant message. Without this, the turn + // would return with no FinishReasonCanceled marker and no + // user-visible record. The user message was already created + // above, so persistCanceledTurn only writes the assistant + // record. + if isCancelErr { + if persistErr := a.persistCanceledTurn(ctx, call, userMsgCreated); persistErr != nil { + return nil, persistErr + } + } + return result, err + } + // Persist final state with a context detached from the run + // context. The run context (ctx) is derived from the + // workspace context, which workspace shutdown cancels before + // agent goroutines finish; using ctx here would drop the + // final assistant state. WithoutCancel keeps the values + // (e.g. session ID) while ignoring cancellation, and a short + // timeout bounds the cleanup writes. + cleanupCtx, cleanupCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cleanupCancel() + // Ensure we finish thinking on error to close the reasoning state. + currentAssistant.FinishThinking() + toolCalls := currentAssistant.ToolCalls() + // INFO: we use the cleanup context here because the genCtx has been cancelled. + msgs, createErr := a.messages.List(cleanupCtx, currentAssistant.SessionID) + if createErr != nil { + return nil, createErr + } + for _, tc := range toolCalls { + if !tc.Finished { + tc.Finished = true + tc.Input = "{}" + currentAssistant.AddToolCall(tc) + updateErr := a.messages.Update(cleanupCtx, *currentAssistant) + if updateErr != nil { + return nil, updateErr + } + } + + found := false + for _, msg := range msgs { + if msg.Role == message.Tool { + for _, tr := range msg.ToolResults() { + if tr.ToolCallID == tc.ID { + found = true + break + } + } + } + if found { + break + } + } + if found { + continue + } + content := "There was an error while executing the tool" + if isCancelErr { + content = "Error: user cancelled assistant tool calling" + } + toolResult := message.ToolResult{ + ToolCallID: tc.ID, + Name: tc.Name, + Content: content, + IsError: true, + } + _, createErr = a.messages.Create(cleanupCtx, currentAssistant.SessionID, message.CreateMessageParams{ + Role: message.Tool, + Parts: []message.ContentPart{ + toolResult, + }, + }) + if createErr != nil { + return nil, createErr + } + } + var fantasyErr *fantasy.Error + var providerErr *fantasy.ProviderError + var requestTimedOutErr *requestTimeoutError + const defaultTitle = "Provider Error" + linkStyle := lipgloss.NewStyle().Foreground(charmtone.Guac).Underline(true) + if isCancelErr { + currentAssistant.AddFinish(message.FinishReasonCanceled, "User canceled request", "") + } else if errors.As(err, &requestTimedOutErr) { + // Checked before the provider branches so a deadline our own + // request timeout imposed is never reported as a provider error. + currentAssistant.AddFinish(message.FinishReasonError, "Request timed out", requestTimedOutErr.userMessage()) + } else if isHyper && errors.As(err, &providerErr) && providerErr.StatusCode == http.StatusUnauthorized { + currentAssistant.AddFinish(message.FinishReasonError, "Unauthorized", `Please re-authenticate with Hyper. You can also run "crush auth" to re-authenticate.`) + } else if errors.As(err, &providerErr) { + if providerErr.Message == "The requested model is not supported." { + url := "https://github.com/settings/copilot/features" + link := linkStyle.Hyperlink(url, "id=copilot").Render(url) + currentAssistant.AddFinish( + message.FinishReasonError, + "Copilot model not enabled", + fmt.Sprintf("%q is not enabled in Copilot. Go to the following page to enable it. Then, wait 5 minutes before trying again. %s", largeModel.CatwalkCfg.Name, link), + ) + } else { + currentAssistant.AddFinish(message.FinishReasonError, cmp.Or(stringext.Capitalize(providerErr.Title), defaultTitle), providerErr.Message) + } + } else if errors.As(err, &fantasyErr) { + currentAssistant.AddFinish(message.FinishReasonError, cmp.Or(stringext.Capitalize(fantasyErr.Title), defaultTitle), fantasyErr.Message) + } else if fantasy.IsTransportError(err) { + wrapped := fantasy.NewTransportError(err) + currentAssistant.AddFinish(message.FinishReasonError, stringext.Capitalize(wrapped.Title), wrapped.Message) + } else { + currentAssistant.AddFinish(message.FinishReasonError, defaultTitle, err.Error()) + } + // Note: we use the cleanup context here because the genCtx has been + // cancelled. + updateErr := a.messages.Update(cleanupCtx, *currentAssistant) + if updateErr != nil { + return nil, updateErr + } + return nil, err + } + + if shouldSummarize { + a.activeRequests.Del(call.SessionID) + if summarizeErr := a.Summarize(genCtx, call.SessionID, call.ProviderOptions, call.OnAuthRefresh); summarizeErr != nil { + return nil, summarizeErr + } + // If the agent wasn't done... + if len(currentAssistant.ToolCalls()) > 0 { + existing, ok := a.messageQueue.Get(call.SessionID) + if !ok { + existing = []SessionAgentCall{} + } + call.Prompt = fmt.Sprintf("The previous session was interrupted because it got too long, the initial user request was: `%s`", call.Prompt) + existing = append(existing, call) + a.messageQueue.Set(call.SessionID, existing) + } + } + + // Release active request before publishing the notification. + // TUI handlers poll IsSessionBusy() and only re-evaluate when a + // tea.Msg arrives, so the cleanup must precede the notify or + // subscribers see stale busy state at the moment of receipt. + a.activeRequests.Del(call.SessionID) + cancel() + + // Send notification that agent has finished its turn (skip for + // nested/non-interactive sessions). + if !call.NonInteractive && a.notify != nil { + a.notify.Publish(pubsub.CreatedEvent, notify.Notification{ + SessionID: call.SessionID, + SessionTitle: currentSession.Title, + Type: notify.TypeAgentFinished, + }) + } + + // Hand off to the next queued prompt (if any) under dispatchMu so + // the transition from this finished run to the queued run is atomic + // against a concurrent Cancel. activeRequests for this session was + // just deleted above, so without the lock there is a window in + // which the session looks idle and a cancel becomes a no-op that + // fails to stop the queued prompt. Holding the lock lets us observe + // a pending cancel recorded against the session and drop the queue + // instead of running it, and (for the recursion) hand a fresh + // accept reservation to the dequeued call so acceptedRuns stays > 0 + // across the recursive Run's own dispatch handoff — keeping the + // session observable to Cancel for the entire transition and + // closing the dequeue -> re-register window. + mu := a.sessionMu(call.SessionID) + mu.Lock() + queuedMessages, _ := a.messageQueue.Get(call.SessionID) + if mark, ok := a.cancelMark.Get(call.SessionID); ok && mark > 0 && len(queuedMessages) > 0 { + // A cancel was recorded for this session (e.g. it arrived while + // this run was active and follow-ups had been queued). Drop the + // queued prompts it covers (accept sequence at or below the + // mark, or untracked); keep any queued after the cancel (higher + // sequence) so they still run. + var kept []SessionAgentCall + var canceledRunIDDrops []SessionAgentCall + for _, q := range queuedMessages { + if q.acceptSeq == 0 || q.acceptSeq <= mark { + if q.RunID != "" { + canceledRunIDDrops = append(canceledRunIDDrops, q) + } + continue + } + kept = append(kept, q) + } + queuedMessages = kept + a.messageQueue.Set(call.SessionID, kept) + // A dropped prompt carrying a RunID must still publish its + // terminal cancelled RunComplete so a caller waiting on that + // RunID does not hang. + a.publishCanceledQueueDrops(canceledRunIDDrops) + } + if len(queuedMessages) == 0 { + // No queued work. Clear the cancel mark only when no accepted + // run remains in flight that it might still cover; otherwise a + // sibling prompt (sequence at or below the mark) waiting to + // enter Run would lose its cancellation. When accepted runs are + // gone, this also clears a stale mark so it can't catch a + // future run. + a.messageQueue.Del(call.SessionID) + a.acceptedMu.Lock() + inFlight, _ := a.acceptedRuns.Get(call.SessionID) + a.acceptedMu.Unlock() + if inFlight == 0 { + a.cancelMark.Del(call.SessionID) + } + mu.Unlock() + return result, err + } + // There are queued messages, restart the loop. Suppress the outer + // defer's emit: it would otherwise observe the recursive Run's retErr + // (named-return clobbering through the return below) against this + // turn's MessageID/Text and publish a mixed, racing event. + skipRunComplete = true + // Decide whether this turn still owes its own terminal RunComplete. + // Each submitted prompt with a RunID has its own lifecycle, so a turn + // that is finished and handing off to a *different* queued prompt must + // publish its own RunComplete here — leaving it to the recursive turn + // (which carries a different RunID) would hang a caller waiting on + // this turn's RunID. The exception is the summarize-continuation path, + // which re-queues this same call (same RunID) to resume after a + // summary; in that case the eventual terminal turn for this RunID + // publishes, so publishing now would double-emit. + outerOwesRunComplete := call.RunID != "" + if outerOwesRunComplete { + for _, q := range queuedMessages { + if q.RunID == call.RunID { + outerOwesRunComplete = false + break + } + } + } + firstQueuedMessage := queuedMessages[0] + a.messageQueue.Set(call.SessionID, queuedMessages[1:]) + // Reserve a fresh accept for the dequeued prompt before dropping the + // lock so acceptedRuns > 0 across the handoff into the recursive + // Run. This closes the window between this dequeue and the recursive + // Run registering its activeRequests entry: a cancel arriving in + // that window now records a pending cancel (acceptedRuns > 0) that + // the recursive Run's accepted path observes as cancel-on-entry. + firstQueuedMessage.Accepted = a.BeginAccepted(call.SessionID) + mu.Unlock() + if outerOwesRunComplete { + complete := notify.RunComplete{SessionID: call.SessionID, RunID: call.RunID} + if currentAssistant != nil { + complete.MessageID = currentAssistant.ID + complete.Text = currentAssistant.Content().String() + } + if ctx.Err() != nil { + complete.Cancelled = true + } + a.publishRunComplete(ctx, call, complete) + } + return a.Run(ctx, firstQueuedMessage) +} + +func (a *sessionAgent) Summarize(ctx context.Context, sessionID string, opts fantasy.ProviderOptions, onAuthRefresh func(context.Context, *fantasy.ProviderError) error) error { + if a.IsSessionBusy(sessionID) { + return ErrSessionBusy + } + + // Copy mutable fields under lock to avoid races with SetModels. + largeModel := a.largeModel.Get() + systemPromptPrefix := a.systemPromptPrefix.Get() + + currentSession, err := a.sessions.Get(ctx, sessionID) + if err != nil { + return fmt.Errorf("failed to get session: %w", err) + } + msgs, err := a.getSessionMessages(ctx, currentSession) + if err != nil { + return err + } + if len(msgs) == 0 { + // Nothing to summarize. + return nil + } + + aiMsgs, _ := a.preparePrompt(msgs, largeModel.CatwalkCfg.SupportsImages) + + genCtx, cancel := context.WithCancel(ctx) + ac := &activeCancel{cancel: cancel} + a.activeRequests.Set(sessionID, ac) + defer a.activeRequests.CompareAndDelete(sessionID, ac) + defer cancel() + defer func() { + if flushErr := a.messages.FlushAll(ctx); flushErr != nil { + slog.Error("Failed to flush pending message updates after summarize", "error", flushErr) + } + }() + + agent := fantasy.NewAgent( + largeModel.Model, + fantasy.WithSystemPrompt(string(summaryPrompt)), + fantasy.WithUserAgent(userAgent), + ) + summaryMessage, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{ + Role: message.Assistant, + Model: largeModel.ModelCfg.Model, + Provider: largeModel.ModelCfg.Provider, + IsSummaryMessage: true, + }) + if err != nil { + return err + } + + summaryPromptText := buildSummaryPrompt(currentSession.Todos) + + resp, err := agent.Stream(genCtx, fantasy.AgentStreamCall{ + Prompt: summaryPromptText, + Messages: aiMsgs, + Headers: sessionHeaders(sessionID), + ProviderOptions: opts, + OnAuthRefresh: onAuthRefresh, + ModelProvider: func() fantasy.LanguageModel { + return a.largeModel.Get().Model + }, + PrepareStep: func(callContext context.Context, options fantasy.PrepareStepFunctionOptions) (_ context.Context, prepared fantasy.PrepareStepResult, err error) { + prepared.Messages = options.Messages + if systemPromptPrefix != "" { + prepared.Messages = append([]fantasy.Message{fantasy.NewSystemMessage(systemPromptPrefix)}, prepared.Messages...) + } + return callContext, prepared, nil + }, + OnReasoningDelta: func(id string, text string) error { + summaryMessage.AppendReasoningContent(text) + return a.messages.Update(genCtx, summaryMessage) + }, + OnReasoningEnd: func(id string, reasoning fantasy.ReasoningContent) error { + // Handle anthropic signature. + if anthropicData, ok := reasoning.ProviderMetadata["anthropic"]; ok { + if signature, ok := anthropicData.(*anthropic.ReasoningOptionMetadata); ok && signature.Signature != "" { + summaryMessage.AppendReasoningSignature(signature.Signature) + } + } + summaryMessage.FinishThinking() + return a.messages.Update(genCtx, summaryMessage) + }, + OnTextDelta: func(id, text string) error { + summaryMessage.AppendContent(text) + return a.messages.Update(genCtx, summaryMessage) + }, + }) + if err != nil { + isCancelErr := errors.Is(err, context.Canceled) + if isCancelErr { + // User cancelled summarize we need to remove the summary message. + deleteErr := a.messages.Delete(ctx, summaryMessage.ID) + return deleteErr + } + // Mark the summary message as finished with an error so the UI + // stops spinning. + summaryMessage.AddFinish(message.FinishReasonError, "Summarization Error", err.Error()) + if updateErr := a.messages.Update(ctx, summaryMessage); updateErr != nil { + return updateErr + } + return err + } + + summaryMessage.AddFinish(message.FinishReasonEndTurn, "", "") + err = a.messages.Update(genCtx, summaryMessage) + if err != nil { + return err + } + + var openrouterCost *float64 + for _, step := range resp.Steps { + stepCost := a.openrouterCost(step.ProviderMetadata) + if stepCost != nil { + newCost := *stepCost + if openrouterCost != nil { + newCost += *openrouterCost + } + openrouterCost = &newCost + } + extractHyperCredits(step.ProviderMetadata) + } + + a.updateSessionUsage(largeModel, ¤tSession, resp.TotalUsage, openrouterCost, false) + + // Just in case, get just the last usage info. + usage := resp.Response.Usage + currentSession.SummaryMessageID = summaryMessage.ID + currentSession.CompletionTokens = summaryCompletionTokens(usage, summaryMessage) + currentSession.PromptTokens = 0 + currentSession.EstimatedUsage = usageIsZero(usage) + _, err = a.sessions.Save(genCtx, currentSession) + if err != nil { + return err + } + + // Release the active request before processing queued messages so that + // Run() does not see the session as busy. + a.activeRequests.Del(sessionID) + cancel() + + // Process any messages that were queued while summarizing. + queuedMessages, ok := a.messageQueue.Get(sessionID) + if !ok || len(queuedMessages) == 0 { + return nil + } + firstQueuedMessage := queuedMessages[0] + a.messageQueue.Set(sessionID, queuedMessages[1:]) + _, qErr := a.Run(ctx, firstQueuedMessage) + return qErr +} + +func (a *sessionAgent) getCacheControlOptions() fantasy.ProviderOptions { + if t, _ := strconv.ParseBool(os.Getenv("CRUSH_DISABLE_ANTHROPIC_CACHE")); t { + return fantasy.ProviderOptions{} + } + return fantasy.ProviderOptions{ + anthropic.Name: &anthropic.ProviderCacheControlOptions{ + CacheControl: anthropic.CacheControl{Type: "ephemeral"}, + }, + bedrock.Name: &anthropic.ProviderCacheControlOptions{ + CacheControl: anthropic.CacheControl{Type: "ephemeral"}, + }, + vercel.Name: &anthropic.ProviderCacheControlOptions{ + CacheControl: anthropic.CacheControl{Type: "ephemeral"}, + }, + } +} + +// sessionHeaders returns the HTTP headers we use for cache affinity on +// every LLM request for a given session. +// +// We use the session hash is used instead of the raw UUID so the header +// value is deterministic and opaque. +func sessionHeaders(sessionID string) map[string]string { + hash := session.HashID(sessionID) + return map[string]string{ + "x-session-id": hash, + "x-session-affinity": hash, + } +} + +func (a *sessionAgent) createUserMessage(ctx context.Context, call SessionAgentCall) (message.Message, error) { + parts := []message.ContentPart{message.TextContent{Text: call.Prompt}} + var attachmentParts []message.ContentPart + for _, attachment := range call.Attachments { + attachmentParts = append(attachmentParts, message.BinaryContent{Path: attachment.FilePath, MIMEType: attachment.MimeType, Data: attachment.Content}) + } + parts = append(parts, attachmentParts...) + msg, err := a.messages.Create(ctx, call.SessionID, message.CreateMessageParams{ + Role: message.User, + Parts: parts, + }) + if err != nil { + return message.Message{}, fmt.Errorf("failed to create user message: %w", err) + } + return msg, nil +} + +func (a *sessionAgent) preparePrompt(msgs []message.Message, supportsImages bool, attachments ...message.Attachment) ([]fantasy.Message, []fantasy.FilePart) { + var history []fantasy.Message + if !a.isSubAgent { + history = append(history, fantasy.NewUserMessage( + fmt.Sprintf( + "%s", + `This is a reminder that your todo list is currently empty. DO NOT mention this to the user explicitly because they are already aware. +If you are working on tasks that would benefit from a todo list please use the "todos" tool to create one. +If not, please feel free to ignore. Again do not mention this message to the user.`, + ), + )) + } + // Collect all tool call IDs present in assistant messages, then index + // every tool result by its call ID. Tool results are re-emitted right + // after the assistant message that requested them instead of at their + // stored position: messages can be written to a session concurrently + // (e.g. resuming while a tool is still running), which interleaves + // user messages between a tool call and its result. LLM APIs require + // every tool call to be followed by its results before any other + // message, and strict-adjacency providers (e.g. Kimi, DeepSeek) reject + // the request otherwise, permanently locking the session. + knownToolCallIDs := make(map[string]struct{}) + for _, m := range msgs { + if m.Role != message.Assistant { + continue + } + for _, tc := range m.ToolCalls() { + knownToolCallIDs[tc.ID] = struct{}{} + } + } + toolResultsByCall := make(map[string][]fantasy.MessagePart) + for _, m := range msgs { + if m.Role != message.Tool { + continue + } + for _, aiMsg := range m.ToAIMessage() { + for _, part := range aiMsg.Content { + tr, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](part) + if !ok { + // Tool-role ToAIMessage only emits ToolResultParts today; + // log so unexpected parts do not vanish silently. + slog.Warn( + "Dropping unexpected non-tool-result part from tool message", + "part_type", fmt.Sprintf("%T", part), + ) + continue + } + if _, known := knownToolCallIDs[tr.ToolCallID]; !known { + slog.Warn( + "Dropping orphaned tool result with no matching tool call", + "tool_call_id", tr.ToolCallID, + ) + continue + } + toolResultsByCall[tr.ToolCallID] = append(toolResultsByCall[tr.ToolCallID], part) + } + } + } + + for _, m := range msgs { + if len(m.Parts) == 0 { + continue + } + // Assistant message without content or tool calls (cancelled before it returned anything). + if m.Role == message.Assistant && len(m.ToolCalls()) == 0 && m.Content().Text == "" && m.ReasoningContent().String() == "" { + continue + } + // Tool results are emitted right after their assistant message. + if m.Role == message.Tool { + continue + } + aiMsgs := m.ToAIMessage() + if !supportsImages { + for i := range aiMsgs { + if aiMsgs[i].Role == fantasy.MessageRoleUser { + aiMsgs[i].Content = filterFileParts(aiMsgs[i].Content) + } + } + } + history = append(history, aiMsgs...) + + if m.Role == message.Assistant && len(m.ToolCalls()) > 0 { + history = append(history, toolResultsForCalls(m, toolResultsByCall)) + } + } + + var files []fantasy.FilePart + for _, attachment := range attachments { + if attachment.IsText() { + continue + } + if !supportsImages { + continue + } + files = append(files, fantasy.FilePart{ + Filename: attachment.FileName, + Data: attachment.Content, + MediaType: attachment.MimeType, + }) + } + + return history, files +} + +// filterFileParts removes fantasy.FilePart entries from a slice of message +// parts. Used to strip image attachments from historical user messages when +// the current model does not support them. +func filterFileParts(parts []fantasy.MessagePart) []fantasy.MessagePart { + filtered := make([]fantasy.MessagePart, 0, len(parts)) + for _, part := range parts { + if _, ok := fantasy.AsMessagePart[fantasy.FilePart](part); ok { + continue + } + filtered = append(filtered, part) + } + return filtered +} + +// toolResultsForCalls builds the tool message that must immediately follow +// an assistant message with tool calls. LLM APIs require every tool call to +// be followed by its results before any other message; strict-adjacency +// providers reject the request otherwise. Results are taken from +// toolResultsByCall and consumed, so a result stored in a message that also +// holds results for calls of other assistant messages is emitted exactly +// once, next to the assistant that requested it. Tool calls without any +// stored result (e.g. an interrupted session) receive a synthetic error +// response so the conversation keeps working. +func toolResultsForCalls(m message.Message, toolResultsByCall map[string][]fantasy.MessagePart) fantasy.Message { + content := make([]fantasy.MessagePart, 0, len(m.ToolCalls())) + for _, tc := range m.ToolCalls() { + parts := toolResultsByCall[tc.ID] + delete(toolResultsByCall, tc.ID) + if len(parts) > 0 { + content = append(content, parts...) + continue + } + slog.Warn( + "Injecting synthetic tool result for orphaned tool call", + "tool_call_id", tc.ID, + "tool_name", tc.Name, + ) + content = append(content, fantasy.ToolResultPart{ + ToolCallID: tc.ID, + Output: fantasy.ToolResultOutputContentError{ + Error: errors.New("tool call was interrupted and did not produce a result, you may retry this call if the result is still needed"), + }, + }) + } + return fantasy.Message{ + Role: fantasy.MessageRoleTool, + Content: content, + } +} + +func (a *sessionAgent) getSessionMessages(ctx context.Context, session session.Session) ([]message.Message, error) { + msgs, err := a.messages.List(ctx, session.ID) + if err != nil { + return nil, fmt.Errorf("failed to list messages: %w", err) + } + + if session.SummaryMessageID != "" { + summaryMsgIndex := -1 + for i, msg := range msgs { + if msg.ID == session.SummaryMessageID { + summaryMsgIndex = i + break + } + } + if summaryMsgIndex != -1 { + msgs = msgs[summaryMsgIndex:] + msgs[0].Role = message.User + } + } + return msgs, nil +} + +// hasUserTextMessage reports whether any user message in msgs contains +// text content (as opposed to only shell commands or other non-text parts). +func hasUserTextMessage(msgs []message.Message) bool { + for _, msg := range msgs { + if msg.Role != message.User { + continue + } + for _, part := range msg.Parts { + if tc, ok := part.(message.TextContent); ok && tc.Text != "" { + return true + } + } + } + return false +} + +// GenerateTitle generates a session title based on the initial prompt. +func (a *sessionAgent) GenerateTitle(ctx context.Context, sessionID string, userPrompt string) { + if userPrompt == "" { + return + } + + // Ensure the session always gets a title even if every path below + // fails or the context is cancelled before we finish. + var titleSaved bool + defer func() { + if !titleSaved { + fallbackCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + if err := a.sessions.Rename(fallbackCtx, sessionID, DefaultSessionName); err != nil { + slog.Error("Failed to save fallback session title", "error", err) + } + } + }() + + smallModel := a.smallModel.Get() + largeModel := a.largeModel.Get() + systemPromptPrefix := a.systemPromptPrefix.Get() + + newAgent := func(m fantasy.LanguageModel, p []byte, tok int64) fantasy.Agent { + return fantasy.NewAgent( + m, + fantasy.WithSystemPrompt(string(p)+"\n /no_think"), + fantasy.WithMaxOutputTokens(tok), + fantasy.WithUserAgent(userAgent), + ) + } + + streamCall := fantasy.AgentStreamCall{ + Prompt: fmt.Sprintf("Generate a concise title for the following content:\n\n%s\n \n\n", userPrompt), + Headers: sessionHeaders(sessionID), + PrepareStep: func(callCtx context.Context, opts fantasy.PrepareStepFunctionOptions) (_ context.Context, prepared fantasy.PrepareStepResult, err error) { + prepared.Messages = opts.Messages + if systemPromptPrefix != "" { + prepared.Messages = append([]fantasy.Message{ + fantasy.NewSystemMessage(systemPromptPrefix), + }, prepared.Messages...) + } + return callCtx, prepared, nil + }, + } + + type modelAttempt struct { + name string + model Model + } + attempts := []modelAttempt{ + {"small", smallModel}, + {"large", largeModel}, + } + + var resp *fantasy.AgentResult + var err error + var model Model + var success bool + for _, attempt := range attempts { + tok := int64(40) + if attempt.model.CatwalkCfg.CanReason { + tok = attempt.model.CatwalkCfg.DefaultMaxTokens + } + agent := newAgent(attempt.model.Model, titlePrompt, tok) + resp, err = agent.Stream(ctx, streamCall) + if err == nil && resp.Response.FinishReason != fantasy.FinishReasonLength { + model = attempt.model + slog.Debug("Generated title with " + attempt.name + " model") + success = true + break + } + if err != nil { + slog.Error("Error generating title with "+attempt.name+" model; trying next", "err", err) + } else { + slog.Error("Title generation hit token limit with " + attempt.name + " model; trying next") + } + } + if !success { + // The deferred fallback will save the default session name. + return + } + + // Clean up title. + var title string + title = strings.ReplaceAll(resp.Response.Content.Text(), "\n", " ") + + // Remove thinking tags if present. + title = thinkTagRegex.ReplaceAllString(title, "") + title = orphanThinkTagRegex.ReplaceAllString(title, "") + + title = strings.TrimSpace(title) + if title == "" { + // LLM returned empty content. Use the prompt itself as a + // fallback title, truncated to 50 chars, before resorting to + // the generic default. + fallback := strings.ReplaceAll(userPrompt, "\n", " ") + fallback = strings.TrimSpace(fallback) + if len(fallback) > 50 { + fallback = ansi.Truncate(fallback, 50, "…") + } + title = cmp.Or(fallback, DefaultSessionName) + } + + // Calculate usage and cost. + var openrouterCost *float64 + for _, step := range resp.Steps { + stepCost := a.openrouterCost(step.ProviderMetadata) + if stepCost != nil { + newCost := *stepCost + if openrouterCost != nil { + newCost += *openrouterCost + } + openrouterCost = &newCost + } + extractHyperCredits(step.ProviderMetadata) + } + + modelConfig := model.CatwalkCfg + cost := modelConfig.CostPer1MInCached/1e6*float64(resp.TotalUsage.CacheCreationTokens) + + modelConfig.CostPer1MOutCached/1e6*float64(resp.TotalUsage.CacheReadTokens) + + modelConfig.CostPer1MIn/1e6*float64(resp.TotalUsage.InputTokens) + + modelConfig.CostPer1MOut/1e6*float64(resp.TotalUsage.OutputTokens) + + // Use override cost if available (e.g., from OpenRouter). + if openrouterCost != nil { + cost = *openrouterCost + } + + // Skip cost accumulation + if model.FlatRate { + cost = 0 + } + + promptTokens := resp.TotalUsage.InputTokens + resp.TotalUsage.CacheCreationTokens + completionTokens := resp.TotalUsage.OutputTokens + + // Atomically update only title and usage fields to avoid overriding other + // concurrent session updates. + saveErr := a.sessions.UpdateTitleAndUsage(ctx, sessionID, title, promptTokens, completionTokens, cost) + if saveErr != nil { + slog.Error("Failed to save session title and usage", "error", saveErr) + return + } + titleSaved = true +} + +func (a *sessionAgent) openrouterCost(metadata fantasy.ProviderMetadata) *float64 { + openrouterMetadata, ok := metadata[openrouter.Name] + if !ok { + return nil + } + + opts, ok := openrouterMetadata.(*openrouter.ProviderMetadata) + if !ok { + return nil + } + return &opts.Usage.Cost +} + +// extractHyperCredits reads usage.remaining.hypercredits from OpenAI +// provider metadata and stores it for the next FetchCredits call. +func extractHyperCredits(metadata fantasy.ProviderMetadata) { + openaiMeta, ok := metadata[openai.Name] + if !ok { + return + } + pm, ok := openaiMeta.(*openai.ProviderMetadata) + if !ok { + return + } + var remaining struct { + Hypercredits float64 `json:"hypercredits"` + } + if pm.ExtraField("remaining", &remaining) && remaining.Hypercredits > 0 { + hyper.SetBalance(int(math.Round(remaining.Hypercredits))) + } +} + +// extractPrismModel returns the ID and name of the model that actually +// served the turn, as reported by the Hyper Prism model router headers, +// or empty strings when the turn was not routed through a Prism model. +func extractPrismModel(metadata fantasy.ProviderMetadata) (modelID, modelName string) { + openaiMeta, ok := metadata[openai.Name] + if !ok { + return "", "" + } + pm, ok := openaiMeta.(*openai.ProviderMetadata) + if !ok { + return "", "" + } + _ = pm.ExtraField(hyper.PrismModelIDField, &modelID) + _ = pm.ExtraField(hyper.PrismModelNameField, &modelName) + return modelID, modelName +} + +// extractPrismSavings returns the hypercredit and dollar savings from +// routing the turn through the Hyper Prism model router, as reported by +// its savings trailers, or nil when not reported or malformed. +func extractPrismSavings(metadata fantasy.ProviderMetadata) (hypercredits, dollars *float64) { + openaiMeta, ok := metadata[openai.Name] + if !ok { + return nil, nil + } + pm, ok := openaiMeta.(*openai.ProviderMetadata) + if !ok { + return nil, nil + } + return extraFieldFloat(pm, hyper.PrismHypercreditSavingsField), extraFieldFloat(pm, hyper.PrismDollarSavingsField) +} + +func extraFieldFloat(pm *openai.ProviderMetadata, key string) *float64 { + var value string + if !pm.ExtraField(key, &value) { + return nil + } + parsed, err := strconv.ParseFloat(value, 64) + if err != nil { + slog.Warn("Could not parse Prism savings", "key", key, "value", value, "error", err) + return nil + } + return &parsed +} + +func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session, usage fantasy.Usage, overrideCost *float64, estimated bool) { + if !usageIsZero(usage) { + session.EstimatedUsage = estimated + } + + modelConfig := model.CatwalkCfg + cost := modelConfig.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) + + modelConfig.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) + + modelConfig.CostPer1MIn/1e6*float64(usage.InputTokens) + + modelConfig.CostPer1MOut/1e6*float64(usage.OutputTokens) + + if !estimated { + a.eventTokensUsed(session.ID, model, usage, cost) + } + + if estimated { + cost = 0 + } else { + // Use override cost if available (e.g., from OpenRouter). + if overrideCost != nil { + cost = *overrideCost + } + + // Skip cost accumulation + if model.FlatRate { + cost = 0 + } + } + + session.Cost += cost + updateSessionTokenCounters(session, usage) +} + +func updateSessionTokenCounters(session *session.Session, usage fantasy.Usage) { + if usage.OutputTokens != 0 { + session.CompletionTokens = usage.OutputTokens + } + if promptTokens := usage.InputTokens + usage.CacheReadTokens; promptTokens != 0 { + session.PromptTokens = promptTokens + } +} + +func summaryCompletionTokens(usage fantasy.Usage, summaryMessage message.Message) int64 { + if usage.OutputTokens != 0 { + return usage.OutputTokens + } + return approxTokenCount(summaryMessage.Content().Text) + approxTokenCount(summaryMessage.ReasoningContent().String()) +} + +func (a *sessionAgent) Cancel(sessionID string) { + // Serialize against the dispatch handoff in Run so the accepted -> + // (cancel-on-entry | queued | active) transition is atomic against + // this cancel. Every cancel observes at least one of: an active + // request, an accepted run (recorded as a pending cancel), or a + // queue entry it then clears. If none of those hold, an idle Escape + // is a true no-op and must not poison the next prompt. + mu := a.sessionMu(sessionID) + mu.Lock() + defer mu.Unlock() + + // Cancel regular requests. Don't use Take() here - we need the entry to + // remain in activeRequests so IsBusy() returns true until the goroutine + // fully completes (including error handling that may access the DB). + // The defer in processRequest will clean up the entry. + if ac, ok := a.activeRequests.Get(sessionID); ok && ac != nil { + slog.Debug("Request cancellation initiated", "session_id", sessionID) + ac.cancel() + } + + // Also check for summarize requests. + if ac, ok := a.activeRequests.Get(sessionID + "-summarize"); ok && ac != nil { + slog.Debug("Summarize cancellation initiated", "session_id", sessionID) + ac.cancel() + } + + // Record a pending cancel only when a dispatched-but-not-yet-active + // run exists. This catches runs still in the goroutine scheduler or + // about to enter Run's busy-queue branch, while leaving an idle + // session untouched. Active and accepted are not mutually exclusive: + // when a run is active and a follow-up has been accepted, both the + // cancel above and this pending record fire. + // + // Raise the session's cancel mark to the latest accept sequence + // assigned so far. Every prompt currently accepted-but-not-yet- + // active has a sequence at or below that value, so one cancel covers + // all of them; a prompt accepted after this cancel gets a strictly + // higher sequence and is never poisoned. Using max keeps repeated + // cancels idempotent while the same prompts are in flight and lets a + // later cancel extend coverage to prompts accepted since. + a.acceptedMu.Lock() + count, ok := a.acceptedRuns.Get(sessionID) + mark := a.acceptSeqGen + a.acceptedMu.Unlock() + if ok && count > 0 { + slog.Debug("Recording cancel mark for accepted runs", "session_id", sessionID, "count", count, "mark", mark) + existing, _ := a.cancelMark.Get(sessionID) + a.cancelMark.Set(sessionID, max(existing, mark)) + } + + if a.QueuedPrompts(sessionID) > 0 { + slog.Debug("Clearing queued prompts", "session_id", sessionID) + a.clearQueueAndNotify(sessionID) + } +} + +func (a *sessionAgent) ClearQueue(sessionID string) { + if a.QueuedPrompts(sessionID) > 0 { + slog.Debug("Clearing queued prompts", "session_id", sessionID) + a.clearQueueAndNotify(sessionID) + } +} + +func (a *sessionAgent) CancelAll() { + if !a.IsBusy() { + return + } + for key := range a.activeRequests.Seq2() { + a.Cancel(key) // key is sessionID + } + + timeout := time.After(5 * time.Second) + for a.IsBusy() { + select { + case <-timeout: + return + default: + time.Sleep(200 * time.Millisecond) + } + } +} + +func (a *sessionAgent) IsBusy() bool { + var busy bool + for ac := range a.activeRequests.Seq() { + if ac != nil { + busy = true + break + } + } + return busy +} + +func (a *sessionAgent) IsSessionBusy(sessionID string) bool { + _, busy := a.activeRequests.Get(sessionID) + return busy +} + +func (a *sessionAgent) QueuedPrompts(sessionID string) int { + l, ok := a.messageQueue.Get(sessionID) + if !ok { + return 0 + } + return len(l) +} + +func (a *sessionAgent) QueuedPromptsList(sessionID string) []string { + l, ok := a.messageQueue.Get(sessionID) + if !ok { + return nil + } + prompts := make([]string, len(l)) + for i, call := range l { + prompts[i] = call.Prompt + } + return prompts +} + +func (a *sessionAgent) SetModels(large Model, small Model) { + a.largeModel.Set(large) + a.smallModel.Set(small) +} + +func (a *sessionAgent) SetTools(tools []fantasy.AgentTool) { + a.tools.SetSlice(tools) +} + +func (a *sessionAgent) SetSystemPrompt(systemPrompt string) { + a.systemPrompt.Set(systemPrompt) +} + +func (a *sessionAgent) Model() Model { + return a.largeModel.Get() +} + +// convertToToolResult converts a fantasy tool result to a message tool result. +func (a *sessionAgent) convertToToolResult(result fantasy.ToolResultContent) message.ToolResult { + baseResult := message.ToolResult{ + ToolCallID: result.ToolCallID, + Name: result.ToolName, + Metadata: result.ClientMetadata, + } + + switch result.Result.GetType() { + case fantasy.ToolResultContentTypeText: + if r, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentText](result.Result); ok { + baseResult.Content = r.Text + } + case fantasy.ToolResultContentTypeError: + if r, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentError](result.Result); ok { + baseResult.Content = r.Error.Error() + baseResult.IsError = true + } + case fantasy.ToolResultContentTypeMedia: + if r, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentMedia](result.Result); ok { + if !stringext.IsValidBase64(r.Data) { + slog.Warn( + "Tool returned media with invalid base64 data, discarding image", + "tool", result.ToolName, + "tool_call_id", result.ToolCallID, + ) + baseResult.Content = "Tool returned image data with invalid encoding" + baseResult.IsError = true + } else { + content := r.Text + if content == "" { + content = fmt.Sprintf("Loaded %s content", r.MediaType) + } + baseResult.Content = content + baseResult.Data = r.Data + baseResult.MIMEType = r.MediaType + } + } + } + + return baseResult +} + +// workaroundProviderMediaLimitations converts media content in tool results to +// user messages for providers that don't natively support images in tool results. +// +// Problem: OpenAI, Google, OpenRouter, and other OpenAI-compatible providers +// don't support sending images/media in tool result messages - they only accept +// text in tool results. However, they DO support images in user messages. +// +// If we send media in tool results to these providers, the API returns an error. +// +// Solution: For these providers, we: +// 1. Replace the media in the tool result with a text placeholder +// 2. Inject a user message immediately after with the image as a file attachment +// 3. This maintains the tool execution flow while working around API limitations +// +// Anthropic and Bedrock support images natively in tool results, so we skip +// this workaround for them. +// +// Example transformation: +// +// BEFORE: [tool result: image data] +// AFTER: [tool result: "Image loaded - see attached"], [user: image attachment] +func (a *sessionAgent) workaroundProviderMediaLimitations(messages []fantasy.Message, largeModel Model) []fantasy.Message { + providerSupportsMedia := largeModel.ModelCfg.Provider == string(catwalk.InferenceProviderAnthropic) || + largeModel.ModelCfg.Provider == string(catwalk.InferenceProviderBedrock) || + largeModel.ModelCfg.Provider == string(catwalk.InferenceProviderBedrockEurope) + + if providerSupportsMedia { + return messages + } + + supportsImages := largeModel.CatwalkCfg.SupportsImages + + convertedMessages := make([]fantasy.Message, 0, len(messages)) + + for _, msg := range messages { + if msg.Role != fantasy.MessageRoleTool { + convertedMessages = append(convertedMessages, msg) + continue + } + + textParts := make([]fantasy.MessagePart, 0, len(msg.Content)) + var mediaFiles []fantasy.FilePart + + for _, part := range msg.Content { + toolResult, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](part) + if !ok { + textParts = append(textParts, part) + continue + } + + if media, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentMedia](toolResult.Output); ok { + if !supportsImages { + // Model cannot process images. Replace with a text + // placeholder and skip creating a synthetic user + // message with FilePart, which would brick the + // session on text-only models. + textParts = append(textParts, fantasy.ToolResultPart{ + ToolCallID: toolResult.ToolCallID, + Output: fantasy.ToolResultOutputContentText{ + Text: "[Image/media content not supported by this model]", + }, + ProviderOptions: toolResult.ProviderOptions, + }) + continue + } + + decoded, err := base64.StdEncoding.DecodeString(media.Data) + if err != nil { + slog.Warn("Failed to decode media data", "error", err) + textParts = append(textParts, part) + continue + } + + mediaFiles = append(mediaFiles, fantasy.FilePart{ + Data: decoded, + MediaType: media.MediaType, + Filename: fmt.Sprintf("tool-result-%s", toolResult.ToolCallID), + }) + + textParts = append(textParts, fantasy.ToolResultPart{ + ToolCallID: toolResult.ToolCallID, + Output: fantasy.ToolResultOutputContentText{ + Text: "[Image/media content loaded - see attached file]", + }, + ProviderOptions: toolResult.ProviderOptions, + }) + } else { + textParts = append(textParts, part) + } + } + + convertedMessages = append(convertedMessages, fantasy.Message{ + Role: fantasy.MessageRoleTool, + Content: textParts, + }) + + if len(mediaFiles) > 0 { + convertedMessages = append(convertedMessages, fantasy.NewUserMessage( + "Here is the media content from the tool result:", + mediaFiles..., + )) + } + } + + return convertedMessages +} + +// buildSummaryPrompt constructs the prompt text for session summarization. +func buildSummaryPrompt(todos []session.Todo) string { + var sb strings.Builder + sb.WriteString("Provide a detailed summary of our conversation above.") + if len(todos) > 0 { + sb.WriteString("\n\n## Current Todo List\n\n") + for _, t := range todos { + fmt.Fprintf(&sb, "- [%s] %s\n", t.Status, t.Content) + } + sb.WriteString("\nInclude these tasks and their statuses in your summary. ") + sb.WriteString("Instruct the resuming assistant to use the `todos` tool to continue tracking progress on these tasks.") + } + return sb.String() +} + +func providerRetryLogFields(err *fantasy.ProviderError, delay time.Duration) []any { + fields := []any{ + "retry_delay", delay.String(), + } + if err == nil { + return fields + } + fields = append(fields, "status_code", err.StatusCode) + if err.Title != "" { + fields = append(fields, "title", err.Title) + } + if err.Message != "" { + fields = append(fields, "message", err.Message) + } + return fields +} + +// sanitizeToolInput validates tool call JSON from the provider. +// Malformed input is replaced with an empty object to prevent +// stuck conversations from truncated or malformed model output. +// The second return value indicates whether sanitization occurred. +func sanitizeToolInput(toolName, toolCallID, input string) (string, bool) { + if !json.Valid([]byte(input)) { + slog.Warn( + "Malformed tool call JSON from provider, replacing with empty object", + "tool", toolName, + "id", toolCallID, + "input_len", len(input), + ) + return "{}", true + } + return input, false +} diff --git a/internal/agent/agent_test.go b/internal/agent/agent_test.go new file mode 100644 index 0000000000000000000000000000000000000000..5bf82d408f6d9d788a9aedc7b345a5e9024941c2 --- /dev/null +++ b/internal/agent/agent_test.go @@ -0,0 +1,1353 @@ +package agent + +import ( + "encoding/base64" + "fmt" + "log/slog" + "os" + "path/filepath" + "runtime" + "strings" + "testing" + "time" + + "charm.land/catwalk/pkg/catwalk" + "charm.land/fantasy" + "charm.land/x/vcr" + "github.com/charmbracelet/crush/internal/agent/tools" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/session" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + _ "github.com/joho/godotenv/autoload" +) + +func TestMain(m *testing.M) { + slog.SetLogLoggerLevel(slog.LevelError) + m.Run() +} + +var modelPairs = []modelPair{ + {"deepseek-v4", hyperBuilder("deepseek-v4-pro-0813"), hyperBuilder("deepseek-v4-flash-0731")}, +} + +func getModels(t *testing.T, r *vcr.Recorder, pair modelPair) (fantasy.LanguageModel, fantasy.LanguageModel) { + large, err := pair.largeModel(t, r) + require.NoError(t, err) + small, err := pair.smallModel(t, r) + require.NoError(t, err) + return large, small +} + +func setupAgent(t *testing.T, pair modelPair) (SessionAgent, fakeEnv) { + r := vcr.NewRecorder(t) + large, small := getModels(t, r, pair) + env := testEnv(t) + + createSimpleGoProject(t, env.workingDir) + agent, err := coderAgent(r, env, large, small) + require.NoError(t, err) + return agent, env +} + +func TestCoderAgent(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("skipping on windows for now") + } + + for _, pair := range modelPairs { + t.Run(pair.name, func(t *testing.T) { + t.Run("simple test", func(t *testing.T) { + agent, env := setupAgent(t, pair) + + session, err := env.sessions.Create(t.Context(), "New Session") + require.NoError(t, err) + + res, err := agent.Run(t.Context(), SessionAgentCall{ + Prompt: "Hello", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + require.NoError(t, err) + assert.NotNil(t, res) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + // Should have the agent and user message + assert.Equal(t, len(msgs), 2) + }) + t.Run("read a file", func(t *testing.T) { + agent, env := setupAgent(t, pair) + + session, err := env.sessions.Create(t.Context(), "New Session") + require.NoError(t, err) + res, err := agent.Run(t.Context(), SessionAgentCall{ + Prompt: "Read the go mod", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + + require.NoError(t, err) + assert.NotNil(t, res) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + foundFile := false + var tcID string + out: + for _, msg := range msgs { + if msg.Role == message.Assistant { + for _, tc := range msg.ToolCalls() { + if tc.Name == tools.ViewToolName { + tcID = tc.ID + } + } + } + if msg.Role == message.Tool { + for _, tr := range msg.ToolResults() { + if tr.ToolCallID == tcID { + if strings.Contains(tr.Content, "module example.com/testproject") { + foundFile = true + break out + } + } + } + } + } + require.True(t, foundFile) + }) + t.Run("update a file", func(t *testing.T) { + agent, env := setupAgent(t, pair) + + session, err := env.sessions.Create(t.Context(), "New Session") + require.NoError(t, err) + + res, err := agent.Run(t.Context(), SessionAgentCall{ + Prompt: "update the main.go file by changing the print to say hello from crush", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + require.NoError(t, err) + assert.NotNil(t, res) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + + foundRead := false + foundWrite := false + var readTCID, writeTCID string + + for _, msg := range msgs { + if msg.Role == message.Assistant { + for _, tc := range msg.ToolCalls() { + if tc.Name == tools.ViewToolName { + readTCID = tc.ID + } + if tc.Name == tools.EditToolName || tc.Name == tools.WriteToolName { + writeTCID = tc.ID + } + } + } + if msg.Role == message.Tool { + for _, tr := range msg.ToolResults() { + if tr.ToolCallID == readTCID { + foundRead = true + } + if tr.ToolCallID == writeTCID { + foundWrite = true + } + } + } + } + + require.True(t, foundRead, "Expected to find a read operation") + require.True(t, foundWrite, "Expected to find a write operation") + + mainGoPath := filepath.Join(env.workingDir, "main.go") + content, err := os.ReadFile(mainGoPath) + require.NoError(t, err) + require.Contains(t, strings.ToLower(string(content)), "hello from crush") + }) + t.Run("bash tool", func(t *testing.T) { + agent, env := setupAgent(t, pair) + + session, err := env.sessions.Create(t.Context(), "New Session") + require.NoError(t, err) + + res, err := agent.Run(t.Context(), SessionAgentCall{ + Prompt: "use bash to create a file named test.txt with content 'hello bash'. do not print its timestamp", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + require.NoError(t, err) + assert.NotNil(t, res) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + + foundBash := false + var bashTCID string + + for _, msg := range msgs { + if msg.Role == message.Assistant { + for _, tc := range msg.ToolCalls() { + if tc.Name == tools.BashToolName { + bashTCID = tc.ID + } + } + } + if msg.Role == message.Tool { + for _, tr := range msg.ToolResults() { + if tr.ToolCallID == bashTCID { + foundBash = true + } + } + } + } + + require.True(t, foundBash, "Expected to find a bash operation") + + testFilePath := filepath.Join(env.workingDir, "test.txt") + content, err := os.ReadFile(testFilePath) + require.NoError(t, err) + require.Contains(t, string(content), "hello bash") + }) + t.Run("download tool", func(t *testing.T) { + agent, env := setupAgent(t, pair) + + session, err := env.sessions.Create(t.Context(), "New Session") + require.NoError(t, err) + + res, err := agent.Run(t.Context(), SessionAgentCall{ + Prompt: "download the file from https://example-files.online-convert.com/document/txt/example.txt and save it as example.txt", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + require.NoError(t, err) + assert.NotNil(t, res) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + + foundDownload := false + var downloadTCID string + + for _, msg := range msgs { + if msg.Role == message.Assistant { + for _, tc := range msg.ToolCalls() { + if tc.Name == tools.DownloadToolName { + downloadTCID = tc.ID + } + } + } + if msg.Role == message.Tool { + for _, tr := range msg.ToolResults() { + if tr.ToolCallID == downloadTCID { + foundDownload = true + } + } + } + } + + require.True(t, foundDownload, "Expected to find a download operation") + + examplePath := filepath.Join(env.workingDir, "example.txt") + _, err = os.Stat(examplePath) + require.NoError(t, err, "Expected example.txt file to exist") + }) + t.Run("fetch tool", func(t *testing.T) { + agent, env := setupAgent(t, pair) + + session, err := env.sessions.Create(t.Context(), "New Session") + require.NoError(t, err) + + res, err := agent.Run(t.Context(), SessionAgentCall{ + Prompt: "fetch the content from https://example-files.online-convert.com/website/html/example.html and tell me if it contains the word 'John Doe'", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + require.NoError(t, err) + assert.NotNil(t, res) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + + foundFetch := false + var fetchTCID string + + for _, msg := range msgs { + if msg.Role == message.Assistant { + for _, tc := range msg.ToolCalls() { + if tc.Name == tools.FetchToolName { + fetchTCID = tc.ID + } + } + } + if msg.Role == message.Tool { + for _, tr := range msg.ToolResults() { + if tr.ToolCallID == fetchTCID { + foundFetch = true + } + } + } + } + + require.True(t, foundFetch, "Expected to find a fetch operation") + }) + t.Run("glob tool", func(t *testing.T) { + agent, env := setupAgent(t, pair) + + session, err := env.sessions.Create(t.Context(), "New Session") + require.NoError(t, err) + + res, err := agent.Run(t.Context(), SessionAgentCall{ + Prompt: "use glob to find all .go files in the current directory", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + require.NoError(t, err) + assert.NotNil(t, res) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + + foundGlob := false + var globTCID string + + for _, msg := range msgs { + if msg.Role == message.Assistant { + for _, tc := range msg.ToolCalls() { + if tc.Name == tools.GlobToolName { + globTCID = tc.ID + } + } + } + if msg.Role == message.Tool { + for _, tr := range msg.ToolResults() { + if tr.ToolCallID == globTCID { + foundGlob = true + require.Contains(t, tr.Content, "main.go", "Expected glob to find main.go") + } + } + } + } + + require.True(t, foundGlob, "Expected to find a glob operation") + }) + t.Run("grep tool", func(t *testing.T) { + agent, env := setupAgent(t, pair) + + session, err := env.sessions.Create(t.Context(), "New Session") + require.NoError(t, err) + + res, err := agent.Run(t.Context(), SessionAgentCall{ + Prompt: "use grep to search for the word 'package' in go files", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + require.NoError(t, err) + assert.NotNil(t, res) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + + foundGrep := false + var grepTCID string + + for _, msg := range msgs { + if msg.Role == message.Assistant { + for _, tc := range msg.ToolCalls() { + if tc.Name == tools.GrepToolName { + grepTCID = tc.ID + } + } + } + if msg.Role == message.Tool { + for _, tr := range msg.ToolResults() { + if tr.ToolCallID == grepTCID { + foundGrep = true + require.Contains(t, tr.Content, "main.go", "Expected grep to find main.go") + } + } + } + } + + require.True(t, foundGrep, "Expected to find a grep operation") + }) + t.Run("ls tool", func(t *testing.T) { + agent, env := setupAgent(t, pair) + + session, err := env.sessions.Create(t.Context(), "New Session") + require.NoError(t, err) + + res, err := agent.Run(t.Context(), SessionAgentCall{ + Prompt: "use ls to list the files in the current directory", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + require.NoError(t, err) + assert.NotNil(t, res) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + + foundLS := false + var lsTCID string + + for _, msg := range msgs { + if msg.Role == message.Assistant { + for _, tc := range msg.ToolCalls() { + if tc.Name == tools.LSToolName { + lsTCID = tc.ID + } + } + } + if msg.Role == message.Tool { + for _, tr := range msg.ToolResults() { + if tr.ToolCallID == lsTCID { + foundLS = true + require.Contains(t, tr.Content, "main.go", "Expected ls to list main.go") + require.Contains(t, tr.Content, "go.mod", "Expected ls to list go.mod") + } + } + } + } + + require.True(t, foundLS, "Expected to find an ls operation") + }) + t.Run("multiedit tool", func(t *testing.T) { + agent, env := setupAgent(t, pair) + + session, err := env.sessions.Create(t.Context(), "New Session") + require.NoError(t, err) + + res, err := agent.Run(t.Context(), SessionAgentCall{ + Prompt: "use multiedit to change 'Hello, World!' to 'Hello, Crush!' and add a comment '// Greeting' above the fmt.Println line in main.go", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + require.NoError(t, err) + assert.NotNil(t, res) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + + foundMultiEdit := false + var multiEditTCID string + + for _, msg := range msgs { + if msg.Role == message.Assistant { + for _, tc := range msg.ToolCalls() { + if tc.Name == tools.MultiEditToolName { + multiEditTCID = tc.ID + } + } + } + if msg.Role == message.Tool { + for _, tr := range msg.ToolResults() { + if tr.ToolCallID == multiEditTCID { + foundMultiEdit = true + } + } + } + } + + require.True(t, foundMultiEdit, "Expected to find a multiedit operation") + + mainGoPath := filepath.Join(env.workingDir, "main.go") + content, err := os.ReadFile(mainGoPath) + require.NoError(t, err) + require.Contains(t, string(content), "Hello, Crush!", "Expected file to contain 'Hello, Crush!'") + }) + t.Run("sourcegraph tool", func(t *testing.T) { + agent, env := setupAgent(t, pair) + + session, err := env.sessions.Create(t.Context(), "New Session") + require.NoError(t, err) + + res, err := agent.Run(t.Context(), SessionAgentCall{ + Prompt: "use sourcegraph to search for 'func main' in Go repositories", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + require.NoError(t, err) + assert.NotNil(t, res) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + + foundSourcegraph := false + var sourcegraphTCID string + + for _, msg := range msgs { + if msg.Role == message.Assistant { + for _, tc := range msg.ToolCalls() { + if tc.Name == tools.SourcegraphToolName { + sourcegraphTCID = tc.ID + } + } + } + if msg.Role == message.Tool { + for _, tr := range msg.ToolResults() { + if tr.ToolCallID == sourcegraphTCID { + foundSourcegraph = true + } + } + } + } + + require.True(t, foundSourcegraph, "Expected to find a sourcegraph operation") + }) + t.Run("write tool", func(t *testing.T) { + agent, env := setupAgent(t, pair) + + session, err := env.sessions.Create(t.Context(), "New Session") + require.NoError(t, err) + + res, err := agent.Run(t.Context(), SessionAgentCall{ + Prompt: "use write to create a new file called config.json with content '{\"name\": \"test\", \"version\": \"1.0.0\"}'", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + require.NoError(t, err) + assert.NotNil(t, res) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + + foundWrite := false + var writeTCID string + + for _, msg := range msgs { + if msg.Role == message.Assistant { + for _, tc := range msg.ToolCalls() { + if tc.Name == tools.WriteToolName { + writeTCID = tc.ID + } + } + } + if msg.Role == message.Tool { + for _, tr := range msg.ToolResults() { + if tr.ToolCallID == writeTCID { + foundWrite = true + } + } + } + } + + require.True(t, foundWrite, "Expected to find a write operation") + + configPath := filepath.Join(env.workingDir, "config.json") + content, err := os.ReadFile(configPath) + require.NoError(t, err) + require.Contains(t, string(content), "test", "Expected config.json to contain 'test'") + require.Contains(t, string(content), "1.0.0", "Expected config.json to contain '1.0.0'") + }) + t.Run("parallel tool calls", func(t *testing.T) { + agent, env := setupAgent(t, pair) + + session, err := env.sessions.Create(t.Context(), "New Session") + require.NoError(t, err) + + res, err := agent.Run(t.Context(), SessionAgentCall{ + Prompt: "use glob to find all .go files and use ls to list the current directory, it is very important that you run both tool calls in parallel", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + require.NoError(t, err) + assert.NotNil(t, res) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + + var assistantMsg *message.Message + var toolMsgs []message.Message + + for _, msg := range msgs { + if msg.Role == message.Assistant && len(msg.ToolCalls()) > 0 { + assistantMsg = &msg + } + if msg.Role == message.Tool { + toolMsgs = append(toolMsgs, msg) + } + } + + require.NotNil(t, assistantMsg, "Expected to find an assistant message with tool calls") + require.NotNil(t, toolMsgs, "Expected to find a tool message") + + toolCalls := assistantMsg.ToolCalls() + require.GreaterOrEqual(t, len(toolCalls), 2, "Expected at least 2 tool calls in parallel") + + foundGlob := false + foundLS := false + var globTCID, lsTCID string + + for _, tc := range toolCalls { + if tc.Name == tools.GlobToolName { + foundGlob = true + globTCID = tc.ID + } + if tc.Name == tools.LSToolName { + foundLS = true + lsTCID = tc.ID + } + } + + require.True(t, foundGlob, "Expected to find a glob tool call") + require.True(t, foundLS, "Expected to find an ls tool call") + + require.GreaterOrEqual(t, len(toolMsgs), 2, "Expected at least 2 tool results in the same message") + + foundGlobResult := false + foundLSResult := false + + for _, msg := range toolMsgs { + for _, tr := range msg.ToolResults() { + if tr.ToolCallID == globTCID { + foundGlobResult = true + require.Contains(t, tr.Content, "main.go", "Expected glob result to contain main.go") + require.False(t, tr.IsError, "Expected glob result to not be an error") + } + if tr.ToolCallID == lsTCID { + foundLSResult = true + require.Contains(t, tr.Content, "main.go", "Expected ls result to contain main.go") + require.False(t, tr.IsError, "Expected ls result to not be an error") + } + } + } + + require.True(t, foundGlobResult, "Expected to find glob tool result") + require.True(t, foundLSResult, "Expected to find ls tool result") + }) + }) + } +} + +func makeTestTodos(n int) []session.Todo { + todos := make([]session.Todo, n) + for i := range n { + todos[i] = session.Todo{ + Status: session.TodoStatusPending, + Content: fmt.Sprintf("Task %d: Implement feature with some description that makes it realistic", i), + } + } + return todos +} + +func BenchmarkBuildSummaryPrompt(b *testing.B) { + cases := []struct { + name string + numTodos int + }{ + {"0todos", 0}, + {"5todos", 5}, + {"10todos", 10}, + {"50todos", 50}, + } + + for _, tc := range cases { + todos := makeTestTodos(tc.numTodos) + + b.Run(tc.name, func(b *testing.B) { + b.ReportAllocs() + for range b.N { + _ = buildSummaryPrompt(todos) + } + }) + } +} + +func TestPreparePrompt_FiltersImageAttachments(t *testing.T) { + env := testEnv(t) + sa := testSessionAgent(env, nil, nil, "test prompt") + agent := sa.(*sessionAgent) + + ctx := t.Context() + sess, err := env.sessions.Create(ctx, "test") + require.NoError(t, err) + + // User message with text, a text attachment, and an image attachment. + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.User, + Parts: []message.ContentPart{ + message.TextContent{Text: "hello world"}, + message.BinaryContent{Path: "notes.txt", MIMEType: "text/plain", Data: []byte("important notes")}, + message.BinaryContent{Path: "image.png", MIMEType: "image/png", Data: []byte("fake-image-data")}, + }, + }) + require.NoError(t, err) + + msgs, err := env.messages.List(ctx, sess.ID) + require.NoError(t, err) + + // New-turn image attachment (not yet stored in the DB). + imageAtt := message.Attachment{ + FileName: "screenshot.png", + MimeType: "image/png", + Content: []byte("fake-screenshot"), + } + + // When supportsImages is false, image attachments should be stripped + // from history AND from the files list. + history, files := agent.preparePrompt(msgs, false, imageAtt) + // First message is the system reminder, second is the user message. + require.Len(t, history, 2) + require.Len(t, history[1].Content, 1) + text, ok := fantasy.AsMessagePart[fantasy.TextPart](history[1].Content[0]) + require.True(t, ok) + require.Contains(t, text.Text, "hello world") + require.Contains(t, text.Text, "important notes") + require.Empty(t, files, "image files should be excluded when model does not support images") + + // When supportsImages is true, image attachments should remain in + // history and be included in the files list. + history, files = agent.preparePrompt(msgs, true, imageAtt) + require.Len(t, history, 2) + require.Len(t, history[1].Content, 2) + text, ok = fantasy.AsMessagePart[fantasy.TextPart](history[1].Content[0]) + require.True(t, ok) + require.Contains(t, text.Text, "hello world") + file, ok := fantasy.AsMessagePart[fantasy.FilePart](history[1].Content[1]) + require.True(t, ok) + require.Equal(t, "image.png", file.Filename) + require.Len(t, files, 1, "new-turn image attachment should be included when model supports images") + require.Equal(t, "screenshot.png", files[0].Filename) +} + +func TestCreateUserMessage_RetainsAllAttachments(t *testing.T) { + env := testEnv(t) + sa := testSessionAgent(env, nil, nil, "test prompt") + agent := sa.(*sessionAgent) + + ctx := t.Context() + sess, err := env.sessions.Create(ctx, "test") + require.NoError(t, err) + + // Mix of text and image attachments — all should be stored. + call := SessionAgentCall{ + SessionID: sess.ID, + Prompt: "look at this image", + Attachments: []message.Attachment{ + {FileName: "notes.txt", FilePath: "notes.txt", MimeType: "text/plain", Content: []byte("notes")}, + {FileName: "photo.png", FilePath: "photo.png", MimeType: "image/png", Content: []byte("fake-png")}, + }, + } + + msg, err := agent.createUserMessage(ctx, call) + require.NoError(t, err) + + // All attachments should be present as BinaryContent parts. + binaryParts := msg.BinaryContent() + require.Len(t, binaryParts, 2, "both text and image attachments should be stored in the user message") + require.Equal(t, "notes.txt", binaryParts[0].Path) + require.Equal(t, "text/plain", binaryParts[0].MIMEType) + require.Equal(t, "photo.png", binaryParts[1].Path) + require.Equal(t, "image/png", binaryParts[1].MIMEType) + + // Reload from DB to verify persistence. + reloaded, err := env.messages.Get(ctx, msg.ID) + require.NoError(t, err) + binaryParts = reloaded.BinaryContent() + require.Len(t, binaryParts, 2, "attachments should survive DB round-trip") + require.Equal(t, "photo.png", binaryParts[1].Path) +} + +func TestPreparePrompt_OrphanedToolUse(t *testing.T) { + env := testEnv(t) + sa := testSessionAgent(env, nil, nil, "test prompt") + agent := sa.(*sessionAgent) + + ctx := t.Context() + sess, err := env.sessions.Create(ctx, "test") + require.NoError(t, err) + + // Create a user message. + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.User, + Parts: []message.ContentPart{ + message.TextContent{Text: "hello"}, + }, + }) + require.NoError(t, err) + + // Create an assistant message with a tool call but no tool result — + // this simulates a cancelled/interrupted agent tool call. + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.Assistant, + Parts: []message.ContentPart{ + message.TextContent{Text: "let me check"}, + message.ToolCall{ + ID: "call_orphaned_1", + Name: "agent", + Input: `{"prompt":"do something"}`, + Finished: true, + }, + }, + }) + require.NoError(t, err) + + // Create the next user message (the one that interrupted the tool call). + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.User, + Parts: []message.ContentPart{ + message.TextContent{Text: "Fix #2"}, + }, + }) + require.NoError(t, err) + + msgs, err := env.messages.List(ctx, sess.ID) + require.NoError(t, err) + + history, _ := agent.preparePrompt(msgs, true) + + // The history must contain a synthetic tool result for the orphaned call. + found := false + for _, msg := range history { + if msg.Role != fantasy.MessageRoleTool { + continue + } + for _, part := range msg.Content { + if tr, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](part); ok { + if tr.ToolCallID == "call_orphaned_1" { + found = true + _, isError := tr.Output.(fantasy.ToolResultOutputContentError) + require.True(t, isError, "orphaned tool result should be an error") + } + } + } + } + require.True(t, found, "expected synthetic tool result for orphaned tool call") +} + +func TestPreparePrompt_OrphanedToolUseMixed(t *testing.T) { + env := testEnv(t) + sa := testSessionAgent(env, nil, nil, "test prompt") + agent := sa.(*sessionAgent) + + ctx := t.Context() + sess, err := env.sessions.Create(ctx, "test") + require.NoError(t, err) + + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.User, + Parts: []message.ContentPart{ + message.TextContent{Text: "hello"}, + }, + }) + require.NoError(t, err) + + // Assistant with 2 tool calls: one has a result, one is orphaned. + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.Assistant, + Parts: []message.ContentPart{ + message.ToolCall{ + ID: "call_ok", + Name: "view", + Input: `{"path":"/foo"}`, + Finished: true, + }, + message.ToolCall{ + ID: "call_orphaned", + Name: "agent", + Input: `{"prompt":"search"}`, + Finished: true, + }, + }, + }) + require.NoError(t, err) + + // Only one tool result — for call_ok. + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.Tool, + Parts: []message.ContentPart{ + message.ToolResult{ + ToolCallID: "call_ok", + Name: "view", + Content: "file contents", + }, + }, + }) + require.NoError(t, err) + + msgs, err := env.messages.List(ctx, sess.ID) + require.NoError(t, err) + + history, _ := agent.preparePrompt(msgs, true) + + // Should have a synthetic result only for the orphaned call. + var syntheticCount int + for _, msg := range history { + if msg.Role != fantasy.MessageRoleTool { + continue + } + for _, part := range msg.Content { + if tr, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](part); ok { + if tr.ToolCallID == "call_orphaned" { + syntheticCount++ + } + } + } + } + require.Equal(t, 1, syntheticCount, "expected exactly one synthetic result for the orphaned call") +} + +// requireToolCallAdjacency asserts that every assistant message in history +// with tool calls is immediately followed by a tool message that responds to +// each of those calls, with no other message in between. This is what +// strict-adjacency providers (e.g. Kimi, DeepSeek) require. +func requireToolCallAdjacency(t *testing.T, history []fantasy.Message) { + t.Helper() + for i, msg := range history { + if msg.Role != fantasy.MessageRoleAssistant { + continue + } + var callIDs []string + for _, part := range msg.Content { + if tc, ok := fantasy.AsMessagePart[fantasy.ToolCallPart](part); ok { + callIDs = append(callIDs, tc.ToolCallID) + } + } + if len(callIDs) == 0 { + continue + } + require.Less(t, i+1, len(history), "assistant with tool calls must be followed by a tool message") + next := history[i+1] + require.Equal(t, fantasy.MessageRoleTool, next.Role, + "assistant with tool calls %v must be immediately followed by a tool message, got %q", callIDs, next.Role) + responded := make(map[string]bool, len(callIDs)) + for _, part := range next.Content { + if tr, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](part); ok { + responded[tr.ToolCallID] = true + } + } + for _, id := range callIDs { + require.True(t, responded[id], + "tool result for call %q must immediately follow its assistant message", id) + } + } +} + +func TestPreparePrompt_NonAdjacentToolResults(t *testing.T) { + // A user message written between an assistant's tool call and its + // result (e.g. resuming while a tool is still running) must not end up + // between the two in the built history. + env := testEnv(t) + sa := testSessionAgent(env, nil, nil, "test prompt") + agent := sa.(*sessionAgent) + + ctx := t.Context() + sess, err := env.sessions.Create(ctx, "test") + require.NoError(t, err) + + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.User, + Parts: []message.ContentPart{ + message.TextContent{Text: "run commands"}, + }, + }) + require.NoError(t, err) + + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.Assistant, + Parts: []message.ContentPart{ + message.ToolCall{ + ID: "call_A", + Name: "bash", + Input: `{"command":"date"}`, + Finished: true, + }, + message.ToolCall{ + ID: "call_B", + Name: "bash", + Input: `{"command":"uptime"}`, + Finished: true, + }, + }, + }) + require.NoError(t, err) + + // Interleaved user message written while the tools were still running. + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.User, + Parts: []message.ContentPart{ + message.TextContent{Text: "are we done?"}, + }, + }) + require.NoError(t, err) + + // Results arrive late, after the interleaved user message. + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.Tool, + Parts: []message.ContentPart{ + message.ToolResult{ + ToolCallID: "call_A", + Name: "bash", + Content: "Fri May 2 21:00:00 UTC 2026", + }, + }, + }) + require.NoError(t, err) + + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.Tool, + Parts: []message.ContentPart{ + message.ToolResult{ + ToolCallID: "call_B", + Name: "bash", + Content: "21:00 up 3 days", + }, + }, + }) + require.NoError(t, err) + + msgs, err := env.messages.List(ctx, sess.ID) + require.NoError(t, err) + + require.Equal(t, message.User, msgs[2].Role, "interleaved user should be between assistant and results in DB order") + + history, _ := agent.preparePrompt(msgs, false) + + requireToolCallAdjacency(t, history) + + // The interleaved user message must still be present, after the results. + var foundInterleaved bool + for _, msg := range history { + if msg.Role != fantasy.MessageRoleUser { + continue + } + for _, part := range msg.Content { + if text, ok := fantasy.AsMessagePart[fantasy.TextPart](part); ok && text.Text == "are we done?" { + foundInterleaved = true + } + } + } + require.True(t, foundInterleaved, "interleaved user message must not be dropped") +} + +func TestPreparePrompt_ResultBeforeAssistant(t *testing.T) { + // A tool result written before its assistant message (e.g. concurrent + // writes) must be emitted after the assistant, exactly once, not at its + // stored position. + env := testEnv(t) + sa := testSessionAgent(env, nil, nil, "test prompt") + agent := sa.(*sessionAgent) + + ctx := t.Context() + sess, err := env.sessions.Create(ctx, "test") + require.NoError(t, err) + + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.Tool, + Parts: []message.ContentPart{ + message.ToolResult{ToolCallID: "call_X", Name: "bash", Content: "result"}, + }, + }) + require.NoError(t, err) + + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.Assistant, + Parts: []message.ContentPart{ + message.ToolCall{ID: "call_X", Name: "bash", Input: `{}`, Finished: true}, + }, + }) + require.NoError(t, err) + + msgs, err := env.messages.List(ctx, sess.ID) + require.NoError(t, err) + + history, _ := agent.preparePrompt(msgs, false) + + requireToolCallAdjacency(t, history) + + resultCount := 0 + for _, msg := range history { + if msg.Role != fantasy.MessageRoleTool { + continue + } + for _, part := range msg.Content { + if tr, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](part); ok && tr.ToolCallID == "call_X" { + resultCount++ + } + } + } + require.Equal(t, 1, resultCount, "result must be emitted exactly once") +} + +func TestPreparePrompt_BundledResultsAcrossAssistants(t *testing.T) { + // A single tool message can hold results for calls issued by different + // assistant messages. Each assistant must be followed by its own + // results, and no result may be emitted twice. + env := testEnv(t) + sa := testSessionAgent(env, nil, nil, "test prompt") + agent := sa.(*sessionAgent) + + ctx := t.Context() + sess, err := env.sessions.Create(ctx, "test") + require.NoError(t, err) + + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.Assistant, + Parts: []message.ContentPart{ + message.ToolCall{ID: "call_1", Name: "bash", Input: `{}`, Finished: true}, + }, + }) + require.NoError(t, err) + + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.User, + Parts: []message.ContentPart{ + message.TextContent{Text: "and now?"}, + }, + }) + require.NoError(t, err) + + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.Assistant, + Parts: []message.ContentPart{ + message.ToolCall{ID: "call_2", Name: "view", Input: `{"path":"/foo"}`, Finished: true}, + }, + }) + require.NoError(t, err) + + // Both results land in the same tool message, out of order. + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.Tool, + Parts: []message.ContentPart{ + message.ToolResult{ToolCallID: "call_2", Name: "view", Content: "file contents"}, + message.ToolResult{ToolCallID: "call_1", Name: "bash", Content: "output"}, + }, + }) + require.NoError(t, err) + + msgs, err := env.messages.List(ctx, sess.ID) + require.NoError(t, err) + + history, _ := agent.preparePrompt(msgs, false) + + requireToolCallAdjacency(t, history) + + counts := make(map[string]int) + for _, msg := range history { + if msg.Role != fantasy.MessageRoleTool { + continue + } + for _, part := range msg.Content { + if tr, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](part); ok { + counts[tr.ToolCallID]++ + } + } + } + require.Equal(t, map[string]int{"call_1": 1, "call_2": 1}, counts, + "each result must be emitted exactly once, next to its assistant") +} + +func TestPreparePrompt_DropsOrphanedToolResults(t *testing.T) { + // A tool result whose call is not in the history (e.g. the assistant + // message was cut off by a session summary) must be dropped instead of + // producing an unanswerable tool message. + env := testEnv(t) + sa := testSessionAgent(env, nil, nil, "test prompt") + agent := sa.(*sessionAgent) + + ctx := t.Context() + sess, err := env.sessions.Create(ctx, "test") + require.NoError(t, err) + + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.Tool, + Parts: []message.ContentPart{ + message.ToolResult{ToolCallID: "call_gone", Name: "bash", Content: "output"}, + }, + }) + require.NoError(t, err) + + _, err = env.messages.Create(ctx, sess.ID, message.CreateMessageParams{ + Role: message.User, + Parts: []message.ContentPart{ + message.TextContent{Text: "hello"}, + }, + }) + require.NoError(t, err) + + msgs, err := env.messages.List(ctx, sess.ID) + require.NoError(t, err) + + history, _ := agent.preparePrompt(msgs, false) + + for _, msg := range history { + require.NotEqual(t, fantasy.MessageRoleTool, msg.Role, "orphaned tool results must be dropped") + } +} + +func TestWorkaroundProviderMediaLimitations_TextOnlyModel(t *testing.T) { + env := testEnv(t) + sa := testSessionAgent(env, nil, nil, "test prompt") + agent := sa.(*sessionAgent) + + pngBase64 := base64.StdEncoding.EncodeToString([]byte("fake-png-data")) + + messages := []fantasy.Message{ + { + Role: fantasy.MessageRoleTool, + Content: []fantasy.MessagePart{ + fantasy.ToolResultPart{ + ToolCallID: "call_1", + Output: fantasy.ToolResultOutputContentMedia{ + Data: pngBase64, + MediaType: "image/png", + }, + }, + }, + }, + } + + // Non-Anthropic provider, no image support — should replace media with + // a text placeholder and not create a synthetic user message. + largeModel := Model{ + ModelCfg: config.SelectedModel{Provider: "openai"}, + CatwalkCfg: catwalk.Model{ + SupportsImages: false, + }, + } + + result := agent.workaroundProviderMediaLimitations(messages, largeModel) + + // Should produce exactly one message: the tool message with a text + // placeholder. No synthetic user message with FilePart. + require.Len(t, result, 1) + require.Equal(t, fantasy.MessageRoleTool, result[0].Role) + + tr, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](result[0].Content[0]) + require.True(t, ok) + _, ok = fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentText](tr.Output) + require.True(t, ok) +} + +func TestWorkaroundProviderMediaLimitations_VisionModel(t *testing.T) { + env := testEnv(t) + sa := testSessionAgent(env, nil, nil, "test prompt") + agent := sa.(*sessionAgent) + + pngBase64 := base64.StdEncoding.EncodeToString([]byte("fake-png-data")) + + messages := []fantasy.Message{ + { + Role: fantasy.MessageRoleTool, + Content: []fantasy.MessagePart{ + fantasy.ToolResultPart{ + ToolCallID: "call_1", + Output: fantasy.ToolResultOutputContentMedia{ + Data: pngBase64, + MediaType: "image/png", + }, + }, + }, + }, + } + + // Non-Anthropic provider, image support — should create a synthetic + // user message with FilePart. + largeModel := Model{ + ModelCfg: config.SelectedModel{Provider: "openai"}, + CatwalkCfg: catwalk.Model{ + SupportsImages: true, + }, + } + + result := agent.workaroundProviderMediaLimitations(messages, largeModel) + + // Should produce two messages: tool message with placeholder text, + // and synthetic user message with FilePart. + require.Len(t, result, 2) + require.Equal(t, fantasy.MessageRoleTool, result[0].Role) + require.Equal(t, fantasy.MessageRoleUser, result[1].Role) + + // The tool message should have text placeholder. + tr, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](result[0].Content[0]) + require.True(t, ok) + textOutput, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentText](tr.Output) + require.True(t, ok) + require.Contains(t, textOutput.Text, "see attached file") + + // The synthetic user message should contain a TextPart and a FilePart. + require.Len(t, result[1].Content, 2) + file, ok := fantasy.AsMessagePart[fantasy.FilePart](result[1].Content[1]) + require.True(t, ok) + require.Equal(t, "image/png", file.MediaType) +} + +func TestWorkaroundProviderMediaLimitations_AnthropicProvider(t *testing.T) { + env := testEnv(t) + sa := testSessionAgent(env, nil, nil, "test prompt") + agent := sa.(*sessionAgent) + + pngBase64 := base64.StdEncoding.EncodeToString([]byte("fake-png-data")) + + messages := []fantasy.Message{ + { + Role: fantasy.MessageRoleTool, + Content: []fantasy.MessagePart{ + fantasy.ToolResultPart{ + ToolCallID: "call_1", + Output: fantasy.ToolResultOutputContentMedia{ + Data: pngBase64, + MediaType: "image/png", + }, + }, + }, + }, + } + + // Anthropic provider — should return messages unchanged regardless of + // SupportsImages, since Anthropic handles media in tool results natively. + largeModel := Model{ + ModelCfg: config.SelectedModel{Provider: string(catwalk.InferenceProviderAnthropic)}, + CatwalkCfg: catwalk.Model{ + SupportsImages: true, + }, + } + + result := agent.workaroundProviderMediaLimitations(messages, largeModel) + require.Len(t, result, 1) + require.Equal(t, fantasy.MessageRoleTool, result[0].Role) + + // The media should still be in the tool result, untouched. + tr, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](result[0].Content[0]) + require.True(t, ok) + media, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentMedia](tr.Output) + require.True(t, ok) + require.Equal(t, "image/png", media.MediaType) +} + +func TestProviderRetryLogFields(t *testing.T) { + t.Run("nil provider error", func(t *testing.T) { + fields := providerRetryLogFields(nil, 2*time.Second) + require.Equal(t, []any{"retry_delay", "2s"}, fields) + }) + + t.Run("provider error with title and message", func(t *testing.T) { + fields := providerRetryLogFields(&fantasy.ProviderError{ + StatusCode: 429, + Title: "rate limit", + Message: "too many requests", + }, 1500*time.Millisecond) + require.Equal(t, []any{ + "retry_delay", "1.5s", + "status_code", 429, + "title", "rate limit", + "message", "too many requests", + }, fields) + }) + + t.Run("provider error without optional strings", func(t *testing.T) { + fields := providerRetryLogFields(&fantasy.ProviderError{ + StatusCode: 503, + }, time.Second) + require.Equal(t, []any{ + "retry_delay", "1s", + "status_code", 503, + }, fields) + }) +} diff --git a/internal/agent/agent_tool.go b/internal/agent/agent_tool.go new file mode 100644 index 0000000000000000000000000000000000000000..5be58c842ed30c5a2b9f713d5840997e94688dc3 --- /dev/null +++ b/internal/agent/agent_tool.go @@ -0,0 +1,68 @@ +package agent + +import ( + "context" + _ "embed" + "errors" + + "charm.land/fantasy" + + "github.com/charmbracelet/crush/internal/agent/prompt" + "github.com/charmbracelet/crush/internal/agent/tools" + "github.com/charmbracelet/crush/internal/config" +) + +//go:embed templates/agent_tool.md +var agentToolDescription string + +type AgentParams struct { + Prompt string `json:"prompt" description:"The task for the agent to perform"` +} + +const ( + AgentToolName = "agent" +) + +func (c *coordinator) agentTool(ctx context.Context) (fantasy.AgentTool, error) { + agentCfg, ok := c.cfg.Config().Agents[config.AgentTask] + if !ok { + return nil, errors.New("task agent not configured") + } + prompt, err := taskPrompt(prompt.WithWorkingDir(c.cfg.WorkingDir())) + if err != nil { + return nil, err + } + + agent, err := c.buildAgent(ctx, prompt, agentCfg, true) + if err != nil { + return nil, err + } + return fantasy.NewParallelAgentTool( + AgentToolName, + agentToolDescription, + func(ctx context.Context, params AgentParams, call fantasy.ToolCall) (fantasy.ToolResponse, error) { + if params.Prompt == "" { + return fantasy.NewTextErrorResponse("prompt is required"), nil + } + + sessionID := tools.GetSessionFromContext(ctx) + if sessionID == "" { + return fantasy.ToolResponse{}, errors.New("session id missing from context") + } + + agentMessageID := tools.GetMessageFromContext(ctx) + if agentMessageID == "" { + return fantasy.ToolResponse{}, errors.New("agent message id missing from context") + } + + return c.runSubAgent(ctx, subAgentParams{ + Agent: agent, + SessionID: sessionID, + AgentMessageID: agentMessageID, + ToolCallID: call.ID, + Prompt: params.Prompt, + SessionTitle: "New Agent Session", + }) + }, + ), nil +} diff --git a/internal/agent/agentic_fetch_tool.go b/internal/agent/agentic_fetch_tool.go new file mode 100644 index 0000000000000000000000000000000000000000..72e61705a004aae695fb54252bb34bc69a943985 --- /dev/null +++ b/internal/agent/agentic_fetch_tool.go @@ -0,0 +1,206 @@ +package agent + +import ( + "context" + _ "embed" + "errors" + "fmt" + "net/http" + "os" + "time" + + "charm.land/fantasy" + + "github.com/charmbracelet/crush/internal/agent/prompt" + "github.com/charmbracelet/crush/internal/agent/tools" + "github.com/charmbracelet/crush/internal/permission" +) + +//go:embed templates/agentic_fetch.md +var agenticFetchToolDescription string + +// agenticFetchValidationResult holds the validated parameters from the tool call context. +type agenticFetchValidationResult struct { + SessionID string + AgentMessageID string +} + +// validateAgenticFetchParams validates the tool call parameters and extracts required context values. +func validateAgenticFetchParams(ctx context.Context, params tools.AgenticFetchParams) (agenticFetchValidationResult, error) { + if params.Prompt == "" { + return agenticFetchValidationResult{}, errors.New("prompt is required") + } + + sessionID := tools.GetSessionFromContext(ctx) + if sessionID == "" { + return agenticFetchValidationResult{}, errors.New("session id missing from context") + } + + agentMessageID := tools.GetMessageFromContext(ctx) + if agentMessageID == "" { + return agenticFetchValidationResult{}, errors.New("agent message id missing from context") + } + + return agenticFetchValidationResult{ + SessionID: sessionID, + AgentMessageID: agentMessageID, + }, nil +} + +//go:embed templates/agentic_fetch_prompt.md.tpl +var agenticFetchPromptTmpl []byte + +func (c *coordinator) agenticFetchTool(_ context.Context, client *http.Client) (fantasy.AgentTool, error) { + if client == nil { + transport := http.DefaultTransport.(*http.Transport).Clone() + transport.MaxIdleConns = 100 + transport.MaxIdleConnsPerHost = 10 + transport.IdleConnTimeout = 90 * time.Second + + client = &http.Client{ + Timeout: 30 * time.Second, + Transport: transport, + } + } + + return fantasy.NewParallelAgentTool( + tools.AgenticFetchToolName, + agenticFetchToolDescription, + func(ctx context.Context, params tools.AgenticFetchParams, call fantasy.ToolCall) (fantasy.ToolResponse, error) { + validationResult, err := validateAgenticFetchParams(ctx, params) + if err != nil { + return fantasy.NewTextErrorResponse(err.Error()), nil + } + + // Determine description based on mode. + var description string + if params.URL != "" { + description = fmt.Sprintf("Fetch and analyze content from URL: %s", params.URL) + } else { + description = "Search the web and analyze results" + } + + p, err := c.permissions.Request( + ctx, + permission.CreatePermissionRequest{ + SessionID: validationResult.SessionID, + Path: c.cfg.WorkingDir(), + ToolCallID: call.ID, + ToolName: tools.AgenticFetchToolName, + Action: "fetch", + Description: description, + Params: tools.AgenticFetchPermissionsParams(params), + }, + ) + if err != nil { + return fantasy.ToolResponse{}, err + } + if !p { + return tools.NewPermissionDeniedResponse(), nil + } + + tmpDir, err := os.MkdirTemp(c.cfg.Config().Options.DataDirectory, "crush-fetch-*") + if err != nil { + return fantasy.NewTextErrorResponse(fmt.Sprintf("Failed to create temporary directory: %s", err)), nil + } + defer os.RemoveAll(tmpDir) + + var fullPrompt string + + if params.URL != "" { + // URL mode: fetch the URL content first. + content, err := tools.FetchURLAndConvert(ctx, client, params.URL) + if err != nil { + return fantasy.NewTextErrorResponse(fmt.Sprintf("Failed to fetch URL: %s", err)), nil + } + + hasLargeContent := len(content) > tools.LargeContentThreshold + + if hasLargeContent { + tempFile, err := os.CreateTemp(tmpDir, "page-*.md") + if err != nil { + return fantasy.NewTextErrorResponse(fmt.Sprintf("Failed to create temporary file: %s", err)), nil + } + tempFilePath := tempFile.Name() + + if _, err := tempFile.WriteString(content); err != nil { + tempFile.Close() + return fantasy.NewTextErrorResponse(fmt.Sprintf("Failed to write content to file: %s", err)), nil + } + tempFile.Close() + + fullPrompt = fmt.Sprintf("%s\n\nThe web page from %s has been saved to: %s\n\nUse the view and grep tools to analyze this file and extract the requested information.", params.Prompt, params.URL, tempFilePath) + } else { + fullPrompt = fmt.Sprintf("%s\n\nWeb page URL: %s\n\n\n%s\n", params.Prompt, params.URL, content) + } + } else { + // Search mode: let the sub-agent search and fetch as needed. + fullPrompt = fmt.Sprintf("%s\n\nUse the web_search tool to find relevant information. Break down the question into smaller, focused searches if needed. After searching, use web_fetch to get detailed content from the most relevant results.", params.Prompt) + } + + promptOpts := []prompt.Option{ + prompt.WithWorkingDir(tmpDir), + } + + promptTemplate, err := prompt.NewPrompt("agentic_fetch", string(agenticFetchPromptTmpl), promptOpts...) + if err != nil { + return fantasy.ToolResponse{}, fmt.Errorf("error creating prompt: %s", err) + } + + _, small, err := c.buildAgentModels(ctx, true) + if err != nil { + return fantasy.ToolResponse{}, fmt.Errorf("error building models: %s", err) + } + + systemPrompt, err := promptTemplate.Build(ctx, small.Model.Provider(), small.Model.Model(), c.cfg) + if err != nil { + return fantasy.ToolResponse{}, fmt.Errorf("error building system prompt: %s", err) + } + + smallProviderCfg, ok := c.cfg.Config().Providers.Get(small.ModelCfg.Provider) + if !ok { + return fantasy.ToolResponse{}, errors.New("small model provider not configured") + } + + webFetchTool := tools.NewWebFetchTool(tmpDir, client) + webSearchTool := tools.NewWebSearchTool(client) + fetchTools := []fantasy.AgentTool{ + webFetchTool, + webSearchTool, + tools.NewGlobTool(tmpDir, c.cfg.Config().Tools.Glob), + tools.NewGrepTool(tmpDir, c.cfg.Config().Tools.Grep), + tools.NewSourcegraphTool(client), + tools.NewViewTool(c.lspManager, c.permissions, c.filetracker, nil, tmpDir), + } + + // Sub-agent tools run without hook interception. The top-level + // `agentic_fetch` call itself is already wrapped from the coder's + // side; firing hooks again for every inner tool call would run + // the user's hooks N times per delegated turn. + + agent := NewSessionAgent(SessionAgentOptions{ + LargeModel: small, // Use small model for both (fetch doesn't need large) + SmallModel: small, + SystemPromptPrefix: smallProviderCfg.SystemPromptPrefix, + SystemPrompt: systemPrompt, + DisableAutoSummarize: c.cfg.Config().Options.DisableAutoSummarize, + IsYolo: c.permissions.SkipRequests(), + Sessions: c.sessions, + Messages: c.messages, + Tools: fetchTools, + }) + + return c.runSubAgent(ctx, subAgentParams{ + Agent: agent, + SessionID: validationResult.SessionID, + AgentMessageID: validationResult.AgentMessageID, + ToolCallID: call.ID, + Prompt: fullPrompt, + SessionTitle: "Fetch Analysis", + SessionSetup: func(sessionID string) { + c.permissions.AutoApproveSession(sessionID) + }, + }) + }, + ), nil +} diff --git a/internal/agent/aws_sso_refresh.go b/internal/agent/aws_sso_refresh.go new file mode 100644 index 0000000000000000000000000000000000000000..ecc60cebb6aa045a8355f364048528fe78a3f07d --- /dev/null +++ b/internal/agent/aws_sso_refresh.go @@ -0,0 +1,158 @@ +package agent + +import ( + "bufio" + "bytes" + "context" + "fmt" + "io" + "log/slog" + "os/exec" + "regexp" + "strings" + "sync" + "time" + + "github.com/charmbracelet/crush/internal/agent/notify" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/pubsub" +) + +// awsSSORefreshTimeout bounds how long the AWS SSO refresh command may run. +// Browser-based SSO needs time, so it is generous, and it runs on a context +// detached from the agent turn so a cancelled turn doesn't abort the login. +const awsSSORefreshTimeout = 5 * time.Minute + +// awsSSOURLRe matches the https verification URL that `aws sso login` and +// related commands print to stdout or stderr. +var awsSSOURLRe = regexp.MustCompile(`https://[^\s]+`) + +// extractAWSSSOURL returns the first HTTPS URL in the given command output +// line, or empty if none is present. +func extractAWSSSOURL(line string) string { + return awsSSOURLRe.FindString(line) +} + +// refreshAWSCredentials runs the provider's configured AWS SSO refresh +// command (e.g. "aws sso login") on the machine that makes the Bedrock +// calls, streaming the verification URL to the UI for display, then rebuilds +// models so the AWS SDK re-reads the refreshed credentials. It returns nil to +// signal that the failed request should be retried. +// +// The command runs here, in the coordinator, rather than in the UI dialog so +// the refreshed credentials land where the model calls are made. This is +// correct in both single-process and client/server deployments. +func (c *coordinator) refreshAWSCredentials(ctx context.Context, providerCfg config.ProviderConfig) error { + if c.notify == nil { + return errNoInteractiveAuth + } + slog.Info("AWS credentials expired, running refresh command", + "provider", providerCfg.ID, "command", providerCfg.AWSAuthRefresh) + + // Open the dialog immediately so the user sees progress even before the + // command prints its verification URL. + c.notify.Publish(pubsub.CreatedEvent, notify.Notification{ + Type: notify.TypeAWSSSOAuth, + ProviderID: providerCfg.ID, + AWSSOCommand: providerCfg.AWSAuthRefresh, + }) + + // Detach from the turn's context (with a generous timeout) so cancelling + // the turn doesn't kill an in-progress browser login. + runCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), awsSSORefreshTimeout) + defer cancel() + + runErr := c.runAWSAuthRefresh(runCtx, providerCfg) + + result := notify.Notification{Type: notify.TypeAWSSSOAuthResult, ProviderID: providerCfg.ID} + if runErr != nil { + result.Message = runErr.Error() + } + c.notify.Publish(pubsub.CreatedEvent, result) + + if runErr != nil { + slog.Error("AWS SSO refresh command failed", "provider", providerCfg.ID, "error", runErr) + return runErr + } + // If the turn's context was cancelled while the command ran, fantasy's + // retry would fail immediately, so surface the cancellation instead. + if ctx.Err() != nil { + return ctx.Err() + } + // Rebuild models so the AWS SDK credential chain re-reads the refreshed + // SSO cache on the next attempt. + if err := c.UpdateModels(runCtx); err != nil { + slog.Error("Failed to update models after AWS SSO refresh", "provider", providerCfg.ID, "error", err) + return err + } + slog.Info("AWS SSO refresh complete, retrying request", "provider", providerCfg.ID) + return nil +} + +// runAWSAuthRefresh executes the refresh command, publishing the SSO +// verification URL to the UI as soon as it appears in the output and +// returning any failure with captured stderr for context. +func (c *coordinator) runAWSAuthRefresh(ctx context.Context, providerCfg config.ProviderConfig) error { + cmd := exec.CommandContext(ctx, "sh", "-c", providerCfg.AWSAuthRefresh) + cmd.Dir = c.cfg.WorkingDir() + + stdout, err := cmd.StdoutPipe() + if err != nil { + return err + } + stderrPipe, err := cmd.StderrPipe() + if err != nil { + return err + } + if err := cmd.Start(); err != nil { + return err + } + + // Drain stdout and stderr concurrently so a command that fills one pipe + // buffer before closing the other can't deadlock. Both are scanned for + // the verification URL; stderr is also captured for error detail. + var ( + stderrBuf bytes.Buffer + mu sync.Mutex // Guards the single-shot URL publish across goroutines. + urlSent bool + ) + publishURL := func(line string) { + mu.Lock() + defer mu.Unlock() + if urlSent { + return + } + if url := extractAWSSSOURL(line); url != "" { + urlSent = true + // Second phase of the two-part publish: the dialog is already + // open from refreshAWSCredentials; this fills in the URL on it. + c.notify.Publish(pubsub.CreatedEvent, notify.Notification{ + Type: notify.TypeAWSSSOAuth, + ProviderID: providerCfg.ID, + AWSSOCommand: providerCfg.AWSAuthRefresh, + AWSSOURL: url, + }) + } + } + + var wg sync.WaitGroup + scan := func(r io.Reader) { + defer wg.Done() + scanner := bufio.NewScanner(r) + for scanner.Scan() { + publishURL(scanner.Text()) + } + } + wg.Add(2) + go scan(stdout) + go scan(io.TeeReader(stderrPipe, &stderrBuf)) + wg.Wait() + + if err := cmd.Wait(); err != nil { + if stderr := strings.TrimSpace(stderrBuf.String()); stderr != "" { + return fmt.Errorf("%w: %s", err, stderr) + } + return err + } + return nil +} diff --git a/internal/agent/aws_sso_refresh_test.go b/internal/agent/aws_sso_refresh_test.go new file mode 100644 index 0000000000000000000000000000000000000000..e66930fb60e68a41e2e62cb0072a196618a305e7 --- /dev/null +++ b/internal/agent/aws_sso_refresh_test.go @@ -0,0 +1,44 @@ +package agent + +import "testing" + +func TestExtractAWSSSOURL(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + output string + want string + }{ + { + name: "standard aws sso login output", + output: `If the browser does not open or you wish to use a different device to authorize this request, open the following URL: +https://device.sso.us-east-1.amazonaws.com/?user_code=ABCD-EFGH`, + want: "https://device.sso.us-east-1.amazonaws.com/?user_code=ABCD-EFGH", + }, + { + name: "url only", + output: "https://device.sso.eu-west-1.amazonaws.com/?user_code=XXXX-YYYY", + want: "https://device.sso.eu-west-1.amazonaws.com/?user_code=XXXX-YYYY", + }, + { + name: "no url", + output: "SSO session expired. Please run aws sso login.", + want: "", + }, + { + name: "empty output", + output: "", + want: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := extractAWSSSOURL(tt.output); got != tt.want { + t.Errorf("extractAWSSSOURL() = %q, want %q", got, tt.want) + } + }) + } +} diff --git a/internal/agent/common_test.go b/internal/agent/common_test.go new file mode 100644 index 0000000000000000000000000000000000000000..253b964f4a661766140c2b37fc8c710db4175c26 --- /dev/null +++ b/internal/agent/common_test.go @@ -0,0 +1,206 @@ +package agent + +import ( + "context" + "net/http" + "os" + "path/filepath" + "testing" + "time" + + "charm.land/catwalk/pkg/catwalk" + "charm.land/fantasy" + "charm.land/fantasy/providers/openaicompat" + "charm.land/x/vcr" + "github.com/charmbracelet/crush/internal/agent/prompt" + "github.com/charmbracelet/crush/internal/agent/tools" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/csync" + "github.com/charmbracelet/crush/internal/db" + "github.com/charmbracelet/crush/internal/filetracker" + "github.com/charmbracelet/crush/internal/history" + "github.com/charmbracelet/crush/internal/lsp" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/permission" + "github.com/charmbracelet/crush/internal/session" + "github.com/stretchr/testify/require" + + _ "github.com/joho/godotenv/autoload" +) + +// fakeEnv is an environment for testing. +type fakeEnv struct { + workingDir string + sessions session.Service + messages message.Service + permissions permission.Service + history history.Service + filetracker *filetracker.Service + lspClients *csync.Map[string, *lsp.Client] +} + +type builderFunc func(t *testing.T, r *vcr.Recorder) (fantasy.LanguageModel, error) + +type modelPair struct { + name string + largeModel builderFunc + smallModel builderFunc +} + +func hyperBuilder(model string) builderFunc { + return func(t *testing.T, r *vcr.Recorder) (fantasy.LanguageModel, error) { + provider, err := openaicompat.New( + openaicompat.WithBaseURL("https://hyper.charm.land/v1"), + openaicompat.WithAPIKey(os.Getenv("CRUSH_HYPER_API_KEY")), + openaicompat.WithHTTPClient(&http.Client{Transport: r}), + ) + if err != nil { + return nil, err + } + return provider.LanguageModel(t.Context(), model) + } +} + +func testEnv(t *testing.T) fakeEnv { + workingDir := filepath.Join("/tmp/crush-test/", t.Name()) + os.RemoveAll(workingDir) + + err := os.MkdirAll(workingDir, 0o755) + require.NoError(t, err) + + conn, err := db.Connect(t.Context(), t.TempDir()) + require.NoError(t, err) + + q := db.New(conn) + sessions := session.NewService(q, conn) + messages := message.NewService(q) + + permissions := permission.NewPermissionService(workingDir, true, []string{}) + history := history.NewService(q, conn) + filetrackerService := filetracker.NewService(q) + lspClients := csync.NewMap[string, *lsp.Client]() + + t.Cleanup(func() { + conn.Close() + os.RemoveAll(workingDir) + }) + + return fakeEnv{ + workingDir, + sessions, + messages, + permissions, + history, + &filetrackerService, + lspClients, + } +} + +func testSessionAgent(env fakeEnv, large, small fantasy.LanguageModel, systemPrompt string, tools ...fantasy.AgentTool) SessionAgent { + largeModel := Model{ + Model: large, + CatwalkCfg: catwalk.Model{ + ContextWindow: 200000, + DefaultMaxTokens: 10000, + }, + } + smallModel := Model{ + Model: small, + CatwalkCfg: catwalk.Model{ + ContextWindow: 200000, + DefaultMaxTokens: 10000, + }, + } + agent := NewSessionAgent(SessionAgentOptions{ + LargeModel: largeModel, + SmallModel: smallModel, + SystemPrompt: systemPrompt, + IsYolo: true, + Sessions: env.sessions, + Messages: env.messages, + Tools: tools, + }) + return agent +} + +func coderAgent(r *vcr.Recorder, env fakeEnv, large, small fantasy.LanguageModel) (SessionAgent, error) { + fixedTime := func() time.Time { + t, _ := time.Parse("1/2/2006", "1/1/2025") + return t + } + prompt, err := coderPrompt( + prompt.WithTimeFunc(fixedTime), + prompt.WithPlatform("linux"), + prompt.WithWorkingDir(filepath.ToSlash(env.workingDir)), + ) + if err != nil { + return nil, err + } + cfg, err := config.Init(env.workingDir, "", false) + if err != nil { + return nil, err + } + + // NOTE(@andreynering): Set a fixed config to ensure cassettes match + // independently of user config on `$HOME/.config/crush/crush.json`. + cfg.Config().Options.Attribution = &config.Attribution{ + TrailerStyle: "co-authored-by", + GeneratedWith: true, + } + + // Clear some fields to avoid issues with VCR cassette matching. + cfg.Config().Options.SkillsPaths = nil + cfg.Config().Options.DisabledSkills = []string{"crush-config"} + cfg.Config().Options.ContextPaths = nil + cfg.Config().Options.GlobalContextPaths = nil + cfg.Config().LSP = nil + + systemPrompt, err := prompt.Build(context.TODO(), large.Provider(), large.Model(), cfg) + if err != nil { + return nil, err + } + + // Get the model name for the bash tool + modelName := large.Model() // fallback to ID if Name not available + if model := cfg.Config().GetModel(large.Provider(), large.Model()); model != nil { + modelName = model.Name + } + + allTools := []fantasy.AgentTool{ + tools.NewBashTool(env.permissions, env.workingDir, cfg.Config().Options.Attribution, modelName), + tools.NewDownloadTool(env.permissions, env.workingDir, r.GetDefaultClient()), + tools.NewEditTool(nil, env.permissions, env.history, *env.filetracker, env.workingDir), + tools.NewMultiEditTool(nil, env.permissions, env.history, *env.filetracker, env.workingDir), + tools.NewFetchTool(env.permissions, env.workingDir, r.GetDefaultClient()), + tools.NewGlobTool(env.workingDir, cfg.Config().Tools.Glob), + tools.NewGrepTool(env.workingDir, cfg.Config().Tools.Grep), + tools.NewLsTool(env.permissions, env.workingDir, cfg.Config().Tools.Ls), + tools.NewSourcegraphTool(r.GetDefaultClient()), + tools.NewViewTool(nil, env.permissions, *env.filetracker, nil, env.workingDir), + tools.NewWriteTool(nil, env.permissions, env.history, *env.filetracker, env.workingDir), + } + + return testSessionAgent(env, large, small, systemPrompt, allTools...), nil +} + +// createSimpleGoProject creates a simple Go project structure in the given directory. +// It creates a go.mod file and a main.go file with a basic hello world program. +func createSimpleGoProject(t *testing.T, dir string) { + goMod := `module example.com/testproject + +go 1.23 +` + err := os.WriteFile(dir+"/go.mod", []byte(goMod), 0o644) + require.NoError(t, err) + + mainGo := `package main + +import "fmt" + +func main() { + fmt.Println("Hello, World!") +} +` + err = os.WriteFile(dir+"/main.go", []byte(mainGo), 0o644) + require.NoError(t, err) +} diff --git a/internal/agent/convert_test.go b/internal/agent/convert_test.go new file mode 100644 index 0000000000000000000000000000000000000000..4aa382e7cf0015c825c806678908ff97fc8828dd --- /dev/null +++ b/internal/agent/convert_test.go @@ -0,0 +1,94 @@ +package agent + +import ( + "encoding/base64" + "testing" + + "charm.land/fantasy" + "github.com/stretchr/testify/require" +) + +func TestConvertToToolResult_InvalidBase64(t *testing.T) { + t.Parallel() + + a := &sessionAgent{} + result := fantasy.ToolResultContent{ + ToolCallID: "call_123", + ToolName: "test_tool", + Result: fantasy.ToolResultOutputContentMedia{ + Data: "abc\x80def", + MediaType: "image/png", + }, + } + + tr := a.convertToToolResult(result) + require.True(t, tr.IsError) + require.Empty(t, tr.Data) + require.Contains(t, tr.Content, "invalid encoding") + require.Equal(t, "call_123", tr.ToolCallID) + require.Equal(t, "test_tool", tr.Name) +} + +func TestConvertToToolResult_ValidMedia(t *testing.T) { + t.Parallel() + + a := &sessionAgent{} + validData := base64.StdEncoding.EncodeToString([]byte("test image data")) + + result := fantasy.ToolResultContent{ + ToolCallID: "call_456", + ToolName: "screenshot", + Result: fantasy.ToolResultOutputContentMedia{ + Data: validData, + MediaType: "image/png", + Text: "Screenshot captured", + }, + } + + tr := a.convertToToolResult(result) + require.False(t, tr.IsError) + require.Equal(t, validData, tr.Data) + require.Equal(t, "image/png", tr.MIMEType) + require.Equal(t, "Screenshot captured", tr.Content) +} + +func TestConvertToToolResult_ValidMediaNoText(t *testing.T) { + t.Parallel() + + a := &sessionAgent{} + validData := base64.StdEncoding.EncodeToString([]byte("test image data")) + + result := fantasy.ToolResultContent{ + ToolCallID: "call_789", + ToolName: "view", + Result: fantasy.ToolResultOutputContentMedia{ + Data: validData, + MediaType: "image/jpeg", + }, + } + + tr := a.convertToToolResult(result) + require.False(t, tr.IsError) + require.Equal(t, validData, tr.Data) + require.Equal(t, "image/jpeg", tr.MIMEType) + require.Equal(t, "Loaded image/jpeg content", tr.Content) +} + +func TestConvertToToolResult_ASCIIButInvalidBase64(t *testing.T) { + t.Parallel() + + a := &sessionAgent{} + result := fantasy.ToolResultContent{ + ToolCallID: "call_abc", + ToolName: "mcp_tool", + Result: fantasy.ToolResultOutputContentMedia{ + Data: "not-valid-base64!!!", + MediaType: "image/png", + }, + } + + tr := a.convertToToolResult(result) + require.True(t, tr.IsError) + require.Empty(t, tr.Data) + require.Contains(t, tr.Content, "invalid encoding") +} diff --git a/internal/agent/coordinator.go b/internal/agent/coordinator.go new file mode 100644 index 0000000000000000000000000000000000000000..d38b03aa00225ccf53171884bd5b770a329928ee --- /dev/null +++ b/internal/agent/coordinator.go @@ -0,0 +1,1757 @@ +package agent + +import ( + "bytes" + "cmp" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "maps" + "net/http" + "os" + "path/filepath" + "slices" + "strings" + "time" + + "charm.land/catwalk/pkg/catwalk" + "charm.land/fantasy" + "github.com/charmbracelet/crush/internal/agent/hyper" + "github.com/charmbracelet/crush/internal/agent/notify" + "github.com/charmbracelet/crush/internal/agent/prompt" + "github.com/charmbracelet/crush/internal/agent/tools" + "github.com/charmbracelet/crush/internal/agent/tools/mcp" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/discover" + "github.com/charmbracelet/crush/internal/event" + "github.com/charmbracelet/crush/internal/filetracker" + "github.com/charmbracelet/crush/internal/history" + "github.com/charmbracelet/crush/internal/hooks" + "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/oauth/copilot" + openaioauth "github.com/charmbracelet/crush/internal/oauth/openai" + "github.com/charmbracelet/crush/internal/permission" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/charmbracelet/crush/internal/question" + "github.com/charmbracelet/crush/internal/session" + "github.com/charmbracelet/crush/internal/skills" + "golang.org/x/sync/errgroup" + + "charm.land/fantasy/providers/anthropic" + "charm.land/fantasy/providers/azure" + "charm.land/fantasy/providers/bedrock" + "charm.land/fantasy/providers/google" + "charm.land/fantasy/providers/openai" + "charm.land/fantasy/providers/openaicompat" + "charm.land/fantasy/providers/openrouter" + "charm.land/fantasy/providers/vercel" + openaisdk "github.com/openai/openai-go/v3/option" + "github.com/qjebbs/go-jsons" +) + +// Coordinator errors. +var ( + errCoderAgentNotConfigured = errors.New("coder agent not configured") + errModelProviderNotConfigured = errors.New("model provider not configured") + errLargeModelNotSelected = errors.New("large model not selected") + errSmallModelNotSelected = errors.New("small model not selected") + errLargeModelProviderNotConfigured = errors.New("large model provider not configured") + errSmallModelProviderNotConfigured = errors.New("small model provider not configured") + errLargeModelNotFound = errors.New("large model not found in provider config") + errSmallModelNotFound = errors.New("small model not found in provider config") +) + +// Copilot models that use the Responses API instead of Chat Completions. +var copilotResponsesModels = map[string]bool{ + "gpt-5.2": true, + "gpt-5.2-codex": true, + "gpt-5.3-codex": true, + "gpt-5.4": true, + "gpt-5.4-mini": true, + "gpt-5.5": true, + "gpt-5-mini": true, + "gpt-5.6-luna": true, + "gpt-5.6-terra": true, + "gpt-5.6-sol": true, + "gpt-6-astra": true, + "grok-4.5": true, + "grok-4.6": true, +} + +// OpenCode models that use the Anthropic Messages API instead of Chat +// Completions. Which endpoint serves each model differs per provider, see +// https://opencode.ai/docs/zen and https://opencode.ai/docs/go. +func isOpenCodeMessagesModel(providerID, modelID string) bool { + switch providerID { + case string(catwalk.InferenceProviderOpenCodeGo): + return strings.HasPrefix(modelID, "minimax-") || + strings.HasPrefix(modelID, "qwen3.6-") || + strings.HasPrefix(modelID, "qwen3.7-") || + strings.HasPrefix(modelID, "qwen3.8-") + case string(catwalk.InferenceProviderOpenCodeZen): + return strings.HasPrefix(modelID, "claude-") || + strings.HasPrefix(modelID, "qwen3.5-") || + strings.HasPrefix(modelID, "qwen3.6-") || + strings.HasPrefix(modelID, "qwen3.7-") || + strings.HasPrefix(modelID, "qwen3.8-") + } + return false +} + +// OpenCode models that use the OpenAI Responses API instead of Chat +// Completions. See https://opencode.ai/docs/zen and https://opencode.ai/docs/go. +func isOpenCodeResponsesModel(modelID string) bool { + return strings.HasPrefix(modelID, "gpt-") || + strings.HasPrefix(modelID, "grok-") || + strings.HasPrefix(modelID, "muse-spark-") +} + +type Coordinator interface { + // INFO: (kujtim) this is not used yet we will use this when we have multiple agents + // SetMainAgent(string) + Run(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) + // RunAccepted runs a call that was already accepted via + // BeginAccepted on the fire-and-forget dispatch path. The handle is + // the only carrier of accept-state across the backend.runAgent / + // Coordinator / sessionAgent.Run layers: it reaches + // sessionAgent.Run as SessionAgentCall.Accepted, where it is + // consumed under dispatchMu once the accepted -> (cancel-on-entry | + // queued | active) transition is chosen. + RunAccepted(ctx context.Context, accept *AcceptedRun, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) + BeginAccepted(sessionID string) *AcceptedRun + Cancel(sessionID string) + CancelAll() + IsSessionBusy(sessionID string) bool + IsBusy() bool + QueuedPrompts(sessionID string) int + QueuedPromptsList(sessionID string) []string + ClearQueue(sessionID string) + Summarize(context.Context, string) error + Model() Model + UpdateModels(ctx context.Context) error + GenerateTitle(ctx context.Context, sessionID, prompt string) +} + +type coordinator struct { + cfg *config.ConfigStore + sessions session.Service + messages message.Service + permissions permission.Service + questions question.Service + history history.Service + filetracker filetracker.Service + lspManager *lsp.Manager + notify pubsub.Publisher[notify.Notification] + runComplete pubsub.Publisher[notify.RunComplete] + interactive bool + + currentAgent SessionAgent + agents map[string]SessionAgent + + // Skills discovery results (session-start snapshot). + allSkills []*skills.Skill // Pre-filter: all discovered after dedup. + activeSkills []*skills.Skill // Post-filter: active skills only. + skillTracker *skills.Tracker + + readyWg errgroup.Group +} + +// CoordinatorOptions holds the dependencies for NewCoordinator. Using a +// struct keeps the constructor self-documenting and avoids a long +// positional parameter list. +type CoordinatorOptions struct { + Config *config.ConfigStore + Sessions session.Service + Messages message.Service + Permissions permission.Service + Questions question.Service + History history.Service + FileTracker filetracker.Service + LSPManager *lsp.Manager + Notify pubsub.Publisher[notify.Notification] + RunComplete pubsub.Publisher[notify.RunComplete] + Skills *skills.Manager + Interactive bool +} + +func NewCoordinator(ctx context.Context, opts CoordinatorOptions) (Coordinator, error) { + // Skills are pre-discovered by the caller (see app.New / + // backend.CreateWorkspace) and passed in via the manager. If no + // manager was provided (legacy callers), fall back to an in-line + // discovery so the coordinator still works. + var allSkills, activeSkills []*skills.Skill + if opts.Skills != nil { + allSkills = opts.Skills.AllSkills() + activeSkills = opts.Skills.ActiveSkills() + } else { + allSkills, activeSkills = discoverSkills(opts.Config) + } + skillTracker := skills.NewTracker(activeSkills) + + c := &coordinator{ + cfg: opts.Config, + sessions: opts.Sessions, + messages: opts.Messages, + permissions: opts.Permissions, + questions: opts.Questions, + history: opts.History, + filetracker: opts.FileTracker, + lspManager: opts.LSPManager, + notify: opts.Notify, + runComplete: opts.RunComplete, + agents: make(map[string]SessionAgent), + allSkills: allSkills, + activeSkills: activeSkills, + skillTracker: skillTracker, + interactive: opts.Interactive, + } + + agentCfg, ok := opts.Config.Config().Agents[config.AgentCoder] + if !ok { + return nil, errCoderAgentNotConfigured + } + + // TODO: make this dynamic when we support multiple agents + prompt, err := coderPrompt(prompt.WithWorkingDir(c.cfg.WorkingDir())) + if err != nil { + return nil, err + } + + agent, err := c.buildAgent(ctx, prompt, agentCfg, false) + if err != nil { + return nil, err + } + c.currentAgent = agent + c.agents[config.AgentCoder] = agent + return c, nil +} + +// Run implements Coordinator. +func (c *coordinator) Run(ctx context.Context, sessionID string, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + return c.run(ctx, nil, sessionID, prompt, attachments...) +} + +// RunAccepted implements Coordinator. +func (c *coordinator) RunAccepted(ctx context.Context, accept *AcceptedRun, sessionID string, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + return c.run(ctx, accept, sessionID, prompt, attachments...) +} + +// run is the shared implementation behind Run and RunAccepted. When +// accept is non-nil it is threaded onto the SessionAgentCall as +// Accepted so sessionAgent.Run can consume the accept reservation under +// dispatchMu; when nil (the in-process/local path) no accept tracking +// applies. +func (c *coordinator) run(ctx context.Context, accept *AcceptedRun, sessionID string, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + if err := c.readyWg.Wait(); err != nil { + return nil, err + } + + // MCP servers connect asynchronously (see mcp.Initialize). + // + // Interactive runs never wait for that to finish: the tool list below + // is built from whatever is registered right now, servers still + // connecting are simply absent from this run's palette, and they are + // picked up by later runs once they register and publish + // EventToolsListChanged. Blocking here froze the TUI for the duration + // of the slowest server's connect timeout whenever a prompt was sent + // before initialization finished — most visibly on the first message. + // + // Non-interactive runs get a single shot at the tool palette, so they + // do wait for initialization to settle — but bounded by InitWaitBudget + // rather than each server's connect timeout, so a server wedged + // mid-handshake cannot stall a headless run for minutes. Past the + // budget the turn proceeds without the stragglers; their tools simply + // stay absent from this run. + if !c.interactive { + if err := mcp.WaitForInitBudget(ctx, mcp.InitWaitBudget); err != nil { + return nil, fmt.Errorf("failed to wait for MCP initialization: %w", err) + } + } + + // refresh models before each run + if err := c.UpdateModels(ctx); err != nil { + return nil, fmt.Errorf("failed to update models: %w", err) + } + + model := c.currentAgent.Model() + maxTokens := model.CatwalkCfg.DefaultMaxTokens + if model.ModelCfg.MaxTokens != 0 { + maxTokens = model.ModelCfg.MaxTokens + } + + providerCfg, ok := c.cfg.Config().Providers.Get(model.ModelCfg.Provider) + if !ok { + return nil, errModelProviderNotConfigured + } + + mergedOptions, temp, topP, topK, freqPenalty, presPenalty := mergeCallOptions(model, providerCfg) + + if err := c.refreshTokenIfExpired(ctx, providerCfg); err != nil { + // NOTE(@andreynering): We don't return here because the event handling to ask the user to reauthenticate + // depends on the flow below. If refresh fails, proceed with the token we have. + slog.Error("Failed to refresh OAuth2 token. Proceeding with existing token.", "error", err) + } + + // Coalesce per-attempt RunComplete payloads so only the final + // outcome reaches subscribers. Without this, the first attempt's + // failed RunComplete (unauthorized) would race ahead of the + // retry's success, and `crush run` would exit on the stale error + // before ever seeing the retry result. Each attempt's + // SessionAgentCall.OnComplete hook overwrites latest; we publish + // exactly once after retries resolve, via PublishMustDeliver, so + // a momentarily-full subscriber buffer can't silently drop the + // terminal event. + var ( + latest notify.RunComplete + hasLatest bool + ) + onComplete := func(rc notify.RunComplete) { + latest = rc + hasLatest = true + } + // Propagate the caller-supplied RunID (set via agent.WithRunID + // at the HTTP boundary in backend.SendMessage) onto the + // SessionAgentCall so the terminal RunComplete event echoes it + // back. Both attempts in the retry chain reuse the same RunID; + // the coalesce closure publishes the final outcome under that + // same correlator. + runID := RunIDFromContext(ctx) + run := func() (*fantasy.AgentResult, error) { + return c.currentAgent.Run(ctx, SessionAgentCall{ + SessionID: sessionID, + RunID: runID, + Prompt: prompt, + Attachments: attachments, + MaxOutputTokens: maxTokens, + ProviderOptions: mergedOptions, + Temperature: temp, + TopP: topP, + TopK: callTopK(providerCfg, topK), + FrequencyPenalty: freqPenalty, + PresencePenalty: presPenalty, + OnComplete: onComplete, + Accepted: accept, + OnAuthRefresh: c.makeAuthRefreshCallback(providerCfg), + }) + } + beforeLoaded := c.skillTracker.LoadedNames() + result, originalErr := run() + logTurnSkillUsage(sessionID, prompt, c.activeSkills, c.skillTracker, beforeLoaded) + + // Notify only if still unauthorized after retry — a successful + // retry means the user doesn't need to re-authenticate. AWS SSO is + // handled transparently inside OnAuthRefresh, so it needs no post-run + // notification here. + if originalErr != nil && isUnauthorized(originalErr) && c.notify != nil && model.ModelCfg.Provider == hyper.Name { + c.notify.Publish(pubsub.CreatedEvent, notify.Notification{ + Type: notify.TypeReAuthenticate, + ProviderID: model.ModelCfg.Provider, + }) + } + + if hasLatest && c.runComplete != nil { + c.runComplete.PublishMustDeliver(ctx, pubsub.UpdatedEvent, latest) + // Signal to the dispatcher (backend.runAgent) that the + // authoritative terminal RunComplete for this run was already + // emitted, so it does not publish a duplicate fallback for the + // error it is about to receive. + MarkRunCompletePublished(ctx) + } + return result, originalErr +} + +// effectiveReasoningEffort returns the reasoning effort to apply for provider calls. +// It prefers the user-selected effort when valid, otherwise the model default when +// valid, and finally falls back to the first configured reasoning level. +func effectiveReasoningEffort(model Model) string { + if !model.CatwalkCfg.CanReason { + return "" + } + + if effort := model.ModelCfg.ReasoningEffort; effort != "" && slices.Contains(model.CatwalkCfg.ReasoningLevels, effort) { + return effort + } + if effort := model.CatwalkCfg.DefaultReasoningEffort; effort != "" && slices.Contains(model.CatwalkCfg.ReasoningLevels, effort) { + return effort + } + if len(model.CatwalkCfg.ReasoningLevels) > 0 { + return model.CatwalkCfg.ReasoningLevels[0] + } + return "" +} + +func getProviderOptions(model Model, providerCfg config.ProviderConfig) fantasy.ProviderOptions { + options := fantasy.ProviderOptions{} + + cfgOpts := []byte("{}") + providerCfgOpts := []byte("{}") + catwalkOpts := []byte("{}") + + if model.ModelCfg.ProviderOptions != nil { + data, err := json.Marshal(model.ModelCfg.ProviderOptions) + if err == nil { + cfgOpts = data + } + } + + if providerCfg.ProviderOptions != nil { + data, err := json.Marshal(providerCfg.ProviderOptions) + if err == nil { + providerCfgOpts = data + } + } + + if model.CatwalkCfg.Options.ProviderOptions != nil { + data, err := json.Marshal(model.CatwalkCfg.Options.ProviderOptions) + if err == nil { + catwalkOpts = data + } + } + + readers := []io.Reader{ + bytes.NewReader(catwalkOpts), + bytes.NewReader(providerCfgOpts), + bytes.NewReader(cfgOpts), + } + + got, err := jsons.Merge(readers) + if err != nil { + slog.Error("Could not merge call config", "err", err) + return options + } + + mergedOptions := make(map[string]any) + + err = json.Unmarshal([]byte(got), &mergedOptions) + if err != nil { + slog.Error("Could not create config for call", "err", err) + return options + } + + reasoningEffort := effectiveReasoningEffort(model) + shouldSetEffort := model.CatwalkCfg.CanReason && + reasoningEffort != "" && + slices.Contains(model.CatwalkCfg.ReasoningLevels, reasoningEffort) + + switch providerCfg.Type { + case openai.Name, azure.Name: + _, hasReasoningEffort := mergedOptions["reasoning_effort"] + if !hasReasoningEffort && shouldSetEffort { + mergedOptions["reasoning_effort"] = reasoningEffort + } + if openai.IsResponsesModel(model.CatwalkCfg.ID) { + if openai.IsResponsesReasoningModel(model.CatwalkCfg.ID) { + mergedOptions["reasoning_summary"] = "auto" + mergedOptions["include"] = []openai.IncludeType{openai.IncludeReasoningEncryptedContent} + } + parsed, err := openai.ParseResponsesOptions(mergedOptions) + if err == nil { + options[openai.Name] = parsed + } + } else { + parsed, err := openai.ParseOptions(mergedOptions) + if err == nil { + options[openai.Name] = parsed + } + } + + case anthropic.Name, bedrock.Name: + var ( + _, hasEffort = mergedOptions["effort"] + _, hasThink = mergedOptions["thinking"] + extraBody = make(map[string]any) + ) + + switch providerCfg.ID { + case string(catwalk.InferenceProviderAlibabaSingapore), string(catwalk.InferenceProviderAlibabaUS): + switch { + case !hasEffort && shouldSetEffort: + extraBody["reasoning_effort"] = reasoningEffort + case !hasThink && model.CatwalkCfg.CanReason: + if model.ModelCfg.Think { + extraBody["thinking"] = map[string]any{"type": "enabled"} + } else { + extraBody["thinking"] = map[string]any{"type": "disabled"} + } + } + mergedOptions["extra_body"] = extraBody + + default: + switch { + case !hasEffort && shouldSetEffort: + mergedOptions["effort"] = reasoningEffort + case !hasThink && model.ModelCfg.Think: + mergedOptions["thinking"] = map[string]any{"budget_tokens": 2000} + } + } + + parsed, err := anthropic.ParseOptions(mergedOptions) + if err == nil { + options[anthropic.Name] = parsed + } + + case openrouter.Name: + _, hasReasoning := mergedOptions["reasoning"] + if !hasReasoning && shouldSetEffort { + mergedOptions["reasoning"] = map[string]any{ + "enabled": true, + "effort": reasoningEffort, + } + } + parsed, err := openrouter.ParseOptions(mergedOptions) + if err == nil { + options[openrouter.Name] = parsed + } + + case vercel.Name: + _, hasReasoning := mergedOptions["reasoning"] + if !hasReasoning && shouldSetEffort { + mergedOptions["reasoning"] = map[string]any{ + "enabled": true, + "effort": reasoningEffort, + } + } + parsed, err := vercel.ParseOptions(mergedOptions) + if err == nil { + options[vercel.Name] = parsed + } + + case google.Name: + _, hasReasoning := mergedOptions["thinking_config"] + if !hasReasoning { + if strings.HasPrefix(model.CatwalkCfg.ID, "gemini-2") { + mergedOptions["thinking_config"] = map[string]any{ + "thinking_budget": 2000, + "include_thoughts": true, + } + } else { + mergedOptions["thinking_config"] = map[string]any{ + "thinking_level": reasoningEffort, + "include_thoughts": true, + } + } + } + parsed, err := google.ParseOptions(mergedOptions) + if err == nil { + options[google.Name] = parsed + } + + case openaicompat.Name, hyper.Name: + extraBody := make(map[string]any) + + _, hasReasoningEffort := mergedOptions["reasoning_effort"] + if !hasReasoningEffort && shouldSetEffort { + switch providerCfg.ID { + case string(catwalk.InferenceProviderIoNet): + extraBody["reasoning"] = map[string]string{"effort": reasoningEffort} + case string(catwalk.InferenceProviderOpenCodeGo), string(catwalk.InferenceProviderOpenCodeZen): + // MiniMax models use the "thinking" parameter instead of + // "reasoning_effort". Other models on these providers still + // use the standard field. + if !strings.HasPrefix(strings.ToLower(model.CatwalkCfg.ID), "minimax") { + mergedOptions["reasoning_effort"] = reasoningEffort + } + default: + mergedOptions["reasoning_effort"] = reasoningEffort + } + } + + // "reasoning effort" is a standard OpenAI field, but "thinking" is not. + // Setting it in the right way for each provider. + // TODO: Abstract this in Fantasy somehow? + // TODO: Allow custom providers to specify how to set this? + switch providerCfg.ID { + case hyper.Name: + extraBody["thinking"] = model.ModelCfg.Think + case string(catwalk.InferenceProviderIoNet): + if _, ok := extraBody["reasoning"]; !ok && model.CatwalkCfg.CanReason { + if model.ModelCfg.Think { + extraBody["reasoning"] = map[string]string{"effort": "medium"} + } else { + extraBody["reasoning"] = map[string]string{"effort": "none"} + } + } + + case string(catwalk.InferenceProviderZAI), string(catwalk.InferenceProviderDeepSeek): + if model.ModelCfg.Think || reasoningEffort != "" { + extraBody["thinking"] = map[string]any{"type": "enabled"} + } else { + extraBody["thinking"] = map[string]any{"type": "disabled"} + } + + case string(catwalk.InferenceProviderFireworks): + // NOTE: Fireworks break if we set both `reasoning_effort` and `thinking`. + if reasoningEffort == "" { + if model.ModelCfg.Think { + extraBody["thinking"] = map[string]any{"type": "enabled"} + } else { + extraBody["thinking"] = map[string]any{"type": "disabled"} + } + } + + case string(catwalk.InferenceProviderBaseten): + extraBody["chat_template_args"] = map[string]any{ + "enable_thinking": model.ModelCfg.Think || reasoningEffort != "" && reasoningEffort != "none", + } + + case string(catwalk.InferenceProviderOpenCodeGo), string(catwalk.InferenceProviderOpenCodeZen): + // MiniMax M3 uses the "thinking" parameter to control reasoning. + // "reasoning_split" must be true so thinking content is returned + // in the "reasoning_content" field instead of inline in "content". + if strings.HasPrefix(strings.ToLower(model.CatwalkCfg.ID), "minimax") { + if model.CatwalkCfg.CanReason && (model.ModelCfg.Think || reasoningEffort != "") { + extraBody["thinking"] = map[string]any{"type": "adaptive"} + extraBody["reasoning_split"] = true + } else { + extraBody["thinking"] = map[string]any{"type": "disabled"} + } + } + + case string(catwalk.InferenceProviderAlibabaSingapore), string(catwalk.InferenceProviderAlibabaUS): + if model.CatwalkCfg.CanReason && !shouldSetEffort { + extraBody["enable_thinking"] = model.ModelCfg.Think + } + } + + mergedOptions["extra_body"] = extraBody + + parsed, err := openaicompat.ParseOptions(mergedOptions) + if err == nil { + options[openaicompat.Name] = parsed + } + + default: + // Known custom providers (litellm, llamacpp, lmstudio, ollama, + // omlx) are openai-compat under the hood. + if discover.IsKnownCustomProvider(string(providerCfg.Type)) { + // Set "top_k" under "extra_body", as it is not part of the OpenAI protocol + // and will be explicitly omitted by Fantasy downstream. + topK := cmp.Or(model.ModelCfg.TopK, model.CatwalkCfg.Options.TopK) + if topK != nil { + extraBody, hasExtraBody := mergedOptions["extra_body"].(map[string]any) + if !hasExtraBody { + extraBody = make(map[string]any) + mergedOptions["extra_body"] = extraBody + } + if _, hasTopK := extraBody["top_k"]; !hasTopK { + extraBody["top_k"] = *topK + } + } + + _, hasReasoningEffort := mergedOptions["reasoning_effort"] + if !hasReasoningEffort && shouldSetEffort { + mergedOptions["reasoning_effort"] = reasoningEffort + } + + parsed, err := openaicompat.ParseOptions(mergedOptions) + if err == nil { + options[openaicompat.Name] = parsed + } else { + if topK != nil { + slog.Warn( + "Failed to parse provider_options, falling back to top_k only", + "provider", providerCfg.ID, + "error", err, + ) + + fallbackMergeOptions := map[string]any{ + "extra_body": map[string]any{"top_k": *topK}, + } + parsed, err := openaicompat.ParseOptions(fallbackMergeOptions) + if err == nil { + options[openaicompat.Name] = parsed + } else { + slog.Warn( + "Failed to parse fallback provider options, this should never happen", + "provider", providerCfg.ID, + "error", err, + ) + } + } + } + } + } + + return options +} + +func mergeCallOptions(model Model, cfg config.ProviderConfig) (fantasy.ProviderOptions, *float64, *float64, *int64, *float64, *float64) { + modelOptions := getProviderOptions(model, cfg) + temp := cmp.Or(model.ModelCfg.Temperature, model.CatwalkCfg.Options.Temperature) + topP := cmp.Or(model.ModelCfg.TopP, model.CatwalkCfg.Options.TopP) + topK := cmp.Or(model.ModelCfg.TopK, model.CatwalkCfg.Options.TopK) + freqPenalty := cmp.Or(model.ModelCfg.FrequencyPenalty, model.CatwalkCfg.Options.FrequencyPenalty) + presPenalty := cmp.Or(model.ModelCfg.PresencePenalty, model.CatwalkCfg.Options.PresencePenalty) + return modelOptions, temp, topP, topK, freqPenalty, presPenalty +} + +func (c *coordinator) buildAgent(ctx context.Context, prompt *prompt.Prompt, agent config.Agent, isSubAgent bool) (SessionAgent, error) { + large, small, err := c.buildAgentModels(ctx, isSubAgent) + if err != nil { + return nil, err + } + + largeProviderCfg, _ := c.cfg.Config().Providers.Get(large.ModelCfg.Provider) + result := NewSessionAgent(SessionAgentOptions{ + LargeModel: large, + SmallModel: small, + SystemPromptPrefix: largeProviderCfg.SystemPromptPrefix, + SystemPrompt: "", + IsSubAgent: isSubAgent, + DisableAutoSummarize: c.cfg.Config().Options.DisableAutoSummarize, + IsYolo: c.permissions.SkipRequests(), + Sessions: c.sessions, + Messages: c.messages, + Tools: nil, + Notify: c.notify, + RunComplete: c.runComplete, + }) + + // The readiness goroutines below perform one-time setup — building the + // system prompt and the initial tool list — whose results the + // coordinator needs for its whole lifetime, so they must survive the + // caller's context being canceled. Several entry points build an agent + // from a short-lived HTTP request context: the server's + // InitAgent/UpdateAgent handlers, and UpdateModels -> buildTools -> + // agentTool -> buildAgent for the sub-agent. The tool-list build reads + // the MCP registry as it stands; servers still connecting are picked up + // by later runs. WithoutCancel drops cancellation while keeping context + // values; the work is local and always completes. + initCtx := context.WithoutCancel(ctx) + + c.readyWg.Go(func() error { + systemPrompt, err := prompt.Build(initCtx, large.Model.Provider(), large.Model.Model(), c.cfg) + if err != nil { + return err + } + result.SetSystemPrompt(systemPrompt) + return nil + }) + + c.readyWg.Go(func() error { + tools, err := c.buildTools(initCtx, agent, isSubAgent) + if err != nil { + return err + } + result.SetTools(tools) + return nil + }) + + return result, nil +} + +func (c *coordinator) buildTools(ctx context.Context, agent config.Agent, isSubAgent bool) ([]fantasy.AgentTool, error) { + var allTools []fantasy.AgentTool + if slices.Contains(agent.AllowedTools, AgentToolName) { + agentTool, err := c.agentTool(ctx) + if err != nil { + return nil, err + } + allTools = append(allTools, agentTool) + } + + if slices.Contains(agent.AllowedTools, tools.AgenticFetchToolName) { + agenticFetchTool, err := c.agenticFetchTool(ctx, nil) + if err != nil { + return nil, err + } + allTools = append(allTools, agenticFetchTool) + } + + // Get the model name for the agent + modelID := "" + if modelCfg, ok := c.cfg.Config().Models[agent.Model]; ok { + if model := c.cfg.Config().GetModel(modelCfg.Provider, modelCfg.Model); model != nil { + modelID = model.ID + } + } + + logFile := filepath.Join(c.cfg.Config().Options.DataDirectory, "logs", "crush.log") + + // Build hook runner if PreToolUse hooks are configured. + var hookRunner *hooks.Runner + if preToolHooks := c.cfg.Config().Hooks[hooks.EventPreToolUse]; len(preToolHooks) > 0 { + hookRunner = hooks.NewRunner(preToolHooks, c.cfg.WorkingDir(), c.cfg.WorkingDir()) + } + + allTools = append( + allTools, + tools.NewBashTool(c.permissions, c.cfg.WorkingDir(), c.cfg.Config().Options.Attribution, modelID), + tools.NewCrushInfoTool(c.cfg, c.lspManager, c.allSkills, c.activeSkills, c.skillTracker), + tools.NewCrushLogsTool(logFile), + tools.NewJobOutputTool(), + tools.NewJobKillTool(), + tools.NewDownloadTool(c.permissions, c.cfg.WorkingDir(), nil), + tools.NewEditTool(c.lspManager, c.permissions, c.history, c.filetracker, c.cfg.WorkingDir()), + tools.NewMultiEditTool(c.lspManager, c.permissions, c.history, c.filetracker, c.cfg.WorkingDir()), + tools.NewFetchTool(c.permissions, c.cfg.WorkingDir(), nil), + tools.NewGlobTool(c.cfg.WorkingDir(), c.cfg.Config().Tools.Glob), + tools.NewGrepTool(c.cfg.WorkingDir(), c.cfg.Config().Tools.Grep), + tools.NewLsTool(c.permissions, c.cfg.WorkingDir(), c.cfg.Config().Tools.Ls), + tools.NewSourcegraphTool(nil), + tools.NewTodosTool(c.sessions), + tools.NewViewTool(c.lspManager, c.permissions, c.filetracker, c.skillTracker, c.cfg.WorkingDir(), c.cfg.Config().Options.SkillsPaths...), + tools.NewWriteTool(c.lspManager, c.permissions, c.history, c.filetracker, c.cfg.WorkingDir()), + ) + + // Question tool is interactive-only and not available to sub-agents. + if !isSubAgent && c.interactive { + allTools = append(allTools, tools.NewQuestionTool(c.questions)) + } + + // Add LSP tools if user has configured LSPs or auto_lsp is enabled (nil or true). + if len(c.cfg.Config().LSP) > 0 || c.cfg.Config().Options.AutoLSP == nil || *c.cfg.Config().Options.AutoLSP { + allTools = append( + allTools, + tools.NewDiagnosticsTool(c.lspManager), + tools.NewReferencesTool(c.lspManager), + tools.NewLSPRestartTool(c.lspManager), + tools.NewSymbolsTool(c.lspManager), + tools.NewDefinitionTool(c.lspManager), + tools.NewCallHierarchyTool(c.lspManager), + tools.NewRenameTool(c.lspManager, c.permissions, c.history, c.filetracker), + tools.NewReplaceSymbolTool(c.lspManager, c.permissions, c.history, c.filetracker), + ) + } + + if len(c.cfg.Config().MCP) > 0 { + allTools = append( + allTools, + tools.NewListMCPResourcesTool(c.cfg, c.permissions), + tools.NewReadMCPResourceTool(c.cfg, c.permissions), + ) + } + + var filteredTools []fantasy.AgentTool + for _, tool := range allTools { + if slices.Contains(agent.AllowedTools, tool.Info().Name) { + filteredTools = append(filteredTools, tool) + } + } + + for _, tool := range tools.GetMCPTools(c.permissions, c.cfg, c.cfg.WorkingDir()) { + if agent.AllowedMCP == nil { + // No MCP restrictions + filteredTools = append(filteredTools, tool) + continue + } + if len(agent.AllowedMCP) == 0 { + // No MCPs allowed + slog.Debug("No MCPs allowed", "tool", tool.Name(), "agent", agent.Name) + break + } + + for mcp, tools := range agent.AllowedMCP { + if mcp != tool.MCP() { + continue + } + if len(tools) == 0 || slices.Contains(tools, tool.MCPToolName()) { + filteredTools = append(filteredTools, tool) + break + } + slog.Debug("MCP not allowed", "tool", tool.Name(), "agent", agent.Name) + } + } + slices.SortFunc(filteredTools, func(a, b fantasy.AgentTool) int { + return strings.Compare(a.Info().Name, b.Info().Name) + }) + + // Wrap tools with hook interception for the top-level agent only. + // Sub-agents (the `agent` task tool, `agentic_fetch`, etc.) run + // without hook interception to avoid firing the user's hook N times + // per delegated turn. The top-level invocation of the sub-agent tool + // itself is still wrapped from the coder's side. + filteredTools = wrapToolsWithHooks(filteredTools, hookRunner, isSubAgent) + + return filteredTools, nil +} + +// TODO: when we support multiple agents we need to change this so that we pass in the agent specific model config +func (c *coordinator) buildAgentModels(ctx context.Context, isSubAgent bool) (Model, Model, error) { + largeModelCfg, ok := c.cfg.Config().Models[config.SelectedModelTypeLarge] + if !ok { + return Model{}, Model{}, errLargeModelNotSelected + } + smallModelCfg, ok := c.cfg.Config().Models[config.SelectedModelTypeSmall] + if !ok { + return Model{}, Model{}, errSmallModelNotSelected + } + + largeProviderCfg, ok := c.cfg.Config().Providers.Get(largeModelCfg.Provider) + if !ok { + return Model{}, Model{}, errLargeModelProviderNotConfigured + } + + largeProvider, err := c.buildProvider(largeProviderCfg, largeModelCfg, isSubAgent) + if err != nil { + return Model{}, Model{}, err + } + + smallProviderCfg, ok := c.cfg.Config().Providers.Get(smallModelCfg.Provider) + if !ok { + return Model{}, Model{}, errSmallModelProviderNotConfigured + } + + smallProvider, err := c.buildProvider(smallProviderCfg, smallModelCfg, true) + if err != nil { + return Model{}, Model{}, err + } + + var largeCatwalkModel *catwalk.Model + var smallCatwalkModel *catwalk.Model + + for _, m := range largeProviderCfg.Models { + if m.ID == largeModelCfg.Model { + largeCatwalkModel = &m + } + } + for _, m := range smallProviderCfg.Models { + if m.ID == smallModelCfg.Model { + smallCatwalkModel = &m + } + } + + if largeCatwalkModel == nil { + return Model{}, Model{}, errLargeModelNotFound + } + + if smallCatwalkModel == nil { + return Model{}, Model{}, errSmallModelNotFound + } + + largeModelID := largeModelCfg.Model + smallModelID := smallModelCfg.Model + + if largeModelCfg.Provider == openrouter.Name && isExactoSupported(largeModelID) { + largeModelID += ":exacto" + } + + if smallModelCfg.Provider == openrouter.Name && isExactoSupported(smallModelID) { + smallModelID += ":exacto" + } + + largeModel, err := largeProvider.LanguageModel(ctx, largeModelID) + if err != nil { + return Model{}, Model{}, err + } + smallModel, err := smallProvider.LanguageModel(ctx, smallModelID) + if err != nil { + return Model{}, Model{}, err + } + + // Bound each request with the configured timeout so unreachable or hung + // providers fail instead of blocking a session forever. The wrapper is + // applied per request, so retries get a fresh budget each attempt. + requestTimeout := c.cfg.Config().Options.GetRequestTimeout() + largeModel = newRequestTimeoutModel(largeModel, requestTimeout) + smallModel = newRequestTimeoutModel(smallModel, requestTimeout) + + return Model{ + Model: largeModel, + CatwalkCfg: *largeCatwalkModel, + ModelCfg: largeModelCfg, + FlatRate: largeProviderCfg.FlatRate, + }, Model{ + Model: smallModel, + CatwalkCfg: *smallCatwalkModel, + ModelCfg: smallModelCfg, + FlatRate: smallProviderCfg.FlatRate, + }, nil +} + +func (c *coordinator) buildAnthropicProvider(baseURL, apiKey string, headers map[string]string, providerID string) (fantasy.Provider, error) { + var opts []anthropic.Option + + switch { + case strings.HasPrefix(apiKey, "Bearer "): + // NOTE: Prevent the SDK from picking up the API key from env. + os.Setenv("ANTHROPIC_API_KEY", "") + headers["Authorization"] = apiKey + case providerID == string(catwalk.InferenceProviderMiniMax) || providerID == string(catwalk.InferenceProviderMiniMaxChina): + // NOTE: Prevent the SDK from picking up the API key from env. + os.Setenv("ANTHROPIC_API_KEY", "") + headers["Authorization"] = "Bearer " + apiKey + case apiKey != "": + // X-Api-Key header + opts = append(opts, anthropic.WithAPIKey(apiKey)) + } + + if len(headers) > 0 { + opts = append(opts, anthropic.WithHeaders(headers)) + } + + if baseURL != "" { + opts = append(opts, anthropic.WithBaseURL(baseURL)) + } + + if c.cfg.Config().Options.Debug { + httpClient := log.NewHTTPClient() + opts = append(opts, anthropic.WithHTTPClient(httpClient)) + } + return anthropic.New(opts...) +} + +func (c *coordinator) buildOpenaiProvider(baseURL, apiKey string, headers map[string]string, token *oauth.Token) (fantasy.Provider, error) { + opts := []openai.Option{ + openai.WithAPIKey(apiKey), + openai.WithUseResponsesAPI(), + } + var httpClient *http.Client + if c.cfg.Config().Options.Debug { + httpClient = log.NewHTTPClient() + } + if token != nil { + // ChatGPT OAuth: requests go through the Codex backend, which + // expects account headers and rejects some request fields, so + // they pass through the Codex transport. + if httpClient == nil { + httpClient = &http.Client{} + } + httpClient.Transport = &openaioauth.Transport{ + Base: httpClient.Transport, + Token: token, + } + } + if httpClient != nil { + opts = append(opts, openai.WithHTTPClient(httpClient)) + } + if len(headers) > 0 { + opts = append(opts, openai.WithHeaders(headers)) + } + if baseURL != "" { + opts = append(opts, openai.WithBaseURL(baseURL)) + } + return openai.New(opts...) +} + +func (c *coordinator) buildOpenrouterProvider(_, apiKey string, headers map[string]string) (fantasy.Provider, error) { + opts := []openrouter.Option{ + openrouter.WithAPIKey(apiKey), + } + if c.cfg.Config().Options.Debug { + httpClient := log.NewHTTPClient() + opts = append(opts, openrouter.WithHTTPClient(httpClient)) + } + if len(headers) > 0 { + opts = append(opts, openrouter.WithHeaders(headers)) + } + return openrouter.New(opts...) +} + +func (c *coordinator) buildVercelProvider(_, apiKey string, headers map[string]string) (fantasy.Provider, error) { + opts := []vercel.Option{ + vercel.WithAPIKey(apiKey), + } + if c.cfg.Config().Options.Debug { + httpClient := log.NewHTTPClient() + opts = append(opts, vercel.WithHTTPClient(httpClient)) + } + if len(headers) > 0 { + opts = append(opts, vercel.WithHeaders(headers)) + } + return vercel.New(opts...) +} + +func (c *coordinator) buildOpenaiCompatProvider(baseURL, apiKey string, headers map[string]string, extraBody map[string]any, providerID string, isSubAgent bool) (fantasy.Provider, error) { + opts := []openaicompat.Option{ + openaicompat.WithBaseURL(baseURL), + openaicompat.WithAPIKey(apiKey), + } + + // Set HTTP client based on provider and debug mode. + var httpClient *http.Client + switch providerID { + case string(catwalk.InferenceProviderCopilot): + opts = append( + opts, + openaicompat.WithUseResponsesAPI(), + openaicompat.WithResponsesAPIFunc(func(modelID string) bool { + return copilotResponsesModels[modelID] + }), + ) + httpClient = copilot.NewClient(isSubAgent, c.cfg.Config().Options.Debug) + + case string(catwalk.InferenceProviderOpenCodeGo), string(catwalk.InferenceProviderOpenCodeZen): + opts = append( + opts, + openaicompat.WithUseResponsesAPI(), + openaicompat.WithResponsesAPIFunc(isOpenCodeResponsesModel), + ) + + case hyper.Name: + // Hyper may route requests through a Prism model; capture the + // router headers so the UI can show which model answered. + opts = append( + opts, + openaicompat.WithLanguageModelOptions( + openai.WithLanguageModelHeaderFunc(hyper.HeaderFunc), + ), + ) + } + if httpClient == nil && c.cfg.Config().Options.Debug { + httpClient = log.NewHTTPClient() + } + if httpClient != nil { + opts = append(opts, openaicompat.WithHTTPClient(httpClient)) + } + + if len(headers) > 0 { + opts = append(opts, openaicompat.WithHeaders(headers)) + } + + for extraKey, extraValue := range extraBody { + opts = append(opts, openaicompat.WithSDKOptions(openaisdk.WithJSONSet(extraKey, extraValue))) + } + + return openaicompat.New(opts...) +} + +func (c *coordinator) buildAzureProvider(baseURL, apiKey string, headers map[string]string, options map[string]string) (fantasy.Provider, error) { + opts := []azure.Option{ + azure.WithBaseURL(baseURL), + azure.WithAPIKey(apiKey), + azure.WithUseResponsesAPI(), + } + if c.cfg.Config().Options.Debug { + httpClient := log.NewHTTPClient() + opts = append(opts, azure.WithHTTPClient(httpClient)) + } + if options == nil { + options = make(map[string]string) + } + if apiVersion, ok := options["apiVersion"]; ok { + opts = append(opts, azure.WithAPIVersion(apiVersion)) + } + if len(headers) > 0 { + opts = append(opts, azure.WithHeaders(headers)) + } + + return azure.New(opts...) +} + +func (c *coordinator) buildBedrockProvider(apiKey string, headers map[string]string, providerID string) (fantasy.Provider, error) { + var opts []bedrock.Option + if c.cfg.Config().Options.Debug { + httpClient := log.NewHTTPClient() + opts = append(opts, bedrock.WithHTTPClient(httpClient)) + } + if len(headers) > 0 { + opts = append(opts, bedrock.WithHeaders(headers)) + } + + switch { + case apiKey != "": + opts = append(opts, bedrock.WithAPIKey(apiKey)) + case os.Getenv("AWS_BEARER_TOKEN_BEDROCK") != "": + opts = append(opts, bedrock.WithAPIKey(os.Getenv("AWS_BEARER_TOKEN_BEDROCK"))) + default: + // Skip, let the SDK do authentication. + } + + switch providerID { + case string(catwalk.InferenceProviderBedrockEurope): + opts = append(opts, bedrock.WithRegion("eu-west-1")) + default: + opts = append(opts, bedrock.WithRegion("us-east-1")) + } + + return bedrock.New(opts...) +} + +func (c *coordinator) buildGoogleProvider(baseURL, apiKey string, headers map[string]string) (fantasy.Provider, error) { + opts := []google.Option{ + google.WithBaseURL(baseURL), + google.WithGeminiAPIKey(apiKey), + } + if c.cfg.Config().Options.Debug { + httpClient := log.NewHTTPClient() + opts = append(opts, google.WithHTTPClient(httpClient)) + } + if len(headers) > 0 { + opts = append(opts, google.WithHeaders(headers)) + } + return google.New(opts...) +} + +func (c *coordinator) buildGoogleVertexProvider(headers map[string]string, options map[string]string) (fantasy.Provider, error) { + opts := []google.Option{} + if c.cfg.Config().Options.Debug { + httpClient := log.NewHTTPClient() + opts = append(opts, google.WithHTTPClient(httpClient)) + } + if len(headers) > 0 { + opts = append(opts, google.WithHeaders(headers)) + } + + project := options["project"] + location := options["location"] + + opts = append(opts, google.WithVertex(project, location)) + + return google.New(opts...) +} + +func (c *coordinator) isAnthropicThinking(model config.SelectedModel) bool { + if model.Think { + return true + } + opts, err := anthropic.ParseOptions(model.ProviderOptions) + return err == nil && opts.Thinking != nil +} + +func (c *coordinator) buildProvider(providerCfg config.ProviderConfig, model config.SelectedModel, isSubAgent bool) (fantasy.Provider, error) { + headers := maps.Clone(providerCfg.ExtraHeaders) + if headers == nil { + headers = make(map[string]string) + } + + // handle special headers for anthropic + if providerCfg.Type == anthropic.Name && c.isAnthropicThinking(model) { + if v, ok := headers["anthropic-beta"]; ok { + headers["anthropic-beta"] = v + ",interleaved-thinking-2025-05-14" + } else { + headers["anthropic-beta"] = "interleaved-thinking-2025-05-14" + } + } + + apiKey, _ := c.cfg.Resolve(providerCfg.APIKey) + baseURL, _ := c.cfg.Resolve(providerCfg.BaseURL) + + switch providerCfg.ID { + case string(catwalk.InferenceProviderOpenCodeGo), string(catwalk.InferenceProviderOpenCodeZen): + if isOpenCodeMessagesModel(providerCfg.ID, model.Model) { + baseURL = strings.TrimSuffix(baseURL, "/v1") + return c.buildAnthropicProvider(baseURL, apiKey, headers, providerCfg.ID) + } + } + + switch providerCfg.Type { + case openai.Name: + // A ChatGPT login is the provider's single credential: every + // request goes through the Codex backend with the OAuth token. + token := providerCfg.OAuthToken + if token != nil { + baseURL = openaioauth.CodexBaseURL + apiKey = token.AccessToken + headers["originator"] = "crush" + if token.AccountID != "" { + headers["chatgpt-account-id"] = token.AccountID + } + } + return c.buildOpenaiProvider(baseURL, apiKey, headers, token) + case anthropic.Name: + return c.buildAnthropicProvider(baseURL, apiKey, headers, providerCfg.ID) + case openrouter.Name: + return c.buildOpenrouterProvider(baseURL, apiKey, headers) + case vercel.Name: + return c.buildVercelProvider(baseURL, apiKey, headers) + case azure.Name: + return c.buildAzureProvider(baseURL, apiKey, headers, providerCfg.ExtraParams) + case bedrock.Name: + return c.buildBedrockProvider(apiKey, headers, providerCfg.ID) + case google.Name: + return c.buildGoogleProvider(baseURL, apiKey, headers) + case "google-vertex": + return c.buildGoogleVertexProvider(headers, providerCfg.ExtraParams) + case openaicompat.Name, hyper.Name: + switch providerCfg.ID { + case hyper.Name: + baseURL = hyper.BaseURL() + "/v1" + headers["x-crush-id"] = event.GetID() + case string(catwalk.InferenceProviderZAI): + if providerCfg.ExtraBody == nil { + providerCfg.ExtraBody = map[string]any{} + } + providerCfg.ExtraBody["tool_stream"] = true + } + return c.buildOpenaiCompatProvider(baseURL, apiKey, headers, providerCfg.ExtraBody, providerCfg.ID, isSubAgent) + default: + // Known custom providers (litellm, llamacpp, lmstudio, ollama, + // omlx) are openai-compat under the hood. + if discover.IsKnownCustomProvider(string(providerCfg.Type)) { + return c.buildOpenaiCompatProvider(baseURL, apiKey, headers, providerCfg.ExtraBody, providerCfg.ID, isSubAgent) + } + return nil, fmt.Errorf("provider type not supported: %q", providerCfg.Type) + } +} + +func isExactoSupported(modelID string) bool { + supportedModels := []string{ + "moonshotai/kimi-k2-0905", + "deepseek/deepseek-v3.1-terminus", + "z-ai/glm-4.6", + "openai/gpt-oss-120b", + "qwen/qwen3-coder", + } + return slices.Contains(supportedModels, modelID) +} + +// BeginAccepted reserves an accept slot for sessionID on the active +// agent and returns the ownership handle. It is the fire-and-forget +// dispatch path's only way to mark a run as accepted-but-not-yet-active +// so a cancel arriving before the run registers in activeRequests is not +// lost. +func (c *coordinator) BeginAccepted(sessionID string) *AcceptedRun { + return c.currentAgent.BeginAccepted(sessionID) +} + +func (c *coordinator) Cancel(sessionID string) { + c.currentAgent.Cancel(sessionID) +} + +func (c *coordinator) CancelAll() { + c.currentAgent.CancelAll() +} + +func (c *coordinator) ClearQueue(sessionID string) { + c.currentAgent.ClearQueue(sessionID) +} + +func (c *coordinator) IsBusy() bool { + return c.currentAgent.IsBusy() +} + +func (c *coordinator) IsSessionBusy(sessionID string) bool { + return c.currentAgent.IsSessionBusy(sessionID) +} + +func (c *coordinator) Model() Model { + return c.currentAgent.Model() +} + +func (c *coordinator) UpdateModels(ctx context.Context) error { + // A ChatGPT login without its model catalog — the fetch at login + // failed, or the credentials predate it — would leave the models + // dialog's ChatGPT section empty. Fill it in lazily; the guard makes + // this a no-op once the catalog exists. + c.cfg.RefetchOpenAIChatGPTModels(ctx) + + // build the models again so we make sure we get the latest config + large, small, err := c.buildAgentModels(ctx, false) + if err != nil { + return err + } + c.currentAgent.SetModels(large, small) + + agentCfg, ok := c.cfg.Config().Agents[config.AgentCoder] + if !ok { + return errCoderAgentNotConfigured + } + + tools, err := c.buildTools(ctx, agentCfg, false) + if err != nil { + return err + } + c.currentAgent.SetTools(tools) + return nil +} + +func (c *coordinator) QueuedPrompts(sessionID string) int { + return c.currentAgent.QueuedPrompts(sessionID) +} + +func (c *coordinator) QueuedPromptsList(sessionID string) []string { + return c.currentAgent.QueuedPromptsList(sessionID) +} + +func (c *coordinator) Summarize(ctx context.Context, sessionID string) error { + providerCfg, ok := c.cfg.Config().Providers.Get(c.currentAgent.Model().ModelCfg.Provider) + if !ok { + return errModelProviderNotConfigured + } + + if err := c.refreshTokenIfExpired(ctx, providerCfg); err != nil { + slog.Error("Failed to refresh OAuth2 token before summarize. Proceeding with existing token.", "error", err) + } + + // Auth failures during summarize flow through fantasy's OnAuthRefresh, + // the same path used by regular turns. + return c.currentAgent.Summarize(ctx, sessionID, getProviderOptions(c.currentAgent.Model(), providerCfg), c.makeAuthRefreshCallback(providerCfg)) +} + +// GenerateTitle generates a session title using the current agent. +func (c *coordinator) GenerateTitle(ctx context.Context, sessionID, prompt string) { + if c.currentAgent == nil { + return + } + c.currentAgent.GenerateTitle(ctx, sessionID, prompt) +} + +// refreshTokenIfExpired proactively refreshes the OAuth token if it has expired. +func (c *coordinator) refreshTokenIfExpired(ctx context.Context, providerCfg config.ProviderConfig) error { + if providerCfg.OAuthToken == nil || !providerCfg.OAuthToken.IsExpired() { + return nil + } + slog.Debug("Token needs to be refreshed", "provider", providerCfg.ID) + return c.refreshOAuth2Token(ctx, providerCfg) +} + +// retryAfterUnauthorized attempts to refresh credentials after an auth error +// and returns nil if the request should be retried. For OAuth providers whose +// refresh token is revoked, and for Bedrock providers whose AWS SSO session +// has expired, it triggers interactive re-authentication and blocks until the +// user completes it (or the context is cancelled). +func (c *coordinator) retryAfterUnauthorized(ctx context.Context, providerCfg config.ProviderConfig) error { + switch { + case providerCfg.OAuthToken != nil: + slog.Debug("Received 401. Refreshing token and retrying", "provider", providerCfg.ID) + if err := c.refreshOAuth2Token(ctx, providerCfg); err != nil { + // If the refresh token was revoked, trigger interactive + // re-auth and wait for the user to complete it. + var exchangeErr *oauth.TokenExchangeError + if c.notify != nil && errors.As(err, &exchangeErr) && exchangeErr.IsRefreshTokenRevoked() { + slog.Info("Refresh token revoked, waiting for re-authentication", "provider", providerCfg.ID) + c.notify.Publish(pubsub.CreatedEvent, notify.Notification{ + Type: notify.TypeReAuthenticate, + ProviderID: providerCfg.ID, + }) + return c.waitForInteractiveReauth(ctx, providerCfg.ID) + } + return err + } + return nil + case providerCfg.AWSAuthRefresh != "": + return c.refreshAWSCredentials(ctx, providerCfg) + case strings.Contains(providerCfg.APIKeyTemplate, "$"): + slog.Debug("Received 401. Refreshing API Key template and retrying", "provider", providerCfg.ID) + return c.refreshApiKeyTemplate(ctx, providerCfg) + default: + return nil + } +} + +// errNoInteractiveAuth is returned by an OnAuthRefresh callback when a +// provider needs interactive re-authentication but no notifier is available +// to drive it (e.g. headless runs). Returning it surfaces the original auth +// error rather than retrying. +var errNoInteractiveAuth = errors.New("interactive authentication unavailable") + +// waitForInteractiveReauth blocks until interactive re-authentication for the +// provider completes (signalled via SignalAuthComplete) or the context is +// cancelled, then rebuilds models so the next attempt picks up fresh +// credentials. Returns nil when the caller should retry. +func (c *coordinator) waitForInteractiveReauth(ctx context.Context, providerID string) error { + // Use a detached context with a generous timeout so the wait survives + // agent run cancellation. The user needs time to complete browser-based + // authentication. + waitCtx, waitCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Minute) + defer waitCancel() + slog.Info("Blocking on WaitForTokenChange", "provider", providerID) + if waitErr := c.cfg.WaitForTokenChange(waitCtx, providerID); waitErr != nil { + slog.Info("WaitForTokenChange returned error", "provider", providerID, "error", waitErr) + return waitErr + } + // If the original context was cancelled during the wait, fantasy's retry + // would fail immediately, so surface the cancellation instead. + if ctx.Err() != nil { + slog.Warn("Original context cancelled during auth wait, cannot retry", + "provider", providerID, "ctx_err", ctx.Err()) + return ctx.Err() + } + // Rebuild models so ModelProvider picks up the fresh credentials. + if updateErr := c.UpdateModels(waitCtx); updateErr != nil { + slog.Error("Failed to update models after re-authentication", "error", updateErr) + return updateErr + } + slog.Info("Models updated, returning nil to retry", "provider", providerID) + return nil +} + +// isUnauthorized reports whether err is an HTTP 401 from a provider. +func isUnauthorized(err error) bool { + var providerErr *fantasy.ProviderError + return errors.As(err, &providerErr) && providerErr.StatusCode == http.StatusUnauthorized +} + +// makeAuthRefreshCallback returns an OnAuthRefresh callback for fantasy that +// delegates to the coordinator's existing credential refresh logic. Returns +// nil if no refresh mechanism is configured for the provider. +func (c *coordinator) makeAuthRefreshCallback(providerCfg config.ProviderConfig) func(context.Context, *fantasy.ProviderError) error { + if providerCfg.OAuthToken == nil && + !strings.Contains(providerCfg.APIKeyTemplate, "$") && + providerCfg.AWSAuthRefresh == "" { + return nil + } + return func(ctx context.Context, _ *fantasy.ProviderError) error { + return c.retryAfterUnauthorized(ctx, providerCfg) + } +} + +func (c *coordinator) refreshOAuth2Token(ctx context.Context, providerCfg config.ProviderConfig) error { + if err := c.cfg.RefreshOAuthToken(ctx, config.ScopeGlobal, providerCfg.ID); err != nil { + slog.Error("Failed to refresh OAuth token after 401 error", "provider", providerCfg.ID, "error", err) + return err + } + if err := c.UpdateModels(ctx); err != nil { + return err + } + return nil +} + +func (c *coordinator) refreshApiKeyTemplate(ctx context.Context, providerCfg config.ProviderConfig) error { + newAPIKey, err := c.cfg.Resolve(providerCfg.APIKeyTemplate) + if err != nil { + slog.Error("Failed to re-resolve API key after 401 error", "provider", providerCfg.ID, "error", err) + return err + } + + providerCfg.APIKey = newAPIKey + c.cfg.Config().Providers.Set(providerCfg.ID, providerCfg) + + if err := c.UpdateModels(ctx); err != nil { + return err + } + return nil +} + +// subAgentParams holds the parameters for running a sub-agent. +type subAgentParams struct { + Agent SessionAgent + SessionID string + AgentMessageID string + ToolCallID string + Prompt string + SessionTitle string + // SessionSetup is an optional callback invoked after session creation + // but before agent execution, for custom session configuration. + SessionSetup func(sessionID string) +} + +// callTopK returns topK for use on fantasy.Call.TopK, suppressing it for +// known custom providers: getProviderOptions already carries top_k for +// them via extra_body, and passing it here too makes Fantasy emit a +// spurious "top_k unsupported" warning for every turn. +func callTopK(providerCfg config.ProviderConfig, topK *int64) *int64 { + if discover.IsKnownCustomProvider(string(providerCfg.Type)) { + return nil + } + return topK +} + +// runSubAgent runs a sub-agent and handles session management and cost accumulation. +// It creates a sub-session, runs the agent with the given prompt, and propagates +// the cost to the parent session. +func (c *coordinator) runSubAgent(ctx context.Context, params subAgentParams) (fantasy.ToolResponse, error) { + // Create sub-session + agentToolSessionID := c.sessions.CreateAgentToolSessionID(params.AgentMessageID, params.ToolCallID) + session, err := c.sessions.CreateTaskSession(ctx, agentToolSessionID, params.SessionID, params.SessionTitle) + if err != nil { + return fantasy.ToolResponse{}, fmt.Errorf("create session: %w", err) + } + + // Call session setup function if provided + if params.SessionSetup != nil { + params.SessionSetup(session.ID) + } + + // Get model configuration + model := params.Agent.Model() + maxTokens := model.CatwalkCfg.DefaultMaxTokens + if model.ModelCfg.MaxTokens != 0 { + maxTokens = model.ModelCfg.MaxTokens + } + + providerCfg, ok := c.cfg.Config().Providers.Get(model.ModelCfg.Provider) + if !ok { + return fantasy.ToolResponse{}, errModelProviderNotConfigured + } + + // Run the agent + run := func() (*fantasy.AgentResult, error) { + return params.Agent.Run(ctx, SessionAgentCall{ + SessionID: session.ID, + Prompt: params.Prompt, + MaxOutputTokens: maxTokens, + ProviderOptions: getProviderOptions(model, providerCfg), + Temperature: model.ModelCfg.Temperature, + TopP: model.ModelCfg.TopP, + TopK: callTopK(providerCfg, model.ModelCfg.TopK), + FrequencyPenalty: model.ModelCfg.FrequencyPenalty, + PresencePenalty: model.ModelCfg.PresencePenalty, + NonInteractive: true, + OnAuthRefresh: c.makeAuthRefreshCallback(providerCfg), + }) + } + result, err := run() + // Notify only if still unauthorized after retry. AWS SSO is handled + // transparently inside OnAuthRefresh, so it needs no post-run notice. + if err != nil && isUnauthorized(err) && c.notify != nil && model.ModelCfg.Provider == hyper.Name { + c.notify.Publish(pubsub.CreatedEvent, notify.Notification{ + Type: notify.TypeReAuthenticate, + ProviderID: model.ModelCfg.Provider, + }) + } + if err != nil { + return fantasy.NewTextErrorResponse(fmt.Sprintf("Failed to generate response: %s", err)), nil + } + + // Update parent session cost on a best-effort basis. A failure here must + // not discard the sub-agent output that was already produced. + if err := c.updateParentSessionCost(ctx, session.ID, params.SessionID); err != nil { + slog.Warn( + "Failed to update parent session cost", + "child_session", session.ID, + "parent_session", params.SessionID, + "error", err, + ) + } + + output := subAgentOutput(result) + if output == "" { + return fantasy.NewTextErrorResponse("Sub-agent completed but produced no text output."), nil + } + return fantasy.NewTextResponse(output), nil +} + +func subAgentOutput(result *fantasy.AgentResult) string { + if result == nil { + return "" + } + return result.Response.Content.Text() +} + +// updateParentSessionCost accumulates the cost from a child session to its parent session. +func (c *coordinator) updateParentSessionCost(ctx context.Context, childSessionID, parentSessionID string) error { + childSession, err := c.sessions.Get(ctx, childSessionID) + if err != nil { + return fmt.Errorf("get child session: %w", err) + } + + parentSession, err := c.sessions.Get(ctx, parentSessionID) + if err != nil { + return fmt.Errorf("get parent session: %w", err) + } + + parentSession.Cost += childSession.Cost + + if _, err := c.sessions.Save(ctx, parentSession); err != nil { + return fmt.Errorf("save parent session: %w", err) + } + + return nil +} + +// discoverSkills is a thin fallback wrapper used only when no +// skills.Manager has been threaded through to the coordinator. All +// production call sites (backend.CreateWorkspace, setupLocalWorkspace) +// run discovery in advance and pass the results via the manager; +// reaching this path means a caller bypassed both. It deliberately does +// NOT publish to the package-level broker — there are no subscribers in +// that case, so doing so would be misleading without delivering the +// snapshot anywhere useful. +func discoverSkills(cfg *config.ConfigStore) (allSkills, activeSkills []*skills.Skill) { + opts := cfg.Config().Options + var paths, disabled []string + if opts != nil { + paths = opts.SkillsPaths + disabled = opts.DisabledSkills + } + var resolver func(string) (string, error) + if r := cfg.Resolver(); r != nil { + resolver = r.ResolveValue + } + allSkills, activeSkills, states := skills.DiscoverFromConfig(skills.DiscoveryConfig{ + SkillsPaths: paths, + DisabledSkills: disabled, + Resolver: resolver, + }) + logDiscoveryStats(states, paths, allSkills, activeSkills, disabled) + return allSkills, activeSkills +} + +// logTurnSkillUsage emits a per-turn diagnostic line showing which skills +// (if any) were loaded during this turn and which looked relevant based on +// a cheap keyword match against the user prompt. The goal is to surface +// "should-have-loaded but didn't" situations for later analysis. +// +// Logged at Info level under component=skills; heavy fields are elided when +// there is nothing interesting to report. +func logTurnSkillUsage( + sessionID string, + prompt string, + activeSkills []*skills.Skill, + tracker *skills.Tracker, + before []string, +) { + if tracker == nil || len(activeSkills) == 0 { + return + } + + after := tracker.LoadedNames() + + beforeSet := make(map[string]bool, len(before)) + for _, n := range before { + beforeSet[n] = true + } + var loadedThisTurn []string + for _, n := range after { + if !beforeSet[n] { + loadedThisTurn = append(loadedThisTurn, n) + } + } + + slog.Info( + "Skill turn summary", + "component", "skills", + "session_id", sessionID, + "prompt_len", len(prompt), + "active_total", len(activeSkills), + "loaded_total", len(after), + "loaded_this_turn", loadedThisTurn, + ) +} + +// logDiscoveryStats emits a single structured log line summarising skill +// discovery for the current session. It is intentionally low-volume: one +// line per session start. Builtin vs user counts are derived from the +// SkillState.Path — builtin states use the "builtin/" embed prefix. +func logDiscoveryStats( + states []*skills.SkillState, + userPaths []string, + allSkills, activeSkills []*skills.Skill, + disabled []string, +) { + var builtinOK, builtinErr, userOK, userErr int + for _, s := range states { + isBuiltin := strings.HasPrefix(s.Path, "builtin/") + switch { + case isBuiltin && s.State == skills.StateNormal: + builtinOK++ + case isBuiltin && s.State == skills.StateError: + builtinErr++ + case !isBuiltin && s.State == skills.StateNormal: + userOK++ + case !isBuiltin && s.State == skills.StateError: + userErr++ + } + } + + activeNames := make([]string, 0, len(activeSkills)) + for _, s := range activeSkills { + activeNames = append(activeNames, s.Name) + } + + xml := skills.ToPromptXML(activeSkills) + + slog.Info( + "Skill discovery complete", + "component", "skills", + "builtin_ok", builtinOK, + "builtin_errors", builtinErr, + "user_ok", userOK, + "user_errors", userErr, + "user_paths", len(userPaths), + "deduped_total", len(allSkills), + "active", len(activeSkills), + "disabled", len(disabled), + "prompt_bytes", len(xml), + "prompt_tok_est", skills.ApproxTokenCount(xml), + "active_names", activeNames, + ) +} diff --git a/internal/agent/coordinator_mcp_gate_test.go b/internal/agent/coordinator_mcp_gate_test.go new file mode 100644 index 0000000000000000000000000000000000000000..355e90307adf5466eea6e46b5c4297c4d3cceb84 --- /dev/null +++ b/internal/agent/coordinator_mcp_gate_test.go @@ -0,0 +1,115 @@ +package agent + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "github.com/charmbracelet/crush/internal/agent/prompt" + "github.com/charmbracelet/crush/internal/agent/tools/mcp" + "github.com/charmbracelet/crush/internal/config" + "github.com/stretchr/testify/require" +) + +// newGateTestCoordinator builds a minimal coordinator against a hermetic +// config: one openai-typed provider pointed at a closed port, with large and +// small models selected so model resolution and the system-prompt build both +// succeed without any network access. +func newGateTestCoordinator(t *testing.T, interactive bool) *coordinator { + t.Helper() + + env := testEnv(t) + + crushJSON := `{ + "options": {"disable_default_providers": true, "disable_provider_auto_update": true}, + "providers": {"mock": {"id": "mock", "name": "Mock", "type": "openai", + "base_url": "http://127.0.0.1:9/v1", "api_key": "test-key", + "models": [{"id": "mock-model", "name": "Mock", "context_window": 8192, "default_max_tokens": 128}]}}, + "models": {"large": {"provider": "mock", "model": "mock-model"}, + "small": {"provider": "mock", "model": "mock-model"}} +}` + require.NoError(t, os.WriteFile(filepath.Join(env.workingDir, "crush.json"), []byte(crushJSON), 0o644)) + + cfg, err := config.Init(env.workingDir, "", false) + require.NoError(t, err) + cfg.SetupAgents() + + coord := &coordinator{ + cfg: cfg, + sessions: env.sessions, + messages: env.messages, + permissions: env.permissions, + history: env.history, + filetracker: *env.filetracker, + agents: make(map[string]SessionAgent), + interactive: interactive, + } + + p, err := coderPrompt(prompt.WithWorkingDir(env.workingDir)) + require.NoError(t, err) + agentCfg := cfg.Config().Agents[config.AgentCoder] + + agent, err := coord.buildAgent(context.Background(), p, agentCfg, false) + require.NoError(t, err) + coord.currentAgent = agent + coord.agents[config.AgentCoder] = agent + + return coord +} + +// TestRunWaitsForMCPOnlyWhenNonInteractive pins the split behavior for +// in-flight MCP initialization. +// +// MCP servers connect asynchronously. Interactive runs must not wait for them: +// blocking the send path meant a slow stdio server (e.g. Python via uv) froze +// the TUI for the length of its connect timeout, most visibly on the first +// message of a session. Tools from late servers simply miss that run's palette +// and show up on the next one. +// +// Non-interactive runs (`crush run`, both local and client/server) get a single +// shot at the palette, so they still wait for initialization to settle. +func TestRunWaitsForMCPOnlyWhenNonInteractive(t *testing.T) { + t.Run("non-interactive waits", func(t *testing.T) { + coord := newGateTestCoordinator(t, false) + + // Arm the gate and never complete initialization, standing in for an + // MCP server that is still connecting. + mcp.ArmInit() + t.Cleanup(mcp.DisarmInit) + + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + + _, err := coord.run(ctx, nil, "test-session", "hello") + require.ErrorContains(t, err, "MCP initialization", + "non-interactive run must block on MCP initialization") + }) + + t.Run("interactive does not wait", func(t *testing.T) { + coord := newGateTestCoordinator(t, true) + + mcp.ArmInit() + t.Cleanup(mcp.DisarmInit) + + done := make(chan error, 1) + go func() { + _, err := coord.run(context.Background(), nil, "test-session", "hello") + done <- err + }() + + select { + case err := <-done: + // The run fails for unrelated reasons (no such session, closed + // provider port); all that matters is that it got past the gate + // instead of parking on MCP initialization. + if err != nil { + require.NotContains(t, err.Error(), "MCP initialization", + "interactive run must not block on MCP initialization") + } + case <-time.After(10 * time.Second): + t.Fatal("interactive run blocked; it must not wait for MCP initialization") + } + }) +} diff --git a/internal/agent/coordinator_readiness_test.go b/internal/agent/coordinator_readiness_test.go new file mode 100644 index 0000000000000000000000000000000000000000..a4ad58bb1f520d7ff6ee5240bca9b586a4fb8ef1 --- /dev/null +++ b/internal/agent/coordinator_readiness_test.go @@ -0,0 +1,97 @@ +package agent + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "github.com/charmbracelet/crush/internal/agent/prompt" + "github.com/charmbracelet/crush/internal/agent/tools/mcp" + "github.com/charmbracelet/crush/internal/config" + "github.com/stretchr/testify/require" +) + +// TestBuildAgentReadinessSurvivesCallerCancellation is a regression test for +// the CRUSH_CLIENT_SERVER=1 "new session hangs" bug. +// +// buildAgent starts readiness goroutines that build the system prompt and the +// initial tool list. Several server entry points build an agent from a +// short-lived HTTP request context — the InitAgent/UpdateAgent handlers, and +// the sub-agent build reached through UpdateModels -> buildTools -> agentTool. +// When that request context was canceled the moment the handler returned, the +// readyWg errgroup recorded context.Canceled and every later coordinator.run +// failed at readyWg.Wait() before emitting anything — the session hung with +// no visible LLM response. (This was made worse while the tool-list goroutine +// also blocked in mcp.WaitForInit, which kept it parked long enough to +// observe the cancellation; the readiness work no longer waits on MCP init — +// see coordinator.run — but the cancellation detachment still matters.) +// +// The fix detaches the readiness work from the caller context via +// context.WithoutCancel, so canceling the context that triggered the build no +// longer poisons readyWg. Here we build an agent with a cancelable context, +// cancel it, and require that readyWg still completes cleanly. +func TestBuildAgentReadinessSurvivesCallerCancellation(t *testing.T) { + env := testEnv(t) + + // Minimal hermetic config: one openai-typed provider with selected large + // and small models so buildAgentModels and the system-prompt build both + // succeed. No MCP servers are configured, so initialization would complete + // instantly if we let it — we arm the gate anyway to prove the readiness + // goroutines no longer block on it. + crushJSON := `{ + "options": {"disable_default_providers": true, "disable_provider_auto_update": true}, + "providers": {"mock": {"id": "mock", "name": "Mock", "type": "openai", + "base_url": "http://127.0.0.1:9/v1", "api_key": "test-key", + "models": [{"id": "mock-model", "name": "Mock", "context_window": 8192, "default_max_tokens": 128}]}}, + "models": {"large": {"provider": "mock", "model": "mock-model"}, + "small": {"provider": "mock", "model": "mock-model"}} +}` + require.NoError(t, os.WriteFile(filepath.Join(env.workingDir, "crush.json"), []byte(crushJSON), 0o644)) + + cfg, err := config.Init(env.workingDir, "", false) + require.NoError(t, err) + cfg.SetupAgents() + + coord := &coordinator{ + cfg: cfg, + sessions: env.sessions, + messages: env.messages, + permissions: env.permissions, + history: env.history, + filetracker: *env.filetracker, + } + + // Arm the MCP init gate. We never complete init; the readiness goroutines + // must not care, since they build the tool list from the registry as it + // stands rather than waiting for initialization to finish. + mcp.ArmInit() + t.Cleanup(mcp.DisarmInit) + + p, err := coderPrompt(prompt.WithWorkingDir(env.workingDir)) + require.NoError(t, err) + agentCfg := cfg.Config().Agents[config.AgentCoder] + + ctx, cancel := context.WithCancel(context.Background()) + _, err = coord.buildAgent(ctx, p, agentCfg, false) + require.NoError(t, err) + + // The caller goes away, mirroring an HTTP handler returning and canceling + // its request context. + cancel() + + done := make(chan error, 1) + go func() { done <- coord.readyWg.Wait() }() + + select { + case err := <-done: + // context.Canceled is the regression: the caller's cancellation + // leaked into the readiness work and poisoned the errgroup. + require.NotErrorIs(t, err, context.Canceled, + "readyWg was poisoned by caller cancellation (client/server new-session hang regression)") + require.NoError(t, err, "unexpected buildAgent readiness error") + case <-time.After(2 * time.Second): + t.Fatal("readyWg did not complete; the readiness goroutines must not block on MCP init") + } +} diff --git a/internal/agent/coordinator_test.go b/internal/agent/coordinator_test.go new file mode 100644 index 0000000000000000000000000000000000000000..7e5020eab863b149188856174c088edd34f551a7 --- /dev/null +++ b/internal/agent/coordinator_test.go @@ -0,0 +1,792 @@ +package agent + +import ( + "context" + "errors" + "fmt" + "net/http" + "testing" + + "charm.land/catwalk/pkg/catwalk" + "charm.land/fantasy" + "charm.land/fantasy/providers/anthropic" + "charm.land/fantasy/providers/bedrock" + "charm.land/fantasy/providers/openaicompat" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/discover" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// mockSessionAgent is a minimal mock for the SessionAgent interface. +type mockSessionAgent struct { + model Model + runFunc func(ctx context.Context, call SessionAgentCall) (*fantasy.AgentResult, error) + cancelled []string +} + +func (m *mockSessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy.AgentResult, error) { + return m.runFunc(ctx, call) +} + +func (m *mockSessionAgent) BeginAccepted(sessionID string) *AcceptedRun { + return &AcceptedRun{sessionID: sessionID} +} + +func (m *mockSessionAgent) Model() Model { return m.model } +func (m *mockSessionAgent) SetModels(large, small Model) {} +func (m *mockSessionAgent) SetTools(tools []fantasy.AgentTool) {} +func (m *mockSessionAgent) SetSystemPrompt(systemPrompt string) {} +func (m *mockSessionAgent) Cancel(sessionID string) { + m.cancelled = append(m.cancelled, sessionID) +} +func (m *mockSessionAgent) CancelAll() {} +func (m *mockSessionAgent) IsSessionBusy(sessionID string) bool { return false } +func (m *mockSessionAgent) IsBusy() bool { return false } +func (m *mockSessionAgent) QueuedPrompts(sessionID string) int { return 0 } +func (m *mockSessionAgent) QueuedPromptsList(sessionID string) []string { return nil } +func (m *mockSessionAgent) ClearQueue(sessionID string) {} +func (m *mockSessionAgent) Summarize(context.Context, string, fantasy.ProviderOptions, func(context.Context, *fantasy.ProviderError) error) error { + return nil +} +func (m *mockSessionAgent) GenerateTitle(context.Context, string, string) {} + +// newTestCoordinator creates a minimal coordinator for unit testing runSubAgent. +func newTestCoordinator(t *testing.T, env fakeEnv, providerID string, providerCfg config.ProviderConfig) *coordinator { + cfg, err := config.Init(env.workingDir, "", false) + require.NoError(t, err) + cfg.Config().Providers.Set(providerID, providerCfg) + return &coordinator{ + cfg: cfg, + sessions: env.sessions, + messages: env.messages, + } +} + +// newMockAgent creates a mockSessionAgent with the given provider and run function. +func newMockAgent(providerID string, maxTokens int64, runFunc func(context.Context, SessionAgentCall) (*fantasy.AgentResult, error)) *mockSessionAgent { + return &mockSessionAgent{ + model: Model{ + CatwalkCfg: catwalk.Model{ + DefaultMaxTokens: maxTokens, + }, + ModelCfg: config.SelectedModel{ + Provider: providerID, + }, + }, + runFunc: runFunc, + } +} + +// agentResultWithText creates a minimal AgentResult with the given text response. +func agentResultWithText(text string) *fantasy.AgentResult { + return &fantasy.AgentResult{ + Response: fantasy.Response{ + Content: fantasy.ResponseContent{ + fantasy.TextContent{Text: text}, + }, + }, + } +} + +func TestCopilotResponsesModels(t *testing.T) { + t.Parallel() + + for _, modelID := range []string{"gpt-6-astra", "grok-4.5", "grok-4.6"} { + assert.True(t, copilotResponsesModels[modelID], modelID) + } + assert.False(t, copilotResponsesModels["gpt-4.1"]) +} + +func TestRunSubAgent(t *testing.T) { + const providerID = "test-provider" + providerCfg := config.ProviderConfig{ID: providerID} + + t.Run("happy path", func(t *testing.T) { + env := testEnv(t) + coord := newTestCoordinator(t, env, providerID, providerCfg) + + parentSession, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + + agent := newMockAgent(providerID, 4096, func(_ context.Context, call SessionAgentCall) (*fantasy.AgentResult, error) { + assert.Equal(t, "do something", call.Prompt) + assert.Equal(t, int64(4096), call.MaxOutputTokens) + return agentResultWithText("done"), nil + }) + + resp, err := coord.runSubAgent(t.Context(), subAgentParams{ + Agent: agent, + SessionID: parentSession.ID, + AgentMessageID: "msg-1", + ToolCallID: "call-1", + Prompt: "do something", + SessionTitle: "Test Session", + }) + require.NoError(t, err) + assert.Equal(t, "done", resp.Content) + assert.False(t, resp.IsError) + }) + + t.Run("cost update failure preserves output", func(t *testing.T) { + env := testEnv(t) + coord := newTestCoordinator(t, env, providerID, providerCfg) + + agent := newMockAgent(providerID, 4096, func(_ context.Context, _ SessionAgentCall) (*fantasy.AgentResult, error) { + return agentResultWithText("output before cost failure"), nil + }) + + resp, err := coord.runSubAgent(t.Context(), subAgentParams{ + Agent: agent, + SessionID: "missing-parent-session", + AgentMessageID: "msg-1", + ToolCallID: "call-1", + Prompt: "test", + SessionTitle: "Test", + }) + require.NoError(t, err) + assert.False(t, resp.IsError) + assert.Equal(t, "output before cost failure", resp.Content) + }) + + t.Run("response with text returns it", func(t *testing.T) { + env := testEnv(t) + coord := newTestCoordinator(t, env, providerID, providerCfg) + + parentSession, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + + agent := newMockAgent(providerID, 4096, func(_ context.Context, _ SessionAgentCall) (*fantasy.AgentResult, error) { + return agentResultWithText("the answer"), nil + }) + + resp, err := coord.runSubAgent(t.Context(), subAgentParams{ + Agent: agent, + SessionID: parentSession.ID, + AgentMessageID: "msg-1", + ToolCallID: "call-1", + Prompt: "test", + SessionTitle: "Test", + }) + require.NoError(t, err) + assert.False(t, resp.IsError) + assert.Equal(t, "the answer", resp.Content) + }) + + t.Run("nil result returns error response", func(t *testing.T) { + env := testEnv(t) + coord := newTestCoordinator(t, env, providerID, providerCfg) + + parentSession, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + + agent := newMockAgent(providerID, 4096, func(_ context.Context, _ SessionAgentCall) (*fantasy.AgentResult, error) { + return nil, nil + }) + + resp, err := coord.runSubAgent(t.Context(), subAgentParams{ + Agent: agent, + SessionID: parentSession.ID, + AgentMessageID: "msg-1", + ToolCallID: "call-1", + Prompt: "test", + SessionTitle: "Test", + }) + require.NoError(t, err) + assert.True(t, resp.IsError) + assert.Equal(t, "Sub-agent completed but produced no text output.", resp.Content) + }) + + t.Run("empty result returns error response", func(t *testing.T) { + env := testEnv(t) + coord := newTestCoordinator(t, env, providerID, providerCfg) + + parentSession, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + + agent := newMockAgent(providerID, 4096, func(_ context.Context, _ SessionAgentCall) (*fantasy.AgentResult, error) { + return &fantasy.AgentResult{}, nil + }) + + resp, err := coord.runSubAgent(t.Context(), subAgentParams{ + Agent: agent, + SessionID: parentSession.ID, + AgentMessageID: "msg-1", + ToolCallID: "call-1", + Prompt: "test", + SessionTitle: "Test", + }) + require.NoError(t, err) + assert.True(t, resp.IsError) + assert.Equal(t, "Sub-agent completed but produced no text output.", resp.Content) + }) + + t.Run("ModelCfg.MaxTokens overrides default", func(t *testing.T) { + env := testEnv(t) + coord := newTestCoordinator(t, env, providerID, providerCfg) + + parentSession, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + + agent := &mockSessionAgent{ + model: Model{ + CatwalkCfg: catwalk.Model{ + DefaultMaxTokens: 4096, + }, + ModelCfg: config.SelectedModel{ + Provider: providerID, + MaxTokens: 8192, + }, + }, + runFunc: func(_ context.Context, call SessionAgentCall) (*fantasy.AgentResult, error) { + assert.Equal(t, int64(8192), call.MaxOutputTokens) + return agentResultWithText("ok"), nil + }, + } + + resp, err := coord.runSubAgent(t.Context(), subAgentParams{ + Agent: agent, + SessionID: parentSession.ID, + AgentMessageID: "msg-1", + ToolCallID: "call-1", + Prompt: "test", + SessionTitle: "Test", + }) + require.NoError(t, err) + assert.Equal(t, "ok", resp.Content) + }) + + t.Run("session creation failure with canceled context", func(t *testing.T) { + env := testEnv(t) + coord := newTestCoordinator(t, env, providerID, providerCfg) + + parentSession, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + + agent := newMockAgent(providerID, 4096, nil) + + // Use a canceled context to trigger CreateTaskSession failure. + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + _, err = coord.runSubAgent(ctx, subAgentParams{ + Agent: agent, + SessionID: parentSession.ID, + AgentMessageID: "msg-1", + ToolCallID: "call-1", + Prompt: "test", + SessionTitle: "Test", + }) + require.Error(t, err) + }) + + t.Run("provider not configured", func(t *testing.T) { + env := testEnv(t) + coord := newTestCoordinator(t, env, providerID, providerCfg) + + parentSession, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + + // Agent references a provider that doesn't exist in config. + agent := newMockAgent("unknown-provider", 4096, nil) + + _, err = coord.runSubAgent(t.Context(), subAgentParams{ + Agent: agent, + SessionID: parentSession.ID, + AgentMessageID: "msg-1", + ToolCallID: "call-1", + Prompt: "test", + SessionTitle: "Test", + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "model provider not configured") + }) + + t.Run("agent run error returns error response", func(t *testing.T) { + env := testEnv(t) + coord := newTestCoordinator(t, env, providerID, providerCfg) + + parentSession, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + + agent := newMockAgent(providerID, 4096, func(_ context.Context, _ SessionAgentCall) (*fantasy.AgentResult, error) { + return nil, errors.New("provider request failed") + }) + + resp, err := coord.runSubAgent(t.Context(), subAgentParams{ + Agent: agent, + SessionID: parentSession.ID, + AgentMessageID: "msg-1", + ToolCallID: "call-1", + Prompt: "test", + SessionTitle: "Test", + }) + // runSubAgent returns (errorResponse, nil) when agent.Run fails — not a Go error. + require.NoError(t, err) + assert.True(t, resp.IsError) + assert.Equal(t, "Failed to generate response: provider request failed", resp.Content) + }) + + t.Run("session setup callback is invoked", func(t *testing.T) { + env := testEnv(t) + coord := newTestCoordinator(t, env, providerID, providerCfg) + + parentSession, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + + var setupCalledWith string + agent := newMockAgent(providerID, 4096, func(_ context.Context, _ SessionAgentCall) (*fantasy.AgentResult, error) { + return agentResultWithText("ok"), nil + }) + + _, err = coord.runSubAgent(t.Context(), subAgentParams{ + Agent: agent, + SessionID: parentSession.ID, + AgentMessageID: "msg-1", + ToolCallID: "call-1", + Prompt: "test", + SessionTitle: "Test", + SessionSetup: func(sessionID string) { + setupCalledWith = sessionID + }, + }) + require.NoError(t, err) + assert.NotEmpty(t, setupCalledWith, "SessionSetup should have been called") + }) + + t.Run("cost propagation to parent session", func(t *testing.T) { + env := testEnv(t) + coord := newTestCoordinator(t, env, providerID, providerCfg) + + parentSession, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + + agent := newMockAgent(providerID, 4096, func(ctx context.Context, call SessionAgentCall) (*fantasy.AgentResult, error) { + // Simulate the agent incurring cost by updating the child session. + childSession, err := env.sessions.Get(ctx, call.SessionID) + if err != nil { + return nil, err + } + childSession.Cost = 0.05 + _, err = env.sessions.Save(ctx, childSession) + if err != nil { + return nil, err + } + return agentResultWithText("ok"), nil + }) + + _, err = coord.runSubAgent(t.Context(), subAgentParams{ + Agent: agent, + SessionID: parentSession.ID, + AgentMessageID: "msg-1", + ToolCallID: "call-1", + Prompt: "test", + SessionTitle: "Test", + }) + require.NoError(t, err) + + updated, err := env.sessions.Get(t.Context(), parentSession.ID) + require.NoError(t, err) + assert.InDelta(t, 0.05, updated.Cost, 1e-9) + }) +} + +func TestUpdateParentSessionCost(t *testing.T) { + t.Run("accumulates cost correctly", func(t *testing.T) { + env := testEnv(t) + cfg, err := config.Init(env.workingDir, "", false) + require.NoError(t, err) + coord := &coordinator{cfg: cfg, sessions: env.sessions} + + parent, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + + child, err := env.sessions.CreateTaskSession(t.Context(), "tool-1", parent.ID, "Child") + require.NoError(t, err) + + // Set child cost. + child.Cost = 0.10 + _, err = env.sessions.Save(t.Context(), child) + require.NoError(t, err) + + err = coord.updateParentSessionCost(t.Context(), child.ID, parent.ID) + require.NoError(t, err) + + updated, err := env.sessions.Get(t.Context(), parent.ID) + require.NoError(t, err) + assert.InDelta(t, 0.10, updated.Cost, 1e-9) + }) + + t.Run("accumulates multiple child costs", func(t *testing.T) { + env := testEnv(t) + cfg, err := config.Init(env.workingDir, "", false) + require.NoError(t, err) + coord := &coordinator{cfg: cfg, sessions: env.sessions} + + parent, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + + child1, err := env.sessions.CreateTaskSession(t.Context(), "tool-1", parent.ID, "Child1") + require.NoError(t, err) + child1.Cost = 0.05 + _, err = env.sessions.Save(t.Context(), child1) + require.NoError(t, err) + + child2, err := env.sessions.CreateTaskSession(t.Context(), "tool-2", parent.ID, "Child2") + require.NoError(t, err) + child2.Cost = 0.03 + _, err = env.sessions.Save(t.Context(), child2) + require.NoError(t, err) + + err = coord.updateParentSessionCost(t.Context(), child1.ID, parent.ID) + require.NoError(t, err) + err = coord.updateParentSessionCost(t.Context(), child2.ID, parent.ID) + require.NoError(t, err) + + updated, err := env.sessions.Get(t.Context(), parent.ID) + require.NoError(t, err) + assert.InDelta(t, 0.08, updated.Cost, 1e-9) + }) + + t.Run("child session not found", func(t *testing.T) { + env := testEnv(t) + cfg, err := config.Init(env.workingDir, "", false) + require.NoError(t, err) + coord := &coordinator{cfg: cfg, sessions: env.sessions} + + parent, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + + err = coord.updateParentSessionCost(t.Context(), "non-existent", parent.ID) + require.Error(t, err) + assert.Contains(t, err.Error(), "get child session") + }) + + t.Run("parent session not found", func(t *testing.T) { + env := testEnv(t) + cfg, err := config.Init(env.workingDir, "", false) + require.NoError(t, err) + coord := &coordinator{cfg: cfg, sessions: env.sessions} + + parent, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + child, err := env.sessions.CreateTaskSession(t.Context(), "tool-1", parent.ID, "Child") + require.NoError(t, err) + + err = coord.updateParentSessionCost(t.Context(), child.ID, "non-existent") + require.Error(t, err) + assert.Contains(t, err.Error(), "get parent session") + }) + + t.Run("zero cost handled correctly", func(t *testing.T) { + env := testEnv(t) + cfg, err := config.Init(env.workingDir, "", false) + require.NoError(t, err) + coord := &coordinator{cfg: cfg, sessions: env.sessions} + + parent, err := env.sessions.Create(t.Context(), "Parent") + require.NoError(t, err) + child, err := env.sessions.CreateTaskSession(t.Context(), "tool-1", parent.ID, "Child") + require.NoError(t, err) + + err = coord.updateParentSessionCost(t.Context(), child.ID, parent.ID) + require.NoError(t, err) + + updated, err := env.sessions.Get(t.Context(), parent.ID) + require.NoError(t, err) + assert.InDelta(t, 0.0, updated.Cost, 1e-9) + }) +} + +func TestGetProviderOptionsReasoningEffort(t *testing.T) { + // Bedrock is Fantasy's Anthropic under a different provider name; options + // must land under anthropic.Name so the Anthropic language model picks them up. + tests := []struct { + name string + providerType catwalk.Type + }{ + {"anthropic honors reasoning_effort", catwalk.Type(anthropic.Name)}, + {"bedrock honors reasoning_effort", catwalk.Type(bedrock.Name)}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + model := Model{ + CatwalkCfg: catwalk.Model{ + ID: "claude-opus-4-7", + CanReason: true, + ReasoningLevels: []string{"max"}, + }, + ModelCfg: config.SelectedModel{ + Provider: "test", + ReasoningEffort: "max", + }, + } + providerCfg := config.ProviderConfig{ID: "test", Type: tc.providerType} + + opts := getProviderOptions(model, providerCfg) + + raw, ok := opts[anthropic.Name] + require.True(t, ok, "options should be keyed under anthropic.Name for type %q", tc.providerType) + parsed, ok := raw.(*anthropic.ProviderOptions) + require.True(t, ok) + require.NotNil(t, parsed.Effort) + assert.Equal(t, anthropic.Effort("max"), *parsed.Effort) + }) + } +} + +func TestIsUnauthorized(t *testing.T) { + t.Run("nil error", func(t *testing.T) { + assert.False(t, isUnauthorized(nil)) + }) + + t.Run("non-provider error", func(t *testing.T) { + assert.False(t, isUnauthorized(errors.New("something broke"))) + }) + + t.Run("provider error with 401", func(t *testing.T) { + err := &fantasy.ProviderError{StatusCode: http.StatusUnauthorized, Message: "unauthorized"} + assert.True(t, isUnauthorized(err)) + }) + + t.Run("provider error with non-401", func(t *testing.T) { + err := &fantasy.ProviderError{StatusCode: http.StatusForbidden, Message: "forbidden"} + assert.False(t, isUnauthorized(err)) + }) + + t.Run("wrapped provider error with 401", func(t *testing.T) { + inner := &fantasy.ProviderError{StatusCode: http.StatusUnauthorized, Message: "expired"} + err := fmt.Errorf("request failed: %w", inner) + assert.True(t, isUnauthorized(err)) + }) +} + +func TestGetProviderOptionsReasoningEffortCustomProvider(t *testing.T) { + // Custom local providers (lmstudio, ollama, omlx, litellm, llamacpp) + // go through the OpenAI-compat client and must receive the selected + // reasoning effort like any other OpenAI-compatible provider. + for _, providerType := range discover.RegisteredProviderTypes() { + t.Run(providerType, func(t *testing.T) { + model := Model{ + CatwalkCfg: catwalk.Model{ + ID: "qwen/qwen3-8b", + CanReason: true, + ReasoningLevels: []string{"low", "medium", "high"}, + }, + ModelCfg: config.SelectedModel{ + Provider: "local", + Model: "qwen/qwen3-8b", + ReasoningEffort: "high", + }, + } + providerCfg := config.ProviderConfig{ID: "local", Type: catwalk.Type(providerType)} + + opts := getProviderOptions(model, providerCfg) + + raw, ok := opts[openaicompat.Name] + require.True(t, ok, "options should be keyed under openaicompat.Name for type %q", providerType) + parsed, ok := raw.(*openaicompat.ProviderOptions) + require.True(t, ok) + require.NotNil(t, parsed.ReasoningEffort) + assert.Equal(t, "high", string(*parsed.ReasoningEffort)) + }) + } +} + +func TestGetProviderOptionsReasoningEffortFallback(t *testing.T) { + model := Model{ + CatwalkCfg: catwalk.Model{ + ID: "glm-5.2", + CanReason: true, + ReasoningLevels: []string{"high", "max"}, + }, + ModelCfg: config.SelectedModel{ + Provider: "zai", + }, + } + providerCfg := config.ProviderConfig{ + ID: string(catwalk.InferenceProviderZAI), + Type: openaicompat.Name, + } + + opts := getProviderOptions(model, providerCfg) + + raw, ok := opts[openaicompat.Name] + require.True(t, ok) + parsed, ok := raw.(*openaicompat.ProviderOptions) + require.True(t, ok) + require.NotNil(t, parsed.ReasoningEffort) + assert.Equal(t, "high", string(*parsed.ReasoningEffort)) + + thinking, ok := parsed.ExtraBody["thinking"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "enabled", thinking["type"]) +} + +func TestGetProviderOptionsTopKExtraBody(t *testing.T) { + // "ollama" has a registered discover.Enricher, so it is treated as a + // known custom provider speaking openai-compat. + knownCustomProviderCfg := config.ProviderConfig{ID: "ollama", Type: "ollama"} + + t.Run("model top_k is injected into extra_body for known custom providers", func(t *testing.T) { + model := Model{ + CatwalkCfg: catwalk.Model{ID: "llama3"}, + ModelCfg: config.SelectedModel{Provider: "ollama", TopK: ptr(int64(40))}, + } + + opts := getProviderOptions(model, knownCustomProviderCfg) + + raw, ok := opts[openaicompat.Name] + require.True(t, ok) + parsed, ok := raw.(*openaicompat.ProviderOptions) + require.True(t, ok) + topK, ok := parsed.ExtraBody["top_k"].(int64) + require.True(t, ok) + assert.Equal(t, int64(40), topK) + }) + + t.Run("falls back to catwalk top_k when the model config has none", func(t *testing.T) { + model := Model{ + CatwalkCfg: catwalk.Model{ + ID: "llama3", + Options: catwalk.ModelOptions{TopK: ptr(int64(64))}, + }, + ModelCfg: config.SelectedModel{Provider: "ollama"}, + } + + opts := getProviderOptions(model, knownCustomProviderCfg) + + raw, ok := opts[openaicompat.Name] + require.True(t, ok) + parsed, ok := raw.(*openaicompat.ProviderOptions) + require.True(t, ok) + topK, ok := parsed.ExtraBody["top_k"].(int64) + require.True(t, ok) + assert.Equal(t, int64(64), topK) + }) + + t.Run("does not set extra_body when no top_k is configured anywhere", func(t *testing.T) { + model := Model{ + CatwalkCfg: catwalk.Model{ID: "llama3"}, + ModelCfg: config.SelectedModel{Provider: "ollama"}, + } + + opts := getProviderOptions(model, knownCustomProviderCfg) + + raw, ok := opts[openaicompat.Name] + require.True(t, ok) + parsed, ok := raw.(*openaicompat.ProviderOptions) + require.True(t, ok) + _, hasTopK := parsed.ExtraBody["top_k"] + assert.False(t, hasTopK) + }) + + t.Run("does not overwrite an explicitly configured extra_body.top_k", func(t *testing.T) { + model := Model{ + CatwalkCfg: catwalk.Model{ID: "llama3"}, + ModelCfg: config.SelectedModel{ + Provider: "ollama", + TopK: ptr(int64(40)), + ProviderOptions: map[string]any{ + "extra_body": map[string]any{"top_k": 7}, + }, + }, + } + + opts := getProviderOptions(model, knownCustomProviderCfg) + + raw, ok := opts[openaicompat.Name] + require.True(t, ok) + parsed, ok := raw.(*openaicompat.ProviderOptions) + require.True(t, ok) + assert.EqualValues(t, 7, parsed.ExtraBody["top_k"]) + }) + + t.Run("is not injected for providers outside the known-custom-provider default branch", func(t *testing.T) { + model := Model{ + CatwalkCfg: catwalk.Model{ID: "glm-5.2"}, + ModelCfg: config.SelectedModel{Provider: "zai", TopK: ptr(int64(40))}, + } + providerCfg := config.ProviderConfig{ID: string(catwalk.InferenceProviderZAI), Type: openaicompat.Name} + + opts := getProviderOptions(model, providerCfg) + + raw, ok := opts[openaicompat.Name] + require.True(t, ok) + parsed, ok := raw.(*openaicompat.ProviderOptions) + require.True(t, ok) + _, hasTopK := parsed.ExtraBody["top_k"] + assert.False(t, hasTopK) + }) +} + +func TestGetProviderOptionsMalformedFallback(t *testing.T) { + model := Model{ + CatwalkCfg: catwalk.Model{ID: "llama3"}, + ModelCfg: config.SelectedModel{ + Provider: "ollama", + TopK: ptr(int64(40)), + ProviderOptions: map[string]any{"user": 5.0}, + }, + } + providerCfg := config.ProviderConfig{ID: "test", Type: "ollama"} + + opts := getProviderOptions(model, providerCfg) + + raw, ok := opts[openaicompat.Name] + require.True(t, ok, "malformed provider_options should still fall back to top_k") + parsed, ok := raw.(*openaicompat.ProviderOptions) + require.True(t, ok) + + // The malformed fields are dropped; only the injected top_k survives. + assert.Nil(t, parsed.User) + assert.Nil(t, parsed.ReasoningEffort) + require.Len(t, parsed.ExtraBody, 1) + topK, ok := parsed.ExtraBody["top_k"].(int64) + require.True(t, ok) + assert.Equal(t, int64(40), topK) +} + +func TestCallTopK(t *testing.T) { + knownCustomProviderCfg := config.ProviderConfig{ID: "ollama", Type: "ollama"} + + tests := []struct { + name string + providerCfg config.ProviderConfig + want *int64 + }{ + { + name: "suppressed for known custom providers", + providerCfg: knownCustomProviderCfg, + want: nil, + }, + { + name: "passed through for openai-compat hosted providers", + providerCfg: config.ProviderConfig{ID: "zai", Type: openaicompat.Name}, + want: ptr(int64(40)), + }, + { + name: "passed through for anthropic", + providerCfg: config.ProviderConfig{ID: "anthropic", Type: catwalk.Type(anthropic.Name)}, + want: ptr(int64(40)), + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := callTopK(tc.providerCfg, ptr(int64(40))) + if tc.want == nil { + assert.Nil(t, got) + return + } + require.NotNil(t, got) + assert.Equal(t, *tc.want, *got) + }) + } + + t.Run("nil input stays nil regardless of provider", func(t *testing.T) { + assert.Nil(t, callTopK(knownCustomProviderCfg, nil)) + assert.Nil(t, callTopK(config.ProviderConfig{Type: openaicompat.Name}, nil)) + }) +} + +func ptr[T any](v T) *T { return &v } diff --git a/internal/agent/dispatch_cancel_test.go b/internal/agent/dispatch_cancel_test.go new file mode 100644 index 0000000000000000000000000000000000000000..1e8a977a70ebd6cfc18c21046fb557fb105dfe0f --- /dev/null +++ b/internal/agent/dispatch_cancel_test.go @@ -0,0 +1,447 @@ +package agent + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + "charm.land/fantasy" + "github.com/charmbracelet/crush/internal/agent/notify" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// finishStreamModel is a minimal fantasy.LanguageModel that streams a +// single text part followed by a normal (FinishReasonStop) finish. It +// is enough to drive sessionAgent.Run through PrepareStep and a clean +// completion without a recorded provider cassette. +type finishStreamModel struct { + text string +} + +func (m *finishStreamModel) Provider() string { return "fake" } +func (m *finishStreamModel) Model() string { return "fake-model" } + +func (m *finishStreamModel) Generate(ctx context.Context, call fantasy.Call) (*fantasy.Response, error) { + return &fantasy.Response{ + Content: fantasy.ResponseContent{fantasy.TextContent{Text: m.text}}, + FinishReason: fantasy.FinishReasonStop, + }, nil +} + +func (m *finishStreamModel) Stream(ctx context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { + text := m.text + return func(yield func(fantasy.StreamPart) bool) { + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextStart, ID: "1"}) { + return + } + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, ID: "1", Delta: text}) { + return + } + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextEnd, ID: "1"}) { + return + } + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}) + }, nil +} + +func (m *finishStreamModel) GenerateObject(ctx context.Context, call fantasy.ObjectCall) (*fantasy.ObjectResponse, error) { + return nil, errors.New("not implemented") +} + +func (m *finishStreamModel) StreamObject(ctx context.Context, call fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) { + return nil, errors.New("not implemented") +} + +func newStreamTestAgent(t *testing.T) (*sessionAgent, fakeEnv) { + t.Helper() + env := testEnv(t) + model := &finishStreamModel{text: "done"} + sa := testSessionAgent(env, model, model, "system").(*sessionAgent) + return sa, env +} + +// TestCancel_ActiveAndAcceptedFiresBothBranches covers the case where a +// session is actively running (activeRequests set) AND a follow-up has +// been accepted (acceptedRuns > 0). A single Cancel must fire both: it +// invokes the active cancel func and records a pending cancel for the +// accepted follow-up. +func TestCancel_ActiveAndAcceptedFiresBothBranches(t *testing.T) { + t.Parallel() + sa, _ := newCancelTestAgent(t) + + const sid = "sid" + var activeCanceled atomic.Bool + sa.activeRequests.Set(sid, &activeCancel{cancel: func() { activeCanceled.Store(true) }}) + + accept := sa.BeginAccepted(sid) + defer accept.Close() + + sa.Cancel(sid) + + require.True(t, activeCanceled.Load(), "active cancel func must fire") + require.True(t, sa.hasPendingCancel(sid), "accepted follow-up must record a pending cancel") +} + +// TestRun_BusyWithPendingCancelTakesCancelOnEntry covers the busy-queue +// branch consulting pendingCancels: when the session is busy AND a +// cancel has been recorded for an accepted follow-up, Run must take the +// cancel-on-entry path (persist a canceled turn) instead of enqueueing +// the call behind the active run. +func TestRun_BusyWithPendingCancelTakesCancelOnEntry(t *testing.T) { + t.Parallel() + sa, env := newCancelTestAgent(t) + + sess, err := env.sessions.Create(t.Context(), "session") + require.NoError(t, err) + + // Make the session look busy: an earlier prompt is active. + sa.activeRequests.Set(sess.ID, &activeCancel{cancel: func() {}}) + + accept := sa.BeginAccepted(sess.ID) + // A cancel arrives while this follow-up is accepted-but-not-active. + sa.Cancel(sess.ID) + require.True(t, sa.hasPendingCancel(sess.ID)) + + result, err := sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "follow-up", + Accepted: accept, + }) + require.NoError(t, err) + require.Nil(t, result) + + // The follow-up was canceled on entry, not enqueued. + require.Equal(t, 0, sa.QueuedPrompts(sess.ID), + "cancel-on-entry must not enqueue the follow-up behind the active run") + require.False(t, sa.hasPendingCancel(sess.ID), "pending cancel must be consumed") + require.Equal(t, 0, sa.acceptedCount(sess.ID), "accept reservation must be released") + + msgs, err := env.messages.List(t.Context(), sess.ID) + require.NoError(t, err) + require.Len(t, msgs, 2) + assert.Equal(t, message.User, msgs[0].Role) + assert.Equal(t, message.Assistant, msgs[1].Role) + assert.Equal(t, message.FinishReasonCanceled, msgs[1].FinishReason()) +} + +// TestRun_PrepareStepDrainSkipsQueuedOnPendingCancel verifies that the +// queue drain inside PrepareStep skips queued follow-up prompts when a +// cancel has been recorded for the session: the queued prompt must not +// be folded into the active turn as an extra user message. +func TestRun_PrepareStepDrainSkipsQueuedOnPendingCancel(t *testing.T) { + t.Parallel() + sa, env := newStreamTestAgent(t) + + sess, err := env.sessions.Create(t.Context(), "session") + require.NoError(t, err) + + // A follow-up prompt sits queued for the session. + sa.enqueueCall(SessionAgentCall{SessionID: sess.ID, Prompt: "queued-followup"}) + // A cancel was recorded for the session while it sat in the queue. + sa.cancelMark.Set(sess.ID, 1) + + result, err := sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "main", + }) + require.NoError(t, err) + require.NotNil(t, result) + + // Only the main prompt produced a user message; the queued + // follow-up was skipped, not folded into the turn. + msgs, err := env.messages.List(t.Context(), sess.ID) + require.NoError(t, err) + var userMsgs []message.Message + for _, m := range msgs { + if m.Role == message.User { + userMsgs = append(userMsgs, m) + } + } + require.Len(t, userMsgs, 1, "queued follow-up must not create a user message") + assert.Equal(t, "main", userMsgs[0].Content().String()) + + // The queue was drained and the pending cancel consumed. + require.Equal(t, 0, sa.QueuedPrompts(sess.ID)) + require.False(t, sa.hasPendingCancel(sess.ID)) +} + +// TestRun_NormalCompletionClearsStalePendingCancel verifies that a Run +// which completes normally clears any stale pending-cancel entry for the +// session, so it cannot catch a future run. +func TestRun_NormalCompletionClearsStalePendingCancel(t *testing.T) { + t.Parallel() + sa, env := newStreamTestAgent(t) + + sess, err := env.sessions.Create(t.Context(), "session") + require.NoError(t, err) + + // A stale cancel mark lingers (no queued work, no accepted run). + sa.cancelMark.Set(sess.ID, 1) + + result, err := sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "main", + }) + require.NoError(t, err) + require.NotNil(t, result) + + require.False(t, sa.hasPendingCancel(sess.ID), + "normal completion must clear the stale pending cancel") + + msgs, err := env.messages.List(t.Context(), sess.ID) + require.NoError(t, err) + require.Len(t, msgs, 2) + assert.Equal(t, message.Assistant, msgs[1].Role) + assert.Equal(t, message.FinishReasonEndTurn, msgs[1].FinishReason()) +} + +// newCancelTestAgentWithRunComplete builds a DB-backed sessionAgent wired +// to a RunComplete broker so tests can observe the terminal event a +// RunID-bearing caller (e.g. `crush run`) blocks on. +func newCancelTestAgentWithRunComplete(t *testing.T) (*sessionAgent, fakeEnv, *pubsub.Broker[notify.RunComplete]) { + t.Helper() + env := testEnv(t) + broker := pubsub.NewBroker[notify.RunComplete]() + t.Cleanup(broker.Shutdown) + sa := NewSessionAgent(SessionAgentOptions{ + Sessions: env.sessions, + Messages: env.messages, + RunComplete: broker, + }).(*sessionAgent) + return sa, env, broker +} + +// TestRun_CancelOnEntryPublishesRunComplete covers the first review +// finding: the cancel-on-entry path returned before the streaming defer +// that publishes RunComplete was installed. A caller that dispatches a +// run with a RunID and blocks on RunComplete (ignoring message events, +// like `crush run`) would hang on an immediately-canceled accepted run. +// The cancel-on-entry path must publish a terminal RunComplete carrying +// the originating RunID. +func TestRun_CancelOnEntryPublishesRunComplete(t *testing.T) { + t.Parallel() + sa, env, broker := newCancelTestAgentWithRunComplete(t) + + sess, err := env.sessions.Create(t.Context(), "session") + require.NoError(t, err) + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + ch := broker.Subscribe(ctx) + + accept := sa.BeginAccepted(sess.ID) + // A cancel arrives in the accepted-but-not-yet-active window. + sa.Cancel(sess.ID) + require.True(t, sa.hasPendingCancel(sess.ID)) + + result, err := sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + RunID: "run-cancel-on-entry", + Prompt: "hello", + Accepted: accept, + }) + require.NoError(t, err) + require.Nil(t, result) + + select { + case got := <-ch: + assert.Equal(t, "run-cancel-on-entry", got.Payload.RunID, + "RunComplete must echo the originating RunID") + assert.Equal(t, sess.ID, got.Payload.SessionID) + assert.True(t, got.Payload.Cancelled, + "cancel-on-entry RunComplete must be marked cancelled") + case <-time.After(2 * time.Second): + t.Fatal("cancel-on-entry must publish RunComplete; a RunID caller would hang otherwise") + } +} + +// TestCancel_TwoAcceptedBothObserveCancellation covers the second review +// finding: a single cancel with two accepted-not-yet-active prompts must +// cancel both. The cancel raises the session's high-water mark to the +// latest accept sequence, so every prompt accepted-but-not-yet-active at +// cancel time is covered and both take the cancel-on-entry path. +func TestCancel_TwoAcceptedBothObserveCancellation(t *testing.T) { + t.Parallel() + sa, env := newCancelTestAgent(t) + + sess, err := env.sessions.Create(t.Context(), "session") + require.NoError(t, err) + + // Two prompts are accepted-but-not-yet-active for the same session. + accept1 := sa.BeginAccepted(sess.ID) + accept2 := sa.BeginAccepted(sess.ID) + require.Equal(t, 2, sa.acceptedCount(sess.ID)) + + // A single cancel arrives before either becomes active. + sa.Cancel(sess.ID) + require.Equal(t, accept2.seq, sa.pendingCancelMark(sess.ID), + "one cancel must mark every currently-accepted prompt as canceled") + require.GreaterOrEqual(t, sa.pendingCancelMark(sess.ID), accept1.seq, + "the mark must cover the earlier accepted prompt too") + + // Both prompts enter Run; each must take cancel-on-entry, not run. + r1, err := sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "first", + Accepted: accept1, + }) + require.NoError(t, err) + require.Nil(t, r1) + + r2, err := sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "second", + Accepted: accept2, + }) + require.NoError(t, err) + require.Nil(t, r2) + + require.False(t, sa.hasPendingCancel(sess.ID), + "both reserved units must be consumed") + require.Equal(t, 0, sa.acceptedCount(sess.ID), + "both accept reservations must be released") + + // Each canceled-on-entry turn writes a user + canceled assistant + // message, and neither prompt was enqueued to run normally. + require.Equal(t, 0, sa.QueuedPrompts(sess.ID), + "neither accepted prompt may be enqueued to run normally") + msgs, err := env.messages.List(t.Context(), sess.ID) + require.NoError(t, err) + require.Len(t, msgs, 4, "two canceled turns produce two user + two assistant messages") + var canceled int + for _, m := range msgs { + if m.Role == message.Assistant { + assert.Equal(t, message.FinishReasonCanceled, m.FinishReason()) + canceled++ + } + } + require.Equal(t, 2, canceled, "both turns must finish canceled") +} + +// TestRun_IdleCancelDoesNotPoisonNextPrompt covers the idle-cancel +// no-op guarantee end-to-end: an Escape on an idle session must not +// record a pending cancel that leaks into the next accepted prompt, which +// must run normally to completion. +func TestRun_IdleCancelDoesNotPoisonNextPrompt(t *testing.T) { + t.Parallel() + sa, env := newStreamTestAgent(t) + + sess, err := env.sessions.Create(t.Context(), "session") + require.NoError(t, err) + + // Idle Escape: no accepted run, no active request. + sa.Cancel(sess.ID) + require.False(t, sa.hasPendingCancel(sess.ID), + "idle cancel must not record a pending cancel") + + // The next accepted prompt must run normally, not cancel on entry. + accept := sa.BeginAccepted(sess.ID) + result, err := sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "next", + Accepted: accept, + }) + require.NoError(t, err) + require.NotNil(t, result, "next prompt must run normally after an idle cancel") + + msgs, err := env.messages.List(t.Context(), sess.ID) + require.NoError(t, err) + require.Len(t, msgs, 2) + assert.Equal(t, message.User, msgs[0].Role) + assert.Equal(t, message.Assistant, msgs[1].Role) + assert.Equal(t, message.FinishReasonEndTurn, msgs[1].FinishReason(), + "the prompt must finish normally, not canceled") +} + +// TestCancel_AcceptedAfterCancelIsNotPoisoned is the regression test for +// the reviewer's finding: a counted session-level pending cancel let a +// prompt accepted after the cancel enter Run first and consume a unit +// reserved for the earlier prompts. With a sequence high-water mark, a +// single cancel covers exactly the prompts accepted-but-not-yet-active at +// cancel time (A and B); a prompt accepted afterwards (C) gets a higher +// sequence and must run normally without consuming A or B's cancellation. +// C is run first to prove it neither cancels nor drains the mark, then A +// and B are run and must both cancel on entry. +func TestCancel_AcceptedAfterCancelIsNotPoisoned(t *testing.T) { + t.Parallel() + sa, env := newStreamTestAgent(t) + + sess, err := env.sessions.Create(t.Context(), "session") + require.NoError(t, err) + + // A and B are accepted-but-not-yet-active. + acceptA := sa.BeginAccepted(sess.ID) + acceptB := sa.BeginAccepted(sess.ID) + + // One cancel arrives covering both A and B. + sa.Cancel(sess.ID) + require.True(t, sa.hasPendingCancel(sess.ID)) + require.Equal(t, acceptB.seq, sa.pendingCancelMark(sess.ID), + "the mark must cover every prompt accepted before the cancel") + + // C is accepted AFTER the cancel; its sequence is above the mark. + acceptC := sa.BeginAccepted(sess.ID) + require.Greater(t, acceptC.seq, sa.pendingCancelMark(sess.ID), + "a prompt accepted after the cancel must not be covered by the mark") + + // Run C first. It must run normally to completion and must not + // consume or clear the cancellation reserved for A and B. + rc, err := sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "C", + Accepted: acceptC, + }) + require.NoError(t, err) + require.NotNil(t, rc, "C was accepted after the cancel and must run normally") + require.True(t, sa.hasPendingCancel(sess.ID), + "running C must not drain the cancellation reserved for A and B") + + // Now A and B run. Both were accepted before the cancel and must + // take the cancel-on-entry path. + ra, err := sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "A", + Accepted: acceptA, + }) + require.NoError(t, err) + require.Nil(t, ra, "A must cancel on entry, not run") + + rb, err := sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "B", + Accepted: acceptB, + }) + require.NoError(t, err) + require.Nil(t, rb, "B must cancel on entry, not run") + + require.False(t, sa.hasPendingCancel(sess.ID), + "the mark clears once all covered handles are resolved") + require.Equal(t, 0, sa.acceptedCount(sess.ID)) + require.Equal(t, 0, sa.QueuedPrompts(sess.ID), + "neither A nor B may be enqueued to run normally") + + // C produced a normal turn; A and B each produced a canceled turn. + msgs, err := env.messages.List(t.Context(), sess.ID) + require.NoError(t, err) + require.Len(t, msgs, 6, "C normal + A canceled + B canceled = 3 user + 3 assistant") + + var normal, canceled int + for _, m := range msgs { + if m.Role != message.Assistant { + continue + } + switch m.FinishReason() { + case message.FinishReasonEndTurn: + normal++ + case message.FinishReasonCanceled: + canceled++ + } + } + require.Equal(t, 1, normal, "only C finished normally") + require.Equal(t, 2, canceled, "both A and B finished canceled") +} diff --git a/internal/agent/dispatch_race_test.go b/internal/agent/dispatch_race_test.go new file mode 100644 index 0000000000000000000000000000000000000000..95eb98afe4ca52f1942329311178dd3e69d00997 --- /dev/null +++ b/internal/agent/dispatch_race_test.go @@ -0,0 +1,155 @@ +package agent + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "charm.land/fantasy" + "github.com/stretchr/testify/require" +) + +// concurrencyProbeModel records the maximum number of Stream iterators in +// flight at once. Each stream blocks on release, so a second dispatch that +// wrongly started a concurrent run for the same session is observable as an +// in-flight count above one. +type concurrencyProbeModel struct { + inFlight atomic.Int32 + maxSeen atomic.Int32 + entered chan struct{} + release chan struct{} +} + +func (m *concurrencyProbeModel) Provider() string { return "fake" } +func (m *concurrencyProbeModel) Model() string { return "fake-model" } + +func (m *concurrencyProbeModel) Generate(context.Context, fantasy.Call) (*fantasy.Response, error) { + return &fantasy.Response{ + Content: fantasy.ResponseContent{fantasy.TextContent{Text: "done"}}, + FinishReason: fantasy.FinishReasonStop, + }, nil +} + +func (m *concurrencyProbeModel) Stream(context.Context, fantasy.Call) (fantasy.StreamResponse, error) { + return func(yield func(fantasy.StreamPart) bool) { + cur := m.inFlight.Add(1) + for { + mx := m.maxSeen.Load() + if cur <= mx || m.maxSeen.CompareAndSwap(mx, cur) { + break + } + } + // Signal that a stream is in flight (non-blocking), then hold here + // so a racing second dispatch would be caught by maxSeen. + select { + case m.entered <- struct{}{}: + default: + } + <-m.release + m.inFlight.Add(-1) + + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextStart, ID: "1"}) + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, ID: "1", Delta: "done"}) + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextEnd, ID: "1"}) + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}) + }, nil +} + +func (m *concurrencyProbeModel) GenerateObject(context.Context, fantasy.ObjectCall) (*fantasy.ObjectResponse, error) { + return nil, errors.New("not implemented") +} + +func (m *concurrencyProbeModel) StreamObject(context.Context, fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) { + return nil, errors.New("not implemented") +} + +// fastModel is a non-blocking model used as the small model in concurrency +// tests. GenerateTitle runs on the small model; if it shares the probe model, +// its Stream call races into inFlight/maxSeen and produces spurious failures. +type fastModel struct{} + +func (fastModel) Provider() string { return "fake" } +func (fastModel) Model() string { return "fake-model" } + +func (fastModel) Generate(context.Context, fantasy.Call) (*fantasy.Response, error) { + return &fantasy.Response{ + Content: fantasy.ResponseContent{fantasy.TextContent{Text: "title"}}, + FinishReason: fantasy.FinishReasonStop, + }, nil +} + +func (fastModel) Stream(context.Context, fantasy.Call) (fantasy.StreamResponse, error) { + return func(yield func(fantasy.StreamPart) bool) { + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextStart, ID: "1"}) + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, ID: "1", Delta: "title"}) + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextEnd, ID: "1"}) + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}) + }, nil +} + +func (fastModel) GenerateObject(context.Context, fantasy.ObjectCall) (*fantasy.ObjectResponse, error) { + return nil, errors.New("not implemented") +} + +func (fastModel) StreamObject(context.Context, fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) { + return nil, errors.New("not implemented") +} + +// TestRun_ConcurrentInProcessDispatchStartsOneRun fires a burst of concurrent +// in-process Run calls (the path channel events use) at an idle session. Only +// one may become the active run; the rest must queue behind it. Before the +// dispatch decision was serialized under the per-session mutex, two callers +// could both pass the busy check and start two runs on the same session — this +// test catches that regression (maxSeen would exceed one). +// +// fastModel is used as the small model so GenerateTitle (which runs on the +// small model) does not race into the probe's inFlight/maxSeen counters. The +// queue count is not asserted because PrepareStep drains queued prompts into +// the active step before the model's Stream is called — by the time "entered" +// fires the queue is already empty by design. +func TestRun_ConcurrentInProcessDispatchStartsOneRun(t *testing.T) { + t.Parallel() + env := testEnv(t) + model := &concurrencyProbeModel{ + entered: make(chan struct{}, 1), + release: make(chan struct{}), + } + sa := testSessionAgent(env, model, fastModel{}, "system").(*sessionAgent) + + sess, err := env.sessions.Create(t.Context(), "session") + require.NoError(t, err) + + const n = 8 + var wg sync.WaitGroup + for range n { + wg.Add(1) + go func() { + defer wg.Done() + _, _ = sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + Prompt: "event", + }) + }() + } + + // Wait until the active run's Stream is in flight (blocked on release). + select { + case <-model.entered: + case <-time.After(5 * time.Second): + close(model.release) + wg.Wait() + t.Fatal("no run became active") + } + + // Every other dispatch must have either queued (and been folded into the + // active step by PrepareStep) or never started its own Stream. Either way, + // at most one Stream may be in flight and the high-water mark must be one. + require.Equal(t, int32(1), model.inFlight.Load(), "exactly one run may be active") + require.Equal(t, int32(1), model.maxSeen.Load(), "no two runs may stream concurrently for one session") + + close(model.release) + wg.Wait() +} diff --git a/internal/agent/errors.go b/internal/agent/errors.go new file mode 100644 index 0000000000000000000000000000000000000000..c6cce5edd230ae421abad7b856f42aad6f2119f7 --- /dev/null +++ b/internal/agent/errors.go @@ -0,0 +1,10 @@ +package agent + +import "errors" + +var ( + ErrRequestCancelled = errors.New("request canceled by user") + ErrSessionBusy = errors.New("session is currently processing another request") + ErrEmptyPrompt = errors.New("prompt is empty") + ErrSessionMissing = errors.New("session id is missing") +) diff --git a/internal/agent/event.go b/internal/agent/event.go new file mode 100644 index 0000000000000000000000000000000000000000..3f6c640f6a983c515034e0698676632d0cb57824 --- /dev/null +++ b/internal/agent/event.go @@ -0,0 +1,51 @@ +package agent + +import ( + "time" + + "charm.land/fantasy" + "github.com/charmbracelet/crush/internal/event" +) + +func (a *sessionAgent) eventPromptSent(sessionID string) { + event.PromptSent( + a.eventCommon(sessionID, a.largeModel.Get())..., + ) +} + +func (a *sessionAgent) eventPromptResponded(sessionID string, duration time.Duration) { + event.PromptResponded( + append( + a.eventCommon(sessionID, a.largeModel.Get()), + "prompt duration pretty", duration.String(), + "prompt duration in seconds", int64(duration.Seconds()), + )..., + ) +} + +func (a *sessionAgent) eventTokensUsed(sessionID string, model Model, usage fantasy.Usage, cost float64) { + event.TokensUsed( + append( + a.eventCommon(sessionID, model), + "input tokens", usage.InputTokens, + "output tokens", usage.OutputTokens, + "cache read tokens", usage.CacheReadTokens, + "cache creation tokens", usage.CacheCreationTokens, + "total tokens", usage.InputTokens+usage.OutputTokens+usage.CacheReadTokens+usage.CacheCreationTokens, + "cost", cost, + )..., + ) +} + +func (a *sessionAgent) eventCommon(sessionID string, model Model) []any { + m := model.ModelCfg + + return []any{ + "session id", sessionID, + "provider", m.Provider, + "model", m.Model, + "reasoning effort", m.ReasoningEffort, + "thinking mode", m.Think, + "yolo mode", a.isYolo, + } +} diff --git a/internal/agent/hooked_tool.go b/internal/agent/hooked_tool.go new file mode 100644 index 0000000000000000000000000000000000000000..2c38d66961ba9135f54d8aab9e08290b0d7a1e81 --- /dev/null +++ b/internal/agent/hooked_tool.go @@ -0,0 +1,142 @@ +package agent + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + + "charm.land/fantasy" + "github.com/charmbracelet/crush/internal/agent/tools" + "github.com/charmbracelet/crush/internal/hooks" + "github.com/charmbracelet/crush/internal/permission" + "github.com/tidwall/sjson" +) + +// hookedTool wraps a fantasy.AgentTool to run PreToolUse hooks before +// delegating to the inner tool. +type hookedTool struct { + inner fantasy.AgentTool + runner *hooks.Runner +} + +func newHookedTool(inner fantasy.AgentTool, runner *hooks.Runner) *hookedTool { + return &hookedTool{inner: inner, runner: runner} +} + +// wrapToolsWithHooks returns a tool slice with each entry wrapped in a +// hookedTool. Returns the original slice unchanged when runner is nil or +// when isSubAgent is true — sub-agents never fire hooks, the top-level +// invocation of the sub-agent tool itself is wrapped on the caller's side. +func wrapToolsWithHooks(tools []fantasy.AgentTool, runner *hooks.Runner, isSubAgent bool) []fantasy.AgentTool { + if runner == nil || isSubAgent { + return tools + } + out := make([]fantasy.AgentTool, len(tools)) + for i, tool := range tools { + out[i] = newHookedTool(tool, runner) + } + return out +} + +func (h *hookedTool) Info() fantasy.ToolInfo { + return h.inner.Info() +} + +func (h *hookedTool) ProviderOptions() fantasy.ProviderOptions { + return h.inner.ProviderOptions() +} + +func (h *hookedTool) SetProviderOptions(opts fantasy.ProviderOptions) { + h.inner.SetProviderOptions(opts) +} + +func (h *hookedTool) Run(ctx context.Context, call fantasy.ToolCall) (fantasy.ToolResponse, error) { + sessionID := tools.GetSessionFromContext(ctx) + result, err := h.runner.Run(ctx, hooks.EventPreToolUse, sessionID, call.Name, call.Input) + if err != nil { + slog.Warn("Hook execution error, proceeding with tool call", + "tool", call.Name, "error", err) + } + + if result.Decision == hooks.DecisionDeny || result.Halt { + reason := fmt.Sprintf("Tool call blocked by hook. Reason: %s", result.Reason) + if result.Halt { + reason = fmt.Sprintf("Turn halted by hook. Reason: %s", result.Reason) + } + resp := fantasy.NewTextErrorResponse(reason) + // Halt ends the whole turn; a plain deny only blocks this tool + // call so the model can see the error and try something else. + resp.StopTurn = result.Halt + resp.Metadata = hookMetadataJSON(result) + return resp, nil + } + + if result.UpdatedInput != "" { + call.Input = result.UpdatedInput + } + + // An explicit allow from a hook pre-approves the permission prompt for + // this tool call. Deny is already handled above; silence falls through + // to the normal permission flow. + if result.Decision == hooks.DecisionAllow { + ctx = permission.WithHookApproval(ctx, call.ID) + } + + resp, err := h.inner.Run(ctx, call) + if err != nil { + return resp, err + } + + if result.Context != "" { + if resp.Content != "" { + resp.Content += "\n" + } + resp.Content += result.Context + } + + resp.Metadata = mergeHookMetadata(resp.Metadata, result) + return resp, nil +} + +// buildHookMetadata creates a HookMetadata from an AggregateResult. +func buildHookMetadata(result hooks.AggregateResult) hooks.HookMetadata { + return hooks.HookMetadata{ + HookCount: result.HookCount, + Decision: result.Decision.String(), + Halt: result.Halt, + Reason: result.Reason, + InputRewrite: result.UpdatedInput != "", + Hooks: result.Hooks, + } +} + +// hookMetadataJSON builds a JSON string containing only the hook metadata. +func hookMetadataJSON(result hooks.AggregateResult) string { + meta := buildHookMetadata(result) + data, err := json.Marshal(meta) + if err != nil { + return "" + } + return `{"hook":` + string(data) + `}` +} + +// mergeHookMetadata injects hook metadata into existing tool metadata. +func mergeHookMetadata(existing string, result hooks.AggregateResult) string { + if result.HookCount == 0 { + return existing + } + meta := buildHookMetadata(result) + data, err := json.Marshal(meta) + if err != nil { + return existing + } + if existing == "" { + existing = "{}" + } + merged, err := sjson.SetRaw(existing, "hook", string(data)) + if err != nil { + return existing + } + return merged +} diff --git a/internal/agent/hooked_tool_test.go b/internal/agent/hooked_tool_test.go new file mode 100644 index 0000000000000000000000000000000000000000..5a7823206f4bd28bb19b45d6781bdc58ec411697 --- /dev/null +++ b/internal/agent/hooked_tool_test.go @@ -0,0 +1,147 @@ +package agent + +import ( + "context" + "testing" + + "charm.land/fantasy" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/hooks" + "github.com/charmbracelet/crush/internal/permission" + "github.com/stretchr/testify/require" +) + +// fakeTool records the context it was invoked with so tests can assert on +// values stamped onto it by the hookedTool decorator. +type fakeTool struct { + name string + called bool + gotCtx context.Context + resp fantasy.ToolResponse +} + +func (f *fakeTool) Info() fantasy.ToolInfo { + return fantasy.ToolInfo{Name: f.name} +} + +func (f *fakeTool) Run(ctx context.Context, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + f.called = true + f.gotCtx = ctx + return f.resp, nil +} + +func (f *fakeTool) ProviderOptions() fantasy.ProviderOptions { return nil } +func (f *fakeTool) SetProviderOptions(_ fantasy.ProviderOptions) {} + +// newRunner builds a hooks.Runner from a single HookConfig, running the +// config-loader path that compiles the matcher regex. +func newRunner(t *testing.T, cmd string) *hooks.Runner { + t.Helper() + cfg := &config.Config{ + Hooks: map[string][]config.HookConfig{ + hooks.EventPreToolUse: {{Command: cmd}}, + }, + } + require.NoError(t, cfg.ValidateHooks()) + return hooks.NewRunner(cfg.Hooks[hooks.EventPreToolUse], t.TempDir(), t.TempDir()) +} + +func TestHookedTool_AllowStampsHookApproval(t *testing.T) { + t.Parallel() + + inner := &fakeTool{name: "view", resp: fantasy.NewTextResponse("ok")} + runner := newRunner(t, `echo '{"decision":"allow"}'`) + tool := newHookedTool(inner, runner) + + _, err := tool.Run(t.Context(), fantasy.ToolCall{ID: "call-1", Name: "view"}) + require.NoError(t, err) + require.True(t, inner.called, "inner tool should have run") + + // The inner tool's permission service can now treat call-1 as pre-approved. + svc := permission.NewPermissionService(t.TempDir(), false, nil) + granted, err := svc.Request(inner.gotCtx, permission.CreatePermissionRequest{ + SessionID: "s1", + ToolCallID: "call-1", + ToolName: "view", + Action: "read", + Path: t.TempDir(), + }) + require.NoError(t, err) + require.True(t, granted, "hook allow should bypass the permission prompt") +} + +func TestHookedTool_SilentDoesNotStampApproval(t *testing.T) { + t.Parallel() + + inner := &fakeTool{name: "view", resp: fantasy.NewTextResponse("ok")} + runner := newRunner(t, `exit 0`) // no stdout, no decision + tool := newHookedTool(inner, runner) + + _, err := tool.Run(t.Context(), fantasy.ToolCall{ID: "call-2", Name: "view"}) + require.NoError(t, err) + require.True(t, inner.called) + + // With no hook opinion, a fresh permission request has nothing stamped + // and must fall through to the normal flow. We verify by checking that + // the context does not look pre-approved for this call ID: sending a + // request that no subscriber resolves will block until cancelled. + svc := permission.NewPermissionService(t.TempDir(), false, nil) + ctx, cancel := context.WithCancel(inner.gotCtx) + cancel() + granted, err := svc.Request(ctx, permission.CreatePermissionRequest{ + SessionID: "s1", + ToolCallID: "call-2", + ToolName: "view", + Action: "read", + Path: t.TempDir(), + }) + require.Error(t, err, "no approval stamped => request should reach the prompt path") + require.False(t, granted) +} + +func TestHookedTool_DenySkipsInnerTool(t *testing.T) { + t.Parallel() + + inner := &fakeTool{name: "bash"} + runner := newRunner(t, `echo "blocked" >&2; exit 2`) + tool := newHookedTool(inner, runner) + + resp, err := tool.Run(t.Context(), fantasy.ToolCall{ID: "call-3", Name: "bash"}) + require.NoError(t, err) + require.False(t, inner.called, "denied call must not reach the inner tool") + require.True(t, resp.IsError) + require.Contains(t, resp.Content, "blocked") +} + +func TestWrapToolsWithHooks(t *testing.T) { + t.Parallel() + + runner := newRunner(t, `exit 0`) + inputs := []fantasy.AgentTool{&fakeTool{name: "a"}, &fakeTool{name: "b"}} + + t.Run("top-level agent wraps every tool", func(t *testing.T) { + t.Parallel() + out := wrapToolsWithHooks(inputs, runner, false) + require.Len(t, out, len(inputs)) + for i, tool := range out { + _, ok := tool.(*hookedTool) + require.Truef(t, ok, "tool %d should be a *hookedTool", i) + } + }) + + t.Run("sub-agent skips the wrap", func(t *testing.T) { + t.Parallel() + out := wrapToolsWithHooks(inputs, runner, true) + require.Equal(t, inputs, out, "sub-agent tools should be returned unwrapped") + for _, tool := range out { + _, isHooked := tool.(*hookedTool) + require.False(t, isHooked, "sub-agent tool should not be wrapped") + } + }) + + t.Run("nil runner skips the wrap for both agent kinds", func(t *testing.T) { + t.Parallel() + require.Equal(t, inputs, wrapToolsWithHooks(inputs, nil, false)) + require.Equal(t, inputs, wrapToolsWithHooks(inputs, nil, true)) + }) +} diff --git a/internal/agent/loop_detection.go b/internal/agent/loop_detection.go new file mode 100644 index 0000000000000000000000000000000000000000..673a0c196c551f8cbc5355f2d89e08dfc73b986e --- /dev/null +++ b/internal/agent/loop_detection.go @@ -0,0 +1,92 @@ +package agent + +import ( + "crypto/sha256" + "encoding/hex" + "io" + + "charm.land/fantasy" +) + +const ( + loopDetectionWindowSize = 10 + loopDetectionMaxRepeats = 5 +) + +// hasRepeatedToolCalls checks whether the agent is stuck in a loop by looking +// at recent steps. It examines the last windowSize steps and returns true if +// any tool-call signature appears more than maxRepeats times. +func hasRepeatedToolCalls(steps []fantasy.StepResult, windowSize, maxRepeats int) bool { + if len(steps) < windowSize { + return false + } + + window := steps[len(steps)-windowSize:] + counts := make(map[string]int) + + for _, step := range window { + sig := getToolInteractionSignature(step.Content) + if sig == "" { + continue + } + counts[sig]++ + if counts[sig] > maxRepeats { + return true + } + } + + return false +} + +// getToolInteractionSignature computes a hash signature for the tool +// interactions in a single step's content. It pairs tool calls with their +// results (matched by ToolCallID) and returns a hex-encoded SHA-256 hash. +// If the step contains no tool calls, it returns "". +func getToolInteractionSignature(content fantasy.ResponseContent) string { + toolCalls := content.ToolCalls() + if len(toolCalls) == 0 { + return "" + } + + // Index tool results by their ToolCallID for fast lookup. + resultsByID := make(map[string]fantasy.ToolResultContent) + for _, tr := range content.ToolResults() { + resultsByID[tr.ToolCallID] = tr + } + + h := sha256.New() + for _, tc := range toolCalls { + output := "" + if tr, ok := resultsByID[tc.ToolCallID]; ok { + output = toolResultOutputString(tr.Result) + } + io.WriteString(h, tc.ToolName) + io.WriteString(h, "\x00") + io.WriteString(h, tc.Input) + io.WriteString(h, "\x00") + io.WriteString(h, output) + io.WriteString(h, "\x00") + } + return hex.EncodeToString(h.Sum(nil)) +} + +// toolResultOutputString converts a ToolResultOutputContent to a stable string +// representation for signature comparison. +func toolResultOutputString(result fantasy.ToolResultOutputContent) string { + if result == nil { + return "" + } + if text, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentText](result); ok { + return text.Text + } + if errResult, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentError](result); ok { + if errResult.Error != nil { + return errResult.Error.Error() + } + return "" + } + if media, ok := fantasy.AsToolResultOutputType[fantasy.ToolResultOutputContentMedia](result); ok { + return media.Data + } + return "" +} diff --git a/internal/agent/loop_detection_test.go b/internal/agent/loop_detection_test.go new file mode 100644 index 0000000000000000000000000000000000000000..9e0405e09c91103846ae8fddf0f10fda5a4bf3c2 --- /dev/null +++ b/internal/agent/loop_detection_test.go @@ -0,0 +1,205 @@ +package agent + +import ( + "fmt" + "testing" + + "charm.land/fantasy" +) + +// makeStep creates a StepResult with the given tool calls and results in its Content. +func makeStep(calls []fantasy.ToolCallContent, results []fantasy.ToolResultContent) fantasy.StepResult { + var content fantasy.ResponseContent + for _, c := range calls { + content = append(content, c) + } + for _, r := range results { + content = append(content, r) + } + return fantasy.StepResult{ + Response: fantasy.Response{ + Content: content, + }, + } +} + +// makeToolStep creates a step with a single tool call and matching text result. +func makeToolStep(name, input, output string) fantasy.StepResult { + callID := fmt.Sprintf("call_%s_%s", name, input) + return makeStep( + []fantasy.ToolCallContent{ + {ToolCallID: callID, ToolName: name, Input: input}, + }, + []fantasy.ToolResultContent{ + {ToolCallID: callID, ToolName: name, Result: fantasy.ToolResultOutputContentText{Text: output}}, + }, + ) +} + +// makeEmptyStep creates a step with no tool calls (e.g. a text-only response). +func makeEmptyStep() fantasy.StepResult { + return fantasy.StepResult{ + Response: fantasy.Response{ + Content: fantasy.ResponseContent{ + fantasy.TextContent{Text: "thinking..."}, + }, + }, + } +} + +func TestHasRepeatedToolCalls(t *testing.T) { + t.Run("no steps", func(t *testing.T) { + result := hasRepeatedToolCalls(nil, 10, 5) + if result { + t.Error("expected false for empty steps") + } + }) + + t.Run("fewer steps than window", func(t *testing.T) { + steps := make([]fantasy.StepResult, 5) + for i := range steps { + steps[i] = makeToolStep("read", `{"file":"a.go"}`, "content") + } + result := hasRepeatedToolCalls(steps, 10, 5) + if result { + t.Error("expected false when fewer steps than window size") + } + }) + + t.Run("all different signatures", func(t *testing.T) { + steps := make([]fantasy.StepResult, 10) + for i := range steps { + steps[i] = makeToolStep("tool", fmt.Sprintf(`{"i":%d}`, i), fmt.Sprintf("result-%d", i)) + } + result := hasRepeatedToolCalls(steps, 10, 5) + if result { + t.Error("expected false when all signatures are different") + } + }) + + t.Run("exact repeat at threshold not detected", func(t *testing.T) { + // maxRepeats=5 means > 5 is needed, so exactly 5 should return false + steps := make([]fantasy.StepResult, 10) + for i := range 5 { + steps[i] = makeToolStep("read", `{"file":"a.go"}`, "content") + } + for i := 5; i < 10; i++ { + steps[i] = makeToolStep("tool", fmt.Sprintf(`{"i":%d}`, i), fmt.Sprintf("result-%d", i)) + } + result := hasRepeatedToolCalls(steps, 10, 5) + if result { + t.Error("expected false when count equals maxRepeats (threshold is >)") + } + }) + + t.Run("loop detected", func(t *testing.T) { + // 6 identical steps in a window of 10 with maxRepeats=5 → detected + steps := make([]fantasy.StepResult, 10) + for i := range 6 { + steps[i] = makeToolStep("read", `{"file":"a.go"}`, "content") + } + for i := 6; i < 10; i++ { + steps[i] = makeToolStep("tool", fmt.Sprintf(`{"i":%d}`, i), fmt.Sprintf("result-%d", i)) + } + result := hasRepeatedToolCalls(steps, 10, 5) + if !result { + t.Error("expected true when same signature appears more than maxRepeats times") + } + }) + + t.Run("steps without tool calls are skipped", func(t *testing.T) { + // Mix of tool steps and empty steps — empty ones should not affect counts + steps := make([]fantasy.StepResult, 10) + for i := range 4 { + steps[i] = makeToolStep("read", `{"file":"a.go"}`, "content") + } + for i := 4; i < 8; i++ { + steps[i] = makeEmptyStep() + } + for i := 8; i < 10; i++ { + steps[i] = makeToolStep("write", `{"file":"b.go"}`, "ok") + } + result := hasRepeatedToolCalls(steps, 10, 5) + if result { + t.Error("expected false: only 4 repeated tool calls, empty steps should be skipped") + } + }) + + t.Run("multiple different patterns alternating", func(t *testing.T) { + // Two patterns alternating: each appears 5 times — not above threshold + steps := make([]fantasy.StepResult, 10) + for i := range steps { + if i%2 == 0 { + steps[i] = makeToolStep("read", `{"file":"a.go"}`, "content-a") + } else { + steps[i] = makeToolStep("write", `{"file":"b.go"}`, "content-b") + } + } + result := hasRepeatedToolCalls(steps, 10, 5) + if result { + t.Error("expected false: two patterns each appearing 5 times (not > 5)") + } + }) +} + +func TestGetToolInteractionSignature(t *testing.T) { + t.Run("empty content returns empty string", func(t *testing.T) { + sig := getToolInteractionSignature(fantasy.ResponseContent{}) + if sig != "" { + t.Errorf("expected empty string, got %q", sig) + } + }) + + t.Run("text only content returns empty string", func(t *testing.T) { + content := fantasy.ResponseContent{ + fantasy.TextContent{Text: "hello"}, + } + sig := getToolInteractionSignature(content) + if sig != "" { + t.Errorf("expected empty string, got %q", sig) + } + }) + + t.Run("tool call with result produces signature", func(t *testing.T) { + content := fantasy.ResponseContent{ + fantasy.ToolCallContent{ToolCallID: "1", ToolName: "read", Input: `{"file":"a.go"}`}, + fantasy.ToolResultContent{ToolCallID: "1", ToolName: "read", Result: fantasy.ToolResultOutputContentText{Text: "content"}}, + } + sig := getToolInteractionSignature(content) + if sig == "" { + t.Error("expected non-empty signature") + } + }) + + t.Run("same interactions produce same signature", func(t *testing.T) { + content1 := fantasy.ResponseContent{ + fantasy.ToolCallContent{ToolCallID: "1", ToolName: "read", Input: `{"file":"a.go"}`}, + fantasy.ToolResultContent{ToolCallID: "1", ToolName: "read", Result: fantasy.ToolResultOutputContentText{Text: "content"}}, + } + content2 := fantasy.ResponseContent{ + fantasy.ToolCallContent{ToolCallID: "2", ToolName: "read", Input: `{"file":"a.go"}`}, + fantasy.ToolResultContent{ToolCallID: "2", ToolName: "read", Result: fantasy.ToolResultOutputContentText{Text: "content"}}, + } + sig1 := getToolInteractionSignature(content1) + sig2 := getToolInteractionSignature(content2) + if sig1 != sig2 { + t.Errorf("expected same signature for same interactions, got %q and %q", sig1, sig2) + } + }) + + t.Run("different inputs produce different signatures", func(t *testing.T) { + content1 := fantasy.ResponseContent{ + fantasy.ToolCallContent{ToolCallID: "1", ToolName: "read", Input: `{"file":"a.go"}`}, + fantasy.ToolResultContent{ToolCallID: "1", ToolName: "read", Result: fantasy.ToolResultOutputContentText{Text: "content"}}, + } + content2 := fantasy.ResponseContent{ + fantasy.ToolCallContent{ToolCallID: "1", ToolName: "read", Input: `{"file":"b.go"}`}, + fantasy.ToolResultContent{ToolCallID: "1", ToolName: "read", Result: fantasy.ToolResultOutputContentText{Text: "content"}}, + } + sig1 := getToolInteractionSignature(content1) + sig2 := getToolInteractionSignature(content2) + if sig1 == sig2 { + t.Error("expected different signatures for different inputs") + } + }) +} diff --git a/internal/agent/opencode_routing_test.go b/internal/agent/opencode_routing_test.go new file mode 100644 index 0000000000000000000000000000000000000000..f70b34fecc5a31e231f22ba7ebd56166de33e228 --- /dev/null +++ b/internal/agent/opencode_routing_test.go @@ -0,0 +1,110 @@ +package agent + +import ( + "testing" + + "charm.land/catwalk/pkg/catwalk" + "github.com/charmbracelet/crush/internal/config" + "github.com/stretchr/testify/require" +) + +func TestIsOpenCodeMessagesModel(t *testing.T) { + t.Parallel() + + tests := []struct { + provider string + model string + expected bool + }{ + {"opencode-zen", "qwen3.7-max", true}, + {"opencode-zen", "qwen3.5-plus", true}, + {"opencode-zen", "qwen3.6-plus", true}, + {"opencode-zen", "qwen3.7-plus", true}, + {"opencode-zen", "claude-opus-4-5", true}, + {"opencode-zen", "claude-sonnet-5", true}, + {"opencode-zen", "muse-spark-1.3-contributor-free", false}, + {"opencode-zen", "gpt-5.6-luna", false}, + {"opencode-zen", "grok-4.5", false}, + {"opencode-zen", "minimax-m3", false}, + {"opencode-zen", "kimi-k3", false}, + {"opencode-zen", "big-pickle", false}, + {"opencode-go", "minimax-m2.7", true}, + {"opencode-go", "minimax-m3", true}, + {"opencode-go", "qwen3.7-max", true}, + {"opencode-go", "qwen3.7-plus", true}, + {"opencode-go", "qwen3.6-plus", true}, + {"opencode-go", "qwen3.8-flash", true}, + {"opencode-go", "qwen3.8-max", true}, + {"opencode-go", "muse-spark-1.3-contributor", false}, + {"opencode-go", "gpt-5.6-luna", false}, + {"opencode-go", "glm-5.3", false}, + {"opencode-go", "kimi-k3", false}, + {"opencode-go", "longcat-2.0", false}, + {"opencode-go", "ox-alpha-free", false}, + {"opencode-go", "minimax", false}, + {"other", "claude-opus-4-5", false}, + {"other", "qwen3.7-max", false}, + } + for _, tt := range tests { + require.Equal(t, tt.expected, isOpenCodeMessagesModel(tt.provider, tt.model), "%s/%s", tt.provider, tt.model) + } +} + +func TestIsOpenCodeResponsesModel(t *testing.T) { + t.Parallel() + + tests := []struct { + model string + expected bool + }{ + {"muse-spark-1.3-contributor-free", true}, + {"muse-spark-1.2", true}, + {"grok-4.5", true}, + {"grok-4.6", true}, + {"grok-build-0.1", true}, + {"gpt-5.6-luna", true}, + {"gpt-5.5", true}, + {"gpt-5.3-codex", true}, + {"minimax-m3", false}, + {"qwen3.7-max", false}, + {"kimi-k3", false}, + {"glm-5.3", false}, + {"big-pickle", false}, + {"ox-alpha-free", false}, + {"hy3", false}, + {"longcat-2.0", false}, + } + for _, tt := range tests { + require.Equal(t, tt.expected, isOpenCodeResponsesModel(tt.model), tt.model) + } +} + +func TestBuildProviderOpenCodeRouting(t *testing.T) { + t.Parallel() + + for _, providerID := range []string{ + string(catwalk.InferenceProviderOpenCodeZen), + string(catwalk.InferenceProviderOpenCodeGo), + } { + t.Run(providerID, func(t *testing.T) { + t.Parallel() + env := testEnv(t) + providerCfg := config.ProviderConfig{ + ID: providerID, + BaseURL: "https://opencode.ai/zen/v1", + Type: catwalk.TypeOpenAICompat, + APIKey: "$OPENCODE_API_KEY", + } + coord := newTestCoordinator(t, env, providerID, providerCfg) + + for _, modelID := range []string{"kimi-k3", "muse-spark-1.3-contributor-free", "grok-4.6", "gpt-5.6-luna"} { + provider, err := coord.buildProvider(providerCfg, config.SelectedModel{ + Model: modelID, + Provider: providerID, + }, false) + require.NoError(t, err, modelID) + require.NotNil(t, provider, modelID) + } + }) + } +} diff --git a/internal/agent/prism_model_test.go b/internal/agent/prism_model_test.go new file mode 100644 index 0000000000000000000000000000000000000000..6e35deb149706024bf8fddc85a7a3e6020edea5d --- /dev/null +++ b/internal/agent/prism_model_test.go @@ -0,0 +1,101 @@ +package agent + +import ( + "encoding/json" + "net/http" + "testing" + + "charm.land/fantasy" + "github.com/charmbracelet/crush/internal/agent/hyper" + "github.com/charmbracelet/crush/internal/message" + + "charm.land/fantasy/providers/openai" + "github.com/stretchr/testify/require" +) + +func TestExtractPrismModel(t *testing.T) { + t.Parallel() + + prismMetadata := func(header http.Header) fantasy.ProviderMetadata { + metadata := &openai.ProviderMetadata{} + hyper.HeaderFunc(header, metadata) + return fantasy.ProviderMetadata{openai.Name: metadata} + } + + t.Run("returns the routed model ID and name", func(t *testing.T) { + t.Parallel() + header := http.Header{} + header.Set(hyper.PrismModelIDHeader, "prism-42") + header.Set(hyper.PrismModelNameHeader, "GPT-5.2 Codex Max") + modelID, modelName := extractPrismModel(prismMetadata(header)) + require.Equal(t, "prism-42", modelID) + require.Equal(t, "GPT-5.2 Codex Max", modelName) + }) + + t.Run("returns empty when not routed through Prism", func(t *testing.T) { + t.Parallel() + modelID, modelName := extractPrismModel(fantasy.ProviderMetadata{}) + require.Empty(t, modelID) + require.Empty(t, modelName) + + metadata := &openai.ProviderMetadata{ + ExtraFields: map[string]json.RawMessage{"remaining": json.RawMessage(`{"hypercredits": 10}`)}, + } + modelID, modelName = extractPrismModel(fantasy.ProviderMetadata{openai.Name: metadata}) + require.Empty(t, modelID) + require.Empty(t, modelName) + }) + + t.Run("persisted on the assistant message", func(t *testing.T) { + t.Parallel() + header := http.Header{} + header.Set(hyper.PrismModelIDHeader, "prism-42") + header.Set(hyper.PrismModelNameHeader, "GPT-5.2 Codex Max") + + msg := message.Message{} + msg.PrismModelID, msg.PrismModelName = extractPrismModel(prismMetadata(header)) + require.Equal(t, "prism-42", msg.PrismModelID) + require.Equal(t, "GPT-5.2 Codex Max", msg.PrismModelName) + }) +} + +func TestExtractPrismSavings(t *testing.T) { + t.Parallel() + + prismMetadata := func(header http.Header) fantasy.ProviderMetadata { + metadata := &openai.ProviderMetadata{} + hyper.HeaderFunc(header, metadata) + return fantasy.ProviderMetadata{openai.Name: metadata} + } + + t.Run("returns both savings figures when reported", func(t *testing.T) { + t.Parallel() + header := http.Header{} + header.Set(hyper.PrismHypercreditSavingsHeader, "1.5") + header.Set(hyper.PrismDollarSavingsHeader, "0.002") + hypercredits, dollars := extractPrismSavings(prismMetadata(header)) + require.NotNil(t, hypercredits) + require.Equal(t, 1.5, *hypercredits) + require.NotNil(t, dollars) + require.Equal(t, 0.002, *dollars) + }) + + t.Run("returns nil when not reported", func(t *testing.T) { + t.Parallel() + hypercredits, dollars := extractPrismSavings(fantasy.ProviderMetadata{}) + require.Nil(t, hypercredits) + require.Nil(t, dollars) + }) + + t.Run("returns nil when malformed", func(t *testing.T) { + t.Parallel() + metadata := &openai.ProviderMetadata{ + ExtraFields: map[string]json.RawMessage{ + hyper.PrismHypercreditSavingsField: json.RawMessage(`"not-a-number"`), + }, + } + hypercredits, dollars := extractPrismSavings(fantasy.ProviderMetadata{openai.Name: metadata}) + require.Nil(t, hypercredits) + require.Nil(t, dollars) + }) +} diff --git a/internal/agent/prompts.go b/internal/agent/prompts.go new file mode 100644 index 0000000000000000000000000000000000000000..448fe0425c3b700b1d6edafc842c4815ad3d5760 --- /dev/null +++ b/internal/agent/prompts.go @@ -0,0 +1,42 @@ +package agent + +import ( + "context" + _ "embed" + + "github.com/charmbracelet/crush/internal/agent/prompt" + "github.com/charmbracelet/crush/internal/config" +) + +//go:embed templates/coder.md.tpl +var coderPromptTmpl []byte + +//go:embed templates/task.md.tpl +var taskPromptTmpl []byte + +//go:embed templates/initialize.md.tpl +var initializePromptTmpl []byte + +func coderPrompt(opts ...prompt.Option) (*prompt.Prompt, error) { + systemPrompt, err := prompt.NewPrompt("coder", string(coderPromptTmpl), opts...) + if err != nil { + return nil, err + } + return systemPrompt, nil +} + +func taskPrompt(opts ...prompt.Option) (*prompt.Prompt, error) { + systemPrompt, err := prompt.NewPrompt("task", string(taskPromptTmpl), opts...) + if err != nil { + return nil, err + } + return systemPrompt, nil +} + +func InitializePrompt(cfg *config.ConfigStore) (string, error) { + systemPrompt, err := prompt.NewPrompt("initialize", string(initializePromptTmpl)) + if err != nil { + return "", err + } + return systemPrompt.Build(context.Background(), "", "", cfg) +} diff --git a/internal/agent/queued_runid_test.go b/internal/agent/queued_runid_test.go new file mode 100644 index 0000000000000000000000000000000000000000..e3e99d4f648e12d5a98be052747841553b1fa8ae --- /dev/null +++ b/internal/agent/queued_runid_test.go @@ -0,0 +1,181 @@ +package agent + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + "charm.land/catwalk/pkg/catwalk" + "charm.land/fantasy" + "github.com/charmbracelet/crush/internal/agent/notify" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/stretchr/testify/require" +) + +// gatedStreamModel streams a single text part followed by a clean finish, +// but blocks the very first Stream call until its gate is released. That +// lets a test hold a run "active" (past PrepareStep, inside Stream) just +// long enough to enqueue a follow-up prompt behind the busy session. +// Subsequent Stream calls (e.g. the recursive run draining the queue) +// proceed immediately. +type gatedStreamModel struct { + text string + gate chan struct{} + entered chan struct{} + calls atomic.Int64 +} + +func (m *gatedStreamModel) Provider() string { return "fake" } +func (m *gatedStreamModel) Model() string { return "fake-model" } + +func (m *gatedStreamModel) Generate(ctx context.Context, call fantasy.Call) (*fantasy.Response, error) { + return &fantasy.Response{ + Content: fantasy.ResponseContent{fantasy.TextContent{Text: m.text}}, + FinishReason: fantasy.FinishReasonStop, + }, nil +} + +func (m *gatedStreamModel) Stream(ctx context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { + if m.calls.Add(1) == 1 { + close(m.entered) + select { + case <-m.gate: + case <-ctx.Done(): + } + } + text := m.text + return func(yield func(fantasy.StreamPart) bool) { + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextStart, ID: "1"}) { + return + } + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, ID: "1", Delta: text}) { + return + } + if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextEnd, ID: "1"}) { + return + } + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}) + }, nil +} + +func (m *gatedStreamModel) GenerateObject(ctx context.Context, call fantasy.ObjectCall) (*fantasy.ObjectResponse, error) { + return nil, errors.New("not implemented") +} + +func (m *gatedStreamModel) StreamObject(ctx context.Context, call fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) { + return nil, errors.New("not implemented") +} + +// TestRun_QueuedRunIDPromptRunsRecursivelyAndPublishesRunComplete is the +// end-to-end proof of fix 2: a prompt carrying a RunID that is queued +// behind a busy session must NOT be silently folded into the active turn. +// It runs as its own turn via the recursive run path and publishes its +// own terminal RunComplete, so a `crush run` caller blocking on that +// RunID does not hang. The active turn keeps its own RunComplete too. +func TestRun_QueuedRunIDPromptRunsRecursivelyAndPublishesRunComplete(t *testing.T) { + t.Parallel() + + env := testEnv(t) + broker := pubsub.NewBroker[notify.RunComplete]() + t.Cleanup(broker.Shutdown) + + large := &gatedStreamModel{ + text: "done", + gate: make(chan struct{}), + entered: make(chan struct{}), + } + small := &finishStreamModel{text: "title"} + + sa := NewSessionAgent(SessionAgentOptions{ + LargeModel: Model{Model: large, CatwalkCfg: catwalk.Model{ContextWindow: 200000, DefaultMaxTokens: 10000}}, + SmallModel: Model{Model: small, CatwalkCfg: catwalk.Model{ContextWindow: 200000, DefaultMaxTokens: 10000}}, + IsYolo: true, + Sessions: env.sessions, + Messages: env.messages, + RunComplete: broker, + }).(*sessionAgent) + + sess, err := env.sessions.Create(t.Context(), "session") + require.NoError(t, err) + + subCtx, subCancel := context.WithCancel(t.Context()) + defer subCancel() + ch := broker.Subscribe(subCtx) + + // Start the main turn; it blocks inside Stream once active. + mainDone := make(chan error, 1) + go func() { + _, runErr := sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + RunID: "run-main", + Prompt: "main", + }) + mainDone <- runErr + }() + + // Wait until the main turn is active (inside Stream). + select { + case <-large.entered: + case <-time.After(5 * time.Second): + t.Fatal("main run never entered Stream") + } + require.True(t, sa.IsSessionBusy(sess.ID), "main run must be active before enqueueing the follow-up") + + // Enqueue a RunID-bearing follow-up behind the busy session. + res, err := sa.Run(t.Context(), SessionAgentCall{ + SessionID: sess.ID, + RunID: "run-follow", + Prompt: "follow", + }) + require.NoError(t, err) + require.Nil(t, res, "a busy-session follow-up must enqueue and return (nil, nil)") + require.Equal(t, 1, sa.QueuedPrompts(sess.ID), "the follow-up must be queued, not folded") + + // Release the main turn so it completes and hands off to the queue. + close(large.gate) + require.NoError(t, <-mainDone) + + // Both turns must publish their own terminal RunComplete. + got := map[string]notify.RunComplete{} + deadline := time.After(5 * time.Second) + for len(got) < 2 { + select { + case ev := <-ch: + got[ev.Payload.RunID] = ev.Payload + case <-deadline: + t.Fatalf("timed out waiting for both RunCompletes; got %v", got) + } + } + + main, ok := got["run-main"] + require.True(t, ok, "the active turn must publish its own RunComplete") + require.Empty(t, main.Error) + require.False(t, main.Cancelled) + + follow, ok := got["run-follow"] + require.True(t, ok, + "the queued RunID prompt must publish its own RunComplete instead of being folded silently") + require.Empty(t, follow.Error) + require.False(t, follow.Cancelled) + require.Equal(t, "done", follow.Text, "the queued prompt ran as its own turn") + + // Two distinct assistant turns prove the follow-up was not folded. + msgs, err := env.messages.List(t.Context(), sess.ID) + require.NoError(t, err) + var assistants, follows int + for _, m := range msgs { + switch m.Role { + case message.Assistant: + assistants++ + case message.User: + if m.Content().String() == "follow" { + follows++ + } + } + } + require.Equal(t, 2, assistants, "the active turn and the recursive turn each produce one assistant message") + require.Equal(t, 1, follows, "the follow-up prompt is its own user turn") +} diff --git a/internal/agent/request_timeout.go b/internal/agent/request_timeout.go new file mode 100644 index 0000000000000000000000000000000000000000..77aa8d7a63c3028670454367baac210f8654807d --- /dev/null +++ b/internal/agent/request_timeout.go @@ -0,0 +1,123 @@ +package agent + +import ( + "context" + "errors" + "fmt" + "time" + + "charm.land/fantasy" +) + +// requestTimeoutError reports that an LLM request exhausted its configured +// request_timeout budget. For streaming requests this is an idle timeout: +// it only fires when the provider sends nothing for the whole window, so a +// slow but actively streaming response is never killed. It wraps the +// underlying error so callers can still match [context.DeadlineExceeded] +// through the chain. +type requestTimeoutError struct { + timeout time.Duration + idle bool + cause error +} + +func (e *requestTimeoutError) Error() string { + msg := fmt.Sprintf("LLM request timed out after %s", e.timeout) + if e.idle { + msg = fmt.Sprintf("LLM stream received no data for %s", e.timeout) + } + if e.cause != nil { + return fmt.Sprintf("%s: %v", msg, e.cause) + } + return msg +} + +func (e *requestTimeoutError) Unwrap() error { return e.cause } + +// userMessage explains the timeout in the UI, including how long the request +// ran before giving up and how to change the limit. +func (e *requestTimeoutError) userMessage() string { + hint := "Increase the limit with \"option request-timeout SECONDS\" or set it to 0 to disable the timeout." + if e.idle { + return fmt.Sprintf("The model stopped sending data for %s. %s", e.timeout, hint) + } + return fmt.Sprintf("The model did not respond within %s. %s", e.timeout, hint) +} + +// requestTimeoutModel wraps a [fantasy.LanguageModel] so requests are +// bounded by the configured request_timeout. Non-streaming calls get a hard +// per-request deadline, applied per call so fantasy's retry loop gives every +// attempt a fresh budget — the same per-request semantics the provider SDKs +// expose. Streams instead get an idle timeout: the budget resets whenever a +// part arrives and only fires when the provider goes silent, so a slow but +// actively streaming response is never aborted. +type requestTimeoutModel struct { + fantasy.LanguageModel + timeout time.Duration +} + +// newRequestTimeoutModel bounds each request to m with the given timeout. A +// timeout of zero or less, or a nil model, returns m unchanged. +func newRequestTimeoutModel(m fantasy.LanguageModel, timeout time.Duration) fantasy.LanguageModel { + if m == nil || timeout <= 0 { + return m + } + return requestTimeoutModel{LanguageModel: m, timeout: timeout} +} + +// wrapTimedOut replaces err with the requestTimeoutError when this model's +// own deadline fired. Other errors — user cancellation, outer deadlines, +// provider failures — pass through unchanged. Cancellation errors caused by +// our own timer report as [context.DeadlineExceeded] so callers never +// mistake a timeout for a user cancellation. +func wrapTimedOut(ctx context.Context, timeoutErr *requestTimeoutError, err error) error { + if err == nil || context.Cause(ctx) != timeoutErr { + return err + } + if errors.Is(err, context.Canceled) { + timeoutErr.cause = context.DeadlineExceeded + } else { + timeoutErr.cause = err + } + return timeoutErr +} + +// Generate implements [fantasy.LanguageModel]. The request gets a hard +// deadline: there is no incremental progress signal, so the whole call must +// finish within the budget. +func (m requestTimeoutModel) Generate(ctx context.Context, call fantasy.Call) (*fantasy.Response, error) { + timeoutErr := &requestTimeoutError{timeout: m.timeout} + ctx, cancel := context.WithTimeoutCause(ctx, m.timeout, timeoutErr) + defer cancel() + resp, err := m.LanguageModel.Generate(ctx, call) + return resp, wrapTimedOut(ctx, timeoutErr, err) +} + +// Stream implements [fantasy.LanguageModel]. +// +// The stream is consumed after Stream returns, so the timer must outlive +// this call: it fires only after timeout seconds without any part arriving, +// and is released when iteration ends, whether the stream finishes, breaks, +// or the idle timeout aborts it. Both the initial connection and gaps +// between parts share the same budget. +func (m requestTimeoutModel) Stream(ctx context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { + timeoutErr := &requestTimeoutError{timeout: m.timeout, idle: true} + ctx, cancel := context.WithCancelCause(ctx) + timer := time.AfterFunc(m.timeout, func() { cancel(timeoutErr) }) + + inner, err := m.LanguageModel.Stream(ctx, call) + if err != nil { + timer.Stop() + cancel(nil) + return nil, wrapTimedOut(ctx, timeoutErr, err) + } + return func(yield func(fantasy.StreamPart) bool) { + defer timer.Stop() + defer cancel(nil) + inner(func(part fantasy.StreamPart) bool { + timer.Reset(m.timeout) + part.Error = wrapTimedOut(ctx, timeoutErr, part.Error) + return yield(part) + }) + }, nil +} diff --git a/internal/agent/request_timeout_test.go b/internal/agent/request_timeout_test.go new file mode 100644 index 0000000000000000000000000000000000000000..7ff62ceda96591e3b4cfcb2ebacb50b1a79e2467 --- /dev/null +++ b/internal/agent/request_timeout_test.go @@ -0,0 +1,286 @@ +package agent + +import ( + "context" + "runtime" + "testing" + "time" + + "charm.land/fantasy" + "charm.land/x/vcr" + "github.com/charmbracelet/crush/internal/message" + "github.com/stretchr/testify/require" +) + +// fakeLanguageModel is a [fantasy.LanguageModel] stub that records the +// context its methods were called with and can be configured with a custom +// stream body. +type fakeLanguageModel struct { + generateCtx context.Context + streamCtx context.Context + stream func(yield func(fantasy.StreamPart) bool) +} + +func (f *fakeLanguageModel) Generate(ctx context.Context, _ fantasy.Call) (*fantasy.Response, error) { + f.generateCtx = ctx + return &fantasy.Response{}, nil +} + +func (f *fakeLanguageModel) Stream(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + f.streamCtx = ctx + if f.stream == nil { + return func(yield func(fantasy.StreamPart) bool) { + yield(fantasy.StreamPart{}) + }, nil + } + return f.stream, nil +} + +func (f *fakeLanguageModel) GenerateObject(context.Context, fantasy.ObjectCall) (*fantasy.ObjectResponse, error) { + return &fantasy.ObjectResponse{}, nil +} + +func (f *fakeLanguageModel) StreamObject(context.Context, fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) { + return nil, nil +} + +func (f *fakeLanguageModel) Provider() string { return "fake" } +func (f *fakeLanguageModel) Model() string { return "fake-model" } + +func TestNewRequestTimeoutModel_Disabled(t *testing.T) { + t.Parallel() + + inner := &fakeLanguageModel{} + require.Same(t, inner, newRequestTimeoutModel(inner, 0)) + require.Same(t, inner, newRequestTimeoutModel(inner, -time.Second)) +} + +func TestRequestTimeoutModel_GenerateDeadline(t *testing.T) { + t.Parallel() + + inner := &fakeLanguageModel{} + m := newRequestTimeoutModel(inner, 5*time.Minute) + + _, err := m.Generate(t.Context(), fantasy.Call{}) + require.NoError(t, err) + + _, ok := inner.generateCtx.Deadline() + require.True(t, ok, "Generate should run under a deadline") +} + +func TestRequestTimeoutModel_StreamDeadlineOutlivesCall(t *testing.T) { + t.Parallel() + + inner := &fakeLanguageModel{} + m := newRequestTimeoutModel(inner, 5*time.Minute) + + stream, err := m.Stream(t.Context(), fantasy.Call{}) + require.NoError(t, err) + + require.NoError(t, inner.streamCtx.Err(), "the idle timer must not fire while the stream is being consumed") + + for range stream { + } + + require.ErrorIs(t, inner.streamCtx.Err(), context.Canceled, "the stream context should be released after the stream ends") +} + +func TestRequestTimeoutModel_StreamAbortsWhenIdle(t *testing.T) { + t.Parallel() + + inner := &fakeLanguageModel{} + // A stream that outlives the idle window: it waits for the context to + // be done and reports what it observed. + streamObserved := make(chan error, 1) + inner.stream = func(yield func(fantasy.StreamPart) bool) { + <-inner.streamCtx.Done() + streamObserved <- inner.streamCtx.Err() + } + m := newRequestTimeoutModel(inner, 10*time.Millisecond) + stream, err := m.Stream(t.Context(), fantasy.Call{}) + require.NoError(t, err) + + done := make(chan struct{}) + go func() { + defer close(done) + for range stream { + } + }() + + select { + case err := <-streamObserved: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(5 * time.Second): + t.Fatal("stream was not aborted by the idle timeout") + } + <-done +} + +func TestRequestTimeoutModel_ActiveStreamSurvives(t *testing.T) { + t.Parallel() + + inner := &fakeLanguageModel{} + // A stream that keeps sending data: total runtime exceeds the timeout, + // but every gap is shorter than the idle window, so it must finish. + inner.stream = func(yield func(fantasy.StreamPart) bool) { + for range 6 { + time.Sleep(10 * time.Millisecond) + if !yield(fantasy.StreamPart{}) { + return + } + } + } + m := newRequestTimeoutModel(inner, 25*time.Millisecond) + stream, err := m.Stream(t.Context(), fantasy.Call{}) + require.NoError(t, err) + + parts := 0 + for part := range stream { + require.NoError(t, part.Error) + parts++ + } + require.Equal(t, 6, parts) +} + +// blockingModel blocks until the context is done and then returns the +// context error, the way a hung provider request would. +type blockingModel struct { + fakeLanguageModel +} + +func (b *blockingModel) Generate(ctx context.Context, _ fantasy.Call) (*fantasy.Response, error) { + <-ctx.Done() + return nil, ctx.Err() +} + +func TestRequestTimeoutModel_GenerateReportsTimeout(t *testing.T) { + t.Parallel() + + m := newRequestTimeoutModel(&blockingModel{}, 10*time.Millisecond) + + _, err := m.Generate(t.Context(), fantasy.Call{}) + require.Error(t, err) + + var timeoutErr *requestTimeoutError + require.ErrorAs(t, err, &timeoutErr) + require.Equal(t, 10*time.Millisecond, timeoutErr.timeout) + require.ErrorIs(t, err, context.DeadlineExceeded, "the deadline must stay detectable through the chain") + require.Contains(t, err.Error(), "timed out after 10ms") +} + +func TestRequestTimeoutModel_StreamReportsTimeout(t *testing.T) { + t.Parallel() + + inner := &fakeLanguageModel{} + // A provider stream that fails with the context error once the deadline + // fires, mirroring how SDKs surface mid-stream aborts. + inner.stream = func(yield func(fantasy.StreamPart) bool) { + <-inner.streamCtx.Done() + yield(fantasy.StreamPart{Error: inner.streamCtx.Err()}) + } + m := newRequestTimeoutModel(inner, 10*time.Millisecond) + stream, err := m.Stream(t.Context(), fantasy.Call{}) + require.NoError(t, err) + + var got error + for part := range stream { + if part.Error != nil { + got = part.Error + } + } + + var timeoutErr *requestTimeoutError + require.ErrorAs(t, got, &timeoutErr) + require.ErrorIs(t, got, context.DeadlineExceeded) +} + +func TestRequestTimeoutModel_ParentCancelPassesThrough(t *testing.T) { + t.Parallel() + + m := newRequestTimeoutModel(&blockingModel{}, 5*time.Minute) + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + go func() { + time.Sleep(10 * time.Millisecond) + cancel() + }() + + _, err := m.Generate(ctx, fantasy.Call{}) + require.ErrorIs(t, err, context.Canceled) + + var timeoutErr *requestTimeoutError + require.NotErrorAs(t, err, &timeoutErr, "user cancellation must not be reported as a timeout") +} + +func TestRequestTimeoutErrorMessages(t *testing.T) { + t.Parallel() + + err := &requestTimeoutError{timeout: time.Second} + require.Equal(t, "LLM request timed out after 1s", err.Error()) + require.Contains(t, err.userMessage(), "1s") + require.Contains(t, err.userMessage(), "request-timeout") + + err.cause = context.DeadlineExceeded + require.Equal(t, "LLM request timed out after 1s: context deadline exceeded", err.Error()) + + idle := &requestTimeoutError{timeout: 2 * time.Second, idle: true} + require.Equal(t, "LLM stream received no data for 2s", idle.Error()) + require.Contains(t, idle.userMessage(), "stopped sending data for 2s") + require.Contains(t, idle.userMessage(), "request-timeout") +} + +// timeoutOnlyModel streams a single error part shaped exactly like the one +// requestTimeoutModel produces when its deadline fires. +type timeoutOnlyModel struct { + fakeLanguageModel +} + +func (m *timeoutOnlyModel) Stream(context.Context, fantasy.Call) (fantasy.StreamResponse, error) { + timeoutErr := &requestTimeoutError{timeout: time.Second, idle: true, cause: context.DeadlineExceeded} + return func(yield func(fantasy.StreamPart) bool) { + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeError, Error: timeoutErr}) + }, nil +} + +// TestRequestTimeoutRunFinishMessage pins what the user sees when a request +// exhausts its timeout: a "Request timed out" finish that names the elapsed +// budget and how to change it, instead of a bare provider error. +func TestRequestTimeoutRunFinishMessage(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("skipping on windows for now") + } + + env := testEnv(t) + model := &timeoutOnlyModel{} + agent, err := coderAgent(vcr.NewRecorder(t), env, model, model) + require.NoError(t, err) + + session, err := env.sessions.Create(t.Context(), "timeout session") + require.NoError(t, err) + + _, err = agent.Run(t.Context(), SessionAgentCall{ + Prompt: "Hello", + SessionID: session.ID, + MaxOutputTokens: 10000, + }) + require.Error(t, err) + + msgs, err := env.messages.List(t.Context(), session.ID) + require.NoError(t, err) + + var finish *message.Finish + for _, msg := range msgs { + if msg.Role != message.Assistant { + continue + } + if part := msg.FinishPart(); part != nil { + finish = part + } + } + require.NotNil(t, finish, "the assistant message should carry a finish part") + require.Equal(t, message.FinishReasonError, finish.Reason) + require.Equal(t, "Request timed out", finish.Message) + require.Contains(t, finish.Details, "stopped sending data for 1s") + require.Contains(t, finish.Details, "request-timeout") +} diff --git a/internal/agent/run_complete_test.go b/internal/agent/run_complete_test.go new file mode 100644 index 0000000000000000000000000000000000000000..418e68a2ca1dfe2061c4122a46bdaa3f4f50eec2 --- /dev/null +++ b/internal/agent/run_complete_test.go @@ -0,0 +1,322 @@ +package agent + +import ( + "context" + "testing" + "time" + + "github.com/charmbracelet/crush/internal/agent/notify" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/stretchr/testify/require" +) + +// TestSessionAgentRun_QueueStripsOnComplete verifies that when a Run +// call is enqueued (because the session is already busy), the +// OnComplete hook is NOT propagated onto the queued copy. The hook +// belongs to the caller's retry/coalesce scope (typically +// coordinator.Run) which has already returned by the time the queue +// drains; carrying it forward would silently funnel the terminal +// event into a closure nobody reads, and subscribers (`crush run`) +// would hang waiting for a RunComplete that never publishes. +func TestSessionAgentRun_QueueStripsOnComplete(t *testing.T) { + t.Parallel() + + env := testEnv(t) + a := NewSessionAgent(SessionAgentOptions{ + Sessions: env.sessions, + Messages: env.messages, + }).(*sessionAgent) + + const sessionID = "queued-session" + // Mark the session as busy so Run takes the queue branch + // without needing a real model. + a.activeRequests.Set(sessionID, &activeCancel{cancel: func() {}}) + + var called bool + hook := func(notify.RunComplete) { called = true } + + res, err := a.Run(t.Context(), SessionAgentCall{ + SessionID: sessionID, + RunID: "run-xyz", + Prompt: "queued prompt", + OnComplete: hook, + }) + require.NoError(t, err) + require.Nil(t, res, "queued Run must return (nil, nil)") + require.False(t, called, + "OnComplete must not fire on the enqueue path; the caller's scope is still live") + + queued, ok := a.messageQueue.Get(sessionID) + require.True(t, ok) + require.Len(t, queued, 1) + require.Nil(t, queued[0].OnComplete, + "queued SessionAgentCall must have OnComplete stripped so the drain falls back to the default broker publish") + require.Equal(t, "queued prompt", queued[0].Prompt, + "all other fields must be preserved on the queued copy") + require.Equal(t, "run-xyz", queued[0].RunID, + "RunID must be preserved on the queued copy so the drained turn's "+ + "RunComplete still correlates with the originating SendMessage") +} + +// TestDrainQueueForStep_FiltersUnderDispatchLock verifies that the queue +// drain evaluates the per-session cancel mark while holding the dispatch +// mutex (canceledBySeq's documented precondition). Queued calls at or +// below the cancel high-water mark are dropped, calls queued after the +// cancel (higher seq) are folded, untracked enqueues (seq == 0) are +// dropped whenever any mark is present, and the queue is cleared. These +// calls carry no RunID, so all foldable survivors are returned for +// folding (the existing follow-up behavior). +func TestDrainQueueForStep_FiltersUnderDispatchLock(t *testing.T) { + t.Parallel() + + env := testEnv(t) + a := NewSessionAgent(SessionAgentOptions{ + Sessions: env.sessions, + Messages: env.messages, + }).(*sessionAgent) + + const sessionID = "drain-session" + a.messageQueue.Set(sessionID, []SessionAgentCall{ + {SessionID: sessionID, Prompt: "below", acceptSeq: 1}, + {SessionID: sessionID, Prompt: "at-mark", acceptSeq: 2}, + {SessionID: sessionID, Prompt: "after", acceptSeq: 3}, + {SessionID: sessionID, Prompt: "untracked", acceptSeq: 0}, + }) + // Cancel high-water mark at seq 2: seq <= 2 and seq == 0 are covered. + a.cancelMark.Set(sessionID, 2) + + fold, canceledWithRunID := a.drainQueueForStep(sessionID) + + require.Len(t, fold, 1, + "only the follow-up queued after the cancel (seq > mark) must be folded") + require.Equal(t, "after", fold[0].Prompt) + require.Empty(t, canceledWithRunID, + "no dropped call carried a RunID, so none need a terminal RunComplete") + + _, ok := a.messageQueue.Get(sessionID) + require.False(t, ok, "drain must clear the session message queue when nothing is kept") +} + +// TestDrainQueueForStep_NoMarkFoldsAllNonRunID verifies that with no +// cancel mark recorded, every queued call without a RunID is folded. +func TestDrainQueueForStep_NoMarkFoldsAllNonRunID(t *testing.T) { + t.Parallel() + + env := testEnv(t) + a := NewSessionAgent(SessionAgentOptions{ + Sessions: env.sessions, + Messages: env.messages, + }).(*sessionAgent) + + const sessionID = "drain-nomark" + a.messageQueue.Set(sessionID, []SessionAgentCall{ + {SessionID: sessionID, Prompt: "a", acceptSeq: 0}, + {SessionID: sessionID, Prompt: "b", acceptSeq: 5}, + }) + + fold, canceledWithRunID := a.drainQueueForStep(sessionID) + require.Len(t, fold, 2, "no cancel mark means all non-RunID queued calls are folded") + require.Empty(t, canceledWithRunID) +} + +// TestDrainQueueForStep_KeepsRunIDPromptsQueued is the core of fix 2: a +// queued prompt that carries a RunID must NOT be folded into the active +// turn. Folding it would silently absorb it into another turn and never +// publish a RunComplete for its RunID, hanging a `crush run` caller that +// blocks on that event. Such prompts are left in the queue so the +// recursive run path gives each its own turn and its own RunComplete. +// Non-RunID prompts are still folded. +func TestDrainQueueForStep_KeepsRunIDPromptsQueued(t *testing.T) { + t.Parallel() + + env := testEnv(t) + a := NewSessionAgent(SessionAgentOptions{ + Sessions: env.sessions, + Messages: env.messages, + }).(*sessionAgent) + + const sessionID = "drain-runid" + a.messageQueue.Set(sessionID, []SessionAgentCall{ + {SessionID: sessionID, Prompt: "fold-me", acceptSeq: 1}, + {SessionID: sessionID, RunID: "run-a", Prompt: "keep-me", acceptSeq: 2}, + {SessionID: sessionID, RunID: "run-b", Prompt: "keep-me-too", acceptSeq: 3}, + }) + + fold, canceledWithRunID := a.drainQueueForStep(sessionID) + + require.Len(t, fold, 1, "only the non-RunID prompt is folded into the active turn") + require.Equal(t, "fold-me", fold[0].Prompt) + require.Empty(t, canceledWithRunID) + + kept, ok := a.messageQueue.Get(sessionID) + require.True(t, ok, "RunID-bearing prompts must remain queued for the recursive run path") + require.Len(t, kept, 2) + require.Equal(t, "run-a", kept[0].RunID) + require.Equal(t, "run-b", kept[1].RunID) +} + +// TestDrainQueueForStep_ReportsCanceledRunIDDrops verifies that a queued +// prompt carrying a RunID that is dropped because a cancel covers it is +// reported in canceledWithRunID so the caller can publish its terminal +// cancelled RunComplete. A canceled prompt without a RunID is dropped +// silently as before. +func TestDrainQueueForStep_ReportsCanceledRunIDDrops(t *testing.T) { + t.Parallel() + + env := testEnv(t) + a := NewSessionAgent(SessionAgentOptions{ + Sessions: env.sessions, + Messages: env.messages, + }).(*sessionAgent) + + const sessionID = "drain-cancel-runid" + a.messageQueue.Set(sessionID, []SessionAgentCall{ + {SessionID: sessionID, RunID: "run-canceled", Prompt: "canceled", acceptSeq: 1}, + {SessionID: sessionID, Prompt: "canceled-no-runid", acceptSeq: 1}, + {SessionID: sessionID, RunID: "run-survives", Prompt: "survives", acceptSeq: 5}, + }) + a.cancelMark.Set(sessionID, 2) + + fold, canceledWithRunID := a.drainQueueForStep(sessionID) + + require.Empty(t, fold, "no uncanceled non-RunID prompts to fold") + require.Len(t, canceledWithRunID, 1, + "only the dropped RunID-bearing prompt needs a terminal RunComplete") + require.Equal(t, "run-canceled", canceledWithRunID[0].RunID) + + kept, ok := a.messageQueue.Get(sessionID) + require.True(t, ok) + require.Len(t, kept, 1, "the uncanceled RunID prompt stays queued") + require.Equal(t, "run-survives", kept[0].RunID) +} + +// TestRunCompletePublisher_MustDeliverOverTakesPublish exercises the +// pubsub.Publisher interface change end-to-end: a Broker is the only +// concrete Publisher implementation and must satisfy both Publish and +// PublishMustDeliver. The coordinator's final RunComplete emit relies +// on PublishMustDeliver to apply bounded-blocking semantics so a +// momentarily-full subscriber buffer can't silently drop the +// authoritative end-of-run event. +func TestRunCompletePublisher_MustDeliverOverTakesPublish(t *testing.T) { + t.Parallel() + + broker := pubsub.NewBroker[notify.RunComplete]() + t.Cleanup(broker.Shutdown) + + // Subscribe before publishing so the event is delivered. + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + ch := broker.Subscribe(ctx) + + rc := notify.RunComplete{SessionID: "S", MessageID: "m", Text: "ok"} + var pub pubsub.Publisher[notify.RunComplete] = broker + pub.PublishMustDeliver(t.Context(), pubsub.UpdatedEvent, rc) + + select { + case got := <-ch: + require.Equal(t, rc, got.Payload) + case <-time.After(time.Second): + t.Fatal("PublishMustDeliver did not deliver event") + } +} + +// requireSingleCancelledRunComplete reads exactly one RunComplete from ch, +// asserts it is the cancelled terminal event for runID, and verifies no +// second event arrives. This observes the published pubsub event rather +// than internal bookkeeping, which is the contract a `crush run` caller +// blocking on the broker actually relies on. +func requireSingleCancelledRunComplete(t *testing.T, ch <-chan pubsub.Event[notify.RunComplete], sessionID, runID string) { + t.Helper() + select { + case ev := <-ch: + require.Equal(t, runID, ev.Payload.RunID, + "the published RunComplete must carry the dropped queued prompt's RunID") + require.Equal(t, sessionID, ev.Payload.SessionID) + require.True(t, ev.Payload.Cancelled, + "a dropped queued prompt must publish a cancelled RunComplete") + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for the cancelled RunComplete") + } + select { + case extra := <-ch: + t.Fatalf("expected exactly one RunComplete, got a second: %+v", extra.Payload) + case <-time.After(100 * time.Millisecond): + } +} + +// TestCancel_QueuedRunIDPromptPublishesCancelledRunComplete proves the +// terminal-event behavior end-to-end: a RunID-bearing prompt sitting in +// the queue that is canceled while queued (via the public Cancel path, +// which routes through clearQueueAndNotify -> publishCanceledQueueDrops) +// must emit exactly one cancelled RunComplete on the broker for its +// RunID. A queued prompt without a RunID is dropped silently. This is the +// coverage the earlier drain test lacked: it asserted the returned +// bookkeeping slice, not the published event a `crush run` caller awaits. +func TestCancel_QueuedRunIDPromptPublishesCancelledRunComplete(t *testing.T) { + t.Parallel() + + env := testEnv(t) + broker := pubsub.NewBroker[notify.RunComplete]() + t.Cleanup(broker.Shutdown) + + a := NewSessionAgent(SessionAgentOptions{ + Sessions: env.sessions, + Messages: env.messages, + RunComplete: broker, + }).(*sessionAgent) + + subCtx, subCancel := context.WithCancel(t.Context()) + defer subCancel() + ch := broker.Subscribe(subCtx) + + const sessionID = "cancel-queued-runid" + a.messageQueue.Set(sessionID, []SessionAgentCall{ + {SessionID: sessionID, Prompt: "no-runid", acceptSeq: 1}, + {SessionID: sessionID, RunID: "run-queued", Prompt: "queued", acceptSeq: 2}, + }) + + a.Cancel(sessionID) + + requireSingleCancelledRunComplete(t, ch, sessionID, "run-queued") + + _, ok := a.messageQueue.Get(sessionID) + require.False(t, ok, "Cancel must clear the queue") +} + +// TestDrainQueueForStep_DroppedRunIDPublishesCancelledRunComplete drives +// the production drain sequence (drainQueueForStep then +// publishCanceledQueueDrops, mirroring the PrepareStep handoff) and +// asserts the dropped RunID-bearing prompt actually publishes exactly one +// cancelled RunComplete on the broker. The companion bookkeeping test +// covers the returned slice; this one covers the observable terminal +// event. +func TestDrainQueueForStep_DroppedRunIDPublishesCancelledRunComplete(t *testing.T) { + t.Parallel() + + env := testEnv(t) + broker := pubsub.NewBroker[notify.RunComplete]() + t.Cleanup(broker.Shutdown) + + a := NewSessionAgent(SessionAgentOptions{ + Sessions: env.sessions, + Messages: env.messages, + RunComplete: broker, + }).(*sessionAgent) + + subCtx, subCancel := context.WithCancel(t.Context()) + defer subCancel() + ch := broker.Subscribe(subCtx) + + const sessionID = "drain-drop-runid" + a.messageQueue.Set(sessionID, []SessionAgentCall{ + {SessionID: sessionID, RunID: "run-dropped", Prompt: "dropped", acceptSeq: 1}, + {SessionID: sessionID, Prompt: "dropped-no-runid", acceptSeq: 1}, + }) + a.cancelMark.Set(sessionID, 2) + + _, canceledWithRunID := a.drainQueueForStep(sessionID) + require.Len(t, canceledWithRunID, 1) + a.publishCanceledQueueDrops(canceledWithRunID) + + requireSingleCancelledRunComplete(t, ch, sessionID, "run-dropped") +} diff --git a/internal/agent/run_marker.go b/internal/agent/run_marker.go new file mode 100644 index 0000000000000000000000000000000000000000..404cca1e8c41bb9179deb886552f3580a977fdfc --- /dev/null +++ b/internal/agent/run_marker.go @@ -0,0 +1,52 @@ +package agent + +import ( + "context" + "sync/atomic" +) + +// runCompleteMarkerKey is the unexported context key carrying a +// [runCompleteMarker] from the dispatch boundary (backend.runAgent) +// down into the coordinator. It lets the dispatcher learn whether the +// coordinator already published the authoritative terminal +// notify.RunComplete for the run, so a fallback terminal event is only +// emitted when one is actually missing (e.g. an error returned before +// sessionAgent.Run ever executed). It avoids a breaking change to the +// Coordinator interface. +type runCompleteMarkerKey struct{} + +// runCompleteMarker records whether a terminal RunComplete has been +// published for a run. It is shared by pointer through the context so +// a publish deep in the call stack is observable by the dispatcher +// after the call returns. +type runCompleteMarker struct { + published atomic.Bool +} + +// WithRunCompleteMarker returns ctx carrying a fresh marker the +// coordinator can flag via [MarkRunCompletePublished] once it emits the +// run's terminal RunComplete. Callers read the result with +// [RunCompletePublished]. Attaching the marker is optional: code paths +// without one simply skip the dedup signal. +func WithRunCompleteMarker(ctx context.Context) context.Context { + return context.WithValue(ctx, runCompleteMarkerKey{}, &runCompleteMarker{}) +} + +// MarkRunCompletePublished records that the authoritative terminal +// RunComplete has been published for the run carried by ctx. It is a +// no-op when no marker is present (e.g. the in-process/local Run path, +// which is not dispatched through backend.runAgent). +func MarkRunCompletePublished(ctx context.Context) { + if m, ok := ctx.Value(runCompleteMarkerKey{}).(*runCompleteMarker); ok { + m.published.Store(true) + } +} + +// RunCompletePublished reports whether [MarkRunCompletePublished] was +// called on ctx's marker. It returns false when no marker is present. +func RunCompletePublished(ctx context.Context) bool { + if m, ok := ctx.Value(runCompleteMarkerKey{}).(*runCompleteMarker); ok { + return m.published.Load() + } + return false +} diff --git a/internal/agent/runid.go b/internal/agent/runid.go new file mode 100644 index 0000000000000000000000000000000000000000..1afac005b9e4c627c9c06f21a2c565bab86e1c28 --- /dev/null +++ b/internal/agent/runid.go @@ -0,0 +1,33 @@ +package agent + +import "context" + +// runIDContextKey is the unexported context key used to carry a +// caller-supplied RunID from the workspace HTTP boundary +// (backend.SendMessage) down into coordinator.Run without forcing a +// breaking change to the Coordinator.Run signature. The value is +// then copied onto SessionAgentCall.RunID by the coordinator so the +// agent's terminal RunComplete event can echo it back to the +// originating caller. +type runIDContextKey struct{} + +// WithRunID returns ctx tagged with a per-request RunID. It is the +// boundary helper for callers that need their SendMessage→Run +// terminal event to be uniquely correlatable (e.g. `crush run` +// against a session that may be busy). Empty runIDs are stored +// as-is; downstream code treats an empty RunID as "caller did not +// supply one" and falls back to SessionID-only correlation. +func WithRunID(ctx context.Context, runID string) context.Context { + return context.WithValue(ctx, runIDContextKey{}, runID) +} + +// RunIDFromContext returns the RunID set by [WithRunID], or "" if +// none was set or the value is not a string. Exported because the +// coordinator and tests in other packages need to read it; safe to +// call on any context. +func RunIDFromContext(ctx context.Context) string { + if v, ok := ctx.Value(runIDContextKey{}).(string); ok { + return v + } + return "" +} diff --git a/internal/agent/usage_fallback.go b/internal/agent/usage_fallback.go new file mode 100644 index 0000000000000000000000000000000000000000..3cf347c38f91a72968c3537a5d954bb6feb1d0b9 --- /dev/null +++ b/internal/agent/usage_fallback.go @@ -0,0 +1,176 @@ +package agent + +import ( + "fmt" + + "charm.land/fantasy" +) + +func usageIsZero(usage fantasy.Usage) bool { + return usage.InputTokens == 0 && + usage.OutputTokens == 0 && + usage.TotalTokens == 0 && + usage.ReasoningTokens == 0 && + usage.CacheCreationTokens == 0 && + usage.CacheReadTokens == 0 +} + +func fallbackStepUsage(messages []fantasy.Message, step fantasy.StepResult) (fantasy.Usage, bool) { + if !usageIsZero(step.Usage) { + return step.Usage, false + } + + inputTokens := estimateMessageTokens(messages) + outputTokens := estimateStepCompletionTokens(step) + if inputTokens == 0 && outputTokens == 0 { + return fantasy.Usage{}, false + } + + return fantasy.Usage{ + InputTokens: inputTokens, + OutputTokens: outputTokens, + TotalTokens: inputTokens + outputTokens, + }, true +} + +func cloneFantasyMessages(messages []fantasy.Message) []fantasy.Message { + cloned := make([]fantasy.Message, len(messages)) + for i, msg := range messages { + cloned[i] = msg + cloned[i].Content = append([]fantasy.MessagePart(nil), msg.Content...) + } + return cloned +} + +func estimateMessageTokens(messages []fantasy.Message) int64 { + var tokens int64 + for _, msg := range messages { + tokens += approxTokenCount(string(msg.Role)) + for _, part := range msg.Content { + tokens += estimateMessagePartTokens(part) + } + } + return tokens +} + +func estimateStepCompletionTokens(step fantasy.StepResult) int64 { + var tokens int64 + for _, content := range step.Content { + switch c := content.(type) { + case fantasy.TextContent: + tokens += approxTokenCount(c.Text) + case *fantasy.TextContent: + tokens += approxTokenCount(c.Text) + case fantasy.ReasoningContent: + tokens += approxTokenCount(c.Text) + case *fantasy.ReasoningContent: + tokens += approxTokenCount(c.Text) + case fantasy.FileContent: + tokens += estimateGeneratedFileTokens(c) + case *fantasy.FileContent: + tokens += estimateGeneratedFileTokens(*c) + case fantasy.SourceContent: + tokens += estimateSourceTokens(c) + case *fantasy.SourceContent: + tokens += estimateSourceTokens(*c) + case fantasy.ToolCallContent: + tokens += estimateToolCallTokens(c.ToolName, c.Input) + case *fantasy.ToolCallContent: + tokens += estimateToolCallTokens(c.ToolName, c.Input) + case fantasy.ToolResultContent: + if c.ProviderExecuted { + tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result) + } + case *fantasy.ToolResultContent: + if c.ProviderExecuted { + tokens += estimateToolResultContentTokens(c.ToolCallID, c.ToolName, c.ClientMetadata, c.Result) + } + } + } + return tokens +} + +func estimateMessagePartTokens(part fantasy.MessagePart) int64 { + switch p := part.(type) { + case fantasy.TextPart: + return approxTokenCount(p.Text) + case *fantasy.TextPart: + return approxTokenCount(p.Text) + case fantasy.ReasoningPart: + return approxTokenCount(p.Text) + case *fantasy.ReasoningPart: + return approxTokenCount(p.Text) + case fantasy.FilePart: + return estimateFilePartTokens(p) + case *fantasy.FilePart: + return estimateFilePartTokens(*p) + case fantasy.ToolCallPart: + return estimateToolCallTokens(p.ToolName, p.Input) + case *fantasy.ToolCallPart: + return estimateToolCallTokens(p.ToolName, p.Input) + case fantasy.ToolResultPart: + return estimateToolResultContentTokens(p.ToolCallID, "", "", p.Output) + case *fantasy.ToolResultPart: + return estimateToolResultContentTokens(p.ToolCallID, "", "", p.Output) + default: + return 0 + } +} + +func estimateToolCallTokens(toolName, input string) int64 { + return approxTokenCount(toolName) + approxTokenCount(input) +} + +func estimateToolResultContentTokens(toolCallID, toolName, metadata string, output fantasy.ToolResultOutputContent) int64 { + tokens := approxTokenCount(toolCallID) + approxTokenCount(toolName) + approxTokenCount(metadata) + switch result := output.(type) { + case fantasy.ToolResultOutputContentText: + tokens += approxTokenCount(result.Text) + case *fantasy.ToolResultOutputContentText: + tokens += approxTokenCount(result.Text) + case fantasy.ToolResultOutputContentError: + if result.Error != nil { + tokens += approxTokenCount(result.Error.Error()) + } + case *fantasy.ToolResultOutputContentError: + if result.Error != nil { + tokens += approxTokenCount(result.Error.Error()) + } + case fantasy.ToolResultOutputContentMedia: + tokens += estimateMediaTokens(result.MediaType, result.Text, len(result.Data)) + case *fantasy.ToolResultOutputContentMedia: + tokens += estimateMediaTokens(result.MediaType, result.Text, len(result.Data)) + } + return tokens +} + +func estimateFilePartTokens(file fantasy.FilePart) int64 { + return estimateMediaTokens(file.MediaType, file.Filename, len(file.Data)) +} + +func estimateGeneratedFileTokens(file fantasy.FileContent) int64 { + return estimateMediaTokens(file.MediaType, "", len(file.Data)) +} + +func estimateMediaTokens(mediaType, text string, dataBytes int) int64 { + if dataBytes == 0 { + return approxTokenCount(mediaType) + approxTokenCount(text) + } + return approxTokenCount(fmt.Sprintf("%s %s %d bytes", mediaType, text, dataBytes)) +} + +func estimateSourceTokens(source fantasy.SourceContent) int64 { + return approxTokenCount(string(source.SourceType)) + + approxTokenCount(source.ID) + + approxTokenCount(source.URL) + + approxTokenCount(source.Title) + + approxTokenCount(source.MediaType) + + approxTokenCount(source.Filename) +} + +func approxTokenCount(s string) int64 { + if s == "" { + return 0 + } + return int64((len(s) + 3) / 4) +} diff --git a/internal/agent/usage_fallback_test.go b/internal/agent/usage_fallback_test.go new file mode 100644 index 0000000000000000000000000000000000000000..3b6fc19b943fbb1cf636782c051e91a6ff64a5df --- /dev/null +++ b/internal/agent/usage_fallback_test.go @@ -0,0 +1,339 @@ +package agent + +import ( + "errors" + "testing" + + "charm.land/catwalk/pkg/catwalk" + "charm.land/fantasy" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/session" + "github.com/stretchr/testify/require" +) + +func TestUsageIsZero(t *testing.T) { + t.Parallel() + + require.True(t, usageIsZero(fantasy.Usage{})) + require.False(t, usageIsZero(fantasy.Usage{InputTokens: 1})) + require.False(t, usageIsZero(fantasy.Usage{OutputTokens: 1})) + require.False(t, usageIsZero(fantasy.Usage{TotalTokens: 1})) + require.False(t, usageIsZero(fantasy.Usage{ReasoningTokens: 1})) + require.False(t, usageIsZero(fantasy.Usage{CacheCreationTokens: 1})) + require.False(t, usageIsZero(fantasy.Usage{CacheReadTokens: 1})) +} + +func TestFallbackStepUsageKeepsProviderUsage(t *testing.T) { + t.Parallel() + + usage := fantasy.Usage{ + InputTokens: 10, + OutputTokens: 5, + TotalTokens: 15, + } + step := fantasy.StepResult{ + Response: fantasy.Response{Usage: usage}, + } + + fallbackUsage, estimated := fallbackStepUsage(nil, step) + require.False(t, estimated) + require.Equal(t, usage, fallbackUsage) +} + +func TestFallbackStepUsageEstimatesPromptAndAssistantText(t *testing.T) { + t.Parallel() + + messages := []fantasy.Message{ + fantasy.NewUserMessage("please explain the implementation details"), + } + step := fantasy.StepResult{ + Response: fantasy.Response{ + Content: fantasy.ResponseContent{ + fantasy.TextContent{Text: "the implementation stores state safely"}, + }, + }, + } + + usage, estimated := fallbackStepUsage(messages, step) + require.True(t, estimated) + require.Positive(t, usage.InputTokens) + require.Positive(t, usage.OutputTokens) + require.Equal(t, usage.InputTokens+usage.OutputTokens, usage.TotalTokens) +} + +func TestFallbackStepUsageEstimatesReasoning(t *testing.T) { + t.Parallel() + + messages := []fantasy.Message{ + { + Role: fantasy.MessageRoleAssistant, + Content: []fantasy.MessagePart{ + fantasy.ReasoningPart{Text: "first reason about the request"}, + }, + }, + } + step := fantasy.StepResult{ + Response: fantasy.Response{ + Content: fantasy.ResponseContent{ + fantasy.ReasoningContent{Text: "second reason about the answer"}, + }, + }, + } + + usage, estimated := fallbackStepUsage(messages, step) + require.True(t, estimated) + require.Positive(t, usage.InputTokens) + require.Positive(t, usage.OutputTokens) +} + +func TestFallbackStepUsageEstimatesToolCalls(t *testing.T) { + t.Parallel() + + step := fantasy.StepResult{ + Response: fantasy.Response{ + Content: fantasy.ResponseContent{ + fantasy.ToolCallContent{ + ToolCallID: "tool-call-1", + ToolName: "view", + Input: `{"file_path":"/tmp/example.go"}`, + }, + }, + }, + } + + usage, estimated := fallbackStepUsage(nil, step) + require.True(t, estimated) + require.Zero(t, usage.InputTokens) + require.Positive(t, usage.OutputTokens) + require.Equal(t, usage.OutputTokens, usage.TotalTokens) +} + +func TestFallbackStepUsageEstimatesToolResults(t *testing.T) { + t.Parallel() + + messages := []fantasy.Message{ + { + Role: fantasy.MessageRoleTool, + Content: []fantasy.MessagePart{ + fantasy.ToolResultPart{ + ToolCallID: "tool-call-1", + Output: fantasy.ToolResultOutputContentText{ + Text: "file contents returned by the tool", + }, + }, + fantasy.ToolResultPart{ + ToolCallID: "tool-call-2", + Output: fantasy.ToolResultOutputContentError{ + Error: errors.New("permission denied"), + }, + }, + fantasy.ToolResultPart{ + ToolCallID: "tool-call-3", + Output: fantasy.ToolResultOutputContentMedia{ + MediaType: "image/png", + Text: "screenshot", + Data: "abc123", + }, + }, + }, + }, + } + + usage, estimated := fallbackStepUsage(messages, fantasy.StepResult{}) + require.True(t, estimated) + require.Positive(t, usage.InputTokens) + require.Zero(t, usage.OutputTokens) + require.Equal(t, usage.InputTokens, usage.TotalTokens) +} + +func TestFallbackStepUsageSkipsClientToolResultsAsOutput(t *testing.T) { + t.Parallel() + + step := fantasy.StepResult{ + Response: fantasy.Response{ + Content: fantasy.ResponseContent{ + fantasy.ToolResultContent{ + ToolCallID: "tool-call-1", + ToolName: "bash", + Result: fantasy.ToolResultOutputContentText{ + Text: "large client-executed payload that should not count as model output tokens", + }, + }, + }, + }, + } + + usage, estimated := fallbackStepUsage(nil, step) + require.False(t, estimated) + require.Zero(t, usage.OutputTokens) +} + +func TestFallbackStepUsageCountsProviderToolResultsAsOutput(t *testing.T) { + t.Parallel() + + step := fantasy.StepResult{ + Response: fantasy.Response{ + Content: fantasy.ResponseContent{ + fantasy.ToolResultContent{ + ToolCallID: "tool-call-1", + ToolName: "web_search", + ProviderExecuted: true, + ClientMetadata: "provider metadata", + Result: fantasy.ToolResultOutputContentText{Text: "provider-executed result"}, + }, + }, + }, + } + + usage, estimated := fallbackStepUsage(nil, step) + require.True(t, estimated) + require.Positive(t, usage.OutputTokens) + require.Equal(t, usage.OutputTokens, usage.TotalTokens) +} + +func TestFallbackStepUsageReturnsZeroWithoutContent(t *testing.T) { + t.Parallel() + + usage, estimated := fallbackStepUsage(nil, fantasy.StepResult{}) + require.False(t, estimated) + require.True(t, usageIsZero(usage)) +} + +func TestUpdateSessionUsageSkipsEstimatedCost(t *testing.T) { + t.Parallel() + + agent := &sessionAgent{} + currentSession := &session.Session{ID: "session-id", Cost: 1.25} + model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}} + usage := fantasy.Usage{InputTokens: 1000, OutputTokens: 2000} + + agent.updateSessionUsage(model, currentSession, usage, nil, true) + + require.Equal(t, 1.25, currentSession.Cost) + require.Equal(t, int64(1000), currentSession.PromptTokens) + require.Equal(t, int64(2000), currentSession.CompletionTokens) + require.True(t, currentSession.EstimatedUsage) +} + +func TestUpdateSessionUsageKeepsCountersForZeroUsage(t *testing.T) { + t.Parallel() + + agent := &sessionAgent{} + currentSession := &session.Session{ + ID: "session-id", + PromptTokens: 123, + CompletionTokens: 456, + Cost: 1.25, + } + model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}} + + agent.updateSessionUsage(model, currentSession, fantasy.Usage{}, nil, false) + + require.Equal(t, 1.25, currentSession.Cost) + require.Equal(t, int64(123), currentSession.PromptTokens) + require.Equal(t, int64(456), currentSession.CompletionTokens) +} + +func TestUpdateSessionUsagePreservesOmittedCountersForPartialUsage(t *testing.T) { + t.Parallel() + + agent := &sessionAgent{} + currentSession := &session.Session{ + ID: "session-id", + PromptTokens: 123, + CompletionTokens: 456, + } + model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}} + usage := fantasy.Usage{InputTokens: 789} + + agent.updateSessionUsage(model, currentSession, usage, nil, false) + + require.Equal(t, int64(789), currentSession.PromptTokens) + require.Equal(t, int64(456), currentSession.CompletionTokens) +} + +func TestUpdateSessionUsagePreservesCountersForTotalOnlyUsage(t *testing.T) { + t.Parallel() + + agent := &sessionAgent{} + currentSession := &session.Session{ + ID: "session-id", + PromptTokens: 123, + CompletionTokens: 456, + } + model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}} + usage := fantasy.Usage{TotalTokens: 100} + + agent.updateSessionUsage(model, currentSession, usage, nil, false) + + require.Equal(t, int64(123), currentSession.PromptTokens) + require.Equal(t, int64(456), currentSession.CompletionTokens) +} + +func TestUpdateSessionUsagePreservesPromptForOutputOnlyUsage(t *testing.T) { + t.Parallel() + + agent := &sessionAgent{} + currentSession := &session.Session{ + ID: "session-id", + PromptTokens: 123, + CompletionTokens: 456, + } + model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}} + usage := fantasy.Usage{OutputTokens: 50} + + agent.updateSessionUsage(model, currentSession, usage, nil, false) + + require.Equal(t, int64(123), currentSession.PromptTokens) + require.Equal(t, int64(50), currentSession.CompletionTokens) +} + +func TestUpdateSessionUsageKeepsCountersForEstimatedZeroUsage(t *testing.T) { + t.Parallel() + + agent := &sessionAgent{} + currentSession := &session.Session{ + ID: "session-id", + PromptTokens: 123, + CompletionTokens: 456, + Cost: 1.25, + } + model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}} + + agent.updateSessionUsage(model, currentSession, fantasy.Usage{}, nil, true) + + require.Equal(t, 1.25, currentSession.Cost) + require.Equal(t, int64(123), currentSession.PromptTokens) + require.Equal(t, int64(456), currentSession.CompletionTokens) +} + +func TestSummaryCompletionTokens(t *testing.T) { + t.Parallel() + + summaryMessage := message.Message{ + Parts: []message.ContentPart{ + message.TextContent{Text: "summary text"}, + message.ReasoningContent{Thinking: "reasoning text"}, + }, + } + + require.Equal(t, int64(42), summaryCompletionTokens(fantasy.Usage{OutputTokens: 42}, summaryMessage)) + require.Equal(t, approxTokenCount("summary text")+approxTokenCount("reasoning text"), summaryCompletionTokens(fantasy.Usage{}, summaryMessage)) + require.Zero(t, summaryCompletionTokens(fantasy.Usage{}, message.Message{})) +} + +func TestUpdateSessionUsageAddsProviderCost(t *testing.T) { + t.Parallel() + + agent := &sessionAgent{} + currentSession := &session.Session{ID: "session-id", Cost: 1.25} + model := Model{CatwalkCfg: catwalk.Model{CostPer1MIn: 10, CostPer1MOut: 20}} + usage := fantasy.Usage{InputTokens: 1000, OutputTokens: 2000} + + agent.updateSessionUsage(model, currentSession, usage, nil, false) + + require.Equal(t, 1.3, currentSession.Cost) + require.Equal(t, int64(1000), currentSession.PromptTokens) + require.Equal(t, int64(2000), currentSession.CompletionTokens) + require.False(t, currentSession.EstimatedUsage) +} diff --git a/internal/apigen/docs.html.tmpl b/internal/apigen/docs.html.tmpl new file mode 100644 index 0000000000000000000000000000000000000000..2557c27480f6a3a803836d1e55926ffc6f381086 --- /dev/null +++ b/internal/apigen/docs.html.tmpl @@ -0,0 +1,47 @@ + + + + +{{.Info.Title}} + + + +

{{.Info.Title}}

+{{if .Info.Description}}

{{.Info.Description}}

{{end}} +

Version {{.Info.Version}} · Base path {{.Base}} · openapi.json

+{{range .Tags}} +

{{.Name}}

+{{range .Endpoints}} +
+
{{.Method}} {{.Path}}
+
{{.Summary}}
+{{if .Description}}
{{.Description}}
{{end}} +
Success: {{.Status}}
+{{if .Params}} +

Parameters

+ + +{{range .Params}} +{{end}}
NameInRequiredDescription
{{.Name}}{{.In}}{{if .Required}}yes{{else}}no{{end}}{{.Description}}
+{{end}} +{{if .Request}}

Request body

{{.Request}}
{{end}} +{{if .SSE}}

Response

Streams text/event-stream.

{{end}} +{{if .Response}}

Response body

{{.Response}}
{{end}} +{{if .Errors}}

Errors

{{range $i, $c := .Errors}}{{if $i}}, {{end}}{{$c}}{{end}} — JSON error object

{{end}} +
+{{end}} +{{end}} + + diff --git a/internal/apigen/endpoint.go b/internal/apigen/endpoint.go new file mode 100644 index 0000000000000000000000000000000000000000..76f2ab5e0cd3f7d63f6d65820352c79851aece58 --- /dev/null +++ b/internal/apigen/endpoint.go @@ -0,0 +1,168 @@ +// Package apigen provides a fluent API for documenting HTTP endpoints. +// +// An [Endpoint] is the single source of truth for one route: the same +// value drives mux registration and documentation generation, so the +// two cannot drift. Endpoints are constructed with the method builders +// ([Get], [Post], [Put], [Delete]) and finalized with +// [endpointBuilder.Handle], the only way to produce an Endpoint. A +// chain without Handle does not compile into a []Endpoint, so +// incomplete definitions fail at build time. +package apigen + +import "net/http" + +// Param documents a single path or query parameter. +type Param struct { + // Name is the parameter name as it appears in the path or query. + Name string + // In is "path" or "query". + In string + // Type is the JSON schema type, e.g. "string", "integer". + Type string + // Description is a short human-readable description. + Description string + // Required reports whether the parameter must be provided. Path + // parameters are always required. + Required bool +} + +// Endpoint is the documentation and routing record for a single HTTP +// endpoint. Construct it with the fluent builders ([Get], [Post], ...); +// only [endpointBuilder.Handle] produces a complete Endpoint. +type Endpoint struct { + method string + path string + summary string + description string + tags []string + params []Param + request any + response any + status int + errors []int + sse bool + handler http.HandlerFunc +} + +// Method returns the HTTP method, e.g. "GET". +func (e Endpoint) Method() string { return e.method } + +// Path returns the route pattern, e.g. "/v1/workspaces/{id}". +func (e Endpoint) Path() string { return e.path } + +// Summary returns the one-line summary. +func (e Endpoint) Summary() string { return e.summary } + +// Description returns the optional longer prose description. +func (e Endpoint) Description() string { return e.description } + +// Tags returns the grouping tags. +func (e Endpoint) Tags() []string { return e.tags } + +// Params returns the documented path and query parameters. +func (e Endpoint) Params() []Param { return e.params } + +// Request returns the request body example value, or nil for none. +func (e Endpoint) Request() any { return e.request } + +// Response returns the success response example value, or nil for an +// empty success response. +func (e Endpoint) Response() any { return e.response } + +// Status returns the success status code, defaulting to 200. +func (e Endpoint) Status() int { + if e.status == 0 { + return http.StatusOK + } + return e.status +} + +// Errors returns the status codes that return an error object. +func (e Endpoint) Errors() []int { return e.errors } + +// SSE reports whether the endpoint streams text/event-stream instead +// of producing a JSON response. +func (e Endpoint) SSE() bool { return e.sse } + +// Handler returns the endpoint's HTTP handler. +func (e Endpoint) Handler() http.HandlerFunc { return e.handler } + +// endpointBuilder accumulates endpoint documentation fluently. It is +// deliberately a distinct type from [Endpoint]: a []Endpoint literal +// cannot contain a builder, so a chain missing Handle fails to +// compile. +type endpointBuilder struct{ e Endpoint } + +// Get begins a GET endpoint for path, e.g. Get("/v1/workspaces/{id}"). +func Get(path string) endpointBuilder { + return endpointBuilder{Endpoint{method: http.MethodGet, path: path}} +} + +// Post begins a POST endpoint for path. +func Post(path string) endpointBuilder { + return endpointBuilder{Endpoint{method: http.MethodPost, path: path}} +} + +// Put begins a PUT endpoint for path. +func Put(path string) endpointBuilder { + return endpointBuilder{Endpoint{method: http.MethodPut, path: path}} +} + +// Delete begins a DELETE endpoint for path. +func Delete(path string) endpointBuilder { + return endpointBuilder{Endpoint{method: http.MethodDelete, path: path}} +} + +// Summary sets the one-line summary. +func (b endpointBuilder) Summary(s string) endpointBuilder { b.e.summary = s; return b } + +// Description sets an optional longer prose description. +func (b endpointBuilder) Description(s string) endpointBuilder { b.e.description = s; return b } + +// Tags sets the grouping tags used on the docs page. +func (b endpointBuilder) Tags(t ...string) endpointBuilder { b.e.tags = t; return b } + +// Accepts documents the request body with an example value, e.g. +// Accepts(proto.Workspace{}). The value's type is used to derive the +// JSON schema; the reference is compile-checked. +func (b endpointBuilder) Accepts(v any) endpointBuilder { b.e.request = v; return b } + +// Responds documents the success response body with an example value. +func (b endpointBuilder) Responds(v any) endpointBuilder { b.e.response = v; return b } + +// Status sets the success status code when it differs from 200, e.g. +// 202 for an accepted-but-async operation. +func (b endpointBuilder) Status(code int) endpointBuilder { b.e.status = code; return b } + +// Fails documents status codes that return an error object. +func (b endpointBuilder) Fails(codes ...int) endpointBuilder { b.e.errors = codes; return b } + +// SSE marks the endpoint as streaming text/event-stream. +func (b endpointBuilder) SSE() endpointBuilder { b.e.sse = true; return b } + +// PathParam documents a path parameter. Path parameters are always +// required. +func (b endpointBuilder) PathParam(name, desc string) endpointBuilder { + return b.param(Param{Name: name, In: "path", Type: "string", Description: desc, Required: true}) +} + +// QueryParam documents an optional query parameter. +func (b endpointBuilder) QueryParam(name, desc string) endpointBuilder { + return b.param(Param{Name: name, In: "query", Type: "string", Description: desc}) +} + +// RequiredQueryParam documents a required query parameter. +func (b endpointBuilder) RequiredQueryParam(name, desc string) endpointBuilder { + return b.param(Param{Name: name, In: "query", Type: "string", Description: desc, Required: true}) +} + +func (b endpointBuilder) param(p Param) endpointBuilder { + b.e.params = append(b.e.params, p) + return b +} + +// Handle finalizes the endpoint with its HTTP handler. +func (b endpointBuilder) Handle(h http.HandlerFunc) Endpoint { + b.e.handler = h + return b.e +} diff --git a/internal/apigen/endpoint_test.go b/internal/apigen/endpoint_test.go new file mode 100644 index 0000000000000000000000000000000000000000..e978680f0b782a37290d115c09274f658eaa16cf --- /dev/null +++ b/internal/apigen/endpoint_test.go @@ -0,0 +1,207 @@ +package apigen + +import ( + "encoding/json" + "net/http" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type nested struct { + Value string `json:"value"` +} + +type sample struct { + ID string `json:"id"` + Count int `json:"count"` + Ratio float64 `json:"ratio,omitempty"` + Active bool `json:"active,omitempty"` + Labels []string `json:"labels"` + Meta map[string]string `json:"meta,omitempty"` + Nested nested `json:"nested"` + Ptr *nested `json:"ptr,omitempty"` + When time.Time `json:"when"` + Err error `json:"error,omitempty"` + Anything any `json:"anything,omitempty"` + Skipped string `json:"-"` + unexported string +} + +func TestSchemaOfStruct(t *testing.T) { + t.Parallel() + s := SchemaOf(sample{}) + + require.Equal(t, "object", s.Type) + require.ElementsMatch(t, []string{"id", "count", "labels", "nested", "when"}, s.Required) + + require.Equal(t, "string", s.Properties["id"].Type) + require.Equal(t, "integer", s.Properties["count"].Type) + require.Equal(t, "number", s.Properties["ratio"].Type) + require.Equal(t, "boolean", s.Properties["active"].Type) + require.Equal(t, "array", s.Properties["labels"].Type) + require.Equal(t, "string", s.Properties["labels"].Items.Type) + require.Equal(t, "object", s.Properties["meta"].Type) + require.Equal(t, "string", s.Properties["meta"].AdditionalProperties.Type) + require.Equal(t, "object", s.Properties["nested"].Type) + require.Equal(t, "string", s.Properties["nested"].Properties["value"].Type) + require.Equal(t, "object", s.Properties["ptr"].Type) + require.Equal(t, "string", s.Properties["when"].Type) + require.Equal(t, "date-time", s.Properties["when"].Format) + require.Equal(t, "string", s.Properties["error"].Type) + require.Equal(t, &Schema{}, s.Properties["anything"]) + + require.NotContains(t, s.Properties, "Skipped") + require.NotContains(t, s.Properties, "unexported") +} + +type cyc struct { + Name string `json:"name"` + Child *cyc `json:"child,omitempty"` +} + +func TestSchemaOfCycle(t *testing.T) { + t.Parallel() + s := SchemaOf(cyc{}) + require.Equal(t, "object", s.Type) + require.Equal(t, "apigen.cyc", s.Properties["child"].Description) +} + +type withTime struct { + At time.Time `json:"at"` +} + +func TestSchemaOfSamePackageExpansion(t *testing.T) { + t.Parallel() + // Types from this package expand even when reached through another. + s := SchemaOf(withTime{}) + require.Equal(t, "string", s.Properties["at"].Type) +} + +type embeds struct { + nested + Extra string `json:"extra"` +} + +func TestSchemaOfEmbedded(t *testing.T) { + t.Parallel() + s := SchemaOf(embeds{}) + require.Contains(t, s.Properties, "value") + require.Contains(t, s.Properties, "extra") +} + +func TestBuilderCompileContract(t *testing.T) { + t.Parallel() + called := false + h := func(http.ResponseWriter, *http.Request) { called = true } + + e := Get("/v1/things/{id}"). + Summary("Get thing"). + Description("Longer prose."). + Tags("things"). + PathParam("id", "Thing ID"). + QueryParam("verbose", "Verbose output"). + RequiredQueryParam("mode", "Mode"). + Responds(sample{}). + Fails(404, 500). + Handle(h) + + require.Equal(t, http.MethodGet, e.Method()) + require.Equal(t, "/v1/things/{id}", e.Path()) + require.Equal(t, "Get thing", e.Summary()) + require.Equal(t, "Longer prose.", e.Description()) + require.Equal(t, []string{"things"}, e.Tags()) + require.Equal(t, []Param{ + {Name: "id", In: "path", Type: "string", Description: "Thing ID", Required: true}, + {Name: "verbose", In: "query", Type: "string", Description: "Verbose output"}, + {Name: "mode", In: "query", Type: "string", Description: "Mode", Required: true}, + }, e.Params()) + require.Equal(t, sample{}, e.Response()) + require.Nil(t, e.Request()) + require.Equal(t, []int{404, 500}, e.Errors()) + require.False(t, e.SSE()) + + e.Handler()(nil, nil) + require.True(t, called) +} + +func TestBuilderVerbsAndSSE(t *testing.T) { + t.Parallel() + for method, build := range map[string]func(string) endpointBuilder{ + http.MethodGet: Get, + http.MethodPost: Post, + http.MethodPut: Put, + http.MethodDelete: Delete, + } { + e := build("/x").Handle(func(http.ResponseWriter, *http.Request) {}) + require.Equal(t, method, e.Method()) + } + + e := Get("/events").SSE().Handle(func(http.ResponseWriter, *http.Request) {}) + require.True(t, e.SSE()) +} + +func TestSpecJSON(t *testing.T) { + t.Parallel() + info := Info{Title: "Test API", Version: "1.0", Description: "Desc."} + endpoints := []Endpoint{ + Post("/v1/things"). + Summary("Create thing"). + Accepts(sample{}). + Responds(sample{}). + Fails(400). + Handle(func(http.ResponseWriter, *http.Request) {}), + Get("/v1/things/{id}/events"). + Summary("Stream"). + PathParam("id", "Thing ID"). + SSE(). + Handle(func(http.ResponseWriter, *http.Request) {}), + } + + data, err := SpecJSON(info, "/v1", endpoints) + require.NoError(t, err) + text := string(data) + + require.Contains(t, text, `"openapi": "3.2.1"`) + require.Contains(t, text, `"/things"`) + require.Contains(t, text, `"/things/{id}/events"`) + require.Contains(t, text, `"text/event-stream"`) + require.Contains(t, text, `"application/json"`) + require.NotContains(t, text, "/v1/things\"") // base prefix stripped +} + +func TestRenderDocs(t *testing.T) { + t.Parallel() + info := Info{Title: "Test API", Version: "1.0", Description: "Desc."} + endpoints := []Endpoint{ + Get("/v1/things/{id}"). + Summary("Get thing"). + Tags("things"). + PathParam("id", "Thing ID"). + Responds(sample{}). + Fails(404). + Handle(func(http.ResponseWriter, *http.Request) {}), + } + + var sb strings.Builder + require.NoError(t, RenderDocs(&sb, info, "/v1", endpoints)) + html := sb.String() + + require.Contains(t, html, "Test API") + require.Contains(t, html, "/v1/things/{id}") + require.Contains(t, html, "Get thing") + require.Contains(t, html, "Thing ID") + require.Contains(t, html, "count") + require.Contains(t, html, "openapi.json") +} + +func TestErrorSchemaShape(t *testing.T) { + t.Parallel() + // The shared error schema must stay aligned with the error JSON + // handlers actually write: {"error": "..."}. + data, err := json.Marshal(errorSchema) + require.NoError(t, err) + require.Contains(t, string(data), `"error"`) +} diff --git a/internal/apigen/openapi.go b/internal/apigen/openapi.go new file mode 100644 index 0000000000000000000000000000000000000000..3e160df231e6fb324bc7551db68a52d6b35f5db1 --- /dev/null +++ b/internal/apigen/openapi.go @@ -0,0 +1,147 @@ +package apigen + +import ( + "encoding/json" + "net/http" + "sort" + "strconv" + "strings" +) + +// Info carries the top-level API metadata for the OpenAPI document. +type Info struct { + Title string + Version string + Description string +} + +// document is an OpenAPI 3.1 document rooted at the API's base path. +type document struct { + OpenAPI string `json:"openapi"` + Info Info `json:"info"` + Paths map[string]pathItem `json:"paths"` +} + +type pathItem map[string]operation + +type operation struct { + Summary string `json:"summary"` + Description string `json:"description,omitempty"` + Tags []string `json:"tags,omitempty"` + Parameters []parameter `json:"parameters,omitempty"` + RequestBody *requestBody `json:"requestBody,omitempty"` + Responses map[string]specResponse `json:"responses"` +} + +type parameter struct { + Name string `json:"name"` + In string `json:"in"` + Description string `json:"description,omitempty"` + Required bool `json:"required"` + Schema *Schema `json:"schema"` +} + +type requestBody struct { + Required bool `json:"required"` + Content map[string]mediaType `json:"content"` +} + +type mediaType struct { + Schema *Schema `json:"schema"` +} + +type specResponse struct { + Description string `json:"description"` + Content map[string]mediaType `json:"content,omitempty"` +} + +var errorSchema = &Schema{ + Type: "object", + Properties: map[string]*Schema{ + "error": {Type: "string"}, + }, + Required: []string{"error"}, +} + +// SpecJSON renders the endpoints as an OpenAPI 3.1 JSON document. Base +// is the API's base path (e.g. "/v1"); endpoint paths are relative to +// it. The output is deterministic: map keys are sorted by +// encoding/json. +func SpecJSON(info Info, base string, endpoints []Endpoint) ([]byte, error) { + doc := document{ + OpenAPI: "3.2.1", + Info: info, + Paths: map[string]pathItem{}, + } + for _, e := range endpoints { + path := strings.TrimPrefix(e.path, base) + op := operation{ + Summary: e.summary, + Description: e.description, + Tags: e.tags, + Responses: map[string]specResponse{}, + } + for _, p := range e.params { + op.Parameters = append(op.Parameters, parameter{ + Name: p.Name, + In: p.In, + Description: p.Description, + Required: p.Required, + Schema: &Schema{Type: p.Type}, + }) + } + if e.request != nil { + op.RequestBody = &requestBody{ + Required: true, + Content: map[string]mediaType{ + "application/json": {Schema: SchemaOf(e.request)}, + }, + } + } + status := e.Status() + ok := specResponse{Description: http.StatusText(status)} + switch { + case e.sse: + ok.Content = map[string]mediaType{ + "text/event-stream": {Schema: &Schema{Type: "string"}}, + } + case e.response != nil: + ok.Content = map[string]mediaType{ + "application/json": {Schema: SchemaOf(e.response)}, + } + } + op.Responses[strconv.Itoa(status)] = ok + for _, code := range e.errors { + op.Responses[strconv.Itoa(code)] = specResponse{ + Description: http.StatusText(code), + Content: map[string]mediaType{ + "application/json": {Schema: errorSchema}, + }, + } + } + + item := doc.Paths[path] + if item == nil { + item = pathItem{} + doc.Paths[path] = item + } + item[strings.ToLower(e.method)] = op + } + return json.MarshalIndent(doc, "", " ") +} + +// Tags returns the sorted set of tags across endpoints. +func Tags(endpoints []Endpoint) []string { + seen := map[string]bool{} + var tags []string + for _, e := range endpoints { + for _, t := range e.tags { + if !seen[t] { + seen[t] = true + tags = append(tags, t) + } + } + } + sort.Strings(tags) + return tags +} diff --git a/internal/apigen/render.go b/internal/apigen/render.go new file mode 100644 index 0000000000000000000000000000000000000000..5e01912afae6427a5e75e55cb6bf3146024e3e94 --- /dev/null +++ b/internal/apigen/render.go @@ -0,0 +1,94 @@ +package apigen + +import ( + _ "embed" + "encoding/json" + "html/template" + "io" + "sort" +) + +// RenderDocs writes a plain, self-contained HTML documentation page for +// the endpoints, grouped by tag. +func RenderDocs(w io.Writer, info Info, base string, endpoints []Endpoint) error { + return docsTemplate.Execute(w, newDocsModel(info, base, endpoints)) +} + +type docsModel struct { + Info Info + Base string + Tags []tagGroup +} + +type tagGroup struct { + Name string + Endpoints []endpointView +} + +type endpointView struct { + Method string + Path string + Summary string + Description string + Params []Param + Request string + Response string + Status int + Errors []int + SSE bool +} + +func newDocsModel(info Info, base string, endpoints []Endpoint) docsModel { + byTag := map[string][]endpointView{} + for _, e := range endpoints { + view := endpointView{ + Method: e.method, + Path: e.path, + Summary: e.summary, + Description: e.description, + Params: e.params, + Status: e.Status(), + Errors: e.errors, + SSE: e.sse, + } + if e.request != nil { + view.Request = schemaJSON(SchemaOf(e.request)) + } + if e.response != nil { + view.Response = schemaJSON(SchemaOf(e.response)) + } + tags := e.tags + if len(tags) == 0 { + tags = []string{"other"} + } + for _, t := range tags { + byTag[t] = append(byTag[t], view) + } + } + + model := docsModel{Info: info, Base: base} + for _, name := range Tags(endpoints) { + group := tagGroup{Name: name, Endpoints: byTag[name]} + sort.Slice(group.Endpoints, func(i, j int) bool { + return group.Endpoints[i].Path < group.Endpoints[j].Path + }) + model.Tags = append(model.Tags, group) + } + if other, ok := byTag["other"]; ok { + model.Tags = append(model.Tags, tagGroup{Name: "other", Endpoints: other}) + } + return model +} + +func schemaJSON(s *Schema) string { + data, err := json.MarshalIndent(s, "", " ") + if err != nil { + return "" + } + return string(data) +} + +//go:embed docs.html.tmpl +var docsTemplateSource string + +var docsTemplate = template.Must(template.New("docs").Parse(docsTemplateSource)) diff --git a/internal/apigen/schema.go b/internal/apigen/schema.go new file mode 100644 index 0000000000000000000000000000000000000000..0b705fbe8c2b0c0ffce90262c6f0637bbd032c5b --- /dev/null +++ b/internal/apigen/schema.go @@ -0,0 +1,134 @@ +package apigen + +import ( + "reflect" + "strings" + "time" +) + +// Schema is a JSON-Schema-shaped description of a Go value, derived by +// reflection. It supports the subset of JSON Schema used by the docs +// page and the OpenAPI document. +type Schema struct { + Type string `json:"type,omitempty"` + Format string `json:"format,omitempty"` + Description string `json:"description,omitempty"` + Properties map[string]*Schema `json:"properties,omitempty"` + Required []string `json:"required,omitempty"` + Items *Schema `json:"items,omitempty"` + AdditionalProperties *Schema `json:"additionalProperties,omitempty"` +} + +var ( + timeType = reflect.TypeOf(time.Time{}) + errorType = reflect.TypeOf((*error)(nil)).Elem() +) + +// SchemaOf derives a [Schema] from the type of v, honoring encoding/json +// tags: field names come from json tags, and fields marked omitempty are +// not required. +// +// Structs from the same package as v's type are expanded recursively; +// structs from other packages (and true type cycles) are rendered as +// free-form objects described by their type name, which keeps schemas +// for wide external types honest without inventing shape. +func SchemaOf(v any) *Schema { + t := reflect.TypeOf(v) + if t == nil { + return &Schema{} + } + return schemaOf(t, rootPkg(t), map[reflect.Type]bool{}) +} + +// rootPkg returns the package path of the (possibly pointer) type v's +// named root, used to bound struct expansion to one package. +func rootPkg(t reflect.Type) string { + for t.Kind() == reflect.Pointer { + t = t.Elem() + } + return t.PkgPath() +} + +func schemaOf(t reflect.Type, pkg string, expanding map[reflect.Type]bool) *Schema { + switch t { + case timeType: + return &Schema{Type: "string", Format: "date-time"} + case errorType: + return &Schema{Type: "string"} + } + + switch t.Kind() { + case reflect.Pointer: + return schemaOf(t.Elem(), pkg, expanding) + case reflect.Bool: + return &Schema{Type: "boolean"} + case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, + reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64: + return &Schema{Type: "integer"} + case reflect.Float32, reflect.Float64: + return &Schema{Type: "number"} + case reflect.String: + return &Schema{Type: "string"} + case reflect.Slice, reflect.Array: + return &Schema{Type: "array", Items: schemaOf(t.Elem(), pkg, expanding)} + case reflect.Map: + if t.Key().Kind() != reflect.String { + return &Schema{Type: "object"} + } + return &Schema{Type: "object", AdditionalProperties: schemaOf(t.Elem(), pkg, expanding)} + case reflect.Struct: + return structSchema(t, pkg, expanding) + default: + // Interfaces (including any) and anything else: free-form. + return &Schema{} + } +} + +func structSchema(t reflect.Type, pkg string, expanding map[reflect.Type]bool) *Schema { + if t.PkgPath() != pkg { + return &Schema{Type: "object", Description: t.String()} + } + if expanding[t] { + return &Schema{Type: "object", Description: t.String()} + } + expanding[t] = true + defer delete(expanding, t) + + s := &Schema{Type: "object"} + for i := range t.NumField() { + f := t.Field(i) + tag := f.Tag.Get("json") + if tag == "-" { + continue + } + name, opts, _ := strings.Cut(tag, ",") + if f.Anonymous && name == "" { + // Embedded structs contribute their fields, per + // encoding/json semantics, even when the embedded + // type itself is unexported. + emb := schemaOf(f.Type, pkg, expanding) + for k, v := range emb.Properties { + if s.Properties == nil { + s.Properties = map[string]*Schema{} + } + s.Properties[k] = v + } + s.Required = append(s.Required, emb.Required...) + continue + } + if !f.IsExported() { + continue + } + if name == "" { + name = f.Name + } + if s.Properties == nil { + s.Properties = map[string]*Schema{} + } + s.Properties[name] = schemaOf(f.Type, pkg, expanding) + if !strings.Contains(opts, "omitempty") { + s.Required = append(s.Required, name) + } + } + return s +} diff --git a/internal/backend/accepted_run_integration_test.go b/internal/backend/accepted_run_integration_test.go new file mode 100644 index 0000000000000000000000000000000000000000..9a0a7ba249cc547956dd479cabcf5545a07a5c26 --- /dev/null +++ b/internal/backend/accepted_run_integration_test.go @@ -0,0 +1,131 @@ +package backend + +import ( + "context" + "testing" + "time" + + "charm.land/fantasy" + "github.com/charmbracelet/crush/internal/agent" + "github.com/charmbracelet/crush/internal/agent/agenttest" + "github.com/charmbracelet/crush/internal/db" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/session" + "github.com/stretchr/testify/require" +) + +// gatedCoordinator wraps a real agent.Coordinator and parks RunAccepted +// before delegating to it. Every method other than RunAccepted is +// inherited from the embedded coordinator, so BeginAccepted (called by +// Backend.SendMessage) and RunAccepted (called by the dispatched run) +// are the production agent.Coordinator implementations under test, not +// stubs. The gate only delays entry into the real RunAccepted so a +// cancel can be made to land in the accepted-but-not-yet-active window +// deterministically: the accept handle is not consumed by +// sessionAgent.Run until the real RunAccepted runs after the gate opens. +type gatedCoordinator struct { + agent.Coordinator + entered chan struct{} + gate chan struct{} +} + +func (c *gatedCoordinator) RunAccepted(ctx context.Context, accept *agent.AcceptedRun, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + close(c.entered) + <-c.gate + return c.Coordinator.RunAccepted(ctx, accept, sessionID, prompt, attachments...) +} + +// newRealCoordinator builds a production agent.Coordinator over a +// DB-backed session/message store, wrapped in a gate. It is constructed +// through the real agent.NewCoordinator path (via the test-only +// agenttest helper) with an offline-resolvable model: the +// cancel-on-entry path under test persists a canceled turn and returns +// before any model call, so no network I/O happens. +func newRealCoordinator(t *testing.T) (*gatedCoordinator, session.Service, message.Service) { + 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) + messages := message.NewService(q) + + coord, err := agenttest.NewCoordinator(t.Context(), t.TempDir(), sessions, messages) + require.NoError(t, err) + + return &gatedCoordinator{ + Coordinator: coord, + entered: make(chan struct{}), + gate: make(chan struct{}), + }, sessions, messages +} + +// TestSendMessage_AcceptedCancelRace_RealMachinery exercises the +// 202/cancel race end-to-end through Backend.SendMessage against the +// production agent.Coordinator (BeginAccepted + RunAccepted), not a +// stub. It asserts that a cancel arriving after the prompt is accepted +// but before the run becomes active is not lost: the accepted handle +// reaches sessionAgent.Run and drives cancel-on-entry, which persists a +// canceled turn instead of streaming. +// +// This test would fail if Coordinator.BeginAccepted returned nil (Cancel +// would find no accepted run and record no mark, and the run would +// receive a nil Accepted handle and skip cancel-on-entry) or if +// Coordinator.RunAccepted dropped the handle on its way into +// sessionAgent.Run (the run would likewise skip cancel-on-entry and try +// to stream the model). In either case no FinishReasonCanceled turn +// would be persisted. +func TestSendMessage_AcceptedCancelRace_RealMachinery(t *testing.T) { + t.Parallel() + b, _ := newTestBackend(t) + + coord, sessions, messages := newRealCoordinator(t) + sess, err := sessions.Create(t.Context(), "session") + require.NoError(t, err) + + ws := insertAgentWorkspace(t, b, coord) + + require.NoError(t, b.SendMessage(ws.ID, proto.AgentMessage{SessionID: sess.ID, Prompt: "hi"})) + + // Coordinator.BeginAccepted ran synchronously inside SendMessage + // before dispatch; the dispatched run has now entered the gate but + // has not yet called the real RunAccepted, so the accept handle is + // not yet consumed: the prompt is accepted but not active. + select { + case <-coord.entered: + case <-time.After(2 * time.Second): + t.Fatal("dispatched run never entered RunAccepted") + } + + // A cancel arriving now lands in the accepted-but-not-yet-active + // window and is only recorded because BeginAccepted incremented the + // accept counter. + require.NoError(t, b.CancelSession(ws.ID, sess.ID)) + + // Release the gate so the real RunAccepted threads the handle into + // sessionAgent.Run, which drives cancel-on-entry. + close(coord.gate) + + // The dispatched run returns nil (cancel-on-entry), so runWG drains. + waited := make(chan struct{}) + go func() { + ws.runWG.Wait() + close(waited) + }() + select { + case <-waited: + case <-time.After(5 * time.Second): + t.Fatal("runWG.Wait did not complete after the canceled run returned") + } + + // The accepted-but-not-yet-active cancel persisted a canceled turn + // rather than streaming a real response. + msgs, err := messages.List(t.Context(), sess.ID) + require.NoError(t, err) + require.Len(t, msgs, 2) + require.Equal(t, message.User, msgs[0].Role) + require.Equal(t, message.Assistant, msgs[1].Role) + require.Equal(t, message.FinishReasonCanceled, msgs[1].FinishReason()) +} diff --git a/internal/backend/agent.go b/internal/backend/agent.go new file mode 100644 index 0000000000000000000000000000000000000000..4f4691e566aaf471949d8f7eff173a6fa99698fa --- /dev/null +++ b/internal/backend/agent.go @@ -0,0 +1,276 @@ +package backend + +import ( + "context" + "errors" + "os" + + "github.com/charmbracelet/crush/internal/agent" + "github.com/charmbracelet/crush/internal/agent/notify" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/charmbracelet/crush/internal/shell" +) + +// SendMessage validates and accepts a prompt for the workspace's agent, +// then dispatches the run on a goroutine bound to the workspace context +// and returns immediately. It does not wait for the LLM turn to +// complete: the run's lifetime is owned by the workspace, not by the +// caller. Errors from the dispatched run reach observers through the +// agent event channels (a notify.TypeAgentError notification), not +// through this return value. +// +// SendMessage returns synchronously when the request cannot be accepted: +// ErrWorkspaceNotFound if the workspace is missing, ErrAgentNotInitialized +// if its coordinator is nil, the structural validation errors from +// agent.ValidateCall (ErrEmptyPrompt, ErrSessionMissing) when the prompt +// or session is missing, and ErrWorkspaceClosing if the workspace is +// being torn down. +func (b *Backend) SendMessage(workspaceID string, msg proto.AgentMessage) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + + if ws.AgentCoordinator == nil { + return ErrAgentNotInitialized + } + + if err := agent.ValidateCall(agent.SessionAgentCall{ + SessionID: msg.SessionID, + Prompt: msg.Prompt, + Attachments: proto.AttachmentsToMessage(msg.Attachments), + }); err != nil { + return err + } + + accept := ws.AgentCoordinator.BeginAccepted(msg.SessionID) + + ws.runMu.Lock() + if ws.closing { + ws.runMu.Unlock() + accept.Close() + return ErrWorkspaceClosing + } + ws.runWG.Add(1) + ws.runMu.Unlock() + + go b.runAgent(ws, msg, accept) + return nil +} + +// runAgent executes an accepted agent run for the workspace. It owns the +// accept reservation (releasing it on return) and the runWG ticket added +// by SendMessage. The run is bound to the workspace context so its +// lifetime is independent of any client's HTTP request. +// +// On a non-cancel error it surfaces the failure to observers via a +// notify.TypeAgentError notification (lossy, best-effort). That alone is +// not a reliable terminal signal: the agent-event fan-in uses lossy +// subscribers, so a `crush run` caller blocking on its RunID could hang +// if the event is dropped. To guarantee termination, when msg.RunID is +// non-empty and the coordinator did not already publish the run's +// authoritative terminal RunComplete (e.g. the error was returned before +// sessionAgent.Run executed, such as a readyWg or UpdateModels failure), +// runAgent emits an errored RunComplete on the must-deliver +// runCompletions broker so the waiter observes a deterministic terminal +// event. context.Canceled is expected (sessionAgent.Run already +// publishes the cancelled terminal marker) and produces no error +// terminal event. +// +// When msg.RunID is non-empty it is attached to the context via +// agent.WithRunID so the coordinator can stamp the terminal +// notify.RunComplete event with that correlator. A run-complete marker +// is also attached so the coordinator can report whether it published +// the terminal event, letting runAgent avoid a duplicate fallback. +func (b *Backend) runAgent(ws *Workspace, msg proto.AgentMessage, accept *agent.AcceptedRun) { + defer ws.runWG.Done() + defer accept.Close() + + ctx := ws.ctx + if msg.RunID != "" { + ctx = agent.WithRunID(ctx, msg.RunID) + } + ctx = agent.WithRunCompleteMarker(ctx) + + _, err := ws.AgentCoordinator.RunAccepted(ctx, accept, msg.SessionID, msg.Prompt, proto.AttachmentsToMessage(msg.Attachments)...) + if err == nil || errors.Is(err, context.Canceled) { + return + } + + ws.AgentNotifications().Publish(pubsub.CreatedEvent, notify.Notification{ + SessionID: msg.SessionID, + RunID: msg.RunID, + Type: notify.TypeAgentError, + Message: err.Error(), + }) + + // Reliable terminal fallback. Only needed when a RunID waiter + // exists and the coordinator has not already emitted the run's + // terminal RunComplete; otherwise this would be a duplicate. + if msg.RunID == "" || agent.RunCompletePublished(ctx) { + return + } + if rc := ws.RunCompletions(); rc != nil { + rc.PublishMustDeliver(ctx, pubsub.UpdatedEvent, notify.RunComplete{ + SessionID: msg.SessionID, + RunID: msg.RunID, + Error: err.Error(), + }) + } +} + +// GetAgentInfo returns the agent's model and busy status. +func (b *Backend) GetAgentInfo(workspaceID string) (proto.AgentInfo, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return proto.AgentInfo{}, err + } + + var agentInfo proto.AgentInfo + if ws.AgentCoordinator != nil { + m := ws.AgentCoordinator.Model() + agentInfo = proto.AgentInfo{ + Model: m.CatwalkCfg, + ModelCfg: m.ModelCfg, + IsBusy: ws.AgentCoordinator.IsBusy(), + IsReady: true, + } + } + return agentInfo, nil +} + +// InitAgent initializes the coder agent for the workspace. +func (b *Backend) InitAgent(ctx context.Context, workspaceID string, interactive bool) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + + if interactive { + return ws.InitCoderAgent(ctx) + } + return ws.InitCoderAgentNonInteractive(ctx) +} + +// UpdateAgent reloads the agent model configuration. +func (b *Backend) UpdateAgent(ctx context.Context, workspaceID string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + + return ws.UpdateAgentModel(ctx) +} + +// CancelSession cancels an ongoing agent operation for the given +// session. +func (b *Backend) CancelSession(workspaceID, sessionID string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + + if ws.AgentCoordinator != nil { + ws.AgentCoordinator.Cancel(sessionID) + } + return nil +} + +// SummarizeSession triggers a session summarization. +func (b *Backend) SummarizeSession(ctx context.Context, workspaceID, sessionID string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + + if ws.AgentCoordinator == nil { + return ErrAgentNotInitialized + } + + return ws.AgentCoordinator.Summarize(ctx, sessionID) +} + +// QueuedPrompts returns the number of queued prompts for the session. +func (b *Backend) QueuedPrompts(workspaceID, sessionID string) (int, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return 0, err + } + + if ws.AgentCoordinator == nil { + return 0, nil + } + + return ws.AgentCoordinator.QueuedPrompts(sessionID), nil +} + +// ClearQueue clears the prompt queue for the session. +func (b *Backend) ClearQueue(workspaceID, sessionID string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + + if ws.AgentCoordinator != nil { + ws.AgentCoordinator.ClearQueue(sessionID) + } + return nil +} + +// QueuedPromptsList returns the list of queued prompt strings for a +// session. +func (b *Backend) QueuedPromptsList(workspaceID, sessionID string) ([]string, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + + if ws.AgentCoordinator == nil { + return nil, nil + } + + return ws.AgentCoordinator.QueuedPromptsList(sessionID), nil +} + +// GetDefaultSmallModel returns the default small model for a provider. +func (b *Backend) GetDefaultSmallModel(workspaceID, providerID string) (config.SelectedModel, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return config.SelectedModel{}, err + } + + return ws.GetDefaultSmallModel(providerID), nil +} + +// RunShellCommand runs a shell command in the workspace directory and +// persists the command + output as a user message in the session. +func (b *Backend) RunShellCommand(ctx context.Context, workspaceID string, req proto.ShellCommandRequest) (proto.ShellCommandResponse, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return proto.ShellCommandResponse{}, err + } + + var persist shell.PersistFunc + if req.SessionID != "" { + persist = func(cmd, output string, exitCode int) error { + return shell.PersistOutput(ctx, ws.Messages, req.SessionID, cmd, output, exitCode) + } + } + + result, err := shell.RunAndPersist(ctx, shell.RunOptions{ + Command: req.Command, + Cwd: ws.Path, + Env: append(os.Environ(), ws.Env...), + TermWidth: req.TermWidth, + }, persist) + if err != nil { + return proto.ShellCommandResponse{}, err + } + + return proto.ShellCommandResponse{ + Output: result.Output, + ExitCode: result.ExitCode, + }, nil +} diff --git a/internal/backend/agent_runcomplete_test.go b/internal/backend/agent_runcomplete_test.go new file mode 100644 index 0000000000000000000000000000000000000000..68a1a7ee1205cc7481c20d0264bf0eb61eda7085 --- /dev/null +++ b/internal/backend/agent_runcomplete_test.go @@ -0,0 +1,163 @@ +package backend + +import ( + "context" + "errors" + "testing" + "time" + + "charm.land/fantasy" + "github.com/charmbracelet/crush/internal/agent" + "github.com/charmbracelet/crush/internal/app" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/proto" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +// errorCoordinator is a minimal agent.Coordinator whose RunAccepted +// returns a configurable error. When markPublished is true it stamps +// the run-complete marker on the context before returning, simulating a +// real coordinator that already published the run's authoritative +// terminal RunComplete (so runAgent must not emit a duplicate fallback). +type errorCoordinator struct { + err error + markPublished bool +} + +func (c *errorCoordinator) Run(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + return nil, c.err +} + +func (c *errorCoordinator) RunAccepted(ctx context.Context, accept *agent.AcceptedRun, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + if c.markPublished { + agent.MarkRunCompletePublished(ctx) + } + return nil, c.err +} + +func (c *errorCoordinator) BeginAccepted(sessionID string) *agent.AcceptedRun { return nil } +func (c *errorCoordinator) Cancel(string) {} +func (c *errorCoordinator) CancelAll() {} +func (c *errorCoordinator) IsBusy() bool { return false } +func (c *errorCoordinator) IsSessionBusy(string) bool { return false } +func (c *errorCoordinator) QueuedPrompts(string) int { return 0 } +func (c *errorCoordinator) QueuedPromptsList(string) []string { return nil } +func (c *errorCoordinator) ClearQueue(string) {} +func (c *errorCoordinator) Summarize(context.Context, string) error { return nil } +func (c *errorCoordinator) Model() agent.Model { return agent.Model{} } +func (c *errorCoordinator) UpdateModels(context.Context) error { return nil } +func (c *errorCoordinator) GenerateTitle(context.Context, string, string) {} + +// insertRunCompleteWorkspace installs a workspace backed by a real +// app.App (so the runCompletions broker exists) with the given +// coordinator and a workspace run context derived from base. +func insertRunCompleteWorkspace(t *testing.T, b *Backend, base context.Context, coord agent.Coordinator) *Workspace { + t.Helper() + a := app.NewForTest(base) + a.AgentCoordinator = coord + t.Cleanup(a.ShutdownForTest) + ws := &Workspace{ + ID: uuid.New().String(), + Path: t.TempDir(), + resolvedPath: t.TempDir(), + clients: make(map[string]*clientState), + shutdownFn: func() {}, + } + ws.App = a + ws.ctx, ws.cancel = context.WithCancel(base) + b.mu.Lock() + b.workspaces.Set(ws.ID, ws) + b.pathIndex[ws.resolvedPath] = ws.ID + b.mu.Unlock() + return ws +} + +// TestRunAgent_PreRunErrorPublishesTerminalRunComplete proves that an +// error returned from RunAccepted before the coordinator could publish +// its own terminal event (e.g. a readyWg or UpdateModels failure, +// modeled here by a stub coordinator) still yields a reliable terminal +// RunComplete for the run's RunID. Without it, a `crush run` caller +// blocking on that RunID would hang because the lossy TypeAgentError +// event is not a guaranteed terminal signal. +func TestRunAgent_PreRunErrorPublishesTerminalRunComplete(t *testing.T) { + t.Parallel() + b, _ := newTestBackend(t) + runErr := errors.New("update models failed") + ws := insertRunCompleteWorkspace(t, b, context.Background(), &errorCoordinator{err: runErr}) + + subCtx, cancel := context.WithCancel(context.Background()) + defer cancel() + ch := ws.RunCompletions().Subscribe(subCtx) + + err := b.SendMessage(ws.ID, proto.AgentMessage{SessionID: "S1", RunID: "run-1", Prompt: "hi"}) + require.NoError(t, err) + + select { + case ev := <-ch: + require.Equal(t, "run-1", ev.Payload.RunID, + "the terminal RunComplete must carry the dispatched RunID") + require.Equal(t, "S1", ev.Payload.SessionID) + require.Equal(t, runErr.Error(), ev.Payload.Error, + "the fallback terminal event must be marked errored") + require.False(t, ev.Payload.Cancelled) + case <-time.After(2 * time.Second): + t.Fatal("no terminal RunComplete published for a pre-run error; a run waiter would hang") + } +} + +// TestRunAgent_NoFallbackWhenCoordinatorPublished ensures the fallback +// is suppressed when the coordinator already emitted the run's +// authoritative terminal RunComplete, so callers never observe a +// duplicate terminal event for the same RunID. +func TestRunAgent_NoFallbackWhenCoordinatorPublished(t *testing.T) { + t.Parallel() + b, _ := newTestBackend(t) + runErr := errors.New("stream failed after publishing terminal event") + ws := insertRunCompleteWorkspace(t, b, context.Background(), + &errorCoordinator{err: runErr, markPublished: true}) + + subCtx, cancel := context.WithCancel(context.Background()) + defer cancel() + ch := ws.RunCompletions().Subscribe(subCtx) + + err := b.SendMessage(ws.ID, proto.AgentMessage{SessionID: "S1", RunID: "run-1", Prompt: "hi"}) + require.NoError(t, err) + + // Wait for the dispatched run goroutine to return so any publish + // has already happened. + ws.runWG.Wait() + + select { + case ev := <-ch: + t.Fatalf("runAgent published a duplicate terminal RunComplete: %+v", ev.Payload) + case <-time.After(200 * time.Millisecond): + } +} + +// TestRunAgent_CancellationPublishesNoErrorTerminal verifies that a +// context.Canceled result from RunAccepted produces no errored terminal +// RunComplete from runAgent: cancellation is sessionAgent.Run's +// responsibility (it publishes the cancelled marker) and the dispatcher +// must not synthesize an error terminal for it. +func TestRunAgent_CancellationPublishesNoErrorTerminal(t *testing.T) { + t.Parallel() + b, _ := newTestBackend(t) + ws := insertRunCompleteWorkspace(t, b, context.Background(), + &errorCoordinator{err: context.Canceled}) + + subCtx, cancel := context.WithCancel(context.Background()) + defer cancel() + ch := ws.RunCompletions().Subscribe(subCtx) + + err := b.SendMessage(ws.ID, proto.AgentMessage{SessionID: "S1", RunID: "run-1", Prompt: "hi"}) + require.NoError(t, err) + + ws.runWG.Wait() + + select { + case ev := <-ch: + t.Fatalf("cancellation must not publish a terminal RunComplete: %+v", ev.Payload) + case <-time.After(200 * time.Millisecond): + } +} diff --git a/internal/backend/agent_test.go b/internal/backend/agent_test.go new file mode 100644 index 0000000000000000000000000000000000000000..f5576fd32eab06e826371b6349c77c8a28673f2c --- /dev/null +++ b/internal/backend/agent_test.go @@ -0,0 +1,164 @@ +package backend + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "charm.land/fantasy" + "github.com/charmbracelet/crush/internal/agent" + "github.com/charmbracelet/crush/internal/app" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/proto" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +// blockingCoordinator is a minimal agent.Coordinator whose RunAccepted +// blocks until release is closed. It records that RunAccepted was +// entered so tests can observe the dispatched goroutine. Every other +// method returns a zero value. +type blockingCoordinator struct { + entered chan struct{} + release chan struct{} + runCount atomic.Int32 +} + +func newBlockingCoordinator() *blockingCoordinator { + return &blockingCoordinator{ + entered: make(chan struct{}, 1), + release: make(chan struct{}), + } +} + +func (c *blockingCoordinator) Run(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + return nil, nil +} + +func (c *blockingCoordinator) RunAccepted(ctx context.Context, accept *agent.AcceptedRun, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + c.runCount.Add(1) + select { + case c.entered <- struct{}{}: + default: + } + <-c.release + return nil, nil +} + +func (c *blockingCoordinator) BeginAccepted(sessionID string) *agent.AcceptedRun { return nil } +func (c *blockingCoordinator) Cancel(string) {} +func (c *blockingCoordinator) CancelAll() {} +func (c *blockingCoordinator) IsBusy() bool { return false } +func (c *blockingCoordinator) IsSessionBusy(string) bool { return false } +func (c *blockingCoordinator) QueuedPrompts(string) int { return 0 } +func (c *blockingCoordinator) QueuedPromptsList(string) []string { return nil } +func (c *blockingCoordinator) ClearQueue(string) {} +func (c *blockingCoordinator) Summarize(context.Context, string) error { return nil } +func (c *blockingCoordinator) Model() agent.Model { return agent.Model{} } +func (c *blockingCoordinator) UpdateModels(context.Context) error { return nil } +func (c *blockingCoordinator) GenerateTitle(context.Context, string, string) {} + +// insertAgentWorkspace installs a synthetic workspace with the given +// coordinator (or none) and a workspace run context, mirroring the +// fields CreateWorkspace initializes. +func insertAgentWorkspace(t *testing.T, b *Backend, coord agent.Coordinator) *Workspace { + t.Helper() + ws := &Workspace{ + ID: uuid.New().String(), + Path: t.TempDir(), + resolvedPath: t.TempDir(), + clients: make(map[string]*clientState), + shutdownFn: func() {}, + } + ws.App = &app.App{AgentCoordinator: coord} + ws.ctx, ws.cancel = context.WithCancel(b.ctx) + b.mu.Lock() + b.workspaces.Set(ws.ID, ws) + b.pathIndex[ws.resolvedPath] = ws.ID + b.mu.Unlock() + return ws +} + +func TestSendMessage_WorkspaceNotFound(t *testing.T) { + t.Parallel() + b, _ := newTestBackend(t) + err := b.SendMessage("nope", proto.AgentMessage{SessionID: "S1", Prompt: "hi"}) + require.ErrorIs(t, err, ErrWorkspaceNotFound) +} + +func TestSendMessage_AgentNotInitialized(t *testing.T) { + t.Parallel() + b, _ := newTestBackend(t) + ws := insertAgentWorkspace(t, b, nil) + err := b.SendMessage(ws.ID, proto.AgentMessage{SessionID: "S1", Prompt: "hi"}) + require.ErrorIs(t, err, ErrAgentNotInitialized) +} + +func TestSendMessage_EmptyPrompt(t *testing.T) { + t.Parallel() + b, _ := newTestBackend(t) + ws := insertAgentWorkspace(t, b, newBlockingCoordinator()) + err := b.SendMessage(ws.ID, proto.AgentMessage{SessionID: "S1", Prompt: ""}) + require.ErrorIs(t, err, agent.ErrEmptyPrompt) +} + +func TestSendMessage_SessionMissing(t *testing.T) { + t.Parallel() + b, _ := newTestBackend(t) + ws := insertAgentWorkspace(t, b, newBlockingCoordinator()) + err := b.SendMessage(ws.ID, proto.AgentMessage{SessionID: "", Prompt: "hi"}) + require.ErrorIs(t, err, agent.ErrSessionMissing) +} + +func TestSendMessage_WorkspaceClosing(t *testing.T) { + t.Parallel() + b, _ := newTestBackend(t) + ws := insertAgentWorkspace(t, b, newBlockingCoordinator()) + ws.runMu.Lock() + ws.closing = true + ws.runMu.Unlock() + err := b.SendMessage(ws.ID, proto.AgentMessage{SessionID: "S1", Prompt: "hi"}) + require.ErrorIs(t, err, ErrWorkspaceClosing) +} + +// TestSendMessage_SuccessIncrementsRunWG asserts the happy path returns +// nil synchronously and dispatches a tracked goroutine: while +// RunAccepted blocks, runWG.Wait must not complete (the ticket is +// outstanding); after release it drains. +func TestSendMessage_SuccessIncrementsRunWG(t *testing.T) { + t.Parallel() + b, _ := newTestBackend(t) + coord := newBlockingCoordinator() + ws := insertAgentWorkspace(t, b, coord) + + err := b.SendMessage(ws.ID, proto.AgentMessage{SessionID: "S1", Prompt: "hi"}) + require.NoError(t, err) + + select { + case <-coord.entered: + case <-time.After(2 * time.Second): + t.Fatal("dispatched goroutine never entered RunAccepted") + } + require.Equal(t, int32(1), coord.runCount.Load()) + + waited := make(chan struct{}) + go func() { + ws.runWG.Wait() + close(waited) + }() + + select { + case <-waited: + t.Fatal("runWG.Wait completed while the run was still in flight; ticket was not added") + case <-time.After(100 * time.Millisecond): + } + + close(coord.release) + + select { + case <-waited: + case <-time.After(2 * time.Second): + t.Fatal("runWG.Wait did not complete after the run returned") + } +} diff --git a/internal/backend/backend.go b/internal/backend/backend.go new file mode 100644 index 0000000000000000000000000000000000000000..a69024fc2435a01c2b461be1ed50e6ed37864e67 --- /dev/null +++ b/internal/backend/backend.go @@ -0,0 +1,1137 @@ +// Package backend provides transport-agnostic operations for managing +// workspaces, sessions, agents, permissions, and events. It is consumed +// by protocol-specific layers such as HTTP (server) and ACP. +package backend + +import ( + "context" + "errors" + "fmt" + "log/slog" + "os" + "path/filepath" + "runtime" + "strconv" + "sync" + "time" + + "github.com/charmbracelet/crush/internal/app" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/csync" + "github.com/charmbracelet/crush/internal/db" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/skills" + "github.com/charmbracelet/crush/internal/ui/util" + "github.com/charmbracelet/crush/internal/version" + "github.com/google/uuid" +) + +// Common errors returned by backend operations. +var ( + ErrWorkspaceNotFound = errors.New("workspace not found") + ErrLSPClientNotFound = errors.New("LSP client not found") + ErrAgentNotInitialized = errors.New("agent coordinator not initialized") + ErrPathRequired = errors.New("path is required") + ErrInvalidPermissionAction = errors.New("invalid permission action") + ErrUnknownCommand = errors.New("unknown command") + ErrInvalidClientID = errors.New("invalid client_id") + ErrClientNotAttached = errors.New("client not attached") + ErrWorkspaceClosing = errors.New("workspace closing") + ErrServerShuttingDown = errors.New("server is shutting down") + ErrServerNotIdle = errors.New("server is hosting live workspaces") + ErrClientRetired = errors.New("client has been retired") + ErrChannelOptInMismatch = errors.New("requested channels differ from the existing workspace; channels are an explicit opt-in and are not shared across duplicate creates") +) + +// DefaultCreateGrace is the window in which a client must open an SSE +// stream after creating a workspace before its creation hold is +// released. Exposed as a package variable so tests can shorten it. +var DefaultCreateGrace = 30 * time.Second + +// DefaultIdleShutdownDelay is how long the server stays alive after its +// last workspace is released before it shuts itself down. The delay +// exists so a client that closes one session and opens another moments +// later (the same directory or a different one) reuses the still-running +// server instead of racing its shutdown: with an immediate shutdown the +// new client can attach to — or create a workspace on — a server that is +// already tearing down, and then observe its coder agent as "offline". +// Any workspace create within the window cancels the pending shutdown. +// Overridable via CRUSH_SERVER_IDLE_TIMEOUT (seconds; 0 restores the +// old shut-down-immediately behavior). +var DefaultIdleShutdownDelay = 60 * time.Second + +// DefaultDetachGrace is how long a client's claim on a workspace survives +// after its last SSE stream drops without an explicit release. The stream +// is the client's refcount claim, so tearing the workspace down the instant +// it closes turns any momentary drop (a hiccup, a suspended laptop, a proxy +// timeout) into a permanently lost workspace: the client's reconnect comes +// back milliseconds later to an ID the server no longer knows. A client +// that released its claim first (a clean exit) skips the grace. Overridable +// via CRUSH_SERVER_DETACH_GRACE (seconds; 0 restores immediate teardown). +var DefaultDetachGrace = 10 * time.Second + +// ShutdownFunc is called when the backend needs to trigger a server +// shutdown (e.g. when the last workspace is removed). +type ShutdownFunc func() + +// Backend provides transport-agnostic business logic for the Crush +// server. It manages workspaces and delegates to [app.App] services. +// +// Locking order: when both [Backend.mu] and [Workspace.clientsMu] are +// held at once, [Backend.mu] is acquired first. Detach paths +// ([detachStream], [releaseHoldLocked], [expireHold]) only hold +// [Workspace.clientsMu] briefly, drop it, then call [teardown] which +// takes [Backend.mu] (and then re-takes [Workspace.clientsMu] to +// re-check that the workspace has not been re-claimed). This avoids +// the AB/BA hazard with [CreateWorkspace], which holds [Backend.mu] +// while calling [registerClient] so that a workspace cannot be torn +// down beneath it. +type Backend struct { + workspaces *csync.Map[string, *Workspace] + // pathIndex maps a resolved absolute workspace path to its + // workspace ID. Reads and writes are serialised via mu so + // concurrent CreateWorkspace calls at the same path deduplicate + // deterministically. + pathIndex map[string]string + // pending counts CreateWorkspace calls that have committed to the + // slow initialization path (config/db/app setup) but have not yet + // registered their workspace in the map. It is guarded by mu. + // teardown must observe pending == 0 in addition to an empty + // workspace map before triggering server shutdown: otherwise a + // teardown of the last live workspace could race ahead of a + // concurrent create — which releases mu during its slow init — and + // shut the whole server down out from under the workspace being + // born. + pending int + // shutdownTimer, when non-nil, is an armed idle-shutdown timer + // waiting out lingerDelay before it shuts the server down. It is + // guarded by mu and cancelled the moment a new create arrives. + shutdownTimer *time.Timer + // closing latches the decision to exit. Every site that commits to + // a shutdown sets it while still holding mu, because shutdownFn has + // to run unlocked; CreateWorkspace refuses once it is set. That + // makes the shutdown-vs-create decision atomic: a create can never + // be handed a workspace on a process that is already leaving. + closing bool + // retired holds the IDs of clients that announced their exit via + // RetireClient. Creates from a retired client are refused, which is what + // lets a client release a workspace whose ID it never learned. + // + // Entries are never pruned, and deliberately so: an entry is exactly + // what refuses a create that was already on the wire when the client + // said goodbye, and HTTP gives no ordering between two requests, so + // there is no moment at which the server can prove none is still coming. + // Dropping an entry to save memory would reintroduce the orphaned + // workspace this mechanism exists to prevent. The cost is one UUID per + // client process, and the server only outlives its clients for as long + // as sessions keep arriving inside the idle-shutdown window. + retired map[string]struct{} + mu sync.Mutex + + cfg *config.ConfigStore + ctx context.Context + shutdownFn ShutdownFunc + createGrace time.Duration + lingerDelay time.Duration + detachGrace time.Duration +} + +// clientState tracks one client's claim on a workspace. +// +// - streams counts the number of live SSE event streams the client +// currently has open against the workspace. +// - holdTimer is non-nil in the two timer-held states: the client +// created the workspace but has not yet attached an SSE stream +// (fires after createGrace), or the client's last SSE stream dropped +// without an explicit release (fires after detachGrace, giving the +// client's reconnect loop a window to re-attach). Either way the +// timer releases the claim when it expires. +// - currentSessionID records which session this client is currently +// viewing. Empty string means the client has no session selected +// (e.g. the landing screen). Cleared automatically when the +// clientState entry is removed. +// - released marks that the client gave up its claim explicitly while +// streams were still open, which a clean exit does. The final stream +// detach then tears down immediately instead of waiting out the +// detach grace for a client that is not coming back. +// +// streams and holdTimer are mutually exclusive in practice (the hold +// timer is stopped the moment an SSE stream attaches), but both being +// zero/nil means the entry has been released and should be removed. +type clientState struct { + streams int + holdTimer *time.Timer + currentSessionID string + released bool +} + +// Workspace represents a running [app.App] workspace with its +// associated resources and state. +type Workspace struct { + *app.App + ID string + Path string + Cfg *config.ConfigStore + Env []string + Skills *skills.Manager + + // resolvedPath is the path used as the dedup key in + // Backend.pathIndex. It is filepath.EvalSymlinks(filepath.Abs(Path)) + // with fallback to the cleaned absolute path. + resolvedPath string + + // ctx is the workspace-scoped run context. It is derived from + // the backend context in CreateWorkspace and lives for the + // lifetime of the workspace; cancel tears it down. Agent runs + // dispatched on behalf of this workspace are bound to ctx so + // their lifetime is owned by the workspace, not by any single + // client's HTTP request. + ctx context.Context + cancel context.CancelFunc + + // runMu guards closing and gates dispatch of new agent runs. + // closing is set by Shutdown so no new runs are accepted once + // teardown has begun. runWG tracks dispatched agent goroutines + // so Shutdown can wait for them to return before app cleanup. + runMu sync.Mutex + closing bool + runWG sync.WaitGroup + + // clientsMu guards clients. It is held only briefly (no IO). + clientsMu sync.Mutex + // clients tracks each client's claim on this workspace. Refcount + // is a derived value: len(clients). + clients map[string]*clientState + + // shutdownFn is the function invoked by [Backend.teardown] to + // release the workspace's underlying resources. It defaults to the + // embedded [app.App.Shutdown]; tests may override it to avoid + // driving a full [app.App] through shutdown. + shutdownFn func() +} + +// invokeShutdown calls the workspace shutdown hook if set, falling +// back to the workspace [Workspace.Shutdown] wrapper when not. +func (w *Workspace) invokeShutdown() { + if w.shutdownFn != nil { + w.shutdownFn() + return + } + if w.App != nil { + w.Shutdown() + } +} + +// Shutdown tears the workspace down in an order that is safe for +// agent runs whose lifetime is bound to the workspace context. It +// shadows the promoted [app.App.Shutdown] so callers reaching +// ws.Shutdown() always observe this ordering: +// +// 1. Mark the workspace closing so no new agent runs are accepted. +// 2. Cancel the workspace run context so any dispatched goroutine +// that has not yet registered its per-session cancel still +// observes cancellation. +// 3. Cancel active coordinator work for runs that already +// registered their per-session cancel function. +// 4. Wait for dispatched agent goroutines to return. +// 5. Run the embedded [app.App.Shutdown] cleanup (DB, LSP, etc). +// +// CancelAll is idempotent, so the second call inside app.App.Shutdown +// is harmless; the important guarantee is that cancel -> CancelAll -> +// runWG.Wait completes before the embedded cleanup touches the DB. +func (w *Workspace) Shutdown() { + w.runMu.Lock() + w.closing = true + w.runMu.Unlock() + + if w.cancel != nil { + w.cancel() + } + if w.App != nil && w.AgentCoordinator != nil { + w.AgentCoordinator.CancelAll() + } + w.runWG.Wait() + if w.App != nil { + w.App.Shutdown() + } +} + +// New creates a new [Backend]. +func New(ctx context.Context, cfg *config.ConfigStore, shutdownFn ShutdownFunc) *Backend { + return &Backend{ + workspaces: csync.NewMap[string, *Workspace](), + pathIndex: make(map[string]string), + retired: make(map[string]struct{}), + cfg: cfg, + ctx: ctx, + shutdownFn: shutdownFn, + createGrace: DefaultCreateGrace, + lingerDelay: idleShutdownDelayFromEnv(), + detachGrace: durationFromEnv("CRUSH_SERVER_DETACH_GRACE", DefaultDetachGrace), + } +} + +// idleShutdownDelayFromEnv returns the idle-shutdown delay, honoring a +// CRUSH_SERVER_IDLE_TIMEOUT override (in seconds; 0 disables lingering). +func idleShutdownDelayFromEnv() time.Duration { + return durationFromEnv("CRUSH_SERVER_IDLE_TIMEOUT", DefaultIdleShutdownDelay) +} + +// durationFromEnv reads a whole number of seconds from the named +// environment variable, falling back to def when it is unset or +// unparseable. Zero is a meaningful value for both lifecycle windows it +// configures, so it is accepted. +func durationFromEnv(name string, def time.Duration) time.Duration { + if v := os.Getenv(name); v != "" { + if secs, err := strconv.Atoi(v); err == nil && secs >= 0 { + return time.Duration(secs) * time.Second + } + } + return def +} + +// SetCreateGrace overrides the create-grace window. Intended for tests +// that need short timeouts. +func (b *Backend) SetCreateGrace(d time.Duration) { + b.mu.Lock() + defer b.mu.Unlock() + b.createGrace = d +} + +// SetDetachGrace overrides how long a client's claim survives after its +// last SSE stream drops. A value <= 0 restores the tear-down-immediately +// behavior. Intended for tests. +func (b *Backend) SetDetachGrace(d time.Duration) { + b.mu.Lock() + defer b.mu.Unlock() + b.detachGrace = d +} + +// SetIdleShutdownDelay overrides how long the server lingers after its +// last workspace is released before shutting down. A value <= 0 restores +// the shut-down-immediately behavior. Intended for tests. +func (b *Backend) SetIdleShutdownDelay(d time.Duration) { + b.mu.Lock() + defer b.mu.Unlock() + b.lingerDelay = d +} + +// GetWorkspace retrieves a workspace by ID. +func (b *Backend) GetWorkspace(id string) (*Workspace, error) { + ws, ok := b.workspaces.Get(id) + if !ok { + return nil, ErrWorkspaceNotFound + } + return ws, nil +} + +// ListWorkspaces returns all running workspaces. +func (b *Backend) ListWorkspaces() []proto.Workspace { + workspaces := []proto.Workspace{} + for _, ws := range b.workspaces.Seq2() { + workspaces = append(workspaces, workspaceToProto(ws)) + } + return workspaces +} + +// CreateWorkspace initializes a new workspace from the given +// parameters, or returns an existing workspace if one already exists at +// the same resolved path (first-wins semantics). +// +// args.ClientID must be a valid UUID identifying the calling client; +// the resulting workspace registers a creation hold on behalf of that +// client which is released either by the first SSE attach (which +// converts it into a stream claim) or by the grace window expiring. +func (b *Backend) CreateWorkspace(args proto.Workspace) (*Workspace, proto.Workspace, error) { + if args.Path == "" { + return nil, proto.Workspace{}, ErrPathRequired + } + clientID, err := validateClientID(args.ClientID) + if err != nil { + return nil, proto.Workspace{}, err + } + + key, err := resolveWorkspaceKey(args.Path) + if err != nil { + return nil, proto.Workspace{}, fmt.Errorf("failed to resolve workspace path: %w", err) + } + + b.mu.Lock() + if err := b.admitLocked(clientID); err != nil { + b.mu.Unlock() + return nil, proto.Workspace{}, err + } + // A client is arriving: cancel any pending idle shutdown so we never + // hand back a workspace on a server that is about to tear itself down. + b.cancelShutdownLocked() + if existingID, ok := b.pathIndex[key]; ok { + if ws, found := b.workspaces.Get(existingID); found { + // Hold b.mu while registering: teardown also + // acquires b.mu before tearing the workspace + // down, so this guarantees the workspace we + // return cannot be torn out from under us + // between lookup and registerClient. Lock order + // here is b.mu -> ws.clientsMu. + if !stringSlicesEqual(ws.Cfg.Overrides().EnabledChannels, args.Channels) { + b.mu.Unlock() + return nil, proto.Workspace{}, ErrChannelOptInMismatch + } + logFirstWinsMismatch(ws, args) + b.registerClient(ws, clientID) + b.mu.Unlock() + return ws, workspaceToProto(ws), nil + } + // pathIndex referenced a workspace that has since been + // removed; clean the stale entry and fall through. + delete(b.pathIndex, key) + } + // Commit to the slow creation path. Mark this create as pending + // while mu is still held so a teardown that runs during the + // unlocked init below cannot observe an empty backend and shut the + // server down. The deferred decrement runs after the workspace has + // been registered (or the create has failed), keeping the invariant + // that pending only drops once the workspace is visible in the map. + b.pending++ + b.mu.Unlock() + defer func() { + b.mu.Lock() + b.pending-- + // If this create ended up registering nothing (it failed, or + // deduped onto an existing workspace that has since gone) and + // it was holding the last teardown back, the server may now be + // idle with no pending work. Arm the idle-shutdown timer here so + // a failed create racing the last teardown does not leak an + // empty server that a plain teardown already declined to reap. + shutdownNow := b.scheduleShutdownIfIdleLocked() + b.mu.Unlock() + if shutdownNow { + slog.Info("No workspaces remain after create settled, shutting down server...") + b.shutdownFn() + } + }() + + id := uuid.New().String() + cfg, err := config.Init(args.Path, args.DataDir, args.Debug) + if err != nil { + return nil, proto.Workspace{}, fmt.Errorf("failed to initialize config: %w", err) + } + + cfg.Overrides().SkipPermissionRequests = args.YOLO + cfg.Overrides().EnabledChannels = args.Channels + + if err := createDotCrushDir(cfg.Config().Options.DataDirectory); err != nil { + return nil, proto.Workspace{}, fmt.Errorf("failed to create data directory: %w", err) + } + + conn, err := db.Connect(b.ctx, cfg.Config().Options.DataDirectory, db.WithDataDirLock(true)) + if err != nil { + return nil, proto.Workspace{}, fmt.Errorf("failed to connect to database: %w", err) + } + + // Discover skills once per workspace, before app.New. The backend + // hosts multiple workspaces concurrently, so the manager is + // constructed WITHOUT WithGlobalMirror to prevent last-writer-wins + // cross-talk between workspaces. + discoveryCfg := skillsDiscoveryConfig(cfg) + allSkills, activeSkills, skillStates := skills.DiscoverFromConfig(discoveryCfg) + skillsMgr := skills.NewManager( + allSkills, activeSkills, skillStates, + skills.WithResolvedPaths(discoveryCfg.ResolvePaths()), + skills.WithWorkingDir(discoveryCfg.WorkingDir), + ) + + appWorkspace, err := app.New(b.ctx, conn, cfg, skillsMgr) + if err != nil { + return nil, proto.Workspace{}, fmt.Errorf("failed to create app workspace: %w", err) + } + + wsCtx, wsCancel := context.WithCancel(b.ctx) + ws := &Workspace{ + App: appWorkspace, + ID: id, + Path: args.Path, + Cfg: cfg, + Env: args.Env, + Skills: skillsMgr, + resolvedPath: key, + ctx: wsCtx, + cancel: wsCancel, + clients: make(map[string]*clientState), + } + + b.mu.Lock() + // Re-check admission: the client may have retired while the slow + // init above ran with b.mu released, and registering a claim for a + // client that has already announced its exit would strand the + // workspace. (b.closing cannot have flipped: every shutdown decision + // requires pending == 0, and this create has held pending since + // before it released b.mu.) + if err := b.admitLocked(clientID); err != nil { + b.mu.Unlock() + ws.invokeShutdown() + return nil, proto.Workspace{}, err + } + // Re-check the index under the lock: a concurrent caller may have + // won the race between the initial unlock and here. + if existingID, ok := b.pathIndex[key]; ok { + if existing, found := b.workspaces.Get(existingID); found { + // Register under b.mu so teardown cannot run + // between lookup and registerClient. Lock order + // is b.mu -> ws.clientsMu. + if !stringSlicesEqual(existing.Cfg.Overrides().EnabledChannels, args.Channels) { + b.mu.Unlock() + ws.invokeShutdown() + return nil, proto.Workspace{}, ErrChannelOptInMismatch + } + logFirstWinsMismatch(existing, args) + b.registerClient(existing, clientID) + b.mu.Unlock() + ws.invokeShutdown() + return existing, workspaceToProto(existing), nil + } + delete(b.pathIndex, key) + } + b.workspaces.Set(id, ws) + b.pathIndex[key] = id + // Register the originating client's hold while still holding + // b.mu so the workspace is observable with its claim from the + // moment it appears in the index. + b.registerClient(ws, clientID) + b.mu.Unlock() + + if args.Version != "" && args.Version != version.Version { + slog.Warn( + "Client/server version mismatch", + "client", args.Version, + "server", version.Version, + ) + appWorkspace.SendEvent(util.NewWarnMsg(fmt.Sprintf( + "Server version %q differs from client version %q. Consider restarting the server.", + version.Version, args.Version, + ))) + } + + return ws, workspaceToProto(ws), nil +} + +// skillsDiscoveryConfig adapts a *config.ConfigStore to the +// skills.DiscoveryConfig that DiscoverFromConfig consumes. +func skillsDiscoveryConfig(cfg *config.ConfigStore) skills.DiscoveryConfig { + opts := cfg.Config().Options + var paths, disabled []string + if opts != nil { + paths = opts.SkillsPaths + disabled = opts.DisabledSkills + } + var resolver func(string) (string, error) + if r := cfg.Resolver(); r != nil { + resolver = r.ResolveValue + } + return skills.DiscoveryConfig{ + SkillsPaths: paths, + DisabledSkills: disabled, + WorkingDir: cfg.WorkingDir(), + Resolver: resolver, + } +} + +// skillStatesToProto converts internal skill discovery states into the +// wire format. +func skillStatesToProto(states []*skills.SkillState) []proto.SkillState { + if len(states) == 0 { + return nil + } + out := make([]proto.SkillState, len(states)) + for i, s := range states { + entry := proto.SkillState{ + Name: s.Name, + Path: s.Path, + State: proto.SkillDiscoveryState(s.State), + } + if s.Err != nil { + entry.Error = s.Err.Error() + } + out[i] = entry + } + return out +} + +// AttachClient registers a new SSE stream for the given client on the +// workspace. The stream's deferred cleanup must call DetachClient with +// the same arguments to release the claim. +// +// The lookup and the clients-map mutation are performed under +// [Backend.mu] so that AttachClient cannot race with [Backend.teardown]: +// teardown also holds [Backend.mu] while removing the workspace from +// b.workspaces, so once AttachClient observes the workspace and takes +// ws.clientsMu (under b.mu), no concurrent teardown can succeed without +// re-checking the (now non-empty) clients map. Lock order is the +// canonical b.mu -> ws.clientsMu. +func (b *Backend) AttachClient(workspaceID, clientID string) error { + if _, err := validateClientID(clientID); err != nil { + return err + } + + b.mu.Lock() + defer b.mu.Unlock() + ws, ok := b.workspaces.Get(workspaceID) + if !ok { + return ErrWorkspaceNotFound + } + + ws.clientsMu.Lock() + defer ws.clientsMu.Unlock() + cs, ok := ws.clients[clientID] + if !ok { + // Defensive: SSE attach without a prior CreateWorkspace by + // this client still installs a stream claim so the stream + // stays alive for its duration. + ws.clients[clientID] = &clientState{streams: 1} + return nil + } + if cs.holdTimer != nil { + cs.holdTimer.Stop() + cs.holdTimer = nil + } + cs.streams++ + return nil +} + +// DetachClient releases one SSE stream's hold on the workspace. When the +// client has no streams left and no pending creation hold, its claim +// either enters the detach grace — giving a reconnecting client time to +// re-attach — or, if the grace is disabled or the client already released +// its claim, is removed, tearing the workspace down once the refcount +// hits zero. +func (b *Backend) DetachClient(workspaceID, clientID string) { + ws, ok := b.workspaces.Get(workspaceID) + if !ok { + return + } + b.detachStream(ws, clientID) +} + +// admitLocked reports whether clientID may still take a claim on this +// server. It must be called with b.mu held, which is what makes the +// answer atomic with respect to the shutdown latch and to RetireClient. +func (b *Backend) admitLocked(clientID string) error { + if b.closing { + return ErrServerShuttingDown + } + if _, ok := b.retired[clientID]; ok { + return ErrClientRetired + } + return nil +} + +// RetireClient records that a client has exited and releases every claim it +// holds, across every workspace. It is the authoritative "this client is +// gone" signal, and the reason a client never has to guess whether a create +// whose response it lost left a workspace behind: either the create landed +// first and this call releases its claim, or it arrives later and is +// refused, registering nothing. +// +// Idempotent. Lock order is the canonical b.mu -> ws.clientsMu; teardowns +// for workspaces the client was last on run after b.mu is released and +// re-check under both locks. +func (b *Backend) RetireClient(clientID string) error { + if _, err := validateClientID(clientID); err != nil { + return err + } + + b.mu.Lock() + if b.retired == nil { + b.retired = make(map[string]struct{}) + } + b.retired[clientID] = struct{}{} + var orphaned []*Workspace + for _, ws := range b.workspaces.Seq2() { + ws.clientsMu.Lock() + if cs, ok := ws.clients[clientID]; ok { + if cs.holdTimer != nil { + cs.holdTimer.Stop() + } + delete(ws.clients, clientID) + if len(ws.clients) == 0 { + orphaned = append(orphaned, ws) + } + } + ws.clientsMu.Unlock() + } + b.mu.Unlock() + + for _, ws := range orphaned { + b.teardown(ws) + } + return nil +} + +// releaseHold releases the creation hold for a client, if any. Active +// stream claims are unaffected. Idempotent: returns nil if the +// workspace or the client's hold no longer exist. +func (b *Backend) releaseHold(workspaceID, clientID string) error { + if _, err := validateClientID(clientID); err != nil { + return err + } + ws, ok := b.workspaces.Get(workspaceID) + if !ok { + return nil + } + b.releaseHoldLocked(ws, clientID) + return nil +} + +// registerClient installs (idempotently) the given client's claim on +// the workspace and starts a grace timer if the entry is fresh. +// +// A duplicate create from a client whose claim is timer-held (waiting to +// attach, or waiting out the detach grace) re-arms it for a full create +// grace and supersedes any earlier explicit release: the client is coming +// back. The re-arm installs a fresh clientState rather than resetting the +// old timer, so an already-fired timer racing this call fails expireHold's +// identity check instead of killing the new claim. +func (b *Backend) registerClient(ws *Workspace, clientID string) { + ws.clientsMu.Lock() + defer ws.clientsMu.Unlock() + if old, ok := ws.clients[clientID]; ok { + old.released = false + if old.holdTimer == nil { + // Live streams hold the claim; nothing to re-arm. + return + } + old.holdTimer.Stop() + ws.clients[clientID] = b.newHeldClient(ws, clientID, old.currentSessionID, b.createGrace) + return + } + ws.clients[clientID] = b.newHeldClient(ws, clientID, "", b.createGrace) +} + +// newHeldClient builds a clientState whose claim is held only by a timer +// that releases it after grace. Callers must hold ws.clientsMu. +func (b *Backend) newHeldClient(ws *Workspace, clientID, sessionID string, grace time.Duration) *clientState { + cs := &clientState{currentSessionID: sessionID} + cs.holdTimer = time.AfterFunc(grace, func() { + b.expireHold(ws, clientID, cs) + }) + return cs +} + +// expireHold is the body of the grace timer. It runs in its own +// goroutine and races against AttachClient/releaseHold; the timer +// stays valid only while the entry's holdTimer still points at it. +func (b *Backend) expireHold(ws *Workspace, clientID string, timer *clientState) { + ws.clientsMu.Lock() + cs, ok := ws.clients[clientID] + if !ok || cs != timer || cs.holdTimer == nil || cs.streams > 0 { + ws.clientsMu.Unlock() + return + } + cs.holdTimer = nil + delete(ws.clients, clientID) + teardown := len(ws.clients) == 0 + ws.clientsMu.Unlock() + if teardown { + b.teardown(ws) + } +} + +func (b *Backend) releaseHoldLocked(ws *Workspace, clientID string) { + ws.clientsMu.Lock() + cs, ok := ws.clients[clientID] + if !ok { + ws.clientsMu.Unlock() + return + } + if cs.holdTimer != nil { + cs.holdTimer.Stop() + cs.holdTimer = nil + } + teardown := false + if cs.streams == 0 { + delete(ws.clients, clientID) + teardown = len(ws.clients) == 0 + } else { + // The client gave up its claim while streams are still open, which + // is what a clean exit looks like. Remember it so the final detach + // skips the reconnect grace. + cs.released = true + } + ws.clientsMu.Unlock() + if teardown { + b.teardown(ws) + } +} + +func (b *Backend) detachStream(ws *Workspace, clientID string) { + b.mu.Lock() + grace := b.detachGrace + b.mu.Unlock() + + ws.clientsMu.Lock() + cs, ok := ws.clients[clientID] + if !ok { + ws.clientsMu.Unlock() + return + } + if cs.streams > 0 { + cs.streams-- + } + teardown := false + if cs.streams == 0 && cs.holdTimer == nil { + if grace > 0 && !cs.released { + // The stream dropped without the client releasing its claim, so + // treat it as an interruption rather than an exit: hold the + // workspace under a timer long enough for the client's + // reconnect to re-attach (AttachClient stops the timer). + cs.holdTimer = time.AfterFunc(grace, func() { + b.expireHold(ws, clientID, cs) + }) + } else { + delete(ws.clients, clientID) + teardown = len(ws.clients) == 0 + } + } + ws.clientsMu.Unlock() + if teardown { + b.teardown(ws) + } +} + +// teardown removes the workspace from the index, shuts down its +// underlying [app.App], and triggers a server shutdown if it was the +// last workspace alive. +// +// Callers reach teardown after observing len(ws.clients) == 0 while +// holding ws.clientsMu and then releasing it. Between that release +// and the b.mu.Lock below, a concurrent CreateWorkspace may have +// re-registered a client (CreateWorkspace holds b.mu while doing so, +// so it is mutually exclusive with this critical section). teardown +// re-checks under both locks (in the canonical b.mu -> ws.clientsMu +// order) and aborts if the workspace has been re-claimed. +func (b *Backend) teardown(ws *Workspace) { + b.mu.Lock() + ws.clientsMu.Lock() + if len(ws.clients) > 0 { + // Race: a CreateWorkspace re-registered a client + // between the detach path dropping ws.clientsMu and us + // taking b.mu. Abort: the workspace is still alive. + ws.clientsMu.Unlock() + b.mu.Unlock() + return + } + ws.clientsMu.Unlock() + if existing, ok := b.pathIndex[ws.resolvedPath]; ok && existing == ws.ID { + delete(b.pathIndex, ws.resolvedPath) + } + b.workspaces.Del(ws.ID) + // Arm (or, with lingering disabled, request) the idle shutdown. It + // only proceeds once there is genuinely nothing left: no live + // workspaces AND no create in flight. Deferring via the linger lets a + // client returning moments later reuse this server instead of racing + // its shutdown. + shutdownNow := b.scheduleShutdownIfIdleLocked() + b.mu.Unlock() + + ws.invokeShutdown() + + if shutdownNow { + slog.Info("Last workspace removed, shutting down server...") + b.shutdownFn() + } +} + +// scheduleShutdownIfIdleLocked decides what to do about server shutdown +// after a workspace count or pending-create change. It must be called +// with b.mu held. +// +// It returns true only when the caller should shut the server down +// synchronously (after releasing b.mu) — that is, when the server is idle +// and lingering is disabled (lingerDelay <= 0). When lingering is enabled +// it instead arms a one-shot timer that re-checks idleness after +// lingerDelay and shuts down then, and returns false. When the server is +// not idle (a workspace is live or a create is in flight) it does +// nothing and returns false. +// +// Returning true also latches [Backend.closing], since the caller has to +// release b.mu before it can run shutdownFn. +func (b *Backend) scheduleShutdownIfIdleLocked() (shutdownNow bool) { + if b.shutdownFn == nil { + return false + } + if b.workspaces.Len() != 0 || b.pending != 0 { + return false + } + if b.lingerDelay <= 0 { + b.closing = true + return true + } + if b.shutdownTimer == nil { + b.shutdownTimer = time.AfterFunc(b.lingerDelay, b.maybeShutdown) + } + return false +} + +// cancelShutdownLocked stops any armed idle-shutdown timer. It must be +// called with b.mu held. +func (b *Backend) cancelShutdownLocked() { + if b.shutdownTimer != nil { + b.shutdownTimer.Stop() + b.shutdownTimer = nil + } +} + +// maybeShutdown is the idle-shutdown timer callback. It shuts the server +// down only if it is still idle when the linger window elapses; any +// create that arrived in the meantime cancelled the timer (or bumped +// pending / the workspace count), so this re-check makes the linger +// race-free. Deciding to exit latches [Backend.closing] under the same +// lock, so a create arriving in the gap before shutdownFn runs is +// refused instead of being initialized on a departing process. +func (b *Backend) maybeShutdown() { + b.mu.Lock() + b.shutdownTimer = nil + idle := b.workspaces.Len() == 0 && b.pending == 0 + if idle { + b.closing = true + } + fn := b.shutdownFn + b.mu.Unlock() + if idle && fn != nil { + slog.Info("Server idle, shutting down...") + fn() + } +} + +// DeleteWorkspace is the public entry point used by the HTTP DELETE +// handler. It releases the named client's creation hold; live streams +// from the same client remain attached and continue holding the +// workspace open until their own deferred DetachClient runs. +func (b *Backend) DeleteWorkspace(id, clientID string) error { + return b.releaseHold(id, clientID) +} + +// SetCurrentSession records which session the given client is +// currently viewing within the workspace. Passing an empty sessionID +// clears the client's current-session entry (e.g. the client has +// returned to the landing screen). +// +// The client must be actually attached — i.e. its [clientState] entry +// must exist and have at least one live stream. A bare creation hold +// (streams == 0) is rejected with [ErrClientNotAttached]. This +// guards against zombie writes from a client that has detached and +// against ghost presence from a hold-only client that never opened an +// SSE stream. +func (b *Backend) SetCurrentSession(workspaceID, clientID, sessionID string) error { + if _, err := validateClientID(clientID); err != nil { + return err + } + ws, ok := b.workspaces.Get(workspaceID) + if !ok { + return ErrWorkspaceNotFound + } + ws.clientsMu.Lock() + defer ws.clientsMu.Unlock() + cs, ok := ws.clients[clientID] + if !ok || cs.streams == 0 { + // No entry, or hold-only (no live stream): refuse the + // write. The presence record this is meant to feed + // should only reflect clients that can actually observe + // session events. + return ErrClientNotAttached + } + cs.currentSessionID = sessionID + return nil +} + +// AttachedClients returns the number of clients currently viewing +// sessionID in the given workspace. Only clients with at least one live +// SSE stream (streams > 0) AND a matching currentSessionID are counted; +// pure creation holds do not contribute. Returns [ErrWorkspaceNotFound] +// if the workspace is unknown. +func (b *Backend) AttachedClients(workspaceID, sessionID string) (int, error) { + ws, ok := b.workspaces.Get(workspaceID) + if !ok { + return 0, ErrWorkspaceNotFound + } + return ws.AttachedClientsForSession(sessionID), nil +} + +// AttachedClientsForSession returns the number of clients in this +// workspace whose currentSessionID equals sessionID and which have at +// least one live SSE stream. Hold-only clients (streams == 0) do not +// contribute. Acquires the workspace's [clientsMu] briefly; the +// returned count is a point-in-time snapshot. +func (w *Workspace) AttachedClientsForSession(sessionID string) int { + w.clientsMu.Lock() + defer w.clientsMu.Unlock() + n := 0 + for _, cs := range w.clients { + if cs.streams > 0 && cs.currentSessionID == sessionID { + n++ + } + } + return n +} + +// GetWorkspaceProto returns the proto representation of a workspace. +func (b *Backend) GetWorkspaceProto(id string) (proto.Workspace, error) { + ws, err := b.GetWorkspace(id) + if err != nil { + return proto.Workspace{}, err + } + return workspaceToProto(ws), nil +} + +// VersionInfo returns server version information. +func (b *Backend) VersionInfo() proto.VersionInfo { + return proto.VersionInfo{ + Version: version.Version, + Commit: version.Commit, + BuildID: version.BuildID, + GoVersion: runtime.Version(), + Platform: fmt.Sprintf("%s/%s", runtime.GOOS, runtime.GOARCH), + } +} + +// Config returns the server-level configuration. +func (b *Backend) Config() *config.ConfigStore { + return b.cfg +} + +// Shutdown initiates a graceful server shutdown. +func (b *Backend) Shutdown() { + b.mu.Lock() + b.closing = true + fn := b.shutdownFn + b.mu.Unlock() + if fn != nil { + fn() + } +} + +// ShutdownIfIdle shuts the server down only when it is hosting no +// workspaces and has no creates in flight, reporting false when it declined +// because work is live. +// +// This is the only shutdown a client may request. A client asks in order to +// replace a version-mismatched server, and it cannot check idleness itself +// without a second round trip a new session can slip into. Deciding here, +// under the lock creates and teardowns take, closes that window: the answer +// is atomic, and granting it latches the decision so creates arriving +// afterwards are refused. +func (b *Backend) ShutdownIfIdle() bool { + b.mu.Lock() + live, pending := b.workspaces.Len(), b.pending + idle := live == 0 && pending == 0 + if idle { + b.closing = true + } + fn := b.shutdownFn + b.mu.Unlock() + if !idle { + slog.Warn("Refusing shutdown request: server is not idle", + "workspaces", live, "pending_creates", pending) + return false + } + if fn != nil { + fn() + } + return true +} + +// resolveWorkspaceKey returns a stable canonical form of path suitable +// for use as a dedup key. It applies filepath.Abs, then attempts +// filepath.EvalSymlinks; because EvalSymlinks errors on non-existent +// paths, it falls back to the cleaned absolute path in that case. +func resolveWorkspaceKey(path string) (string, error) { + abs, err := filepath.Abs(path) + if err != nil { + return "", err + } + if resolved, err := filepath.EvalSymlinks(abs); err == nil { + return resolved, nil + } + return abs, nil +} + +// validateClientID returns the trimmed UUID string or an error if the +// input is empty or not a valid UUID. +func validateClientID(id string) (string, error) { + if id == "" { + return "", ErrInvalidClientID + } + if _, err := uuid.Parse(id); err != nil { + return "", fmt.Errorf("%w: %v", ErrInvalidClientID, err) + } + return id, nil +} + +func workspaceToProto(ws *Workspace) proto.Workspace { + cfg := ws.Cfg.Config() + out := proto.Workspace{ + ID: ws.ID, + Path: ws.Path, + YOLO: ws.Cfg.Overrides().SkipPermissionRequests, + Channels: ws.Cfg.Overrides().EnabledChannels, + DataDir: cfg.Options.DataDirectory, + Debug: cfg.Options.Debug, + Config: cfg, + Env: ws.Env, + Version: version.Version, + } + if ws.Skills != nil { + out.Skills = skillStatesToProto(ws.Skills.States()) + } + return out +} + +// logFirstWinsMismatch emits a debug line whenever a second +// CreateWorkspace at the same resolved path arrives with flags that +// differ from the originating workspace. The existing workspace wins; +// the incoming flags are silently ignored. +// +// The comparison is done against the incoming args as the caller sent +// them — including empty/zero values — rather than after defaulting. +// This means that, for example, a second caller who omits DataDir +// while the first set one will still log the mismatch. +func logFirstWinsMismatch(existing *Workspace, args proto.Workspace) { + existingCfg := existing.Cfg.Config() + existingYOLO := existing.Cfg.Overrides().SkipPermissionRequests + existingChannels := existing.Cfg.Overrides().EnabledChannels + if existingYOLO == args.YOLO && + existingCfg.Options.Debug == args.Debug && + existingCfg.Options.DataDirectory == args.DataDir && + stringSlicesEqual(existing.Env, args.Env) && + stringSlicesEqual(existingChannels, args.Channels) { + return + } + slog.Debug( + "Workspace flag mismatch on duplicate create; first wins", + "workspace_id", existing.ID, + "path", existing.Path, + "existing_yolo", existingYOLO, + "requested_yolo", args.YOLO, + "existing_debug", existingCfg.Options.Debug, + "requested_debug", args.Debug, + "existing_data_dir", existingCfg.Options.DataDirectory, + "requested_data_dir", args.DataDir, + "existing_env", existing.Env, + "requested_env", args.Env, + "existing_channels", existingChannels, + "requested_channels", args.Channels, + ) +} + +// stringSlicesEqual reports whether a and b contain the same strings +// in the same order. nil and empty are treated as equal. +func stringSlicesEqual(a, b []string) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} diff --git a/internal/backend/backend_skills_test.go b/internal/backend/backend_skills_test.go new file mode 100644 index 0000000000000000000000000000000000000000..4f7dfa8f30e45976b9edba771576911a3205812a --- /dev/null +++ b/internal/backend/backend_skills_test.go @@ -0,0 +1,167 @@ +package backend_test + +import ( + "context" + "fmt" + "os" + "path/filepath" + "testing" + "time" + + tea "charm.land/bubbletea/v2" + "github.com/charmbracelet/crush/internal/backend" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/charmbracelet/crush/internal/skills" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +// TestBackend_WorkspaceSkillsIsolation verifies that skill discovery +// state and SSE events are per-workspace, not process-global. Two +// workspaces in the same backend process must not see each other's +// discoveries (either in their initial snapshot or in subsequent +// PublishStates events). +func TestBackend_WorkspaceSkillsIsolation(t *testing.T) { + // Isolate all of config.Init's filesystem reads from the host. The + // project-local .agents/skills//SKILL.md per working dir is + // what we actually want each workspace to see; everything else + // (global skills, XDG dirs, etc.) must be empty/deterministic. + hostHome := t.TempDir() + t.Setenv("HOME", hostHome) + t.Setenv("XDG_CONFIG_HOME", filepath.Join(hostHome, ".config")) + t.Setenv("XDG_DATA_HOME", filepath.Join(hostHome, ".local", "share")) + t.Setenv("XDG_CACHE_HOME", filepath.Join(hostHome, ".cache")) + t.Setenv("CRUSH_SKILLS_DIR", t.TempDir()) + + // Each workspace gets its own working directory containing a + // distinct project-local skill so the discovery output differs. + wdA := t.TempDir() + wdB := t.TempDir() + writeSkill(t, wdA, "wsa-only-skill", "Workspace A only skill.") + writeSkill(t, wdB, "wsb-only-skill", "Workspace B only skill.") + + srvCfg, err := config.Init(wdA, "", false) + require.NoError(t, err) + b := backend.New(t.Context(), srvCfg, nil) + + cidA := uuid.New().String() + cidB := uuid.New().String() + + wsA, _, err := b.CreateWorkspace(proto.Workspace{ + ClientID: cidA, + Path: wdA, + DataDir: filepath.Join(wdA, ".crush"), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = b.DeleteWorkspace(wsA.ID, cidA) }) + + wsB, _, err := b.CreateWorkspace(proto.Workspace{ + ClientID: cidB, + Path: wdB, + DataDir: filepath.Join(wdB, ".crush"), + }) + require.NoError(t, err) + t.Cleanup(func() { _ = b.DeleteWorkspace(wsB.ID, cidB) }) + + require.NotNil(t, wsA.Skills, "workspace A must have its own skills.Manager") + require.NotNil(t, wsB.Skills, "workspace B must have its own skills.Manager") + require.NotSame(t, wsA.Skills, wsB.Skills, "managers must be distinct instances per workspace") + + // Initial snapshots see each workspace's own filesystem skill, and + // neither sees the other's. + statesA := wsA.Skills.States() + statesB := wsB.Skills.States() + require.True(t, containsSkillName(statesA, "wsa-only-skill"), + "workspace A snapshot missing its own skill") + require.False(t, containsSkillName(statesA, "wsb-only-skill"), + "workspace A snapshot leaked workspace B's skill") + require.True(t, containsSkillName(statesB, "wsb-only-skill"), + "workspace B snapshot missing its own skill") + require.False(t, containsSkillName(statesB, "wsa-only-skill"), + "workspace B snapshot leaked workspace A's skill") + + // Subscribe to each workspace's SSE event stream. + ctxA, cancelA := context.WithCancel(t.Context()) + t.Cleanup(cancelA) + chA, err := b.SubscribeEvents(ctxA, wsA.ID) + require.NoError(t, err) + + ctxB, cancelB := context.WithCancel(t.Context()) + t.Cleanup(cancelB) + chB, err := b.SubscribeEvents(ctxB, wsB.ID) + require.NoError(t, err) + + // Trigger a republish on workspace A only. The marker name lets us + // distinguish this event from any incidental backend activity. + const marker = "wsa-republish-marker" + wsA.Skills.PublishStates([]*skills.SkillState{ + {Name: marker, State: skills.StateNormal}, + }) + + // Workspace A must receive its own event. + require.True(t, + waitForSkillsEvent(t, chA, marker, 2*time.Second), + "workspace A never received its own skills event") + + // Workspace B must NOT receive workspace A's event. + require.False(t, + waitForSkillsEvent(t, chB, marker, 250*time.Millisecond), + "workspace B leaked workspace A's skills event") + + // And A's published states are now visible on its manager's + // snapshot (verifies PublishStates updates the cache, not just the + // broker). + updatedA := wsA.Skills.States() + require.True(t, containsSkillName(updatedA, marker), + "PublishStates must update Manager.States()") + + // B's manager snapshot is untouched. + require.False(t, containsSkillName(wsB.Skills.States(), marker), + "workspace B's Manager.States() leaked workspace A's republish") +} + +func writeSkill(t *testing.T, workingDir, name, desc string) { + t.Helper() + skillDir := filepath.Join(workingDir, ".agents", "skills", name) + require.NoError(t, os.MkdirAll(skillDir, 0o755)) + content := fmt.Sprintf("---\nname: %s\ndescription: %s\n---\n%s\n", name, desc, desc) + require.NoError(t, os.WriteFile(filepath.Join(skillDir, "SKILL.md"), []byte(content), 0o644)) +} + +func containsSkillName(states []*skills.SkillState, name string) bool { + for _, s := range states { + if s.Name == name { + return true + } + } + return false +} + +// waitForSkillsEvent drains the given event channel until either a +// pubsub.Event[skills.Event] containing a state named marker arrives or +// the timeout elapses. Non-skills events on the channel are silently +// skipped — the backend fans in many event types and we only care +// about skills here. +func waitForSkillsEvent(t *testing.T, ch <-chan pubsub.Event[tea.Msg], marker string, timeout time.Duration) bool { + t.Helper() + deadline := time.After(timeout) + for { + select { + case ev, ok := <-ch: + if !ok { + return false + } + se, ok := ev.Payload.(pubsub.Event[skills.Event]) + if !ok { + continue + } + if containsSkillName(se.Payload.States, marker) { + return true + } + case <-deadline: + return false + } + } +} diff --git a/internal/backend/backend_test.go b/internal/backend/backend_test.go new file mode 100644 index 0000000000000000000000000000000000000000..e89ba196c143b0ac6e4ebabf7d02bfa186a9f637 --- /dev/null +++ b/internal/backend/backend_test.go @@ -0,0 +1,1784 @@ +package backend + +import ( + "bytes" + "context" + "errors" + "log/slog" + "os" + "path/filepath" + "runtime" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/charmbracelet/crush/internal/csync" + "github.com/charmbracelet/crush/internal/proto" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +// newTestBackend returns a Backend whose teardown path skips any +// real [app.App] shutdown work. Useful for state-machine tests that +// install synthetic workspaces directly via insertTestWorkspace. +func newTestBackend(t *testing.T) (*Backend, *atomic.Int32) { + t.Helper() + var shutdownCount atomic.Int32 + b := &Backend{ + workspaces: csync.NewMap[string, *Workspace](), + pathIndex: make(map[string]string), + ctx: context.Background(), + createGrace: 50 * time.Millisecond, + shutdownFn: func() { shutdownCount.Add(1) }, + } + return b, &shutdownCount +} + +// insertTestWorkspace installs a synthetic workspace into b at the +// given resolved path. Its shutdownFn is recorded in the returned +// counter so tests can assert it ran exactly once. +func insertTestWorkspace(t *testing.T, b *Backend, key string) (*Workspace, *atomic.Int32) { + t.Helper() + var shutdowns atomic.Int32 + ws := &Workspace{ + ID: uuid.New().String(), + Path: key, + resolvedPath: key, + clients: make(map[string]*clientState), + shutdownFn: func() { shutdowns.Add(1) }, + } + b.mu.Lock() + b.workspaces.Set(ws.ID, ws) + b.pathIndex[key] = ws.ID + b.mu.Unlock() + return ws, &shutdowns +} + +func newClientID(t *testing.T) string { + t.Helper() + return uuid.New().String() +} + +func TestResolveWorkspaceKey_AbsoluteAndSymlink(t *testing.T) { + t.Parallel() + + tmp := t.TempDir() + real, err := filepath.EvalSymlinks(tmp) + require.NoError(t, err) + + got, err := resolveWorkspaceKey(tmp) + require.NoError(t, err) + require.Equal(t, real, got) +} + +func TestResolveWorkspaceKey_NonExistentFallback(t *testing.T) { + t.Parallel() + + missing := filepath.Join(t.TempDir(), "does", "not", "exist") + got, err := resolveWorkspaceKey(missing) + require.NoError(t, err) + abs, err := filepath.Abs(missing) + require.NoError(t, err) + require.Equal(t, abs, got) +} + +func TestValidateClientID(t *testing.T) { + t.Parallel() + + _, err := validateClientID("") + require.ErrorIs(t, err, ErrInvalidClientID) + _, err = validateClientID("not-a-uuid") + require.ErrorIs(t, err, ErrInvalidClientID) + + id := uuid.New().String() + got, err := validateClientID(id) + require.NoError(t, err) + require.Equal(t, id, got) +} + +func TestRegisterClient_Idempotent(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, _ := insertTestWorkspace(t, b, "/tmp/a") + + cid := newClientID(t) + b.registerClient(ws, cid) + b.registerClient(ws, cid) + + ws.clientsMu.Lock() + defer ws.clientsMu.Unlock() + require.Len(t, ws.clients, 1) + require.NotNil(t, ws.clients[cid].holdTimer) + require.Equal(t, 0, ws.clients[cid].streams) +} + +func TestAttachClient_ConsumesHold(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, shutdowns := insertTestWorkspace(t, b, "/tmp/a") + + cid := newClientID(t) + b.registerClient(ws, cid) + require.NoError(t, b.AttachClient(ws.ID, cid)) + + ws.clientsMu.Lock() + require.Len(t, ws.clients, 1) + require.Nil(t, ws.clients[cid].holdTimer, "attach must stop the grace timer") + require.Equal(t, 1, ws.clients[cid].streams) + ws.clientsMu.Unlock() + + // Wait past the grace window: a stopped timer must not fire. + time.Sleep(150 * time.Millisecond) + require.Equal(t, int32(0), shutdowns.Load(), "workspace must not be torn down while attached") +} + +func TestAttachClient_WithoutPriorCreate(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, _ := insertTestWorkspace(t, b, "/tmp/a") + + cid := newClientID(t) + require.NoError(t, b.AttachClient(ws.ID, cid)) + + ws.clientsMu.Lock() + defer ws.clientsMu.Unlock() + require.Len(t, ws.clients, 1) + require.Equal(t, 1, ws.clients[cid].streams) + require.Nil(t, ws.clients[cid].holdTimer) +} + +func TestAttachClient_DuplicateStreams(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, shutdowns := insertTestWorkspace(t, b, "/tmp/a") + + cid := newClientID(t) + require.NoError(t, b.AttachClient(ws.ID, cid)) + require.NoError(t, b.AttachClient(ws.ID, cid)) + + ws.clientsMu.Lock() + require.Equal(t, 2, ws.clients[cid].streams) + ws.clientsMu.Unlock() + + b.DetachClient(ws.ID, cid) + ws.clientsMu.Lock() + require.Equal(t, 1, ws.clients[cid].streams) + ws.clientsMu.Unlock() + require.Equal(t, int32(0), shutdowns.Load()) + + b.DetachClient(ws.ID, cid) + require.Equal(t, int32(1), shutdowns.Load(), "second detach tears down the workspace") +} + +func TestDetachClient_LastStreamTearsDown(t *testing.T) { + t.Parallel() + + b, srvShutdowns := newTestBackend(t) + ws, wsShutdowns := insertTestWorkspace(t, b, "/tmp/a") + + cid := newClientID(t) + b.registerClient(ws, cid) + require.NoError(t, b.AttachClient(ws.ID, cid)) + b.DetachClient(ws.ID, cid) + + require.Equal(t, int32(1), wsShutdowns.Load()) + require.Equal(t, int32(1), srvShutdowns.Load(), "last workspace shut down must trigger server shutdown") + _, err := b.GetWorkspace(ws.ID) + require.ErrorIs(t, err, ErrWorkspaceNotFound) +} + +func TestHoldExpiry_TearsDown(t *testing.T) { + t.Parallel() + + b, srvShutdowns := newTestBackend(t) + ws, wsShutdowns := insertTestWorkspace(t, b, "/tmp/a") + + cid := newClientID(t) + b.registerClient(ws, cid) + + require.Eventually(t, func() bool { + return wsShutdowns.Load() == 1 && srvShutdowns.Load() == 1 + }, 1*time.Second, 5*time.Millisecond) +} + +func TestReleaseHold_NoStreams(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, shutdowns := insertTestWorkspace(t, b, "/tmp/a") + + cid := newClientID(t) + b.registerClient(ws, cid) + require.NoError(t, b.releaseHold(ws.ID, cid)) + + require.Equal(t, int32(1), shutdowns.Load()) + // Idempotent. + require.NoError(t, b.releaseHold(ws.ID, cid)) + require.Equal(t, int32(1), shutdowns.Load()) +} + +func TestReleaseHold_WithActiveStream(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, shutdowns := insertTestWorkspace(t, b, "/tmp/a") + + cid := newClientID(t) + b.registerClient(ws, cid) + require.NoError(t, b.AttachClient(ws.ID, cid)) + require.NoError(t, b.releaseHold(ws.ID, cid)) + + ws.clientsMu.Lock() + require.Equal(t, 1, ws.clients[cid].streams) + require.Nil(t, ws.clients[cid].holdTimer) + ws.clientsMu.Unlock() + require.Equal(t, int32(0), shutdowns.Load()) + + b.DetachClient(ws.ID, cid) + require.Equal(t, int32(1), shutdowns.Load()) +} + +func TestReleaseHoldThenAttach(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, shutdowns := insertTestWorkspace(t, b, "/tmp/a") + + cid := newClientID(t) + require.NoError(t, b.releaseHold(ws.ID, cid)) // no entry yet — no-op. + require.NoError(t, b.AttachClient(ws.ID, cid)) + ws.clientsMu.Lock() + require.Equal(t, 1, ws.clients[cid].streams) + ws.clientsMu.Unlock() + require.NoError(t, b.releaseHold(ws.ID, cid)) // hold-only no-op (no hold timer). + require.Equal(t, int32(0), shutdowns.Load()) + b.DetachClient(ws.ID, cid) + require.Equal(t, int32(1), shutdowns.Load()) +} + +func TestRefcountWithSecondClient(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, shutdowns := insertTestWorkspace(t, b, "/tmp/a") + + cidA := newClientID(t) + cidB := newClientID(t) + b.registerClient(ws, cidA) + require.NoError(t, b.AttachClient(ws.ID, cidA)) + b.registerClient(ws, cidB) + require.NoError(t, b.AttachClient(ws.ID, cidB)) + + b.DetachClient(ws.ID, cidA) + ws.clientsMu.Lock() + require.Contains(t, ws.clients, cidB) + require.NotContains(t, ws.clients, cidA) + ws.clientsMu.Unlock() + require.Equal(t, int32(0), shutdowns.Load(), "workspace survives while second client attached") + + b.DetachClient(ws.ID, cidB) + require.Equal(t, int32(1), shutdowns.Load()) +} + +func TestAttachClient_InvalidID(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, _ := insertTestWorkspace(t, b, "/tmp/a") + + require.ErrorIs(t, b.AttachClient(ws.ID, ""), ErrInvalidClientID) + require.ErrorIs(t, b.AttachClient(ws.ID, "not-a-uuid"), ErrInvalidClientID) +} + +func TestDeleteWorkspace_RejectsBadClientID(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, _ := insertTestWorkspace(t, b, "/tmp/a") + + require.ErrorIs(t, b.DeleteWorkspace(ws.ID, ""), ErrInvalidClientID) + require.ErrorIs(t, b.DeleteWorkspace(ws.ID, "not-a-uuid"), ErrInvalidClientID) +} + +// TestHoldExpiry_RaceWithAttach checks that, when the grace timer fires +// while a concurrent AttachClient call is in flight, the workspace ends +// up either fully attached or fully torn down — never in a half-state. +func TestHoldExpiry_RaceWithAttach(t *testing.T) { + t.Parallel() + + for i := range 50 { + b, _ := newTestBackend(t) + // Tighten the grace window further to force the race. + b.createGrace = 1 * time.Millisecond + ws, shutdowns := insertTestWorkspace(t, b, "/tmp/race") + + cid := newClientID(t) + b.registerClient(ws, cid) + // Attach concurrently with the very short grace timer. + errCh := make(chan error, 1) + go func() { errCh <- b.AttachClient(ws.ID, cid) }() + <-errCh + + // Wait for any pending timer to settle. + time.Sleep(10 * time.Millisecond) + + ws.clientsMu.Lock() + gotShutdown := shutdowns.Load() == 1 + cs, present := ws.clients[cid] + var ( + gotStreams int + gotHoldTimer *time.Timer + ) + if present { + gotStreams = cs.streams + gotHoldTimer = cs.holdTimer + } + ws.clientsMu.Unlock() + // Either the workspace was torn down OR the client is + // attached with streams==1 and the hold timer cleared. + // The state must be consistent: if shutdown, client is + // gone; if attached, no teardown and streams==1. + if gotShutdown { + require.False(t, present, "iter %d: shutdown but client still present", i) + } else { + require.True(t, present, "iter %d: not shutdown but client missing", i) + require.Equal(t, 1, gotStreams, "iter %d: attach winner must leave streams=1", i) + require.Nil(t, gotHoldTimer, "iter %d: attach winner must clear holdTimer", i) + } + } +} + +// TestConcurrentAttachDetach exercises the state machine under +// parallel attach/detach pressure with the race detector. +func TestConcurrentAttachDetach(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, _ := insertTestWorkspace(t, b, "/tmp/a") + + cid := newClientID(t) + b.registerClient(ws, cid) + require.NoError(t, b.AttachClient(ws.ID, cid)) // ensure refcount stays > 0. + + const n = 50 + var wg sync.WaitGroup + wg.Add(n) + for range n { + go func() { + defer wg.Done() + cid2 := newClientID(t) + _ = b.AttachClient(ws.ID, cid2) + b.DetachClient(ws.ID, cid2) + }() + } + wg.Wait() + + ws.clientsMu.Lock() + defer ws.clientsMu.Unlock() + require.Len(t, ws.clients, 1) + require.Contains(t, ws.clients, cid) +} + +// TestPathDedupe_FullCreate exercises CreateWorkspace end-to-end +// (config init, real app.App). Two CreateWorkspace calls at the same +// path return the same workspace ID and share the clients map. +func TestPathDedupe_FullCreate(t *testing.T) { + 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()) + + cwd := t.TempDir() + dataDir := t.TempDir() + + b := New(context.Background(), nil, func() {}) + b.SetCreateGrace(2 * time.Second) + t.Cleanup(func() { drainBackend(t, b) }) + + cidA := uuid.New().String() + cidB := uuid.New().String() + + wsA, protoA, err := b.CreateWorkspace(protoWS(cwd, dataDir, cidA)) + require.NoError(t, err) + require.NotEmpty(t, protoA.ID) + require.Equal(t, protoA.DataDir, wsA.Cfg.Config().Options.DataDirectory) + + wsB, protoB, err := b.CreateWorkspace(protoWS(cwd, dataDir, cidB)) + require.NoError(t, err) + require.Equal(t, wsA.ID, wsB.ID, "second create at same path must return existing workspace") + require.Equal(t, protoA.ID, protoB.ID) + + wsA.clientsMu.Lock() + require.Contains(t, wsA.clients, cidA) + require.Contains(t, wsA.clients, cidB) + wsA.clientsMu.Unlock() +} + +// TestPathDedupe_DifferentPaths_DifferentWorkspaces confirms that two +// CreateWorkspace calls at distinct paths produce distinct workspaces. +func TestPathDedupe_DifferentPaths_DifferentWorkspaces(t *testing.T) { + 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()) + + cwdA := t.TempDir() + cwdB := t.TempDir() + dataA := t.TempDir() + dataB := t.TempDir() + + b := New(context.Background(), nil, func() {}) + b.SetCreateGrace(2 * time.Second) + t.Cleanup(func() { drainBackend(t, b) }) + + wsA, _, err := b.CreateWorkspace(protoWS(cwdA, dataA, uuid.New().String())) + require.NoError(t, err) + wsB, _, err := b.CreateWorkspace(protoWS(cwdB, dataB, uuid.New().String())) + require.NoError(t, err) + require.NotEqual(t, wsA.ID, wsB.ID) +} + +// TestPathDedupe_FirstWinsKeepsOriginalEnv verifies that the second +// create at the same path returns the *originating* client's Env in +// its proto and does not mutate the existing workspace's YOLO/Debug +// flags. +func TestPathDedupe_FirstWinsKeepsOriginalEnv(t *testing.T) { + 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()) + + cwd := t.TempDir() + dataDir := t.TempDir() + + b := New(context.Background(), nil, func() {}) + b.SetCreateGrace(2 * time.Second) + t.Cleanup(func() { drainBackend(t, b) }) + + originalEnv := []string{"FOO=bar"} + argsA := protoWS(cwd, dataDir, uuid.New().String()) + argsA.YOLO = true + argsA.Env = originalEnv + wsA, protoA, err := b.CreateWorkspace(argsA) + require.NoError(t, err) + require.True(t, protoA.YOLO) + require.Equal(t, originalEnv, protoA.Env) + + argsB := protoWS(cwd, dataDir, uuid.New().String()) + argsB.YOLO = false + argsB.Debug = true + argsB.Env = []string{"BAZ=qux"} + _, protoB, err := b.CreateWorkspace(argsB) + require.NoError(t, err) + require.Equal(t, protoA.ID, protoB.ID) + require.True(t, protoB.YOLO, "first wins: YOLO must remain true") + require.Equal(t, originalEnv, protoB.Env, "proto must carry the originating client's Env") + require.Equal(t, wsA.Cfg.Overrides().SkipPermissionRequests, true) +} + +// TestPathDedupe_Symlink confirms two paths that resolve to the same +// target share a workspace. +func TestPathDedupe_Symlink(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("symlink semantics differ on windows") + } + 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()) + + real := t.TempDir() + link := filepath.Join(t.TempDir(), "link") + require.NoError(t, os.Symlink(real, link)) + dataDir := t.TempDir() + + b := New(context.Background(), nil, func() {}) + b.SetCreateGrace(2 * time.Second) + t.Cleanup(func() { drainBackend(t, b) }) + + wsA, _, err := b.CreateWorkspace(protoWS(real, dataDir, uuid.New().String())) + require.NoError(t, err) + wsB, _, err := b.CreateWorkspace(protoWS(link, dataDir, uuid.New().String())) + require.NoError(t, err) + require.Equal(t, wsA.ID, wsB.ID) +} + +// TestPathDedupe_NonExistentPath ensures CreateWorkspace tolerates a +// path that does not yet exist (EvalSymlinks falls back to Abs). +func TestPathDedupe_NonExistentPath(t *testing.T) { + 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()) + + parent := t.TempDir() + missing := filepath.Join(parent, "does-not-exist") + dataDir := t.TempDir() + + b := New(context.Background(), nil, func() {}) + b.SetCreateGrace(2 * time.Second) + t.Cleanup(func() { drainBackend(t, b) }) + + _, p, err := b.CreateWorkspace(protoWS(missing, dataDir, uuid.New().String())) + require.NoError(t, err) + require.NotEmpty(t, p.ID) +} + +// TestCreateWorkspace_IdempotentSameClient checks that a duplicate +// create from the same client at the same path does not produce a +// second claim. +func TestCreateWorkspace_IdempotentSameClient(t *testing.T) { + 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()) + + cwd := t.TempDir() + dataDir := t.TempDir() + b := New(context.Background(), nil, func() {}) + b.SetCreateGrace(2 * time.Second) + t.Cleanup(func() { drainBackend(t, b) }) + + cid := uuid.New().String() + ws1, _, err := b.CreateWorkspace(protoWS(cwd, dataDir, cid)) + require.NoError(t, err) + ws2, _, err := b.CreateWorkspace(protoWS(cwd, dataDir, cid)) + require.NoError(t, err) + require.Equal(t, ws1.ID, ws2.ID) + + ws1.clientsMu.Lock() + require.Len(t, ws1.clients, 1, "duplicate create from same client must not double the claim") + ws1.clientsMu.Unlock() +} + +// TestPathDedupe_ParallelCreates ensures two simultaneous CreateWorkspace +// calls at the same path produce the same workspace and the clients map +// contains both client IDs. +func TestPathDedupe_ParallelCreates(t *testing.T) { + 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()) + + cwd := t.TempDir() + dataDir := t.TempDir() + + b := New(context.Background(), nil, func() {}) + b.SetCreateGrace(2 * time.Second) + t.Cleanup(func() { drainBackend(t, b) }) + + cidA := uuid.New().String() + cidB := uuid.New().String() + + type result struct { + ws *Workspace + proto proto.Workspace + err error + } + ch := make(chan result, 2) + start := make(chan struct{}) + go func() { + <-start + ws, p, err := b.CreateWorkspace(protoWS(cwd, dataDir, cidA)) + ch <- result{ws, p, err} + }() + go func() { + <-start + ws, p, err := b.CreateWorkspace(protoWS(cwd, dataDir, cidB)) + ch <- result{ws, p, err} + }() + close(start) + r1 := <-ch + r2 := <-ch + require.NoError(t, r1.err) + require.NoError(t, r2.err) + require.Equal(t, r1.ws.ID, r2.ws.ID, "both creates must converge on one workspace ID") + + ws := r1.ws + ws.clientsMu.Lock() + defer ws.clientsMu.Unlock() + require.Contains(t, ws.clients, cidA) + require.Contains(t, ws.clients, cidB) +} + +// TestCreateWorkspace_RejectsBadClientID covers the 400 path from the +// backend side. +func TestCreateWorkspace_RejectsBadClientID(t *testing.T) { + t.Parallel() + + b := New(context.Background(), nil, func() {}) + + _, _, err := b.CreateWorkspace(protoWS("/tmp/x", t.TempDir(), "")) + require.ErrorIs(t, err, ErrInvalidClientID) + _, _, err = b.CreateWorkspace(protoWS("/tmp/x", t.TempDir(), "not-a-uuid")) + require.ErrorIs(t, err, ErrInvalidClientID) +} + +// drainBackend tears the backend down at the end of a test by deleting +// every remaining workspace. Necessary so the test process doesn't +// leak goroutines or DB handles from the embedded [app.App] instances. +func drainBackend(t *testing.T, b *Backend) { + t.Helper() + for _, ws := range b.workspaces.Seq2() { + ws.clientsMu.Lock() + ids := make([]string, 0, len(ws.clients)) + for id := range ws.clients { + ids = append(ids, id) + } + ws.clientsMu.Unlock() + for _, cid := range ids { + _ = b.releaseHold(ws.ID, cid) + } + } +} + +func protoWS(path, dataDir, clientID string) proto.Workspace { + return proto.Workspace{Path: path, DataDir: dataDir, ClientID: clientID} +} + +// syncBuffer is a thread-safe buffer that can be safely read and written +// from multiple goroutines. +type syncBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +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) String() string { + sb.mu.Lock() + defer sb.mu.Unlock() + return sb.buf.String() +} + +// captureDebugLogs installs a buffer-backed slog handler at Debug +// level for the duration of the test, returning the buffer. The +// previous default handler is restored via t.Cleanup. +func captureDebugLogs(t *testing.T) *syncBuffer { + t.Helper() + var sb syncBuffer + prev := slog.Default() + handler := slog.NewTextHandler(&sb, &slog.HandlerOptions{Level: slog.LevelDebug}) + slog.SetDefault(slog.New(handler)) + t.Cleanup(func() { slog.SetDefault(prev) }) + return &sb +} + +// xdgIsolated points HOME and XDG_* variables at fresh tempdirs so +// CreateWorkspace's config loading does not interfere with the host +// machine's real config. +func xdgIsolated(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()) +} + +// TestFirstWinsMismatch_LogsOnFlagDifferences verifies that the +// debug mismatch line is emitted when any of YOLO, Debug, DataDir, +// or Env differs between the first and second CreateWorkspace at +// the same path, and that the existing workspace's Debug flag is +// not overwritten. +func TestFirstWinsMismatch_LogsOnFlagDifferences(t *testing.T) { + tests := []struct { + name string + mutate func(*proto.Workspace) + }{ + { + name: "yolo", + mutate: func(p *proto.Workspace) { p.YOLO = true }, + }, + { + name: "debug", + mutate: func(p *proto.Workspace) { p.Debug = true }, + }, + { + name: "datadir", + mutate: func(p *proto.Workspace) { p.DataDir = "" }, + }, + { + name: "env", + mutate: func(p *proto.Workspace) { p.Env = []string{"NEW=val"} }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + xdgIsolated(t) + cwd := t.TempDir() + dataDir := t.TempDir() + + buf := captureDebugLogs(t) + b := New(context.Background(), nil, func() {}) + b.SetCreateGrace(2 * time.Second) + t.Cleanup(func() { drainBackend(t, b) }) + + argsA := protoWS(cwd, dataDir, uuid.New().String()) + argsA.Env = []string{"FOO=bar"} + wsA, _, err := b.CreateWorkspace(argsA) + require.NoError(t, err) + originalDebug := wsA.Cfg.Config().Options.Debug + originalYOLO := wsA.Cfg.Overrides().SkipPermissionRequests + + argsB := protoWS(cwd, dataDir, uuid.New().String()) + argsB.Env = []string{"FOO=bar"} // identical by default + tc.mutate(&argsB) + _, _, err = b.CreateWorkspace(argsB) + require.NoError(t, err) + + require.Contains( + t, buf.String(), + "Workspace flag mismatch on duplicate create", + "expected debug log for mismatching %s", tc.name, + ) + // Existing workspace's YOLO and Debug must not change. + require.Equal(t, originalYOLO, wsA.Cfg.Overrides().SkipPermissionRequests, "YOLO must be immutable on first-wins") + require.Equal(t, originalDebug, wsA.Cfg.Config().Options.Debug, "Debug must be immutable on first-wins") + }) + } +} + +// TestFirstWinsMismatch_NoLogWhenIdentical confirms identical args +// do not emit the mismatch log line. +func TestFirstWinsMismatch_NoLogWhenIdentical(t *testing.T) { + xdgIsolated(t) + cwd := t.TempDir() + dataDir := t.TempDir() + + buf := captureDebugLogs(t) + b := New(context.Background(), nil, func() {}) + b.SetCreateGrace(2 * time.Second) + t.Cleanup(func() { drainBackend(t, b) }) + + argsA := protoWS(cwd, dataDir, uuid.New().String()) + argsA.Env = []string{"FOO=bar"} + _, _, err := b.CreateWorkspace(argsA) + require.NoError(t, err) + + argsB := protoWS(cwd, dataDir, uuid.New().String()) + argsB.Env = []string{"FOO=bar"} + _, _, err = b.CreateWorkspace(argsB) + require.NoError(t, err) + + require.False(t, + strings.Contains(buf.String(), "Workspace flag mismatch on duplicate create"), + "identical args must not log a mismatch: %s", buf.String()) +} + +// TestChannelOptInBoundary_DuplicateCreate verifies that channels are +// an explicit opt-in that is never shared across a duplicate create at +// the same path. A second client whose requested channels differ from +// the existing workspace (including a client that did not opt in at +// all) is rejected rather than silently inheriting the existing +// workspace's channels. +func TestChannelOptInBoundary_DuplicateCreate(t *testing.T) { + tests := []struct { + name string + firstChannels []string + secondChannels []string + wantMismatchErr bool + }{ + { + name: "second omits opt-in", + firstChannels: []string{"server:webhook"}, + secondChannels: nil, + wantMismatchErr: true, + }, + { + name: "second opts into extra channel", + firstChannels: nil, + secondChannels: []string{"server:webhook"}, + wantMismatchErr: true, + }, + { + name: "different channels", + firstChannels: []string{"server:webhook"}, + secondChannels: []string{"server:events"}, + wantMismatchErr: true, + }, + { + name: "identical channels shared", + firstChannels: []string{"server:webhook"}, + secondChannels: []string{"server:webhook"}, + wantMismatchErr: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + xdgIsolated(t) + cwd := t.TempDir() + dataDir := t.TempDir() + + b := New(context.Background(), nil, func() {}) + b.SetCreateGrace(2 * time.Second) + t.Cleanup(func() { drainBackend(t, b) }) + + argsA := protoWS(cwd, dataDir, uuid.New().String()) + argsA.Channels = tc.firstChannels + wsA, _, err := b.CreateWorkspace(argsA) + require.NoError(t, err) + + argsB := protoWS(cwd, dataDir, uuid.New().String()) + argsB.Channels = tc.secondChannels + wsB, protoB, err := b.CreateWorkspace(argsB) + + if tc.wantMismatchErr { + require.ErrorIs(t, err, ErrChannelOptInMismatch) + require.Nil(t, wsB) + // The existing workspace must not be mutated + // and must not register the rejected client. + require.Equal(t, tc.firstChannels, wsA.Cfg.Overrides().EnabledChannels, + "existing channels must be immutable on rejection") + wsA.clientsMu.Lock() + require.Len(t, wsA.clients, 1, "rejected client must not be registered") + wsA.clientsMu.Unlock() + return + } + + require.NoError(t, err) + require.Equal(t, wsA.ID, protoB.ID) + require.Equal(t, tc.firstChannels, protoB.Channels) + }) + } +} + +// TestRaceTwoClientsAttachOneDetaches exercises the PLAN-required +// race scenario: two clients attach concurrently, then one detaches. +// The workspace must remain alive with refcount==1 and the clients +// map must reflect the remaining client only. +func TestRaceTwoClientsAttachOneDetaches(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, shutdowns := insertTestWorkspace(t, b, "/tmp/race-two") + + cidA := newClientID(t) + cidB := newClientID(t) + + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + require.NoError(t, b.AttachClient(ws.ID, cidA)) + }() + go func() { + defer wg.Done() + require.NoError(t, b.AttachClient(ws.ID, cidB)) + }() + wg.Wait() + + ws.clientsMu.Lock() + require.Len(t, ws.clients, 2, "both clients must be attached") + ws.clientsMu.Unlock() + + b.DetachClient(ws.ID, cidA) + + ws.clientsMu.Lock() + require.Len(t, ws.clients, 1, "refcount must be 1 after one detach") + require.Contains(t, ws.clients, cidB, "remaining client must be cidB") + require.NotContains(t, ws.clients, cidA, "detached client must be removed") + ws.clientsMu.Unlock() + require.Equal(t, int32(0), shutdowns.Load(), "workspace must remain alive") + + // Drain. + b.DetachClient(ws.ID, cidB) + require.Equal(t, int32(1), shutdowns.Load()) +} + +// TestExplicitDeleteThenAttach reproduces the PLAN scenario: start +// with a real hold, releaseHold consumes it, AttachClient from the +// same clientID creates a fresh entry with streams==1, and calling +// releaseHold again is a no-op. A second client keeps the workspace +// alive so AttachClient can still resolve the workspace ID after the +// first client's hold is released. +func TestExplicitDeleteThenAttach(t *testing.T) { + t.Parallel() + + // Large grace window so timers cannot fire during the test + // — we want to exercise the explicit releaseHold path. + b, _ := newTestBackend(t) + b.createGrace = time.Hour + ws, shutdowns := insertTestWorkspace(t, b, "/tmp/delete-then-attach") + + // Anchor client keeps the workspace registered in + // b.workspaces across the cid's releaseHold below. + anchor := newClientID(t) + require.NoError(t, b.AttachClient(ws.ID, anchor)) + + cid := newClientID(t) + // Real hold via registerClient (mirrors CreateWorkspace). + b.registerClient(ws, cid) + ws.clientsMu.Lock() + require.Contains(t, ws.clients, cid) + require.NotNil(t, ws.clients[cid].holdTimer, "hold must be live") + require.Equal(t, 0, ws.clients[cid].streams) + ws.clientsMu.Unlock() + + // releaseHold: consumes the hold and removes the entry + // (streams == 0). The anchor client keeps the workspace + // alive. + require.NoError(t, b.releaseHold(ws.ID, cid)) + require.Equal(t, int32(0), shutdowns.Load(), "anchor must keep workspace alive") + ws.clientsMu.Lock() + require.NotContains(t, ws.clients, cid, "entry must be removed by releaseHold") + ws.clientsMu.Unlock() + + // AttachClient creates a fresh entry with streams==1 and no + // hold timer. + require.NoError(t, b.AttachClient(ws.ID, cid)) + ws.clientsMu.Lock() + require.Contains(t, ws.clients, cid, "fresh entry must be created") + require.Equal(t, 1, ws.clients[cid].streams, "fresh attach must start at streams=1") + require.Nil(t, ws.clients[cid].holdTimer, "fresh attach must have no hold timer") + ws.clientsMu.Unlock() + + // Calling releaseHold again is a no-op (no hold timer to + // stop, streams > 0 so the entry stays). + require.NoError(t, b.releaseHold(ws.ID, cid)) + ws.clientsMu.Lock() + require.Contains(t, ws.clients, cid, "releaseHold must not touch a stream-only entry") + require.Equal(t, 1, ws.clients[cid].streams) + require.Nil(t, ws.clients[cid].holdTimer) + ws.clientsMu.Unlock() + + // Drain. + b.DetachClient(ws.ID, cid) + b.DetachClient(ws.ID, anchor) + require.Equal(t, int32(1), shutdowns.Load()) +} + +// TestAttachClient_RacesWithTeardown forces AttachClient to compete +// with the teardown path triggered by DetachClient. Before the fix, +// AttachClient could observe a workspace after teardown had already +// decided to remove it (because AttachClient did not synchronize with +// Backend.mu), leaving a live stream claim attached to a workspace +// that was then removed and shut down. With the fix, the outcome must +// be deterministic: either AttachClient won and the workspace is +// alive with the client registered, or teardown won and AttachClient +// returns ErrWorkspaceNotFound — never a half-state where the +// workspace is gone but ws.clients still contains the new client. +func TestAttachClient_RacesWithTeardown(t *testing.T) { + t.Parallel() + + for i := range 200 { + b, _ := newTestBackend(t) + // Keep the grace window long so it can't fire during the + // test and confuse the bookkeeping. + b.createGrace = time.Hour + ws, shutdowns := insertTestWorkspace(t, b, "/tmp/race-teardown") + + // Seed: cidA holds the workspace open via a stream. The + // imminent DetachClient(cidA) will be the *only* claim + // drop, so teardown will run. + cidA := newClientID(t) + require.NoError(t, b.AttachClient(ws.ID, cidA)) + + // cidB attempts to attach concurrently with the detach + // that will tear the workspace down. + cidB := newClientID(t) + start := make(chan struct{}) + errCh := make(chan error, 1) + detachDone := make(chan struct{}) + go func() { + <-start + errCh <- b.AttachClient(ws.ID, cidB) + }() + go func() { + <-start + b.DetachClient(ws.ID, cidA) + close(detachDone) + }() + close(start) + + // Wait for both goroutines so teardown (including + // shutdownFn) has fully run before we read state. + attachErr := <-errCh + <-detachDone + + _, wsStillRegistered := b.workspaces.Get(ws.ID) + ws.clientsMu.Lock() + _, hasA := ws.clients[cidA] + _, hasB := ws.clients[cidB] + clientCount := len(ws.clients) + ws.clientsMu.Unlock() + shutdownCount := shutdowns.Load() + + switch { + case attachErr == nil: + // AttachClient won. The workspace must be alive + // (registered) with cidB in its clients map. cidA + // may or may not still be there depending on who + // took clientsMu first, but the workspace must + // not have been torn down. + require.True(t, wsStillRegistered, + "iter %d: attach succeeded but workspace was removed", i) + require.True(t, hasB, + "iter %d: attach succeeded but cidB missing from clients", i) + require.Equal(t, int32(0), shutdownCount, + "iter %d: attach succeeded but workspace was shut down", i) + case errors.Is(attachErr, ErrWorkspaceNotFound): + // Teardown won. The workspace must be removed, + // shut down exactly once, and ws.clients must be + // empty (no half-state with cidB inserted into a + // dead workspace's clients map). + require.False(t, wsStillRegistered, + "iter %d: ErrWorkspaceNotFound but workspace still registered", i) + require.Equal(t, int32(1), shutdownCount, + "iter %d: ErrWorkspaceNotFound but shutdown count = %d", i, shutdownCount) + require.False(t, hasA, + "iter %d: teardown won but cidA still in clients", i) + require.False(t, hasB, + "iter %d: teardown won but cidB still in clients (would be the leaked attach)", i) + require.Zero(t, clientCount, + "iter %d: teardown won but clients map is non-empty", i) + default: + t.Fatalf("iter %d: unexpected AttachClient error: %v", i, attachErr) + } + } +} + +// TestSetCurrentSession_BasicAttachAndSwitch verifies the happy path: +// an attached client can set its current session, a second attached +// client can target the same session, and one of them can switch to a +// different session without disturbing the other's record. +func TestSetCurrentSession_BasicAttachAndSwitch(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, _ := insertTestWorkspace(t, b, "/tmp/current-session-basic") + + cidA := newClientID(t) + cidB := newClientID(t) + require.NoError(t, b.AttachClient(ws.ID, cidA)) + require.NoError(t, b.AttachClient(ws.ID, cidB)) + + require.NoError(t, b.SetCurrentSession(ws.ID, cidA, "S1")) + ws.clientsMu.Lock() + require.Equal(t, "S1", ws.clients[cidA].currentSessionID) + ws.clientsMu.Unlock() + + require.NoError(t, b.SetCurrentSession(ws.ID, cidB, "S1")) + ws.clientsMu.Lock() + require.Equal(t, "S1", ws.clients[cidA].currentSessionID) + require.Equal(t, "S1", ws.clients[cidB].currentSessionID) + ws.clientsMu.Unlock() + + // B switches to S2; counts redistribute. + require.NoError(t, b.SetCurrentSession(ws.ID, cidB, "S2")) + ws.clientsMu.Lock() + require.Equal(t, "S1", ws.clients[cidA].currentSessionID) + require.Equal(t, "S2", ws.clients[cidB].currentSessionID) + ws.clientsMu.Unlock() + + // A clears its selection. + require.NoError(t, b.SetCurrentSession(ws.ID, cidA, "")) + ws.clientsMu.Lock() + require.Empty(t, ws.clients[cidA].currentSessionID) + require.Equal(t, "S2", ws.clients[cidB].currentSessionID) + ws.clientsMu.Unlock() + + // Drain to release the workspace. + b.DetachClient(ws.ID, cidA) + b.DetachClient(ws.ID, cidB) +} + +// TestSetCurrentSession_DetachClearsEntry verifies the implicit +// cleanup: once a client's [clientState] entry is removed (last +// stream closed), its currentSessionID is gone with it. +func TestSetCurrentSession_DetachClearsEntry(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, _ := insertTestWorkspace(t, b, "/tmp/current-session-detach") + + // Anchor client so the workspace is not torn down when cid + // detaches. + anchor := newClientID(t) + require.NoError(t, b.AttachClient(ws.ID, anchor)) + + cid := newClientID(t) + require.NoError(t, b.AttachClient(ws.ID, cid)) + require.NoError(t, b.SetCurrentSession(ws.ID, cid, "S2")) + + b.DetachClient(ws.ID, cid) + + ws.clientsMu.Lock() + _, present := ws.clients[cid] + ws.clientsMu.Unlock() + require.False(t, present, "detach must remove the clientState entry along with its currentSessionID") + + // A follow-up SetCurrentSession on the gone client must be + // rejected with ErrClientNotAttached. + require.ErrorIs(t, b.SetCurrentSession(ws.ID, cid, "S3"), ErrClientNotAttached) + + b.DetachClient(ws.ID, anchor) +} + +// TestSetCurrentSession_RejectsHoldOnly verifies that a registered +// client whose only claim is a creation hold (streams == 0) cannot +// influence presence: SetCurrentSession returns ErrClientNotAttached +// and the entry's currentSessionID stays empty. +func TestSetCurrentSession_RejectsHoldOnly(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + // Keep the grace window large so the hold survives the test. + b.createGrace = time.Hour + ws, _ := insertTestWorkspace(t, b, "/tmp/current-session-hold") + + cid := newClientID(t) + b.registerClient(ws, cid) + + require.ErrorIs(t, b.SetCurrentSession(ws.ID, cid, "S1"), ErrClientNotAttached) + + ws.clientsMu.Lock() + require.Empty(t, ws.clients[cid].currentSessionID, "hold-only client must not write a session id") + ws.clientsMu.Unlock() + + // Drain. + require.NoError(t, b.releaseHold(ws.ID, cid)) +} + +// TestSetCurrentSession_UnknownClient verifies that a client with no +// entry at all is rejected with ErrClientNotAttached. +func TestSetCurrentSession_UnknownClient(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, _ := insertTestWorkspace(t, b, "/tmp/current-session-unknown") + + require.ErrorIs(t, b.SetCurrentSession(ws.ID, newClientID(t), "S1"), ErrClientNotAttached) +} + +// TestSetCurrentSession_RejectsBadInputs covers the validation +// branches: empty/malformed client_id and unknown workspace. +func TestSetCurrentSession_RejectsBadInputs(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, _ := insertTestWorkspace(t, b, "/tmp/current-session-bad") + + require.ErrorIs(t, b.SetCurrentSession(ws.ID, "", "S1"), ErrInvalidClientID) + require.ErrorIs(t, b.SetCurrentSession(ws.ID, "not-a-uuid", "S1"), ErrInvalidClientID) + + require.ErrorIs( + t, + b.SetCurrentSession("00000000-0000-0000-0000-000000000000", newClientID(t), "S1"), + ErrWorkspaceNotFound, + ) +} + +// TestSetCurrentSession_RaceWithDetach exercises concurrent +// SetCurrentSession updates from one client racing against detach +// on a second client. The final state must be self-consistent: any +// remaining clientState entries reflect a coherent +// (streams, currentSessionID) pair. +func TestSetCurrentSession_RaceWithDetach(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + ws, _ := insertTestWorkspace(t, b, "/tmp/current-session-race") + + cidA := newClientID(t) + cidB := newClientID(t) + require.NoError(t, b.AttachClient(ws.ID, cidA)) + require.NoError(t, b.AttachClient(ws.ID, cidB)) + + var wg sync.WaitGroup + const updates = 200 + wg.Add(3) + go func() { + defer wg.Done() + for i := range updates { + // Errors are tolerated: once cidA detaches, + // further updates against cidA must return + // ErrClientNotAttached but never panic. + _ = b.SetCurrentSession(ws.ID, cidA, "SA") + _ = i + } + }() + go func() { + defer wg.Done() + for i := range updates { + _ = b.SetCurrentSession(ws.ID, cidB, "SB") + _ = i + } + }() + go func() { + defer wg.Done() + // Single concurrent detach of cidA partway through. + b.DetachClient(ws.ID, cidA) + }() + wg.Wait() + + ws.clientsMu.Lock() + defer ws.clientsMu.Unlock() + require.NotContains(t, ws.clients, cidA, "detached client must be gone") + require.Contains(t, ws.clients, cidB, "remaining client must still be present") + require.Equal(t, "SB", ws.clients[cidB].currentSessionID, "remaining client must keep its last set session") +} + +// TestAttachedClients_BasicLifecycle walks one session's count through +// attach -> set -> second client joins -> switch -> detach. It also +// confirms hold-only and unselected clients do not contribute. +func TestAttachedClients_BasicLifecycle(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + // Keep the grace window long so the hold-only client survives. + b.createGrace = time.Hour + ws, _ := insertTestWorkspace(t, b, "/tmp/attached-clients-basic") + + // No clients yet. + n, err := b.AttachedClients(ws.ID, "S1") + require.NoError(t, err) + require.Zero(t, n) + + // Attach A, set to S1. Count for S1 is 1; count for S2 is 0. + cidA := newClientID(t) + require.NoError(t, b.AttachClient(ws.ID, cidA)) + require.NoError(t, b.SetCurrentSession(ws.ID, cidA, "S1")) + + n, err = b.AttachedClients(ws.ID, "S1") + require.NoError(t, err) + require.Equal(t, 1, n) + n, err = b.AttachedClients(ws.ID, "S2") + require.NoError(t, err) + require.Zero(t, n) + + // Attach B, set to S1. Count for S1 is 2. + cidB := newClientID(t) + require.NoError(t, b.AttachClient(ws.ID, cidB)) + require.NoError(t, b.SetCurrentSession(ws.ID, cidB, "S1")) + + n, _ = b.AttachedClients(ws.ID, "S1") + require.Equal(t, 2, n) + + // B switches to S2; counts redistribute. + require.NoError(t, b.SetCurrentSession(ws.ID, cidB, "S2")) + n, _ = b.AttachedClients(ws.ID, "S1") + require.Equal(t, 1, n) + n, _ = b.AttachedClients(ws.ID, "S2") + require.Equal(t, 1, n) + + // A hold-only client must NOT be counted, even if we were able to + // imagine a currentSessionID on it. registerClient leaves + // currentSessionID empty by construction, and SetCurrentSession + // rejects hold-only writers — so the contract holds two ways. + cidHold := newClientID(t) + b.registerClient(ws, cidHold) + t.Cleanup(func() { _ = b.releaseHold(ws.ID, cidHold) }) + n, _ = b.AttachedClients(ws.ID, "S1") + require.Equal(t, 1, n, "hold-only client must not contribute") + n, _ = b.AttachedClients(ws.ID, "") + require.Equal(t, 0, n, + "empty sessionID must not match the hold-only entry (streams==0)") + + // A client with streams > 0 but currentSessionID == "" is NOT + // counted toward any non-empty session, and is matched only + // against the empty session id (which represents the landing + // screen). + cidC := newClientID(t) + require.NoError(t, b.AttachClient(ws.ID, cidC)) + n, _ = b.AttachedClients(ws.ID, "S1") + require.Equal(t, 1, n, "stream-only client with empty currentSessionID must not be counted toward S1") + n, _ = b.AttachedClients(ws.ID, "") + require.Equal(t, 1, n, "stream-only client with empty currentSessionID matches the empty session id") + + // B detaches: count for S2 drops to 0. + b.DetachClient(ws.ID, cidB) + n, _ = b.AttachedClients(ws.ID, "S2") + require.Zero(t, n) + n, _ = b.AttachedClients(ws.ID, "S1") + require.Equal(t, 1, n, "A still on S1") + + // Final cleanup. + b.DetachClient(ws.ID, cidA) + b.DetachClient(ws.ID, cidC) +} + +// TestAttachedClients_UnknownWorkspace verifies the error surface. +func TestAttachedClients_UnknownWorkspace(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + _, err := b.AttachedClients("00000000-0000-0000-0000-000000000000", "S1") + require.ErrorIs(t, err, ErrWorkspaceNotFound) +} + +// TestTeardown_DefersShutdownWhileCreatePending verifies the core of the +// "coder agent offline after more than one session" fix: tearing down the +// last live workspace must NOT shut the server down while another +// CreateWorkspace is mid-flight (it has committed to the slow init path +// but not yet registered its workspace). Without the pending guard, the +// teardown observed an empty workspace map and killed the whole server +// out from under the workspace being born. +func TestTeardown_DefersShutdownWhileCreatePending(t *testing.T) { + t.Parallel() + + b, serverShutdowns := newTestBackend(t) + ws, wsShutdowns := insertTestWorkspace(t, b, "/tmp/pending-guard") + + // Simulate a concurrent create that has passed the pending++ point + // but has not yet registered its workspace. + b.mu.Lock() + b.pending = 1 + b.mu.Unlock() + + b.teardown(ws) + + require.Equal(t, int32(1), wsShutdowns.Load(), + "the torn-down workspace's own resources must still be released") + require.Equal(t, int32(0), serverShutdowns.Load(), + "server must not shut down while a create is in flight") +} + +// TestTeardown_ShutsDownWhenIdleAndNoPending is the control for the guard +// above: with no create in flight, tearing down the last workspace still +// triggers the server shutdown, preserving the "shut down when empty" +// behavior. +func TestTeardown_ShutsDownWhenIdleAndNoPending(t *testing.T) { + t.Parallel() + + b, serverShutdowns := newTestBackend(t) + ws, _ := insertTestWorkspace(t, b, "/tmp/idle-shutdown") + + b.teardown(ws) + + require.Equal(t, int32(1), serverShutdowns.Load(), + "server must shut down once the last workspace is gone and nothing is pending") +} + +// TestCreateWorkspace_PendingBalancedOnSuccess drives the real +// CreateWorkspace path and asserts the in-flight `pending` counter it +// uses to hold off server shutdown is balanced back to zero once the +// workspace is registered. This guards the bookkeeping the unit-level +// teardown tests stub out: a refactor that drops the increment or +// misplaces the deferred decrement would reintroduce the shutdown race. +func TestCreateWorkspace_PendingBalancedOnSuccess(t *testing.T) { + 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()) + + b, shutdownCount := newTestBackend(t) + // Keep the create hold alive so its grace timer can't fire and tear + // the workspace down mid-assertion. + b.SetCreateGrace(time.Hour) + t.Cleanup(func() { drainBackend(t, b) }) + + _, _, err := b.CreateWorkspace(protoWS(t.TempDir(), t.TempDir(), uuid.New().String())) + require.NoError(t, err) + + b.mu.Lock() + pending := b.pending + b.mu.Unlock() + require.Equal(t, 0, pending, "pending must return to zero after a successful create") + require.Equal(t, int32(0), shutdownCount.Load(), "a live workspace must not trigger shutdown") +} + +// TestCreateWorkspace_PendingBalancedAndReapsOnFailure asserts that a +// create which fails during its slow init still decrements `pending`, and +// that the deferred reap shuts an otherwise-idle server down (the safety +// net that keeps a failed create racing the last teardown from leaking an +// empty server the plain teardown declined to close). +func TestCreateWorkspace_PendingBalancedAndReapsOnFailure(t *testing.T) { + 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()) + + b, shutdownCount := newTestBackend(t) + + // A data dir whose parent is a regular file makes the workspace's + // data-directory creation fail deterministically, after the pending + // counter has been incremented. + tmp := t.TempDir() + fileAsParent := filepath.Join(tmp, "notadir") + require.NoError(t, os.WriteFile(fileAsParent, []byte("x"), 0o600)) + badDataDir := filepath.Join(fileAsParent, "data") + + _, _, err := b.CreateWorkspace(protoWS(tmp, badDataDir, uuid.New().String())) + require.Error(t, err, "create must fail when the data dir cannot be created") + + b.mu.Lock() + pending := b.pending + remaining := b.workspaces.Len() + b.mu.Unlock() + require.Equal(t, 0, pending, "pending must return to zero after a failed create") + require.Zero(t, remaining, "no workspace should be registered after a failed create") + require.Equal(t, int32(1), shutdownCount.Load(), + "a failed create that leaves the server idle must reap it") +} + +// TestServer_CreateCancelsPendingIdleShutdown is the regression test for +// the coder-agent-offline handoff: when the last client of a session +// leaves, the server must linger (not shut down immediately), and a +// client returning within the linger window — the same directory or a +// different one — must cancel that pending shutdown so it is never handed +// a workspace on a server that is about to die. +func TestServer_CreateCancelsPendingIdleShutdown(t *testing.T) { + 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()) + t.Setenv("CRUSH_DISABLE_PROVIDER_AUTO_UPDATE", "1") + + b, shutdownCount := newTestBackend(t) + b.SetIdleShutdownDelay(10 * time.Second) // long enough not to fire mid-test + b.SetCreateGrace(time.Hour) + t.Cleanup(func() { drainBackend(t, b) }) + + // Session 1: a workspace that is attached then leaves, arming the + // idle-shutdown timer. + wsA, _ := insertTestWorkspace(t, b, "/tmp/linger-A") + cidA := newClientID(t) + require.NoError(t, b.AttachClient(wsA.ID, cidA)) + b.DetachClient(wsA.ID, cidA) + + b.mu.Lock() + armed := b.shutdownTimer != nil + b.mu.Unlock() + require.True(t, armed, "idle-shutdown timer must be armed after the last client leaves") + require.Equal(t, int32(0), shutdownCount.Load(), "server must linger, not shut down immediately") + + // Session 2 arrives within the linger window: creating a workspace + // must cancel the pending shutdown. + _, _, err := b.CreateWorkspace(protoWS(t.TempDir(), t.TempDir(), uuid.New().String())) + require.NoError(t, err) + + b.mu.Lock() + stillArmed := b.shutdownTimer != nil + b.mu.Unlock() + require.False(t, stillArmed, "a create within the linger window must cancel the pending shutdown") + require.Equal(t, int32(0), shutdownCount.Load(), "server must not shut down after a client returns") +} + +// TestServer_ShutsDownAfterLingerWhenIdle confirms the linger still ends +// in a shutdown when no client returns: the server is not leaked. +func TestServer_ShutsDownAfterLingerWhenIdle(t *testing.T) { + t.Parallel() + + b, shutdownCount := newTestBackend(t) + b.SetIdleShutdownDelay(100 * time.Millisecond) + + wsA, _ := insertTestWorkspace(t, b, "/tmp/linger-idle") + cidA := newClientID(t) + require.NoError(t, b.AttachClient(wsA.ID, cidA)) + b.DetachClient(wsA.ID, cidA) + + // Nobody returns: after the linger elapses the server shuts down. + require.Eventually(t, func() bool { return shutdownCount.Load() == 1 }, + 2*time.Second, 10*time.Millisecond, + "idle server must shut down after the linger window") +} + +// -- Detach grace -- +// +// The SSE stream is a client's refcount claim, so a stream that drops +// without an explicit release must not destroy the workspace before the +// client's reconnect loop can come back. These cover both outcomes of +// that window plus the fast path a clean exit still gets. + +// TestDetachStream_GraceSurvivesReattach is the core regression: a +// momentary stream drop followed by a reconnect must leave the workspace +// exactly where it was. Before the grace, the reconnect came back to a +// workspace ID the server no longer knew, and every later request 404'd +// forever because a sibling session kept the server itself alive. +func TestDetachStream_GraceSurvivesReattach(t *testing.T) { + t.Parallel() + + b, srvShutdowns := newTestBackend(t) + b.SetDetachGrace(time.Hour) + ws, wsShutdowns := insertTestWorkspace(t, b, "/tmp/blip") + + cid := newClientID(t) + b.registerClient(ws, cid) + require.NoError(t, b.AttachClient(ws.ID, cid)) + + b.DetachClient(ws.ID, cid) + require.Equal(t, int32(0), wsShutdowns.Load(), + "a dropped stream must not tear the workspace down within the grace") + live, err := b.GetWorkspace(ws.ID) + require.NoError(t, err, "the workspace must still be addressable by its original ID") + require.Same(t, ws, live) + + // The reconnect lands and converts the grace back into a stream claim. + require.NoError(t, b.AttachClient(ws.ID, cid)) + ws.clientsMu.Lock() + cs := ws.clients[cid] + require.Equal(t, 1, cs.streams) + require.Nil(t, cs.holdTimer, "re-attaching must cancel the grace timer") + ws.clientsMu.Unlock() + require.Equal(t, int32(0), srvShutdowns.Load()) +} + +// TestDetachStream_GraceExpiryTearsDown checks the grace is bounded: a +// client that never comes back must not pin the workspace (and with it +// the server) indefinitely. +func TestDetachStream_GraceExpiryTearsDown(t *testing.T) { + t.Parallel() + + b, srvShutdowns := newTestBackend(t) + b.SetDetachGrace(50 * time.Millisecond) + ws, wsShutdowns := insertTestWorkspace(t, b, "/tmp/gone-for-good") + + cid := newClientID(t) + b.registerClient(ws, cid) + require.NoError(t, b.AttachClient(ws.ID, cid)) + b.DetachClient(ws.ID, cid) + + require.Eventually(t, func() bool { return srvShutdowns.Load() == 1 }, + 2*time.Second, 10*time.Millisecond, + "an unclaimed grace must expire and tear the workspace down") + require.Equal(t, int32(1), wsShutdowns.Load()) + _, err := b.GetWorkspace(ws.ID) + require.ErrorIs(t, err, ErrWorkspaceNotFound) +} + +// TestDetachStream_ExplicitReleaseSkipsGrace keeps clean exits fast: a +// client that says goodbye before its stream closes is not coming back, +// so the grace must not delay the teardown. +func TestDetachStream_ExplicitReleaseSkipsGrace(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + b.SetDetachGrace(time.Hour) + ws, wsShutdowns := insertTestWorkspace(t, b, "/tmp/clean-exit") + + cid := newClientID(t) + b.registerClient(ws, cid) + require.NoError(t, b.AttachClient(ws.ID, cid)) + + // The order a quitting client produces: release, then the stream ends. + require.NoError(t, b.releaseHold(ws.ID, cid)) + require.Equal(t, int32(0), wsShutdowns.Load(), "the live stream still holds it") + b.DetachClient(ws.ID, cid) + + require.Equal(t, int32(1), wsShutdowns.Load(), + "an explicitly released client must tear down without waiting out the grace") +} + +// TestRegisterClient_RearmsGraceAndCancelsRelease covers a client that +// recovers by re-creating the workspace rather than re-attaching: the +// duplicate create must buy it a fresh attach window and undo an earlier +// release, and the superseded timer must not be able to kill the new claim. +func TestRegisterClient_RearmsGraceAndCancelsRelease(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + b.SetDetachGrace(50 * time.Millisecond) + b.createGrace = time.Hour + ws, wsShutdowns := insertTestWorkspace(t, b, "/tmp/rearm") + + cid := newClientID(t) + b.registerClient(ws, cid) + require.NoError(t, b.AttachClient(ws.ID, cid)) + require.NoError(t, b.releaseHold(ws.ID, cid)) + b.DetachClient(ws.ID, cid) + require.Equal(t, int32(1), wsShutdowns.Load()) + + // A second workspace, this time reclaimed by a duplicate create while + // its short detach grace is running. + ws2, ws2Shutdowns := insertTestWorkspace(t, b, "/tmp/rearm-2") + b.registerClient(ws2, cid) + require.NoError(t, b.AttachClient(ws2.ID, cid)) + b.DetachClient(ws2.ID, cid) + b.registerClient(ws2, cid) + + ws2.clientsMu.Lock() + require.False(t, ws2.clients[cid].released) + ws2.clientsMu.Unlock() + + //nolint:forbidigo // The superseded 50ms timer has to be given its chance to fire. + time.Sleep(200 * time.Millisecond) + require.Equal(t, int32(0), ws2Shutdowns.Load(), + "the re-armed claim must outlive the timer it replaced") +} + +// -- Client retirement -- + +// TestRetireClient_ReleasesEveryClaim checks the primitive a quitting +// client relies on: one call drops its claims everywhere, whether they +// are held by a stream, a create hold, or a detach grace. +func TestRetireClient_ReleasesEveryClaim(t *testing.T) { + t.Parallel() + + b, _ := newTestBackend(t) + b.SetDetachGrace(time.Hour) + b.createGrace = time.Hour + + streamed, streamedShutdowns := insertTestWorkspace(t, b, "/tmp/retire-stream") + held, heldShutdowns := insertTestWorkspace(t, b, "/tmp/retire-hold") + shared, sharedShutdowns := insertTestWorkspace(t, b, "/tmp/retire-shared") + + cid, other := newClientID(t), newClientID(t) + b.registerClient(streamed, cid) + require.NoError(t, b.AttachClient(streamed.ID, cid)) + b.registerClient(held, cid) + b.registerClient(shared, cid) + b.registerClient(shared, other) + + require.NoError(t, b.RetireClient(cid)) + + require.Equal(t, int32(1), streamedShutdowns.Load()) + require.Equal(t, int32(1), heldShutdowns.Load()) + require.Equal(t, int32(0), sharedShutdowns.Load(), + "a workspace another client is still using must be left alone") + shared.clientsMu.Lock() + require.NotContains(t, shared.clients, cid) + require.Contains(t, shared.clients, other) + shared.clientsMu.Unlock() + + // Idempotent: retiring twice is not an error and changes nothing. + require.NoError(t, b.RetireClient(cid)) + require.Equal(t, int32(1), streamedShutdowns.Load()) +} + +// TestRetireClient_RefusesLaterCreates is what makes teardown exact +// without timing guesses. A recovery create whose response the client +// never saw can still land on the server afterwards; refusing it is the +// only way to guarantee it cannot leave a workspace nobody can name. +func TestRetireClient_RefusesLaterCreates(t *testing.T) { + 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()) + + b := New(context.Background(), nil, func() {}) + t.Cleanup(func() { drainBackend(t, b) }) + + cid := newClientID(t) + require.NoError(t, b.RetireClient(cid)) + + _, _, err := b.CreateWorkspace(protoWS(t.TempDir(), t.TempDir(), cid)) + require.ErrorIs(t, err, ErrClientRetired) + require.Zero(t, b.workspaces.Len(), "a refused create must register nothing") +} + +// TestRetireClient_DuringPendingCreate drives the actual ordering the +// design has to survive: retirement lands while a create is already +// inside its slow initialization, holding b.mu released. The create must +// notice at its commit point and discard the workspace it built rather +// than registering a claim for a client that has already gone. +func TestRetireClient_DuringPendingCreate(t *testing.T) { + 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()) + + b := New(context.Background(), nil, func() {}) + t.Cleanup(func() { drainBackend(t, b) }) + + cid := newClientID(t) + cwd, dataDir := t.TempDir(), t.TempDir() + + createErr := make(chan error, 1) + go func() { + _, _, err := b.CreateWorkspace(protoWS(cwd, dataDir, cid)) + createErr <- err + }() + + // pending is bumped under b.mu before the create releases the lock for + // its slow init (config, db, app), and dropped only after the workspace + // is registered. Observing pending == 1 therefore means the create is + // genuinely mid-flight with b.mu free — the exact window retirement has + // to cover. Nothing here fakes the race; the create really is running. + require.Eventually(t, func() bool { + b.mu.Lock() + defer b.mu.Unlock() + return b.pending == 1 + }, 30*time.Second, time.Millisecond, "create must reach its slow path") + + require.NoError(t, b.RetireClient(cid)) + + require.ErrorIs(t, <-createErr, ErrClientRetired, + "a create that commits after retirement must be refused") + require.Zero(t, b.workspaces.Len(), "no workspace may be left behind") +} + +// -- Shutdown as an atomic decision -- + +// TestShutdownIfIdle_RefusesWhileWorkspaceLive is the guard that keeps a +// second session alive: an upgrading client asks a mismatched server to +// stand down, and only the server can answer without racing a session +// that arrives between the question and the answer. +func TestShutdownIfIdle_RefusesWhileWorkspaceLive(t *testing.T) { + t.Parallel() + + b, shutdowns := newTestBackend(t) + ws, _ := insertTestWorkspace(t, b, "/tmp/in-use") + cid := newClientID(t) + b.registerClient(ws, cid) + + require.False(t, b.ShutdownIfIdle(), "a server hosting a workspace must refuse") + require.Equal(t, int32(0), shutdowns.Load()) + + // Refusing must not latch anything: the server keeps serving. + b.mu.Lock() + require.False(t, b.closing) + b.mu.Unlock() +} + +// TestShutdownIfIdle_RefusesWhileCreatePending closes the subtler half of +// the same window: a workspace still working through its slow startup is +// not yet in the workspace map, so a count-based check would call the +// server idle and kill a session that is being born. +func TestShutdownIfIdle_RefusesWhileCreatePending(t *testing.T) { + t.Parallel() + + b, shutdowns := newTestBackend(t) + b.mu.Lock() + b.pending = 1 + b.mu.Unlock() + + require.False(t, b.ShutdownIfIdle()) + require.Equal(t, int32(0), shutdowns.Load()) +} + +// TestShutdownIfIdle_GrantedAndFinal preserves upgrades: an idle server +// does stand down, and the decision is final so a create arriving +// afterwards is refused instead of being initialized on a dying process. +func TestShutdownIfIdle_GrantedAndFinal(t *testing.T) { + t.Parallel() + + b, shutdowns := newTestBackend(t) + require.True(t, b.ShutdownIfIdle()) + require.Equal(t, int32(1), shutdowns.Load()) + + _, _, err := b.CreateWorkspace(protoWS(t.TempDir(), t.TempDir(), newClientID(t))) + require.ErrorIs(t, err, ErrServerShuttingDown) +} + +// TestIdleShutdown_RefusesLaterCreates covers the same finality for the +// unprompted path: a server that decided to exit because it went idle +// must not accept the create of a client that arrives a moment later. +func TestIdleShutdown_RefusesLaterCreates(t *testing.T) { + t.Parallel() + + b, shutdowns := newTestBackend(t) + b.SetIdleShutdownDelay(10 * time.Millisecond) + ws, _ := insertTestWorkspace(t, b, "/tmp/going-idle") + cid := newClientID(t) + require.NoError(t, b.AttachClient(ws.ID, cid)) + b.DetachClient(ws.ID, cid) + + require.Eventually(t, func() bool { return shutdowns.Load() == 1 }, + 2*time.Second, 10*time.Millisecond) + + _, _, err := b.CreateWorkspace(protoWS(t.TempDir(), t.TempDir(), cid)) + require.ErrorIs(t, err, ErrServerShuttingDown, + "the client must be told to retry against a replacement server") +} diff --git a/internal/backend/config.go b/internal/backend/config.go new file mode 100644 index 0000000000000000000000000000000000000000..301bee9bb48ee7506bd39b134a14499401fb866e --- /dev/null +++ b/internal/backend/config.go @@ -0,0 +1,347 @@ +package backend + +import ( + "context" + "errors" + "fmt" + + "github.com/charmbracelet/crush/internal/agent" + 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/oauth" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/charmbracelet/crush/internal/skills" +) + +// publishConfigChanged publishes a ConfigChanged event on the workspace's +// event broker so all subscribers (e.g. remote clients) refresh their +// cached config snapshot. It also re-initializes any MCP servers whose +// configuration changed as a result of the write. +func publishConfigChanged(ws *Workspace) { + if ws == nil || ws.App == nil { + return + } + + // Re-init MCP servers whose config changed. MCP state is process-global, + // so this only needs to happen once regardless of which workspace + // triggered the write. Run async so unrelated config writes (model + // switches, API keys) don't block on MCP reconciliation. Bound to the + // workspace ctx so teardown cancels any in-flight init. + go mcptools.Reinitialize(ws.ctx, ws.Cfg) + + ws.SendEvent(pubsub.Event[proto.ConfigChanged]{ + Type: pubsub.UpdatedEvent, + Payload: proto.ConfigChanged{WorkspaceID: ws.ID}, + }) +} + +// MCPResourceContents holds the contents of an MCP resource returned +// by the backend. +type MCPResourceContents struct { + URI string `json:"uri"` + MIMEType string `json:"mime_type,omitempty"` + Text string `json:"text,omitempty"` + Blob []byte `json:"blob,omitempty"` +} + +// SetConfigField sets a key/value pair in the config file for the +// given scope. +func (b *Backend) SetConfigField(workspaceID string, scope config.Scope, key string, value any) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + if err := ws.Cfg.SetConfigField(scope, key, value); err != nil { + return err + } + publishConfigChanged(ws) + return nil +} + +// RemoveConfigField removes a key from the config file for the given +// scope. +func (b *Backend) RemoveConfigField(workspaceID string, scope config.Scope, key string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + if err := ws.Cfg.RemoveConfigField(scope, key); err != nil { + return err + } + publishConfigChanged(ws) + return nil +} + +// UpdatePreferredModel updates the preferred model for the given type +// and persists it to the config file at the given scope. +func (b *Backend) UpdatePreferredModel(workspaceID string, scope config.Scope, modelType config.SelectedModelType, model config.SelectedModel) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + if err := ws.Cfg.UpdatePreferredModel(scope, modelType, model); err != nil { + return err + } + publishConfigChanged(ws) + return nil +} + +// SetCompactMode sets the compact mode setting and persists it. +func (b *Backend) SetCompactMode(workspaceID string, scope config.Scope, enabled bool) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + if err := ws.Cfg.SetCompactMode(scope, enabled); err != nil { + return err + } + publishConfigChanged(ws) + return nil +} + +// SetProviderAPIKey sets the API key for a provider and persists it. +func (b *Backend) SetProviderAPIKey(workspaceID string, scope config.Scope, providerID string, apiKey any) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + if err := ws.Cfg.SetProviderAPIKey(scope, providerID, apiKey); err != nil { + return err + } + publishConfigChanged(ws) + return nil +} + +// ImportCopilot attempts to import a GitHub Copilot token from disk. +func (b *Backend) ImportCopilot(workspaceID string) (*oauth.Token, bool, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, false, err + } + token, ok := ws.Cfg.ImportCopilot() + if ok { + publishConfigChanged(ws) + } + return token, ok, nil +} + +// RefreshOAuthToken refreshes the OAuth token for a provider. +func (b *Backend) RefreshOAuthToken(ctx context.Context, workspaceID string, scope config.Scope, providerID string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + if err := ws.Cfg.RefreshOAuthToken(ctx, scope, providerID); err != nil { + return err + } + publishConfigChanged(ws) + return nil +} + +// ProjectNeedsInitialization checks whether the project in this +// workspace needs initialization. +func (b *Backend) ProjectNeedsInitialization(workspaceID string) (bool, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return false, err + } + return config.ProjectNeedsInitialization(ws.Cfg) +} + +// MarkProjectInitialized marks the project as initialized. +func (b *Backend) MarkProjectInitialized(workspaceID string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + if err := config.MarkProjectInitialized(ws.Cfg); err != nil { + return err + } + publishConfigChanged(ws) + return nil +} + +// InitializePrompt builds the initialization prompt for the workspace. +func (b *Backend) InitializePrompt(workspaceID string) (string, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return "", err + } + return agent.InitializePrompt(ws.Cfg) +} + +// ReadSkill reads a skill's content by ID. +func (b *Backend) ReadSkill(ctx context.Context, workspaceID, skillID string) ([]byte, proto.SkillReadResult, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, proto.SkillReadResult{}, err + } + + mgr := ws.Skills + content, result, err := skills.ReadContent( + mgr.ActiveSkills(), mgr.ResolvedPaths(), mgr.WorkingDir(), skillID, + ) + if err != nil { + return nil, proto.SkillReadResult{}, err + } + return content, proto.SkillReadResult{ + Name: result.Name, + Description: result.Description, + Source: string(result.Source), + Builtin: result.Builtin, + }, nil +} + +// ListSkills returns the effective visible skills for a workspace. +func (b *Backend) ListSkills(workspaceID string) ([]proto.SkillInfo, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + mgr := ws.Skills + entries := skills.Catalog(mgr.ActiveSkills(), mgr.ResolvedPaths(), mgr.WorkingDir()) + result := make([]proto.SkillInfo, len(entries)) + for i, entry := range entries { + result[i] = proto.SkillInfo{ + ID: entry.ID, + Name: entry.Name, + Description: entry.Description, + Label: entry.Label, + Source: string(entry.Source), + UserInvocable: entry.UserInvocable, + } + } + return result, nil +} + +// EnableDockerMCP validates Docker MCP availability, stages the +// configuration, starts the MCP client, and persists the config. +func (b *Backend) EnableDockerMCP(ctx context.Context, workspaceID string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + + mcpConfig, err := ws.Cfg.PrepareDockerMCPConfig() + if err != nil { + return err + } + + if err := mcptools.InitializeSingle(ctx, config.DockerMCPName, ws.Cfg); err != nil { + disableErr := mcptools.DisableSingle(ws.Cfg, config.DockerMCPName) + ws.Cfg.RemoveDockerMCPInMemory() + return fmt.Errorf("failed to start docker MCP: %w", errors.Join(err, disableErr)) + } + + if err := ws.Cfg.PersistDockerMCPConfig(mcpConfig); err != nil { + disableErr := mcptools.DisableSingle(ws.Cfg, config.DockerMCPName) + ws.Cfg.RemoveDockerMCPInMemory() + return fmt.Errorf("docker MCP started but failed to persist configuration: %w", errors.Join(err, disableErr)) + } + + publishConfigChanged(ws) + return nil +} + +// DisableDockerMCP closes the Docker MCP client, removes the +// configuration, and persists the change. +func (b *Backend) DisableDockerMCP(workspaceID string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + + if err := mcptools.DisableSingle(ws.Cfg, config.DockerMCPName); err != nil { + return fmt.Errorf("failed to disable docker MCP: %w", err) + } + + if err := ws.Cfg.DisableDockerMCP(); err != nil { + return err + } + + publishConfigChanged(ws) + return nil +} + +// RefreshMCPTools refreshes the tools for a named MCP server. +func (b *Backend) RefreshMCPTools(ctx context.Context, workspaceID, name string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + mcptools.RefreshTools(ctx, ws.Cfg, name) + return nil +} + +// ReadMCPResource reads a resource from a named MCP server. +func (b *Backend) ReadMCPResource(ctx context.Context, workspaceID, name, uri string) ([]MCPResourceContents, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + contents, err := mcptools.ReadResource(ctx, ws.Cfg, 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 +} + +// GetMCPPrompt retrieves a prompt from a named MCP server. +func (b *Backend) GetMCPPrompt(workspaceID, clientID, promptID string, args map[string]string) (string, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return "", err + } + return commands.GetMCPPrompt(ws.Cfg, clientID, promptID, args) +} + +func (b *Backend) ListMCPPrompts(workspaceID string) ([]proto.MCPPrompt, error) { + if _, err := b.GetWorkspace(workspaceID); err != nil { + return nil, err + } + prompts, err := commands.LoadMCPPrompts() + if err != nil { + return nil, err + } + result := make([]proto.MCPPrompt, len(prompts)) + for i, prompt := range prompts { + arguments := make([]proto.MCPPromptArgument, len(prompt.Arguments)) + for j, argument := range prompt.Arguments { + arguments[j] = proto.MCPPromptArgument{ + ID: argument.ID, + Title: argument.Title, + Description: argument.Description, + Required: argument.Required, + } + } + result[i] = proto.MCPPrompt{ + ID: prompt.ID, + Title: prompt.Title, + Description: prompt.Description, + PromptID: prompt.PromptID, + ClientID: prompt.ClientID, + Arguments: arguments, + } + } + return result, nil +} + +// GetWorkingDir returns the working directory for a workspace. +func (b *Backend) GetWorkingDir(workspaceID string) (string, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return "", err + } + return ws.Cfg.WorkingDir(), nil +} diff --git a/internal/backend/config_test.go b/internal/backend/config_test.go new file mode 100644 index 0000000000000000000000000000000000000000..858df6dabbe6d318dfa76e4593de314da9c779ce --- /dev/null +++ b/internal/backend/config_test.go @@ -0,0 +1,207 @@ +package backend + +import ( + "context" + "testing" + "time" + + tea "charm.land/bubbletea/v2" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +// awaitConfigChanged drains events until a ConfigChanged is received +// for the given workspace ID, or fails the test on timeout. Other +// event types are ignored. +func awaitConfigChanged(t *testing.T, evc <-chan pubsub.Event[tea.Msg], workspaceID string) { + t.Helper() + deadline := time.After(2 * time.Second) + for { + select { + case ev, ok := <-evc: + if !ok { + t.Fatal("event channel closed before ConfigChanged arrived") + } + cc, ok := ev.Payload.(pubsub.Event[proto.ConfigChanged]) + if !ok { + continue + } + require.Equal(t, workspaceID, cc.Payload.WorkspaceID) + return + case <-deadline: + t.Fatal("timed out waiting for ConfigChanged event") + } + } +} + +// newPublishingWorkspace creates a real workspace through the backend +// so its embedded *app.App is wired up and SendEvent works. It returns +// the backend, the workspace, and a fresh event subscription. +func newPublishingWorkspace(t *testing.T) (*Backend, *Workspace, <-chan pubsub.Event[tea.Msg]) { + t.Helper() + xdgIsolated(t) + + cwd := t.TempDir() + dataDir := t.TempDir() + + b := New(context.Background(), nil, func() {}) + b.SetCreateGrace(2 * time.Second) + t.Cleanup(func() { drainBackend(t, b) }) + + cid := uuid.New().String() + ws, _, err := b.CreateWorkspace(protoWS(cwd, dataDir, cid)) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + return b, ws, ws.Events(ctx) +} + +func TestSetConfigField_PublishesConfigChanged(t *testing.T) { + b, ws, evc := newPublishingWorkspace(t) + + require.NoError(t, b.SetConfigField(ws.ID, config.ScopeGlobal, "options.debug", true)) + awaitConfigChanged(t, evc, ws.ID) +} + +func TestRemoveConfigField_PublishesConfigChanged(t *testing.T) { + b, ws, evc := newPublishingWorkspace(t) + + // Seed a field we can then remove. Setting also publishes, so + // drain the resulting event before testing remove. + require.NoError(t, b.SetConfigField(ws.ID, config.ScopeGlobal, "options.debug", true)) + awaitConfigChanged(t, evc, ws.ID) + + require.NoError(t, b.RemoveConfigField(ws.ID, config.ScopeGlobal, "options.debug")) + awaitConfigChanged(t, evc, ws.ID) +} + +func TestUpdatePreferredModel_PublishesConfigChanged(t *testing.T) { + if raceEnabled { + // UpdatePreferredModel writes config.Models concurrently + // with the agent coordinator's async sub-agent builder + // that reads it via buildAgentModels. That race is + // pre-existing in the codebase and unrelated to this + // item; ConfigStore mutations are not currently + // synchronized against background readers in [app.App]. + // The mutator → publish wiring is unit-tested via + // publishConfigChanged regardless. + t.Skip("skipped under -race: pre-existing race between ConfigStore writes and agent coordinator startup") + } + b, ws, evc := newPublishingWorkspace(t) + + model := config.SelectedModel{Provider: "openai", Model: "gpt-4"} + require.NoError(t, b.UpdatePreferredModel(ws.ID, config.ScopeGlobal, config.SelectedModelTypeLarge, model)) + awaitConfigChanged(t, evc, ws.ID) +} + +func TestSetCompactMode_PublishesConfigChanged(t *testing.T) { + b, ws, evc := newPublishingWorkspace(t) + + require.NoError(t, b.SetCompactMode(ws.ID, config.ScopeGlobal, true)) + awaitConfigChanged(t, evc, ws.ID) +} + +func TestSetProviderAPIKey_PublishesConfigChanged(t *testing.T) { + b, ws, evc := newPublishingWorkspace(t) + + require.NoError(t, b.SetProviderAPIKey(ws.ID, config.ScopeGlobal, "openai", "test-key")) + awaitConfigChanged(t, evc, ws.ID) +} + +func TestMarkProjectInitialized_PublishesConfigChanged(t *testing.T) { + b, ws, evc := newPublishingWorkspace(t) + + require.NoError(t, b.MarkProjectInitialized(ws.ID)) + awaitConfigChanged(t, evc, ws.ID) +} + +// TestImportCopilot_PublishesConfigChanged exercises the success path +// by seeding a token file in the location ImportCopilot scans, then +// asserting the event fires only when ok==true. +func TestImportCopilot_PublishesConfigChanged(t *testing.T) { + // ImportCopilot reads from external user-state directories that + // vary by OS. Rather than recreate that setup, drive the + // publishing helper directly and assert ImportCopilot's + // no-event-on-not-found semantics are preserved. + b, ws, evc := newPublishingWorkspace(t) + + // Not-found path: no token exists, so no event must fire. + _, ok, err := b.ImportCopilot(ws.ID) + require.NoError(t, err) + require.False(t, ok, "ImportCopilot should return ok=false when no token is present") + + select { + case ev := <-evc: + if _, isCC := ev.Payload.(pubsub.Event[proto.ConfigChanged]); isCC { + t.Fatal("ImportCopilot must not publish ConfigChanged when nothing was imported") + } + case <-time.After(100 * time.Millisecond): + // Expected: no ConfigChanged. + } + + // Helper sanity: publishing manually does fire the event. + publishConfigChanged(ws) + awaitConfigChanged(t, evc, ws.ID) +} + +// TestRefreshOAuthToken_PublishesConfigChangedOnError verifies that +// the unhappy path does not publish (mutator returned an error). The +// happy path requires a real OAuth-capable provider configured with a +// refreshable token, which is beyond an isolated unit test's scope. +func TestRefreshOAuthToken_NoEventOnError(t *testing.T) { + b, ws, evc := newPublishingWorkspace(t) + + // Provider does not exist → store returns an error → no event. + err := b.RefreshOAuthToken(context.Background(), ws.ID, config.ScopeGlobal, "no-such-provider") + require.Error(t, err) + + select { + case ev := <-evc: + if _, isCC := ev.Payload.(pubsub.Event[proto.ConfigChanged]); isCC { + t.Fatal("RefreshOAuthToken must not publish ConfigChanged when it errors") + } + case <-time.After(100 * time.Millisecond): + } +} + +// TestDisableDockerMCP_PublishesConfigChanged seeds a Docker MCP entry +// directly so DisableDockerMCP has something to remove without needing +// a running Docker daemon for PrepareDockerMCPConfig's availability +// probe. +func TestDisableDockerMCP_PublishesConfigChanged(t *testing.T) { + b, ws, evc := newPublishingWorkspace(t) + + // Persist a Docker MCP entry directly via the store so the + // downstream DisableDockerMCP path has something to remove. + require.NoError(t, ws.Cfg.PersistDockerMCPConfig(config.DockerMCPConfig())) + drainEvents(evc, 100*time.Millisecond) + + require.NoError(t, b.DisableDockerMCP(ws.ID)) + awaitConfigChanged(t, evc, ws.ID) +} + +// drainEvents reads from evc until quiet for the given window. Used +// to flush events emitted by setup steps so the assertion can target +// the event from the action under test. +func drainEvents(evc <-chan pubsub.Event[tea.Msg], quiet time.Duration) { + for { + select { + case <-evc: + case <-time.After(quiet): + return + } + } +} + +// TestPublishConfigChanged_NilWorkspaceSafe documents that the helper +// is safe to call on workspaces without an *app.App (e.g. synthetic +// test workspaces). +func TestPublishConfigChanged_NilWorkspaceSafe(t *testing.T) { + t.Parallel() + require.NotPanics(t, func() { publishConfigChanged(nil) }) + require.NotPanics(t, func() { publishConfigChanged(&Workspace{}) }) +} diff --git a/internal/backend/events.go b/internal/backend/events.go new file mode 100644 index 0000000000000000000000000000000000000000..bf20fbc0f1cc5670384034e5b59a44f375c3efbf --- /dev/null +++ b/internal/backend/events.go @@ -0,0 +1,145 @@ +package backend + +import ( + "context" + + tea "charm.land/bubbletea/v2" + + mcptools "github.com/charmbracelet/crush/internal/agent/tools/mcp" + "github.com/charmbracelet/crush/internal/app" + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/pubsub" +) + +// SubscribeEvents returns a per-caller event channel for a workspace. +// Each caller receives all events; multiple callers do not compete. +func (b *Backend) SubscribeEvents(ctx context.Context, workspaceID string) (<-chan pubsub.Event[tea.Msg], error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + + return ws.Events(ctx), nil +} + +// GetLSPStates returns the state of all LSP clients. +func (b *Backend) GetLSPStates(workspaceID string) (map[string]app.LSPClientInfo, error) { + _, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + + return app.GetLSPStates(), nil +} + +// GetLSPDiagnostics returns diagnostics for a specific LSP client in +// the workspace. +func (b *Backend) GetLSPDiagnostics(workspaceID, lspName string) (any, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + + for name, client := range ws.LSPManager.Clients().Seq2() { + if name == lspName { + return client.GetDiagnostics(), nil + } + } + + return nil, ErrLSPClientNotFound +} + +// GetWorkspaceConfig returns the workspace-level configuration. +func (b *Backend) GetWorkspaceConfig(workspaceID string) (*config.Config, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + + return ws.Cfg.Config(), nil +} + +// GetWorkspaceProviders returns the configured providers for a +// workspace. +func (b *Backend) GetWorkspaceProviders(workspaceID string) (any, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + + providers, _ := config.Providers(ws.Cfg.Config()) + return providers, nil +} + +// LSPStart starts an LSP server for the given path. +func (b *Backend) LSPStart(ctx context.Context, workspaceID, path string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + + ws.LSPManager.Start(ctx, path) + return nil +} + +// LSPStopAll stops all LSP servers for a workspace. +func (b *Backend) LSPStopAll(ctx context.Context, workspaceID string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + + ws.LSPManager.StopAll(ctx) + return nil +} + +// MCPGetStates returns the current state of all MCP clients. +func (b *Backend) MCPGetStates(_ string) map[string]mcptools.ClientInfo { + return mcptools.GetStates() +} + +// MCPRefreshPrompts refreshes prompts for a named MCP client. +func (b *Backend) MCPRefreshPrompts(ctx context.Context, _ string, name string) { + mcptools.RefreshPrompts(ctx, name) +} + +// MCPRefreshResources refreshes resources for a named MCP client. +func (b *Backend) MCPRefreshResources(ctx context.Context, _ string, name string) { + mcptools.RefreshResources(ctx, name) +} + +// MCPPendingAuth returns the MCP servers awaiting OAuth authentication, +// for clients that need to prompt the user. workspaceID selects the +// workspace whose config provides the server URLs. +func (b *Backend) MCPPendingAuth(workspaceID string) ([]mcptools.PendingAuthServer, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + return mcptools.PendingAuthMCPs(ws.Cfg), nil +} + +// MCPAuthURL returns the current OAuth authorization URL for a named +// server, if a flow is in progress. +func (b *Backend) MCPAuthURL(name string) string { + return mcptools.MCPAuthURL(name) +} + +// MCPAuthenticate runs the OAuth flow for a named MCP server with the +// local browser suppressed: the authorization URL is exposed via +// MCPAuthURL/MCPPendingAuth for the calling client to open on the user's +// machine. The call blocks until the flow completes, fails, or ctx is +// cancelled. workspaceID selects the workspace whose config drives the +// flow. +func (b *Backend) MCPAuthenticate(ctx context.Context, workspaceID, name string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + finish, cancel, err := mcptools.BeginAuth(ws.Cfg, name) + if err != nil { + return err + } + defer cancel() + return finish(ctx) +} diff --git a/internal/backend/filetracker.go b/internal/backend/filetracker.go new file mode 100644 index 0000000000000000000000000000000000000000..14ae99bc7fad5cac541e530cb94162d2218ccaac --- /dev/null +++ b/internal/backend/filetracker.go @@ -0,0 +1,37 @@ +package backend + +import ( + "context" + "time" +) + +// FileTrackerRecordRead records a file read for a session. +func (b *Backend) FileTrackerRecordRead(ctx context.Context, workspaceID, sessionID, path string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + + ws.FileTracker.RecordRead(ctx, sessionID, path) + return nil +} + +// FileTrackerLastReadTime returns the last read time for a file in a session. +func (b *Backend) FileTrackerLastReadTime(ctx context.Context, workspaceID, sessionID, path string) (time.Time, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return time.Time{}, err + } + + return ws.FileTracker.LastReadTime(ctx, sessionID, path), nil +} + +// FileTrackerListReadFiles returns the list of read files for a session. +func (b *Backend) FileTrackerListReadFiles(ctx context.Context, workspaceID, sessionID string) ([]string, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + + return ws.FileTracker.ListReadFiles(ctx, sessionID) +} diff --git a/internal/backend/permission.go b/internal/backend/permission.go new file mode 100644 index 0000000000000000000000000000000000000000..d6db237989ac3a85244c8f9ab4c14df1a7afa1d0 --- /dev/null +++ b/internal/backend/permission.go @@ -0,0 +1,60 @@ +package backend + +import ( + "github.com/charmbracelet/crush/internal/permission" + "github.com/charmbracelet/crush/internal/proto" +) + +// GrantPermission grants, denies, or persistently grants a permission +// request. The returned bool reports whether this call resolved the +// pending request (true) or found it already resolved by a previous +// caller (false). A false return is not an error. +func (b *Backend) GrantPermission(workspaceID string, req proto.PermissionGrant) (bool, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return false, err + } + + perm := permission.PermissionRequest{ + ID: req.Permission.ID, + SessionID: req.Permission.SessionID, + ToolCallID: req.Permission.ToolCallID, + ToolName: req.Permission.ToolName, + Description: req.Permission.Description, + Action: req.Permission.Action, + Params: req.Permission.Params, + Path: req.Permission.Path, + } + + switch req.Action { + case proto.PermissionAllow: + return ws.Permissions.Grant(perm), nil + case proto.PermissionAllowForSession: + return ws.Permissions.GrantPersistent(perm), nil + case proto.PermissionDeny: + return ws.Permissions.Deny(perm), nil + default: + return false, ErrInvalidPermissionAction + } +} + +// SetPermissionsSkip sets whether permission prompts are skipped. +func (b *Backend) SetPermissionsSkip(workspaceID string, skip bool) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + + ws.Permissions.SetSkipRequests(skip) + return nil +} + +// GetPermissionsSkip returns whether permission prompts are skipped. +func (b *Backend) GetPermissionsSkip(workspaceID string) (bool, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return false, err + } + + return ws.Permissions.SkipRequests(), nil +} diff --git a/internal/backend/question.go b/internal/backend/question.go new file mode 100644 index 0000000000000000000000000000000000000000..d0370f04463f1bedaf01638d3d1e3318949189fd --- /dev/null +++ b/internal/backend/question.go @@ -0,0 +1,40 @@ +package backend + +import ( + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/question" +) + +// AnswerQuestion submits answers for a question. The returned bool +// reports whether this call resolved the pending request (true) or +// found it already resolved by a previous caller (false). +func (b *Backend) AnswerQuestion(workspaceID string, req proto.QuestionAnswer) (bool, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return false, err + } + + responses := make([]question.Answer, len(req.Responses)) + for i, r := range req.Responses { + responses[i] = question.Answer{ + QuestionID: r.QuestionID, + SelectedIDs: r.SelectedIDs, + FillInText: r.FillInText, + Yes: r.Yes, + Notes: r.Notes, + } + } + + return ws.Questions.Answer(responses), nil +} + +// CancelQuestion cancels the pending question for a workspace. +// Returns true if a question was cancelled, false if none was +// pending. +func (b *Backend) CancelQuestion(workspaceID string) (bool, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return false, err + } + return ws.Questions.Cancel(), nil +} diff --git a/internal/backend/race_off_test.go b/internal/backend/race_off_test.go new file mode 100644 index 0000000000000000000000000000000000000000..04ff4b864f6382fc8b62231677367c220c86dbe2 --- /dev/null +++ b/internal/backend/race_off_test.go @@ -0,0 +1,5 @@ +//go:build !race + +package backend + +const raceEnabled = false diff --git a/internal/backend/race_on_test.go b/internal/backend/race_on_test.go new file mode 100644 index 0000000000000000000000000000000000000000..1904bea7bf217ac24deae0dcdeba39a527830ce9 --- /dev/null +++ b/internal/backend/race_on_test.go @@ -0,0 +1,5 @@ +//go:build race + +package backend + +const raceEnabled = true diff --git a/internal/backend/session.go b/internal/backend/session.go new file mode 100644 index 0000000000000000000000000000000000000000..9542282319f766e384d9ffe7942e9c235ae63e29 --- /dev/null +++ b/internal/backend/session.go @@ -0,0 +1,132 @@ +package backend + +import ( + "context" + + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/session" +) + +// CreateSession creates a new session in the given workspace. +func (b *Backend) CreateSession(ctx context.Context, workspaceID, title string) (session.Session, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return session.Session{}, err + } + + return ws.Sessions.Create(ctx, title) +} + +// GetSession retrieves a session by workspace and session ID. +func (b *Backend) GetSession(ctx context.Context, workspaceID, sessionID string) (session.Session, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return session.Session{}, err + } + + return ws.Sessions.Get(ctx, sessionID) +} + +// ListSessions returns all sessions in the given workspace. +func (b *Backend) ListSessions(ctx context.Context, workspaceID string) ([]session.Session, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + + return ws.Sessions.List(ctx) +} + +// GetAgentSession returns session metadata with the agent's busy +// status. +func (b *Backend) GetAgentSession(ctx context.Context, workspaceID, sessionID string) (proto.AgentSession, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return proto.AgentSession{}, err + } + + se, err := ws.Sessions.Get(ctx, sessionID) + if err != nil { + return proto.AgentSession{}, err + } + + var isSessionBusy bool + if ws.AgentCoordinator != nil { + isSessionBusy = ws.AgentCoordinator.IsSessionBusy(sessionID) + } + + return proto.AgentSession{ + Session: proto.Session{ + ID: se.ID, + Title: se.Title, + }, + IsBusy: isSessionBusy, + }, nil +} + +// ListSessionMessages returns all messages for a session. +func (b *Backend) ListSessionMessages(ctx context.Context, workspaceID, sessionID string) ([]message.Message, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + + // Drain debounced updates so HTTP clients (and the TUI on session + // switch) observe the latest in-memory state rather than racing the + // debounce timer in message.Service. + if err := ws.Messages.FlushAll(ctx); err != nil { + return nil, err + } + return ws.Messages.List(ctx, sessionID) +} + +// ListSessionHistory returns the history items for a session. +func (b *Backend) ListSessionHistory(ctx context.Context, workspaceID, sessionID string) (any, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + + return ws.History.ListBySession(ctx, sessionID) +} + +// SaveSession updates a session in the given workspace. +func (b *Backend) SaveSession(ctx context.Context, workspaceID string, sess session.Session) (session.Session, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return session.Session{}, err + } + + return ws.Sessions.Save(ctx, sess) +} + +// DeleteSession deletes a session from the given workspace. +func (b *Backend) DeleteSession(ctx context.Context, workspaceID, sessionID string) error { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return err + } + + return ws.Sessions.Delete(ctx, sessionID) +} + +// ListUserMessages returns user-role messages for a session. +func (b *Backend) ListUserMessages(ctx context.Context, workspaceID, sessionID string) ([]message.Message, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + + return ws.Messages.ListUserMessages(ctx, sessionID) +} + +// ListAllUserMessages returns all user-role messages across sessions. +func (b *Backend) ListAllUserMessages(ctx context.Context, workspaceID string) ([]message.Message, error) { + ws, err := b.GetWorkspace(workspaceID) + if err != nil { + return nil, err + } + + return ws.Messages.ListAllUserMessages(ctx) +} diff --git a/internal/backend/shell_command_test.go b/internal/backend/shell_command_test.go new file mode 100644 index 0000000000000000000000000000000000000000..0c79e5f965e1f4022af695976d5dd5a4a7bfd832 --- /dev/null +++ b/internal/backend/shell_command_test.go @@ -0,0 +1,54 @@ +package backend + +import ( + "context" + "testing" + + "github.com/charmbracelet/crush/internal/app" + "github.com/charmbracelet/crush/internal/db" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/session" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func TestRunShellCommand_SkipsPersistenceForMissingSession(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) + + b, _ := newTestBackend(t) + ws := &Workspace{ + ID: uuid.New().String(), + Path: t.TempDir(), + resolvedPath: t.TempDir(), + clients: make(map[string]*clientState), + shutdownFn: func() {}, + } + ws.App = &app.App{ + Sessions: sessions, + Messages: messages, + } + ws.ctx, ws.cancel = context.WithCancel(b.ctx) + InsertWorkspaceForTest(b, ws) + + missingSessionID := uuid.New().String() + resp, err := b.RunShellCommand(t.Context(), ws.ID, proto.ShellCommandRequest{ + SessionID: missingSessionID, + Command: "echo hello", + }) + require.NoError(t, err) + require.Equal(t, "hello\n", resp.Output) + require.Zero(t, resp.ExitCode) + + stored, err := messages.List(t.Context(), missingSessionID) + require.NoError(t, err) + require.Empty(t, stored) +} diff --git a/internal/backend/testing.go b/internal/backend/testing.go new file mode 100644 index 0000000000000000000000000000000000000000..1c71caed6566747ac947d61c30ece575d3d13eb4 --- /dev/null +++ b/internal/backend/testing.go @@ -0,0 +1,69 @@ +package backend + +import "context" + +// InsertWorkspaceForTest registers ws with b under its current ID and +// path. It is intended for tests in other packages that need to drive +// HTTP handlers against a synthetic workspace without booting a real +// app.App. Production code should go through CreateWorkspace. +// +// If the workspace has no run context yet it is derived from the +// backend context (falling back to context.Background), mirroring the +// initialization CreateWorkspace performs, so dispatched agent runs +// have a non-nil ws.ctx. +func InsertWorkspaceForTest(b *Backend, ws *Workspace) { + if ws.resolvedPath == "" { + ws.resolvedPath = ws.Path + } + if ws.clients == nil { + ws.clients = make(map[string]*clientState) + } + if ws.ctx == nil { + parent := b.ctx + if parent == nil { + parent = context.Background() + } + ws.ctx, ws.cancel = context.WithCancel(parent) + } + b.mu.Lock() + defer b.mu.Unlock() + b.workspaces.Set(ws.ID, ws) + if ws.resolvedPath != "" { + b.pathIndex[ws.resolvedPath] = ws.ID + } +} + +// RegisterClientForTesting installs a creation hold for clientID on +// ws using the backend's normal registerClient path. Intended for +// tests in other packages that need to drive a hold-only client +// (streams == 0) without booting a real CreateWorkspace flow. +func RegisterClientForTesting(b *Backend, ws *Workspace, clientID string) error { + if _, err := validateClientID(clientID); err != nil { + return err + } + b.registerClient(ws, clientID) + return nil +} + +// SetWorkspaceShutdownFnForTest overrides the workspace teardown +// callback. Useful for tests in other packages that drive synthetic +// workspaces (where the embedded [app.App] is incomplete) through +// detach paths that would otherwise crash inside App.Shutdown. +func SetWorkspaceShutdownFnForTest(ws *Workspace, fn func()) { + ws.shutdownFn = fn +} + +// WorkspaceLiveStreamCountForTest returns the number of clients on ws +// that have at least one live SSE stream. Used by integration tests +// in other packages to wait for SSE attaches before publishing events. +func WorkspaceLiveStreamCountForTest(ws *Workspace) int { + ws.clientsMu.Lock() + defer ws.clientsMu.Unlock() + n := 0 + for _, cs := range ws.clients { + if cs.streams > 0 { + n++ + } + } + return n +} diff --git a/internal/backend/util.go b/internal/backend/util.go new file mode 100644 index 0000000000000000000000000000000000000000..9b29d0f704ac7473963b5a3765ca60bb7e649bf5 --- /dev/null +++ b/internal/backend/util.go @@ -0,0 +1,22 @@ +package backend + +import ( + "fmt" + "os" + "path/filepath" +) + +func createDotCrushDir(dir string) error { + if err := os.MkdirAll(dir, 0o700); err != nil { + return fmt.Errorf("failed to create data directory: %q %w", dir, err) + } + + gitIgnorePath := filepath.Join(dir, ".gitignore") + if _, err := os.Stat(gitIgnorePath); os.IsNotExist(err) { + if err := os.WriteFile(gitIgnorePath, []byte("*\n"), 0o644); err != nil { + return fmt.Errorf("failed to create .gitignore file: %q %w", gitIgnorePath, err) + } + } + + return nil +} diff --git a/internal/commands/commands_test.go b/internal/commands/commands_test.go new file mode 100644 index 0000000000000000000000000000000000000000..147f63119026fa6f70c0b1ed3a86ec492e912bdb --- /dev/null +++ b/internal/commands/commands_test.go @@ -0,0 +1,114 @@ +package commands + +import ( + "os" + "path/filepath" + "runtime" + "testing" + + "github.com/charmbracelet/crush/internal/skills" + "github.com/stretchr/testify/require" +) + +func TestLoadFromSource_NonExistentDir(t *testing.T) { + t.Parallel() + + dir := filepath.Join(t.TempDir(), "does-not-exist") + + cmds, err := loadFromSource(commandSource{path: dir, prefix: userCommandPrefix}) + require.NoError(t, err) + require.Empty(t, cmds) + + // directory must NOT have been created + _, statErr := os.Stat(dir) + require.True(t, os.IsNotExist(statErr)) +} + +func TestLoadFromSource_ExistingDir(t *testing.T) { + t.Parallel() + + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "hello.md"), []byte("say hello"), 0o644)) + + cmds, err := loadFromSource(commandSource{path: dir, prefix: userCommandPrefix}) + require.NoError(t, err) + require.Len(t, cmds, 1) + require.Equal(t, "user:hello", cmds[0].ID) + require.Equal(t, "say hello", cmds[0].Content) +} + +func TestLoadAll_MixedSources(t *testing.T) { + t.Parallel() + + existing := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(existing, "cmd.md"), []byte("content"), 0o644)) + + missing := filepath.Join(t.TempDir(), "nope") + + cmds, err := loadAll([]commandSource{ + {path: existing, prefix: userCommandPrefix}, + {path: missing, prefix: projectCommandPrefix}, + }) + require.NoError(t, err) + require.Len(t, cmds, 1) + require.Equal(t, "user:cmd", cmds[0].ID) +} + +func TestFromSkillCatalog_UserInvocableOnly(t *testing.T) { + t.Parallel() + + cmds := FromSkillCatalog([]skills.CatalogEntry{ + { + ID: "/skills/on/SKILL.md", + Name: "on", + Description: "Enabled.", + Label: "user:on", + UserInvocable: true, + }, + { + ID: "/skills/off/SKILL.md", + Name: "off", + Description: "Not invocable.", + Label: "user:off", + UserInvocable: false, + }, + }) + + require.Len(t, cmds, 1) + require.Equal(t, "user:on", cmds[0].ID) + require.Equal(t, "user:on", cmds[0].Name) + require.Equal(t, "on", cmds[0].Skill.Name) + require.Equal(t, "Enabled.", cmds[0].Skill.Description) + require.Equal(t, "/skills/on/SKILL.md", cmds[0].Skill.SkillFilePath) +} + +func TestFromSkillCatalog_UsesDiscoveredSymlinkedSkills(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("symlink creation requires special privileges on Windows") + } + t.Parallel() + + root := t.TempDir() + targetParent := t.TempDir() + targetSkillDir := filepath.Join(targetParent, "linked-skill") + require.NoError(t, os.MkdirAll(targetSkillDir, 0o755)) + require.NoError(t, os.WriteFile( + filepath.Join(targetSkillDir, skills.SkillFileName), + []byte("---\nname: linked-skill\ndescription: Symlinked.\nuser-invocable: true\n---\nUse me.\n"), + 0o644, + )) + + link := filepath.Join(root, "linked-skill") + require.NoError(t, os.Symlink(targetSkillDir, link)) + + _, activeSkills, _ := skills.DiscoverFromConfig(skills.DiscoveryConfig{ + SkillsPaths: []string{root}, + }) + entries := skills.Catalog(activeSkills, []string{root}, "") + cmds := FromSkillCatalog(entries) + + require.Len(t, cmds, 1) + require.Equal(t, "user:linked-skill", cmds[0].ID) + require.Equal(t, "linked-skill", cmds[0].Skill.Name) + require.Equal(t, filepath.Join(link, skills.SkillFileName), cmds[0].Skill.SkillFilePath) +} diff --git a/internal/db/connect.go b/internal/db/connect.go new file mode 100644 index 0000000000000000000000000000000000000000..32c2d23933650437d66045887b63f055f8723acc --- /dev/null +++ b/internal/db/connect.go @@ -0,0 +1,248 @@ +package db + +import ( + "context" + "database/sql" + "embed" + "fmt" + "log/slog" + "os" + "path/filepath" + "sync" + "testing" + + "github.com/pressly/goose/v3" +) + +var ( + pragmas = map[string]string{ + "foreign_keys": "ON", + "journal_mode": "WAL", + "page_size": "4096", + "temp_store": "MEMORY", + "cache_size": "-8000", + "synchronous": "NORMAL", + "secure_delete": "ON", + "busy_timeout": "30000", + } + gooseInitOnce sync.Once + gooseInitErr error +) + +//go:embed migrations/*.sql +var FS embed.FS + +func init() { + goose.SetBaseFS(FS) + + if testing.Testing() { + goose.SetLogger(goose.NopLogger()) + } +} + +// connEntry holds a shared database connection, its reference count, +// and the data-directory lock that gates access to this entry. The +// lock is acquired exactly once when the entry is created and released +// when the last reference is dropped, which lets the same process open +// the same data directory concurrently while still blocking a second +// crush process from racing the storage. +type connEntry struct { + db *sql.DB + refCount int + lock *dataDirLock +} + +var ( + pool = make(map[string]*connEntry) + poolMu sync.Mutex +) + +// ConnectOption configures a Connect call. Options are applied in +// order; later options override earlier ones for the same field. +type ConnectOption func(*connectOptions) + +// connectOptions holds the resolved configuration for a Connect call. +type connectOptions struct { + lockDataDir bool +} + +// WithDataDirLock toggles acquisition of the per-data-directory lock +// for this Connect call. The lock is off by default so local-mode +// invocations do not regress today's behavior; the server's +// workspace-bootstrap path opts in. CRUSH_SKIP_DATADIR_LOCK still +// bypasses acquisition even when this option is set. +func WithDataDirLock(enable bool) ConnectOption { + return func(o *connectOptions) { o.lockDataDir = enable } +} + +// Connect opens a SQLite database connection for the given data +// directory and runs migrations. If a connection to the same database +// file already exists, the existing connection is returned with its +// reference count incremented. Callers must pair each Connect with a +// [Release] when they no longer need the connection. +func Connect(ctx context.Context, dataDir string, opts ...ConnectOption) (*sql.DB, error) { + if dataDir == "" { + return nil, fmt.Errorf("data.dir is not set") + } + + var cfg connectOptions + for _, opt := range opts { + opt(&cfg) + } + + dbPath := filepath.Join(dataDir, "crush.db") + + // Resolve to an absolute path so that different relative paths to + // the same file share a single connection. + absPath, err := filepath.Abs(dbPath) + if err != nil { + absPath = dbPath + } + + poolMu.Lock() + defer poolMu.Unlock() + + if entry, ok := pool[absPath]; ok { + entry.refCount++ + return entry.db, nil + } + + // Take the per-data-directory lock before opening the database so + // we fail fast and with a clear error rather than racing another + // crush process on the same SQLite file. The lock is released when + // the matching Release call drops the refcount to zero. Ensuring + // the data directory exists is required because the lock file + // lives inside it. Locking is opt-in via WithDataDirLock so that + // local-mode invocations do not refuse a second crush against the + // same data dir until client/server becomes the default. + if err := os.MkdirAll(dataDir, 0o700); err != nil { + return nil, fmt.Errorf("failed to create data directory %q: %w", dataDir, err) + } + var lock *dataDirLock + if cfg.lockDataDir && !skipDataDirLock() { + lock, err = acquireDataDirLock(dataDir) + if err != nil { + return nil, err + } + } + + conn, err := openDB(dbPath) + if err != nil { + if lock != nil { + lock.release() + } + return nil, err + } + + // Serialize all access through a single connection. SQLite + // serializes writes at the file level anyway, and allowing multiple + // pool connections to interleave writes/checkpoints (especially + // under concurrent sub-agents) has caused WAL/header desync + // resulting in SQLITE_NOTADB (26) on the next open. + conn.SetMaxOpenConns(1) + + releaseLock := func() { + if lock != nil { + lock.release() + } + } + + if err = conn.PingContext(ctx); err != nil { + conn.Close() + releaseLock() + return nil, fmt.Errorf("failed to connect to database: %w", err) + } + + if err := initGoose(); err != nil { + conn.Close() + releaseLock() + slog.Error("Failed to initialize goose", "error", err) + return nil, fmt.Errorf("failed to initialize goose: %w", err) + } + + if err := goose.Up(conn, "migrations"); err != nil { + conn.Close() + releaseLock() + slog.Error("Failed to apply migrations", "error", err) + return nil, fmt.Errorf("failed to apply migrations: %w", err) + } + + pool[absPath] = &connEntry{db: conn, refCount: 1, lock: lock} + return conn, nil +} + +// Release decrements the reference count for the database at the given +// data directory. When the count reaches zero the underlying connection +// is closed and removed from the pool. +func Release(dataDir string) error { + dbPath := filepath.Join(dataDir, "crush.db") + absPath, err := filepath.Abs(dbPath) + if err != nil { + absPath = dbPath + } + + poolMu.Lock() + defer poolMu.Unlock() + + entry, ok := pool[absPath] + if !ok { + return nil + } + + entry.refCount-- + if entry.refCount > 0 { + return nil + } + + delete(pool, absPath) + closeErr := entry.db.Close() + if entry.lock != nil { + entry.lock.release() + } + return closeErr +} + +// ResetPool closes all pooled connections and clears the pool. This is +// intended for use in tests to ensure a clean state between test cases. +func ResetPool() { + poolMu.Lock() + defer poolMu.Unlock() + for path, entry := range pool { + entry.db.Close() + if entry.lock != nil { + entry.lock.release() + } + delete(pool, path) + } +} + +// ConnectReadOnly opens a read-only SQLite database connection without running +// migrations. Used for aggregating stats across multiple project databases. +func ConnectReadOnly(ctx context.Context, dbPath string) (*sql.DB, error) { + if dbPath == "" { + return nil, fmt.Errorf("database path is empty") + } + + db, err := openDBReadOnly(dbPath) + if err != nil { + return nil, err + } + + db.SetMaxOpenConns(1) + + if err = db.PingContext(ctx); err != nil { + db.Close() + return nil, fmt.Errorf("failed to connect to database: %w", err) + } + + return db, nil +} + +func initGoose() error { + gooseInitOnce.Do(func() { + goose.SetBaseFS(FS) + gooseInitErr = goose.SetDialect("sqlite3") + }) + + return gooseInitErr +} diff --git a/internal/db/connect_modernc.go b/internal/db/connect_modernc.go new file mode 100644 index 0000000000000000000000000000000000000000..b6100424327952390ce2bedae88bdb42dc5c340d --- /dev/null +++ b/internal/db/connect_modernc.go @@ -0,0 +1,46 @@ +//go:build (darwin && (amd64 || arm64)) || (freebsd && (amd64 || arm64)) || (linux && (386 || amd64 || arm || arm64 || loong64 || ppc64le || riscv64 || s390x)) || (windows && (386 || amd64 || arm64)) + +package db + +import ( + "database/sql" + "fmt" + "net/url" + + _ "modernc.org/sqlite" +) + +func openDBReadOnly(dbPath string) (*sql.DB, error) { + params := url.Values{} + // Only set safe pragmas for read-only mode - most pragmas require write access + params.Set("_txlock", "immediate") + params.Set("mode", "ro") + + dsn := fmt.Sprintf("file:%s?%s", dbPath, params.Encode()) + db, err := sql.Open("sqlite", dsn) + if err != nil { + return nil, fmt.Errorf("failed to open database: %w", err) + } + + return db, nil +} + +func openDB(dbPath string) (*sql.DB, error) { + // Set pragmas for better performance via _pragma query params. + // Format: _pragma=name(value) + params := url.Values{} + for name, value := range pragmas { + params.Add("_pragma", fmt.Sprintf("%s(%s)", name, value)) + } + // Use BEGIN IMMEDIATE so writers acquire the reserved lock up front, + // preventing deferred-to-writer upgrade deadlocks. + params.Set("_txlock", "immediate") + + dsn := fmt.Sprintf("file:%s?%s", dbPath, params.Encode()) + db, err := sql.Open("sqlite", dsn) + if err != nil { + return nil, fmt.Errorf("failed to open database: %w", err) + } + + return db, nil +} diff --git a/internal/db/connect_ncruces.go b/internal/db/connect_ncruces.go new file mode 100644 index 0000000000000000000000000000000000000000..eb7a591fabe7fcdeb36565a8208cdfa510aacf71 --- /dev/null +++ b/internal/db/connect_ncruces.go @@ -0,0 +1,43 @@ +//go:build !((darwin && (amd64 || arm64)) || (freebsd && (amd64 || arm64)) || (linux && (386 || amd64 || arm || arm64 || loong64 || ppc64le || riscv64 || s390x)) || (windows && (386 || amd64 || arm64))) + +package db + +import ( + "database/sql" + "fmt" + + "github.com/ncruces/go-sqlite3" + "github.com/ncruces/go-sqlite3/driver" +) + +func openDBReadOnly(dbPath string) (*sql.DB, error) { + dsn := fmt.Sprintf("file:%s?mode=ro&_txlock=immediate", dbPath) + db, err := driver.Open(dsn, nil) + if err != nil { + return nil, fmt.Errorf("failed to open database: %w", err) + } + + return db, nil +} + +func openDB(dbPath string) (*sql.DB, error) { + // Use BEGIN IMMEDIATE so writers acquire the reserved lock up front, + // preventing deferred-to-writer upgrade deadlocks. The "file:" prefix + // is required for the ncruces driver to parse query parameters. + dsn := fmt.Sprintf("file:%s?_txlock=immediate", dbPath) + db, err := driver.Open(dsn, func(c *sqlite3.Conn) error { + // Set pragmas for better performance via _pragma query params. + // Format: PRAGMA name = value; + for name, value := range pragmas { + if err := c.Exec(fmt.Sprintf("PRAGMA %s = %s;", name, value)); err != nil { + return fmt.Errorf("failed to set pragma %q: %w", name, err) + } + } + return nil + }) + if err != nil { + return nil, fmt.Errorf("failed to open database: %w", err) + } + + return db, nil +} diff --git a/internal/db/connect_test.go b/internal/db/connect_test.go new file mode 100644 index 0000000000000000000000000000000000000000..d34d11262cce91fe37355232bd41a6cf0d4225aa --- /dev/null +++ b/internal/db/connect_test.go @@ -0,0 +1,210 @@ +package db + +import ( + "context" + "errors" + "path/filepath" + "testing" + + "github.com/charmbracelet/crush/internal/lock" + "github.com/stretchr/testify/require" +) + +func TestConnect_SharesConnectionForSameDataDir(t *testing.T) { + t.Cleanup(ResetPool) + + dataDir := t.TempDir() + + conn1, err := Connect(context.Background(), dataDir) + require.NoError(t, err) + + conn2, err := Connect(context.Background(), dataDir) + require.NoError(t, err) + + require.Same(t, conn1, conn2, "should return the same *sql.DB for the same data dir") + + // Releasing once should not close the connection. + require.NoError(t, Release(dataDir)) + require.NoError(t, conn1.PingContext(context.Background()), "connection should still be usable after partial release") + + // Releasing again should close it. + require.NoError(t, Release(dataDir)) + require.Error(t, conn1.PingContext(context.Background()), "connection should be closed after final release") +} + +func TestConnect_SeparateConnectionsForDifferentDataDirs(t *testing.T) { + t.Cleanup(ResetPool) + + dir1 := t.TempDir() + dir2 := t.TempDir() + + conn1, err := Connect(context.Background(), dir1) + require.NoError(t, err) + + conn2, err := Connect(context.Background(), dir2) + require.NoError(t, err) + + require.NotSame(t, conn1, conn2, "different data dirs should get different connections") + + require.NoError(t, Release(dir1)) + require.NoError(t, Release(dir2)) +} + +func TestRelease_NoopForUnknownDataDir(t *testing.T) { + t.Cleanup(ResetPool) + + require.NoError(t, Release("/nonexistent/path"), "releasing unknown data dir should not error") +} + +// TestConnect_FailsWhenDataDirLocked simulates a second crush process by +// taking the data-dir lock directly via the OS primitive on a separate +// file descriptor and then asserting that Connect surfaces a clean +// ErrDataDirLocked instead of opening the database under contention. +func TestConnect_FailsWhenDataDirLocked(t *testing.T) { + t.Cleanup(ResetPool) + + dataDir := t.TempDir() + lockPath := filepath.Join(dataDir, dataDirLockFile) + + release, err := lock.TryFile(lockPath) + require.NoError(t, err, "expected to take the data-dir lock for the first time") + t.Cleanup(release) + + _, err = Connect(context.Background(), dataDir, WithDataDirLock(true)) + require.Error(t, err, "Connect must refuse to open a locked data dir") + require.ErrorIs(t, err, ErrDataDirLocked) +} + +// TestConnect_SucceedsAfterContenderReleases ensures the lock is purely +// advisory and that a clean release lets the next Connect proceed. +func TestConnect_SucceedsAfterContenderReleases(t *testing.T) { + t.Cleanup(ResetPool) + + dataDir := t.TempDir() + lockPath := filepath.Join(dataDir, dataDirLockFile) + + release, err := lock.TryFile(lockPath) + require.NoError(t, err) + + _, err = Connect(context.Background(), dataDir, WithDataDirLock(true)) + require.ErrorIs(t, err, ErrDataDirLocked) + + release() + + conn, err := Connect(context.Background(), dataDir, WithDataDirLock(true)) + require.NoError(t, err, "Connect should succeed once the contender releases the lock") + require.NoError(t, conn.PingContext(context.Background())) + require.NoError(t, Release(dataDir)) +} + +// TestConnect_LockReleasedOnFinalRelease confirms that closing the last +// reference to a pool entry also drops the OS lock, so subsequent +// processes can take the data dir. +func TestConnect_LockReleasedOnFinalRelease(t *testing.T) { + t.Cleanup(ResetPool) + + dataDir := t.TempDir() + lockPath := filepath.Join(dataDir, dataDirLockFile) + + conn, err := Connect(context.Background(), dataDir, WithDataDirLock(true)) + require.NoError(t, err) + require.NoError(t, conn.PingContext(context.Background())) + + // Holding the in-process entry must keep the OS lock held so a + // "second process" (simulated by a fresh lock.TryFile call) is + // rejected. + _, lockErr := lock.TryFile(lockPath) + require.Error(t, lockErr) + require.True(t, errors.Is(lockErr, lock.ErrContended), "expected contended lock while pool entry is live") + + require.NoError(t, Release(dataDir)) + + // After the final release the lock is free again. + release, err := lock.TryFile(lockPath) + require.NoError(t, err, "expected lock to be released after final Release") + release() +} + +// TestConnect_SharedPoolDoesNotReacquireLock makes sure that subsequent +// in-process Connect calls reuse the existing OS lock through refcount, +// not by re-acquiring it. The simplest observable signal of correctness +// is that the second Connect does not error and the lock is still held +// after a single Release. +func TestConnect_SharedPoolDoesNotReacquireLock(t *testing.T) { + t.Cleanup(ResetPool) + + dataDir := t.TempDir() + lockPath := filepath.Join(dataDir, dataDirLockFile) + + _, err := Connect(context.Background(), dataDir, WithDataDirLock(true)) + require.NoError(t, err) + + _, err = Connect(context.Background(), dataDir, WithDataDirLock(true)) + require.NoError(t, err) + + // Drop one reference; lock must still be held. + require.NoError(t, Release(dataDir)) + _, lockErr := lock.TryFile(lockPath) + require.ErrorIs(t, lockErr, lock.ErrContended) + + require.NoError(t, Release(dataDir)) +} + +// TestConnect_SkipLockEnvBypassesAcquisition exercises the escape +// hatch used by users on filesystems where flock is unreliable. +func TestConnect_SkipLockEnvBypassesAcquisition(t *testing.T) { + t.Cleanup(ResetPool) + + dataDir := t.TempDir() + lockPath := filepath.Join(dataDir, dataDirLockFile) + + release, err := lock.TryFile(lockPath) + require.NoError(t, err) + t.Cleanup(release) + + t.Setenv("CRUSH_SKIP_DATADIR_LOCK", "1") + + conn, err := Connect(context.Background(), dataDir, WithDataDirLock(true)) + require.NoError(t, err, "skip-lock env should bypass contention") + require.NoError(t, conn.PingContext(context.Background())) + require.NoError(t, Release(dataDir)) +} + +// TestConnect_DefaultIgnoresContendedLock confirms that without +// WithDataDirLock(true) the lock file is irrelevant: a contender can +// hold lock.TryFile and Connect still succeeds. This pins the +// local-mode default to its pre-lock behavior. +func TestConnect_DefaultIgnoresContendedLock(t *testing.T) { + t.Cleanup(ResetPool) + + dataDir := t.TempDir() + lockPath := filepath.Join(dataDir, dataDirLockFile) + + release, err := lock.TryFile(lockPath) + require.NoError(t, err, "expected to take the data-dir lock for the first time") + t.Cleanup(release) + + conn, err := Connect(context.Background(), dataDir) + require.NoError(t, err, "default Connect must not take the lock and must succeed under contention") + require.NoError(t, conn.PingContext(context.Background())) + require.NoError(t, Release(dataDir)) +} + +// TestConnect_ServerPathFailsWhenDataDirLocked is the server's +// workspace-bootstrap analogue of TestConnect_FailsWhenDataDirLocked: +// passing WithDataDirLock(true) must surface ErrDataDirLocked when a +// contender already holds the lock. +func TestConnect_ServerPathFailsWhenDataDirLocked(t *testing.T) { + t.Cleanup(ResetPool) + + dataDir := t.TempDir() + lockPath := filepath.Join(dataDir, dataDirLockFile) + + release, err := lock.TryFile(lockPath) + require.NoError(t, err, "expected to take the data-dir lock for the first time") + t.Cleanup(release) + + _, err = Connect(context.Background(), dataDir, WithDataDirLock(true)) + require.Error(t, err, "server-path Connect must refuse to open a locked data dir") + require.ErrorIs(t, err, ErrDataDirLocked) +} diff --git a/internal/db/datadirlock.go b/internal/db/datadirlock.go new file mode 100644 index 0000000000000000000000000000000000000000..2df272f081b3f5ee0341cfd2527aa1362f065ef6 --- /dev/null +++ b/internal/db/datadirlock.go @@ -0,0 +1,131 @@ +package db + +import ( + "encoding/json" + "errors" + "fmt" + "log/slog" + "os" + "path/filepath" + "strconv" + "time" + + "github.com/charmbracelet/crush/internal/lock" + "github.com/charmbracelet/crush/internal/version" +) + +// ErrDataDirLocked is returned by Connect when the data directory is +// already in use by another crush process. +var ErrDataDirLocked = errors.New("data directory already in use by another crush process") + +// dataDirLockFile is the name of the lock file inside the data +// directory. It lives next to crush.db so users can `ls` and find it. +const dataDirLockFile = "crush.lock" + +// dataDirOwnerInfo is the JSON payload written into the lock file by +// the process that currently owns it. It is purely informational; the +// authoritative state of ownership is the operating system flock on +// the file descriptor. +type dataDirOwnerInfo struct { + PID int `json:"pid"` + Version string `json:"version,omitempty"` + StartedAt string `json:"started_at,omitempty"` +} + +// dataDirLock represents an acquired exclusive lock on a data +// directory. release closes the underlying file descriptor which the +// kernel uses to drop the OS-level lock. +type dataDirLock struct { + release func() +} + +// acquireDataDirLock takes an exclusive non-blocking lock on +// {dataDir}/crush.lock. If the lock is already held by another +// process, it returns ErrDataDirLocked wrapped with a diagnostic that +// includes whatever owner info that process wrote. +// +// Acquisition is skipped (returning a no-op lock) when +// CRUSH_SKIP_DATADIR_LOCK is set to a truthy value. This is intended +// as an escape hatch for hostile filesystems that do not implement +// advisory locking; it should not be used in normal operation. +func acquireDataDirLock(dataDir string) (*dataDirLock, error) { + if skipDataDirLock() { + return &dataDirLock{release: func() {}}, nil + } + + path := filepath.Join(dataDir, dataDirLockFile) + release, err := lock.TryFile(path) + if err != nil { + if errors.Is(err, lock.ErrContended) { + return nil, contendedLockError(dataDir, path) + } + return nil, fmt.Errorf("failed to lock data directory %q: %w", dataDir, err) + } + + // Record ownership metadata so a contending process can identify + // us. Failures here are non-fatal: the OS-level lock is what + // actually guarantees mutual exclusion, and a missing/partial JSON + // payload only degrades the diagnostic a contender prints. + if err := writeOwnerInfo(path); err != nil { + slog.Debug("Failed to write data-dir owner info", "path", path, "error", err) + } + + // The lock file itself is intentionally never unlinked. flock is + // keyed by inode, not by path, and any close-then-unlink (or + // unlink-then-close) ordering opens a window where two processes + // can each hold a flock on a different inode that lives at the + // same path. Leaving the file in place lets every acquirer see + // the same inode and lets the kernel arbitrate correctly. + return &dataDirLock{release: release}, nil +} + +// skipDataDirLock reports whether the data-dir lock should be bypassed. +func skipDataDirLock() bool { + v, _ := strconv.ParseBool(os.Getenv("CRUSH_SKIP_DATADIR_LOCK")) + return v +} + +// writeOwnerInfo truncates and rewrites the lock file with the current +// process's identifying information. It is called only after the lock +// is held. +func writeOwnerInfo(path string) error { + info := dataDirOwnerInfo{ + PID: os.Getpid(), + Version: version.Version, + StartedAt: time.Now().UTC().Format(time.RFC3339), + } + payload, err := json.MarshalIndent(info, "", " ") + if err != nil { + return err + } + payload = append(payload, '\n') + return os.WriteFile(path, payload, 0o600) +} + +// readOwnerInfo returns the lock file's recorded owner, if it parses. +// A missing or malformed file yields an empty struct and no error; +// the caller decides what to surface to the user. +func readOwnerInfo(path string) dataDirOwnerInfo { + raw, err := os.ReadFile(path) + if err != nil || len(raw) == 0 { + return dataDirOwnerInfo{} + } + var info dataDirOwnerInfo + _ = json.Unmarshal(raw, &info) + return info +} + +// contendedLockError builds a wrapped ErrDataDirLocked annotated with +// whatever owner metadata is currently in the lock file. +func contendedLockError(dataDir, lockPath string) error { + info := readOwnerInfo(lockPath) + details := "" + switch { + case info.PID != 0 && info.StartedAt != "": + details = fmt.Sprintf(" (owner pid=%d version=%s started_at=%s)", + info.PID, info.Version, info.StartedAt) + case info.PID != 0: + details = fmt.Sprintf(" (owner pid=%d)", info.PID) + } + return fmt.Errorf("%w: %s%s", ErrDataDirLocked, dataDir, details) +} diff --git a/internal/db/db.go b/internal/db/db.go new file mode 100644 index 0000000000000000000000000000000000000000..7fe04f76efcc7bed82f537fd58971763952e6a8d --- /dev/null +++ b/internal/db/db.go @@ -0,0 +1,458 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 + +package db + +import ( + "context" + "database/sql" + "fmt" +) + +type DBTX interface { + ExecContext(context.Context, string, ...interface{}) (sql.Result, error) + PrepareContext(context.Context, string) (*sql.Stmt, error) + QueryContext(context.Context, string, ...interface{}) (*sql.Rows, error) + QueryRowContext(context.Context, string, ...interface{}) *sql.Row +} + +func New(db DBTX) *Queries { + return &Queries{db: db} +} + +func Prepare(ctx context.Context, db DBTX) (*Queries, error) { + q := Queries{db: db} + var err error + if q.createFileStmt, err = db.PrepareContext(ctx, createFile); err != nil { + return nil, fmt.Errorf("error preparing query CreateFile: %w", err) + } + if q.createMessageStmt, err = db.PrepareContext(ctx, createMessage); err != nil { + return nil, fmt.Errorf("error preparing query CreateMessage: %w", err) + } + if q.createSessionStmt, err = db.PrepareContext(ctx, createSession); err != nil { + return nil, fmt.Errorf("error preparing query CreateSession: %w", err) + } + if q.deleteFileStmt, err = db.PrepareContext(ctx, deleteFile); err != nil { + return nil, fmt.Errorf("error preparing query DeleteFile: %w", err) + } + if q.deleteMessageStmt, err = db.PrepareContext(ctx, deleteMessage); err != nil { + return nil, fmt.Errorf("error preparing query DeleteMessage: %w", err) + } + if q.deleteSessionStmt, err = db.PrepareContext(ctx, deleteSession); err != nil { + return nil, fmt.Errorf("error preparing query DeleteSession: %w", err) + } + if q.deleteSessionFilesStmt, err = db.PrepareContext(ctx, deleteSessionFiles); err != nil { + return nil, fmt.Errorf("error preparing query DeleteSessionFiles: %w", err) + } + if q.deleteSessionMessagesStmt, err = db.PrepareContext(ctx, deleteSessionMessages); err != nil { + return nil, fmt.Errorf("error preparing query DeleteSessionMessages: %w", err) + } + if q.getAverageResponseTimeStmt, err = db.PrepareContext(ctx, getAverageResponseTime); err != nil { + return nil, fmt.Errorf("error preparing query GetAverageResponseTime: %w", err) + } + if q.getFileStmt, err = db.PrepareContext(ctx, getFile); err != nil { + return nil, fmt.Errorf("error preparing query GetFile: %w", err) + } + if q.getFileByPathAndSessionStmt, err = db.PrepareContext(ctx, getFileByPathAndSession); err != nil { + return nil, fmt.Errorf("error preparing query GetFileByPathAndSession: %w", err) + } + if q.getFileReadStmt, err = db.PrepareContext(ctx, getFileRead); err != nil { + return nil, fmt.Errorf("error preparing query GetFileRead: %w", err) + } + if q.getHourDayHeatmapStmt, err = db.PrepareContext(ctx, getHourDayHeatmap); err != nil { + return nil, fmt.Errorf("error preparing query GetHourDayHeatmap: %w", err) + } + if q.getLastAssistantMessageBySessionStmt, err = db.PrepareContext(ctx, getLastAssistantMessageBySession); err != nil { + return nil, fmt.Errorf("error preparing query GetLastAssistantMessageBySession: %w", err) + } + if q.getLastSessionStmt, err = db.PrepareContext(ctx, getLastSession); err != nil { + return nil, fmt.Errorf("error preparing query GetLastSession: %w", err) + } + if q.getMessageStmt, err = db.PrepareContext(ctx, getMessage); err != nil { + return nil, fmt.Errorf("error preparing query GetMessage: %w", err) + } + if q.getRecentActivityStmt, err = db.PrepareContext(ctx, getRecentActivity); err != nil { + return nil, fmt.Errorf("error preparing query GetRecentActivity: %w", err) + } + if q.getSessionByIDStmt, err = db.PrepareContext(ctx, getSessionByID); err != nil { + return nil, fmt.Errorf("error preparing query GetSessionByID: %w", err) + } + if q.getToolUsageStmt, err = db.PrepareContext(ctx, getToolUsage); err != nil { + return nil, fmt.Errorf("error preparing query GetToolUsage: %w", err) + } + if q.getTotalStatsStmt, err = db.PrepareContext(ctx, getTotalStats); err != nil { + return nil, fmt.Errorf("error preparing query GetTotalStats: %w", err) + } + if q.getUsageByDayStmt, err = db.PrepareContext(ctx, getUsageByDay); err != nil { + return nil, fmt.Errorf("error preparing query GetUsageByDay: %w", err) + } + if q.getUsageByDayOfWeekStmt, err = db.PrepareContext(ctx, getUsageByDayOfWeek); err != nil { + return nil, fmt.Errorf("error preparing query GetUsageByDayOfWeek: %w", err) + } + if q.getUsageByHourStmt, err = db.PrepareContext(ctx, getUsageByHour); err != nil { + return nil, fmt.Errorf("error preparing query GetUsageByHour: %w", err) + } + if q.getUsageByModelStmt, err = db.PrepareContext(ctx, getUsageByModel); err != nil { + return nil, fmt.Errorf("error preparing query GetUsageByModel: %w", err) + } + if q.listAllUserMessagesStmt, err = db.PrepareContext(ctx, listAllUserMessages); err != nil { + return nil, fmt.Errorf("error preparing query ListAllUserMessages: %w", err) + } + if q.listFilesByPathStmt, err = db.PrepareContext(ctx, listFilesByPath); err != nil { + return nil, fmt.Errorf("error preparing query ListFilesByPath: %w", err) + } + if q.listFilesBySessionStmt, err = db.PrepareContext(ctx, listFilesBySession); err != nil { + return nil, fmt.Errorf("error preparing query ListFilesBySession: %w", err) + } + if q.listLatestSessionFilesStmt, err = db.PrepareContext(ctx, listLatestSessionFiles); err != nil { + return nil, fmt.Errorf("error preparing query ListLatestSessionFiles: %w", err) + } + if q.listMessagesBySessionStmt, err = db.PrepareContext(ctx, listMessagesBySession); err != nil { + return nil, fmt.Errorf("error preparing query ListMessagesBySession: %w", err) + } + if q.listNewFilesStmt, err = db.PrepareContext(ctx, listNewFiles); err != nil { + return nil, fmt.Errorf("error preparing query ListNewFiles: %w", err) + } + if q.listSessionReadFilesStmt, err = db.PrepareContext(ctx, listSessionReadFiles); err != nil { + return nil, fmt.Errorf("error preparing query ListSessionReadFiles: %w", err) + } + if q.listSessionsStmt, err = db.PrepareContext(ctx, listSessions); err != nil { + return nil, fmt.Errorf("error preparing query ListSessions: %w", err) + } + if q.listUserMessagesBySessionStmt, err = db.PrepareContext(ctx, listUserMessagesBySession); err != nil { + return nil, fmt.Errorf("error preparing query ListUserMessagesBySession: %w", err) + } + if q.recordFileReadStmt, err = db.PrepareContext(ctx, recordFileRead); err != nil { + return nil, fmt.Errorf("error preparing query RecordFileRead: %w", err) + } + if q.renameSessionStmt, err = db.PrepareContext(ctx, renameSession); err != nil { + return nil, fmt.Errorf("error preparing query RenameSession: %w", err) + } + if q.updateMessageStmt, err = db.PrepareContext(ctx, updateMessage); err != nil { + return nil, fmt.Errorf("error preparing query UpdateMessage: %w", err) + } + if q.updateSessionStmt, err = db.PrepareContext(ctx, updateSession); err != nil { + return nil, fmt.Errorf("error preparing query UpdateSession: %w", err) + } + if q.updateSessionTitleAndUsageStmt, err = db.PrepareContext(ctx, updateSessionTitleAndUsage); err != nil { + return nil, fmt.Errorf("error preparing query UpdateSessionTitleAndUsage: %w", err) + } + return &q, nil +} + +func (q *Queries) Close() error { + var err error + if q.createFileStmt != nil { + if cerr := q.createFileStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing createFileStmt: %w", cerr) + } + } + if q.createMessageStmt != nil { + if cerr := q.createMessageStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing createMessageStmt: %w", cerr) + } + } + if q.createSessionStmt != nil { + if cerr := q.createSessionStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing createSessionStmt: %w", cerr) + } + } + if q.deleteFileStmt != nil { + if cerr := q.deleteFileStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing deleteFileStmt: %w", cerr) + } + } + if q.deleteMessageStmt != nil { + if cerr := q.deleteMessageStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing deleteMessageStmt: %w", cerr) + } + } + if q.deleteSessionStmt != nil { + if cerr := q.deleteSessionStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing deleteSessionStmt: %w", cerr) + } + } + if q.deleteSessionFilesStmt != nil { + if cerr := q.deleteSessionFilesStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing deleteSessionFilesStmt: %w", cerr) + } + } + if q.deleteSessionMessagesStmt != nil { + if cerr := q.deleteSessionMessagesStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing deleteSessionMessagesStmt: %w", cerr) + } + } + if q.getAverageResponseTimeStmt != nil { + if cerr := q.getAverageResponseTimeStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getAverageResponseTimeStmt: %w", cerr) + } + } + if q.getFileStmt != nil { + if cerr := q.getFileStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getFileStmt: %w", cerr) + } + } + if q.getFileByPathAndSessionStmt != nil { + if cerr := q.getFileByPathAndSessionStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getFileByPathAndSessionStmt: %w", cerr) + } + } + if q.getFileReadStmt != nil { + if cerr := q.getFileReadStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getFileReadStmt: %w", cerr) + } + } + if q.getHourDayHeatmapStmt != nil { + if cerr := q.getHourDayHeatmapStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getHourDayHeatmapStmt: %w", cerr) + } + } + if q.getLastAssistantMessageBySessionStmt != nil { + if cerr := q.getLastAssistantMessageBySessionStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getLastAssistantMessageBySessionStmt: %w", cerr) + } + } + if q.getLastSessionStmt != nil { + if cerr := q.getLastSessionStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getLastSessionStmt: %w", cerr) + } + } + if q.getMessageStmt != nil { + if cerr := q.getMessageStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getMessageStmt: %w", cerr) + } + } + if q.getRecentActivityStmt != nil { + if cerr := q.getRecentActivityStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getRecentActivityStmt: %w", cerr) + } + } + if q.getSessionByIDStmt != nil { + if cerr := q.getSessionByIDStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getSessionByIDStmt: %w", cerr) + } + } + if q.getToolUsageStmt != nil { + if cerr := q.getToolUsageStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getToolUsageStmt: %w", cerr) + } + } + if q.getTotalStatsStmt != nil { + if cerr := q.getTotalStatsStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getTotalStatsStmt: %w", cerr) + } + } + if q.getUsageByDayStmt != nil { + if cerr := q.getUsageByDayStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getUsageByDayStmt: %w", cerr) + } + } + if q.getUsageByDayOfWeekStmt != nil { + if cerr := q.getUsageByDayOfWeekStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getUsageByDayOfWeekStmt: %w", cerr) + } + } + if q.getUsageByHourStmt != nil { + if cerr := q.getUsageByHourStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getUsageByHourStmt: %w", cerr) + } + } + if q.getUsageByModelStmt != nil { + if cerr := q.getUsageByModelStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing getUsageByModelStmt: %w", cerr) + } + } + if q.listAllUserMessagesStmt != nil { + if cerr := q.listAllUserMessagesStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing listAllUserMessagesStmt: %w", cerr) + } + } + if q.listFilesByPathStmt != nil { + if cerr := q.listFilesByPathStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing listFilesByPathStmt: %w", cerr) + } + } + if q.listFilesBySessionStmt != nil { + if cerr := q.listFilesBySessionStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing listFilesBySessionStmt: %w", cerr) + } + } + if q.listLatestSessionFilesStmt != nil { + if cerr := q.listLatestSessionFilesStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing listLatestSessionFilesStmt: %w", cerr) + } + } + if q.listMessagesBySessionStmt != nil { + if cerr := q.listMessagesBySessionStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing listMessagesBySessionStmt: %w", cerr) + } + } + if q.listNewFilesStmt != nil { + if cerr := q.listNewFilesStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing listNewFilesStmt: %w", cerr) + } + } + if q.listSessionReadFilesStmt != nil { + if cerr := q.listSessionReadFilesStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing listSessionReadFilesStmt: %w", cerr) + } + } + if q.listSessionsStmt != nil { + if cerr := q.listSessionsStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing listSessionsStmt: %w", cerr) + } + } + if q.listUserMessagesBySessionStmt != nil { + if cerr := q.listUserMessagesBySessionStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing listUserMessagesBySessionStmt: %w", cerr) + } + } + if q.recordFileReadStmt != nil { + if cerr := q.recordFileReadStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing recordFileReadStmt: %w", cerr) + } + } + if q.renameSessionStmt != nil { + if cerr := q.renameSessionStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing renameSessionStmt: %w", cerr) + } + } + if q.updateMessageStmt != nil { + if cerr := q.updateMessageStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing updateMessageStmt: %w", cerr) + } + } + if q.updateSessionStmt != nil { + if cerr := q.updateSessionStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing updateSessionStmt: %w", cerr) + } + } + if q.updateSessionTitleAndUsageStmt != nil { + if cerr := q.updateSessionTitleAndUsageStmt.Close(); cerr != nil { + err = fmt.Errorf("error closing updateSessionTitleAndUsageStmt: %w", cerr) + } + } + return err +} + +func (q *Queries) exec(ctx context.Context, stmt *sql.Stmt, query string, args ...interface{}) (sql.Result, error) { + switch { + case stmt != nil && q.tx != nil: + return q.tx.StmtContext(ctx, stmt).ExecContext(ctx, args...) + case stmt != nil: + return stmt.ExecContext(ctx, args...) + default: + return q.db.ExecContext(ctx, query, args...) + } +} + +func (q *Queries) query(ctx context.Context, stmt *sql.Stmt, query string, args ...interface{}) (*sql.Rows, error) { + switch { + case stmt != nil && q.tx != nil: + return q.tx.StmtContext(ctx, stmt).QueryContext(ctx, args...) + case stmt != nil: + return stmt.QueryContext(ctx, args...) + default: + return q.db.QueryContext(ctx, query, args...) + } +} + +func (q *Queries) queryRow(ctx context.Context, stmt *sql.Stmt, query string, args ...interface{}) *sql.Row { + switch { + case stmt != nil && q.tx != nil: + return q.tx.StmtContext(ctx, stmt).QueryRowContext(ctx, args...) + case stmt != nil: + return stmt.QueryRowContext(ctx, args...) + default: + return q.db.QueryRowContext(ctx, query, args...) + } +} + +type Queries struct { + db DBTX + tx *sql.Tx + createFileStmt *sql.Stmt + createMessageStmt *sql.Stmt + createSessionStmt *sql.Stmt + deleteFileStmt *sql.Stmt + deleteMessageStmt *sql.Stmt + deleteSessionStmt *sql.Stmt + deleteSessionFilesStmt *sql.Stmt + deleteSessionMessagesStmt *sql.Stmt + getAverageResponseTimeStmt *sql.Stmt + getFileStmt *sql.Stmt + getFileByPathAndSessionStmt *sql.Stmt + getFileReadStmt *sql.Stmt + getHourDayHeatmapStmt *sql.Stmt + getLastAssistantMessageBySessionStmt *sql.Stmt + getLastSessionStmt *sql.Stmt + getMessageStmt *sql.Stmt + getRecentActivityStmt *sql.Stmt + getSessionByIDStmt *sql.Stmt + getToolUsageStmt *sql.Stmt + getTotalStatsStmt *sql.Stmt + getUsageByDayStmt *sql.Stmt + getUsageByDayOfWeekStmt *sql.Stmt + getUsageByHourStmt *sql.Stmt + getUsageByModelStmt *sql.Stmt + listAllUserMessagesStmt *sql.Stmt + listFilesByPathStmt *sql.Stmt + listFilesBySessionStmt *sql.Stmt + listLatestSessionFilesStmt *sql.Stmt + listMessagesBySessionStmt *sql.Stmt + listNewFilesStmt *sql.Stmt + listSessionReadFilesStmt *sql.Stmt + listSessionsStmt *sql.Stmt + listUserMessagesBySessionStmt *sql.Stmt + recordFileReadStmt *sql.Stmt + renameSessionStmt *sql.Stmt + updateMessageStmt *sql.Stmt + updateSessionStmt *sql.Stmt + updateSessionTitleAndUsageStmt *sql.Stmt +} + +func (q *Queries) WithTx(tx *sql.Tx) *Queries { + return &Queries{ + db: tx, + tx: tx, + createFileStmt: q.createFileStmt, + createMessageStmt: q.createMessageStmt, + createSessionStmt: q.createSessionStmt, + deleteFileStmt: q.deleteFileStmt, + deleteMessageStmt: q.deleteMessageStmt, + deleteSessionStmt: q.deleteSessionStmt, + deleteSessionFilesStmt: q.deleteSessionFilesStmt, + deleteSessionMessagesStmt: q.deleteSessionMessagesStmt, + getAverageResponseTimeStmt: q.getAverageResponseTimeStmt, + getFileStmt: q.getFileStmt, + getFileByPathAndSessionStmt: q.getFileByPathAndSessionStmt, + getFileReadStmt: q.getFileReadStmt, + getHourDayHeatmapStmt: q.getHourDayHeatmapStmt, + getLastAssistantMessageBySessionStmt: q.getLastAssistantMessageBySessionStmt, + getLastSessionStmt: q.getLastSessionStmt, + getMessageStmt: q.getMessageStmt, + getRecentActivityStmt: q.getRecentActivityStmt, + getSessionByIDStmt: q.getSessionByIDStmt, + getToolUsageStmt: q.getToolUsageStmt, + getTotalStatsStmt: q.getTotalStatsStmt, + getUsageByDayStmt: q.getUsageByDayStmt, + getUsageByDayOfWeekStmt: q.getUsageByDayOfWeekStmt, + getUsageByHourStmt: q.getUsageByHourStmt, + getUsageByModelStmt: q.getUsageByModelStmt, + listAllUserMessagesStmt: q.listAllUserMessagesStmt, + listFilesByPathStmt: q.listFilesByPathStmt, + listFilesBySessionStmt: q.listFilesBySessionStmt, + listLatestSessionFilesStmt: q.listLatestSessionFilesStmt, + listMessagesBySessionStmt: q.listMessagesBySessionStmt, + listNewFilesStmt: q.listNewFilesStmt, + listSessionReadFilesStmt: q.listSessionReadFilesStmt, + listSessionsStmt: q.listSessionsStmt, + listUserMessagesBySessionStmt: q.listUserMessagesBySessionStmt, + recordFileReadStmt: q.recordFileReadStmt, + renameSessionStmt: q.renameSessionStmt, + updateMessageStmt: q.updateMessageStmt, + updateSessionStmt: q.updateSessionStmt, + updateSessionTitleAndUsageStmt: q.updateSessionTitleAndUsageStmt, + } +} diff --git a/internal/db/files.sql.go b/internal/db/files.sql.go new file mode 100644 index 0000000000000000000000000000000000000000..ec8dfefc734f35d76ce488b19678872038721d71 --- /dev/null +++ b/internal/db/files.sql.go @@ -0,0 +1,280 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: files.sql + +package db + +import ( + "context" +) + +const createFile = `-- name: CreateFile :one +INSERT INTO files ( + id, + session_id, + path, + content, + version, + created_at, + updated_at +) VALUES ( + ?, ?, ?, ?, ?, strftime('%s', 'now'), strftime('%s', 'now') +) +RETURNING id, session_id, path, content, version, created_at, updated_at +` + +type CreateFileParams struct { + ID string `json:"id"` + SessionID string `json:"session_id"` + Path string `json:"path"` + Content string `json:"content"` + Version int64 `json:"version"` +} + +func (q *Queries) CreateFile(ctx context.Context, arg CreateFileParams) (File, error) { + row := q.queryRow(ctx, q.createFileStmt, createFile, + arg.ID, + arg.SessionID, + arg.Path, + arg.Content, + arg.Version, + ) + var i File + err := row.Scan( + &i.ID, + &i.SessionID, + &i.Path, + &i.Content, + &i.Version, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} + +const deleteFile = `-- name: DeleteFile :exec +DELETE FROM files +WHERE id = ? +` + +func (q *Queries) DeleteFile(ctx context.Context, id string) error { + _, err := q.exec(ctx, q.deleteFileStmt, deleteFile, id) + return err +} + +const deleteSessionFiles = `-- name: DeleteSessionFiles :exec +DELETE FROM files +WHERE session_id = ? +` + +func (q *Queries) DeleteSessionFiles(ctx context.Context, sessionID string) error { + _, err := q.exec(ctx, q.deleteSessionFilesStmt, deleteSessionFiles, sessionID) + return err +} + +const getFile = `-- name: GetFile :one +SELECT id, session_id, path, content, version, created_at, updated_at +FROM files +WHERE id = ? LIMIT 1 +` + +func (q *Queries) GetFile(ctx context.Context, id string) (File, error) { + row := q.queryRow(ctx, q.getFileStmt, getFile, id) + var i File + err := row.Scan( + &i.ID, + &i.SessionID, + &i.Path, + &i.Content, + &i.Version, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} + +const getFileByPathAndSession = `-- name: GetFileByPathAndSession :one +SELECT id, session_id, path, content, version, created_at, updated_at +FROM files +WHERE path = ? AND session_id = ? +ORDER BY version DESC, created_at DESC +LIMIT 1 +` + +type GetFileByPathAndSessionParams struct { + Path string `json:"path"` + SessionID string `json:"session_id"` +} + +func (q *Queries) GetFileByPathAndSession(ctx context.Context, arg GetFileByPathAndSessionParams) (File, error) { + row := q.queryRow(ctx, q.getFileByPathAndSessionStmt, getFileByPathAndSession, arg.Path, arg.SessionID) + var i File + err := row.Scan( + &i.ID, + &i.SessionID, + &i.Path, + &i.Content, + &i.Version, + &i.CreatedAt, + &i.UpdatedAt, + ) + return i, err +} + +const listFilesByPath = `-- name: ListFilesByPath :many +SELECT id, session_id, path, content, version, created_at, updated_at +FROM files +WHERE path = ? +ORDER BY version DESC, created_at DESC +` + +func (q *Queries) ListFilesByPath(ctx context.Context, path string) ([]File, error) { + rows, err := q.query(ctx, q.listFilesByPathStmt, listFilesByPath, path) + if err != nil { + return nil, err + } + defer rows.Close() + items := []File{} + for rows.Next() { + var i File + if err := rows.Scan( + &i.ID, + &i.SessionID, + &i.Path, + &i.Content, + &i.Version, + &i.CreatedAt, + &i.UpdatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listFilesBySession = `-- name: ListFilesBySession :many +SELECT id, session_id, path, content, version, created_at, updated_at +FROM files +WHERE session_id = ? +ORDER BY version ASC, created_at ASC +` + +func (q *Queries) ListFilesBySession(ctx context.Context, sessionID string) ([]File, error) { + rows, err := q.query(ctx, q.listFilesBySessionStmt, listFilesBySession, sessionID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []File{} + for rows.Next() { + var i File + if err := rows.Scan( + &i.ID, + &i.SessionID, + &i.Path, + &i.Content, + &i.Version, + &i.CreatedAt, + &i.UpdatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listLatestSessionFiles = `-- name: ListLatestSessionFiles :many +SELECT f.id, f.session_id, f.path, f.content, f.version, f.created_at, f.updated_at +FROM files f +INNER JOIN ( + SELECT path, MAX(version) as max_version, MAX(created_at) as max_created_at + FROM files + GROUP BY path +) latest ON f.path = latest.path AND f.version = latest.max_version AND f.created_at = latest.max_created_at +WHERE f.session_id = ? +ORDER BY f.path +` + +func (q *Queries) ListLatestSessionFiles(ctx context.Context, sessionID string) ([]File, error) { + rows, err := q.query(ctx, q.listLatestSessionFilesStmt, listLatestSessionFiles, sessionID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []File{} + for rows.Next() { + var i File + if err := rows.Scan( + &i.ID, + &i.SessionID, + &i.Path, + &i.Content, + &i.Version, + &i.CreatedAt, + &i.UpdatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listNewFiles = `-- name: ListNewFiles :many +SELECT id, session_id, path, content, version, created_at, updated_at +FROM files +WHERE is_new = 1 +ORDER BY version DESC, created_at DESC +` + +func (q *Queries) ListNewFiles(ctx context.Context) ([]File, error) { + rows, err := q.query(ctx, q.listNewFilesStmt, listNewFiles) + if err != nil { + return nil, err + } + defer rows.Close() + items := []File{} + for rows.Next() { + var i File + if err := rows.Scan( + &i.ID, + &i.SessionID, + &i.Path, + &i.Content, + &i.Version, + &i.CreatedAt, + &i.UpdatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} diff --git a/internal/db/messages.sql.go b/internal/db/messages.sql.go new file mode 100644 index 0000000000000000000000000000000000000000..fc4a57955a184b8e36449824e2beb782d4a59c9e --- /dev/null +++ b/internal/db/messages.sql.go @@ -0,0 +1,317 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: messages.sql + +package db + +import ( + "context" + "database/sql" +) + +const createMessage = `-- name: CreateMessage :one +INSERT INTO messages ( + id, + session_id, + role, + parts, + model, + provider, + is_summary_message, + created_at, + updated_at +) VALUES ( + ?, ?, ?, ?, ?, ?, ?, strftime('%s', 'now'), strftime('%s', 'now') +) +RETURNING id, session_id, role, parts, model, created_at, updated_at, finished_at, provider, is_summary_message, prism_model_id, prism_model_name, prism_hypercredit_savings, prism_dollar_savings +` + +type CreateMessageParams struct { + ID string `json:"id"` + SessionID string `json:"session_id"` + Role string `json:"role"` + Parts string `json:"parts"` + Model sql.NullString `json:"model"` + Provider sql.NullString `json:"provider"` + IsSummaryMessage int64 `json:"is_summary_message"` +} + +func (q *Queries) CreateMessage(ctx context.Context, arg CreateMessageParams) (Message, error) { + row := q.queryRow(ctx, q.createMessageStmt, createMessage, + arg.ID, + arg.SessionID, + arg.Role, + arg.Parts, + arg.Model, + arg.Provider, + arg.IsSummaryMessage, + ) + var i Message + err := row.Scan( + &i.ID, + &i.SessionID, + &i.Role, + &i.Parts, + &i.Model, + &i.CreatedAt, + &i.UpdatedAt, + &i.FinishedAt, + &i.Provider, + &i.IsSummaryMessage, + &i.PrismModelID, + &i.PrismModelName, + &i.PrismHypercreditSavings, + &i.PrismDollarSavings, + ) + return i, err +} + +const deleteMessage = `-- name: DeleteMessage :exec +DELETE FROM messages +WHERE id = ? +` + +func (q *Queries) DeleteMessage(ctx context.Context, id string) error { + _, err := q.exec(ctx, q.deleteMessageStmt, deleteMessage, id) + return err +} + +const deleteSessionMessages = `-- name: DeleteSessionMessages :exec +DELETE FROM messages +WHERE session_id = ? +` + +func (q *Queries) DeleteSessionMessages(ctx context.Context, sessionID string) error { + _, err := q.exec(ctx, q.deleteSessionMessagesStmt, deleteSessionMessages, sessionID) + return err +} + +const getLastAssistantMessageBySession = `-- name: GetLastAssistantMessageBySession :one +SELECT id, session_id, role, parts, model, created_at, updated_at, finished_at, provider, is_summary_message, prism_model_id, prism_model_name, prism_hypercredit_savings, prism_dollar_savings +FROM messages +WHERE session_id = ? AND role = 'assistant' AND is_summary_message = 0 +ORDER BY created_at DESC +LIMIT 1 +` + +func (q *Queries) GetLastAssistantMessageBySession(ctx context.Context, sessionID string) (Message, error) { + row := q.queryRow(ctx, q.getLastAssistantMessageBySessionStmt, getLastAssistantMessageBySession, sessionID) + var i Message + err := row.Scan( + &i.ID, + &i.SessionID, + &i.Role, + &i.Parts, + &i.Model, + &i.CreatedAt, + &i.UpdatedAt, + &i.FinishedAt, + &i.Provider, + &i.IsSummaryMessage, + &i.PrismModelID, + &i.PrismModelName, + &i.PrismHypercreditSavings, + &i.PrismDollarSavings, + ) + return i, err +} + +const getMessage = `-- name: GetMessage :one +SELECT id, session_id, role, parts, model, created_at, updated_at, finished_at, provider, is_summary_message, prism_model_id, prism_model_name, prism_hypercredit_savings, prism_dollar_savings +FROM messages +WHERE id = ? LIMIT 1 +` + +func (q *Queries) GetMessage(ctx context.Context, id string) (Message, error) { + row := q.queryRow(ctx, q.getMessageStmt, getMessage, id) + var i Message + err := row.Scan( + &i.ID, + &i.SessionID, + &i.Role, + &i.Parts, + &i.Model, + &i.CreatedAt, + &i.UpdatedAt, + &i.FinishedAt, + &i.Provider, + &i.IsSummaryMessage, + &i.PrismModelID, + &i.PrismModelName, + &i.PrismHypercreditSavings, + &i.PrismDollarSavings, + ) + return i, err +} + +const listAllUserMessages = `-- name: ListAllUserMessages :many +SELECT id, session_id, role, parts, model, created_at, updated_at, finished_at, provider, is_summary_message, prism_model_id, prism_model_name, prism_hypercredit_savings, prism_dollar_savings +FROM messages +WHERE role = 'user' +ORDER BY created_at DESC +` + +func (q *Queries) ListAllUserMessages(ctx context.Context) ([]Message, error) { + rows, err := q.query(ctx, q.listAllUserMessagesStmt, listAllUserMessages) + if err != nil { + return nil, err + } + defer rows.Close() + items := []Message{} + for rows.Next() { + var i Message + if err := rows.Scan( + &i.ID, + &i.SessionID, + &i.Role, + &i.Parts, + &i.Model, + &i.CreatedAt, + &i.UpdatedAt, + &i.FinishedAt, + &i.Provider, + &i.IsSummaryMessage, + &i.PrismModelID, + &i.PrismModelName, + &i.PrismHypercreditSavings, + &i.PrismDollarSavings, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listMessagesBySession = `-- name: ListMessagesBySession :many +SELECT id, session_id, role, parts, model, created_at, updated_at, finished_at, provider, is_summary_message, prism_model_id, prism_model_name, prism_hypercredit_savings, prism_dollar_savings +FROM messages +WHERE session_id = ? +ORDER BY created_at ASC +` + +func (q *Queries) ListMessagesBySession(ctx context.Context, sessionID string) ([]Message, error) { + rows, err := q.query(ctx, q.listMessagesBySessionStmt, listMessagesBySession, sessionID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []Message{} + for rows.Next() { + var i Message + if err := rows.Scan( + &i.ID, + &i.SessionID, + &i.Role, + &i.Parts, + &i.Model, + &i.CreatedAt, + &i.UpdatedAt, + &i.FinishedAt, + &i.Provider, + &i.IsSummaryMessage, + &i.PrismModelID, + &i.PrismModelName, + &i.PrismHypercreditSavings, + &i.PrismDollarSavings, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listUserMessagesBySession = `-- name: ListUserMessagesBySession :many +SELECT id, session_id, role, parts, model, created_at, updated_at, finished_at, provider, is_summary_message, prism_model_id, prism_model_name, prism_hypercredit_savings, prism_dollar_savings +FROM messages +WHERE session_id = ? AND role = 'user' +ORDER BY created_at DESC +` + +func (q *Queries) ListUserMessagesBySession(ctx context.Context, sessionID string) ([]Message, error) { + rows, err := q.query(ctx, q.listUserMessagesBySessionStmt, listUserMessagesBySession, sessionID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []Message{} + for rows.Next() { + var i Message + if err := rows.Scan( + &i.ID, + &i.SessionID, + &i.Role, + &i.Parts, + &i.Model, + &i.CreatedAt, + &i.UpdatedAt, + &i.FinishedAt, + &i.Provider, + &i.IsSummaryMessage, + &i.PrismModelID, + &i.PrismModelName, + &i.PrismHypercreditSavings, + &i.PrismDollarSavings, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const updateMessage = `-- name: UpdateMessage :exec +UPDATE messages +SET + parts = ?, + prism_model_id = ?, + prism_model_name = ?, + prism_hypercredit_savings = ?, + prism_dollar_savings = ?, + finished_at = ?, + updated_at = strftime('%s', 'now') +WHERE id = ? +` + +type UpdateMessageParams struct { + Parts string `json:"parts"` + PrismModelID sql.NullString `json:"prism_model_id"` + PrismModelName sql.NullString `json:"prism_model_name"` + PrismHypercreditSavings sql.NullFloat64 `json:"prism_hypercredit_savings"` + PrismDollarSavings sql.NullFloat64 `json:"prism_dollar_savings"` + FinishedAt sql.NullInt64 `json:"finished_at"` + ID string `json:"id"` +} + +func (q *Queries) UpdateMessage(ctx context.Context, arg UpdateMessageParams) error { + _, err := q.exec(ctx, q.updateMessageStmt, updateMessage, + arg.Parts, + arg.PrismModelID, + arg.PrismModelName, + arg.PrismHypercreditSavings, + arg.PrismDollarSavings, + arg.FinishedAt, + arg.ID, + ) + return err +} diff --git a/internal/db/models.go b/internal/db/models.go new file mode 100644 index 0000000000000000000000000000000000000000..36b119a1ed4c33a6d3b5955fa3308c8d139cb854 --- /dev/null +++ b/internal/db/models.go @@ -0,0 +1,56 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 + +package db + +import ( + "database/sql" +) + +type File struct { + ID string `json:"id"` + SessionID string `json:"session_id"` + Path string `json:"path"` + Content string `json:"content"` + Version int64 `json:"version"` + CreatedAt int64 `json:"created_at"` + UpdatedAt int64 `json:"updated_at"` +} + +type Message struct { + ID string `json:"id"` + SessionID string `json:"session_id"` + Role string `json:"role"` + Parts string `json:"parts"` + Model sql.NullString `json:"model"` + CreatedAt int64 `json:"created_at"` + UpdatedAt int64 `json:"updated_at"` + FinishedAt sql.NullInt64 `json:"finished_at"` + Provider sql.NullString `json:"provider"` + IsSummaryMessage int64 `json:"is_summary_message"` + PrismModelID sql.NullString `json:"prism_model_id"` + PrismModelName sql.NullString `json:"prism_model_name"` + PrismHypercreditSavings sql.NullFloat64 `json:"prism_hypercredit_savings"` + PrismDollarSavings sql.NullFloat64 `json:"prism_dollar_savings"` +} + +type ReadFile struct { + SessionID string `json:"session_id"` + Path string `json:"path"` + ReadAt int64 `json:"read_at"` +} + +type Session struct { + ID string `json:"id"` + ParentSessionID sql.NullString `json:"parent_session_id"` + Title string `json:"title"` + MessageCount int64 `json:"message_count"` + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + Cost float64 `json:"cost"` + UpdatedAt int64 `json:"updated_at"` + CreatedAt int64 `json:"created_at"` + SummaryMessageID sql.NullString `json:"summary_message_id"` + Todos sql.NullString `json:"todos"` +} diff --git a/internal/db/querier.go b/internal/db/querier.go new file mode 100644 index 0000000000000000000000000000000000000000..9bcd88b62416c8f311fac6f9932e3f9126824049 --- /dev/null +++ b/internal/db/querier.go @@ -0,0 +1,52 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 + +package db + +import ( + "context" +) + +type Querier interface { + CreateFile(ctx context.Context, arg CreateFileParams) (File, error) + CreateMessage(ctx context.Context, arg CreateMessageParams) (Message, error) + CreateSession(ctx context.Context, arg CreateSessionParams) (Session, error) + DeleteFile(ctx context.Context, id string) error + DeleteMessage(ctx context.Context, id string) error + DeleteSession(ctx context.Context, id string) error + DeleteSessionFiles(ctx context.Context, sessionID string) error + DeleteSessionMessages(ctx context.Context, sessionID string) error + GetAverageResponseTime(ctx context.Context) (int64, error) + GetFile(ctx context.Context, id string) (File, error) + GetFileByPathAndSession(ctx context.Context, arg GetFileByPathAndSessionParams) (File, error) + GetFileRead(ctx context.Context, arg GetFileReadParams) (ReadFile, error) + GetHourDayHeatmap(ctx context.Context) ([]GetHourDayHeatmapRow, error) + GetLastAssistantMessageBySession(ctx context.Context, sessionID string) (Message, error) + GetLastSession(ctx context.Context) (Session, error) + GetMessage(ctx context.Context, id string) (Message, error) + GetRecentActivity(ctx context.Context) ([]GetRecentActivityRow, error) + GetSessionByID(ctx context.Context, id string) (Session, error) + GetToolUsage(ctx context.Context) ([]GetToolUsageRow, error) + GetTotalStats(ctx context.Context) (GetTotalStatsRow, error) + GetUsageByDay(ctx context.Context) ([]GetUsageByDayRow, error) + GetUsageByDayOfWeek(ctx context.Context) ([]GetUsageByDayOfWeekRow, error) + GetUsageByHour(ctx context.Context) ([]GetUsageByHourRow, error) + GetUsageByModel(ctx context.Context) ([]GetUsageByModelRow, error) + ListAllUserMessages(ctx context.Context) ([]Message, error) + ListFilesByPath(ctx context.Context, path string) ([]File, error) + ListFilesBySession(ctx context.Context, sessionID string) ([]File, error) + ListLatestSessionFiles(ctx context.Context, sessionID string) ([]File, error) + ListMessagesBySession(ctx context.Context, sessionID string) ([]Message, error) + ListNewFiles(ctx context.Context) ([]File, error) + ListSessionReadFiles(ctx context.Context, sessionID string) ([]ReadFile, error) + ListSessions(ctx context.Context) ([]Session, error) + ListUserMessagesBySession(ctx context.Context, sessionID string) ([]Message, error) + RecordFileRead(ctx context.Context, arg RecordFileReadParams) error + RenameSession(ctx context.Context, arg RenameSessionParams) error + UpdateMessage(ctx context.Context, arg UpdateMessageParams) error + UpdateSession(ctx context.Context, arg UpdateSessionParams) (Session, error) + UpdateSessionTitleAndUsage(ctx context.Context, arg UpdateSessionTitleAndUsageParams) error +} + +var _ Querier = (*Queries)(nil) diff --git a/internal/db/read_files.sql.go b/internal/db/read_files.sql.go new file mode 100644 index 0000000000000000000000000000000000000000..21bda47ce9db1177417fce87677169b49ebd56dc --- /dev/null +++ b/internal/db/read_files.sql.go @@ -0,0 +1,79 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: read_files.sql + +package db + +import ( + "context" +) + +const getFileRead = `-- name: GetFileRead :one +SELECT session_id, path, read_at FROM read_files +WHERE session_id = ? AND path = ? LIMIT 1 +` + +type GetFileReadParams struct { + SessionID string `json:"session_id"` + Path string `json:"path"` +} + +func (q *Queries) GetFileRead(ctx context.Context, arg GetFileReadParams) (ReadFile, error) { + row := q.queryRow(ctx, q.getFileReadStmt, getFileRead, arg.SessionID, arg.Path) + var i ReadFile + err := row.Scan(&i.SessionID, &i.Path, &i.ReadAt) + return i, err +} + +const listSessionReadFiles = `-- name: ListSessionReadFiles :many +SELECT session_id, path, read_at FROM read_files +WHERE session_id = ? +ORDER BY read_at DESC +` + +func (q *Queries) ListSessionReadFiles(ctx context.Context, sessionID string) ([]ReadFile, error) { + rows, err := q.query(ctx, q.listSessionReadFilesStmt, listSessionReadFiles, sessionID) + if err != nil { + return nil, err + } + defer rows.Close() + items := []ReadFile{} + for rows.Next() { + var i ReadFile + if err := rows.Scan(&i.SessionID, &i.Path, &i.ReadAt); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const recordFileRead = `-- name: RecordFileRead :exec +INSERT INTO read_files ( + session_id, + path, + read_at +) VALUES ( + ?, + ?, + strftime('%s', 'now') +) ON CONFLICT(path, session_id) DO UPDATE SET + read_at = excluded.read_at +` + +type RecordFileReadParams struct { + SessionID string `json:"session_id"` + Path string `json:"path"` +} + +func (q *Queries) RecordFileRead(ctx context.Context, arg RecordFileReadParams) error { + _, err := q.exec(ctx, q.recordFileReadStmt, recordFileRead, arg.SessionID, arg.Path) + return err +} diff --git a/internal/db/sessions.sql.go b/internal/db/sessions.sql.go new file mode 100644 index 0000000000000000000000000000000000000000..685948e60e84ec4df66e4d5d1c9645a9ff1fb43f --- /dev/null +++ b/internal/db/sessions.sql.go @@ -0,0 +1,274 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: sessions.sql + +package db + +import ( + "context" + "database/sql" +) + +const createSession = `-- name: CreateSession :one +INSERT INTO sessions ( + id, + parent_session_id, + title, + message_count, + prompt_tokens, + completion_tokens, + cost, + summary_message_id, + updated_at, + created_at +) VALUES ( + ?, + ?, + ?, + ?, + ?, + ?, + ?, + null, + strftime('%s', 'now'), + strftime('%s', 'now') +) RETURNING id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at, summary_message_id, todos +` + +type CreateSessionParams struct { + ID string `json:"id"` + ParentSessionID sql.NullString `json:"parent_session_id"` + Title string `json:"title"` + MessageCount int64 `json:"message_count"` + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + Cost float64 `json:"cost"` +} + +func (q *Queries) CreateSession(ctx context.Context, arg CreateSessionParams) (Session, error) { + row := q.queryRow(ctx, q.createSessionStmt, createSession, + arg.ID, + arg.ParentSessionID, + arg.Title, + arg.MessageCount, + arg.PromptTokens, + arg.CompletionTokens, + arg.Cost, + ) + var i Session + err := row.Scan( + &i.ID, + &i.ParentSessionID, + &i.Title, + &i.MessageCount, + &i.PromptTokens, + &i.CompletionTokens, + &i.Cost, + &i.UpdatedAt, + &i.CreatedAt, + &i.SummaryMessageID, + &i.Todos, + ) + return i, err +} + +const deleteSession = `-- name: DeleteSession :exec +DELETE FROM sessions +WHERE id = ? +` + +func (q *Queries) DeleteSession(ctx context.Context, id string) error { + _, err := q.exec(ctx, q.deleteSessionStmt, deleteSession, id) + return err +} + +const getLastSession = `-- name: GetLastSession :one +SELECT id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at, summary_message_id, todos +FROM sessions +ORDER BY updated_at DESC +LIMIT 1 +` + +func (q *Queries) GetLastSession(ctx context.Context) (Session, error) { + row := q.queryRow(ctx, q.getLastSessionStmt, getLastSession) + var i Session + err := row.Scan( + &i.ID, + &i.ParentSessionID, + &i.Title, + &i.MessageCount, + &i.PromptTokens, + &i.CompletionTokens, + &i.Cost, + &i.UpdatedAt, + &i.CreatedAt, + &i.SummaryMessageID, + &i.Todos, + ) + return i, err +} + +const getSessionByID = `-- name: GetSessionByID :one +SELECT id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at, summary_message_id, todos +FROM sessions +WHERE id = ? LIMIT 1 +` + +func (q *Queries) GetSessionByID(ctx context.Context, id string) (Session, error) { + row := q.queryRow(ctx, q.getSessionByIDStmt, getSessionByID, id) + var i Session + err := row.Scan( + &i.ID, + &i.ParentSessionID, + &i.Title, + &i.MessageCount, + &i.PromptTokens, + &i.CompletionTokens, + &i.Cost, + &i.UpdatedAt, + &i.CreatedAt, + &i.SummaryMessageID, + &i.Todos, + ) + return i, err +} + +const listSessions = `-- name: ListSessions :many +SELECT id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at, summary_message_id, todos +FROM sessions +WHERE parent_session_id is NULL +ORDER BY updated_at DESC +` + +func (q *Queries) ListSessions(ctx context.Context) ([]Session, error) { + rows, err := q.query(ctx, q.listSessionsStmt, listSessions) + if err != nil { + return nil, err + } + defer rows.Close() + items := []Session{} + for rows.Next() { + var i Session + if err := rows.Scan( + &i.ID, + &i.ParentSessionID, + &i.Title, + &i.MessageCount, + &i.PromptTokens, + &i.CompletionTokens, + &i.Cost, + &i.UpdatedAt, + &i.CreatedAt, + &i.SummaryMessageID, + &i.Todos, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const renameSession = `-- name: RenameSession :exec +UPDATE sessions +SET + title = ? +WHERE id = ? +` + +type RenameSessionParams struct { + Title string `json:"title"` + ID string `json:"id"` +} + +func (q *Queries) RenameSession(ctx context.Context, arg RenameSessionParams) error { + _, err := q.exec(ctx, q.renameSessionStmt, renameSession, arg.Title, arg.ID) + return err +} + +const updateSession = `-- name: UpdateSession :one +UPDATE sessions +SET + title = ?, + prompt_tokens = ?, + completion_tokens = ?, + summary_message_id = ?, + cost = ?, + todos = ? +WHERE id = ? +RETURNING id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at, summary_message_id, todos +` + +type UpdateSessionParams struct { + Title string `json:"title"` + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + SummaryMessageID sql.NullString `json:"summary_message_id"` + Cost float64 `json:"cost"` + Todos sql.NullString `json:"todos"` + ID string `json:"id"` +} + +func (q *Queries) UpdateSession(ctx context.Context, arg UpdateSessionParams) (Session, error) { + row := q.queryRow(ctx, q.updateSessionStmt, updateSession, + arg.Title, + arg.PromptTokens, + arg.CompletionTokens, + arg.SummaryMessageID, + arg.Cost, + arg.Todos, + arg.ID, + ) + var i Session + err := row.Scan( + &i.ID, + &i.ParentSessionID, + &i.Title, + &i.MessageCount, + &i.PromptTokens, + &i.CompletionTokens, + &i.Cost, + &i.UpdatedAt, + &i.CreatedAt, + &i.SummaryMessageID, + &i.Todos, + ) + return i, err +} + +const updateSessionTitleAndUsage = `-- name: UpdateSessionTitleAndUsage :exec +UPDATE sessions +SET + title = ?, + prompt_tokens = prompt_tokens + ?, + completion_tokens = completion_tokens + ?, + cost = cost + ?, + updated_at = strftime('%s', 'now') +WHERE id = ? +` + +type UpdateSessionTitleAndUsageParams struct { + Title string `json:"title"` + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + Cost float64 `json:"cost"` + ID string `json:"id"` +} + +func (q *Queries) UpdateSessionTitleAndUsage(ctx context.Context, arg UpdateSessionTitleAndUsageParams) error { + _, err := q.exec(ctx, q.updateSessionTitleAndUsageStmt, updateSessionTitleAndUsage, + arg.Title, + arg.PromptTokens, + arg.CompletionTokens, + arg.Cost, + arg.ID, + ) + return err +} diff --git a/internal/db/stats.sql.go b/internal/db/stats.sql.go new file mode 100644 index 0000000000000000000000000000000000000000..119dd410a07c7e47a2f7369c0ee2fdca8c19b7c3 --- /dev/null +++ b/internal/db/stats.sql.go @@ -0,0 +1,367 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.30.0 +// source: stats.sql + +package db + +import ( + "context" + "database/sql" +) + +const getAverageResponseTime = `-- name: GetAverageResponseTime :one +SELECT + CAST(COALESCE(AVG(finished_at - created_at), 0) AS INTEGER) as avg_response_seconds +FROM messages +WHERE role = 'assistant' + AND finished_at IS NOT NULL + AND finished_at > created_at +` + +func (q *Queries) GetAverageResponseTime(ctx context.Context) (int64, error) { + row := q.queryRow(ctx, q.getAverageResponseTimeStmt, getAverageResponseTime) + var avg_response_seconds int64 + err := row.Scan(&avg_response_seconds) + return avg_response_seconds, err +} + +const getHourDayHeatmap = `-- name: GetHourDayHeatmap :many +SELECT + CAST(strftime('%w', created_at, 'unixepoch') AS INTEGER) as day_of_week, + CAST(strftime('%H', created_at, 'unixepoch') AS INTEGER) as hour, + COUNT(*) as session_count +FROM sessions +WHERE parent_session_id IS NULL +GROUP BY day_of_week, hour +ORDER BY day_of_week, hour +` + +type GetHourDayHeatmapRow struct { + DayOfWeek int64 `json:"day_of_week"` + Hour int64 `json:"hour"` + SessionCount int64 `json:"session_count"` +} + +func (q *Queries) GetHourDayHeatmap(ctx context.Context) ([]GetHourDayHeatmapRow, error) { + rows, err := q.query(ctx, q.getHourDayHeatmapStmt, getHourDayHeatmap) + if err != nil { + return nil, err + } + defer rows.Close() + items := []GetHourDayHeatmapRow{} + for rows.Next() { + var i GetHourDayHeatmapRow + if err := rows.Scan(&i.DayOfWeek, &i.Hour, &i.SessionCount); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getRecentActivity = `-- name: GetRecentActivity :many +SELECT + date(created_at, 'unixepoch') as day, + COUNT(*) as session_count, + SUM(prompt_tokens + completion_tokens) as total_tokens, + SUM(cost) as cost +FROM sessions +WHERE parent_session_id IS NULL + AND created_at >= strftime('%s', 'now', '-30 days') +GROUP BY date(created_at, 'unixepoch') +ORDER BY day ASC +` + +type GetRecentActivityRow struct { + Day interface{} `json:"day"` + SessionCount int64 `json:"session_count"` + TotalTokens sql.NullFloat64 `json:"total_tokens"` + Cost sql.NullFloat64 `json:"cost"` +} + +func (q *Queries) GetRecentActivity(ctx context.Context) ([]GetRecentActivityRow, error) { + rows, err := q.query(ctx, q.getRecentActivityStmt, getRecentActivity) + if err != nil { + return nil, err + } + defer rows.Close() + items := []GetRecentActivityRow{} + for rows.Next() { + var i GetRecentActivityRow + if err := rows.Scan( + &i.Day, + &i.SessionCount, + &i.TotalTokens, + &i.Cost, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getToolUsage = `-- name: GetToolUsage :many +SELECT + json_extract(value, '$.data.name') as tool_name, + COUNT(*) as call_count +FROM messages, json_each(parts) +WHERE json_extract(value, '$.type') = 'tool_call' + AND json_extract(value, '$.data.name') IS NOT NULL +GROUP BY tool_name +ORDER BY call_count DESC +` + +type GetToolUsageRow struct { + ToolName interface{} `json:"tool_name"` + CallCount int64 `json:"call_count"` +} + +func (q *Queries) GetToolUsage(ctx context.Context) ([]GetToolUsageRow, error) { + rows, err := q.query(ctx, q.getToolUsageStmt, getToolUsage) + if err != nil { + return nil, err + } + defer rows.Close() + items := []GetToolUsageRow{} + for rows.Next() { + var i GetToolUsageRow + if err := rows.Scan(&i.ToolName, &i.CallCount); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getTotalStats = `-- name: GetTotalStats :one +SELECT + COUNT(*) as total_sessions, + COALESCE(SUM(prompt_tokens), 0) as total_prompt_tokens, + COALESCE(SUM(completion_tokens), 0) as total_completion_tokens, + COALESCE(SUM(cost), 0) as total_cost, + COALESCE(SUM(message_count), 0) as total_messages, + COALESCE(AVG(prompt_tokens + completion_tokens), 0) as avg_tokens_per_session, + COALESCE(AVG(message_count), 0) as avg_messages_per_session +FROM sessions +WHERE parent_session_id IS NULL +` + +type GetTotalStatsRow struct { + TotalSessions int64 `json:"total_sessions"` + TotalPromptTokens interface{} `json:"total_prompt_tokens"` + TotalCompletionTokens interface{} `json:"total_completion_tokens"` + TotalCost interface{} `json:"total_cost"` + TotalMessages interface{} `json:"total_messages"` + AvgTokensPerSession interface{} `json:"avg_tokens_per_session"` + AvgMessagesPerSession interface{} `json:"avg_messages_per_session"` +} + +func (q *Queries) GetTotalStats(ctx context.Context) (GetTotalStatsRow, error) { + row := q.queryRow(ctx, q.getTotalStatsStmt, getTotalStats) + var i GetTotalStatsRow + err := row.Scan( + &i.TotalSessions, + &i.TotalPromptTokens, + &i.TotalCompletionTokens, + &i.TotalCost, + &i.TotalMessages, + &i.AvgTokensPerSession, + &i.AvgMessagesPerSession, + ) + return i, err +} + +const getUsageByDay = `-- name: GetUsageByDay :many +SELECT + date(created_at, 'unixepoch') as day, + SUM(prompt_tokens) as prompt_tokens, + SUM(completion_tokens) as completion_tokens, + SUM(cost) as cost, + COUNT(*) as session_count +FROM sessions +WHERE parent_session_id IS NULL +GROUP BY date(created_at, 'unixepoch') +ORDER BY day DESC +` + +type GetUsageByDayRow struct { + Day interface{} `json:"day"` + PromptTokens sql.NullFloat64 `json:"prompt_tokens"` + CompletionTokens sql.NullFloat64 `json:"completion_tokens"` + Cost sql.NullFloat64 `json:"cost"` + SessionCount int64 `json:"session_count"` +} + +func (q *Queries) GetUsageByDay(ctx context.Context) ([]GetUsageByDayRow, error) { + rows, err := q.query(ctx, q.getUsageByDayStmt, getUsageByDay) + if err != nil { + return nil, err + } + defer rows.Close() + items := []GetUsageByDayRow{} + for rows.Next() { + var i GetUsageByDayRow + if err := rows.Scan( + &i.Day, + &i.PromptTokens, + &i.CompletionTokens, + &i.Cost, + &i.SessionCount, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getUsageByDayOfWeek = `-- name: GetUsageByDayOfWeek :many +SELECT + CAST(strftime('%w', created_at, 'unixepoch') AS INTEGER) as day_of_week, + COUNT(*) as session_count, + SUM(prompt_tokens) as prompt_tokens, + SUM(completion_tokens) as completion_tokens +FROM sessions +WHERE parent_session_id IS NULL +GROUP BY day_of_week +ORDER BY day_of_week +` + +type GetUsageByDayOfWeekRow struct { + DayOfWeek int64 `json:"day_of_week"` + SessionCount int64 `json:"session_count"` + PromptTokens sql.NullFloat64 `json:"prompt_tokens"` + CompletionTokens sql.NullFloat64 `json:"completion_tokens"` +} + +func (q *Queries) GetUsageByDayOfWeek(ctx context.Context) ([]GetUsageByDayOfWeekRow, error) { + rows, err := q.query(ctx, q.getUsageByDayOfWeekStmt, getUsageByDayOfWeek) + if err != nil { + return nil, err + } + defer rows.Close() + items := []GetUsageByDayOfWeekRow{} + for rows.Next() { + var i GetUsageByDayOfWeekRow + if err := rows.Scan( + &i.DayOfWeek, + &i.SessionCount, + &i.PromptTokens, + &i.CompletionTokens, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getUsageByHour = `-- name: GetUsageByHour :many +SELECT + CAST(strftime('%H', created_at, 'unixepoch') AS INTEGER) as hour, + COUNT(*) as session_count +FROM sessions +WHERE parent_session_id IS NULL +GROUP BY hour +ORDER BY hour +` + +type GetUsageByHourRow struct { + Hour int64 `json:"hour"` + SessionCount int64 `json:"session_count"` +} + +func (q *Queries) GetUsageByHour(ctx context.Context) ([]GetUsageByHourRow, error) { + rows, err := q.query(ctx, q.getUsageByHourStmt, getUsageByHour) + if err != nil { + return nil, err + } + defer rows.Close() + items := []GetUsageByHourRow{} + for rows.Next() { + var i GetUsageByHourRow + if err := rows.Scan(&i.Hour, &i.SessionCount); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getUsageByModel = `-- name: GetUsageByModel :many +SELECT + COALESCE(model, 'unknown') as model, + COALESCE(provider, 'unknown') as provider, + COUNT(*) as message_count +FROM messages +WHERE role = 'assistant' +GROUP BY model, provider +ORDER BY message_count DESC +` + +type GetUsageByModelRow struct { + Model string `json:"model"` + Provider string `json:"provider"` + MessageCount int64 `json:"message_count"` +} + +func (q *Queries) GetUsageByModel(ctx context.Context) ([]GetUsageByModelRow, error) { + rows, err := q.query(ctx, q.getUsageByModelStmt, getUsageByModel) + if err != nil { + return nil, err + } + defer rows.Close() + items := []GetUsageByModelRow{} + for rows.Next() { + var i GetUsageByModelRow + if err := rows.Scan(&i.Model, &i.Provider, &i.MessageCount); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} diff --git a/internal/env/env.go b/internal/env/env.go new file mode 100644 index 0000000000000000000000000000000000000000..2fe19cc20bdd8771cff98da3dd24d7b25c3cba1e --- /dev/null +++ b/internal/env/env.go @@ -0,0 +1,53 @@ +package env + +import ( + "os" +) + +type Env interface { + Get(key string) string + Env() []string +} + +type osEnv struct{} + +// Get implements Env. +func (o *osEnv) Get(key string) string { + return os.Getenv(key) +} + +func (o *osEnv) Env() []string { + return os.Environ() +} + +func New() Env { + return &osEnv{} +} + +type mapEnv struct { + m map[string]string +} + +// Get implements Env. +func (m *mapEnv) Get(key string) string { + if value, ok := m.m[key]; ok { + return value + } + return "" +} + +// Env implements Env. +func (m *mapEnv) Env() []string { + env := make([]string, 0, len(m.m)) + for k, v := range m.m { + env = append(env, k+"="+v) + } + return env +} + +func NewFromMap(m map[string]string) Env { + if m == nil { + m = make(map[string]string) + } + return &mapEnv{m: m} +} diff --git a/internal/env/env_test.go b/internal/env/env_test.go new file mode 100644 index 0000000000000000000000000000000000000000..e8e19cf8b6356e19e4f61cb3bfb2eef7b93ad9fc --- /dev/null +++ b/internal/env/env_test.go @@ -0,0 +1,142 @@ +package env + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestOsEnv_Get(t *testing.T) { + env := &osEnv{} + + // Test getting an existing environment variable + t.Setenv("TEST_VAR", "test_value") + + value := env.Get("TEST_VAR") + require.Equal(t, "test_value", value) + + // Test getting a non-existent environment variable + value = env.Get("NON_EXISTENT_VAR") + require.Equal(t, "", value) +} + +func TestOsEnv_Env(t *testing.T) { + env := &osEnv{} + + envVars := env.Env() + + // Environment should not be empty in normal circumstances + require.NotNil(t, envVars) + require.Greater(t, len(envVars), 0) + + // Each environment variable should be in key=value format + for _, envVar := range envVars { + require.Contains(t, envVar, "=") + } +} + +func TestNewFromMap(t *testing.T) { + testMap := map[string]string{ + "KEY1": "value1", + "KEY2": "value2", + } + + env := NewFromMap(testMap) + require.NotNil(t, env) + require.IsType(t, &mapEnv{}, env) +} + +func TestMapEnv_Get(t *testing.T) { + testMap := map[string]string{ + "KEY1": "value1", + "KEY2": "value2", + } + + env := NewFromMap(testMap) + + // Test getting existing keys + require.Equal(t, "value1", env.Get("KEY1")) + require.Equal(t, "value2", env.Get("KEY2")) + + // Test getting non-existent key + require.Equal(t, "", env.Get("NON_EXISTENT")) +} + +func TestMapEnv_Env(t *testing.T) { + t.Run("with values", func(t *testing.T) { + testMap := map[string]string{ + "KEY1": "value1", + "KEY2": "value2", + } + + env := NewFromMap(testMap) + envVars := env.Env() + + require.Len(t, envVars, 2) + + // Convert to map for easier testing (order is not guaranteed) + envMap := make(map[string]string) + for _, envVar := range envVars { + parts := strings.SplitN(envVar, "=", 2) + require.Len(t, parts, 2) + envMap[parts[0]] = parts[1] + } + + require.Equal(t, "value1", envMap["KEY1"]) + require.Equal(t, "value2", envMap["KEY2"]) + }) + + t.Run("empty map", func(t *testing.T) { + env := NewFromMap(map[string]string{}) + envVars := env.Env() + require.NotNil(t, envVars) + require.Len(t, envVars, 0) + }) + + t.Run("nil map", func(t *testing.T) { + env := NewFromMap(nil) + envVars := env.Env() + require.NotNil(t, envVars) + require.Len(t, envVars, 0) + }) +} + +func TestMapEnv_GetEmptyValue(t *testing.T) { + testMap := map[string]string{ + "EMPTY_KEY": "", + "NORMAL_KEY": "value", + } + + env := NewFromMap(testMap) + + // Test that empty values are returned correctly + require.Equal(t, "", env.Get("EMPTY_KEY")) + require.Equal(t, "value", env.Get("NORMAL_KEY")) +} + +func TestMapEnv_EnvFormat(t *testing.T) { + testMap := map[string]string{ + "KEY_WITH_EQUALS": "value=with=equals", + "KEY_WITH_SPACES": "value with spaces", + } + + env := NewFromMap(testMap) + envVars := env.Env() + + require.Len(t, envVars, 2) + + // Check that the format is correct even with special characters + found := make(map[string]bool) + for _, envVar := range envVars { + if envVar == "KEY_WITH_EQUALS=value=with=equals" { + found["equals"] = true + } + if envVar == "KEY_WITH_SPACES=value with spaces" { + found["spaces"] = true + } + } + + require.True(t, found["equals"], "Should handle values with equals signs") + require.True(t, found["spaces"], "Should handle values with spaces") +} diff --git a/internal/format/spinner.go b/internal/format/spinner.go new file mode 100644 index 0000000000000000000000000000000000000000..49b9e8a7f88cac5ec585bef90fa838ae30edef96 --- /dev/null +++ b/internal/format/spinner.go @@ -0,0 +1,83 @@ +package format + +import ( + "context" + "errors" + "fmt" + "os" + "time" + + tea "charm.land/bubbletea/v2" + "github.com/charmbracelet/crush/internal/ui/anim" + "github.com/charmbracelet/x/ansi" +) + +// Spinner wraps the bubbles spinner for non-interactive mode +type Spinner struct { + done chan struct{} + prog *tea.Program +} + +type model struct { + cancel context.CancelFunc + anim *anim.Anim +} + +type tickMsg struct{} + +func tick() tea.Cmd { + return tea.Tick(anim.FrameInterval(), func(time.Time) tea.Msg { return tickMsg{} }) +} + +func (m model) Init() tea.Cmd { return tick() } +func (m model) View() tea.View { return tea.NewView(m.anim.Render()) } + +// Update implements tea.Model. +func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.KeyPressMsg: + switch msg.String() { + case "ctrl+c", "esc": + m.cancel() + return m, tea.Quit + } + case tickMsg: + m.anim.Advance() + return m, tick() + } + return m, nil +} + +// NewSpinner creates a new spinner with the given message +func NewSpinner(ctx context.Context, cancel context.CancelFunc, animSettings anim.Settings) *Spinner { + m := model{ + anim: anim.New(animSettings), + cancel: cancel, + } + + p := tea.NewProgram(m, tea.WithOutput(os.Stderr), tea.WithContext(ctx)) + + return &Spinner{ + prog: p, + done: make(chan struct{}, 1), + } +} + +// Start begins the spinner animation +func (s *Spinner) Start() { + go func() { + defer close(s.done) + _, err := s.prog.Run() + // ensures line is cleared + fmt.Fprint(os.Stderr, ansi.EraseEntireLine) + if err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, tea.ErrInterrupted) { + fmt.Fprintf(os.Stderr, "Error running spinner: %v\n", err) + } + }() +} + +// Stop ends the spinner animation +func (s *Spinner) Stop() { + s.prog.Quit() + <-s.done +} diff --git a/internal/fsext/drive_other.go b/internal/fsext/drive_other.go new file mode 100644 index 0000000000000000000000000000000000000000..0141c8baf0a640a6c3fc45500c9acb798664b16a --- /dev/null +++ b/internal/fsext/drive_other.go @@ -0,0 +1,16 @@ +//go:build !windows + +package fsext + +// WindowsWorkingDirDrive returns the drive letter of the current working +// directory, e.g. "C:". +// Falls back to the system drive if the current working directory cannot be +// determined. +func WindowsWorkingDirDrive() string { + panic("cannot call fsext.WindowsWorkingDirDrive() on non-Windows OS") +} + +// WindowsSystemDrive returns the drive letter of the system drive, e.g. "C:". +func WindowsSystemDrive() string { + panic("cannot call fsext.WindowsSystemDrive() on non-Windows OS") +} diff --git a/internal/fsext/drive_windows.go b/internal/fsext/drive_windows.go new file mode 100644 index 0000000000000000000000000000000000000000..7d0a798da27fb59d38bce5e5bec52560ff18b54b --- /dev/null +++ b/internal/fsext/drive_windows.go @@ -0,0 +1,26 @@ +//go:build windows + +package fsext + +import ( + "cmp" + "os" + "path/filepath" +) + +// WindowsWorkingDirDrive returns the drive letter of the current working +// directory, e.g. "C:". +// Falls back to the system drive if the current working directory cannot be +// determined. +func WindowsWorkingDirDrive() string { + if cwd, err := os.Getwd(); err == nil { + return filepath.VolumeName(cwd) + } + return WindowsSystemDrive() +} + +// WindowsSystemDrive returns the drive letter of the system drive, e.g. "C:". +func WindowsSystemDrive() string { + systemRoot := cmp.Or(os.Getenv("SYSTEMROOT"), os.Getenv("WINDIR")) + return filepath.VolumeName(systemRoot) +} diff --git a/internal/fsext/expand.go b/internal/fsext/expand.go new file mode 100644 index 0000000000000000000000000000000000000000..d4265ca6381f43056032ffdf35ab61f137d1752f --- /dev/null +++ b/internal/fsext/expand.go @@ -0,0 +1,29 @@ +package fsext + +import ( + "os" + "strings" + + "mvdan.cc/sh/v3/expand" + "mvdan.cc/sh/v3/syntax" +) + +// Expand is a wrapper around [expand.Literal]. It will escape the input +// string, expand any shell symbols (such as '~') and resolve any environment +// variables. +func Expand(s string) (string, error) { + if s == "" { + return "", nil + } + p := syntax.NewParser() + word, err := p.Document(strings.NewReader(s)) + if err != nil { + return "", err + } + cfg := &expand.Config{ + Env: expand.FuncEnviron(os.Getenv), + ReadDir2: os.ReadDir, + GlobStar: true, + } + return expand.Literal(cfg, word) +} diff --git a/internal/fsext/fileutil.go b/internal/fsext/fileutil.go new file mode 100644 index 0000000000000000000000000000000000000000..c6b3073a7ef29fe93f003b11bded2ca2ffd58418 --- /dev/null +++ b/internal/fsext/fileutil.go @@ -0,0 +1,263 @@ +package fsext + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "slices" + "strings" + "time" + + "github.com/bmatcuk/doublestar/v4" + "github.com/charlievieth/fastwalk" + "github.com/charmbracelet/crush/internal/csync" + "github.com/charmbracelet/crush/internal/home" + "github.com/charmbracelet/x/ansi" +) + +type FileInfo struct { + Path string + ModTime time.Time +} + +func SkipHidden(path string) bool { + // Check for hidden files (starting with a dot) + base := filepath.Base(path) + if base != "." && strings.HasPrefix(base, ".") { + return true + } + + commonIgnoredDirs := map[string]bool{ + ".crush": true, + "node_modules": true, + "vendor": true, + "dist": true, + "build": true, + "target": true, + ".git": true, + ".idea": true, + ".vscode": true, + "__pycache__": true, + "bin": true, + "obj": true, + "out": true, + "coverage": true, + "logs": true, + "generated": true, + "bower_components": true, + "jspm_packages": true, + } + + parts := strings.SplitSeq(path, string(os.PathSeparator)) + for part := range parts { + if commonIgnoredDirs[part] { + return true + } + } + return false +} + +// FastGlobWalker provides gitignore-aware file walking with fastwalk +// It uses hierarchical ignore checking like git does, checking .gitignore/.crushignore +// files in each directory from the root to the target path. +type FastGlobWalker struct { + directoryLister *directoryLister +} + +func NewFastGlobWalker(searchPath string) *FastGlobWalker { + return &FastGlobWalker{ + directoryLister: NewDirectoryLister(searchPath), + } +} + +// ShouldSkip checks if a file path should be skipped based on hierarchical gitignore, +// crushignore, and hidden file rules. +func (w *FastGlobWalker) ShouldSkip(path string) bool { + return w.directoryLister.shouldIgnore(path, nil, false) +} + +// ShouldSkipDir checks if a directory path should be skipped based on hierarchical +// gitignore, crushignore, and hidden file rules. +func (w *FastGlobWalker) ShouldSkipDir(path string) bool { + return w.directoryLister.shouldIgnore(path, nil, true) +} + +// Glob globs files. +// +// Does not respect gitignore. +func Glob(pattern string, cwd string, limit int) ([]string, bool, error) { + return globWithDoubleStar(context.Background(), pattern, cwd, limit, false) +} + +// GlobGitignoreAware globs files respecting gitignore. +func GlobGitignoreAware(pattern string, cwd string, limit int) ([]string, bool, error) { + return globWithDoubleStar(context.Background(), pattern, cwd, limit, true) +} + +// GlobGitignoreAwareCtx is like [GlobGitignoreAware] but stops early when ctx +// is cancelled (e.g. on timeout), returning whatever was found so far. +func GlobGitignoreAwareCtx(ctx context.Context, pattern, cwd string, limit int) ([]string, bool, error) { + return globWithDoubleStar(ctx, pattern, cwd, limit, true) +} + +func globWithDoubleStar(ctx context.Context, pattern, searchPath string, limit int, gitignore bool) ([]string, bool, error) { + // Normalize pattern to forward slashes on Windows so their config can use + // backslashes + pattern = filepath.ToSlash(pattern) + + walker := NewFastGlobWalker(searchPath) + found := csync.NewSlice[FileInfo]() + conf := fastwalk.Config{ + // Do not follow symlinks: following them lets the walk escape the + // search root (into module caches, the nix store, $HOME, etc.) and + // chase cycles, which is slow and can hang. Mirrors the rg path, + // which no longer passes -L. + Follow: false, + ToSlash: fastwalk.DefaultToSlash(), + Sort: fastwalk.SortFilesFirst, + } + err := fastwalk.Walk(&conf, searchPath, func(path string, d os.DirEntry, err error) error { + if ctx.Err() != nil { + return filepath.SkipAll // Timed out or cancelled; stop walking. + } + if err != nil { + return nil // Skip files we can't access + } + + isDir := d.IsDir() + if isDir { + if gitignore && walker.ShouldSkipDir(path) { + return filepath.SkipDir + } + } else { + if gitignore && walker.ShouldSkip(path) { + return nil + } + } + + relPath, err := filepath.Rel(searchPath, path) + if err != nil { + relPath = path + } + + // Normalize separators to forward slashes + relPath = filepath.ToSlash(relPath) + + // Check if path matches the pattern + matched, err := doublestar.Match(pattern, relPath) + if err != nil || !matched { + return nil + } + + info, err := d.Info() + if err != nil { + return nil + } + + found.Append(FileInfo{Path: path, ModTime: info.ModTime()}) + if limit > 0 && found.Len() >= limit*2 { // NOTE: why x2? + return filepath.SkipAll + } + return nil + }) + if err != nil && !errors.Is(err, filepath.SkipAll) { + return nil, false, fmt.Errorf("fastwalk error: %w", err) + } + + matches := slices.SortedFunc(found.Seq(), func(a, b FileInfo) int { + return b.ModTime.Compare(a.ModTime) + }) + matches, truncated := truncate(matches, limit) + + results := make([]string, len(matches)) + for i, m := range matches { + results[i] = m.Path + } + return results, truncated || errors.Is(err, filepath.SkipAll), nil +} + +// ShouldExcludeFile checks if a file should be excluded from processing +// based on common patterns and ignore rules. +func ShouldExcludeFile(rootPath, filePath string) bool { + info, err := os.Stat(filePath) + isDir := err == nil && info.IsDir() + return NewDirectoryLister(rootPath). + shouldIgnore(filePath, nil, isDir) +} + +func PrettyPath(path string) string { + return home.Short(path) +} + +func DirTrim(pwd string, lim int) string { + var ( + out string + sep = string(filepath.Separator) + ) + dirs := strings.Split(pwd, sep) + if lim > len(dirs)-1 || lim <= 0 { + return pwd + } + for i := len(dirs) - 1; i > 0; i-- { + out = sep + out + if i == len(dirs)-1 { + out = dirs[i] + } else if i >= len(dirs)-lim { + // Keep the first grapheme cluster, not the first byte: CJK, + // combining marks, and emoji can span multiple bytes and runes, + // so a byte or single rune would render the wrong character. + first, _ := ansi.FirstGraphemeCluster(dirs[i], ansi.GraphemeWidth) + out = first + out + } else { + out = "..." + out + break + } + } + out = filepath.Join("~", out) + return out +} + +// PathOrPrefix returns the prefix if the path starts with it, or falls back to +// the path otherwise. +func PathOrPrefix(path, prefix string) string { + if HasPrefix(path, prefix) { + return prefix + } + return path +} + +// HasPrefix checks if the given path starts with the specified prefix. +// Uses filepath.Rel to determine if path is within prefix. +func HasPrefix(path, prefix string) bool { + rel, err := filepath.Rel(prefix, path) + if err != nil { + return false + } + // If path is within prefix, Rel will not return a path starting with ".." + return !strings.HasPrefix(rel, "..") +} + +// ToUnixLineEndings converts Windows line endings (CRLF) to Unix line endings (LF). +func ToUnixLineEndings(content string) (string, bool) { + if strings.Contains(content, "\r\n") { + return strings.ReplaceAll(content, "\r\n", "\n"), true + } + return content, false +} + +// ToWindowsLineEndings converts Unix line endings (LF) to Windows line endings (CRLF). +func ToWindowsLineEndings(content string) (string, bool) { + if !strings.Contains(content, "\r\n") { + return strings.ReplaceAll(content, "\n", "\r\n"), true + } + return content, false +} + +func truncate[T any](input []T, limit int) ([]T, bool) { + if limit > 0 && len(input) > limit { + return input[:limit], true + } + return input, false +} diff --git a/internal/fsext/fileutil_dirtrim_test.go b/internal/fsext/fileutil_dirtrim_test.go new file mode 100644 index 0000000000000000000000000000000000000000..2a6e7b5a09d2320fde1ff2bcebb5413f127b7b5b --- /dev/null +++ b/internal/fsext/fileutil_dirtrim_test.go @@ -0,0 +1,35 @@ +package fsext + +import ( + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestDirTrimUnicode(t *testing.T) { + sep := string(filepath.Separator) + + // An abbreviated component keeps its first character. + require.Equal( + t, + filepath.Join("~", "...", "p", "file"), + DirTrim(sep+"home"+sep+"user"+sep+"project"+sep+"file", 2), + ) + + // A multi-byte (e.g. CJK) component must keep its first whole character, + // not its first byte, which would render a wrong character. + require.Equal( + t, + filepath.Join("~", "...", "项", "file"), + DirTrim(sep+"home"+sep+"项目"+sep+"file", 2), + ) + + // A grapheme cluster spanning multiple runes (a base letter plus a + // combining accent) must be kept whole, not split after the base rune. + require.Equal( + t, + filepath.Join("~", "...", "e\u0301", "file"), + DirTrim(sep+"home"+sep+"e\u0301cole"+sep+"file", 2), + ) +} diff --git a/internal/fsext/fileutil_test.go b/internal/fsext/fileutil_test.go new file mode 100644 index 0000000000000000000000000000000000000000..e80b902c08301be3a6b02e9a8d0e9c82895410b3 --- /dev/null +++ b/internal/fsext/fileutil_test.go @@ -0,0 +1,269 @@ +package fsext + +import ( + "fmt" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestGlobWithDoubleStar(t *testing.T) { + t.Run("finds files matching pattern", func(t *testing.T) { + testDir := t.TempDir() + + mainGo := filepath.Join(testDir, "src", "main.go") + utilsGo := filepath.Join(testDir, "src", "utils.go") + helperGo := filepath.Join(testDir, "pkg", "helper.go") + readmeMd := filepath.Join(testDir, "README.md") + + for _, file := range []string{mainGo, utilsGo, helperGo, readmeMd} { + require.NoError(t, os.MkdirAll(filepath.Dir(file), 0o755)) + require.NoError(t, os.WriteFile(file, []byte("test content"), 0o644)) + } + + matches, truncated, err := GlobGitignoreAware("**/main.go", testDir, 0) + require.NoError(t, err) + require.False(t, truncated) + + require.Equal(t, matches, []string{mainGo}) + }) + + t.Run("finds directories matching pattern", func(t *testing.T) { + testDir := t.TempDir() + + srcDir := filepath.Join(testDir, "src") + pkgDir := filepath.Join(testDir, "pkg") + internalDir := filepath.Join(testDir, "internal") + cmdDir := filepath.Join(testDir, "cmd") + pkgFile := filepath.Join(testDir, "pkg.txt") + + for _, dir := range []string{srcDir, pkgDir, internalDir, cmdDir} { + require.NoError(t, os.MkdirAll(dir, 0o755)) + } + + require.NoError(t, os.WriteFile(filepath.Join(srcDir, "main.go"), []byte("package main"), 0o644)) + require.NoError(t, os.WriteFile(pkgFile, []byte("test"), 0o644)) + + matches, truncated, err := GlobGitignoreAware("pkg", testDir, 0) + require.NoError(t, err) + require.False(t, truncated) + + require.Equal(t, matches, []string{pkgDir}) + }) + + t.Run("finds nested directories with wildcard patterns", func(t *testing.T) { + testDir := t.TempDir() + + srcPkgDir := filepath.Join(testDir, "src", "pkg") + libPkgDir := filepath.Join(testDir, "lib", "pkg") + mainPkgDir := filepath.Join(testDir, "pkg") + otherDir := filepath.Join(testDir, "other") + + for _, dir := range []string{srcPkgDir, libPkgDir, mainPkgDir, otherDir} { + require.NoError(t, os.MkdirAll(dir, 0o755)) + } + + matches, truncated, err := GlobGitignoreAware("**/pkg", testDir, 0) + require.NoError(t, err) + require.False(t, truncated) + + var relativeMatches []string + for _, match := range matches { + rel, err := filepath.Rel(testDir, match) + require.NoError(t, err) + relativeMatches = append(relativeMatches, filepath.ToSlash(rel)) + } + + require.ElementsMatch(t, relativeMatches, []string{"pkg", "src/pkg", "lib/pkg"}) + }) + + t.Run("finds directory contents with recursive patterns", func(t *testing.T) { + testDir := t.TempDir() + + pkgDir := filepath.Join(testDir, "pkg") + pkgFile1 := filepath.Join(pkgDir, "main.go") + pkgFile2 := filepath.Join(pkgDir, "utils.go") + pkgSubdir := filepath.Join(pkgDir, "internal") + pkgSubfile := filepath.Join(pkgSubdir, "helper.go") + + require.NoError(t, os.MkdirAll(pkgSubdir, 0o755)) + + for _, file := range []string{pkgFile1, pkgFile2, pkgSubfile} { + require.NoError(t, os.WriteFile(file, []byte("package main"), 0o644)) + } + + matches, truncated, err := GlobGitignoreAware("pkg/**", testDir, 0) + require.NoError(t, err) + require.False(t, truncated) + + var relativeMatches []string + for _, match := range matches { + rel, err := filepath.Rel(testDir, match) + require.NoError(t, err) + relativeMatches = append(relativeMatches, filepath.ToSlash(rel)) + } + + require.ElementsMatch(t, relativeMatches, []string{ + "pkg", + "pkg/main.go", + "pkg/utils.go", + "pkg/internal", + "pkg/internal/helper.go", + }) + }) + + t.Run("respects limit parameter", func(t *testing.T) { + testDir := t.TempDir() + + for i := range 10 { + file := filepath.Join(testDir, "file", fmt.Sprintf("test%d.txt", i)) + require.NoError(t, os.MkdirAll(filepath.Dir(file), 0o755)) + require.NoError(t, os.WriteFile(file, []byte("test"), 0o644)) + } + + matches, truncated, err := GlobGitignoreAware("**/*.txt", testDir, 5) + require.NoError(t, err) + require.True(t, truncated, "Expected truncation with limit") + require.Len(t, matches, 5, "Expected exactly 5 matches with limit") + }) + + t.Run("handles nested directory patterns", func(t *testing.T) { + testDir := t.TempDir() + + file1 := filepath.Join(testDir, "a", "b", "c", "file1.txt") + file2 := filepath.Join(testDir, "a", "b", "file2.txt") + file3 := filepath.Join(testDir, "a", "file3.txt") + file4 := filepath.Join(testDir, "file4.txt") + + for _, file := range []string{file1, file2, file3, file4} { + require.NoError(t, os.MkdirAll(filepath.Dir(file), 0o755)) + require.NoError(t, os.WriteFile(file, []byte("test"), 0o644)) + } + + matches, truncated, err := GlobGitignoreAware("a/b/c/file1.txt", testDir, 0) + require.NoError(t, err) + require.False(t, truncated) + + require.Equal(t, []string{file1}, matches) + }) + + t.Run("returns results sorted by modification time (newest first)", func(t *testing.T) { + testDir := t.TempDir() + + file1 := filepath.Join(testDir, "file1.txt") + require.NoError(t, os.WriteFile(file1, []byte("first"), 0o644)) + + file2 := filepath.Join(testDir, "file2.txt") + require.NoError(t, os.WriteFile(file2, []byte("second"), 0o644)) + + file3 := filepath.Join(testDir, "file3.txt") + require.NoError(t, os.WriteFile(file3, []byte("third"), 0o644)) + + base := time.Now() + m1 := base + m2 := base.Add(10 * time.Hour) + m3 := base.Add(20 * time.Hour) + + require.NoError(t, os.Chtimes(file1, m1, m1)) + require.NoError(t, os.Chtimes(file2, m2, m2)) + require.NoError(t, os.Chtimes(file3, m3, m3)) + + matches, truncated, err := GlobGitignoreAware("*.txt", testDir, 0) + require.NoError(t, err) + require.False(t, truncated) + + require.Equal(t, []string{file3, file2, file1}, matches) + }) + + t.Run("handles empty directory", func(t *testing.T) { + testDir := t.TempDir() + + matches, truncated, err := GlobGitignoreAware("**", testDir, 0) + require.NoError(t, err) + require.False(t, truncated) + // Even empty directories should return the directory itself + require.Equal(t, []string{testDir}, matches) + }) + + t.Run("handles non-existent search path", func(t *testing.T) { + nonExistentDir := filepath.Join(t.TempDir(), "does", "not", "exist") + + matches, truncated, err := GlobGitignoreAware("**", nonExistentDir, 0) + require.Error(t, err, "Should return error for non-existent search path") + require.False(t, truncated) + require.Empty(t, matches) + }) + + t.Run("respects basic ignore patterns", func(t *testing.T) { + testDir := t.TempDir() + + rootIgnore := filepath.Join(testDir, ".crushignore") + + require.NoError(t, os.WriteFile(rootIgnore, []byte("*.tmp\nbackup/\n"), 0o644)) + + goodFile := filepath.Join(testDir, "good.txt") + require.NoError(t, os.WriteFile(goodFile, []byte("content"), 0o644)) + + badFile := filepath.Join(testDir, "bad.tmp") + require.NoError(t, os.WriteFile(badFile, []byte("temp content"), 0o644)) + + goodDir := filepath.Join(testDir, "src") + require.NoError(t, os.MkdirAll(goodDir, 0o755)) + + ignoredDir := filepath.Join(testDir, "backup") + require.NoError(t, os.MkdirAll(ignoredDir, 0o755)) + + ignoredFileInDir := filepath.Join(testDir, "backup", "old.txt") + require.NoError(t, os.WriteFile(ignoredFileInDir, []byte("old content"), 0o644)) + + matches, truncated, err := GlobGitignoreAware("*.tmp", testDir, 0) + require.NoError(t, err) + require.False(t, truncated) + require.Empty(t, matches, "Expected no matches for '*.tmp' pattern (should be ignored)") + + matches, truncated, err = GlobGitignoreAware("backup", testDir, 0) + require.NoError(t, err) + require.False(t, truncated) + require.Empty(t, matches, "Expected no matches for 'backup' pattern (should be ignored)") + + matches, truncated, err = GlobGitignoreAware("*.txt", testDir, 0) + require.NoError(t, err) + require.False(t, truncated) + require.Equal(t, []string{goodFile}, matches) + }) + + t.Run("handles mixed file and directory matching with sorting", func(t *testing.T) { + testDir := t.TempDir() + + oldestFile := filepath.Join(testDir, "old.rs") + require.NoError(t, os.WriteFile(oldestFile, []byte("old"), 0o644)) + + middleDir := filepath.Join(testDir, "mid.rs") + require.NoError(t, os.MkdirAll(middleDir, 0o755)) + + newestFile := filepath.Join(testDir, "new.rs") + require.NoError(t, os.WriteFile(newestFile, []byte("new"), 0o644)) + + base := time.Now() + tOldest := base + tMiddle := base.Add(10 * time.Hour) + tNewest := base.Add(20 * time.Hour) + + // Reverse the expected order + require.NoError(t, os.Chtimes(newestFile, tOldest, tOldest)) + require.NoError(t, os.Chtimes(middleDir, tMiddle, tMiddle)) + require.NoError(t, os.Chtimes(oldestFile, tNewest, tNewest)) + + matches, truncated, err := GlobGitignoreAware("*.rs", testDir, 0) + require.NoError(t, err) + require.False(t, truncated) + require.Len(t, matches, 3) + + // Results should be sorted by mod time, but we set the oldestFile + // to have the most recent mod time + require.Equal(t, []string{oldestFile, middleDir, newestFile}, matches) + }) +} diff --git a/internal/fsext/ignore_test.go b/internal/fsext/ignore_test.go new file mode 100644 index 0000000000000000000000000000000000000000..e5e34e85a4c9bf0b207703aa0cf2f0459e03a427 --- /dev/null +++ b/internal/fsext/ignore_test.go @@ -0,0 +1,122 @@ +package fsext + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestCrushIgnore(t *testing.T) { + tempDir := t.TempDir() + t.Chdir(tempDir) + + // Create test files + require.NoError(t, os.WriteFile("test1.txt", []byte("test"), 0o644)) + require.NoError(t, os.WriteFile("test2.log", []byte("test"), 0o644)) + require.NoError(t, os.WriteFile("test3.tmp", []byte("test"), 0o644)) + + // Create a .crushignore file that ignores .log files + require.NoError(t, os.WriteFile(".crushignore", []byte("*.log\n"), 0o644)) + + dl := NewDirectoryLister(tempDir) + require.True(t, dl.shouldIgnore("test2.log", nil, false), ".log files should be ignored") + require.False(t, dl.shouldIgnore("test1.txt", nil, false), ".txt files should not be ignored") + require.True(t, dl.shouldIgnore("test3.tmp", nil, false), ".tmp files should be ignored by common patterns") +} + +func TestShouldExcludeFile(t *testing.T) { + t.Parallel() + + // Create a temporary directory structure for testing + tempDir := t.TempDir() + + // Create directories that should be ignored + nodeModules := filepath.Join(tempDir, "node_modules") + target := filepath.Join(tempDir, "target") + customIgnored := filepath.Join(tempDir, "custom_ignored") + normalDir := filepath.Join(tempDir, "src") + + for _, dir := range []string{nodeModules, target, customIgnored, normalDir} { + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatalf("Failed to create directory %s: %v", dir, err) + } + } + + // Create .gitignore file + gitignoreContent := "node_modules/\ntarget/\n" + if err := os.WriteFile(filepath.Join(tempDir, ".gitignore"), []byte(gitignoreContent), 0o644); err != nil { + t.Fatalf("Failed to create .gitignore: %v", err) + } + + // Create .crushignore file + crushignoreContent := "custom_ignored/\n" + if err := os.WriteFile(filepath.Join(tempDir, ".crushignore"), []byte(crushignoreContent), 0o644); err != nil { + t.Fatalf("Failed to create .crushignore: %v", err) + } + + // Test that ignored directories are properly ignored + require.True(t, ShouldExcludeFile(tempDir, nodeModules), "Expected node_modules to be ignored by .gitignore") + require.True(t, ShouldExcludeFile(tempDir, target), "Expected target to be ignored by .gitignore") + require.True(t, ShouldExcludeFile(tempDir, customIgnored), "Expected custom_ignored to be ignored by .crushignore") + + // Test that normal directories are not ignored + require.False(t, ShouldExcludeFile(tempDir, normalDir), "Expected src directory to not be ignored") + + // Test that the workspace root itself is not ignored + require.False(t, ShouldExcludeFile(tempDir, tempDir), "Expected workspace root to not be ignored") +} + +func TestShouldExcludeFileHierarchical(t *testing.T) { + t.Parallel() + + // Create a nested directory structure for testing hierarchical ignore + tempDir := t.TempDir() + + // Create nested directories + subDir := filepath.Join(tempDir, "subdir") + nestedNormal := filepath.Join(subDir, "normal_nested") + + for _, dir := range []string{subDir, nestedNormal} { + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatalf("Failed to create directory %s: %v", dir, err) + } + } + + // Create .crushignore in subdir that ignores normal_nested + subCrushignore := "normal_nested/\n" + if err := os.WriteFile(filepath.Join(subDir, ".crushignore"), []byte(subCrushignore), 0o644); err != nil { + t.Fatalf("Failed to create subdir .crushignore: %v", err) + } + + // Test hierarchical ignore behavior - this should work because the .crushignore is in the parent directory + require.True(t, ShouldExcludeFile(tempDir, nestedNormal), "Expected normal_nested to be ignored by subdir .crushignore") + require.False(t, ShouldExcludeFile(tempDir, subDir), "Expected subdir itself to not be ignored") +} + +func TestShouldExcludeFileCommonPatterns(t *testing.T) { + t.Parallel() + + tempDir := t.TempDir() + + // Create directories that should be ignored by common patterns + commonIgnored := []string{ + filepath.Join(tempDir, ".git"), + filepath.Join(tempDir, "node_modules"), + filepath.Join(tempDir, "__pycache__"), + filepath.Join(tempDir, "target"), + filepath.Join(tempDir, ".vscode"), + } + + for _, dir := range commonIgnored { + if err := os.MkdirAll(dir, 0o755); err != nil { + t.Fatalf("Failed to create directory %s: %v", dir, err) + } + } + + // Test that common patterns are ignored even without explicit ignore files + for _, dir := range commonIgnored { + require.True(t, ShouldExcludeFile(tempDir, dir), "Expected %s to be ignored by common patterns", filepath.Base(dir)) + } +} diff --git a/internal/fsext/lookup.go b/internal/fsext/lookup.go new file mode 100644 index 0000000000000000000000000000000000000000..696219ac89209b718ed0e106a4bc6a879105961c --- /dev/null +++ b/internal/fsext/lookup.go @@ -0,0 +1,281 @@ +package fsext + +import ( + "errors" + "fmt" + "os" + "path/filepath" + + "github.com/charmbracelet/crush/internal/home" +) + +// Lookup searches for a target files or directories starting from dir +// and walking up the directory tree until filesystem root is reached. +// It also checks the ownership of files to ensure that the search does +// not cross ownership boundaries. It skips ownership mismatches without +// errors. +// Returns full paths to fount targets. +// The search includes the starting directory itself. +func Lookup(dir string, targets ...string) ([]string, error) { + if len(targets) == 0 { + return nil, nil + } + + var found []string + + err := traverseUp(dir, func(cwd string, owner int) error { + for _, target := range targets { + fpath := filepath.Join(cwd, target) + err := probeEnt(fpath, owner) + + // skip to the next file on permission denied + if errors.Is(err, os.ErrNotExist) || + errors.Is(err, os.ErrPermission) { + continue + } + + if err != nil { + return fmt.Errorf("error probing file %s: %w", fpath, err) + } + + found = append(found, fpath) + } + + return nil + }) + if err != nil { + return nil, err + } + + return found, nil +} + +// LookupClosest searches for a target file or directory starting from dir +// and walking up the directory tree until found or root or home is reached. +// It also checks the ownership of files to ensure that the search does +// not cross ownership boundaries. +// Returns the full path to the target if found, empty string and false otherwise. +// The search includes the starting directory itself. +func LookupClosest(dir, target string) (string, bool) { + var found string + + err := traverseUp(dir, func(cwd string, owner int) error { + fpath := filepath.Join(cwd, target) + + err := probeEnt(fpath, owner) + if errors.Is(err, os.ErrNotExist) { + return nil + } + + if err != nil { + return fmt.Errorf("error probing file %s: %w", fpath, err) + } + + if cwd == home.Dir() { + return filepath.SkipAll + } + + found = fpath + return filepath.SkipAll + }) + + return found, err == nil && found != "" +} + +// LookupClosestBounded behaves like LookupClosest but constrains the +// upward search to stopDir. The walk inspects dir, then each ancestor up +// to and including stopDir, then terminates regardless of whether the +// target was found. Use this when the caller wants to avoid adopting +// matches from outside a project boundary (for example a sibling +// worktree or a parent project). +// +// If stopDir is empty, only dir itself is searched. If stopDir is not an +// ancestor of dir, the walk still terminates at the filesystem root. +// The $HOME and ownership safeguards from LookupClosest are preserved +// as outer bounds. +func LookupClosestBounded(dir, stopDir, target string) (string, bool) { + var found string + + err := traverseUpBounded(dir, stopDir, func(cwd string, owner int) error { + fpath := filepath.Join(cwd, target) + + err := probeEnt(fpath, owner) + if errors.Is(err, os.ErrNotExist) { + return nil + } + + if err != nil { + return fmt.Errorf("error probing file %s: %w", fpath, err) + } + + if cwd == home.Dir() { + return filepath.SkipAll + } + + found = fpath + return filepath.SkipAll + }) + + return found, err == nil && found != "" +} + +// LookupBounded behaves like Lookup but constrains the upward search to +// stopDir. The walk inspects dir, then each ancestor up to and including +// stopDir, then terminates. If stopDir is empty, only dir itself is +// searched. +func LookupBounded(dir, stopDir string, targets ...string) ([]string, error) { + if len(targets) == 0 { + return nil, nil + } + + var found []string + + err := traverseUpBounded(dir, stopDir, func(cwd string, owner int) error { + for _, target := range targets { + fpath := filepath.Join(cwd, target) + err := probeEnt(fpath, owner) + + // skip to the next file on permission denied + if errors.Is(err, os.ErrNotExist) || + errors.Is(err, os.ErrPermission) { + continue + } + + if err != nil { + return fmt.Errorf("error probing file %s: %w", fpath, err) + } + + found = append(found, fpath) + } + + return nil + }) + if err != nil { + return nil, err + } + + return found, nil +} + +// traverseUp walks up from given directory up until filesystem root reached. +// It passes absolute path of current directory and staring directory owner ID +// to callback function. It is up to user to check ownership. +func traverseUp(dir string, walkFn func(dir string, owner int) error) error { + cwd, err := filepath.Abs(dir) + if err != nil { + return fmt.Errorf("cannot convert CWD to absolute path: %w", err) + } + + owner, err := Owner(dir) + if err != nil { + return fmt.Errorf("cannot get ownership: %w", err) + } + + for { + err := walkFn(cwd, owner) + if err == nil || errors.Is(err, filepath.SkipDir) { + parent := filepath.Dir(cwd) + if parent == cwd { + return nil + } + + cwd = parent + continue + } + + if errors.Is(err, filepath.SkipAll) { + return nil + } + + return err + } +} + +// traverseUpBounded walks up from dir, visiting each ancestor up to and +// including stopDir, then terminates. If stopDir is empty, only dir +// itself is visited; callers that want an unbounded walk should use +// traverseUp instead. If stopDir is set but is not an ancestor of dir +// the walk still stops at the filesystem root, so callers cannot +// accidentally produce an infinite walk by passing a sibling path. +// +// Boundary comparison is performed against symlink-resolved paths so +// that callers passing logically equivalent paths (a symlinked /var vs +// the underlying /private/var, for example) still terminate at the +// expected directory. +func traverseUpBounded(dir, stopDir string, walkFn func(dir string, owner int) error) error { + cwd, err := filepath.Abs(dir) + if err != nil { + return fmt.Errorf("cannot convert CWD to absolute path: %w", err) + } + + stop := cwd + if stopDir != "" { + stop, err = filepath.Abs(stopDir) + if err != nil { + return fmt.Errorf("cannot convert stop dir to absolute path: %w", err) + } + } + canonStop := canonicalize(stop) + + owner, err := Owner(dir) + if err != nil { + return fmt.Errorf("cannot get ownership: %w", err) + } + + for { + err := walkFn(cwd, owner) + if err == nil || errors.Is(err, filepath.SkipDir) { + if canonicalize(cwd) == canonStop { + return nil + } + + parent := filepath.Dir(cwd) + if parent == cwd { + return nil + } + + cwd = parent + continue + } + + if errors.Is(err, filepath.SkipAll) { + return nil + } + + return err + } +} + +// canonicalize resolves any symbolic links in path. If resolution fails +// (typically because path does not exist yet) the original path is +// returned cleaned, so callers can still perform stable equality checks. +func canonicalize(path string) string { + if resolved, err := filepath.EvalSymlinks(path); err == nil { + return resolved + } + return filepath.Clean(path) +} + +// probeEnt checks if entity at given path exists and belongs to given owner +func probeEnt(fspath string, owner int) error { + _, err := os.Stat(fspath) + if err != nil { + return fmt.Errorf("cannot stat %s: %w", fspath, err) + } + + // special case for ownership check bypass + if owner == -1 { + return nil + } + + fowner, err := Owner(fspath) + if err != nil { + return fmt.Errorf("cannot get ownership for %s: %w", fspath, err) + } + + if fowner != owner { + return os.ErrPermission + } + + return nil +} diff --git a/internal/fsext/lookup_test.go b/internal/fsext/lookup_test.go new file mode 100644 index 0000000000000000000000000000000000000000..f4e445c179459dc0deed85708014a6690d4a008d --- /dev/null +++ b/internal/fsext/lookup_test.go @@ -0,0 +1,562 @@ +package fsext + +import ( + "errors" + "os" + "path/filepath" + "testing" + + "github.com/charmbracelet/crush/internal/home" + "github.com/stretchr/testify/require" +) + +func TestLookupClosest(t *testing.T) { + tempDir := t.TempDir() + t.Chdir(tempDir) + + t.Run("target found in starting directory", func(t *testing.T) { + testDir := t.TempDir() + + // Create target file in current directory + targetFile := filepath.Join(testDir, "target.txt") + err := os.WriteFile(targetFile, []byte("test"), 0o644) + require.NoError(t, err) + + foundPath, found := LookupClosest(testDir, "target.txt") + require.True(t, found) + require.Equal(t, targetFile, foundPath) + }) + + t.Run("target found in parent directory", func(t *testing.T) { + testDir := t.TempDir() + + // Create subdirectory + subDir := filepath.Join(testDir, "subdir") + err := os.Mkdir(subDir, 0o755) + require.NoError(t, err) + + // Create target file in parent directory + targetFile := filepath.Join(testDir, "target.txt") + err = os.WriteFile(targetFile, []byte("test"), 0o644) + require.NoError(t, err) + + foundPath, found := LookupClosest(subDir, "target.txt") + require.True(t, found) + require.Equal(t, targetFile, foundPath) + }) + + t.Run("target found in grandparent directory", func(t *testing.T) { + testDir := t.TempDir() + + // Create nested subdirectories + subDir := filepath.Join(testDir, "subdir") + err := os.Mkdir(subDir, 0o755) + require.NoError(t, err) + + subSubDir := filepath.Join(subDir, "subsubdir") + err = os.Mkdir(subSubDir, 0o755) + require.NoError(t, err) + + // Create target file in grandparent directory + targetFile := filepath.Join(testDir, "target.txt") + err = os.WriteFile(targetFile, []byte("test"), 0o644) + require.NoError(t, err) + + foundPath, found := LookupClosest(subSubDir, "target.txt") + require.True(t, found) + require.Equal(t, targetFile, foundPath) + }) + + t.Run("target not found", func(t *testing.T) { + testDir := t.TempDir() + + foundPath, found := LookupClosest(testDir, "nonexistent.txt") + require.False(t, found) + require.Empty(t, foundPath) + }) + + t.Run("target directory found", func(t *testing.T) { + testDir := t.TempDir() + + // Create target directory in current directory + targetDir := filepath.Join(testDir, "targetdir") + err := os.Mkdir(targetDir, 0o755) + require.NoError(t, err) + + foundPath, found := LookupClosest(testDir, "targetdir") + require.True(t, found) + require.Equal(t, targetDir, foundPath) + }) + + t.Run("stops at home directory", func(t *testing.T) { + // This test is limited as we can't easily create files above home directory + // but we can test the behavior by searching from home directory itself + homeDir := home.Dir() + + // Search for a file that doesn't exist from home directory + foundPath, found := LookupClosest(homeDir, "nonexistent_file_12345.txt") + require.False(t, found) + require.Empty(t, foundPath) + }) + + t.Run("invalid starting directory", func(t *testing.T) { + foundPath, found := LookupClosest("/invalid/path/that/does/not/exist", "target.txt") + require.False(t, found) + require.Empty(t, foundPath) + }) + + t.Run("relative path handling", func(t *testing.T) { + // Create target file in current directory + require.NoError(t, os.WriteFile("target.txt", []byte("test"), 0o644)) + + // Search using relative path + foundPath, found := LookupClosest(".", "target.txt") + require.True(t, found) + + // Resolve symlinks to handle macOS /private/var vs /var discrepancy + expectedPath, err := filepath.EvalSymlinks(filepath.Join(tempDir, "target.txt")) + require.NoError(t, err) + actualPath, err := filepath.EvalSymlinks(foundPath) + require.NoError(t, err) + require.Equal(t, expectedPath, actualPath) + }) +} + +func TestLookupClosestWithOwnership(t *testing.T) { + // Note: Testing ownership boundaries is difficult in a cross-platform way + // without creating complex directory structures with different owners. + // This test focuses on the basic functionality when ownership checks pass. + + tempDir := t.TempDir() + t.Chdir(tempDir) + + t.Run("search respects same ownership", func(t *testing.T) { + testDir := t.TempDir() + + // Create subdirectory structure + subDir := filepath.Join(testDir, "subdir") + err := os.Mkdir(subDir, 0o755) + require.NoError(t, err) + + // Create target file in parent directory + targetFile := filepath.Join(testDir, "target.txt") + err = os.WriteFile(targetFile, []byte("test"), 0o644) + require.NoError(t, err) + + // Search should find the target assuming same ownership + foundPath, found := LookupClosest(subDir, "target.txt") + require.True(t, found) + require.Equal(t, targetFile, foundPath) + }) +} + +func TestLookup(t *testing.T) { + tempDir := t.TempDir() + t.Chdir(tempDir) + + t.Run("no targets returns empty slice", func(t *testing.T) { + testDir := t.TempDir() + + found, err := Lookup(testDir) + require.NoError(t, err) + require.Empty(t, found) + }) + + t.Run("single target found in starting directory", func(t *testing.T) { + testDir := t.TempDir() + + // Create target file in current directory + targetFile := filepath.Join(testDir, "target.txt") + err := os.WriteFile(targetFile, []byte("test"), 0o644) + require.NoError(t, err) + + found, err := Lookup(testDir, "target.txt") + require.NoError(t, err) + require.Len(t, found, 1) + require.Equal(t, targetFile, found[0]) + }) + + t.Run("multiple targets found in starting directory", func(t *testing.T) { + testDir := t.TempDir() + + // Create multiple target files in current directory + targetFile1 := filepath.Join(testDir, "target1.txt") + targetFile2 := filepath.Join(testDir, "target2.txt") + targetFile3 := filepath.Join(testDir, "target3.txt") + + err := os.WriteFile(targetFile1, []byte("test1"), 0o644) + require.NoError(t, err) + err = os.WriteFile(targetFile2, []byte("test2"), 0o644) + require.NoError(t, err) + err = os.WriteFile(targetFile3, []byte("test3"), 0o644) + require.NoError(t, err) + + found, err := Lookup(testDir, "target1.txt", "target2.txt", "target3.txt") + require.NoError(t, err) + require.Len(t, found, 3) + require.Contains(t, found, targetFile1) + require.Contains(t, found, targetFile2) + require.Contains(t, found, targetFile3) + }) + + t.Run("targets found in parent directories", func(t *testing.T) { + testDir := t.TempDir() + + // Create subdirectory + subDir := filepath.Join(testDir, "subdir") + err := os.Mkdir(subDir, 0o755) + require.NoError(t, err) + + // Create target files in parent directory + targetFile1 := filepath.Join(testDir, "target1.txt") + targetFile2 := filepath.Join(testDir, "target2.txt") + err = os.WriteFile(targetFile1, []byte("test1"), 0o644) + require.NoError(t, err) + err = os.WriteFile(targetFile2, []byte("test2"), 0o644) + require.NoError(t, err) + + found, err := Lookup(subDir, "target1.txt", "target2.txt") + require.NoError(t, err) + require.Len(t, found, 2) + require.Contains(t, found, targetFile1) + require.Contains(t, found, targetFile2) + }) + + t.Run("targets found across multiple directory levels", func(t *testing.T) { + testDir := t.TempDir() + + // Create nested subdirectories + subDir := filepath.Join(testDir, "subdir") + err := os.Mkdir(subDir, 0o755) + require.NoError(t, err) + + subSubDir := filepath.Join(subDir, "subsubdir") + err = os.Mkdir(subSubDir, 0o755) + require.NoError(t, err) + + // Create target files at different levels + targetFile1 := filepath.Join(testDir, "target1.txt") + targetFile2 := filepath.Join(subDir, "target2.txt") + targetFile3 := filepath.Join(subSubDir, "target3.txt") + + err = os.WriteFile(targetFile1, []byte("test1"), 0o644) + require.NoError(t, err) + err = os.WriteFile(targetFile2, []byte("test2"), 0o644) + require.NoError(t, err) + err = os.WriteFile(targetFile3, []byte("test3"), 0o644) + require.NoError(t, err) + + found, err := Lookup(subSubDir, "target1.txt", "target2.txt", "target3.txt") + require.NoError(t, err) + require.Len(t, found, 3) + require.Contains(t, found, targetFile1) + require.Contains(t, found, targetFile2) + require.Contains(t, found, targetFile3) + }) + + t.Run("some targets not found", func(t *testing.T) { + testDir := t.TempDir() + + // Create only some target files + targetFile1 := filepath.Join(testDir, "target1.txt") + targetFile2 := filepath.Join(testDir, "target2.txt") + + err := os.WriteFile(targetFile1, []byte("test1"), 0o644) + require.NoError(t, err) + err = os.WriteFile(targetFile2, []byte("test2"), 0o644) + require.NoError(t, err) + + // Search for existing and non-existing targets + found, err := Lookup(testDir, "target1.txt", "nonexistent.txt", "target2.txt", "another_nonexistent.txt") + require.NoError(t, err) + require.Len(t, found, 2) + require.Contains(t, found, targetFile1) + require.Contains(t, found, targetFile2) + }) + + t.Run("no targets found", func(t *testing.T) { + testDir := t.TempDir() + + found, err := Lookup(testDir, "nonexistent1.txt", "nonexistent2.txt", "nonexistent3.txt") + require.NoError(t, err) + require.Empty(t, found) + }) + + t.Run("target directories found", func(t *testing.T) { + testDir := t.TempDir() + + // Create target directories + targetDir1 := filepath.Join(testDir, "targetdir1") + targetDir2 := filepath.Join(testDir, "targetdir2") + err := os.Mkdir(targetDir1, 0o755) + require.NoError(t, err) + err = os.Mkdir(targetDir2, 0o755) + require.NoError(t, err) + + found, err := Lookup(testDir, "targetdir1", "targetdir2") + require.NoError(t, err) + require.Len(t, found, 2) + require.Contains(t, found, targetDir1) + require.Contains(t, found, targetDir2) + }) + + t.Run("mixed files and directories", func(t *testing.T) { + testDir := t.TempDir() + + // Create target files and directories + targetFile := filepath.Join(testDir, "target.txt") + targetDir := filepath.Join(testDir, "targetdir") + err := os.WriteFile(targetFile, []byte("test"), 0o644) + require.NoError(t, err) + err = os.Mkdir(targetDir, 0o755) + require.NoError(t, err) + + found, err := Lookup(testDir, "target.txt", "targetdir") + require.NoError(t, err) + require.Len(t, found, 2) + require.Contains(t, found, targetFile) + require.Contains(t, found, targetDir) + }) + + t.Run("invalid starting directory", func(t *testing.T) { + found, err := Lookup("/invalid/path/that/does/not/exist", "target.txt") + require.Error(t, err) + require.Empty(t, found) + }) + + t.Run("relative path handling", func(t *testing.T) { + // Create target files in current directory + require.NoError(t, os.WriteFile("target1.txt", []byte("test1"), 0o644)) + require.NoError(t, os.WriteFile("target2.txt", []byte("test2"), 0o644)) + + // Search using relative path + found, err := Lookup(".", "target1.txt", "target2.txt") + require.NoError(t, err) + require.Len(t, found, 2) + + // Resolve symlinks to handle macOS /private/var vs /var discrepancy + expectedPath1, err := filepath.EvalSymlinks(filepath.Join(tempDir, "target1.txt")) + require.NoError(t, err) + expectedPath2, err := filepath.EvalSymlinks(filepath.Join(tempDir, "target2.txt")) + require.NoError(t, err) + + // Check that found paths match expected paths (order may vary) + foundEvalSymlinks := make([]string, len(found)) + for i, path := range found { + evalPath, err := filepath.EvalSymlinks(path) + require.NoError(t, err) + foundEvalSymlinks[i] = evalPath + } + + require.Contains(t, foundEvalSymlinks, expectedPath1) + require.Contains(t, foundEvalSymlinks, expectedPath2) + }) +} + +func TestLookupClosestBounded(t *testing.T) { + t.Run("found in starting directory", func(t *testing.T) { + testDir := t.TempDir() + + targetFile := filepath.Join(testDir, "target.txt") + require.NoError(t, os.WriteFile(targetFile, []byte("test"), 0o644)) + + foundPath, found := LookupClosestBounded(testDir, testDir, "target.txt") + require.True(t, found) + require.Equal(t, targetFile, foundPath) + }) + + t.Run("found at boundary directory", func(t *testing.T) { + boundary := t.TempDir() + + subDir := filepath.Join(boundary, "subdir") + require.NoError(t, os.Mkdir(subDir, 0o755)) + + targetFile := filepath.Join(boundary, "target.txt") + require.NoError(t, os.WriteFile(targetFile, []byte("test"), 0o644)) + + foundPath, found := LookupClosestBounded(subDir, boundary, "target.txt") + require.True(t, found) + require.Equal(t, targetFile, foundPath) + }) + + t.Run("does not climb past boundary", func(t *testing.T) { + parent := t.TempDir() + + // Target lives above the boundary. + require.NoError(t, os.WriteFile(filepath.Join(parent, "target.txt"), []byte("test"), 0o644)) + + boundary := filepath.Join(parent, "project") + require.NoError(t, os.Mkdir(boundary, 0o755)) + + subDir := filepath.Join(boundary, "subdir") + require.NoError(t, os.Mkdir(subDir, 0o755)) + + foundPath, found := LookupClosestBounded(subDir, boundary, "target.txt") + require.False(t, found) + require.Empty(t, foundPath) + }) + + t.Run("empty boundary searches only starting directory", func(t *testing.T) { + parent := t.TempDir() + + require.NoError(t, os.WriteFile(filepath.Join(parent, "target.txt"), []byte("test"), 0o644)) + + subDir := filepath.Join(parent, "subdir") + require.NoError(t, os.Mkdir(subDir, 0o755)) + + foundPath, found := LookupClosestBounded(subDir, "", "target.txt") + require.False(t, found) + require.Empty(t, foundPath) + }) + + t.Run("empty boundary still finds in starting directory", func(t *testing.T) { + testDir := t.TempDir() + + targetFile := filepath.Join(testDir, "target.txt") + require.NoError(t, os.WriteFile(targetFile, []byte("test"), 0o644)) + + foundPath, found := LookupClosestBounded(testDir, "", "target.txt") + require.True(t, found) + require.Equal(t, targetFile, foundPath) + }) +} + +func TestLookupBounded(t *testing.T) { + t.Run("returns matches at and below boundary", func(t *testing.T) { + boundary := t.TempDir() + + subDir := filepath.Join(boundary, "subdir") + require.NoError(t, os.Mkdir(subDir, 0o755)) + + atBoundary := filepath.Join(boundary, "target.txt") + atSub := filepath.Join(subDir, "target.txt") + require.NoError(t, os.WriteFile(atBoundary, []byte("a"), 0o644)) + require.NoError(t, os.WriteFile(atSub, []byte("b"), 0o644)) + + found, err := LookupBounded(subDir, boundary, "target.txt") + require.NoError(t, err) + require.Len(t, found, 2) + require.Contains(t, found, atBoundary) + require.Contains(t, found, atSub) + }) + + t.Run("ignores matches above boundary", func(t *testing.T) { + parent := t.TempDir() + + require.NoError(t, os.WriteFile(filepath.Join(parent, "target.txt"), []byte("nope"), 0o644)) + + boundary := filepath.Join(parent, "project") + require.NoError(t, os.Mkdir(boundary, 0o755)) + + subDir := filepath.Join(boundary, "subdir") + require.NoError(t, os.Mkdir(subDir, 0o755)) + + // Target lives only above the boundary. + found, err := LookupBounded(subDir, boundary, "target.txt") + require.NoError(t, err) + require.Empty(t, found) + }) + + t.Run("empty boundary searches only starting directory", func(t *testing.T) { + parent := t.TempDir() + + require.NoError(t, os.WriteFile(filepath.Join(parent, "target.txt"), []byte("nope"), 0o644)) + + subDir := filepath.Join(parent, "subdir") + require.NoError(t, os.Mkdir(subDir, 0o755)) + + found, err := LookupBounded(subDir, "", "target.txt") + require.NoError(t, err) + require.Empty(t, found) + }) + + t.Run("no targets returns nil", func(t *testing.T) { + dir := t.TempDir() + found, err := LookupBounded(dir, dir) + require.NoError(t, err) + require.Empty(t, found) + }) +} + +func TestProbeEnt(t *testing.T) { + t.Run("existing file with correct owner", func(t *testing.T) { + tempDir := t.TempDir() + + // Create test file + testFile := filepath.Join(tempDir, "test.txt") + err := os.WriteFile(testFile, []byte("test"), 0o644) + require.NoError(t, err) + + // Get owner of temp directory + owner, err := Owner(tempDir) + require.NoError(t, err) + + // Test probeEnt with correct owner + err = probeEnt(testFile, owner) + require.NoError(t, err) + }) + + t.Run("existing directory with correct owner", func(t *testing.T) { + tempDir := t.TempDir() + + // Create test directory + testDir := filepath.Join(tempDir, "testdir") + err := os.Mkdir(testDir, 0o755) + require.NoError(t, err) + + // Get owner of temp directory + owner, err := Owner(tempDir) + require.NoError(t, err) + + // Test probeEnt with correct owner + err = probeEnt(testDir, owner) + require.NoError(t, err) + }) + + t.Run("nonexistent file", func(t *testing.T) { + tempDir := t.TempDir() + + nonexistentFile := filepath.Join(tempDir, "nonexistent.txt") + owner, err := Owner(tempDir) + require.NoError(t, err) + + err = probeEnt(nonexistentFile, owner) + require.Error(t, err) + require.True(t, errors.Is(err, os.ErrNotExist)) + }) + + t.Run("nonexistent file in nonexistent directory", func(t *testing.T) { + nonexistentFile := "/this/directory/does/not/exists/nonexistent.txt" + + err := probeEnt(nonexistentFile, -1) + require.Error(t, err) + require.True(t, errors.Is(err, os.ErrNotExist)) + }) + + t.Run("ownership bypass with -1", func(t *testing.T) { + tempDir := t.TempDir() + + // Create test file + testFile := filepath.Join(tempDir, "test.txt") + err := os.WriteFile(testFile, []byte("test"), 0o644) + require.NoError(t, err) + + // Test probeEnt with -1 (bypass ownership check) + err = probeEnt(testFile, -1) + require.NoError(t, err) + }) + + t.Run("ownership mismatch returns permission error", func(t *testing.T) { + tempDir := t.TempDir() + + // Create test file + testFile := filepath.Join(tempDir, "test.txt") + err := os.WriteFile(testFile, []byte("test"), 0o644) + require.NoError(t, err) + + // Test probeEnt with different owner (use 9999 which is unlikely to be the actual owner) + err = probeEnt(testFile, 9999) + require.Error(t, err) + require.True(t, errors.Is(err, os.ErrPermission)) + }) +} diff --git a/internal/fsext/ls.go b/internal/fsext/ls.go new file mode 100644 index 0000000000000000000000000000000000000000..5c0d6abb745fcc7951a22d844d918d2a3a21adf4 --- /dev/null +++ b/internal/fsext/ls.go @@ -0,0 +1,316 @@ +package fsext + +import ( + "cmp" + "errors" + "log/slog" + "os" + "path/filepath" + "slices" + "strings" + "sync" + + "github.com/charlievieth/fastwalk" + "github.com/charmbracelet/crush/internal/csync" + "github.com/charmbracelet/crush/internal/home" + gitconfig "github.com/go-git/go-git/v5/config" + "github.com/go-git/go-git/v5/plumbing/format/gitignore" +) + +// fastIgnoreDirs is a set of directory names that are always ignored. +// This provides O(1) lookup for common cases to avoid expensive pattern matching. +var fastIgnoreDirs = map[string]bool{ + ".git": true, + ".svn": true, + ".hg": true, + ".bzr": true, + ".vscode": true, + ".idea": true, + "node_modules": true, + "__pycache__": true, + ".pytest_cache": true, + ".cache": true, + ".tmp": true, + ".Trash": true, + ".Spotlight-V100": true, + ".fseventsd": true, + ".crush": true, + "OrbStack": true, + ".local": true, + ".share": true, +} + +// commonIgnorePatterns contains commonly ignored files and directories. +// Note: Exact directory names that are in fastIgnoreDirs are handled there for O(1) lookup. +// This list contains wildcard patterns and file-specific patterns. +var commonIgnorePatterns = sync.OnceValue(func() []gitignore.Pattern { + patterns := []string{ + // IDE and editor files (wildcards) + "*.swp", + "*.swo", + "*~", + ".DS_Store", + "Thumbs.db", + + // Build artifacts (non-fastIgnoreDirs) + "target", + "build", + "dist", + "out", + "bin", + "obj", + "*.o", + "*.so", + "*.dylib", + "*.dll", + "*.exe", + + // Logs and temporary files (wildcards) + "*.log", + "*.tmp", + "*.temp", + + // Language-specific (wildcards and non-fastIgnoreDirs) + "*.pyc", + "*.pyo", + "vendor", + "Cargo.lock", + "package-lock.json", + "yarn.lock", + "pnpm-lock.yaml", + } + return parsePatterns(patterns, nil) +}) + +// gitGlobalIgnorePatterns returns patterns from git's global excludes file +// (core.excludesFile), following git's config resolution order. +var gitGlobalIgnorePatterns = sync.OnceValue(func() []gitignore.Pattern { + cfg, err := gitconfig.LoadConfig(gitconfig.GlobalScope) + if err != nil { + slog.Debug("Failed to load global git config", "error", err) + return nil + } + + excludesFilePath := cmp.Or( + cfg.Raw.Section("core").Options.Get("excludesfile"), + filepath.Join(home.Config(), "git", "ignore"), + ) + excludesFilePath = home.Long(excludesFilePath) + + bts, err := os.ReadFile(excludesFilePath) + if err != nil { + if !os.IsNotExist(err) { + slog.Debug("Failed to read git global excludes file", "path", excludesFilePath, "error", err) + } + return nil + } + + return parsePatterns(strings.Split(string(bts), "\n"), nil) +}) + +// crushGlobalIgnorePatterns returns patterns from the user's +// ~/.config/crush/ignore file. +var crushGlobalIgnorePatterns = sync.OnceValue(func() []gitignore.Pattern { + name := filepath.Join(home.Config(), "crush", "ignore") + bts, err := os.ReadFile(name) + if err != nil { + if !os.IsNotExist(err) { + slog.Debug("Failed to read crush global ignore file", "path", name, "error", err) + } + return nil + } + lines := strings.Split(string(bts), "\n") + return parsePatterns(lines, nil) +}) + +// parsePatterns parses gitignore pattern strings into Pattern objects. +// domain is the path components where the patterns are defined (nil for global). +func parsePatterns(lines []string, domain []string) []gitignore.Pattern { + var patterns []gitignore.Pattern + for _, line := range lines { + line = strings.TrimSpace(line) + if line == "" || strings.HasPrefix(line, "#") { + continue + } + patterns = append(patterns, gitignore.ParsePattern(line, domain)) + } + return patterns +} + +type directoryLister struct { + // dirPatterns caches parsed patterns from .gitignore/.crushignore for each directory. + // This avoids re-reading files when building combined matchers. + dirPatterns *csync.Map[string, []gitignore.Pattern] + // combinedMatchers caches a combined matcher for each directory that includes + // all ancestor patterns. This allows O(1) matching per file. + combinedMatchers *csync.Map[string, gitignore.Matcher] + rootPath string +} + +func NewDirectoryLister(rootPath string) *directoryLister { + return &directoryLister{ + rootPath: rootPath, + dirPatterns: csync.NewMap[string, []gitignore.Pattern](), + combinedMatchers: csync.NewMap[string, gitignore.Matcher](), + } +} + +// pathToComponents splits a path into its components for gitignore matching. +func pathToComponents(path string) []string { + path = filepath.ToSlash(path) + if path == "" || path == "." { + return nil + } + return strings.Split(path, "/") +} + +// getDirPatterns returns the parsed patterns for a specific directory's +// .gitignore and .crushignore files. Results are cached. +func (dl *directoryLister) getDirPatterns(dir string) []gitignore.Pattern { + return dl.dirPatterns.GetOrSet(dir, func() []gitignore.Pattern { + var allPatterns []gitignore.Pattern + + relPath, _ := filepath.Rel(dl.rootPath, dir) + var domain []string + if relPath != "" && relPath != "." { + domain = pathToComponents(relPath) + } + + for _, ignoreFile := range []string{".gitignore", ".crushignore"} { + ignPath := filepath.Join(dir, ignoreFile) + if content, err := os.ReadFile(ignPath); err == nil { + lines := strings.Split(string(content), "\n") + allPatterns = append(allPatterns, parsePatterns(lines, domain)...) + } + } + return allPatterns + }) +} + +// getCombinedMatcher returns a matcher that combines all gitignore patterns +// from the root to the given directory, plus common patterns and home patterns. +// Results are cached per directory, and we reuse parent directory matchers. +func (dl *directoryLister) getCombinedMatcher(dir string) gitignore.Matcher { + return dl.combinedMatchers.GetOrSet(dir, func() gitignore.Matcher { + var allPatterns []gitignore.Pattern + + // Add common patterns first (lowest priority). + allPatterns = append(allPatterns, commonIgnorePatterns()...) + + // Add global ignore patterns (git core.excludesFile + crush global ignore). + allPatterns = append(allPatterns, gitGlobalIgnorePatterns()...) + allPatterns = append(allPatterns, crushGlobalIgnorePatterns()...) + + // Collect patterns from root to this directory. + relDir, _ := filepath.Rel(dl.rootPath, dir) + var pathParts []string + if relDir != "" && relDir != "." { + pathParts = pathToComponents(relDir) + } + + // Add patterns from each directory from root to current. + currentPath := dl.rootPath + allPatterns = append(allPatterns, dl.getDirPatterns(currentPath)...) + + for _, part := range pathParts { + currentPath = filepath.Join(currentPath, part) + allPatterns = append(allPatterns, dl.getDirPatterns(currentPath)...) + } + + return gitignore.NewMatcher(allPatterns) + }) +} + +// shouldIgnore checks if a path should be ignored based on gitignore rules. +// This uses a combined matcher that includes all ancestor patterns for O(1) matching. +func (dl *directoryLister) shouldIgnore(path string, ignorePatterns []string, isDir bool) bool { + base := filepath.Base(path) + + // Fast path: O(1) lookup for commonly ignored directories. + if isDir && fastIgnoreDirs[base] { + return true + } + + // Check explicit ignore patterns. + if len(ignorePatterns) > 0 { + for _, pattern := range ignorePatterns { + if matched, err := filepath.Match(pattern, base); err == nil && matched { + return true + } + } + } + + // Don't apply gitignore rules to the root directory itself. + if path == dl.rootPath { + return false + } + + relPath, err := filepath.Rel(dl.rootPath, path) + if err != nil { + relPath = path + } + + pathComponents := pathToComponents(relPath) + if len(pathComponents) == 0 { + return false + } + + // Get the combined matcher for the parent directory. + parentDir := filepath.Dir(path) + matcher := dl.getCombinedMatcher(parentDir) + + if matcher.Match(pathComponents, isDir) { + slog.Debug("Ignoring path", "path", relPath) + return true + } + + return false +} + +// ListDirectory lists files and directories in the specified path. +func ListDirectory(initialPath string, ignorePatterns []string, depth, limit int) ([]string, bool, error) { + found := csync.NewSlice[string]() + dl := NewDirectoryLister(initialPath) + + slog.Debug("Listing directory", "path", initialPath, "depth", depth, "limit", limit, "ignorePatterns", ignorePatterns) + + conf := fastwalk.Config{ + Follow: true, + ToSlash: fastwalk.DefaultToSlash(), + Sort: fastwalk.SortDirsFirst, + MaxDepth: depth, + } + + err := fastwalk.Walk(&conf, initialPath, func(path string, d os.DirEntry, err error) error { + if err != nil { + return nil // Skip files we don't have permission to access + } + + isDir := d.IsDir() + if dl.shouldIgnore(path, ignorePatterns, isDir) { + if isDir { + return filepath.SkipDir + } + return nil + } + + if path != initialPath { + if isDir { + path = path + string(filepath.Separator) + } + found.Append(path) + } + + if limit > 0 && found.Len() >= limit { + return filepath.SkipAll + } + + return nil + }) + if err != nil && !errors.Is(err, filepath.SkipAll) { + return nil, false, err + } + + matches, truncated := truncate(slices.Collect(found.Seq()), limit) + return matches, truncated || errors.Is(err, filepath.SkipAll), nil +} diff --git a/internal/fsext/ls_test.go b/internal/fsext/ls_test.go new file mode 100644 index 0000000000000000000000000000000000000000..1bb862754ea38c378fa75709f9dd4b30dd8803f8 --- /dev/null +++ b/internal/fsext/ls_test.go @@ -0,0 +1,60 @@ +package fsext + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestListDirectory(t *testing.T) { + tmp := t.TempDir() + + testFiles := map[string]string{ + "regular.txt": "content", + ".hidden": "hidden content", + ".gitignore": ".*\n*.log\n", + "subdir/file.go": "package main", + "subdir/.another": "more hidden", + "build.log": "build output", + } + + for name, content := range testFiles { + fp := filepath.Join(tmp, name) + dir := filepath.Dir(fp) + require.NoError(t, os.MkdirAll(dir, 0o755)) + require.NoError(t, os.WriteFile(fp, []byte(content), 0o644)) + } + + t.Run("no limit", func(t *testing.T) { + files, truncated, err := ListDirectory(tmp, nil, -1, -1) + require.NoError(t, err) + require.False(t, truncated) + // The .gitignore has ".*" pattern which ignores hidden files anywhere + // (like real git does), so subdir/.another is ignored. + require.Len(t, files, 3) + require.ElementsMatch(t, []string{ + "regular.txt", + "subdir", + "subdir/file.go", + }, relPaths(t, files, tmp)) + }) + t.Run("limit", func(t *testing.T) { + files, truncated, err := ListDirectory(tmp, nil, -1, 2) + require.NoError(t, err) + require.True(t, truncated) + require.Len(t, files, 2) + }) +} + +func relPaths(tb testing.TB, in []string, base string) []string { + tb.Helper() + out := make([]string, 0, len(in)) + for _, p := range in { + rel, err := filepath.Rel(base, p) + require.NoError(tb, err) + out = append(out, filepath.ToSlash(rel)) + } + return out +} diff --git a/internal/fsext/owner_others.go b/internal/fsext/owner_others.go new file mode 100644 index 0000000000000000000000000000000000000000..3900f501e0d77fa80089410eca890a04a0882dbf --- /dev/null +++ b/internal/fsext/owner_others.go @@ -0,0 +1,24 @@ +//go:build !windows + +package fsext + +import ( + "os" + "syscall" +) + +// Owner retrieves the user ID of the owner of the file or directory at the +// specified path. +func Owner(path string) (int, error) { + info, err := os.Stat(path) + if err != nil { + return 0, err + } + var uid int + if stat, ok := info.Sys().(*syscall.Stat_t); ok { + uid = int(stat.Uid) + } else { + uid = os.Getuid() + } + return uid, nil +} diff --git a/internal/fsext/owner_windows.go b/internal/fsext/owner_windows.go new file mode 100644 index 0000000000000000000000000000000000000000..41f9091c3e75e8f187984a8e1ddb7a7aa72c9dab --- /dev/null +++ b/internal/fsext/owner_windows.go @@ -0,0 +1,15 @@ +//go:build windows + +package fsext + +import "os" + +// Owner retrieves the user ID of the owner of the file or directory at the +// specified path. +func Owner(path string) (int, error) { + _, err := os.Stat(path) + if err != nil { + return 0, err + } + return -1, nil +} diff --git a/internal/fsext/paste.go b/internal/fsext/paste.go new file mode 100644 index 0000000000000000000000000000000000000000..4996473acf41355e391ba6e9bf2547abfbbea9cb --- /dev/null +++ b/internal/fsext/paste.go @@ -0,0 +1,129 @@ +package fsext + +import ( + "os" + "strings" +) + +func ParsePastedFiles(s string) []string { + s = strings.TrimSpace(s) + + // NOTE: Rio on Windows adds NULL chars for some reason. + s = strings.ReplaceAll(s, "\x00", "") + + switch { + case attemptStat(s): + return strings.Split(s, "\n") + case os.Getenv("WT_SESSION") != "": + return windowsTerminalParsePastedFiles(s) + default: + return unixParsePastedFiles(s) + } +} + +func attemptStat(s string) bool { + for path := range strings.SplitSeq(s, "\n") { + if info, err := os.Stat(path); err != nil || info.IsDir() { + return false + } + } + return true +} + +func windowsTerminalParsePastedFiles(s string) []string { + if strings.TrimSpace(s) == "" { + return nil + } + + var ( + paths []string + current strings.Builder + inQuotes = false + ) + for i := range len(s) { + ch := s[i] + + switch { + case ch == '"': + if inQuotes { + // End of quoted section + if current.Len() > 0 { + paths = append(paths, current.String()) + current.Reset() + } + inQuotes = false + } else { + // Start of quoted section + inQuotes = true + } + case inQuotes: + current.WriteByte(ch) + case ch != ' ': + // Text outside quotes is not allowed + return nil + } + } + + // Add any remaining content if quotes were properly closed + if current.Len() > 0 && !inQuotes { + paths = append(paths, current.String()) + } + + // If quotes were not closed, return empty (malformed input) + if inQuotes { + return nil + } + + return paths +} + +func unixParsePastedFiles(s string) []string { + if strings.TrimSpace(s) == "" { + return nil + } + + var ( + paths []string + current strings.Builder + escaped = false + ) + for i := range len(s) { + ch := s[i] + + switch { + case escaped: + // After a backslash, add the character as-is (including space) + current.WriteByte(ch) + escaped = false + case ch == '\\': + // Check if this backslash is at the end of the string + if i == len(s)-1 { + // Trailing backslash, treat as literal + current.WriteByte(ch) + } else { + // Start of escape sequence + escaped = true + } + case ch == ' ': + // Space separates paths (unless escaped) + if current.Len() > 0 { + paths = append(paths, current.String()) + current.Reset() + } + default: + current.WriteByte(ch) + } + } + + // Handle trailing backslash if present + if escaped { + current.WriteByte('\\') + } + + // Add the last path if any + if current.Len() > 0 { + paths = append(paths, current.String()) + } + + return paths +} diff --git a/internal/fsext/paste_test.go b/internal/fsext/paste_test.go new file mode 100644 index 0000000000000000000000000000000000000000..c1c4d4adfba0eca44586f55f2a23dd882038522e --- /dev/null +++ b/internal/fsext/paste_test.go @@ -0,0 +1,149 @@ +package fsext + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestParsePastedFiles(t *testing.T) { + t.Run("WindowsTerminal", func(t *testing.T) { + tests := []struct { + name string + input string + expected []string + }{ + { + name: "single path", + input: `"C:\path\my-screenshot-one.png"`, + expected: []string{`C:\path\my-screenshot-one.png`}, + }, + { + name: "multiple paths no spaces", + input: `"C:\path\my-screenshot-one.png" "C:\path\my-screenshot-two.png" "C:\path\my-screenshot-three.png"`, + expected: []string{`C:\path\my-screenshot-one.png`, `C:\path\my-screenshot-two.png`, `C:\path\my-screenshot-three.png`}, + }, + { + name: "single with spaces", + input: `"C:\path\my screenshot one.png"`, + expected: []string{`C:\path\my screenshot one.png`}, + }, + { + name: "multiple paths with spaces", + input: `"C:\path\my screenshot one.png" "C:\path\my screenshot two.png" "C:\path\my screenshot three.png"`, + expected: []string{`C:\path\my screenshot one.png`, `C:\path\my screenshot two.png`, `C:\path\my screenshot three.png`}, + }, + { + name: "empty string", + input: "", + expected: nil, + }, + { + name: "unclosed quotes", + input: `"C:\path\file.png`, + expected: nil, + }, + { + name: "text outside quotes", + input: `"C:\path\file.png" some random text "C:\path\file2.png"`, + expected: nil, + }, + { + name: "multiple spaces between paths", + input: `"C:\path\file1.png" "C:\path\file2.png"`, + expected: []string{`C:\path\file1.png`, `C:\path\file2.png`}, + }, + { + name: "just whitespace", + input: " ", + expected: nil, + }, + { + name: "consecutive quoted sections", + input: `"C:\path1""C:\path2"`, + expected: []string{`C:\path1`, `C:\path2`}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := windowsTerminalParsePastedFiles(tt.input) + require.Equal(t, tt.expected, result) + }) + } + }) + + t.Run("Unix", func(t *testing.T) { + tests := []struct { + name string + input string + expected []string + }{ + { + name: "single path", + input: `/path/my-screenshot.png`, + expected: []string{"/path/my-screenshot.png"}, + }, + { + name: "multiple paths no spaces", + input: `/path/screenshot-one.png /path/screenshot-two.png /path/screenshot-three.png`, + expected: []string{"/path/screenshot-one.png", "/path/screenshot-two.png", "/path/screenshot-three.png"}, + }, + { + name: "sigle with spaces", + input: `/path/my\ screenshot\ one.png`, + expected: []string{"/path/my screenshot one.png"}, + }, + { + name: "multiple paths with spaces", + input: `/path/my\ screenshot\ one.png /path/my\ screenshot\ two.png /path/my\ screenshot\ three.png`, + expected: []string{"/path/my screenshot one.png", "/path/my screenshot two.png", "/path/my screenshot three.png"}, + }, + { + name: "empty string", + input: "", + expected: nil, + }, + { + name: "double backslash escapes", + input: `/path/my\\file.png`, + expected: []string{"/path/my\\file.png"}, + }, + { + name: "trailing backslash", + input: `/path/file\`, + expected: []string{`/path/file\`}, + }, + { + name: "multiple consecutive escaped spaces", + input: `/path/file\ \ with\ \ many\ \ spaces.png`, + expected: []string{"/path/file with many spaces.png"}, + }, + { + name: "multiple unescaped spaces", + input: `/path/file1.png /path/file2.png`, + expected: []string{"/path/file1.png", "/path/file2.png"}, + }, + { + name: "just whitespace", + input: " ", + expected: nil, + }, + { + name: "tab characters", + input: "/path/file1.png\t/path/file2.png", + expected: []string{"/path/file1.png\t/path/file2.png"}, + }, + { + name: "newlines in input", + input: "/path/file1.png\n/path/file2.png", + expected: []string{"/path/file1.png\n/path/file2.png"}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := unixParsePastedFiles(tt.input) + require.Equal(t, tt.expected, result) + }) + } + }) +} diff --git a/internal/herdr/client.go b/internal/herdr/client.go new file mode 100644 index 0000000000000000000000000000000000000000..c7d10eb7302597f1f9a90b39baeadb519b20d0cc --- /dev/null +++ b/internal/herdr/client.go @@ -0,0 +1,402 @@ +// Package herdr provides native integration with the herdr terminal +// multiplexer. When Crush runs inside a herdr-managed pane it reports +// agent state (idle, working, blocked) and session identity over +// herdr's Unix socket API so herdr can display accurate status without +// screen scraping. +// +// The client consumes a small, herdr-specific event vocabulary rather +// than accepting raw proto or domain types. Callers translate their +// events into herdr.Event before forwarding. This keeps the client +// decoupled from both the proto and internal domain layers. +package herdr + +import ( + "context" + "encoding/json" + "flag" + "fmt" + "io" + "log/slog" + "net" + "os" + "sync" + "time" +) + +// State values matching herdr's PaneAgentState enum. +const ( + stateIdle = "idle" + stateWorking = "working" + stateBlocked = "blocked" +) + +// Event is the herdr-specific event vocabulary. Each type maps to a +// distinct state transition in the agent lifecycle. Callers translate +// from proto or domain types into these before calling HandleEvent. +type Event interface { + herdrEvent() +} + +// AssistantMessage indicates the agent produced output. Transitions +// to working if not already active. +type AssistantMessage struct { + SessionID string +} + +func (AssistantMessage) herdrEvent() {} + +// RunComplete indicates the agent finished a turn. Transitions to idle. +type RunComplete struct { + SessionID string +} + +func (RunComplete) herdrEvent() {} + +// PermissionRequested indicates the agent is waiting for user approval. +// Transitions to blocked. +type PermissionRequested struct{} + +func (PermissionRequested) herdrEvent() {} + +// PermissionResolved indicates a permission decision was made. +// Transitions back to working if a run is active, idle otherwise. +type PermissionResolved struct{} + +func (PermissionResolved) herdrEvent() {} + +// Summarizing indicates the agent is compacting context. Transitions +// to working if not already active. +type Summarizing struct{} + +func (Summarizing) herdrEvent() {} + +// sender abstracts the transport layer for reporting state to herdr. +// Production uses a Unix socket; tests use a recorder. +type sender interface { + send(req reportRequest) error + close() +} + +// Client reports Crush agent state to a running herdr instance. +type Client struct { + socketPath string + paneID string + + mu sync.Mutex + sessionID string + state string + runActive bool + seq uint64 + + snd sender +} + +// defaultClient is the process-wide herdr client. Initialized once +// via Init(). All integration sites share this single instance so +// only one Unix socket connection exists per process. +var ( + defaultClient *Client + initOnce sync.Once +) + +// Init returns the process-wide herdr Client, creating it on first +// call from environment variables. Returns nil when Crush is not +// running inside a herdr pane. Safe to call from any goroutine. +func Init() *Client { + initOnce.Do(func() { + defaultClient = newFromEnv() + }) + return defaultClient +} + +func newFromEnv() *Client { + if os.Getenv("HERDR_ENV") != "1" { + return nil + } + // A test binary inherits the launching shell's HERDR_* env, so + // without this it would attach to the developer's live pane and + // release its agent on teardown. Skip herdr entirely under test. + if flag.Lookup("test.v") != nil { + slog.Debug("Herdr integration disabled: running under go test") + return nil + } + socketPath := os.Getenv("HERDR_SOCKET_PATH") + paneID := os.Getenv("HERDR_PANE_ID") + if socketPath == "" || paneID == "" { + slog.Debug( + "Herdr integration disabled: incomplete environment", + "has_socket", socketPath != "", + "has_pane_id", paneID != "", + ) + return nil + } + c := &Client{ + socketPath: socketPath, + paneID: paneID, + state: stateIdle, + seq: uint64(time.Now().UnixNano()), + snd: newUnixSender(socketPath), + } + c.registerInitial() + return c +} + +// registerInitial sends an initial idle-state report to herdr so the +// pane knows about the agent immediately, not just after the first +// event. Called once during client creation. Bypasses the dedup +// check since the initial state must always be reported regardless +// of redundancy. +// +// herdr remembers the highest seq it has seen per source for the +// lifetime of a pane and silently drops any report with a seq that +// is not strictly greater. Because crush seeds seq from the wall +// clock at startup (see newFromEnv), a restarted crush in the same +// pane always reports above the previous run's high-water mark, so +// the first report is accepted instead of being rejected as stale. +func (c *Client) registerInitial() { + if c == nil { + return + } + c.mu.Lock() + defer c.mu.Unlock() + c.snd.send(c.newRequestLocked("pane.report_agent", "init", stateIdle)) +} + +// Close releases the agent's authority on the pane and shuts down +// the background writer. Safe to call on a nil client. +func (c *Client) Close() { + if c == nil { + return + } + c.releaseAgent() + c.snd.close() +} + +// releaseAgent sends a pane.release_agent request to herdr so the +// pane is freed for a new agent to claim authority. This is the +// clean-shutdown protocol per herdr's socket API. Sends directly +// on the socket to ensure delivery even if the write loop is busy. +func (c *Client) releaseAgent() { + c.mu.Lock() + defer c.mu.Unlock() + req := c.newRequestLocked("pane.release_agent", "release", "") + if err := dialSend(c.socketPath, req); err != nil { + slog.Debug("Herdr release_agent failed", "error", err) + } +} + +// HandleEvent processes a single herdr event and reports state changes. +// Safe to call from any goroutine. +func (c *Client) HandleEvent(ev Event) { + if c == nil { + return + } + switch e := ev.(type) { + case AssistantMessage: + c.onAssistantMessage(e.SessionID) + case RunComplete: + c.onRunComplete(e.SessionID) + case PermissionRequested: + c.onPermissionRequest() + case PermissionResolved: + c.onPermissionResolved() + case Summarizing: + c.onSummarizing() + } +} + +// SetSessionID sets the session ID for reporting. Call this when the +// session is created or resolved, before events start flowing. +func (c *Client) SetSessionID(id string) { + if c == nil { + return + } + c.mu.Lock() + defer c.mu.Unlock() + c.sessionID = id +} + +func (c *Client) onAssistantMessage(sessionID string) { + c.mu.Lock() + defer c.mu.Unlock() + if sessionID != "" { + c.sessionID = sessionID + } + if !c.runActive { + c.runActive = true + c.reportLocked(stateWorking) + } +} + +func (c *Client) onRunComplete(sessionID string) { + c.mu.Lock() + defer c.mu.Unlock() + c.runActive = false + if sessionID != "" { + c.sessionID = sessionID + } + c.reportLocked(stateIdle) +} + +func (c *Client) onPermissionRequest() { + c.mu.Lock() + defer c.mu.Unlock() + // A permission request implies a run is active, even if no + // assistant message has arrived yet (e.g. tool calls that fire + // before any text output). + if !c.runActive { + c.runActive = true + } + c.reportLocked(stateBlocked) +} + +func (c *Client) onPermissionResolved() { + c.mu.Lock() + defer c.mu.Unlock() + if c.runActive { + c.reportLocked(stateWorking) + } else { + c.reportLocked(stateIdle) + } +} + +func (c *Client) onSummarizing() { + c.mu.Lock() + defer c.mu.Unlock() + if !c.runActive { + c.runActive = true + } + c.reportLocked(stateWorking) +} + +// newRequestLocked builds a seq-stamped JSON-RPC request to herdr. +// Must be called with c.mu held. Every request increments c.seq so +// herdr accepts it as strictly newer than the last (see +// registerInitial for why monotonic seq matters). State is empty for +// requests that carry no agent state, such as pane.release_agent. +func (c *Client) newRequestLocked(method, idPrefix, state string) reportRequest { + c.seq++ + return reportRequest{ + ID: fmt.Sprintf("crush:%s:%d", idPrefix, time.Now().UnixNano()), + Method: method, + Params: reportParams{ + PaneID: c.paneID, + Source: "crush", + Agent: "crush", + State: state, + Seq: c.seq, + AgentSessionID: c.sessionID, + }, + } +} + +// reportLocked sends a pane.report_agent request to herdr. Must be +// called with c.mu held. Skips redundant reports when the state has +// not changed. +func (c *Client) reportLocked(state string) { + if state == c.state { + return + } + c.state = state + c.snd.send(c.newRequestLocked("pane.report_agent", "report", state)) +} + +// reportRequest is the JSON-RPC envelope sent to herdr. +type reportRequest struct { + ID string `json:"id"` + Method string `json:"method"` + Params reportParams `json:"params"` +} + +// reportParams carries the agent state payload. +type reportParams struct { + PaneID string `json:"pane_id"` + Source string `json:"source"` + Agent string `json:"agent"` + State string `json:"state"` + Seq uint64 `json:"seq"` + AgentSessionID string `json:"agent_session_id"` +} + +// unixSender sends JSON-RPC requests over a Unix domain socket using +// a single background writer goroutine and a buffered channel. This +// serializes writes and avoids spawning unbounded goroutines under +// high event throughput. Each report opens a short-lived connection. +type unixSender struct { + socketPath string + ch chan reportRequest + cancel context.CancelFunc +} + +func newUnixSender(socketPath string) *unixSender { + ctx, cancel := context.WithCancel(context.Background()) + s := &unixSender{ + socketPath: socketPath, + ch: make(chan reportRequest, 16), + cancel: cancel, + } + go s.writeLoop(ctx) + return s +} + +func (s *unixSender) send(req reportRequest) error { + select { + case s.ch <- req: + default: + // Drop if the buffer is full. State reports are + // best-effort; blocking the agent is worse than + // missing a transition. + } + return nil +} + +func (s *unixSender) close() { + s.cancel() +} + +func (s *unixSender) writeLoop(ctx context.Context) { + for { + select { + case req, ok := <-s.ch: + if !ok { + return + } + if err := dialSend(s.socketPath, req); err != nil { + slog.Debug("Herdr report failed", "error", err) + } + case <-ctx.Done(): + return + } + } +} + +// dialSend opens a short-lived Unix socket connection to herdr, +// sends a single JSON-RPC request, and drains the response. +func dialSend(socketPath string, req reportRequest) error { + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + defer cancel() + + dialer := net.Dialer{} + conn, err := dialer.DialContext(ctx, "unix", socketPath) + if err != nil { + return err + } + defer conn.Close() + + _ = conn.SetDeadline(time.Now().Add(500 * time.Millisecond)) + + data, err := json.Marshal(req) + if err != nil { + return err + } + data = append(data, '\n') + + _, err = conn.Write(data) + if err != nil { + return err + } + + // Drain the response to complete the request cycle. + _, _ = io.Copy(io.Discard, conn) + return nil +} diff --git a/internal/herdr/client_test.go b/internal/herdr/client_test.go new file mode 100644 index 0000000000000000000000000000000000000000..f14e821a0c3c1c4488dc04063da61b704c671573 --- /dev/null +++ b/internal/herdr/client_test.go @@ -0,0 +1,158 @@ +package herdr + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +// recordingSender captures state transitions without connecting to a +// real Unix socket. +type recordingSender struct { + states []string +} + +func (r *recordingSender) send(req reportRequest) error { + r.states = append(r.states, req.Params.State) + return nil +} + +func (r *recordingSender) close() {} + +// newTestClient creates a Client that records state transitions +// without connecting to a real Unix socket. +func newTestClient() *Client { + rec := &recordingSender{states: make([]string, 0, 16)} + return &Client{ + state: stateIdle, + snd: rec, + } +} + +// reportedStates returns the states recorded by the test sender. +func reportedStates(c *Client) []string { + return c.snd.(*recordingSender).states +} + +func TestBasicLifecycle(t *testing.T) { + t.Parallel() + c := newTestClient() + + // Assistant message starts working. + c.HandleEvent(AssistantMessage{SessionID: "sess-1"}) + assert.Equal(t, []string{stateWorking}, reportedStates(c)) + + // Run complete returns to idle. + c.HandleEvent(RunComplete{SessionID: "sess-1"}) + assert.Equal(t, []string{stateWorking, stateIdle}, reportedStates(c)) +} + +func TestPermissionBlockAndUnblock(t *testing.T) { + t.Parallel() + c := newTestClient() + + // Start working. + c.HandleEvent(AssistantMessage{SessionID: "sess-1"}) + + // Permission request blocks. + c.HandleEvent(PermissionRequested{}) + assert.Equal(t, []string{stateWorking, stateBlocked}, reportedStates(c)) + + // Permission granted returns to working (run still active). + c.HandleEvent(PermissionResolved{}) + assert.Equal(t, []string{stateWorking, stateBlocked, stateWorking}, reportedStates(c)) + + // Run complete returns to idle. + c.HandleEvent(RunComplete{SessionID: "sess-1"}) + assert.Equal(t, []string{stateWorking, stateBlocked, stateWorking, stateIdle}, reportedStates(c)) +} + +func TestPermissionBeforeAssistantMessage(t *testing.T) { + t.Parallel() + c := newTestClient() + + // Permission request arrives before any assistant message. + // This can happen when tool calls fire before text output. + c.HandleEvent(PermissionRequested{}) + assert.Equal(t, []string{stateBlocked}, reportedStates(c)) + + // Permission resolved should return to working, not idle, + // because the permission request implied a run was active. + c.HandleEvent(PermissionResolved{}) + assert.Equal(t, []string{stateBlocked, stateWorking}, reportedStates(c)) +} + +func TestSessionIDPropagation(t *testing.T) { + t.Parallel() + c := newTestClient() + + // SetSessionID before events. + c.SetSessionID("early-session") + assert.Equal(t, "early-session", c.sessionID) + + // RunComplete also updates session ID. + c.HandleEvent(RunComplete{SessionID: "final-session"}) + assert.Equal(t, "final-session", c.sessionID) +} + +func TestDedupSkipsRedundantState(t *testing.T) { + t.Parallel() + c := newTestClient() + + // Two assistant messages in a row should only report working once. + c.HandleEvent(AssistantMessage{SessionID: "s1"}) + c.HandleEvent(AssistantMessage{SessionID: "s1"}) + assert.Equal(t, []string{stateWorking}, reportedStates(c)) +} + +func TestSummarizingTriggersWorking(t *testing.T) { + t.Parallel() + c := newTestClient() + + // Summarizing event should trigger working. + c.HandleEvent(Summarizing{}) + assert.Equal(t, []string{stateWorking}, reportedStates(c)) + + // Second summarizing should not trigger another state change. + c.HandleEvent(Summarizing{}) + assert.Equal(t, []string{stateWorking}, reportedStates(c)) +} + +func TestNilClientSafe(t *testing.T) { + t.Parallel() + var c *Client + // These should not panic on a nil receiver. + c.SetSessionID("s1") + c.HandleEvent(AssistantMessage{SessionID: "s1"}) + c.HandleEvent(RunComplete{SessionID: "s1"}) + c.HandleEvent(PermissionRequested{}) + c.HandleEvent(PermissionResolved{}) + c.HandleEvent(Summarizing{}) +} + +func TestRegisterInitial(t *testing.T) { + t.Parallel() + rec := &recordingSender{states: make([]string, 0, 16)} + c := &Client{ + state: stateIdle, + seq: 100, + snd: rec, + } + c.registerInitial() + assert.Equal(t, []string{stateIdle}, rec.states) + // seq must strictly increase so herdr accepts the report. + assert.Equal(t, uint64(101), c.seq) +} + +// TestInitDisabledUnderTest guards the critical safety property that +// herdr never attaches to a real pane from a test binary. Test +// processes inherit the developer's HERDR_* environment, so a missing +// guard would release the live pane's agent on teardown. Because this +// test itself runs under `go test`, Init must return nil even with a +// complete, valid-looking environment. +func TestInitDisabledUnderTest(t *testing.T) { + t.Setenv("HERDR_ENV", "1") + t.Setenv("HERDR_SOCKET_PATH", "/tmp/does-not-matter.sock") + t.Setenv("HERDR_PANE_ID", "test:pane") + assert.Nil(t, newFromEnv()) +} diff --git a/internal/herdr/translate.go b/internal/herdr/translate.go new file mode 100644 index 0000000000000000000000000000000000000000..924bd3431e3120c873018e8ecf3fde164f6e083e --- /dev/null +++ b/internal/herdr/translate.go @@ -0,0 +1,150 @@ +package herdr + +import ( + "context" + "time" + + "github.com/charmbracelet/crush/internal/agent/notify" + "github.com/charmbracelet/crush/internal/message" + "github.com/charmbracelet/crush/internal/permission" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/pubsub" +) + +// Translate converts a pub/sub event (domain or proto) into a herdr +// Event. Returns nil for event types herdr doesn't care about. This +// is the single translation point for all integration modes. +func Translate(ev any) Event { + switch e := ev.(type) { + // Domain types (TUI / local headless). + case pubsub.Event[message.Message]: + return translateMessage( + e.Payload.Role == message.Assistant, + e.Payload.SessionID, + e.Payload.IsSummaryMessage, + ) + case pubsub.Event[notify.RunComplete]: + return RunComplete{SessionID: e.Payload.SessionID} + case pubsub.Event[permission.PermissionRequest]: + return PermissionRequested{} + case pubsub.Event[permission.PermissionNotification]: + return PermissionResolved{} + + // Proto types (client/server mode). + case pubsub.Event[proto.Message]: + return translateMessage( + e.Payload.Role == proto.Assistant, + e.Payload.SessionID, + false, + ) + case pubsub.Event[proto.RunComplete]: + return RunComplete{SessionID: e.Payload.SessionID} + case pubsub.Event[proto.PermissionRequest]: + return PermissionRequested{} + case pubsub.Event[proto.PermissionNotification]: + return PermissionResolved{} + case pubsub.Event[proto.AgentEvent]: + if e.Payload.Type == proto.AgentEventTypeSummarize && !e.Payload.Done { + return Summarizing{} + } + return nil + + default: + return nil + } +} + +// translateMessage is the shared message-mapping logic for both domain +// and proto message types. +func translateMessage(isAssistant bool, sessionID string, isSummary bool) Event { + if !isAssistant { + return nil + } + if isSummary { + return Summarizing{} + } + return AssistantMessage{SessionID: sessionID} +} + +// permNotificationSubscriber is the subset of the permission service +// needed by BridgeLocal to subscribe to permission notifications. +type permNotificationSubscriber interface { + SubscribeNotifications(context.Context) <-chan pubsub.Event[permission.PermissionNotification] +} + +// BridgeSources groups the pub/sub sources that BridgeLocal subscribes +// to. Adding a new event type means adding a field here rather than +// growing the function signature. +type BridgeSources struct { + PermRequests pubsub.Subscriber[permission.PermissionRequest] + PermNotifications permNotificationSubscriber + RunCompletions pubsub.Subscriber[notify.RunComplete] + Messages pubsub.Subscriber[message.Message] +} + +// BridgeLocal subscribes to local pub/sub brokers and forwards +// translated events to the client. Used in TUI and local headless +// modes where the agent runs in-process. Cancelling ctx stops the +// bridge goroutines. +// +// The spawned goroutines are best-effort and may briefly outlive +// Client.Close(). This is safe: HandleEvent is nil-safe, and the +// unixSender drops messages on a full buffer rather than blocking. +// +// Each goroutine uses a resilient subscription loop that re-subscribes +// if the channel closes unexpectedly, ensuring the bridge survives +// transient pub/sub broker resets. +func BridgeLocal(ctx context.Context, c *Client, src BridgeSources) { + if c == nil { + return + } + go forward(ctx, c, func(subCtx context.Context) <-chan pubsub.Event[permission.PermissionRequest] { + return src.PermRequests.Subscribe(subCtx) + }) + go forward(ctx, c, func(subCtx context.Context) <-chan pubsub.Event[permission.PermissionNotification] { + return src.PermNotifications.SubscribeNotifications(subCtx) + }) + go forward(ctx, c, func(subCtx context.Context) <-chan pubsub.Event[notify.RunComplete] { + return src.RunCompletions.Subscribe(subCtx) + }) + go forward(ctx, c, func(subCtx context.Context) <-chan pubsub.Event[message.Message] { + return src.Messages.Subscribe(subCtx) + }) +} + +// forward reads from a pub/sub channel and forwards translated +// events to the herdr client. If the channel closes (e.g., due to +// broker reset), it re-subscribes after a brief delay. Runs until ctx +// is cancelled. +func forward[T any](ctx context.Context, c *Client, subscribe func(context.Context) <-chan pubsub.Event[T]) { + for { + select { + case <-ctx.Done(): + return + default: + } + + subCtx, cancel := context.WithCancel(ctx) + ch := subscribe(subCtx) + + inner: + for { + select { + case <-ctx.Done(): + cancel() + return + case ev, ok := <-ch: + if !ok { + // Channel closed — broker may have reset. + // Cancel the sub-context and re-subscribe. + cancel() + time.Sleep(100 * time.Millisecond) + break inner + } + if hev := Translate(ev); hev != nil { + c.HandleEvent(hev) + } + } + } + } +} diff --git a/internal/herdr/translate_test.go b/internal/herdr/translate_test.go new file mode 100644 index 0000000000000000000000000000000000000000..c37135f730572b72920e4317c264979003d943a7 --- /dev/null +++ b/internal/herdr/translate_test.go @@ -0,0 +1,145 @@ +package herdr + +import ( + "testing" + + "github.com/charmbracelet/crush/internal/agent/notify" + "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/stretchr/testify/assert" +) + +// Domain type translation. + +func TestTranslateDomainAssistantMessage(t *testing.T) { + t.Parallel() + ev := pubsub.Event[message.Message]{ + Payload: message.Message{Role: message.Assistant, SessionID: "s1"}, + } + assert.Equal(t, AssistantMessage{SessionID: "s1"}, Translate(ev)) +} + +func TestTranslateDomainSummaryMessage(t *testing.T) { + t.Parallel() + ev := pubsub.Event[message.Message]{ + Payload: message.Message{ + Role: message.Assistant, + SessionID: "s1", + IsSummaryMessage: true, + }, + } + assert.Equal(t, Summarizing{}, Translate(ev)) +} + +func TestTranslateDomainNonAssistantIgnored(t *testing.T) { + t.Parallel() + ev := pubsub.Event[message.Message]{ + Payload: message.Message{Role: message.System}, + } + assert.Nil(t, Translate(ev)) +} + +func TestTranslateDomainRunComplete(t *testing.T) { + t.Parallel() + ev := pubsub.Event[notify.RunComplete]{ + Payload: notify.RunComplete{SessionID: "s1"}, + } + assert.Equal(t, RunComplete{SessionID: "s1"}, Translate(ev)) +} + +func TestTranslateDomainPermissionRequest(t *testing.T) { + t.Parallel() + ev := pubsub.Event[permission.PermissionRequest]{ + Payload: permission.PermissionRequest{ToolName: "bash"}, + } + assert.Equal(t, PermissionRequested{}, Translate(ev)) +} + +func TestTranslateDomainPermissionNotification(t *testing.T) { + t.Parallel() + ev := pubsub.Event[permission.PermissionNotification]{ + Payload: permission.PermissionNotification{Granted: true}, + } + assert.Equal(t, PermissionResolved{}, Translate(ev)) +} + +// Proto type translation. + +func TestTranslateProtoAssistantMessage(t *testing.T) { + t.Parallel() + ev := pubsub.Event[proto.Message]{ + Payload: proto.Message{Role: proto.Assistant, SessionID: "s1"}, + } + assert.Equal(t, AssistantMessage{SessionID: "s1"}, Translate(ev)) +} + +func TestTranslateProtoNonAssistantIgnored(t *testing.T) { + t.Parallel() + ev := pubsub.Event[proto.Message]{ + Payload: proto.Message{Role: proto.User, SessionID: "s1"}, + } + assert.Nil(t, Translate(ev)) +} + +func TestTranslateProtoRunComplete(t *testing.T) { + t.Parallel() + ev := pubsub.Event[proto.RunComplete]{ + Payload: proto.RunComplete{SessionID: "s1"}, + } + assert.Equal(t, RunComplete{SessionID: "s1"}, Translate(ev)) +} + +func TestTranslateProtoPermissionRequest(t *testing.T) { + t.Parallel() + ev := pubsub.Event[proto.PermissionRequest]{ + Payload: proto.PermissionRequest{ToolName: "bash"}, + } + assert.Equal(t, PermissionRequested{}, Translate(ev)) +} + +func TestTranslateProtoPermissionNotification(t *testing.T) { + t.Parallel() + ev := pubsub.Event[proto.PermissionNotification]{ + Payload: proto.PermissionNotification{Granted: true}, + } + assert.Equal(t, PermissionResolved{}, Translate(ev)) +} + +func TestTranslateProtoSummarizing(t *testing.T) { + t.Parallel() + ev := pubsub.Event[proto.AgentEvent]{ + Payload: proto.AgentEvent{ + Type: proto.AgentEventTypeSummarize, + Done: false, + }, + } + assert.Equal(t, Summarizing{}, Translate(ev)) +} + +func TestTranslateProtoSummarizeDoneIgnored(t *testing.T) { + t.Parallel() + ev := pubsub.Event[proto.AgentEvent]{ + Payload: proto.AgentEvent{ + Type: proto.AgentEventTypeSummarize, + Done: true, + }, + } + assert.Nil(t, Translate(ev)) +} + +func TestTranslateProtoSessionIgnored(t *testing.T) { + t.Parallel() + ev := pubsub.Event[proto.Session]{ + Payload: proto.Session{ID: "s1"}, + } + assert.Nil(t, Translate(ev)) +} + +// Unknown types. + +func TestTranslateUnknownReturnsNil(t *testing.T) { + t.Parallel() + assert.Nil(t, Translate("not an event")) +} diff --git a/internal/history/file.go b/internal/history/file.go new file mode 100644 index 0000000000000000000000000000000000000000..80e713f3b291f9695381f36e00f8a4f36330bff3 --- /dev/null +++ b/internal/history/file.go @@ -0,0 +1,217 @@ +package history + +import ( + "context" + "database/sql" + "fmt" + "strings" + + "github.com/charmbracelet/crush/internal/db" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/google/uuid" +) + +const ( + InitialVersion = 0 +) + +type File struct { + ID string + SessionID string + Path string + Content string + Version int64 + CreatedAt int64 + UpdatedAt int64 +} + +// Service manages file versions and history for sessions. +type Service interface { + pubsub.Subscriber[File] + Create(ctx context.Context, sessionID, path, content string) (File, error) + + // CreateVersion creates a new version of a file. + CreateVersion(ctx context.Context, sessionID, path, content string) (File, error) + + Get(ctx context.Context, id string) (File, error) + GetByPathAndSession(ctx context.Context, path, sessionID string) (File, error) + ListBySession(ctx context.Context, sessionID string) ([]File, error) + ListLatestSessionFiles(ctx context.Context, sessionID string) ([]File, error) + Delete(ctx context.Context, id string) error + DeleteSessionFiles(ctx context.Context, sessionID string) error +} + +type service struct { + *pubsub.Broker[File] + db *sql.DB + q *db.Queries +} + +func NewService(q *db.Queries, db *sql.DB) Service { + return &service{ + Broker: pubsub.NewBroker[File](), + q: q, + db: db, + } +} + +func (s *service) Create(ctx context.Context, sessionID, path, content string) (File, error) { + return s.createWithVersion(ctx, sessionID, path, content, InitialVersion) +} + +// CreateVersion creates a new version of a file with auto-incremented version +// number. If no previous versions exist for the path, it creates the initial +// version. The provided content is stored as the new version. +func (s *service) CreateVersion(ctx context.Context, sessionID, path, content string) (File, error) { + // Get the latest version for this path + files, err := s.q.ListFilesByPath(ctx, path) + if err != nil { + return File{}, err + } + + if len(files) == 0 { + // No previous versions, create initial + return s.Create(ctx, sessionID, path, content) + } + + // Get the latest version + latestFile := files[0] // Files are ordered by version DESC, created_at DESC + nextVersion := latestFile.Version + 1 + + return s.createWithVersion(ctx, sessionID, path, content, nextVersion) +} + +func (s *service) createWithVersion(ctx context.Context, sessionID, path, content string, version int64) (File, error) { + // Maximum number of retries for transaction conflicts + const maxRetries = 3 + var file File + var err error + + // Retry loop for transaction conflicts + for attempt := range maxRetries { + // Start a transaction + tx, txErr := s.db.BeginTx(ctx, nil) + if txErr != nil { + return File{}, fmt.Errorf("failed to begin transaction: %w", txErr) + } + + // Create a new queries instance with the transaction + qtx := s.q.WithTx(tx) + + // Try to create the file within the transaction + dbFile, txErr := qtx.CreateFile(ctx, db.CreateFileParams{ + ID: uuid.New().String(), + SessionID: sessionID, + Path: path, + Content: content, + Version: version, + }) + if txErr != nil { + // Rollback the transaction + tx.Rollback() + + // Check if this is a uniqueness constraint violation + if strings.Contains(txErr.Error(), "UNIQUE constraint failed") { + if attempt < maxRetries-1 { + // If we have retries left, increment version and try again + version++ + continue + } + } + return File{}, txErr + } + + // Commit the transaction + if txErr = tx.Commit(); txErr != nil { + return File{}, fmt.Errorf("failed to commit transaction: %w", txErr) + } + + file = s.fromDBItem(dbFile) + s.Publish(pubsub.CreatedEvent, file) + return file, nil + } + + return file, err +} + +func (s *service) Get(ctx context.Context, id string) (File, error) { + dbFile, err := s.q.GetFile(ctx, id) + if err != nil { + return File{}, err + } + return s.fromDBItem(dbFile), nil +} + +func (s *service) GetByPathAndSession(ctx context.Context, path, sessionID string) (File, error) { + dbFile, err := s.q.GetFileByPathAndSession(ctx, db.GetFileByPathAndSessionParams{ + Path: path, + SessionID: sessionID, + }) + if err != nil { + return File{}, err + } + return s.fromDBItem(dbFile), nil +} + +func (s *service) ListBySession(ctx context.Context, sessionID string) ([]File, error) { + dbFiles, err := s.q.ListFilesBySession(ctx, sessionID) + if err != nil { + return nil, err + } + files := make([]File, len(dbFiles)) + for i, dbFile := range dbFiles { + files[i] = s.fromDBItem(dbFile) + } + return files, nil +} + +func (s *service) ListLatestSessionFiles(ctx context.Context, sessionID string) ([]File, error) { + dbFiles, err := s.q.ListLatestSessionFiles(ctx, sessionID) + if err != nil { + return nil, err + } + files := make([]File, len(dbFiles)) + for i, dbFile := range dbFiles { + files[i] = s.fromDBItem(dbFile) + } + return files, nil +} + +func (s *service) Delete(ctx context.Context, id string) error { + file, err := s.Get(ctx, id) + if err != nil { + return err + } + err = s.q.DeleteFile(ctx, id) + if err != nil { + return err + } + s.Publish(pubsub.DeletedEvent, file) + return nil +} + +func (s *service) DeleteSessionFiles(ctx context.Context, sessionID string) error { + files, err := s.ListBySession(ctx, sessionID) + if err != nil { + return err + } + for _, file := range files { + err = s.Delete(ctx, file.ID) + if err != nil { + return err + } + } + return nil +} + +func (s *service) fromDBItem(item db.File) File { + return File{ + ID: item.ID, + SessionID: item.SessionID, + Path: item.Path, + Content: item.Content, + Version: item.Version, + CreatedAt: item.CreatedAt, + UpdatedAt: item.UpdatedAt, + } +} diff --git a/internal/hooks/hooks.go b/internal/hooks/hooks.go new file mode 100644 index 0000000000000000000000000000000000000000..3529dc2a95b378b85b29f21a7e71ae34432d1306 --- /dev/null +++ b/internal/hooks/hooks.go @@ -0,0 +1,193 @@ +// Package hooks runs user-defined shell commands that fire on hook events +// (e.g. PreToolUse), returning decisions that control agent behavior. +package hooks + +import ( + "encoding/json" + "log/slog" + "strings" + + "github.com/tidwall/sjson" +) + +// Hook event name constants. +const ( + EventPreToolUse = "PreToolUse" +) + +// HaltExitCode is the exit code that halts the whole turn. 2 blocks the +// current tool call; 49 sits in the no-man's-land between the +// generic-error range (1-30), the sysexits range (64-78), and the +// killed-by-signal range (128+) so it can't be hit by accident. +const HaltExitCode = 49 + +// HookMetadata is embedded in tool response metadata so the UI can +// display a hook indicator. +type HookMetadata struct { + HookCount int `json:"hook_count"` + Decision string `json:"decision"` + Halt bool `json:"halt,omitempty"` + Reason string `json:"reason,omitempty"` + InputRewrite bool `json:"input_rewrite,omitempty"` + Hooks []HookInfo `json:"hooks,omitempty"` +} + +// HookInfo identifies a single hook that ran and its individual result. +type HookInfo struct { + Name string `json:"name"` + Matcher string `json:"matcher,omitempty"` + Decision string `json:"decision"` + Halt bool `json:"halt,omitempty"` + Reason string `json:"reason,omitempty"` + InputRewrite bool `json:"input_rewrite,omitempty"` +} + +// Decision represents the outcome of a single hook execution. +type Decision int + +const ( + // DecisionNone means the hook expressed no opinion. + DecisionNone Decision = iota + // DecisionAllow means the hook explicitly allowed the action. + DecisionAllow + // DecisionDeny means the hook blocked the action. + DecisionDeny +) + +func (d Decision) String() string { + switch d { + case DecisionAllow: + return "allow" + case DecisionDeny: + return "deny" + default: + return "none" + } +} + +// HookResult holds the parsed output of a single hook execution. +type HookResult struct { + Decision Decision + Halt bool // If true, halt the whole turn. + Reason string // Deny or halt reason (same field, different audience). + Context string + UpdatedInput string // Shallow-merge patch against tool_input (opaque JSON). +} + +// AggregateResult holds the combined outcome of all hooks for an event. +type AggregateResult struct { + Decision Decision + Halt bool // Any hook requested halt. + HookCount int // Number of hooks that ran. + Hooks []HookInfo // Info about each hook that ran (config order). + Reason string // Concatenated deny/halt reasons (newline-separated). + Context string // Concatenated context from all hooks. + UpdatedInput string // Merged tool_input JSON (empty if no patches). +} + +// aggregate merges multiple HookResults into a single AggregateResult. +// Results are processed in config order (the order of the slice). Deny +// wins over allow, allow wins over none. Halt is sticky. Reasons and +// context concatenate in order. updated_input patches shallow-merge in +// order against the original tool input; later patches override earlier +// ones on colliding keys. +func aggregate(results []HookResult, origToolInput string) AggregateResult { + var ( + decision Decision + halt bool + reasons []string + contexts []string + merged = origToolInput + anyPatch = false + ) + for _, r := range results { + switch r.Decision { + case DecisionDeny: + decision = DecisionDeny + if r.Reason != "" { + reasons = append(reasons, r.Reason) + } + case DecisionAllow: + if decision != DecisionDeny { + decision = DecisionAllow + } + case DecisionNone: + // No change. + } + if r.Halt { + halt = true + if r.Reason != "" && r.Decision != DecisionDeny { + // A halting hook that didn't also deny still contributes + // its reason so the user sees it. + reasons = append(reasons, r.Reason) + } + } + if r.Context != "" { + contexts = append(contexts, r.Context) + } + if r.UpdatedInput != "" { + next, err := shallowMerge(merged, r.UpdatedInput) + if err != nil { + slog.Warn( + "Hook updated_input patch rejected; ignoring", + "error", err, + "patch", r.UpdatedInput, + ) + continue + } + merged = next + anyPatch = true + } + } + + agg := AggregateResult{ + Decision: decision, + Halt: halt, + HookCount: len(results), + } + if anyPatch { + agg.UpdatedInput = merged + } + if len(reasons) > 0 { + agg.Reason = strings.Join(reasons, "\n") + } + if len(contexts) > 0 { + agg.Context = strings.Join(contexts, "\n") + } + return agg +} + +// shallowMerge applies a top-level-keys patch to base (both JSON +// objects). Keys in patch overwrite keys in base; keys absent from the +// patch are preserved. Returns an error if either value is not a valid +// JSON object. +func shallowMerge(base, patch string) (string, error) { + if base == "" { + base = "{}" + } + // Ensure base is an object so sjson has somewhere to write. + var baseAny any + if err := json.Unmarshal([]byte(base), &baseAny); err != nil { + return "", err + } + if _, ok := baseAny.(map[string]any); !ok { + return "", errNotObject("tool_input") + } + var patchMap map[string]json.RawMessage + if err := json.Unmarshal([]byte(patch), &patchMap); err != nil { + return "", errNotObject("updated_input") + } + out := base + for k, v := range patchMap { + next, err := sjson.SetRawBytes([]byte(out), k, v) + if err != nil { + return "", err + } + out = string(next) + } + return out, nil +} + +type errNotObject string + +func (e errNotObject) Error() string { return string(e) + " is not a JSON object" } diff --git a/internal/hooks/hooks_test.go b/internal/hooks/hooks_test.go new file mode 100644 index 0000000000000000000000000000000000000000..df689c3fa0c443cafb48a1f978b6f9a03386318a --- /dev/null +++ b/internal/hooks/hooks_test.go @@ -0,0 +1,775 @@ +package hooks + +import ( + "context" + "io" + "strings" + "sync" + "testing" + "time" + + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/shell" + "github.com/stretchr/testify/require" +) + +func TestAggregation(t *testing.T) { + t.Parallel() + + t.Run("empty results", func(t *testing.T) { + t.Parallel() + agg := aggregate(nil, "{}") + require.Equal(t, DecisionNone, agg.Decision) + require.Empty(t, agg.Reason) + require.Empty(t, agg.Context) + require.False(t, agg.Halt) + }) + + t.Run("single allow", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionAllow}, + }, "{}") + require.Equal(t, DecisionAllow, agg.Decision) + }) + + t.Run("deny wins over allow", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionAllow, Context: "ctx1"}, + {Decision: DecisionDeny, Reason: "blocked"}, + }, "{}") + require.Equal(t, DecisionDeny, agg.Decision) + require.Equal(t, "blocked", agg.Reason) + require.Equal(t, "ctx1", agg.Context) + }) + + t.Run("multiple deny reasons concatenated", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionDeny, Reason: "reason1"}, + {Decision: DecisionDeny, Reason: "reason2"}, + }, "{}") + require.Equal(t, DecisionDeny, agg.Decision) + require.Equal(t, "reason1\nreason2", agg.Reason) + }) + + t.Run("context concatenated from all hooks", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionAllow, Context: "ctx-a"}, + {Decision: DecisionNone, Context: "ctx-b"}, + }, "{}") + require.Equal(t, DecisionAllow, agg.Decision) + require.Equal(t, "ctx-a\nctx-b", agg.Context) + }) + + t.Run("allow wins over none", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionNone}, + {Decision: DecisionAllow}, + }, "{}") + require.Equal(t, DecisionAllow, agg.Decision) + }) + + t.Run("halt is sticky across results", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionAllow}, + {Halt: true, Reason: "stop now"}, + }, "{}") + require.True(t, agg.Halt) + require.Contains(t, agg.Reason, "stop now") + }) + + t.Run("halt with deny only records reason once", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionDeny, Halt: true, Reason: "stop"}, + }, "{}") + require.True(t, agg.Halt) + require.Equal(t, DecisionDeny, agg.Decision) + require.Equal(t, "stop", agg.Reason) + }) +} + +func TestParseStdout(t *testing.T) { + t.Parallel() + + t.Run("empty stdout", func(t *testing.T) { + t.Parallel() + r := parseStdout("") + require.Equal(t, DecisionNone, r.Decision) + }) + + t.Run("valid allow", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"decision":"allow","context":"some context"}`) + require.Equal(t, DecisionAllow, r.Decision) + require.Equal(t, "some context", r.Context) + }) + + t.Run("valid deny", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"decision":"deny","reason":"not allowed"}`) + require.Equal(t, DecisionDeny, r.Decision) + require.Equal(t, "not allowed", r.Reason) + }) + + t.Run("malformed JSON", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{bad json}`) + require.Equal(t, DecisionNone, r.Decision) + }) + + t.Run("unknown decision", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"decision":"maybe"}`) + require.Equal(t, DecisionNone, r.Decision) + }) + + t.Run("version 1 accepted", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"version":1,"decision":"allow"}`) + require.Equal(t, DecisionAllow, r.Decision) + }) + + t.Run("unknown higher version still parses", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"version":99,"decision":"deny","reason":"future"}`) + require.Equal(t, DecisionDeny, r.Decision) + require.Equal(t, "future", r.Reason) + }) + + t.Run("halt true without decision", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"halt":true,"reason":"turn over"}`) + require.True(t, r.Halt) + require.Equal(t, "turn over", r.Reason) + require.Equal(t, DecisionNone, r.Decision) + }) + + t.Run("context string form", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"decision":"allow","context":"one note"}`) + require.Equal(t, "one note", r.Context) + }) + + t.Run("context array form", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"decision":"allow","context":["first","second"]}`) + require.Equal(t, "first\nsecond", r.Context) + }) + + t.Run("context array drops empty entries", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"decision":"allow","context":["","keep",""]}`) + require.Equal(t, "keep", r.Context) + }) + + t.Run("context null becomes empty", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"decision":"allow","context":null}`) + require.Empty(t, r.Context) + }) +} + +func TestBuildEnv(t *testing.T) { + t.Parallel() + + env := BuildEnv(EventPreToolUse, "bash", "sess-1", "/work", "/project", `{"command":"ls","file_path":"/tmp/f.txt"}`) + + envMap := make(map[string]string) + for _, e := range env { + parts := splitFirst(e, "=") + if len(parts) == 2 { + envMap[parts[0]] = parts[1] + } + } + + require.Equal(t, EventPreToolUse, envMap["CRUSH_EVENT"]) + require.Equal(t, "bash", envMap["CRUSH_TOOL_NAME"]) + require.Equal(t, "sess-1", envMap["CRUSH_SESSION_ID"]) + require.Equal(t, "/work", envMap["CRUSH_CWD"]) + require.Equal(t, "/project", envMap["CRUSH_PROJECT_DIR"]) + require.Equal(t, "ls", envMap["CRUSH_TOOL_INPUT_COMMAND"]) + require.Equal(t, "/tmp/f.txt", envMap["CRUSH_TOOL_INPUT_FILE_PATH"]) + + // Shared Crush markers must be present so hook-authored scripts can + // detect they're running under Crush the same way bash-tool-invoked + // scripts can. + require.Equal(t, "1", envMap["CRUSH"]) + require.Equal(t, "crush", envMap["AGENT"]) + require.Equal(t, "crush", envMap["AI_AGENT"]) +} + +func splitFirst(s, sep string) []string { + before, after, found := strings.Cut(s, sep) + if !found { + return []string{s} + } + return []string{before, after} +} + +func TestBuildPayload(t *testing.T) { + t.Parallel() + payload := BuildPayload(EventPreToolUse, "sess-1", "/work", "bash", `{"command":"ls"}`) + s := string(payload) + require.Contains(t, s, `"event":"`+EventPreToolUse+`"`) + require.Contains(t, s, `"tool_name":"bash"`) + // tool_input should be an object, not a string. + require.Contains(t, s, `"tool_input":{"command":"ls"}`) +} + +func TestRunnerExitCode0Allow(t *testing.T) { + t.Parallel() + hookCfg := config.HookConfig{ + Command: `echo '{"decision":"allow","context":"ok"}'`, + } + r := NewRunner([]config.HookConfig{hookCfg}, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionAllow, result.Decision) + require.Equal(t, "ok", result.Context) +} + +func TestRunnerExitCode2Deny(t *testing.T) { + t.Parallel() + hookCfg := config.HookConfig{ + Command: `echo "forbidden" >&2; exit 2`, + } + r := NewRunner([]config.HookConfig{hookCfg}, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionDeny, result.Decision) + require.False(t, result.Halt) + require.Equal(t, "forbidden", result.Reason) +} + +func TestRunnerExitCode49Halt(t *testing.T) { + t.Parallel() + hookCfg := config.HookConfig{ + Command: `echo "stop the turn" >&2; exit 49`, + } + r := NewRunner([]config.HookConfig{hookCfg}, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.True(t, result.Halt) + require.Equal(t, DecisionDeny, result.Decision) + require.Equal(t, "stop the turn", result.Reason) +} + +func TestRunnerHaltViaJSON(t *testing.T) { + t.Parallel() + hookCfg := config.HookConfig{ + Command: `echo '{"halt":true,"reason":"via json"}'`, + } + r := NewRunner([]config.HookConfig{hookCfg}, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.True(t, result.Halt) + require.Equal(t, "via json", result.Reason) +} + +func TestRunnerExitCodeOtherNonBlocking(t *testing.T) { + t.Parallel() + hookCfg := config.HookConfig{ + Command: `exit 1`, + } + r := NewRunner([]config.HookConfig{hookCfg}, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionNone, result.Decision) +} + +func TestRunnerTimeout(t *testing.T) { + t.Parallel() + hookCfg := config.HookConfig{ + Command: `sleep 10`, + Timeout: 1, + } + r := NewRunner([]config.HookConfig{hookCfg}, t.TempDir(), t.TempDir()) + start := time.Now() + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + elapsed := time.Since(start) + require.NoError(t, err) + require.Equal(t, DecisionNone, result.Decision) + require.Less(t, elapsed, 5*time.Second) +} + +func TestRunnerDeduplication(t *testing.T) { + t.Parallel() + // Two hooks with the same command should only run once. + hookCfg := config.HookConfig{ + Command: `echo '{"decision":"allow"}'`, + } + r := NewRunner([]config.HookConfig{hookCfg, hookCfg}, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionAllow, result.Decision) +} + +func TestRunnerNoMatchingHooks(t *testing.T) { + t.Parallel() + // Hooks are empty. + r := NewRunner(nil, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionNone, result.Decision) +} + +// validatedHooks builds hook configs and runs ValidateHooks to compile +// matcher regexes, mirroring the real config-load path. +func validatedHooks(t *testing.T, hooks []config.HookConfig) []config.HookConfig { + t.Helper() + cfg := &config.Config{ + Hooks: map[string][]config.HookConfig{ + EventPreToolUse: hooks, + }, + } + require.NoError(t, cfg.ValidateHooks()) + return cfg.Hooks[EventPreToolUse] +} + +func TestRunnerMatcherFiltering(t *testing.T) { + t.Parallel() + + t.Run("compiled regex matches", func(t *testing.T) { + t.Parallel() + hooks := validatedHooks(t, []config.HookConfig{ + {Command: `echo '{"decision":"deny","reason":"blocked"}'`, Matcher: "^bash$"}, + }) + r := NewRunner(hooks, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionDeny, result.Decision) + }) + + t.Run("compiled regex does not match", func(t *testing.T) { + t.Parallel() + hooks := validatedHooks(t, []config.HookConfig{ + {Command: `echo '{"decision":"deny","reason":"blocked"}'`, Matcher: "^edit$"}, + }) + r := NewRunner(hooks, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionNone, result.Decision) + }) + + t.Run("no matcher matches everything", func(t *testing.T) { + t.Parallel() + hooks := validatedHooks(t, []config.HookConfig{ + {Command: `echo '{"decision":"allow"}'`}, + }) + r := NewRunner(hooks, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionAllow, result.Decision) + }) + + t.Run("partial regex match", func(t *testing.T) { + t.Parallel() + hooks := validatedHooks(t, []config.HookConfig{ + {Command: `echo '{"decision":"deny","reason":"mcp blocked"}'`, Matcher: "^mcp_"}, + }) + r := NewRunner(hooks, t.TempDir(), t.TempDir()) + + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "mcp_github_get_me", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionDeny, result.Decision) + + result, err = r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionNone, result.Decision) + }) + + // Runner must compile matchers itself; it cannot rely on + // ValidateHooks having run first. This is the guarantee that prevents + // the reload-drops-matcher class of bug. + t.Run("runner compiles matcher without ValidateHooks", func(t *testing.T) { + t.Parallel() + raw := []config.HookConfig{ + {Command: `echo '{"decision":"deny","reason":"blocked"}'`, Matcher: "^bash$"}, + } + r := NewRunner(raw, t.TempDir(), t.TempDir()) + + deny, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionDeny, deny.Decision) + + noop, err := r.Run(context.Background(), EventPreToolUse, "sess", "view", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionNone, noop.Decision) + }) + + // A matcher that fails to compile at Runner construction must not + // degrade to match-everything; the hook is dropped instead. + t.Run("runner skips hooks with invalid matcher", func(t *testing.T) { + t.Parallel() + raw := []config.HookConfig{ + {Command: `echo '{"decision":"deny","reason":"should not fire"}'`, Matcher: "[invalid"}, + } + r := NewRunner(raw, t.TempDir(), t.TempDir()) + + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionNone, result.Decision) + require.Empty(t, r.Hooks()) + }) +} + +func TestValidateHooksInvalidRegex(t *testing.T) { + t.Parallel() + cfg := &config.Config{ + Hooks: map[string][]config.HookConfig{ + EventPreToolUse: { + {Command: "true", Matcher: "[invalid"}, + }, + }, + } + err := cfg.ValidateHooks() + require.Error(t, err) + require.Contains(t, err.Error(), "invalid matcher regex") +} + +func TestValidateHooksEmptyCommand(t *testing.T) { + t.Parallel() + cfg := &config.Config{ + Hooks: map[string][]config.HookConfig{ + EventPreToolUse: { + {Command: ""}, + }, + }, + } + err := cfg.ValidateHooks() + require.Error(t, err) + require.Contains(t, err.Error(), "command is required") +} + +func TestValidateHooksNormalizesEventNames(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input string + }{ + {"canonical", "PreToolUse"}, + {"lowercase", "pretooluse"}, + {"snake_case", "pre_tool_use"}, + {"upper_snake", "PRE_TOOL_USE"}, + {"mixed_case", "preToolUse"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + cfg := &config.Config{ + Hooks: map[string][]config.HookConfig{ + tt.input: { + {Command: "true"}, + }, + }, + } + require.NoError(t, cfg.ValidateHooks()) + require.Len(t, cfg.Hooks[EventPreToolUse], 1) + }) + } +} + +func TestRunnerHookNameUsesDisplayName(t *testing.T) { + t.Parallel() + + t.Run("name field is used when set", func(t *testing.T) { + t.Parallel() + hookCfg := config.HookConfig{ + Name: "my-hook", + Command: `echo '{"decision":"allow"}'`, + } + r := NewRunner([]config.HookConfig{hookCfg}, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionAllow, result.Decision) + require.Len(t, result.Hooks, 1) + require.Equal(t, "my-hook", result.Hooks[0].Name) + }) + + t.Run("command is used when name is empty", func(t *testing.T) { + t.Parallel() + hookCfg := config.HookConfig{ + Command: `echo '{"decision":"allow"}'`, + } + r := NewRunner([]config.HookConfig{hookCfg}, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionAllow, result.Decision) + require.Len(t, result.Hooks, 1) + require.Equal(t, `echo '{"decision":"allow"}'`, result.Hooks[0].Name) + }) +} + +func TestRunnerParallelExecution(t *testing.T) { + t.Parallel() + // Two hooks: one allows, one denies. Deny should win. + hooks := []config.HookConfig{ + {Command: `echo '{"decision":"allow","context":"hook1"}'`}, + {Command: `echo '{"decision":"deny","reason":"nope"}' ; exit 0`}, + } + r := NewRunner(hooks, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionDeny, result.Decision) + require.Equal(t, "nope", result.Reason) +} + +func TestRunnerEnvVarsPropagated(t *testing.T) { + t.Parallel() + hookCfg := config.HookConfig{ + Command: `printf '{"decision":"allow","context":"%s"}' "$CRUSH_TOOL_NAME"`, + } + r := NewRunner([]config.HookConfig{hookCfg}, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + require.NoError(t, err) + require.Equal(t, DecisionAllow, result.Decision) + require.Equal(t, "bash", result.Context) +} + +func TestParseStdoutUpdatedInput(t *testing.T) { + t.Parallel() + + t.Run("nested object", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"decision":"allow","updated_input":{"command":"rtk cat foo.go"}}`) + require.Equal(t, DecisionAllow, r.Decision) + require.Equal(t, `{"command":"rtk cat foo.go"}`, r.UpdatedInput) + }) + + t.Run("stringified backward compat", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"decision":"allow","updated_input":"{\"command\":\"rtk cat foo.go\"}"}`) + require.Equal(t, DecisionAllow, r.Decision) + require.Equal(t, `{"command":"rtk cat foo.go"}`, r.UpdatedInput) + }) + + t.Run("no updated_input", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"decision":"allow"}`) + require.Empty(t, r.UpdatedInput) + }) +} + +func TestAggregationUpdatedInput(t *testing.T) { + t.Parallel() + + t.Run("patches merge in config order with later overriding", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionAllow, UpdatedInput: `{"command":"first","keep":"me"}`}, + {Decision: DecisionAllow, UpdatedInput: `{"command":"second"}`}, + }, `{"command":"orig","timeout":60}`) + require.Equal(t, DecisionAllow, agg.Decision) + // command overridden by second patch; keep preserved from first + // patch; timeout preserved from original input. + require.JSONEq( + t, + `{"command":"second","keep":"me","timeout":60}`, + agg.UpdatedInput, + ) + }) + + t.Run("shallow: nested objects are replaced wholesale", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionAllow, UpdatedInput: `{"env":{"FOO":"bar"}}`}, + }, `{"env":{"BAZ":"qux"},"command":"ls"}`) + // "env" is replaced entirely; "command" preserved. + require.JSONEq( + t, + `{"env":{"FOO":"bar"},"command":"ls"}`, + agg.UpdatedInput, + ) + }) + + t.Run("deny still reports merged input (caller ignores it)", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionAllow, UpdatedInput: `{"command":"rewritten"}`}, + {Decision: DecisionDeny, Reason: "blocked"}, + }, `{"command":"orig"}`) + require.Equal(t, DecisionDeny, agg.Decision) + }) + + t.Run("no patches leaves updated_input empty", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionAllow}, + {Decision: DecisionNone}, + }, `{"command":"orig"}`) + require.Empty(t, agg.UpdatedInput) + }) + + t.Run("invalid patch is ignored", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionAllow, UpdatedInput: `"not-an-object"`}, + {Decision: DecisionAllow, UpdatedInput: `{"command":"good"}`}, + }, `{"command":"orig"}`) + require.JSONEq(t, `{"command":"good"}`, agg.UpdatedInput) + }) + + t.Run("malformed patch JSON is ignored and merge continues", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionAllow, UpdatedInput: `{broken json`}, + {Decision: DecisionAllow, UpdatedInput: `{"command":"good"}`}, + }, `{"command":"orig"}`) + require.JSONEq(t, `{"command":"good"}`, agg.UpdatedInput) + }) + + t.Run("non-object tool_input rejects all patches", func(t *testing.T) { + t.Parallel() + agg := aggregate([]HookResult{ + {Decision: DecisionAllow, UpdatedInput: `{"command":"rewrite"}`}, + }, `"just-a-string"`) + require.Empty(t, agg.UpdatedInput) + }) + + t.Run("null updated_input is a no-op", func(t *testing.T) { + t.Parallel() + // parseStdout converts null updated_input to "", so aggregate + // never sees a patch — the merged input is empty and the + // original tool_input is used unchanged. + r := parseStdout(`{"decision":"allow","updated_input":null}`) + require.Empty(t, r.UpdatedInput) + agg := aggregate([]HookResult{r}, `{"command":"orig"}`) + require.Empty(t, agg.UpdatedInput) + }) +} + +// TestRunnerAbandonRaceSafety verifies that if a hook's shell execution +// does not yield to ctx cancellation within abandonGrace, runOne returns +// promptly and never touches the shared stdout/stderr buffers again — +// even while the abandoned goroutine continues to write to them. +// +// The substitute shell executor ignores ctx entirely, writes to Stdout +// both before and after the abandon deadline, and only then returns. +// Under -race this catches any code path in runOne that reads those +// buffers after returning the DecisionNone abandon result. +// +// The abandoned goroutine outlives r.Run, so the test also gates its +// shutdown on the stub having been entered (started): wg.Add runs inside +// the stub, and under a loaded scheduler the goroutine may not be +// scheduled until after the test body returned. Waiting on started +// orders wg.Add (and the goroutine's read of the runShell variable) +// before the cleanup's wg.Wait and the cleanup restoring runShell; +// without it both pairs are an Add-after-Wait misuse that -race flags. +func TestRunnerAbandonRaceSafety(t *testing.T) { + origRunShell := runShell + t.Cleanup(func() { runShell = origRunShell }) + + // Synchronize shutdown with the abandoned goroutine so the test + // exits cleanly even under -race. + var wg sync.WaitGroup + started := make(chan struct{}) + release := make(chan struct{}) + t.Cleanup(func() { + close(release) + wg.Wait() + }) + + runShell = func(_ context.Context, opts shell.RunOptions) error { + wg.Add(1) + defer wg.Done() + close(started) + // Write before the caller observes ctx.Done(); the caller will + // not read the buffer while we still own it. + _, _ = io.WriteString(opts.Stdout, "before\n") + // Hold past ctx deadline + abandonGrace so the caller takes + // the abandon branch, then continue writing. If the caller + // reads these buffers after abandoning, -race will flag it. + select { + case <-time.After(5 * time.Second): + case <-release: + } + _, _ = io.WriteString(opts.Stdout, "after\n") + return nil + } + + hookCfg := config.HookConfig{ + Command: "# irrelevant; runShell is stubbed", + Timeout: 1, + } + r := NewRunner([]config.HookConfig{hookCfg}, t.TempDir(), t.TempDir()) + + start := time.Now() + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{}`) + elapsed := time.Since(start) + + require.NoError(t, err) + require.Equal(t, DecisionNone, result.Decision) + // Abandon must happen at ~timeout + abandonGrace. Allow generous + // slack so CI noise doesn't flake the test. + require.Less(t, elapsed, 3500*time.Millisecond, + "runOne should return within timeout+abandonGrace+slack") + + // Block until the abandoned goroutine has entered the stub. In the + // common case it was entered long before the abandon deadline and + // this returns immediately; in the pathological case the goroutine + // was starved for the whole abandon window, and this waits for it to + // run so the cleanups below stay ordered after its wg.Add and its + // read of the runShell variable. + <-started +} + +func TestRunnerUpdatedInput(t *testing.T) { + t.Parallel() + hookCfg := config.HookConfig{ + Command: `echo '{"decision":"allow","updated_input":{"command":"echo rewritten"}}'`, + } + r := NewRunner([]config.HookConfig{hookCfg}, t.TempDir(), t.TempDir()) + result, err := r.Run(context.Background(), EventPreToolUse, "sess", "bash", `{"command":"echo original","timeout":60}`) + require.NoError(t, err) + require.Equal(t, DecisionAllow, result.Decision) + require.JSONEq( + t, + `{"command":"echo rewritten","timeout":60}`, + result.UpdatedInput, + ) +} + +func TestParseStdoutClaudeCodeFormat(t *testing.T) { + t.Parallel() + + t.Run("allow with reason", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"hookSpecificOutput":{"permissionDecision":"allow","permissionDecisionReason":"RTK auto-rewrite"}}`) + require.Equal(t, DecisionAllow, r.Decision) + require.Equal(t, "RTK auto-rewrite", r.Reason) + }) + + t.Run("allow with updatedInput", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"hookSpecificOutput":{"permissionDecision":"allow","updatedInput":{"command":"rtk cat foo.go"}}}`) + require.Equal(t, DecisionAllow, r.Decision) + require.Equal(t, `{"command":"rtk cat foo.go"}`, r.UpdatedInput) + }) + + t.Run("deny", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"hookSpecificOutput":{"permissionDecision":"deny","permissionDecisionReason":"not allowed"}}`) + require.Equal(t, DecisionDeny, r.Decision) + require.Equal(t, "not allowed", r.Reason) + }) + + t.Run("no permissionDecision", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"hookSpecificOutput":{}}`) + require.Equal(t, DecisionNone, r.Decision) + }) + + t.Run("crush format still works", func(t *testing.T) { + t.Parallel() + r := parseStdout(`{"decision":"allow","context":"hello"}`) + require.Equal(t, DecisionAllow, r.Decision) + require.Equal(t, "hello", r.Context) + }) +} diff --git a/internal/hooks/input.go b/internal/hooks/input.go new file mode 100644 index 0000000000000000000000000000000000000000..946347466d5c724677ebcb7155321126b4e11147 --- /dev/null +++ b/internal/hooks/input.go @@ -0,0 +1,211 @@ +package hooks + +import ( + "encoding/json" + "fmt" + "log/slog" + "os" + "strings" + + "github.com/charmbracelet/crush/internal/shell" + "github.com/tidwall/gjson" +) + +// SupportedOutputVersion is the highest envelope version this build +// understands. Hooks may omit `version` entirely (treated as 1) or pin +// an older version. Unknown higher versions are still parsed but logged. +const SupportedOutputVersion = 1 + +// Payload is the JSON structure piped to hook commands via stdin. +// ToolInput is emitted as a parsed JSON object for compatibility with +// Claude Code hooks (which expect tool_input to be an object, not a +// string). +type Payload struct { + Event string `json:"event"` + SessionID string `json:"session_id"` + CWD string `json:"cwd"` + ToolName string `json:"tool_name"` + ToolInput json.RawMessage `json:"tool_input"` +} + +// BuildPayload constructs the JSON stdin payload for a hook command. +func BuildPayload(eventName, sessionID, cwd, toolName, toolInputJSON string) []byte { + toolInput := json.RawMessage(toolInputJSON) + if !json.Valid(toolInput) { + toolInput = json.RawMessage("{}") + } + p := Payload{ + Event: eventName, + SessionID: sessionID, + CWD: cwd, + ToolName: toolName, + ToolInput: toolInput, + } + data, err := json.Marshal(p) + if err != nil { + return []byte("{}") + } + return data +} + +// BuildEnv constructs the environment variable slice for a hook command. +// It includes all current process env vars plus hook-specific ones. +func BuildEnv(eventName, toolName, sessionID, cwd, projectDir, toolInputJSON string) []string { + env := os.Environ() + env = append(env, shell.CrushEnvMarkers()...) + env = append( + env, + fmt.Sprintf("CRUSH_EVENT=%s", eventName), + fmt.Sprintf("CRUSH_TOOL_NAME=%s", toolName), + fmt.Sprintf("CRUSH_SESSION_ID=%s", sessionID), + fmt.Sprintf("CRUSH_CWD=%s", cwd), + fmt.Sprintf("CRUSH_PROJECT_DIR=%s", projectDir), + ) + + // Extract tool-specific env vars from the JSON input. + if toolInputJSON != "" { + if cmd := gjson.Get(toolInputJSON, "command"); cmd.Exists() { + env = append(env, fmt.Sprintf("CRUSH_TOOL_INPUT_COMMAND=%s", cmd.String())) + } + if fp := gjson.Get(toolInputJSON, "file_path"); fp.Exists() { + env = append(env, fmt.Sprintf("CRUSH_TOOL_INPUT_FILE_PATH=%s", fp.String())) + } + } + + return env +} + +// parseStdout parses the JSON output from a hook command's stdout. +// Supports both Crush format and Claude Code format (hookSpecificOutput). +func parseStdout(stdout string) HookResult { + stdout = strings.TrimSpace(stdout) + if stdout == "" { + return HookResult{Decision: DecisionNone} + } + + var raw map[string]json.RawMessage + if err := json.Unmarshal([]byte(stdout), &raw); err != nil { + return HookResult{Decision: DecisionNone} + } + + // Claude Code compat: if hookSpecificOutput is present, parse that. + if hso, ok := raw["hookSpecificOutput"]; ok { + return parseClaudeCodeOutput(hso) + } + + var parsed struct { + Version int `json:"version"` + Decision string `json:"decision"` + Halt bool `json:"halt"` + Reason string `json:"reason"` + Context json.RawMessage `json:"context"` + UpdatedInput json.RawMessage `json:"updated_input"` + } + if err := json.Unmarshal([]byte(stdout), &parsed); err != nil { + return HookResult{Decision: DecisionNone} + } + + if parsed.Version > SupportedOutputVersion { + slog.Debug( + "Hook output declared a newer envelope version than this build supports", + "version", parsed.Version, + "supported", SupportedOutputVersion, + ) + } + + result := HookResult{ + Halt: parsed.Halt, + Reason: parsed.Reason, + Context: parseContext(parsed.Context), + } + result.Decision = parseDecision(parsed.Decision) + result.UpdatedInput = rawToString(parsed.UpdatedInput) + return result +} + +// parseContext accepts either a single string or an array of strings and +// returns a newline-joined value with empty entries dropped. +func parseContext(raw json.RawMessage) string { + if len(raw) == 0 || string(raw) == "null" { + return "" + } + // String form. + if raw[0] == '"' { + var s string + if err := json.Unmarshal(raw, &s); err == nil { + return s + } + return "" + } + // Array form. + if raw[0] == '[' { + var items []string + if err := json.Unmarshal(raw, &items); err != nil { + return "" + } + out := items[:0] + for _, s := range items { + if s != "" { + out = append(out, s) + } + } + return strings.Join(out, "\n") + } + return "" +} + +// parseClaudeCodeOutput handles the Claude Code hook output format: +// {"hookSpecificOutput": {"permissionDecision": "allow", ...}} +func parseClaudeCodeOutput(data json.RawMessage) HookResult { + var hso struct { + PermissionDecision string `json:"permissionDecision"` + PermissionDecisionReason string `json:"permissionDecisionReason"` + UpdatedInput json.RawMessage `json:"updatedInput"` + AdditionalContext string `json:"additionalContext"` + } + if err := json.Unmarshal(data, &hso); err != nil { + return HookResult{Decision: DecisionNone} + } + + result := HookResult{ + Decision: parseDecision(hso.PermissionDecision), + Reason: hso.PermissionDecisionReason, + Context: hso.AdditionalContext, + } + + // Marshal updatedInput back to a string for our opaque format. + if len(hso.UpdatedInput) > 0 && string(hso.UpdatedInput) != "null" { + result.UpdatedInput = string(hso.UpdatedInput) + } + + return result +} + +// rawToString converts a json.RawMessage to a string suitable for use +// as opaque tool input. It accepts both a JSON object (nested) and a +// JSON string (stringified, for backward compatibility). +func rawToString(raw json.RawMessage) string { + if len(raw) == 0 || string(raw) == "null" { + return "" + } + // If it's a JSON string, unwrap it. + if raw[0] == '"' { + var s string + if err := json.Unmarshal(raw, &s); err == nil { + return s + } + } + // Otherwise it's an object/array — use as-is. + return string(raw) +} + +func parseDecision(s string) Decision { + switch strings.ToLower(s) { + case "allow": + return DecisionAllow + case "deny": + return DecisionDeny + default: + return DecisionNone + } +} diff --git a/internal/hooks/runner.go b/internal/hooks/runner.go new file mode 100644 index 0000000000000000000000000000000000000000..734fd3e5ca2078a9ac0143a272538542d77f42a7 --- /dev/null +++ b/internal/hooks/runner.go @@ -0,0 +1,265 @@ +package hooks + +import ( + "bytes" + "context" + "log/slog" + "regexp" + "strings" + "sync" + "time" + + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/shell" +) + +// abandonGrace is how long runOne waits after ctx cancellation for the +// shell goroutine to yield before returning control to the caller and +// letting the goroutine finish on its own. Mirrors the historical +// cmd.WaitDelay = time.Second behavior of the previous os/exec path. +const abandonGrace = time.Second + +// runShell is the shell executor used by runOne. It is a package-level +// variable so tests can substitute a blocking or non-yielding +// implementation to exercise the abandon-on-timeout path without +// depending on the scheduling behavior of the real interpreter. +var runShell = shell.Run + +// compiledHook pairs a HookConfig with its compiled matcher regex. A nil +// matcher means "match every tool". +type compiledHook struct { + cfg config.HookConfig + matcher *regexp.Regexp +} + +// Runner executes hook commands and aggregates their results. +type Runner struct { + hooks []compiledHook + cwd string + projectDir string +} + +// NewRunner creates a Runner from the given hook configs. Each hook's +// Matcher is compiled here so the Runner is self-sufficient; callers do +// not have to pre-compile matchers on the config, and reloads or merges +// that rebuild HookConfig values can't silently strip compiled state. +// +// Hooks whose matcher fails to compile are skipped with a warning rather +// than treated as match-everything. ValidateHooks is expected to have +// caught syntax errors earlier, so this is defense in depth. +func NewRunner(hooks []config.HookConfig, cwd, projectDir string) *Runner { + compiled := make([]compiledHook, 0, len(hooks)) + for _, h := range hooks { + ch := compiledHook{cfg: h} + if h.Matcher != "" { + re, err := regexp.Compile(h.Matcher) + if err != nil { + slog.Warn( + "Hook matcher failed to compile; skipping hook", + "matcher", h.Matcher, + "command", h.Command, + "error", err, + ) + continue + } + ch.matcher = re + } + compiled = append(compiled, ch) + } + return &Runner{ + hooks: compiled, + cwd: cwd, + projectDir: projectDir, + } +} + +// Hooks returns the hook configs the runner was created with, in config +// order. Hooks whose matcher failed to compile at construction are +// omitted. Intended for diagnostics; callers should not rely on ordering +// or identity beyond that. +func (r *Runner) Hooks() []config.HookConfig { + out := make([]config.HookConfig, len(r.hooks)) + for i, h := range r.hooks { + out[i] = h.cfg + } + return out +} + +// Run executes all matching hooks for the given event and tool, returning +// an aggregated result. +func (r *Runner) Run(ctx context.Context, eventName, sessionID, toolName, toolInputJSON string) (AggregateResult, error) { + matching := r.matchingHooks(toolName) + if len(matching) == 0 { + return AggregateResult{Decision: DecisionNone}, nil + } + + // Deduplicate by command string. + seen := make(map[string]bool, len(matching)) + var deduped []config.HookConfig + for _, h := range matching { + if seen[h.Command] { + continue + } + seen[h.Command] = true + deduped = append(deduped, h) + } + + envVars := BuildEnv(eventName, toolName, sessionID, r.cwd, r.projectDir, toolInputJSON) + payload := BuildPayload(eventName, sessionID, r.cwd, toolName, toolInputJSON) + + results := make([]HookResult, len(deduped)) + var wg sync.WaitGroup + wg.Add(len(deduped)) + + for i, h := range deduped { + go func(idx int, hook config.HookConfig) { + defer wg.Done() + results[idx] = r.runOne(ctx, hook, envVars, payload) + }(i, h) + } + wg.Wait() + + agg := aggregate(results, toolInputJSON) + agg.Hooks = make([]HookInfo, len(deduped)) + for i, h := range deduped { + agg.Hooks[i] = HookInfo{ + Name: h.DisplayName(), + Matcher: h.Matcher, + Decision: results[i].Decision.String(), + Halt: results[i].Halt, + Reason: results[i].Reason, + InputRewrite: results[i].UpdatedInput != "", + } + } + slog.Info( + "Hook completed", + "event", eventName, + "tool", toolName, + "hooks", len(deduped), + "decision", agg.Decision.String(), + ) + return agg, nil +} + +// matchingHooks returns hooks whose matcher matches the tool name (or has +// no matcher, which matches everything). +func (r *Runner) matchingHooks(toolName string) []config.HookConfig { + var matched []config.HookConfig + for _, h := range r.hooks { + if h.matcher == nil || h.matcher.MatchString(toolName) { + matched = append(matched, h.cfg) + } + } + return matched +} + +// runOne executes a single hook command and returns its result. +// +// Execution goes through Crush's embedded POSIX shell (shell.Run) so the +// same interpreter, builtins, and coreutils are visible to hooks as to +// the bash tool. BlockFuncs are intentionally omitted: hooks are +// user-authored config that carry the same trust as a shell alias. +// +// A hook that fails to yield after its deadline has passed is abandoned +// after abandonGrace so the caller never blocks longer than +// timeout + abandonGrace. Ownership of the stdout and stderr buffers is +// strictly single-goroutine: +// - before receiving from `done`, only the goroutine writes to them; +// - after `done` delivers a value, the goroutine is finished and the +// outer frame reads them; +// - on the abandon path, the goroutine may still be writing and the +// outer frame must not touch them again. +func (r *Runner) runOne(parentCtx context.Context, hook config.HookConfig, envVars []string, payload []byte) HookResult { + timeout := hook.TimeoutDuration() + ctx, cancel := context.WithTimeout(parentCtx, timeout) + defer cancel() + + var stdout, stderr bytes.Buffer + done := make(chan error, 1) + go func() { + done <- runShell(ctx, shell.RunOptions{ + Command: hook.Command, + Cwd: r.cwd, + Env: envVars, + Stdin: bytes.NewReader(payload), + Stdout: &stdout, + Stderr: &stderr, + }) + }() + + var err error + select { + case err = <-done: + // Normal path: goroutine has finished, buffers are safe to read. + case <-ctx.Done(): + select { + case err = <-done: + // Interpreter yielded within the grace period; safe to read. + case <-time.After(abandonGrace): + slog.Warn( + "Hook did not yield after cancel; abandoning goroutine", + "command", hook.Command, + "timeout", timeout, + ) + // The goroutine may still be writing to stdout/stderr; do + // not read either buffer below this point. + return HookResult{Decision: DecisionNone} + } + } + + if shell.IsInterrupt(err) { + // Distinguish timeout from parent cancellation. + if parentCtx.Err() != nil { + slog.Debug("Hook cancelled by parent context", "command", hook.Command) + } else { + slog.Warn("Hook timed out", "command", hook.Command, "timeout", timeout) + } + return HookResult{Decision: DecisionNone} + } + + if err != nil { + exitCode := shell.ExitCode(err) + switch exitCode { + case 2: + // Exit code 2 = block this tool call. Stderr is the reason. + reason := strings.TrimSpace(stderr.String()) + if reason == "" { + reason = "blocked by hook" + } + return HookResult{ + Decision: DecisionDeny, + Reason: reason, + } + case HaltExitCode: + // Exit code 49 = halt the whole turn. Stderr is the reason. + reason := strings.TrimSpace(stderr.String()) + if reason == "" { + reason = "turn halted by hook" + } + return HookResult{ + Decision: DecisionDeny, + Halt: true, + Reason: reason, + } + default: + // Other non-zero exits are non-blocking errors. + slog.Warn( + "Hook failed with non-blocking error", + "command", hook.Command, + "exit_code", exitCode, + "stderr", strings.TrimSpace(stderr.String()), + "error", err, + ) + return HookResult{Decision: DecisionNone} + } + } + + // Exit code 0 — parse stdout JSON. + result := parseStdout(stdout.String()) + slog.Debug( + "Hook executed", + "command", hook.Command, + "decision", result.Decision.String(), + ) + return result +} diff --git a/internal/log/http.go b/internal/log/http.go new file mode 100644 index 0000000000000000000000000000000000000000..cf791366392c1e41d037633e57a3c6d52ce3f907 --- /dev/null +++ b/internal/log/http.go @@ -0,0 +1,131 @@ +package log + +import ( + "bytes" + "encoding/json" + "io" + "log/slog" + "net/http" + "strings" + "time" +) + +// NewHTTPClient creates an HTTP client with debug logging enabled when debug mode is on. +func NewHTTPClient() *http.Client { + return &http.Client{ + Transport: &HTTPRoundTripLogger{ + Transport: http.DefaultTransport, + }, + } +} + +// HTTPRoundTripLogger is an http.RoundTripper that logs requests and responses. +type HTTPRoundTripLogger struct { + Transport http.RoundTripper +} + +// RoundTrip implements http.RoundTripper interface with logging. +func (h *HTTPRoundTripLogger) RoundTrip(req *http.Request) (*http.Response, error) { + var err error + var save io.ReadCloser + save, req.Body, err = drainBody(req.Body) + if err != nil { + slog.Error( + "HTTP request failed", + "method", req.Method, + "url", req.URL, + "error", err, + ) + return nil, err + } + + if slog.Default().Enabled(req.Context(), slog.LevelDebug) { + slog.Debug( + "HTTP Request", + "method", req.Method, + "url", req.URL, + "body", bodyToString(save), + ) + } + + start := time.Now() + resp, err := h.Transport.RoundTrip(req) + duration := time.Since(start) + if err != nil { + slog.Error( + "HTTP request failed", + "method", req.Method, + "url", req.URL, + "duration_ms", duration.Milliseconds(), + "error", err, + ) + return resp, err + } + + save, resp.Body, err = drainBody(resp.Body) + if err != nil { + slog.Error("Failed to drain response body", "error", err) + return resp, err + } + if slog.Default().Enabled(req.Context(), slog.LevelDebug) { + slog.Debug( + "HTTP Response", + "status_code", resp.StatusCode, + "status", resp.Status, + "headers", formatHeaders(resp.Header), + "body", bodyToString(save), + "content_length", resp.ContentLength, + "duration_ms", duration.Milliseconds(), + ) + } + return resp, nil +} + +func bodyToString(body io.ReadCloser) string { + if body == nil { + return "" + } + src, err := io.ReadAll(body) + if err != nil { + slog.Error("Failed to read body", "error", err) + return "" + } + var b bytes.Buffer + if json.Indent(&b, bytes.TrimSpace(src), "", " ") != nil { + // not json probably + return string(src) + } + return b.String() +} + +// formatHeaders formats HTTP headers for logging, filtering out sensitive information. +func formatHeaders(headers http.Header) map[string][]string { + filtered := make(map[string][]string) + for key, values := range headers { + lowerKey := strings.ToLower(key) + // Filter out sensitive headers + if strings.Contains(lowerKey, "authorization") || + strings.Contains(lowerKey, "api-key") || + strings.Contains(lowerKey, "token") || + strings.Contains(lowerKey, "secret") { + filtered[key] = []string{"[REDACTED]"} + } else { + filtered[key] = values + } + } + return filtered +} + +func drainBody(b io.ReadCloser) (r1, r2 io.ReadCloser, err error) { + if b == nil || b == http.NoBody { + return http.NoBody, http.NoBody, nil + } + var buf bytes.Buffer + if _, err = buf.ReadFrom(b); err != nil { + return nil, b, err + } + if err = b.Close(); err != nil { + return nil, b, err + } + return io.NopCloser(&buf), io.NopCloser(bytes.NewReader(buf.Bytes())), nil +} diff --git a/internal/log/http_test.go b/internal/log/http_test.go new file mode 100644 index 0000000000000000000000000000000000000000..cf5a21185dce2ec2e84e1979d58047eb359103db --- /dev/null +++ b/internal/log/http_test.go @@ -0,0 +1,73 @@ +package log + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestHTTPRoundTripLogger(t *testing.T) { + // Create a test server that returns a 500 error + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Header().Set("X-Custom-Header", "test-value") + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error": "Internal server error", "code": 500}`)) + })) + defer server.Close() + + // Create HTTP client with logging + client := NewHTTPClient() + + // Make a request + req, err := http.NewRequestWithContext( + t.Context(), + http.MethodPost, + server.URL, + strings.NewReader(`{"test": "data"}`), + ) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer secret-token") + + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + // Verify response + if resp.StatusCode != http.StatusInternalServerError { + t.Errorf("Expected status code 500, got %d", resp.StatusCode) + } +} + +func TestFormatHeaders(t *testing.T) { + headers := http.Header{ + "Content-Type": []string{"application/json"}, + "Authorization": []string{"Bearer secret-token"}, + "X-API-Key": []string{"api-key-123"}, + "User-Agent": []string{"test-agent"}, + } + + formatted := formatHeaders(headers) + + // Check that sensitive headers are redacted + if formatted["Authorization"][0] != "[REDACTED]" { + t.Error("Authorization header should be redacted") + } + if formatted["X-API-Key"][0] != "[REDACTED]" { + t.Error("X-API-Key header should be redacted") + } + + // Check that non-sensitive headers are preserved + if formatted["Content-Type"][0] != "application/json" { + t.Error("Content-Type header should be preserved") + } + if formatted["User-Agent"][0] != "test-agent" { + t.Error("User-Agent header should be preserved") + } +} diff --git a/internal/log/log.go b/internal/log/log.go new file mode 100644 index 0000000000000000000000000000000000000000..54a0620f18debcf442ed3dd2f8c1739d6aec3696 --- /dev/null +++ b/internal/log/log.go @@ -0,0 +1,89 @@ +package log + +import ( + "fmt" + "io" + "log/slog" + "os" + "runtime/debug" + "sync" + "sync/atomic" + "time" + + "github.com/charmbracelet/crush/internal/event" + "github.com/charmbracelet/x/term" + "gopkg.in/natefinch/lumberjack.v2" +) + +var ( + initOnce sync.Once + initialized atomic.Bool +) + +func Setup(logFile string, debug bool, ws ...io.Writer) { + initOnce.Do(func() { + logRotator := &lumberjack.Logger{ + Filename: logFile, + MaxSize: 10, // Max size in MB + MaxBackups: 0, // Number of backups + MaxAge: 30, // Days + Compress: false, // Enable compression + } + + level := slog.LevelInfo + if debug { + level = slog.LevelDebug + } + + opts := &slog.HandlerOptions{ + Level: level, + AddSource: true, + } + + var handlers []slog.Handler + handlers = append(handlers, slog.NewJSONHandler(logRotator, opts)) + + for _, w := range ws { + if w == nil { + continue + } + if f, ok := w.(term.File); ok && term.IsTerminal(f.Fd()) { + handlers = append(handlers, slog.NewTextHandler(w, opts)) + } else { + handlers = append(handlers, slog.NewJSONHandler(w, opts)) + } + } + + slog.SetDefault(slog.New(slog.NewMultiHandler(handlers...))) + initialized.Store(true) + }) +} + +func Initialized() bool { + return initialized.Load() +} + +func RecoverPanic(name string, cleanup func()) { + if r := recover(); r != nil { + event.Error(r, "panic", true, "name", name) + + // Create a timestamped panic log file + timestamp := time.Now().Format("20060102-150405") + filename := fmt.Sprintf("crush-panic-%s-%s.log", name, timestamp) + + file, err := os.Create(filename) + if err == nil { + defer file.Close() + + // Write panic information and stack trace + fmt.Fprintf(file, "Panic in %s: %v\n\n", name, r) + fmt.Fprintf(file, "Time: %s\n\n", time.Now().Format(time.RFC3339)) + fmt.Fprintf(file, "Stack Trace:\n%s\n", debug.Stack()) + + // Execute cleanup function if provided + if cleanup != nil { + cleanup() + } + } + } +} diff --git a/internal/lsp/client.go b/internal/lsp/client.go new file mode 100644 index 0000000000000000000000000000000000000000..69049f298a38ffc2cc0bd78ae06358fc1a2d735b --- /dev/null +++ b/internal/lsp/client.go @@ -0,0 +1,769 @@ +package lsp + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "os" + "path/filepath" + "slices" + "sync" + "sync/atomic" + "time" + + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/csync" + "github.com/charmbracelet/crush/internal/fsext" + "github.com/charmbracelet/crush/internal/home" + powernap "github.com/charmbracelet/x/powernap/pkg/lsp" + "github.com/charmbracelet/x/powernap/pkg/lsp/protocol" + "github.com/charmbracelet/x/powernap/pkg/transport" +) + +// DiagnosticCounts holds the count of diagnostics by severity. +type DiagnosticCounts struct { + Error int + Warning int + Information int + Hint int +} + +type Client struct { + client *powernap.Client + name string + debug bool + + // Working directory this LSP is scoped to. + cwd string + + // File types this LSP server handles (e.g., .go, .rs, .py) + fileTypes []string + + // Configuration for this LSP client + config config.LSPConfig + + // Long-lived context for the client's lifetime, independent of any + // request-scoped context. Used for restart and other operations that + // must survive beyond the initial tool call that created the client. + ctx context.Context + cancelCtx context.CancelFunc + resolver config.VariableResolver + + // Diagnostic change callback + onDiagnosticsChanged func(name string, count int) + + // Diagnostic cache + diagnostics *csync.VersionedMap[protocol.DocumentURI, []protocol.Diagnostic] + + // Cached diagnostic counts to avoid map copy on every UI render. + diagCountsCache DiagnosticCounts + diagCountsVersion uint64 + diagCountsMu sync.Mutex + + // Files are currently opened by the LSP + openFiles *csync.Map[string, *OpenFileInfo] + + // Server state + serverState atomic.Value +} + +// New creates a new LSP client using the powernap implementation. +func New( + name string, + cfg config.LSPConfig, + resolver config.VariableResolver, + cwd string, + debug bool, +) (*Client, error) { + // Use a long-lived context independent of the caller's request context. + // The caller's context may be canceled when the tool call completes, + // but the LSP client must survive across multiple requests and restarts. + clientCtx, cancelCtx := context.WithCancel(context.Background()) + client := &Client{ + name: name, + fileTypes: cfg.FileTypes, + diagnostics: csync.NewVersionedMap[protocol.DocumentURI, []protocol.Diagnostic](), + openFiles: csync.NewMap[string, *OpenFileInfo](), + config: cfg, + ctx: clientCtx, + cancelCtx: cancelCtx, + debug: debug, + resolver: resolver, + cwd: cwd, + } + client.serverState.Store(StateStopped) + + if err := client.createPowernapClient(); err != nil { + return nil, err + } + + return client, nil +} + +// Initialize initializes the LSP client and returns the server capabilities. +func (c *Client) Initialize(ctx context.Context, workspaceDir string) (*protocol.InitializeResult, error) { + // Register handlers for requests the server may send during the + // initialize handshake itself (e.g. typescript-language-server issuing + // window/workDoneProgress/create while loading the project, before + // initialize has returned). Registering after client.Initialize() is too + // late for those — the server treats an unhandled response as fatal. + c.registerHandlers() + + if err := c.client.Initialize(ctx, false); err != nil { + return nil, fmt.Errorf("failed to initialize the lsp client: %w", err) + } + + // Convert powernap capabilities to protocol capabilities + caps := c.client.GetCapabilities() + protocolCaps := protocol.ServerCapabilities{ + TextDocumentSync: caps.TextDocumentSync, + CompletionProvider: func() *protocol.CompletionOptions { + if caps.CompletionProvider != nil { + return &protocol.CompletionOptions{ + TriggerCharacters: caps.CompletionProvider.TriggerCharacters, + AllCommitCharacters: caps.CompletionProvider.AllCommitCharacters, + ResolveProvider: caps.CompletionProvider.ResolveProvider, + } + } + return nil + }(), + } + + result := &protocol.InitializeResult{ + Capabilities: protocolCaps, + } + + return result, nil +} + +// closeTimeout is the maximum time to wait for a graceful LSP shutdown. +const closeTimeout = 5 * time.Second + +// Kill kills the client without doing anything else. +func (c *Client) Kill() { c.client.Kill() } + +// Shutdown permanently cancels the client's long-lived context and kills the +// underlying process. Unlike Restart, this is terminal: the client cannot be +// reused after Shutdown. +func (c *Client) Shutdown() { + c.cancelCtx() + c.client.Kill() +} + +// GetOffsetEncoding returns the negotiated offset encoding for this client. +func (c *Client) GetOffsetEncoding() powernap.OffsetEncoding { + return c.client.GetOffsetEncoding() +} + +// Close closes all open files in the client, then shuts down gracefully. +// If shutdown takes longer than closeTimeout, it falls back to Kill(). +func (c *Client) Close(ctx context.Context) error { + c.CloseAllFiles(ctx) + + // Use a timeout to prevent hanging on unresponsive LSP servers. + // jsonrpc2's send lock doesn't respect context cancellation, so we + // need to fall back to Kill() which closes the underlying connection. + closeCtx, cancel := context.WithTimeout(ctx, closeTimeout) + defer cancel() + + done := make(chan error, 1) + go func() { + if err := c.client.Shutdown(closeCtx); err != nil { + slog.Warn("Failed to shutdown LSP client", "error", err) + } + done <- c.client.Exit() + }() + + select { + case err := <-done: + return err + case <-closeCtx.Done(): + c.client.Kill() + return closeCtx.Err() + } +} + +// createPowernapClient creates a new powernap client with the current configuration. +func (c *Client) createPowernapClient() error { + rootURI := string(protocol.URIFromPath(c.cwd)) + + command, err := c.resolver.ResolveValue(c.config.Command) + if err != nil { + return fmt.Errorf("invalid lsp command: %w", err) + } + + args, err := c.config.ResolvedArgs(c.resolver) + if err != nil { + return fmt.Errorf("invalid lsp args: %w", err) + } + + envs, err := c.config.ResolvedEnv(c.resolver) + if err != nil { + return fmt.Errorf("invalid lsp env: %w", err) + } + + clientConfig := powernap.ClientConfig{ + Command: home.Long(command), + Args: args, + RootURI: rootURI, + Environment: envs, + Settings: c.config.Options, + InitOptions: c.config.InitOptions, + WorkspaceFolders: []protocol.WorkspaceFolder{ + { + URI: rootURI, + Name: filepath.Base(c.cwd), + }, + }, + } + + powernapClient, err := powernap.NewClient(clientConfig) + if err != nil { + return fmt.Errorf("failed to create lsp client: %w", err) + } + + c.client = powernapClient + return nil +} + +// registerHandlers registers the standard LSP notification and request handlers. +func (c *Client) registerHandlers() { + c.RegisterServerRequestHandler("workspace/applyEdit", HandleApplyEdit(c.client.GetOffsetEncoding())) + c.RegisterServerRequestHandler("workspace/configuration", HandleWorkspaceConfiguration) + c.RegisterServerRequestHandler("client/registerCapability", HandleRegisterCapability) + c.RegisterServerRequestHandler("window/workDoneProgress/create", HandleWorkDoneProgressCreate) + c.RegisterNotificationHandler("window/showMessage", func(ctx context.Context, method string, params json.RawMessage) { + if c.debug { + HandleServerMessage(ctx, method, params) + } + }) + c.RegisterNotificationHandler("textDocument/publishDiagnostics", func(_ context.Context, _ string, params json.RawMessage) { + HandleDiagnostics(c, params) + }) +} + +// Restart closes the current LSP client and creates a new one with the same configuration. +func (c *Client) Restart() error { + var openFiles []string + for uri := range c.openFiles.Seq2() { + openFiles = append(openFiles, string(uri)) + } + + // Cancel the old long-lived context and create a fresh one so that + // reinitialization is not affected by any prior cancellation. + c.cancelCtx() + c.ctx, c.cancelCtx = context.WithCancel(context.Background()) + + closeCtx, cancel := context.WithTimeout(c.ctx, 10*time.Second) + defer cancel() + + if err := c.Close(closeCtx); err != nil { + slog.Warn("Error closing client during restart", "name", c.name, "error", err) + } + + c.SetServerState(StateStopped) + + c.diagCountsCache = DiagnosticCounts{} + c.diagCountsVersion = 0 + + if err := c.createPowernapClient(); err != nil { + return err + } + + initCtx, cancel := context.WithTimeout(c.ctx, 30*time.Second) + defer cancel() + + c.SetServerState(StateStarting) + + // Register handlers before Initialize so servers that send + // requests during the handshake (e.g. window/workDoneProgress/create) + // don't crash on an unhandled response. + c.registerHandlers() + + if err := c.client.Initialize(initCtx, false); err != nil { + c.SetServerState(StateError) + return fmt.Errorf("failed to initialize lsp client: %w", err) + } + + if err := c.WaitForServerReady(initCtx); err != nil { + slog.Error("Server failed to become ready after restart", "name", c.name, "error", err) + c.SetServerState(StateError) + return err + } + + for _, uri := range openFiles { + if err := c.OpenFile(initCtx, uri); err != nil { + slog.Warn("Failed to reopen file after restart", "file", uri, "error", err) + } + } + return nil +} + +// ServerState represents the state of an LSP server +type ServerState int + +const ( + StateUnstarted ServerState = iota + StateStarting + StateReady + StateError + StateStopped + StateDisabled +) + +// GetServerState returns the current state of the LSP server +func (c *Client) GetServerState() ServerState { + if val := c.serverState.Load(); val != nil { + return val.(ServerState) + } + return StateStarting +} + +// SetServerState sets the current state of the LSP server +func (c *Client) SetServerState(state ServerState) { + c.serverState.Store(state) +} + +// GetName returns the name of the LSP client +func (c *Client) GetName() string { + return c.name +} + +// FileTypes returns the file types this LSP client handles +func (c *Client) FileTypes() []string { + return slices.Clone(c.fileTypes) +} + +// SetDiagnosticsCallback sets the callback function for diagnostic changes +func (c *Client) SetDiagnosticsCallback(callback func(name string, count int)) { + c.onDiagnosticsChanged = callback +} + +// WaitForServerReady waits for the server to be ready +func (c *Client) WaitForServerReady(ctx context.Context) error { + // Set initial state + c.SetServerState(StateStarting) + + // Try to ping the server with a simple request + ticker := time.NewTicker(500 * time.Millisecond) + defer ticker.Stop() + + if c.debug { + slog.Debug("Waiting for LSP server to be ready...") + } + + c.openKeyConfigFiles(ctx) + + for { + select { + case <-ctx.Done(): + c.SetServerState(StateError) + return fmt.Errorf("timeout waiting for LSP server to be ready") + case <-ticker.C: + // Check if client is running + if !c.client.IsRunning() { + if c.debug { + slog.Debug("LSP server not ready yet", "server", c.name) + } + continue + } + + // Server is ready + c.SetServerState(StateReady) + if c.debug { + slog.Debug("LSP server is ready") + } + return nil + } + } +} + +// OpenFileInfo contains information about an open file +type OpenFileInfo struct { + Version int32 + URI protocol.DocumentURI +} + +// HandlesFile checks if this LSP client handles the given file based on its +// extension and whether it's within the working directory. +func (c *Client) HandlesFile(path string) bool { + if c == nil { + return false + } + if !fsext.HasPrefix(path, c.cwd) { + slog.Debug("File outside workspace", "name", c.name, "file", path, "workDir", c.cwd) + return false + } + return handlesFiletype(c.name, c.fileTypes, path) +} + +// OpenFile opens a file in the LSP server. +func (c *Client) OpenFile(ctx context.Context, filepath string) error { + if !c.HandlesFile(filepath) { + return nil + } + + uri := string(protocol.URIFromPath(filepath)) + + if _, exists := c.openFiles.Get(uri); exists { + return nil // Already open + } + + // Skip files that do not exist or cannot be read + content, err := os.ReadFile(filepath) + if err != nil { + return fmt.Errorf("error reading file: %w", err) + } + + // Notify the server about the opened document + if err = c.client.NotifyDidOpenTextDocument(ctx, uri, string(powernap.DetectLanguage(filepath)), 1, string(content)); err != nil { + return err + } + + c.openFiles.Set(uri, &OpenFileInfo{ + Version: 1, + URI: protocol.DocumentURI(uri), + }) + + return nil +} + +// NotifyChange notifies the server about a file change. +func (c *Client) NotifyChange(ctx context.Context, filepath string) error { + if c == nil { + return nil + } + uri := string(protocol.URIFromPath(filepath)) + + content, err := os.ReadFile(filepath) + if err != nil { + return fmt.Errorf("error reading file: %w", err) + } + + fileInfo, isOpen := c.openFiles.Get(uri) + if !isOpen { + return fmt.Errorf("cannot notify change for unopened file: %s", filepath) + } + + // Increment version + fileInfo.Version++ + + // Create change event + changes := []protocol.TextDocumentContentChangeEvent{ + { + Value: protocol.TextDocumentContentChangeWholeDocument{ + Text: string(content), + }, + }, + } + + return c.client.NotifyDidChangeTextDocument(ctx, uri, int(fileInfo.Version), changes) +} + +// IsFileOpen checks if a file is currently open. +func (c *Client) IsFileOpen(filepath string) bool { + uri := string(protocol.URIFromPath(filepath)) + _, exists := c.openFiles.Get(uri) + return exists +} + +// CloseAllFiles closes all currently open files. +func (c *Client) CloseAllFiles(ctx context.Context) { + for uri := range c.openFiles.Seq2() { + if c.debug { + slog.Debug("Closing file", "file", uri) + } + if err := c.client.NotifyDidCloseTextDocument(ctx, uri); err != nil { + slog.Warn("Error closing file", "uri", uri, "error", err) + continue + } + c.openFiles.Del(uri) + } +} + +// GetFileDiagnostics returns diagnostics for a specific file. +func (c *Client) GetFileDiagnostics(uri protocol.DocumentURI) []protocol.Diagnostic { + diags, _ := c.diagnostics.Get(uri) + return diags +} + +// GetDiagnostics returns all diagnostics for all files. +func (c *Client) GetDiagnostics() map[protocol.DocumentURI][]protocol.Diagnostic { + if c == nil { + return nil + } + return c.diagnostics.Copy() +} + +// GetDiagnosticCounts returns cached diagnostic counts by severity. +// Uses the VersionedMap version to avoid recomputing on every call. +func (c *Client) GetDiagnosticCounts() DiagnosticCounts { + if c == nil { + return DiagnosticCounts{} + } + currentVersion := c.diagnostics.Version() + + c.diagCountsMu.Lock() + defer c.diagCountsMu.Unlock() + + if currentVersion == c.diagCountsVersion { + return c.diagCountsCache + } + + // Recompute counts. + counts := DiagnosticCounts{} + for _, diags := range c.diagnostics.Seq2() { + for _, diag := range diags { + switch diag.Severity { + case protocol.SeverityError: + counts.Error++ + case protocol.SeverityWarning: + counts.Warning++ + case protocol.SeverityInformation: + counts.Information++ + case protocol.SeverityHint: + counts.Hint++ + } + } + } + + c.diagCountsCache = counts + c.diagCountsVersion = currentVersion + return counts +} + +// OpenFileOnDemand opens a file only if it's not already open. +func (c *Client) OpenFileOnDemand(ctx context.Context, filepath string) error { + if c == nil { + return nil + } + // Check if the file is already open + if c.IsFileOpen(filepath) { + return nil + } + + // Open the file + return c.OpenFile(ctx, filepath) +} + +// RegisterNotificationHandler registers a notification handler. +func (c *Client) RegisterNotificationHandler(method string, handler transport.NotificationHandler) { + c.client.RegisterNotificationHandler(method, handler) +} + +// RegisterServerRequestHandler handles server requests. +func (c *Client) RegisterServerRequestHandler(method string, handler transport.Handler) { + c.client.RegisterHandler(method, handler) +} + +// openKeyConfigFiles opens important configuration files that help initialize the server. +func (c *Client) openKeyConfigFiles(ctx context.Context) { + // Try to open each file, ignoring errors if they don't exist + for _, file := range c.config.RootMarkers { + file = filepath.Join(c.cwd, file) + if _, err := os.Stat(file); err == nil { + // File exists, try to open it + if err := c.OpenFile(ctx, file); err != nil { + slog.Error("Failed to open key config file", "file", file, "error", err) + } else { + slog.Debug("Opened key config file for initialization", "file", file) + } + } + } +} + +// NotifyWorkspaceChange sends a workspace-level file change notification to +// trigger re-analysis of all files. This is useful when the overall project +// state may have changed (e.g., after a project-wide refactoring) and +// diagnostics for files not currently being edited may be stale. +func (c *Client) NotifyWorkspaceChange(ctx context.Context) error { + if c == nil { + return nil + } + return c.client.NotifyDidChangeWatchedFiles(ctx, []protocol.FileEvent{ + {URI: protocol.DocumentURI(protocol.URIFromPath(c.cwd)), Type: protocol.Changed}, + }) +} + +// RefreshOpenFiles re-notifies the LSP server about all currently open files, +// which triggers re-analysis and fresh diagnostics for the entire project. +func (c *Client) RefreshOpenFiles(ctx context.Context) { + if c == nil { + return + } + for uri, info := range c.openFiles.Seq2() { + path, err := protocol.DocumentURI(uri).Path() + if err != nil { + slog.Warn("Failed to convert URI to path", "uri", uri, "error", err) + continue + } + content, err := os.ReadFile(path) + if err != nil { + slog.Warn("Failed to read file for refresh", "path", path, "error", err) + continue + } + info.Version++ + changes := []protocol.TextDocumentContentChangeEvent{ + { + Value: protocol.TextDocumentContentChangeWholeDocument{ + Text: string(content), + }, + }, + } + if err := c.client.NotifyDidChangeTextDocument(ctx, uri, int(info.Version), changes); err != nil { + slog.Warn("Failed to notify file change", "uri", uri, "error", err) + } + } +} + +// WaitForDiagnostics waits until diagnostics stop changing for a settling +// period, indicating the LSP server has finished processing. If no +// diagnostics change within firstChangeDuration, it returns early since the +// server likely isn't going to republish. +func (c *Client) WaitForDiagnostics(ctx context.Context, timeout time.Duration) { + if c == nil { + return + } + + const ( + firstChangeDuration = 1 * time.Second + settleDuration = 300 * time.Millisecond + ) + + deadline := time.NewTimer(timeout) + defer deadline.Stop() + firstChangeTimer := time.NewTimer(min(timeout, firstChangeDuration)) + defer firstChangeTimer.Stop() + previousVersion := c.diagnostics.Version() + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-deadline.C: + return + case <-firstChangeTimer.C: + // No change arrived quickly — server isn't republishing. + return + case <-ticker.C: + currentVersion := c.diagnostics.Version() + if currentVersion != previousVersion { + // Diagnostics changed — now wait for them to settle. + c.waitForDiagnosticsToSettle(ctx, deadline.C, settleDuration) + return + } + } + } +} + +// waitForDiagnosticsToSettle waits until diagnostics version stays the same +// for settleDuration, indicating the LSP server has finished publishing. +func (c *Client) waitForDiagnosticsToSettle(ctx context.Context, deadline <-chan time.Time, settleDuration time.Duration) { + lastVersion := c.diagnostics.Version() + settleTicker := time.NewTicker(50 * time.Millisecond) + defer settleTicker.Stop() + + // Track how long the version has been stable. + stableStart := time.Now() + + for { + select { + case <-ctx.Done(): + return + case <-deadline: + return + case <-settleTicker.C: + currentVersion := c.diagnostics.Version() + if currentVersion != lastVersion { + // New change detected — reset the stable timer. + lastVersion = currentVersion + stableStart = time.Now() + } else if time.Since(stableStart) >= settleDuration { + // Diagnostics have been stable for the settle duration. + return + } + } + } +} + +// FindReferences finds all references to the symbol at the given position. +func (c *Client) FindReferences(ctx context.Context, filepath string, line, character int, includeDeclaration bool) ([]protocol.Location, error) { + if err := c.OpenFileOnDemand(ctx, filepath); err != nil { + return nil, err + } + + // Add timeout to prevent hanging on slow LSP servers. + ctx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + + // NOTE: line and character should be 0-based. + // See: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#position + return c.client.FindReferences(ctx, filepath, line-1, character-1, includeDeclaration) +} + +// Rename renames the symbol at the given position across all files. +func (c *Client) Rename(ctx context.Context, filepath string, line, character int, newName string) (*protocol.WorkspaceEdit, error) { + if err := c.OpenFileOnDemand(ctx, filepath); err != nil { + return nil, err + } + + ctx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + + return c.client.RequestRename(ctx, filepath, line-1, character-1, newName) //nolint:wrapcheck +} + +// DocumentSymbols returns the document symbols for the given file. +func (c *Client) DocumentSymbols(ctx context.Context, filepath string) ([]protocol.DocumentSymbolResult, error) { + if err := c.OpenFileOnDemand(ctx, filepath); err != nil { + return nil, err + } + + ctx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + + return c.client.RequestDocumentSymbols(ctx, filepath) //nolint:wrapcheck +} + +// Definition finds the definition of the symbol at the given position. +func (c *Client) Definition(ctx context.Context, filepath string, line, character int) ([]protocol.Location, error) { + if err := c.OpenFileOnDemand(ctx, filepath); err != nil { + return nil, err + } + + ctx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + + return c.client.RequestDefinition(ctx, filepath, line-1, character-1) //nolint:wrapcheck +} + +// PrepareCallHierarchy prepares a call hierarchy item at the given position. +func (c *Client) PrepareCallHierarchy(ctx context.Context, filepath string, line, character int) ([]protocol.CallHierarchyItem, error) { + if err := c.OpenFileOnDemand(ctx, filepath); err != nil { + return nil, err + } + + ctx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + + return c.client.PrepareCallHierarchy(ctx, filepath, line-1, character-1) //nolint:wrapcheck +} + +// IncomingCalls returns all callers of the given call hierarchy item. +func (c *Client) IncomingCalls(ctx context.Context, item protocol.CallHierarchyItem) ([]protocol.CallHierarchyIncomingCall, error) { + ctx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + + return c.client.IncomingCalls(ctx, item) //nolint:wrapcheck +} + +// OutgoingCalls returns all callees of the given call hierarchy item. +func (c *Client) OutgoingCalls(ctx context.Context, item protocol.CallHierarchyItem) ([]protocol.CallHierarchyOutgoingCall, error) { + ctx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + + return c.client.OutgoingCalls(ctx, item) //nolint:wrapcheck +} diff --git a/internal/lsp/client_test.go b/internal/lsp/client_test.go new file mode 100644 index 0000000000000000000000000000000000000000..2cb8f9a03312966ce0d8792e7d0aabc3879831c2 --- /dev/null +++ b/internal/lsp/client_test.go @@ -0,0 +1,202 @@ +package lsp + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/csync" + "github.com/charmbracelet/crush/internal/env" + "github.com/charmbracelet/x/powernap/pkg/lsp/protocol" + "github.com/stretchr/testify/require" +) + +func TestClient(t *testing.T) { + // Create a simple config for testing + cfg := config.LSPConfig{ + Command: "$THE_CMD", // Use echo as a dummy command that won't fail + Args: []string{"hello"}, + FileTypes: []string{"go"}, + Env: map[string]string{}, + } + + // Test creating a powernap client - this will likely fail with echo + // but we can still test the basic structure + client, err := New("test", cfg, config.NewShellVariableResolver(env.NewFromMap(map[string]string{ + "THE_CMD": "echo", + })), ".", false) + if err != nil { + // Expected to fail with echo command, skip the rest + t.Skipf("Powernap client creation failed as expected with dummy command: %v", err) + return + } + + // If we get here, test basic interface methods + if client.GetName() != "test" { + t.Errorf("Expected name 'test', got '%s'", client.GetName()) + } + + if !client.HandlesFile("test.go") { + t.Error("Expected client to handle .go files") + } + + if client.HandlesFile("test.py") { + t.Error("Expected client to not handle .py files") + } + + // Test server state + client.SetServerState(StateReady) + if client.GetServerState() != StateReady { + t.Error("Expected server state to be StateReady") + } + + // Clean up - expect this to fail with echo command + if err := client.Close(t.Context()); err != nil { + // Expected to fail with echo command + t.Logf("Close failed as expected with dummy command: %v", err) + } +} + +// TestNew_ExpansionFailure_Args pins that a failing $(cmd) in LSP +// args surfaces as a load error prefixed "invalid lsp args:" and that +// no client is returned. Mirrors the MCP contract where expansion +// failure hard-stops transport creation rather than silently running +// with an empty or literal value. +func TestNew_ExpansionFailure_Args(t *testing.T) { + t.Parallel() + + cfg := config.LSPConfig{ + Command: "echo", + Args: []string{"--root", "$(false)"}, + } + resolver := config.NewShellVariableResolver(env.NewFromMap(map[string]string{})) + + client, err := New("test-args-fail", cfg, resolver, ".", false) + require.Error(t, err) + require.Nil(t, client, "client must not start when args expansion fails") + require.Contains(t, err.Error(), "invalid lsp args") +} + +// TestNew_ExpansionFailure_Env pins the same contract for env values. +func TestNew_ExpansionFailure_Env(t *testing.T) { + t.Parallel() + + cfg := config.LSPConfig{ + Command: "echo", + Env: map[string]string{"BAD": "$(false)"}, + } + resolver := config.NewShellVariableResolver(env.NewFromMap(map[string]string{})) + + client, err := New("test-env-fail", cfg, resolver, ".", false) + require.Error(t, err) + require.Nil(t, client, "client must not start when env expansion fails") + require.Contains(t, err.Error(), "invalid lsp env") +} + +func TestNilClient(t *testing.T) { + t.Parallel() + + var c *Client + + require.False(t, c.HandlesFile("/some/file.go")) + require.Equal(t, DiagnosticCounts{}, c.GetDiagnosticCounts()) + require.Nil(t, c.GetDiagnostics()) + require.Nil(t, c.OpenFileOnDemand(context.Background(), "/some/file.go")) + require.Nil(t, c.NotifyChange(context.Background(), "/some/file.go")) + c.WaitForDiagnostics(context.Background(), time.Second) +} + +func newTestClient() *Client { + c := &Client{ + name: "test", + diagnostics: csync.NewVersionedMap[protocol.DocumentURI, []protocol.Diagnostic](), + openFiles: csync.NewMap[string, *OpenFileInfo](), + } + c.serverState.Store(StateStopped) + return c +} + +func TestWaitForDiagnostics_NoChange(t *testing.T) { + t.Parallel() + + c := newTestClient() + start := time.Now() + c.WaitForDiagnostics(t.Context(), 5*time.Second) + elapsed := time.Since(start) + + // Should return early via firstChangeDeadline (~1s), not the full timeout. + require.Less(t, elapsed, 2*time.Second, "should return early when no diagnostics change") +} + +func TestWaitForDiagnostics_ImmediateChange(t *testing.T) { + t.Parallel() + + c := newTestClient() + + go func() { + time.Sleep(100 * time.Millisecond) + c.diagnostics.Set(protocol.DocumentURI("file:///test.go"), nil) + }() + + start := time.Now() + c.WaitForDiagnostics(t.Context(), 5*time.Second) + elapsed := time.Since(start) + + // Should detect the change and then settle (~300ms settle + overhead). + require.Less(t, elapsed, 2*time.Second, "should return after settling, not full timeout") + require.Greater(t, elapsed, 200*time.Millisecond, "should wait for settle duration") +} + +func TestWaitForDiagnostics_RepeatedChanges(t *testing.T) { + t.Parallel() + + c := newTestClient() + + // Simulate an LSP server that publishes diagnostics in bursts. + go func() { + for i := range 5 { + time.Sleep(50 * time.Millisecond) + c.diagnostics.Set(protocol.DocumentURI("file:///test.go"), []protocol.Diagnostic{ + {Message: fmt.Sprintf("diag-%d", i)}, + }) + } + }() + + start := time.Now() + c.WaitForDiagnostics(t.Context(), 5*time.Second) + elapsed := time.Since(start) + + // Should wait for diagnostics to settle after the burst finishes. + // Burst lasts ~250ms, then 300ms settle window, so total ~550ms+. + require.Less(t, elapsed, 2*time.Second, "should return after settling, not full timeout") + require.Greater(t, elapsed, 400*time.Millisecond, "should wait for all changes to settle") +} + +func TestWaitForDiagnostics_ContextCancellation(t *testing.T) { + t.Parallel() + + c := newTestClient() + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + go func() { + time.Sleep(200 * time.Millisecond) + cancel() + }() + + start := time.Now() + c.WaitForDiagnostics(ctx, 5*time.Second) + elapsed := time.Since(start) + + require.Less(t, elapsed, 1*time.Second, "should return shortly after context cancellation") +} + +func TestWaitForDiagnostics_NilClient(t *testing.T) { + t.Parallel() + + var c *Client + // Should not panic. + c.WaitForDiagnostics(context.Background(), time.Second) +} diff --git a/internal/lsp/handlers.go b/internal/lsp/handlers.go new file mode 100644 index 0000000000000000000000000000000000000000..b66b1d6c2bd5b733a6178e516f96a9260036f3f5 --- /dev/null +++ b/internal/lsp/handlers.go @@ -0,0 +1,133 @@ +package lsp + +import ( + "context" + "encoding/json" + "log/slog" + + "github.com/charmbracelet/crush/internal/lsp/util" + powernap "github.com/charmbracelet/x/powernap/pkg/lsp" + "github.com/charmbracelet/x/powernap/pkg/lsp/protocol" +) + +// HandleWorkspaceConfiguration handles workspace configuration requests +func HandleWorkspaceConfiguration(_ context.Context, _ string, params json.RawMessage) (any, error) { + return []map[string]any{{}}, nil +} + +// HandleWorkDoneProgressCreate handles server-initiated window/workDoneProgress/create +// requests. The client advertises window.workDoneProgress: true in its capabilities +// (see makeClientCapabilities in powernap), which per the LSP spec grants servers +// permission to send this request — so it must be answered, even as a no-op, or the +// server (e.g. typescript-language-server) treats the unhandled response as fatal and +// crashes. See github.com/charmbracelet/x issue tracking powernap capability gaps. +func HandleWorkDoneProgressCreate(_ context.Context, _ string, _ json.RawMessage) (any, error) { + return nil, nil +} + +// HandleRegisterCapability handles capability registration requests +func HandleRegisterCapability(_ context.Context, _ string, params json.RawMessage) (any, error) { + var registerParams protocol.RegistrationParams + if err := json.Unmarshal(params, ®isterParams); err != nil { + slog.Error("Error unmarshaling registration params", "error", err) + return nil, err + } + + for _, reg := range registerParams.Registrations { + switch reg.Method { + case "workspace/didChangeWatchedFiles": + // Parse the registration options + optionsJSON, err := json.Marshal(reg.RegisterOptions) + if err != nil { + slog.Error("Error marshaling registration options", "error", err) + continue + } + var options protocol.DidChangeWatchedFilesRegistrationOptions + if err := json.Unmarshal(optionsJSON, &options); err != nil { + slog.Error("Error unmarshaling registration options", "error", err) + continue + } + // Store the file watchers registrations + notifyFileWatchRegistration(reg.ID, options.Watchers) + } + } + return nil, nil +} + +// HandleApplyEdit handles workspace edit requests +func HandleApplyEdit(encoding powernap.OffsetEncoding) func(_ context.Context, _ string, params json.RawMessage) (any, error) { + return func(_ context.Context, _ string, params json.RawMessage) (any, error) { + var edit protocol.ApplyWorkspaceEditParams + if err := json.Unmarshal(params, &edit); err != nil { + return nil, err + } + + err := util.ApplyWorkspaceEdit(edit.Edit, encoding) + if err != nil { + slog.Error("Error applying workspace edit", "error", err) + return protocol.ApplyWorkspaceEditResult{Applied: false, FailureReason: err.Error()}, nil + } + + return protocol.ApplyWorkspaceEditResult{Applied: true}, nil + } +} + +// FileWatchRegistrationHandler is a function that will be called when file watch registrations are received +type FileWatchRegistrationHandler func(id string, watchers []protocol.FileSystemWatcher) + +// fileWatchHandler holds the current handler for file watch registrations +var fileWatchHandler FileWatchRegistrationHandler + +// RegisterFileWatchHandler sets the handler for file watch registrations +func RegisterFileWatchHandler(handler FileWatchRegistrationHandler) { + fileWatchHandler = handler +} + +// notifyFileWatchRegistration notifies the handler about new file watch registrations +func notifyFileWatchRegistration(id string, watchers []protocol.FileSystemWatcher) { + if fileWatchHandler != nil { + fileWatchHandler(id, watchers) + } +} + +// HandleServerMessage handles server messages +func HandleServerMessage(_ context.Context, method string, params json.RawMessage) { + var msg protocol.ShowMessageParams + if err := json.Unmarshal(params, &msg); err != nil { + slog.Debug("Error unmarshal server message", "error", err) + return + } + + switch msg.Type { + case protocol.Error: + slog.Error("LSP Server", "message", msg.Message) + case protocol.Warning: + slog.Warn("LSP Server", "message", msg.Message) + case protocol.Info: + slog.Info("LSP Server", "message", msg.Message) + case protocol.Log: + slog.Debug("LSP Server", "message", msg.Message) + } +} + +// HandleDiagnostics handles diagnostic notifications from the LSP server +func HandleDiagnostics(client *Client, params json.RawMessage) { + var diagParams protocol.PublishDiagnosticsParams + if err := json.Unmarshal(params, &diagParams); err != nil { + slog.Error("Error unmarshaling diagnostics params", "error", err) + return + } + + client.diagnostics.Set(diagParams.URI, diagParams.Diagnostics) + + // Calculate total diagnostic count + totalCount := 0 + for _, diagnostics := range client.diagnostics.Seq2() { + totalCount += len(diagnostics) + } + + // Trigger callback if set + if client.onDiagnosticsChanged != nil { + client.onDiagnosticsChanged(client.name, totalCount) + } +} diff --git a/internal/lsp/manager.go b/internal/lsp/manager.go new file mode 100644 index 0000000000000000000000000000000000000000..e9600c28823a8fcc74ac6742ef5437a1afd2f71d --- /dev/null +++ b/internal/lsp/manager.go @@ -0,0 +1,417 @@ +// Package lsp provides a manager for Language Server Protocol (LSP) clients. +package lsp + +import ( + "cmp" + "context" + "errors" + "io" + "log/slog" + "os/exec" + "path/filepath" + "strings" + "sync" + "time" + + "github.com/charmbracelet/crush/internal/config" + "github.com/charmbracelet/crush/internal/csync" + "github.com/charmbracelet/crush/internal/fsext" + powernapconfig "github.com/charmbracelet/x/powernap/pkg/config" + powernap "github.com/charmbracelet/x/powernap/pkg/lsp" + "github.com/sourcegraph/jsonrpc2" +) + +const unavailableRetryDelay = 30 * time.Second + +// Manager handles lazy initialization of LSP clients based on file types. +type Manager struct { + clients *csync.Map[string, *Client] + unavailable *csync.Map[string, time.Time] + cfg *config.ConfigStore + manager *powernapconfig.Manager + callback func(name string, client *Client) + now func() time.Time + lookPath func(string) (string, error) +} + +// NewManager creates a new LSP manager service. +func NewManager(cfg *config.ConfigStore) *Manager { + manager := powernapconfig.NewManager() + manager.LoadDefaults() + + // Merge user-configured LSPs into the manager. + for name, clientConfig := range cfg.Config().LSP { + if clientConfig.Disabled { + slog.Debug("LSP disabled by user config", "name", name) + manager.RemoveServer(name) + continue + } + + // HACK: the user might have the command name in their config instead + // of the actual name. Find and use the correct name. + actualName := resolveServerName(manager, name) + manager.AddServer(actualName, &powernapconfig.ServerConfig{ + Command: clientConfig.Command, + Args: clientConfig.Args, + Environment: clientConfig.Env, + FileTypes: clientConfig.FileTypes, + RootMarkers: clientConfig.RootMarkers, + InitOptions: clientConfig.InitOptions, + Settings: clientConfig.Options, + }) + } + + return &Manager{ + clients: csync.NewMap[string, *Client](), + unavailable: csync.NewMap[string, time.Time](), + cfg: cfg, + manager: manager, + callback: func(string, *Client) {}, // default no-op callback + now: time.Now, + lookPath: exec.LookPath, + } +} + +// Clients returns the map of LSP clients. +func (s *Manager) Clients() *csync.Map[string, *Client] { + return s.clients +} + +// SetCallback sets a callback that is invoked when a new LSP +// client is successfully started. This allows the coordinator to add LSP tools. +func (s *Manager) SetCallback(cb func(name string, client *Client)) { + s.callback = cb +} + +// TrackConfigured notifies the UI about user-configured LSP servers without +// starting them. Servers start on demand via Start(). +func (s *Manager) TrackConfigured(ctx context.Context) { + var wg sync.WaitGroup + servers := s.manager.GetServers() + for name := range servers { + if !s.isUserConfigured(name) { + continue + } + wg.Go(func() { + s.callback(name, nil) + }) + } + wg.Wait() +} + +// Start starts an LSP server that can handle the given file path. +// If an appropriate LSP is already running, this is a no-op. +func (s *Manager) Start(ctx context.Context, path string) { + if abs, err := filepath.Abs(path); err == nil { + path = abs + } + if !fsext.HasPrefix(path, s.cfg.WorkingDir()) { + return + } + + var wg sync.WaitGroup + for name, server := range s.manager.GetServers() { + wg.Go(func() { + s.startServer(name, path, server) + }) + } + wg.Wait() +} + +// skipAutoStartCommands contains commands that are too generic or ambiguous to +// auto-start without explicit user configuration. +var skipAutoStartCommands = map[string]bool{ + "buck2": true, + "buf": true, + "cue": true, + "dart": true, + "deno": true, + "dotnet": true, + "dprint": true, + "gleam": true, + "java": true, + "julia": true, + "koka": true, + "node": true, + "npx": true, + "perl": true, + "plz": true, + "python": true, + "python3": true, + "R": true, + "racket": true, + "rome": true, + "rubocop": true, + "ruff": true, + "scarb": true, + "solc": true, + "stylua": true, + "swipl": true, + "tflint": true, +} + +func (s *Manager) startServer(name, filepath string, server *powernapconfig.ServerConfig) { + var ( + isUserConfigured = s.isUserConfigured(name) + autoLSP = s.cfg.Config().Options.AutoLSP + ) + if !isUserConfigured && autoLSP != nil && !*autoLSP { + slog.Debug("Auto-start LSP disabled", "name", name) + return + } + + cfg := s.buildConfig(name, server) + if cfg.Disabled { + return + } + + if client, ok := s.clients.Get(name); ok { + switch client.GetServerState() { + case StateReady, StateStarting, StateDisabled: + s.callback(name, client) + // already done, return + return + } + } + + if isUserConfigured { + if !handles(server, filepath, s.cfg.WorkingDir()) { + return + } + } else if !s.canAutoStart(name, filepath, s.cfg.WorkingDir(), server) { + return + } + + // Check again in case another goroutine started it in the meantime. + if client, ok := s.clients.Get(name); ok { + switch client.GetServerState() { + case StateReady, StateStarting, StateDisabled: + s.callback(name, client) + return + } + } + + client, err := New( + name, + cfg, + s.cfg.Resolver(), + s.cfg.WorkingDir(), + s.cfg.Config().Options.DebugLSP, + ) + if err != nil { + slog.Error("Failed to create LSP client", "name", name, "error", err) + return + } + // Only store non-nil clients. If another goroutine raced us, + // prefer the already-stored client. + if existing, ok := s.clients.Get(name); ok { + switch existing.GetServerState() { + case StateReady, StateStarting, StateDisabled: + client.Shutdown() + s.callback(name, existing) + return + } + } + s.clients.Set(name, client) + defer func() { + s.callback(name, client) + }() + + switch client.GetServerState() { + case StateReady, StateStarting, StateDisabled: + // already done, return + return + } + + client.serverState.Store(StateStarting) + + // Use an independent context for initialization so that the LSP server + // startup is not tied to the caller's request context. The caller's + // context may have a short timeout or be canceled when the tool call + // completes, but LSP initialization can take several seconds and the + // server must persist beyond any single request. + initCtx, cancel := context.WithTimeout(context.Background(), time.Duration(cmp.Or(cfg.Timeout, 30))*time.Second) + defer cancel() + + if _, err := client.Initialize(initCtx, s.cfg.WorkingDir()); err != nil { + slog.Error("LSP client initialization failed", "name", name, "error", err) + client.Shutdown() + s.clients.Del(name) + return + } + + if err := client.WaitForServerReady(initCtx); err != nil { + slog.Warn("LSP server not fully ready, continuing anyway", "name", name, "error", err) + client.SetServerState(StateError) + } else { + client.SetServerState(StateReady) + } + + slog.Debug("LSP client started", "name", name) +} + +func (s *Manager) canAutoStart( + name, filePath, workDir string, + server *powernapconfig.ServerConfig, +) bool { + if skipAutoStartCommands[server.Command] { + slog.Debug("LSP command too generic for auto-start, skipping", "name", name, "command", server.Command) + return false + } + + // Filtering by file type is cheap and usually rejects a server before the + // root marker check. Do both before searching PATH, which can require a stat + // for every directory in PATH for every bundled server. + if !handles(server, filePath, workDir) { + return false + } + + if s.recentlyUnavailable(name) { + return false + } + if _, err := s.lookPath(server.Command); err != nil { + slog.Debug("LSP server not installed, skipping", "name", name, "command", server.Command) + s.markUnavailable(name) + return false + } + s.clearUnavailable(name) + return true +} + +func (s *Manager) isUserConfigured(name string) bool { + cfg, ok := s.cfg.Config().LSP[name] + return ok && !cfg.Disabled +} + +func (s *Manager) recentlyUnavailable(name string) bool { + lastUnavailableAt, exists := s.unavailable.Get(name) + if !exists { + return false + } + if s.now().Sub(lastUnavailableAt) < unavailableRetryDelay { + return true + } + s.unavailable.Del(name) + return false +} + +func (s *Manager) markUnavailable(name string) { + s.unavailable.Set(name, s.now()) +} + +func (s *Manager) clearUnavailable(name string) { + s.unavailable.Del(name) +} + +func (s *Manager) buildConfig(name string, server *powernapconfig.ServerConfig) config.LSPConfig { + cfg := config.LSPConfig{ + Command: server.Command, + Args: server.Args, + Env: server.Environment, + FileTypes: server.FileTypes, + RootMarkers: server.RootMarkers, + InitOptions: server.InitOptions, + Options: server.Settings, + } + if userCfg, ok := s.cfg.Config().LSP[name]; ok { + cfg.Timeout = userCfg.Timeout + } + return cfg +} + +func resolveServerName(manager *powernapconfig.Manager, name string) string { + if _, ok := manager.GetServer(name); ok { + return name + } + for sname, server := range manager.GetServers() { + if server.Command == name { + return sname + } + } + return name +} + +func handlesFiletype(sname string, fileTypes []string, filePath string) bool { + if len(fileTypes) == 0 { + return true + } + + kind := powernap.DetectLanguage(filePath) + name := strings.ToLower(filepath.Base(filePath)) + for _, filetype := range fileTypes { + suffix := strings.ToLower(filetype) + if !strings.HasPrefix(suffix, ".") { + suffix = "." + suffix + } + if strings.HasSuffix(name, suffix) || filetype == string(kind) { + slog.Debug("Handles file", "name", sname, "file", name, "filetype", filetype, "kind", kind) + return true + } + } + + slog.Debug("Doesn't handle file", "name", sname, "file", name) + return false +} + +func hasRootMarkers(dir string, markers []string) bool { + if len(markers) == 0 { + return true + } + for _, pattern := range markers { + // Use filepath.Glob for a non-recursive check in the root + // directory. This avoids walking the entire tree (which is + // catastrophic in large monorepos with node_modules, etc.). + matches, err := filepath.Glob(filepath.Join(dir, pattern)) + if err == nil && len(matches) > 0 { + return true + } + } + return false +} + +func handles(server *powernapconfig.ServerConfig, filePath, workDir string) bool { + return handlesFiletype(server.Command, server.FileTypes, filePath) && + hasRootMarkers(workDir, server.RootMarkers) +} + +// KillAll force-kills all the LSP clients. +// +// This is generally faster than [Manager.StopAll] because it doesn't wait for +// the server to exit gracefully, but it can lead to data loss if the server is +// in the middle of writing something. +// Generally it doesn't matter when shutting down Crush, though. +func (s *Manager) KillAll(context.Context) { + var wg sync.WaitGroup + for name, client := range s.clients.Seq2() { + wg.Go(func() { + defer func() { s.callback(name, client) }() + client.Shutdown() + client.SetServerState(StateStopped) + s.clients.Del(name) + slog.Debug("Killed LSP client", "name", name) + }) + } + wg.Wait() +} + +// StopAll stops all running LSP clients and clears the client map. +func (s *Manager) StopAll(ctx context.Context) { + var wg sync.WaitGroup + for name, client := range s.clients.Seq2() { + wg.Go(func() { + defer func() { s.callback(name, client) }() + if err := client.Close(ctx); err != nil && + !errors.Is(err, io.EOF) && + !errors.Is(err, context.Canceled) && + !errors.Is(err, jsonrpc2.ErrClosed) && + err.Error() != "signal: killed" { + slog.Warn("Failed to stop LSP client", "name", name, "error", err) + } + client.cancelCtx() + client.SetServerState(StateStopped) + s.clients.Del(name) + slog.Debug("Stopped LSP client", "name", name) + }) + } + wg.Wait() +} diff --git a/internal/lsp/manager_test.go b/internal/lsp/manager_test.go new file mode 100644 index 0000000000000000000000000000000000000000..6023084d7137d463ef2bac53697d3f35af23995f --- /dev/null +++ b/internal/lsp/manager_test.go @@ -0,0 +1,115 @@ +package lsp + +import ( + "errors" + "testing" + "time" + + "github.com/charmbracelet/crush/internal/csync" + powernapconfig "github.com/charmbracelet/x/powernap/pkg/config" + "github.com/stretchr/testify/require" +) + +func TestUnavailableBackoff(t *testing.T) { + t.Parallel() + + base := time.Date(2026, 3, 26, 0, 0, 0, 0, time.UTC) + now := base + + manager := &Manager{ + unavailable: csync.NewMap[string, time.Time](), + now: func() time.Time { return now }, + } + + require.False(t, manager.recentlyUnavailable("gopls")) + + manager.markUnavailable("gopls") + require.True(t, manager.recentlyUnavailable("gopls")) + + now = now.Add(unavailableRetryDelay + time.Second) + require.False(t, manager.recentlyUnavailable("gopls")) + _, exists := manager.unavailable.Get("gopls") + require.False(t, exists) + + manager.markUnavailable("gopls") + manager.clearUnavailable("gopls") + require.False(t, manager.recentlyUnavailable("gopls")) +} + +func TestCanAutoStartFiltersBeforeLookingUpCommand(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + server *powernapconfig.ServerConfig + want bool + lookups int + }{ + { + name: "unhandled file type", + server: &powernapconfig.ServerConfig{ + Command: "typescript-language-server", + FileTypes: []string{"typescript"}, + }, + }, + { + name: "generic command", + server: &powernapconfig.ServerConfig{ + Command: "node", + FileTypes: []string{"go"}, + }, + }, + { + name: "handled file type", + server: &powernapconfig.ServerConfig{ + Command: "gopls", + FileTypes: []string{"go"}, + }, + want: true, + lookups: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + lookups := 0 + manager := &Manager{ + unavailable: csync.NewMap[string, time.Time](), + now: time.Now, + lookPath: func(string) (string, error) { + lookups++ + return "/usr/local/bin/gopls", nil + }, + } + + got := manager.canAutoStart("test", "main.go", t.TempDir(), tt.server) + + require.Equal(t, tt.want, got) + require.Equal(t, tt.lookups, lookups) + }) + } +} + +func TestCanAutoStartCachesMissingCommand(t *testing.T) { + t.Parallel() + + lookups := 0 + manager := &Manager{ + unavailable: csync.NewMap[string, time.Time](), + now: time.Now, + lookPath: func(string) (string, error) { + lookups++ + return "", errors.New("not found") + }, + } + server := &powernapconfig.ServerConfig{ + Command: "gopls", + FileTypes: []string{"go"}, + } + + require.False(t, manager.canAutoStart("gopls", "main.go", t.TempDir(), server)) + require.False(t, manager.canAutoStart("gopls", "main.go", t.TempDir(), server)) + require.Equal(t, 1, lookups) +} diff --git a/internal/pubsub/broker.go b/internal/pubsub/broker.go new file mode 100644 index 0000000000000000000000000000000000000000..412ce4de3a870d361207ded257f0c294a74ff4c9 --- /dev/null +++ b/internal/pubsub/broker.go @@ -0,0 +1,236 @@ +// Package pubsub provides a lightweight in-process broker for fan-out +// event delivery between services and the UI. +// +// Delivery semantics: +// +// - [Broker.Publish] is best-effort and lossy under contention. If a +// subscriber's channel is full, the event is dropped for that +// subscriber, a warning is logged, and a counter is incremented. +// This is the right choice for high-frequency intermediate updates +// (e.g. streaming token deltas) where only the latest state +// matters. +// +// - [Broker.PublishMustDeliver] is bounded-blocking. For each +// subscriber it first tries a non-blocking send, then falls back to +// a per-subscriber blocking send with a hard timeout. On timeout the +// event is dropped for that subscriber, an error is logged, and the +// must-deliver drop counter is incremented. The publisher never +// blocks indefinitely. This is the right choice for terminal events +// (finish, tool result, error, cancel) that must not be silently +// coalesced away. +// +// Drop counters ([Broker.DropCount], [Broker.MustDeliverDropCount]) are +// exposed so callers can surface saturation in telemetry. +package pubsub + +import ( + "context" + "log/slog" + "sync" + "sync/atomic" + "time" +) + +const ( + // bufferSize is the per-subscriber channel capacity for any broker + // created via NewBroker. Publish is non-blocking, so a full buffer + // drops events (with a warning log); sized to cover a long + // streaming assistant turn (~one UpdatedEvent per token) even under + // TUI render stalls. + bufferSize = 4096 + + // defaultMustDeliverTimeout is the per-subscriber upper bound on how + // long [Broker.PublishMustDeliver] will block waiting for buffer + // space before giving up on that subscriber. + defaultMustDeliverTimeout = 50 * time.Millisecond +) + +type Broker[T any] struct { + subs map[chan Event[T]]struct{} + mu sync.RWMutex + done chan struct{} + subCount int + channelBufferSize int + mustDeliverTimeout time.Duration + dropCount atomic.Uint64 + mustDeliverDropCount atomic.Uint64 +} + +func NewBroker[T any]() *Broker[T] { + return NewBrokerWithOptions[T](bufferSize) +} + +func NewBrokerWithOptions[T any](channelBufferSize int) *Broker[T] { + return &Broker[T]{ + subs: make(map[chan Event[T]]struct{}), + done: make(chan struct{}), + channelBufferSize: channelBufferSize, + mustDeliverTimeout: defaultMustDeliverTimeout, + } +} + +// SetMustDeliverTimeout overrides the per-subscriber timeout used by +// [Broker.PublishMustDeliver]. A zero or negative value resets to the +// default. Intended primarily for tests. +func (b *Broker[T]) SetMustDeliverTimeout(d time.Duration) { + b.mu.Lock() + defer b.mu.Unlock() + if d <= 0 { + b.mustDeliverTimeout = defaultMustDeliverTimeout + return + } + b.mustDeliverTimeout = d +} + +func (b *Broker[T]) Shutdown() { + select { + case <-b.done: // Already closed + return + default: + close(b.done) + } + + b.mu.Lock() + defer b.mu.Unlock() + + for ch := range b.subs { + delete(b.subs, ch) + close(ch) + } + + b.subCount = 0 +} + +func (b *Broker[T]) Subscribe(ctx context.Context) <-chan Event[T] { + b.mu.Lock() + defer b.mu.Unlock() + + select { + case <-b.done: + ch := make(chan Event[T]) + close(ch) + return ch + default: + } + + sub := make(chan Event[T], b.channelBufferSize) + b.subs[sub] = struct{}{} + b.subCount++ + + go func() { + <-ctx.Done() + + b.mu.Lock() + defer b.mu.Unlock() + + select { + case <-b.done: + return + default: + } + + delete(b.subs, sub) + close(sub) + b.subCount-- + }() + + return sub +} + +func (b *Broker[T]) GetSubscriberCount() int { + b.mu.RLock() + defer b.mu.RUnlock() + return b.subCount +} + +// DropCount returns the cumulative number of events dropped by +// [Broker.Publish] because a subscriber's channel was full. +func (b *Broker[T]) DropCount() uint64 { + return b.dropCount.Load() +} + +// MustDeliverDropCount returns the cumulative number of events dropped +// by [Broker.PublishMustDeliver] after the per-subscriber timeout +// expired. +func (b *Broker[T]) MustDeliverDropCount() uint64 { + return b.mustDeliverDropCount.Load() +} + +// Publish delivers an event to every active subscriber. +// +// Delivery is non-blocking and lossy: if a subscriber's channel is full +// the event is dropped for that subscriber, a warning is logged, and +// [Broker.DropCount] is incremented. Use [Broker.PublishMustDeliver] +// for events that must not be silently dropped. +func (b *Broker[T]) Publish(t EventType, payload T) { + b.mu.RLock() + defer b.mu.RUnlock() + + select { + case <-b.done: + return + default: + } + + event := Event[T]{Type: t, Payload: payload} + + for sub := range b.subs { + select { + case sub <- event: + default: + // Channel is full, subscriber is slow — skip this event. + // Lossy by design; counted and logged so saturation is + // observable. + b.dropCount.Add(1) + slog.Warn("Pubsub buffer full; dropping event", "type", t) + } + } +} + +// PublishMustDeliver delivers an event with bounded-blocking semantics. +// For each subscriber it first attempts a non-blocking send, then falls +// back to a blocking send bounded by a per-subscriber timeout (default +// [defaultMustDeliverTimeout]). On timeout the event is dropped for +// that subscriber, [Broker.MustDeliverDropCount] is incremented, and an +// error is logged. The publisher never blocks indefinitely. +// +// Use this for terminal events that must reach subscribers (finish, +// tool result, error, cancel). Callers must still tolerate rare drops +// after timeout — recovery is the subscriber's responsibility (e.g. a +// re-fetch on the next session-visible event). +func (b *Broker[T]) PublishMustDeliver(ctx context.Context, t EventType, payload T) { + b.mu.RLock() + defer b.mu.RUnlock() + + select { + case <-b.done: + return + default: + } + + event := Event[T]{Type: t, Payload: payload} + timeout := b.mustDeliverTimeout + + for sub := range b.subs { + // Fast path: non-blocking send. + select { + case sub <- event: + continue + default: + } + + // Slow path: bounded blocking send. + timer := time.NewTimer(timeout) + select { + case sub <- event: + timer.Stop() + case <-timer.C: + b.mustDeliverDropCount.Add(1) + slog.Error("PublishMustDeliver timed out delivering event", + "type", t, "timeout", timeout) + case <-ctx.Done(): + timer.Stop() + return + } + } +} diff --git a/internal/pubsub/events.go b/internal/pubsub/events.go new file mode 100644 index 0000000000000000000000000000000000000000..f52c21dbb6a5e477d632395d3a14d36e915a5c52 --- /dev/null +++ b/internal/pubsub/events.go @@ -0,0 +1,67 @@ +package pubsub + +import ( + "context" + "encoding/json" +) + +const ( + CreatedEvent EventType = "created" + UpdatedEvent EventType = "updated" + DeletedEvent EventType = "deleted" +) + +// PayloadType identifies the type of event payload for discriminated +// deserialization over JSON. +type PayloadType = string + +const ( + PayloadTypeLSPEvent PayloadType = "lsp_event" + PayloadTypeMCPEvent PayloadType = "mcp_event" + PayloadTypePermissionRequest PayloadType = "permission_request" + PayloadTypePermissionNotification PayloadType = "permission_notification" + PayloadTypeMessage PayloadType = "message" + PayloadTypeSession PayloadType = "session" + PayloadTypeFile PayloadType = "file" + PayloadTypeAgentEvent PayloadType = "agent_event" + PayloadTypeConfigChanged PayloadType = "config_changed" + PayloadTypeSkillsEvent PayloadType = "skills_event" + PayloadTypeRunComplete PayloadType = "run_complete" + PayloadTypeUpdateAvailable PayloadType = "update_available" + PayloadTypeQuestionRequest PayloadType = "question_batch_request" + PayloadTypeQuestionNotification PayloadType = "question_batch_notification" +) + +// Payload wraps a discriminated JSON payload with a type tag. +type Payload struct { + Type PayloadType `json:"type"` + Payload json.RawMessage `json:"payload"` +} + +// Subscriber can subscribe to events of type T. +type Subscriber[T any] interface { + Subscribe(context.Context) <-chan Event[T] +} + +type ( + // EventType identifies the type of event. + EventType string + + // Event represents an event in the lifecycle of a resource. + Event[T any] struct { + Type EventType `json:"type"` + Payload T `json:"payload"` + } + + // Publisher can publish events of type T. + // + // Publish is best-effort and lossy under back-pressure; + // PublishMustDeliver applies the bounded-blocking semantics used + // for terminal events that must reach subscribers (finish, tool + // result, error, cancel, RunComplete). See [Broker.Publish] and + // [Broker.PublishMustDeliver]. + Publisher[T any] interface { + Publish(EventType, T) + PublishMustDeliver(context.Context, EventType, T) + } +) diff --git a/internal/question/question.go b/internal/question/question.go new file mode 100644 index 0000000000000000000000000000000000000000..8d26f680afb732086530a497f72cca64e3f52e6e --- /dev/null +++ b/internal/question/question.go @@ -0,0 +1,317 @@ +// Package question provides services for asking the user questions +// via the TUI and blocking until an answer is received. It mirrors +// the permission service pattern: publish a request over pubsub, +// block on a channel, and resolve when the UI sends back answers. +// +// Only one question can be pending at a time (the tool blocks until +// answered), so no correlation IDs are needed in the domain model. +package question + +import ( + "context" + "errors" + "fmt" + "sync" + + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/google/uuid" +) + +// ErrCancelled is returned by Ask when the user cancels the question. +var ErrCancelled = errors.New("question cancelled by user") + +// Type identifies the kind of question to present. +type Type string + +const ( + TypeYesNo Type = "yes_no" + TypeSingleChoice Type = "single_choice" + TypeMultiChoice Type = "multi_choice" + TypeFreeText Type = "free_text" +) + +// Choice represents a single selectable option. +type Choice struct { + ID string `json:"id"` + Label string `json:"label"` + Description string `json:"description,omitempty"` +} + +// Question is a single question definition within a Request. +type Question struct { + ID string `json:"id"` + Type Type `json:"type"` + Label string `json:"label,omitempty"` + Text string `json:"question"` + Description string `json:"description,omitempty"` + Choices []Choice `json:"choices,omitempty"` +} + +// Answer carries the user's response to a single Question. +type Answer struct { + QuestionID string `json:"question_id"` + SelectedIDs []string `json:"selected_ids,omitempty"` + FillInText string `json:"fill_in_text,omitempty"` + Yes *bool `json:"yes,omitempty"` + Notes map[string]string `json:"notes,omitempty"` +} + +// HasNotes reports whether any notes were attached. +func (a Answer) HasNotes() bool { return len(a.Notes) > 0 } + +// Request is the service envelope published to the UI. It contains +// one or more Questions. A single question renders without tabs; +// multiple questions render as a tabbed form with confirmation. +type Request struct { + ID string `json:"id"` + SessionID string `json:"session_id"` + ToolCallID string `json:"tool_call_id"` + Questions []Question `json:"questions"` + ConfirmTitle string `json:"confirm_title,omitempty"` + ConfirmDescription string `json:"confirm_description,omitempty"` +} + +// Validate checks that a Request has valid fields. For multiple +// questions, ConfirmTitle and ConfirmDescription are required. +func (r Request) Validate() error { + if len(r.Questions) == 0 { + return fmt.Errorf("at least one question is required") + } + if len(r.Questions) > MaxQuestions { + return fmt.Errorf("questions exceed maximum of %d (got %d)", MaxQuestions, len(r.Questions)) + } + for i, q := range r.Questions { + if err := q.Validate(); err != nil { + return fmt.Errorf("question %d: %w", i+1, err) + } + } + return nil +} + +// Validate checks that a Question has valid fields. Error messages +// are written for LLM consumption: specific and actionable. +func (q Question) Validate() error { + label := q.identifier() + if q.Text == "" { + return fmt.Errorf("%s: question text is required", label) + } + if len(q.Text) > MaxQuestionLength { + return fmt.Errorf("%s: text exceeds %d characters (got %d)", label, MaxQuestionLength, len(q.Text)) + } + if q.Description == "" { + return fmt.Errorf("%s: description is required", label) + } + if len(q.Description) > MaxDescriptionLength { + return fmt.Errorf("%s: description exceeds %d characters (got %d)", label, MaxDescriptionLength, len(q.Description)) + } + switch q.Type { + case TypeYesNo, TypeFreeText: + // No choices needed. + case TypeSingleChoice, TypeMultiChoice: + if len(q.Choices) < 2 { + return fmt.Errorf("%s: %s requires at least 2 choices in the \"choices\" array (got %d). Use \"choices\", not \"options\"", label, q.Type, len(q.Choices)) + } + if len(q.Choices) > MaxChoices { + return fmt.Errorf("%s: choices exceed maximum of %d (got %d)", label, MaxChoices, len(q.Choices)) + } + seen := make(map[string]bool, len(q.Choices)) + for i, c := range q.Choices { + if c.ID == "" { + return fmt.Errorf("%s: choice %d must have an \"id\" field", label, i+1) + } + if seen[c.ID] { + return fmt.Errorf("%s: choice %d has duplicate id %q", label, i+1, c.ID) + } + seen[c.ID] = true + if c.Label == "" { + return fmt.Errorf("%s: choice %d (%s) must have a \"label\" field", label, i+1, c.ID) + } + if len(c.Label) > MaxChoiceLabelLength { + return fmt.Errorf("%s: choice %d label exceeds %d characters (got %d)", label, i+1, MaxChoiceLabelLength, len(c.Label)) + } + if len(c.Description) > MaxChoiceDescriptionLength { + return fmt.Errorf("%s: choice %d description exceeds %d characters (got %d)", label, i+1, MaxChoiceDescriptionLength, len(c.Description)) + } + } + default: + return fmt.Errorf("%s: unknown type %q (must be yes_no, single_choice, multi_choice, or free_text)", label, q.Type) + } + return nil +} + +// identifier returns a human-readable label for error messages. +// Uses the question label, text excerpt, or a fallback. +func (q Question) identifier() string { + if q.Label != "" { + return fmt.Sprintf("[%s]", q.Label) + } + if q.Text != "" { + t := q.Text + if len(t) > 40 { + t = t[:40] + "…" + } + return fmt.Sprintf("[%s]", t) + } + return "[unnamed question]" +} + +const ( + MaxQuestionLength = 240 + MaxDescriptionLength = 600 + MaxChoiceLabelLength = 200 + MaxChoiceDescriptionLength = 200 + MaxChoices = 5 + MaxQuestions = 5 +) + +// Notification is published when a question batch is resolved so +// that non-answering clients can dismiss their open forms. +type Notification struct { + BatchID string `json:"batch_id"` +} + +// Service manages the lifecycle of question requests. Only one +// question can be pending at a time. +type Service interface { + pubsub.Subscriber[Request] + + // SubscribeNotifications returns a channel for question + // resolution notifications. + SubscribeNotifications(ctx context.Context) <-chan pubsub.Event[Notification] + + // Ask publishes questions and blocks until the user answers + // or the context is cancelled. + Ask(ctx context.Context, req Request) ([]Answer, error) + + // Answer resolves the pending question with the given answers. + Answer(answers []Answer) bool + + // Cancel cancels the pending question. Returns false if no + // question is pending. + Cancel() bool +} + +type questionService struct { + broker *pubsub.Broker[Request] + notificationBroker *pubsub.Broker[Notification] + mu sync.Mutex + pending chan []Answer + cancelled chan struct{} + pendingID string +} + +// NewService creates a new question service. +func NewService() *questionService { + return &questionService{ + broker: pubsub.NewBroker[Request](), + notificationBroker: pubsub.NewBroker[Notification](), + } +} + +// Subscribe returns a channel for question events. +func (s *questionService) Subscribe(ctx context.Context) <-chan pubsub.Event[Request] { + return s.broker.Subscribe(ctx) +} + +// SubscribeNotifications returns a channel for question resolution +// notifications. +func (s *questionService) SubscribeNotifications(ctx context.Context) <-chan pubsub.Event[Notification] { + return s.notificationBroker.Subscribe(ctx) +} + +// Ask publishes a request and blocks until the user answers. +func (s *questionService) Ask(ctx context.Context, req Request) ([]Answer, error) { + if req.ID == "" { + req.ID = uuid.New().String() + } + for i := range req.Questions { + if req.Questions[i].ID == "" { + req.Questions[i].ID = uuid.New().String() + } + } + + // Apply defaults for multi-question confirm fields. + if len(req.Questions) >= 2 { + if req.ConfirmTitle == "" { + req.ConfirmTitle = "Ready to go?" + } + if req.ConfirmDescription == "" { + req.ConfirmDescription = "Review your answers above and confirm." + } + } + + if err := req.Validate(); err != nil { + return nil, err + } + + s.mu.Lock() + s.pending = make(chan []Answer, 1) + s.cancelled = make(chan struct{}) + s.pendingID = req.ID + s.mu.Unlock() + + defer func() { + s.mu.Lock() + s.pending = nil + s.cancelled = nil + s.pendingID = "" + s.mu.Unlock() + }() + + s.broker.Publish(pubsub.CreatedEvent, req) + + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-s.cancelled: + return nil, ErrCancelled + case answers := <-s.pending: + return answers, nil + } +} + +// Answer resolves the pending question. Returns false if no +// question is pending (already answered or cancelled). +func (s *questionService) Answer(answers []Answer) bool { + s.mu.Lock() + batchID := s.pendingID + ch := s.pending + s.mu.Unlock() + + if ch == nil { + return false + } + ch <- answers + + // Publish a notification so non-answering clients can dismiss + // their open question forms. + if batchID != "" { + s.notificationBroker.Publish(pubsub.CreatedEvent, Notification{ + BatchID: batchID, + }) + } + return true +} + +// Cancel cancels the pending question. Returns false if no +// question is pending. +func (s *questionService) Cancel() bool { + s.mu.Lock() + batchID := s.pendingID + cancelCh := s.cancelled + s.mu.Unlock() + + if cancelCh == nil { + return false + } + close(cancelCh) + + // Publish a notification so non-answering clients can dismiss + // their open question forms. + if batchID != "" { + s.notificationBroker.Publish(pubsub.CreatedEvent, Notification{ + BatchID: batchID, + }) + } + return true +} diff --git a/internal/server/docs.go b/internal/server/docs.go new file mode 100644 index 0000000000000000000000000000000000000000..ff9219ecaeec533336b1268a59c3ebd623cb618d --- /dev/null +++ b/internal/server/docs.go @@ -0,0 +1,41 @@ +package server + +import ( + "net/http" + + "github.com/charmbracelet/crush/internal/apigen" +) + +// apiInfo is the top-level metadata for the Crush API documentation. +var apiInfo = apigen.Info{ + Title: "Crush API", + Version: "1.0", +} + +// apiBase is the base path shared by all v1 endpoints. +const apiBase = "/v1" + +// handleDocsIndex serves the plain HTML documentation page. +func (c *controllerV1) handleDocsIndex(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != apiBase+"/docs/" { + http.Redirect(w, r, apiBase+"/docs/", http.StatusMovedPermanently) + return + } + w.Header().Set("Content-Type", "text/html; charset=utf-8") + if err := apigen.RenderDocs(w, apiInfo, apiBase, c.endpoints()); err != nil { + c.server.logError(r, "Failed to render docs", "error", err) + jsonError(w, http.StatusInternalServerError, "failed to render docs") + } +} + +// handleDocsSpec serves the OpenAPI 3.1 specification as JSON. +func (c *controllerV1) handleDocsSpec(w http.ResponseWriter, r *http.Request) { + data, err := apigen.SpecJSON(apiInfo, apiBase, c.endpoints()) + if err != nil { + c.server.logError(r, "Failed to render spec", "error", err) + jsonError(w, http.StatusInternalServerError, "failed to render spec") + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(data) +} diff --git a/internal/server/e2e_agent_test.go b/internal/server/e2e_agent_test.go new file mode 100644 index 0000000000000000000000000000000000000000..a4d9e5e2ccc24a6d37192868be39013e02960c58 --- /dev/null +++ b/internal/server/e2e_agent_test.go @@ -0,0 +1,742 @@ +package server + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "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/charmbracelet/crush/internal/pubsub" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +// scriptedCoordinator is an agent.Coordinator stub that mimics the +// externally-observable contract of a real run over the SSE pipeline +// without booting a real model, database, or scheduler. It publishes a +// user message when a run begins and an assistant message (with the +// appropriate FinishReason) when the run ends, exactly the way the real +// sessionAgent.Run surfaces a turn to SSE subscribers. +// +// A run blocks until either its per-session context is canceled (via +// Cancel, mirroring the explicit cancel endpoint) or the test releases +// it. On cancel it emits a FinishReasonCanceled assistant message and +// returns context.Canceled (which backend.runAgent swallows, so no +// AgentEvent error is published). On normal release it emits a +// FinishReasonEndTurn assistant message and returns nil. +// +// The internal scheduler signal points the PLAN's e2e cases reference +// (e.g. "before registration in activeRequests", "between +// activeRequests.Set and assistant create") are not exposed by the +// codebase, so this stub reproduces the documented black-box outcome by +// controlling run timing directly through blockEntered / release. +type scriptedCoordinator struct { + app *app.App + + // blockEntered, when non-nil, is signaled (once) right after a run + // is entered and before the user message is emitted, letting a test + // interleave a cancel with the dispatched goroutine. + blockEntered chan struct{} + + mu sync.Mutex + // cancels holds the cancel func for every in-flight run, keyed by a + // monotonic id so concurrent runs for the same session each get their + // own entry (a map keyed only by sessionID would let a second run + // overwrite the first's cancel func and leak it). + cancels map[int64]sessionCancel + // pendingCancels counts cancels that arrived for a session while a run + // was in flight; a run for that session consumes one on entry and + // cancels itself, modeling the cancel-on-entry path a follow-up takes. + pendingCancels map[string]int + nextRunID int64 + // entered carries the monotonic run id assigned to each run as it is + // entered, so a test can correlate a later assistant message back to a + // specific run (run 1 vs an accepted follow-up). + entered chan int64 + runStarts atomic.Int32 + + release chan struct{} +} + +type sessionCancel struct { + sessionID string + cancel context.CancelFunc +} + +func newScriptedCoordinator(a *app.App) *scriptedCoordinator { + return &scriptedCoordinator{ + app: a, + cancels: make(map[int64]sessionCancel), + pendingCancels: make(map[string]int), + entered: make(chan int64, 8), + release: make(chan struct{}), + } +} + +func (c *scriptedCoordinator) emitUser(sessionID, id string) { + c.app.SendEvent(pubsub.Event[message.Message]{ + Type: pubsub.CreatedEvent, + Payload: message.Message{ + ID: id, + SessionID: sessionID, + Role: message.User, + Parts: []message.ContentPart{message.TextContent{Text: "hi"}}, + }, + }) +} + +func (c *scriptedCoordinator) emitAssistant(sessionID, id string, reason message.FinishReason) { + c.app.SendEvent(pubsub.Event[message.Message]{ + Type: pubsub.CreatedEvent, + Payload: message.Message{ + ID: id, + SessionID: sessionID, + Role: message.Assistant, + Parts: []message.ContentPart{message.Finish{Reason: reason}}, + }, + }) +} + +func (c *scriptedCoordinator) Run(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + c.runStarts.Add(1) + runCtx, cancel := context.WithCancel(ctx) + + c.mu.Lock() + id := c.nextRunID + c.nextRunID++ + c.cancels[id] = sessionCancel{sessionID: sessionID, cancel: cancel} + // Cancel-on-entry: if a cancel for this session arrived while this + // run was still being dispatched (no run yet in flight to receive + // it), consume the pending cancel now so the run takes the canceled + // path instead of streaming output. + if c.pendingCancels[sessionID] > 0 { + c.pendingCancels[sessionID]-- + cancel() + } + c.mu.Unlock() + + select { + case c.entered <- id: + default: + } + + if c.blockEntered != nil { + select { + case <-c.blockEntered: + case <-runCtx.Done(): + } + } + + defer func() { + c.mu.Lock() + delete(c.cancels, id) + c.mu.Unlock() + cancel() + }() + + // Qualify the emitted message ids with the run id so a test can + // attribute an assistant message to the exact run that produced it + // (run 1 vs an accepted follow-up sharing the same session). + userID := fmt.Sprintf("u-%s-%d", sessionID, id) + asstID := fmt.Sprintf("a-%s-%d", sessionID, id) + + c.emitUser(sessionID, userID) + + // Cancellation takes priority: if the run was already canceled it + // must take the canceled path even when release is closed, so a + // canceled run never races into a normal FinishReasonEndTurn. + select { + case <-runCtx.Done(): + c.emitAssistant(sessionID, asstID, message.FinishReasonCanceled) + return nil, context.Canceled + default: + } + + select { + case <-c.release: + c.emitAssistant(sessionID, asstID, message.FinishReasonEndTurn) + return nil, nil + case <-runCtx.Done(): + c.emitAssistant(sessionID, asstID, message.FinishReasonCanceled) + return nil, context.Canceled + } +} + +func (c *scriptedCoordinator) RunAccepted(ctx context.Context, accept *agent.AcceptedRun, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + return c.Run(ctx, sessionID, prompt, attachments...) +} + +func (c *scriptedCoordinator) BeginAccepted(string) *agent.AcceptedRun { return nil } + +func (c *scriptedCoordinator) Cancel(sessionID string) { + c.mu.Lock() + defer c.mu.Unlock() + // Cancel every in-flight run for this session. Concurrent runs for + // the same session (an active run plus an accepted follow-up still + // dispatching) each hold their own entry, so all of them are torn + // down by a single per-session cancel. + var canceled int + for _, sc := range c.cancels { + if sc.sessionID == sessionID { + sc.cancel() + canceled++ + } + } + // If at least one run was in flight, arm a pending cancel so a + // follow-up that has been accepted but not yet entered Run takes the + // cancel-on-entry path. With no run in flight this is a no-op, + // mirroring the production guarantee that an idle cancel does not arm + // a pending cancel against the next prompt. + if canceled > 0 { + c.pendingCancels[sessionID]++ + } +} + +func (c *scriptedCoordinator) CancelAll() { + c.mu.Lock() + defer c.mu.Unlock() + for _, sc := range c.cancels { + sc.cancel() + } +} + +func (c *scriptedCoordinator) IsBusy() bool { return false } +func (c *scriptedCoordinator) IsSessionBusy(string) bool { return false } +func (c *scriptedCoordinator) QueuedPrompts(string) int { return 0 } +func (c *scriptedCoordinator) QueuedPromptsList(string) []string { return nil } +func (c *scriptedCoordinator) ClearQueue(string) {} +func (c *scriptedCoordinator) Summarize(context.Context, string) error { return nil } +func (c *scriptedCoordinator) Model() agent.Model { return agent.Model{} } +func (c *scriptedCoordinator) UpdateModels(context.Context) error { return nil } +func (c *scriptedCoordinator) GenerateTitle(context.Context, string, string) {} + +// agentE2EHarness extends the SSE harness with a scripted coordinator +// wired into the workspace's embedded app.App, so POST /agent drives a +// real backend.SendMessage dispatch whose emitted user/assistant +// messages fan out over the same SSE pipeline production uses. +type agentE2EHarness struct { + *e2eHarness + coord *scriptedCoordinator +} + +func newAgentE2EHarness(t *testing.T) *agentE2EHarness { + t.Helper() + + h := &e2eHarness{} + + appCtx, cancel := context.WithCancel(context.Background()) + a := app.NewForTest(appCtx) + coord := newScriptedCoordinator(a) + a.AgentCoordinator = coord + t.Cleanup(func() { + cancel() + a.ShutdownForTest() + }) + + h.installServer(t) + + ws := &backend.Workspace{ + ID: uuid.New().String(), + Path: t.TempDir(), + App: a, + } + backend.SetWorkspaceShutdownFnForTest(ws, func() {}) + backend.InsertWorkspaceForTest(h.backend, ws) + + h.workspace = ws + h.app = a + return &agentE2EHarness{e2eHarness: h, coord: coord} +} + +// postAgentHTTP drives POST /v1/workspaces/{id}/agent over the harness's +// httptest server and returns the status code. +func (h *agentE2EHarness) postAgentHTTP(t *testing.T, ctx context.Context, sessionID string) int { + t.Helper() + body, err := json.Marshal(proto.AgentMessage{SessionID: sessionID, Prompt: "hi"}) + require.NoError(t, err) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + h.httpSrv.URL+"/v1/workspaces/"+h.workspace.ID+"/agent", 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) + _, _ = io.Copy(io.Discard, resp.Body) + resp.Body.Close() + return resp.StatusCode +} + +// cancelAgentHTTP drives POST /v1/workspaces/{id}/agent/sessions/{sid}/cancel. +func (h *agentE2EHarness) cancelAgentHTTP(t *testing.T, ctx context.Context, sessionID string) int { + t.Helper() + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + h.httpSrv.URL+"/v1/workspaces/"+h.workspace.ID+"/agent/sessions/"+sessionID+"/cancel", nil) + require.NoError(t, err) + resp, err := h.httpSrv.Client().Do(req) + require.NoError(t, err) + _, _ = io.Copy(io.Discard, resp.Body) + resp.Body.Close() + return resp.StatusCode +} + +// waitForRunEntered blocks until a dispatched run for any session has +// been entered by the scripted coordinator, or fails the test. It +// returns the monotonic run id assigned to that run so a caller can +// correlate it with a later assistant message; callers that don't need +// the id can ignore the return value. +func (h *agentE2EHarness) waitForRunEntered(t *testing.T) int64 { + t.Helper() + select { + case id := <-h.coord.entered: + return id + case <-time.After(2 * time.Second): + t.Fatal("dispatched run was never entered") + return 0 + } +} + +// finishReason extracts the assistant message's FinishReason, if any. +func finishReason(m proto.Message) (proto.FinishReason, bool) { + for _, p := range m.Parts { + if f, ok := p.(proto.Finish); ok { + return f.Reason, true + } + } + return "", false +} + +// TestE2E_CancelByOtherClientDoesNotErrorPrompter covers PLAN Tests -> +// New end-to-end coverage item 1: a second client canceling a run does +// not surface a server error to the prompter; the run ends with a +// FinishReasonCanceled assistant message and no AgentEvent carries a +// non-nil Error. +func TestE2E_CancelByOtherClientDoesNotErrorPrompter(t *testing.T) { + t.Parallel() + h := newAgentE2EHarness(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) + + const sid = "s-cancel-other" + + // A posts a long-running prompt; the handler must return 202 + // immediately (the run blocks in the coordinator). + require.Equal(t, http.StatusAccepted, h.postAgentHTTP(t, ctx, sid)) + h.waitForRunEntered(t) + + // B cancels. + require.Equal(t, http.StatusOK, h.cancelAgentHTTP(t, ctx, sid)) + + // A's SSE stream receives the FinishReasonCanceled assistant + // message. + pickCtx, pickCancel := context.WithTimeout(ctx, 3*time.Second) + defer pickCancel() + got, ok := drainUntil(pickCtx, evcA, func(e pubsub.Event[proto.Message]) bool { + r, has := finishReason(e.Payload) + return e.Payload.Role == proto.Assistant && has && r == proto.FinishReasonCanceled + }) + require.True(t, ok, "client A must observe a FinishReasonCanceled assistant message") + require.Equal(t, sid, got.Payload.SessionID) + + // No AgentEvent error reaches A (cancel is not a server error). + errCtx, errCancel := context.WithTimeout(ctx, 250*time.Millisecond) + defer errCancel() + _, gotErrA := drainUntil(errCtx, evcA, func(e pubsub.Event[proto.AgentEvent]) bool { + return e.Payload.Type == proto.AgentEventTypeError && e.Payload.Error != nil + }) + require.False(t, gotErrA, "cancel must not surface an AgentEvent error to the prompter") + + // And no AgentEvent error reaches the canceling client B either; the + // PLAN requires that *no* client observes a non-nil Error. + errCtxB, errCancelB := context.WithTimeout(ctx, 250*time.Millisecond) + defer errCancelB() + _, gotErrB := drainUntil(errCtxB, evcB, func(e pubsub.Event[proto.AgentEvent]) bool { + return e.Payload.Type == proto.AgentEventTypeError && e.Payload.Error != nil + }) + require.False(t, gotErrB, "cancel must not surface an AgentEvent error to any client") +} + +// TestE2E_CancelImmediatelyAfter202IsNotLost covers PLAN item 1a: a +// cancel that races a freshly-dispatched run (before it would emit any +// output) is not lost. The run takes the cancel-on-entry path and emits +// a user message followed by a FinishReasonCanceled assistant message +// rather than streaming model output. +func TestE2E_CancelImmediatelyAfter202IsNotLost(t *testing.T) { + t.Parallel() + h := newAgentE2EHarness(t) + // Gate the run on a signal the test controls so the cancel can be + // observed while the dispatched goroutine is parked at entry. + h.coord.blockEntered = make(chan struct{}) + + ctx, cancel := context.WithCancel(t.Context()) + t.Cleanup(cancel) + + cid := uuid.New().String() + evc, cancelSSE := h.subscribeSSE(t, ctx, h.workspace.ID, cid) + t.Cleanup(cancelSSE) + h.waitForAttached(t, 1) + + const sid = "s-race-cancel" + require.Equal(t, http.StatusAccepted, h.postAgentHTTP(t, ctx, sid)) + h.waitForRunEntered(t) + + // Cancel while the run is still blocked at entry, then release it. + require.Equal(t, http.StatusOK, h.cancelAgentHTTP(t, ctx, sid)) + close(h.coord.blockEntered) + + pickCtx, pickCancel := context.WithTimeout(ctx, 3*time.Second) + defer pickCancel() + + gotUser, okUser := drainUntil(pickCtx, evc, func(e pubsub.Event[proto.Message]) bool { + return e.Payload.Role == proto.User && e.Payload.SessionID == sid + }) + require.True(t, okUser, "the canceled turn must still record a user message") + require.Equal(t, sid, gotUser.Payload.SessionID) + + gotAsst, okAsst := drainUntil(pickCtx, evc, func(e pubsub.Event[proto.Message]) bool { + r, has := finishReason(e.Payload) + return e.Payload.Role == proto.Assistant && has && r == proto.FinishReasonCanceled + }) + require.True(t, okAsst, "the canceled turn must end with a FinishReasonCanceled assistant message") + require.Equal(t, sid, gotAsst.Payload.SessionID) +} + +// TestE2E_IdleCancelDoesNotPoisonNextPrompt covers PLAN item 1b: an +// idle cancel (no active run) must not poison the next prompt. With the +// scripted coordinator the cancel records a pending entry only if a run +// is in flight; an idle cancel records one, but the documented +// guarantee is that the *next* prompt's outcome is observable. Here we +// assert the regression-relevant external behavior: after an idle +// cancel, a subsequent normal prompt is able to run and emit output. +// +// NOTE: This is a simplified version. The real "idle Escape must not +// poison" guarantee lives inside sessionAgent.Cancel's acceptedRuns +// gating, which is covered by the agent unit tests; the e2e stub cannot +// distinguish "truly idle" from "accepted but not yet running" without +// the internal acceptedRuns signal. See test summary. +func TestE2E_IdleCancelDoesNotPoisonNextPrompt(t *testing.T) { + t.Parallel() + h := newAgentE2EHarness(t) + ctx, cancel := context.WithCancel(t.Context()) + t.Cleanup(cancel) + + cid := uuid.New().String() + evc, cancelSSE := h.subscribeSSE(t, ctx, h.workspace.ID, cid) + t.Cleanup(cancelSSE) + h.waitForAttached(t, 1) + + const sid = "s-idle-cancel" + + // Idle cancel: no run in flight. The scripted coordinator drops it + // (no pending cancel recorded for a session that has no run), which + // models the production guarantee that an idle Escape does not arm + // a cancel against the next prompt. + require.Equal(t, http.StatusOK, h.cancelAgentHTTP(t, ctx, sid)) + + // Now a normal prompt; release it so it finishes successfully. + require.Equal(t, http.StatusAccepted, h.postAgentHTTP(t, ctx, sid)) + h.waitForRunEntered(t) + close(h.coord.release) + + pickCtx, pickCancel := context.WithTimeout(ctx, 3*time.Second) + defer pickCancel() + got, ok := drainUntil(pickCtx, evc, func(e pubsub.Event[proto.Message]) bool { + r, has := finishReason(e.Payload) + return e.Payload.Role == proto.Assistant && has && r == proto.FinishReasonEndTurn + }) + require.True(t, ok, "the next prompt after an idle cancel must run to FinishReasonEndTurn") + require.Equal(t, sid, got.Payload.SessionID) + + // And it must not be marked canceled. + canCtx, canCancel := context.WithTimeout(ctx, 200*time.Millisecond) + defer canCancel() + _, gotCanceled := drainUntil(canCtx, evc, func(e pubsub.Event[proto.Message]) bool { + r, has := finishReason(e.Payload) + return e.Payload.Role == proto.Assistant && has && r == proto.FinishReasonCanceled + }) + require.False(t, gotCanceled, "an idle cancel must not produce a FinishReasonCanceled marker on the next prompt") +} + +// TestE2E_CancelBetweenActiveSetAndAssistantCreate covers PLAN item 1d: +// a cancel that arrives after the run has begun but before it would +// create the assistant message must still produce a user message and a +// FinishReasonCanceled assistant message, never a silent return. The +// blockEntered gate parks the run after entry (modeling the window +// between activeRequests.Set and assistant creation). +func TestE2E_CancelBetweenActiveSetAndAssistantCreate(t *testing.T) { + t.Parallel() + h := newAgentE2EHarness(t) + h.coord.blockEntered = make(chan struct{}) + + ctx, cancel := context.WithCancel(t.Context()) + t.Cleanup(cancel) + + cid := uuid.New().String() + evc, cancelSSE := h.subscribeSSE(t, ctx, h.workspace.ID, cid) + t.Cleanup(cancelSSE) + h.waitForAttached(t, 1) + + const sid = "s-mid-window" + require.Equal(t, http.StatusAccepted, h.postAgentHTTP(t, ctx, sid)) + h.waitForRunEntered(t) + + // Cancel while parked at entry; then release so the run proceeds + // into its cancel branch (runCtx already canceled). + require.Equal(t, http.StatusOK, h.cancelAgentHTTP(t, ctx, sid)) + close(h.coord.blockEntered) + + pickCtx, pickCancel := context.WithTimeout(ctx, 3*time.Second) + defer pickCancel() + + _, okUser := drainUntil(pickCtx, evc, func(e pubsub.Event[proto.Message]) bool { + return e.Payload.Role == proto.User && e.Payload.SessionID == sid + }) + require.True(t, okUser, "a user message must be recorded for the canceled turn") + + gotAsst, okAsst := drainUntil(pickCtx, evc, func(e pubsub.Event[proto.Message]) bool { + r, has := finishReason(e.Payload) + return e.Payload.Role == proto.Assistant && has && r == proto.FinishReasonCanceled + }) + require.True(t, okAsst, "the run must not return silently; it must emit a FinishReasonCanceled assistant message") + require.Equal(t, sid, gotAsst.Payload.SessionID) + + // No AgentEvent error is published: a cancel in the + // activeRequests.Set -> assistant-create window is not a server + // error. + errCtx, errCancel := context.WithTimeout(ctx, 250*time.Millisecond) + defer errCancel() + _, gotErr := drainUntil(errCtx, evc, func(e pubsub.Event[proto.AgentEvent]) bool { + return e.Payload.Type == proto.AgentEventTypeError && e.Payload.Error != nil + }) + require.False(t, gotErr, "no AgentEvent error must be published for the canceled turn") +} + +// TestE2E_PromptRequestContextDoesNotOwnRun covers PLAN item 2: the +// prompting client's HTTP request context does not own the run. A POST +// with a very short request-context timeout still returns 202 before +// that context would expire, and the run keeps going (observed via SSE +// finishing normally after release). +func TestE2E_PromptRequestContextDoesNotOwnRun(t *testing.T) { + t.Parallel() + h := newAgentE2EHarness(t) + streamCtx, streamCancel := context.WithCancel(t.Context()) + t.Cleanup(streamCancel) + + cid := uuid.New().String() + evc, cancelSSE := h.subscribeSSE(t, streamCtx, h.workspace.ID, cid) + t.Cleanup(cancelSSE) + h.waitForAttached(t, 1) + + const sid = "s-short-req" + + // The POST request context times out almost immediately. The + // handler must still return 202 (fire-and-forget) and the run must + // survive past the request-context deadline. + reqCtx, reqCancel := context.WithTimeout(t.Context(), 50*time.Millisecond) + defer reqCancel() + require.Equal(t, http.StatusAccepted, h.postAgentHTTP(t, reqCtx, sid)) + h.waitForRunEntered(t) + + // Let the request context expire, then release the run. + <-reqCtx.Done() + close(h.coord.release) + + pickCtx, pickCancel := context.WithTimeout(streamCtx, 3*time.Second) + defer pickCancel() + got, ok := drainUntil(pickCtx, evc, func(e pubsub.Event[proto.Message]) bool { + r, has := finishReason(e.Payload) + return e.Payload.Role == proto.Assistant && has && r == proto.FinishReasonEndTurn + }) + require.True(t, ok, "the run must finish normally even after the prompting request context expired") + require.Equal(t, sid, got.Payload.SessionID) +} + +// TestE2E_AgentRunSurvivesAcrossWorkspaceClaims covers PLAN item 3: a +// run started by client A survives A detaching as long as another +// client (B) keeps the workspace alive; B observes the run finish via +// SSE. +func TestE2E_AgentRunSurvivesAcrossWorkspaceClaims(t *testing.T) { + t.Parallel() + h := newAgentE2EHarness(t) + + ctxA, cancelA := context.WithCancel(t.Context()) + ctxB, cancelB := context.WithCancel(t.Context()) + t.Cleanup(cancelB) + + 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) + + const sid = "s-survive" + // A is the poster; the run must outlive A detaching as long as B + // keeps the workspace alive. + require.Equal(t, http.StatusAccepted, h.postAgentHTTP(t, ctxA, sid)) + h.waitForRunEntered(t) + + // A detaches; B is still attached so the workspace stays alive. + cancelA() + killA() + require.Eventually(t, func() bool { + return backend.WorkspaceLiveStreamCountForTest(h.workspace) == 1 + }, 3*time.Second, 10*time.Millisecond, + "A detaching must leave B as the sole attached client") + require.False(t, h.shutdownHit.Load(), "workspace must stay alive while B is attached") + + // Release the run; B must still observe it finish. + close(h.coord.release) + pickCtx, pickCancel := context.WithTimeout(ctxB, 3*time.Second) + defer pickCancel() + got, ok := drainUntil(pickCtx, evcB, func(e pubsub.Event[proto.Message]) bool { + r, has := finishReason(e.Payload) + return e.Payload.Role == proto.Assistant && has && r == proto.FinishReasonEndTurn + }) + require.True(t, ok, "B must observe the run finish after A detaches") + require.Equal(t, sid, got.Payload.SessionID) +} + +// TestE2E_CancelOfActiveRunAlsoCancelsAcceptedFollowUp covers PLAN item +// 1c at the externally-observable level: while session sid has an active +// run, a second prompt for sid is accepted; a cancel for sid must cancel +// the active run and must not let the follow-up stream a normal +// FinishReasonEndTurn. +// +// The sequence follows the PLAN exactly: prompt 1 becomes the active +// run, prompt 2 for the same sid is accepted, then a cancel for sid +// fires, and only afterwards are any signals released. The scripted +// coordinator models the externally-observable contract of the +// busy-queue branch and pendingCancels (which depend on internal +// scheduler signals the codebase does not expose): a per-session cancel +// tears down every in-flight run for sid and arms a cancel-on-entry for +// a follow-up still dispatching. The invariant asserted is the one that +// matters: after the cancel, the active run ends canceled and the +// follow-up never streams a normal FinishReasonEndTurn. +func TestE2E_CancelOfActiveRunAlsoCancelsAcceptedFollowUp(t *testing.T) { + t.Parallel() + h := newAgentE2EHarness(t) + ctx, cancel := context.WithCancel(t.Context()) + t.Cleanup(cancel) + + cid := uuid.New().String() + evc, cancelSSE := h.subscribeSSE(t, ctx, h.workspace.ID, cid) + t.Cleanup(cancelSSE) + h.waitForAttached(t, 1) + + const sid = "s-followup" + + // (a) Prompt 1 for sid becomes the active run. Capture its run id so + // the canceled assistant message below can be attributed to run 1 + // unambiguously. + require.Equal(t, http.StatusAccepted, h.postAgentHTTP(t, ctx, sid)) + run1 := h.waitForRunEntered(t) + + // (b) Prompt 2 for the *same* sid is accepted while the active run + // is still in flight; it is the follow-up the PLAN describes + // (acceptedRuns > 0, either still dispatching or about to enter the + // busy-queue branch). + require.Equal(t, http.StatusAccepted, h.postAgentHTTP(t, ctx, sid)) + run2 := h.waitForRunEntered(t) + require.NotEqual(t, run1, run2, "the follow-up must be a distinct run from the active one") + + // (c) B cancels sid. This tears down every in-flight run for the + // session and arms a pending cancel for any follow-up that has not + // yet entered Run. + require.Equal(t, http.StatusOK, h.cancelAgentHTTP(t, ctx, sid)) + + // (d) Open the coordinator gate so any run that is NOT canceled would + // be free to proceed straight into the normal FinishReasonEndTurn + // branch. The scripted Run checks runCtx.Done() before the release + // select, so a canceled run still takes the canceled path even with + // release closed; only a non-canceled run reaches FinishReasonEndTurn. + // Releasing here is therefore what makes the assertions below + // meaningful: if the cancel had failed to tear down run 1 or arm the + // cancel-on-entry for the follow-up, the freed gate would let that run + // stream a normal FinishReasonEndTurn and the test would fail. + close(h.coord.release) + + pickCtx, pickCancel := context.WithTimeout(ctx, 3*time.Second) + defer pickCancel() + + // (e) Run 1 (the active run) must end with FinishReasonCanceled. The + // assistant message id is qualified with the run id, so matching on + // run1's id proves the cancellation is attributed to the FIRST run + // and not to the follow-up. + // + // The single drain below is also the negative assertion for run 2: + // the match closure inspects every assistant event for sid as it + // scans, and if it ever observes the follow-up (run 2) streaming a + // normal FinishReasonEndTurn it records that violation immediately. + // This is what makes the run-2 check sound: a previous two-phase + // approach could let this very drain consume and discard a run-2 + // EndTurn while still hunting for run 1's canceled message, leaving a + // later no-EndTurn check unable to prove run 2 stayed canceled. + // Folding the negative check into the same scan means a run-2 EndTurn + // can never slip past unobserved, whether it arrives before or after + // run 1's canceled message. + run1AsstID := fmt.Sprintf("a-%s-%d", sid, run1) + run2AsstID := fmt.Sprintf("a-%s-%d", sid, run2) + var followUpEndTurn bool + got, ok := drainUntil(pickCtx, evc, func(e pubsub.Event[proto.Message]) bool { + if e.Payload.SessionID != sid || e.Payload.Role != proto.Assistant { + return false + } + r, has := finishReason(e.Payload) + if !has { + return false + } + // Any normal model output for sid after the cancel is a + // violation. The follow-up (run 2) must never reach the + // FinishReasonEndTurn branch; flag it the moment it is seen so + // the assertion below fails even if this event arrives while we + // are still waiting for run 1's canceled message. + if r == proto.FinishReasonEndTurn { + if e.Payload.ID == run2AsstID || e.Payload.ID != run1AsstID { + followUpEndTurn = true + } + // Stop draining; the EndTurn observation is decisive and the + // require.False below will surface the failure. + return true + } + return e.Payload.ID == run1AsstID && r == proto.FinishReasonCanceled + }) + require.False(t, followUpEndTurn, "the accepted follow-up must not stream a normal FinishReasonEndTurn after the cancel") + require.True(t, ok, "the first (active) run must end with FinishReasonCanceled") + require.Equal(t, run1AsstID, got.Payload.ID, "the canceled message must belong to the first (active) run") + gotReason, gotHas := finishReason(got.Payload) + require.True(t, gotHas) + require.Equal(t, proto.FinishReasonCanceled, gotReason, "the matched run-1 message must be canceled, not a normal end turn") + require.Equal(t, sid, got.Payload.SessionID) + + // Confirm no normal FinishReasonEndTurn for sid is still in flight. + // By this point the scan above has already ruled out a run-2 EndTurn + // arriving before run 1's canceled message; this guards against one + // arriving afterward. + endCtx, endCancel := context.WithTimeout(ctx, 300*time.Millisecond) + defer endCancel() + _, gotEnd := drainUntil(endCtx, evc, func(e pubsub.Event[proto.Message]) bool { + r, has := finishReason(e.Payload) + return e.Payload.SessionID == sid && e.Payload.Role == proto.Assistant && has && r == proto.FinishReasonEndTurn + }) + require.False(t, gotEnd, "the accepted follow-up must not stream model output after the cancel") +} diff --git a/internal/server/endpoints.go b/internal/server/endpoints.go new file mode 100644 index 0000000000000000000000000000000000000000..242365925e6773a6dddeabe8f13f6a9c567b8e3c --- /dev/null +++ b/internal/server/endpoints.go @@ -0,0 +1,599 @@ +package server + +import ( + "github.com/charmbracelet/crush/internal/apigen" + "github.com/charmbracelet/crush/internal/proto" +) + +// endpoints is the single source of truth for the v1 API: every entry +// drives both mux registration and the documentation served under +// /v1/docs/. Add or change an endpoint here, never by hand elsewhere. +func (c *controllerV1) endpoints() []apigen.Endpoint { + return []apigen.Endpoint{ + apigen.Get("/v1/health"). + Summary("Health check"). + Tags("system"). + Handle(c.handleGetHealth), + + apigen.Get("/v1/version"). + Summary("Get server version"). + Tags("system"). + Responds(proto.VersionInfo{}). + Handle(c.handleGetVersion), + + apigen.Post("/v1/control"). + Summary("Send server control command"). + Description("Accepts a control command (e.g. shutdown, shutdown_if_idle). "+ + "Shutdown is conditional: only the backend can rule on idleness without "+ + "racing a session that arrives between a client's own check and its request."). + Tags("system"). + Accepts(proto.ServerControl{}). + Fails(400, 409). + Handle(c.handlePostControl), + + apigen.Get("/v1/config"). + Summary("Get server config"). + Tags("system"). + Handle(c.handleGetConfig), + + apigen.Delete("/v1/clients/{client_id}"). + Summary("Retire a client"). + Description("Releases every claim the client holds."). + Tags("system"). + PathParam("client_id", "Client ID (UUID)"). + Fails(400). + Handle(c.handleDeleteClient), + + apigen.Get("/v1/workspaces"). + Summary("List workspaces"). + Tags("workspaces"). + Responds([]proto.Workspace{}). + Handle(c.handleGetWorkspaces), + + apigen.Post("/v1/workspaces"). + Summary("Create workspace"). + Tags("workspaces"). + Accepts(proto.Workspace{}). + Responds(proto.Workspace{}). + Fails(400, 500). + Handle(c.handlePostWorkspaces), + + apigen.Get("/v1/workspaces/{id}"). + Summary("Get workspace"). + Tags("workspaces"). + PathParam("id", "Workspace ID"). + Responds(proto.Workspace{}). + Fails(404, 500). + Handle(c.handleGetWorkspace), + + apigen.Delete("/v1/workspaces/{id}"). + Summary("Delete workspace"). + Tags("workspaces"). + PathParam("id", "Workspace ID"). + Fails(404). + Handle(c.handleDeleteWorkspaces), + + apigen.Post("/v1/workspaces/{id}/current-session"). + Summary("Set current session for a client"). + Description("An empty session_id clears the entry (e.g. the client "+ + "is on the landing screen)."). + Tags("workspaces"). + PathParam("id", "Workspace ID"). + RequiredQueryParam("client_id", "Client ID (UUID)"). + Accepts(proto.CurrentSession{}). + Fails(400, 404). + Handle(c.handlePostWorkspaceCurrentSession), + + apigen.Get("/v1/workspaces/{id}/config"). + Summary("Get workspace config"). + Description("Returns the resolved workspace configuration (config.Config)."). + Tags("workspaces"). + PathParam("id", "Workspace ID"). + Fails(404, 500). + Handle(c.handleGetWorkspaceConfig), + + apigen.Get("/v1/workspaces/{id}/providers"). + Summary("Get workspace providers"). + Description("Lists the providers available to the workspace."). + Tags("workspaces"). + PathParam("id", "Workspace ID"). + Fails(404, 500). + Handle(c.handleGetWorkspaceProviders), + + apigen.Get("/v1/workspaces/{id}/events"). + Summary("Stream workspace events (SSE)"). + Description("Streams workspace events as Server-Sent Events. The client must "+ + "identify itself with the client_id query parameter; the stream attaches "+ + "the client to the workspace until it disconnects."). + Tags("workspaces"). + PathParam("id", "Workspace ID"). + RequiredQueryParam("client_id", "Client ID (UUID)"). + SSE(). + Fails(400, 404, 500). + Handle(c.handleGetWorkspaceEvents), + + apigen.Get("/v1/workspaces/{id}/sessions"). + Summary("List sessions"). + Tags("sessions"). + PathParam("id", "Workspace ID"). + Responds([]proto.Session{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceSessions), + + apigen.Post("/v1/workspaces/{id}/sessions"). + Summary("Create session"). + Tags("sessions"). + PathParam("id", "Workspace ID"). + Accepts(proto.Session{}). + Responds(proto.Session{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceSessions), + + apigen.Get("/v1/workspaces/{id}/sessions/{sid}"). + Summary("Get session"). + Tags("sessions"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Responds(proto.Session{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceSession), + + apigen.Put("/v1/workspaces/{id}/sessions/{sid}"). + Summary("Update session"). + Tags("sessions"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Accepts(proto.Session{}). + Responds(proto.Session{}). + Fails(400, 404, 500). + Handle(c.handlePutWorkspaceSession), + + apigen.Delete("/v1/workspaces/{id}/sessions/{sid}"). + Summary("Delete session"). + Tags("sessions"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Fails(404, 500). + Handle(c.handleDeleteWorkspaceSession), + + apigen.Get("/v1/workspaces/{id}/sessions/{sid}/history"). + Summary("Get session history"). + Tags("sessions"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Fails(404, 500). + Handle(c.handleGetWorkspaceSessionHistory), + + apigen.Get("/v1/workspaces/{id}/sessions/{sid}/messages"). + Summary("Get session messages"). + Tags("sessions"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Responds([]proto.Message{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceSessionMessages), + + apigen.Get("/v1/workspaces/{id}/sessions/{sid}/messages/user"). + Summary("Get user messages for session"). + Tags("sessions"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Responds([]proto.Message{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceSessionUserMessages), + + apigen.Get("/v1/workspaces/{id}/messages/user"). + Summary("Get all user messages for workspace"). + Tags("workspaces"). + PathParam("id", "Workspace ID"). + Responds([]proto.Message{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceAllUserMessages), + + apigen.Get("/v1/workspaces/{id}/sessions/{sid}/filetracker/files"). + Summary("List tracked files for session"). + Tags("filetracker"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Responds([]string{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceSessionFileTrackerFiles), + + apigen.Post("/v1/workspaces/{id}/filetracker/read"). + Summary("Record file read"). + Tags("filetracker"). + PathParam("id", "Workspace ID"). + Accepts(proto.FileTrackerReadRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceFileTrackerRead), + + apigen.Get("/v1/workspaces/{id}/filetracker/lastread"). + Summary("Get last read time for file"). + Tags("filetracker"). + PathParam("id", "Workspace ID"). + QueryParam("session_id", "Session ID"). + RequiredQueryParam("path", "File path"). + Fails(404, 500). + Handle(c.handleGetWorkspaceFileTrackerLastRead), + + apigen.Get("/v1/workspaces/{id}/lsps"). + Summary("List LSP clients"). + Tags("lsp"). + PathParam("id", "Workspace ID"). + Responds(map[string]proto.LSPClientInfo{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceLSPs), + + apigen.Get("/v1/workspaces/{id}/lsps/{lsp}/diagnostics"). + Summary("Get LSP diagnostics"). + Tags("lsp"). + PathParam("id", "Workspace ID"). + PathParam("lsp", "LSP client name"). + Fails(404, 500). + Handle(c.handleGetWorkspaceLSPDiagnostics), + + apigen.Post("/v1/workspaces/{id}/lsps/start"). + Summary("Start LSP server"). + Tags("lsp"). + PathParam("id", "Workspace ID"). + Accepts(proto.LSPStartRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceLSPStart), + + apigen.Post("/v1/workspaces/{id}/lsps/stop"). + Summary("Stop all LSP servers"). + Tags("lsp"). + PathParam("id", "Workspace ID"). + Fails(404, 500). + Handle(c.handlePostWorkspaceLSPStopAll), + + apigen.Get("/v1/workspaces/{id}/permissions/skip"). + Summary("Get skip permissions status"). + Tags("permissions"). + PathParam("id", "Workspace ID"). + Responds(proto.PermissionSkipRequest{}). + Fails(404, 500). + Handle(c.handleGetWorkspacePermissionsSkip), + + apigen.Post("/v1/workspaces/{id}/permissions/skip"). + Summary("Set skip permissions"). + Tags("permissions"). + PathParam("id", "Workspace ID"). + Accepts(proto.PermissionSkipRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspacePermissionsSkip), + + apigen.Post("/v1/workspaces/{id}/permissions/grant"). + Summary("Grant permission"). + Tags("permissions"). + PathParam("id", "Workspace ID"). + Accepts(proto.PermissionGrant{}). + Responds(proto.PermissionGrantResponse{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspacePermissionsGrant), + + apigen.Post("/v1/workspaces/{id}/questions/answer"). + Summary("Answer question batch"). + Tags("questions"). + PathParam("id", "Workspace ID"). + Accepts(proto.QuestionAnswer{}). + Responds(proto.QuestionAnswerResponse{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceQuestionsAnswer), + + apigen.Post("/v1/workspaces/{id}/questions/cancel"). + Summary("Cancel question batch"). + Tags("questions"). + PathParam("id", "Workspace ID"). + Responds(proto.QuestionAnswerResponse{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceQuestionsCancel), + + apigen.Get("/v1/workspaces/{id}/agent"). + Summary("Get agent info"). + Tags("agent"). + PathParam("id", "Workspace ID"). + Responds(proto.AgentInfo{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceAgent), + + apigen.Post("/v1/workspaces/{id}/agent"). + Summary("Send message to agent"). + Description("Validates and accepts the prompt, then dispatches the run "+ + "detached from the requesting HTTP connection: the run survives client "+ + "disconnects and is only ended by the explicit cancel endpoint."). + Tags("agent"). + PathParam("id", "Workspace ID"). + Accepts(proto.AgentMessage{}). + Status(202). + Fails(400, 404, 409, 500). + Handle(c.handlePostWorkspaceAgent), + + apigen.Post("/v1/workspaces/{id}/agent/init"). + Summary("Initialize agent"). + Description("The request body is optional."). + Tags("agent"). + PathParam("id", "Workspace ID"). + Accepts(proto.AgentInitRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceAgentInit), + + apigen.Post("/v1/workspaces/{id}/agent/update"). + Summary("Update agent"). + Tags("agent"). + PathParam("id", "Workspace ID"). + Fails(404, 500). + Handle(c.handlePostWorkspaceAgentUpdate), + + apigen.Get("/v1/workspaces/{id}/agent/sessions/{sid}"). + Summary("Get agent session"). + Tags("agent"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Responds(proto.AgentSession{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceAgentSession), + + apigen.Post("/v1/workspaces/{id}/agent/sessions/{sid}/cancel"). + Summary("Cancel agent session"). + Tags("agent"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Fails(404, 500). + Handle(c.handlePostWorkspaceAgentSessionCancel), + + apigen.Get("/v1/workspaces/{id}/agent/sessions/{sid}/prompts/queued"). + Summary("Get queued prompt status"). + Description("Returns the number of queued prompts for the session."). + Tags("agent"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Responds(0). + Fails(404, 500). + Handle(c.handleGetWorkspaceAgentSessionPromptQueued), + + apigen.Get("/v1/workspaces/{id}/agent/sessions/{sid}/prompts/list"). + Summary("List queued prompts"). + Tags("agent"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Responds([]string{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceAgentSessionPromptList), + + apigen.Post("/v1/workspaces/{id}/agent/sessions/{sid}/prompts/clear"). + Summary("Clear prompt queue"). + Tags("agent"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Fails(404, 500). + Handle(c.handlePostWorkspaceAgentSessionPromptClear), + + apigen.Post("/v1/workspaces/{id}/agent/sessions/{sid}/summarize"). + Summary("Summarize session"). + Tags("agent"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Fails(404, 500). + Handle(c.handlePostWorkspaceAgentSessionSummarize), + + apigen.Post("/v1/workspaces/{id}/agent/sessions/{sid}/shell"). + Summary("Run shell command"). + Description("Runs a shell command in the workspace on behalf of the session."). + Tags("agent"). + PathParam("id", "Workspace ID"). + PathParam("sid", "Session ID"). + Accepts(proto.ShellCommandRequest{}). + Responds(proto.ShellCommandResponse{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceAgentSessionShell), + + apigen.Get("/v1/workspaces/{id}/agent/default-small-model"). + Summary("Get default small model"). + Description("Returns the default small model (config.SelectedModel) for the provider."). + Tags("agent"). + PathParam("id", "Workspace ID"). + QueryParam("provider_id", "Provider ID"). + Fails(404, 500). + Handle(c.handleGetWorkspaceAgentDefaultSmallModel), + + apigen.Post("/v1/workspaces/{id}/config/set"). + Summary("Set a config field"). + Tags("config"). + PathParam("id", "Workspace ID"). + Accepts(proto.ConfigSetRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceConfigSet), + + apigen.Post("/v1/workspaces/{id}/config/remove"). + Summary("Remove a config field"). + Tags("config"). + PathParam("id", "Workspace ID"). + Accepts(proto.ConfigRemoveRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceConfigRemove), + + apigen.Post("/v1/workspaces/{id}/config/model"). + Summary("Set the preferred model"). + Tags("config"). + PathParam("id", "Workspace ID"). + Accepts(proto.ConfigModelRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceConfigModel), + + apigen.Post("/v1/workspaces/{id}/config/compact"). + Summary("Set compact mode"). + Tags("config"). + PathParam("id", "Workspace ID"). + Accepts(proto.ConfigCompactRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceConfigCompact), + + apigen.Post("/v1/workspaces/{id}/config/provider-key"). + Summary("Set provider API key"). + Tags("config"). + PathParam("id", "Workspace ID"). + Accepts(proto.ConfigProviderKeyRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceConfigProviderKey), + + apigen.Post("/v1/workspaces/{id}/config/import-copilot"). + Summary("Import Copilot credentials"). + Tags("config"). + PathParam("id", "Workspace ID"). + Responds(proto.ImportCopilotResponse{}). + Fails(404, 500). + Handle(c.handlePostWorkspaceConfigImportCopilot), + + apigen.Post("/v1/workspaces/{id}/config/refresh-oauth"). + Summary("Refresh OAuth token"). + Tags("config"). + PathParam("id", "Workspace ID"). + Accepts(proto.ConfigRefreshOAuthRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceConfigRefreshOAuth), + + apigen.Get("/v1/workspaces/{id}/project/needs-init"). + Summary("Check if project needs initialization"). + Tags("project"). + PathParam("id", "Workspace ID"). + Responds(proto.ProjectNeedsInitResponse{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceProjectNeedsInit), + + apigen.Post("/v1/workspaces/{id}/project/init"). + Summary("Mark project as initialized"). + Tags("project"). + PathParam("id", "Workspace ID"). + Fails(404, 500). + Handle(c.handlePostWorkspaceProjectInit), + + apigen.Get("/v1/workspaces/{id}/project/init-prompt"). + Summary("Get project initialization prompt"). + Tags("project"). + PathParam("id", "Workspace ID"). + Responds(proto.ProjectInitPromptResponse{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceProjectInitPrompt), + + apigen.Get("/v1/workspaces/{id}/skills"). + Summary("List visible skills"). + Tags("skills"). + PathParam("id", "Workspace ID"). + Responds([]proto.SkillInfo{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceSkills), + + apigen.Post("/v1/workspaces/{id}/skills/read"). + Summary("Read skill content"). + Tags("skills"). + PathParam("id", "Workspace ID"). + Accepts(proto.ReadSkillRequest{}). + Responds(proto.ReadSkillResponse{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceSkillRead), + + apigen.Post("/v1/workspaces/{id}/mcp/refresh-tools"). + Summary("Refresh MCP tools"). + Tags("mcp"). + PathParam("id", "Workspace ID"). + Accepts(proto.MCPNameRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceMCPRefreshTools), + + apigen.Post("/v1/workspaces/{id}/mcp/read-resource"). + Summary("Read MCP resource"). + Tags("mcp"). + PathParam("id", "Workspace ID"). + Accepts(proto.MCPReadResourceRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceMCPReadResource), + + apigen.Get("/v1/workspaces/{id}/mcp/prompts"). + Summary("Get MCP prompts"). + Tags("mcp"). + PathParam("id", "Workspace ID"). + Responds([]proto.MCPPrompt{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceMCPPrompts), + + apigen.Post("/v1/workspaces/{id}/mcp/get-prompt"). + Summary("Get MCP prompt"). + Tags("mcp"). + PathParam("id", "Workspace ID"). + Accepts(proto.MCPGetPromptRequest{}). + Responds(proto.MCPGetPromptResponse{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceMCPGetPrompt), + + apigen.Get("/v1/workspaces/{id}/mcp/states"). + Summary("Get MCP client states"). + Tags("mcp"). + PathParam("id", "Workspace ID"). + Responds(map[string]proto.MCPClientInfo{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceMCPStates), + + apigen.Get("/v1/workspaces/{id}/mcp/pending-auth"). + Summary("Get MCP servers pending OAuth"). + Tags("mcp"). + PathParam("id", "Workspace ID"). + Responds([]proto.MCPPendingAuthServer{}). + Fails(404, 500). + Handle(c.handleGetWorkspaceMCPPendingAuth), + + apigen.Get("/v1/workspaces/{id}/mcp/auth-url"). + Summary("Get MCP OAuth authorization URL"). + Description("Returns the current OAuth authorization URL for a named MCP "+ + "server, if a flow is in progress."). + Tags("mcp"). + PathParam("id", "Workspace ID"). + RequiredQueryParam("name", "MCP server name"). + Responds(proto.MCPAuthResponse{}). + Fails(400). + Handle(c.handleGetWorkspaceMCPAuthURL), + + apigen.Post("/v1/workspaces/{id}/mcp/auth"). + Summary("Authenticate an MCP server"). + Description("Runs the OAuth flow for a named MCP server. The local browser "+ + "is suppressed on the server; the client polls pending-auth and auth-url "+ + "to surface the authorization URL. The call blocks until the flow "+ + "completes or the request context is cancelled."). + Tags("mcp"). + PathParam("id", "Workspace ID"). + Accepts(proto.MCPNameRequest{}). + Responds(proto.MCPAuthResponse{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceMCPAuth), + + apigen.Post("/v1/workspaces/{id}/mcp/refresh-prompts"). + Summary("Refresh MCP prompts"). + Tags("mcp"). + PathParam("id", "Workspace ID"). + Accepts(proto.MCPNameRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceMCPRefreshPrompts), + + apigen.Post("/v1/workspaces/{id}/mcp/refresh-resources"). + Summary("Refresh MCP resources"). + Tags("mcp"). + PathParam("id", "Workspace ID"). + Accepts(proto.MCPNameRequest{}). + Fails(400, 404, 500). + Handle(c.handlePostWorkspaceMCPRefreshResources), + + apigen.Post("/v1/workspaces/{id}/mcp/docker/enable"). + Summary("Enable Docker MCP"). + Tags("mcp"). + PathParam("id", "Workspace ID"). + Fails(404, 500). + Handle(c.handlePostWorkspaceMCPEnableDocker), + + apigen.Post("/v1/workspaces/{id}/mcp/docker/disable"). + Summary("Disable Docker MCP"). + Tags("mcp"). + PathParam("id", "Workspace ID"). + Fails(404, 500). + Handle(c.handlePostWorkspaceMCPDisableDocker), + } +} diff --git a/internal/server/endpoints_test.go b/internal/server/endpoints_test.go new file mode 100644 index 0000000000000000000000000000000000000000..1afba49999e5f1d95ced3dc6566f73ed0da9a499 --- /dev/null +++ b/internal/server/endpoints_test.go @@ -0,0 +1,110 @@ +package server + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "regexp" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +var braceParam = regexp.MustCompile(`\{(\w+)\}`) + +// TestEndpointsPathParams verifies every {param} in a route pattern is +// documented by a PathParam of the same name, and vice versa. +func TestEndpointsPathParams(t *testing.T) { + t.Parallel() + c := &controllerV1{} + for _, e := range c.endpoints() { + declared := map[string]bool{} + for _, p := range e.Params() { + if p.In == "path" { + declared[p.Name] = true + } + } + seen := map[string]bool{} + for _, m := range braceParam.FindAllStringSubmatch(e.Path(), -1) { + seen[m[1]] = true + require.True(t, declared[m[1]], + "%s %s: path param {%s} missing PathParam documentation", e.Method(), e.Path(), m[1]) + } + for name := range declared { + require.True(t, seen[name], + "%s %s: PathParam %q has no matching {%s} in path", e.Method(), e.Path(), name, name) + } + } +} + +// TestEndpointsUnique verifies no method+path pair is registered twice. +func TestEndpointsUnique(t *testing.T) { + t.Parallel() + c := &controllerV1{} + seen := map[string]bool{} + for _, e := range c.endpoints() { + key := e.Method() + " " + e.Path() + require.False(t, seen[key], "duplicate endpoint %s", key) + seen[key] = true + } +} + +// TestEndpointsComplete verifies the registry covers every handler +// method on controllerV1 and that handlers are always set. +func TestEndpointsHandlers(t *testing.T) { + t.Parallel() + c := &controllerV1{} + for _, e := range c.endpoints() { + require.NotNil(t, e.Handler(), "%s %s: nil handler", e.Method(), e.Path()) + require.NotEmpty(t, e.Summary(), "%s %s: missing summary", e.Method(), e.Path()) + } +} + +// TestDocsEndpoints serves the docs routes from a bare server and +// checks the HTML page and OpenAPI spec. +func TestDocsEndpoints(t *testing.T) { + t.Parallel() + srv := &Server{} + srv.installHandler() + hs := httptest.NewServer(srv.Handler()) + t.Cleanup(hs.Close) + + get := func(path string) *http.Response { + t.Helper() + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, hs.URL+path, nil) + require.NoError(t, err) + resp, err := hs.Client().Do(req) + require.NoError(t, err) + return resp + } + + resp := get("/v1/docs/") + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + require.Contains(t, resp.Header.Get("Content-Type"), "text/html") + + resp = get("/v1/docs/openapi.json") + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + var spec struct { + OpenAPI string `json:"openapi"` + Paths map[string]map[string]json.RawMessage `json:"paths"` + } + require.NoError(t, json.NewDecoder(resp.Body).Decode(&spec)) + require.Equal(t, "3.2.1", spec.OpenAPI) + + // One spec operation per registered endpoint (paths with multiple + // methods contribute one operation each). + c := &controllerV1{} + count := 0 + for _, e := range c.endpoints() { + path := strings.TrimPrefix(e.Path(), apiBase) + item, ok := spec.Paths[path] + require.True(t, ok, "spec missing path %s", path) + require.Contains(t, item, strings.ToLower(e.Method()), "spec missing %s %s", e.Method(), path) + count++ + } + require.Len(t, c.endpoints(), count) +} diff --git a/internal/server/events.go b/internal/server/events.go new file mode 100644 index 0000000000000000000000000000000000000000..505aa9f90012027e5a463b96f6dbceffda3513ce --- /dev/null +++ b/internal/server/events.go @@ -0,0 +1,388 @@ +package server + +import ( + "encoding/json" + "errors" + "fmt" + "log/slog" + + "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/backend" + "github.com/charmbracelet/crush/internal/history" + "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/question" + "github.com/charmbracelet/crush/internal/session" + "github.com/charmbracelet/crush/internal/skills" +) + +// wrapEvent converts a raw tea.Msg (a pubsub.Event[T] from the app +// event fan-in) into a pubsub.Payload envelope with the correct +// PayloadType discriminator and a proto-typed inner payload that has +// proper JSON tags. Returns nil if the event type is unrecognized. +func wrapEvent(ev any) *pubsub.Payload { + switch e := ev.(type) { + case pubsub.Event[app.LSPEvent]: + return envelope(pubsub.PayloadTypeLSPEvent, pubsub.Event[proto.LSPEvent]{ + Type: e.Type, + Payload: proto.LSPEvent{ + Type: proto.LSPEventType(e.Payload.Type), + Name: e.Payload.Name, + State: e.Payload.State, + Error: e.Payload.Error, + DiagnosticCount: e.Payload.DiagnosticCount, + }, + }) + case pubsub.Event[mcp.Event]: + pt := mcpEventTypeToProto(e.Payload.Type) + if pt == "" { + // Unsupported MCP event type (e.g. EventChannelMessage, which + // has no proto representation until session delivery is wired + // up). Drop it instead of fabricating a state_changed event. + slog.Debug("Dropping unsupported MCP event type for SSE", "type", e.Payload.Type) + return nil + } + return envelope(pubsub.PayloadTypeMCPEvent, pubsub.Event[proto.MCPEvent]{ + Type: e.Type, + Payload: proto.MCPEvent{ + Type: pt, + Name: e.Payload.Name, + State: proto.MCPState(e.Payload.State), + Error: e.Payload.Error, + ToolCount: e.Payload.Counts.Tools, + }, + }) + case pubsub.Event[permission.PermissionRequest]: + return envelope(pubsub.PayloadTypePermissionRequest, pubsub.Event[proto.PermissionRequest]{ + Type: e.Type, + Payload: proto.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[permission.PermissionNotification]: + return envelope(pubsub.PayloadTypePermissionNotification, pubsub.Event[proto.PermissionNotification]{ + Type: e.Type, + Payload: proto.PermissionNotification{ + ToolCallID: e.Payload.ToolCallID, + Granted: e.Payload.Granted, + Denied: e.Payload.Denied, + }, + }) + case pubsub.Event[question.Request]: + slog.Info("Wrapping question batch event for SSE", "id", e.Payload.ID, "questions", len(e.Payload.Questions)) + return envelope(pubsub.PayloadTypeQuestionRequest, pubsub.Event[proto.QuestionRequest]{ + Type: e.Type, + Payload: proto.QuestionRequest{ + ID: e.Payload.ID, + SessionID: e.Payload.SessionID, + ToolCallID: e.Payload.ToolCallID, + Questions: questionsToProto(e.Payload.Questions), + ConfirmTitle: e.Payload.ConfirmTitle, + ConfirmDescription: e.Payload.ConfirmDescription, + }, + }) + case pubsub.Event[question.Notification]: + return envelope(pubsub.PayloadTypeQuestionNotification, pubsub.Event[proto.QuestionNotification]{ + Type: e.Type, + Payload: proto.QuestionNotification{ + BatchID: e.Payload.BatchID, + }, + }) + case pubsub.Event[message.Message]: + return envelope(pubsub.PayloadTypeMessage, pubsub.Event[proto.Message]{ + Type: e.Type, + Payload: messageToProto(e.Payload), + }) + case pubsub.Event[session.Session]: + return envelope(pubsub.PayloadTypeSession, pubsub.Event[proto.Session]{ + Type: e.Type, + Payload: sessionToProto(e.Payload), + }) + case pubsub.Event[history.File]: + return envelope(pubsub.PayloadTypeFile, pubsub.Event[proto.File]{ + Type: e.Type, + Payload: fileToProto(e.Payload), + }) + case pubsub.Event[notify.Notification]: + payload := proto.AgentEvent{ + SessionID: e.Payload.SessionID, + SessionTitle: e.Payload.SessionTitle, + RunID: e.Payload.RunID, + Type: proto.AgentEventType(e.Payload.Type), + AWSSOCommand: e.Payload.AWSSOCommand, + AWSSOURL: e.Payload.AWSSOURL, + } + // Carry any human-readable message across the wire; the client + // maps Error back into Notification.Message. + if e.Payload.Message != "" { + payload.Error = errors.New(e.Payload.Message) + } + if e.Payload.Type == notify.TypeAgentError { + payload.Type = proto.AgentEventTypeError + } + return envelope(pubsub.PayloadTypeAgentEvent, pubsub.Event[proto.AgentEvent]{ + Type: e.Type, + Payload: payload, + }) + case pubsub.Event[notify.RunComplete]: + return envelope(pubsub.PayloadTypeRunComplete, pubsub.Event[proto.RunComplete]{ + Type: e.Type, + Payload: proto.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.ConfigChanged]: + return envelope(pubsub.PayloadTypeConfigChanged, e) + case app.UpdateAvailableMsg: + return envelope(pubsub.PayloadTypeUpdateAvailable, pubsub.Event[proto.UpdateAvailable]{ + Type: pubsub.UpdatedEvent, + Payload: proto.UpdateAvailable{ + CurrentVersion: e.CurrentVersion, + LatestVersion: e.LatestVersion, + IsDevelopment: e.IsDevelopment, + }, + }) + case pubsub.Event[skills.Event]: + return envelope(pubsub.PayloadTypeSkillsEvent, pubsub.Event[proto.SkillsEvent]{ + Type: e.Type, + Payload: skillsEventToProto(e.Payload), + }) + default: + slog.Warn("Unrecognized event type for SSE wrapping", "type", fmt.Sprintf("%T", ev)) + return nil + } +} + +// envelope marshals the inner event and wraps it in a pubsub.Payload. +func envelope(payloadType pubsub.PayloadType, inner any) *pubsub.Payload { + raw, err := json.Marshal(inner) + if err != nil { + slog.Error("Failed to marshal event payload", "error", err) + return nil + } + return &pubsub.Payload{ + Type: payloadType, + Payload: raw, + } +} + +func mcpEventTypeToProto(t mcp.EventType) proto.MCPEventType { + switch t { + case mcp.EventStateChanged: + return proto.MCPEventStateChanged + case mcp.EventToolsListChanged: + return proto.MCPEventToolsListChanged + case mcp.EventPromptsListChanged: + return proto.MCPEventPromptsListChanged + case mcp.EventResourcesListChanged: + return proto.MCPEventResourcesListChanged + default: + // Unsupported type (e.g. EventChannelMessage). Return empty so + // callers can drop it rather than coercing to state_changed. + return "" + } +} + +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, + } +} + +// isSessionBusy reports whether the given workspace has an in-flight +// agent run for sessionID. It tolerates a nil workspace (treating it as +// "not busy") so REST handlers can pass GetWorkspace's result through +// unconditionally — the workspace lookup error is already surfaced by +// the prior ListSessions/GetSession call when relevant. +func isSessionBusy(ws *backend.Workspace, sessionID string) bool { + if ws == nil || ws.App == nil || ws.AgentCoordinator == nil { + return false + } + return ws.AgentCoordinator.IsSessionBusy(sessionID) +} + +// attachedClients returns the number of clients currently viewing +// sessionID in ws. Hold-only clients (streams == 0) do not contribute. +// A nil workspace is treated as zero so handlers can pass GetWorkspace's +// result through without an extra guard. +func attachedClients(ws *backend.Workspace, sessionID string) int { + if ws == nil { + return 0 + } + return ws.AttachedClientsForSession(sessionID) +} + +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 fileToProto(f history.File) proto.File { + return proto.File{ + ID: f.ID, + SessionID: f.SessionID, + Path: f.Path, + Content: f.Content, + Version: f.Version, + CreatedAt: f.CreatedAt, + UpdatedAt: f.UpdatedAt, + } +} + +func messageToProto(m message.Message) proto.Message { + msg := proto.Message{ + ID: m.ID, + SessionID: m.SessionID, + Role: proto.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 message.TextContent: + msg.Parts = append(msg.Parts, proto.TextContent{Text: v.Text}) + case message.ReasoningContent: + msg.Parts = append(msg.Parts, proto.ReasoningContent{ + Thinking: v.Thinking, + Signature: v.Signature, + StartedAt: v.StartedAt, + FinishedAt: v.FinishedAt, + }) + case message.ToolCall: + msg.Parts = append(msg.Parts, proto.ToolCall{ + ID: v.ID, + Name: v.Name, + Input: v.Input, + Finished: v.Finished, + }) + case message.ToolResult: + msg.Parts = append(msg.Parts, proto.ToolResult{ + ToolCallID: v.ToolCallID, + Name: v.Name, + Content: v.Content, + Data: v.Data, + MIMEType: v.MIMEType, + Metadata: v.Metadata, + IsError: v.IsError, + }) + case message.Finish: + msg.Parts = append(msg.Parts, proto.Finish{ + Reason: proto.FinishReason(v.Reason), + Time: v.Time, + Message: v.Message, + Details: v.Details, + }) + case message.ImageURLContent: + msg.Parts = append(msg.Parts, proto.ImageURLContent{URL: v.URL, Detail: v.Detail}) + case message.BinaryContent: + msg.Parts = append(msg.Parts, proto.BinaryContent{Path: v.Path, MIMEType: v.MIMEType, Data: v.Data}) + case message.ShellCommand: + msg.Parts = append(msg.Parts, proto.ShellCommand{ + Command: v.Command, + Output: v.Output, + ExitCode: v.ExitCode, + }) + } + } + + return msg +} + +// skillsEventToProto converts a skills.Event into its wire form. Errors +// are flattened to strings because error does not round-trip over JSON. +func skillsEventToProto(e skills.Event) proto.SkillsEvent { + if len(e.States) == 0 { + return proto.SkillsEvent{} + } + out := proto.SkillsEvent{States: make([]proto.SkillState, len(e.States))} + for i, s := range e.States { + entry := proto.SkillState{ + Name: s.Name, + Path: s.Path, + State: proto.SkillDiscoveryState(s.State), + } + if s.Err != nil { + entry.Error = s.Err.Error() + } + out.States[i] = entry + } + return out +} + +func messagesToProto(msgs []message.Message) []proto.Message { + out := make([]proto.Message, len(msgs)) + for i, m := range msgs { + out[i] = messageToProto(m) + } + return out +} + +func questionsToProto(qs []question.Question) []proto.QuestionItem { + if len(qs) == 0 { + return nil + } + out := make([]proto.QuestionItem, len(qs)) + for i, q := range qs { + choices := make([]proto.QuestionChoice, len(q.Choices)) + for j, c := range q.Choices { + choices[j] = proto.QuestionChoice{ + ID: c.ID, + Label: c.Label, + Description: c.Description, + } + } + out[i] = proto.QuestionItem{ + ID: q.ID, + Type: string(q.Type), + Label: q.Label, + Question: q.Text, + Description: q.Description, + Choices: choices, + } + } + return out +} diff --git a/internal/server/events_test.go b/internal/server/events_test.go new file mode 100644 index 0000000000000000000000000000000000000000..b0e726e0218b3ab0029cca59f7e8a3f1c6cc666b --- /dev/null +++ b/internal/server/events_test.go @@ -0,0 +1,276 @@ +package server + +import ( + "encoding/json" + "errors" + "testing" + + "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/message" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/pubsub" + "github.com/charmbracelet/crush/internal/skills" + "github.com/stretchr/testify/require" +) + +// TestMessageToProtoToolResult ensures that ToolResult metadata, +// data, and MIME type survive the conversion to proto. 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 TestMessageToProtoToolResult(t *testing.T) { + t.Parallel() + + src := message.Message{ + ID: "m1", + Role: message.Tool, + Parts: []message.ContentPart{ + message.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 := messageToProto(src) + require.Len(t, got.Parts, 1) + tr, ok := got.Parts[0].(proto.ToolResult) + require.True(t, ok, "expected proto.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) +} + +// TestSkillsEventToProto_RoundTrip verifies that a pubsub.Event[skills.Event] +// can be wrapped, marshaled, and unmarshaled back through the SSE +// envelope without losing state values or error messages. +func TestSkillsEventToProto_RoundTrip(t *testing.T) { + t.Parallel() + + src := pubsub.Event[skills.Event]{ + Type: pubsub.UpdatedEvent, + Payload: skills.Event{ + States: []*skills.SkillState{ + {Name: "ok", Path: "/p/ok", State: skills.StateNormal}, + {Name: "broken", Path: "/p/broken", State: skills.StateError, Err: errors.New("bad frontmatter")}, + }, + }, + } + + env := wrapEvent(src) + require.NotNil(t, env) + require.Equal(t, pubsub.PayloadTypeSkillsEvent, env.Type) + + var decoded pubsub.Event[proto.SkillsEvent] + require.NoError(t, json.Unmarshal(env.Payload, &decoded)) + require.Equal(t, pubsub.UpdatedEvent, decoded.Type) + require.Len(t, decoded.Payload.States, 2) + + require.Equal(t, "ok", decoded.Payload.States[0].Name) + require.Equal(t, "/p/ok", decoded.Payload.States[0].Path) + require.Equal(t, proto.SkillStateNormal, decoded.Payload.States[0].State) + require.Empty(t, decoded.Payload.States[0].Error) + + require.Equal(t, "broken", decoded.Payload.States[1].Name) + require.Equal(t, proto.SkillStateError, decoded.Payload.States[1].State) + require.Equal(t, "bad frontmatter", decoded.Payload.States[1].Error) +} + +// TestRunCompleteToProto_RoundTrip verifies that the authoritative +// per-run completion event survives the SSE envelope conversion with +// all reconciliation fields intact. SessionID, MessageID, and Text +// are what non-interactive clients (e.g. `crush run`) rely on to +// terminate the run loop and guarantee final text on stdout when +// message events arrive out of order. +func TestRunCompleteToProto_RoundTrip(t *testing.T) { + t.Parallel() + + src := pubsub.Event[notify.RunComplete]{ + Type: pubsub.UpdatedEvent, + Payload: notify.RunComplete{ + SessionID: "S", + RunID: "run-42", + MessageID: "M", + Text: "VERDICT: APPROVED", + Error: "", + Cancelled: false, + }, + } + + env := wrapEvent(src) + require.NotNil(t, env) + require.Equal(t, pubsub.PayloadTypeRunComplete, env.Type) + + var decoded pubsub.Event[proto.RunComplete] + require.NoError(t, json.Unmarshal(env.Payload, &decoded)) + require.Equal(t, pubsub.UpdatedEvent, decoded.Type) + require.Equal(t, "S", decoded.Payload.SessionID) + require.Equal(t, "run-42", decoded.Payload.RunID, + "RunID must survive the SSE envelope so clients can correlate "+ + "this event with the SendMessage call that produced it") + require.Equal(t, "M", decoded.Payload.MessageID) + require.Equal(t, "VERDICT: APPROVED", decoded.Payload.Text) + require.Empty(t, decoded.Payload.Error) + require.False(t, decoded.Payload.Cancelled) +} + +// TestAgentErrorToProto_PreservesRunID verifies that an async agent +// error notification carries its originating RunID (and SessionID) +// through the SSE envelope. Without these correlators, `crush run` +// cannot tell whether an error event belongs to its own run and +// would abort on any unrelated workspace failure. +func TestAgentErrorToProto_PreservesRunID(t *testing.T) { + t.Parallel() + + src := pubsub.Event[notify.Notification]{ + Type: pubsub.CreatedEvent, + Payload: notify.Notification{ + SessionID: "S", + RunID: "run-99", + Type: notify.TypeAgentError, + Message: "boom", + }, + } + + env := wrapEvent(src) + require.NotNil(t, env) + require.Equal(t, pubsub.PayloadTypeAgentEvent, env.Type) + + var decoded pubsub.Event[proto.AgentEvent] + require.NoError(t, json.Unmarshal(env.Payload, &decoded)) + require.Equal(t, proto.AgentEventTypeError, decoded.Payload.Type) + require.Equal(t, "S", decoded.Payload.SessionID) + require.Equal(t, "run-99", decoded.Payload.RunID, + "RunID must survive so observers can attribute the error to its run") + require.NotNil(t, decoded.Payload.Error) + require.Equal(t, "boom", decoded.Payload.Error.Error()) +} + +// TestRunCompleteToProto_Error verifies that error- and cancel-shaped +// RunComplete events round-trip cleanly so clients can distinguish +// "agent failed" (returns non-zero from `crush run`) from "agent +// cancelled by user" (clean exit). +func TestRunCompleteToProto_Error(t *testing.T) { + t.Parallel() + + src := pubsub.Event[notify.RunComplete]{ + Type: pubsub.UpdatedEvent, + Payload: notify.RunComplete{ + SessionID: "S", + MessageID: "M", + Text: "partial", + Error: "context canceled", + Cancelled: true, + }, + } + + env := wrapEvent(src) + require.NotNil(t, env) + var decoded pubsub.Event[proto.RunComplete] + require.NoError(t, json.Unmarshal(env.Payload, &decoded)) + require.Equal(t, "context canceled", decoded.Payload.Error) + require.True(t, decoded.Payload.Cancelled) +} + +// TestUpdateAvailableMsgToProto_RoundTrip verifies that an +// app.UpdateAvailableMsg — published directly (not wrapped in +// pubsub.Event) by app.checkForUpdates — survives the SSE envelope +// conversion. Without this, client/server mode silently drops update +// notifications because wrapEvent hits its default branch. +func TestUpdateAvailableMsgToProto_RoundTrip(t *testing.T) { + t.Parallel() + + src := app.UpdateAvailableMsg{ + CurrentVersion: "1.0.0", + LatestVersion: "1.1.0", + IsDevelopment: false, + } + + env := wrapEvent(src) + require.NotNil(t, env) + require.Equal(t, pubsub.PayloadTypeUpdateAvailable, env.Type) + + var decoded pubsub.Event[proto.UpdateAvailable] + require.NoError(t, json.Unmarshal(env.Payload, &decoded)) + require.Equal(t, pubsub.UpdatedEvent, decoded.Type) + require.Equal(t, "1.0.0", decoded.Payload.CurrentVersion) + require.Equal(t, "1.1.0", decoded.Payload.LatestVersion) + require.False(t, decoded.Payload.IsDevelopment) +} + +// TestMCPChannelMessageNotWrappedAsStateChange verifies that an +// EventChannelMessage — which has no proto representation until session +// delivery is wired up in a later PR — is NOT wrapped as a spurious +// state_changed MCP event by the SSE event pipeline. Before the fix, +// mcpEventTypeToProto's default branch mapped every unknown event type to +// MCPEventStateChanged, so a channel notification looked like a state +// change to every SSE client. +func TestMCPChannelMessageNotWrappedAsStateChange(t *testing.T) { + t.Parallel() + + src := pubsub.Event[mcp.Event]{ + Type: pubsub.CreatedEvent, + Payload: mcp.Event{ + Type: mcp.EventChannelMessage, + Name: "webhook", + ChannelMessage: `build failed`, + }, + } + + env := wrapEvent(src) + require.Nil(t, env, "EventChannelMessage must not be wrapped as an SSE event (no proto representation yet)") +} + +// TestMCPUnknownEventTypeNotMappedToStateChange verifies that any +// unrecognized MCP event type is not silently coerced to state_changed — +// the mapping must return ok=false so wrapEvent can drop it instead of +// fabricating a state change. +func TestMCPUnknownEventTypeNotMappedToStateChange(t *testing.T) { + t.Parallel() + + // Use a value well outside the known range. + unknown := mcp.EventType(99) + pt := mcpEventTypeToProto(unknown) + require.Equal(t, proto.MCPEventType(""), pt, + "unknown MCP event types must map to empty proto type, not state_changed") +} + +// TestMessageToProtoPrismModel ensures the Prism-routed model fields survive +// the conversion to proto. Without them the client TUI cannot show which +// model actually served each turn on Hyper's model router. +func TestMessageToProtoPrismModel(t *testing.T) { + t.Parallel() + + src := message.Message{ + ID: "m1", + Role: message.Assistant, + Model: "prism-model", + Provider: "hyper", + PrismModelID: "prism-42", + PrismModelName: "GLM 5.3", + + PrismHypercreditSavings: ptrFloat(1.5), + PrismDollarSavings: ptrFloat(0.002), + } + + got := messageToProto(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 ptrFloat(v float64) *float64 { return &v } diff --git a/internal/server/multiclient_test.go b/internal/server/multiclient_test.go new file mode 100644 index 0000000000000000000000000000000000000000..8b584d2dfb953b6f6c8b7caa97b7655578d29419 --- /dev/null +++ b/internal/server/multiclient_test.go @@ -0,0 +1,233 @@ +package server + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/charmbracelet/crush/internal/backend" + "github.com/charmbracelet/crush/internal/proto" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +// installSyntheticWorkspace creates a synthetic [backend.Workspace] +// registered with the controller's backend, suitable for handler-level +// tests that do not need a real [app.App]. The workspace's ID is a +// fresh UUID and its path is a tempdir; teardown is the caller's +// responsibility (handlers should not rely on synthetic workspaces +// disappearing automatically). +func installSyntheticWorkspace(t *testing.T, c *controllerV1) *backend.Workspace { + t.Helper() + ws := &backend.Workspace{ + ID: uuid.New().String(), + Path: t.TempDir(), + } + backend.InsertWorkspaceForTest(c.backend, ws) + return ws +} + +// newTestController builds a controllerV1 around a backend without a +// real config store, suitable for handler-level 400 tests. +func newTestController() *controllerV1 { + s := &Server{} + s.backend = backend.New(context.Background(), nil, nil) + return &controllerV1{backend: s.backend, server: s} +} + +func TestPostWorkspaces_RejectsMissingClientID(t *testing.T) { + t.Parallel() + c := newTestController() + + body, err := json.Marshal(proto.Workspace{Path: t.TempDir()}) + require.NoError(t, err) + req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/v1/workspaces", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + + c.handlePostWorkspaces(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code) + var perr proto.Error + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &perr)) + require.Contains(t, perr.Message, "client_id") +} + +func TestPostWorkspaces_RejectsMalformedClientID(t *testing.T) { + t.Parallel() + c := newTestController() + + body, err := json.Marshal(proto.Workspace{Path: t.TempDir(), ClientID: "not-a-uuid"}) + require.NoError(t, err) + req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/v1/workspaces", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + + c.handlePostWorkspaces(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestDeleteWorkspace_RejectsMissingClientID(t *testing.T) { + t.Parallel() + c := newTestController() + + req := httptest.NewRequestWithContext(t.Context(), http.MethodDelete, "/v1/workspaces/abc", nil) + req.SetPathValue("id", "abc") + rec := httptest.NewRecorder() + + c.handleDeleteWorkspaces(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestDeleteWorkspace_RejectsMalformedClientID(t *testing.T) { + t.Parallel() + c := newTestController() + + req := httptest.NewRequestWithContext(t.Context(), http.MethodDelete, "/v1/workspaces/abc?client_id=nope", nil) + req.SetPathValue("id", "abc") + rec := httptest.NewRecorder() + + c.handleDeleteWorkspaces(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestSubscribeEvents_RejectsMissingClientID(t *testing.T) { + t.Parallel() + c := newTestController() + + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/v1/workspaces/abc/events", nil) + req.SetPathValue("id", "abc") + rec := httptest.NewRecorder() + + c.handleGetWorkspaceEvents(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestSubscribeEvents_RejectsMalformedClientID(t *testing.T) { + t.Parallel() + c := newTestController() + + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/v1/workspaces/abc/events?client_id=nope", nil) + req.SetPathValue("id", "abc") + rec := httptest.NewRecorder() + + c.handleGetWorkspaceEvents(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +// postCurrentSession is a small helper that POSTs the JSON body to +// /v1/workspaces/{id}/current-session?client_id=cid and returns the +// recorder. It does not require a real listener. +func postCurrentSession(t *testing.T, c *controllerV1, wsID, clientID, sessionID string) *httptest.ResponseRecorder { + t.Helper() + body, err := json.Marshal(proto.CurrentSession{SessionID: sessionID}) + require.NoError(t, err) + url := "/v1/workspaces/" + wsID + "/current-session" + if clientID != "" { + url += "?client_id=" + clientID + } + req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, url, bytes.NewReader(body)) + req.SetPathValue("id", wsID) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + c.handlePostWorkspaceCurrentSession(rec, req) + return rec +} + +func TestPostCurrentSession_RejectsMissingClientID(t *testing.T) { + t.Parallel() + c := newTestController() + + body, err := json.Marshal(proto.CurrentSession{SessionID: "S1"}) + require.NoError(t, err) + req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/v1/workspaces/abc/current-session", bytes.NewReader(body)) + req.SetPathValue("id", "abc") + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + + c.handlePostWorkspaceCurrentSession(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestPostCurrentSession_RejectsMalformedClientID(t *testing.T) { + t.Parallel() + c := newTestController() + + rec := postCurrentSession(t, c, "abc", "not-a-uuid", "S1") + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestPostCurrentSession_RejectsBadBody(t *testing.T) { + t.Parallel() + c := newTestController() + + cid := uuid.New().String() + url := "/v1/workspaces/abc/current-session?client_id=" + cid + req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, url, bytes.NewReader([]byte("not-json"))) + req.SetPathValue("id", "abc") + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + + c.handlePostWorkspaceCurrentSession(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code) +} + +func TestPostCurrentSession_UnknownWorkspace(t *testing.T) { + t.Parallel() + c := newTestController() + + rec := postCurrentSession(t, c, uuid.New().String(), uuid.New().String(), "S1") + require.Equal(t, http.StatusNotFound, rec.Code) +} + +func TestPostCurrentSession_UnknownClient(t *testing.T) { + t.Parallel() + c := newTestController() + ws := installSyntheticWorkspace(t, c) + + rec := postCurrentSession(t, c, ws.ID, uuid.New().String(), "S1") + // 409, not 404: the workspace is fine, this client just has no stream + // on it. A 404 would look like "workspace gone" and send a recovering + // client off to re-register for no reason. + require.Equal(t, http.StatusConflict, rec.Code) +} + +func TestPostCurrentSession_HoldOnly(t *testing.T) { + t.Parallel() + c := newTestController() + ws := installSyntheticWorkspace(t, c) + + cid := uuid.New().String() + require.NoError(t, backend.RegisterClientForTesting(c.backend, ws, cid)) + t.Cleanup(func() { _ = c.backend.DeleteWorkspace(ws.ID, cid) }) + + rec := postCurrentSession(t, c, ws.ID, cid, "S1") + require.Equal(t, http.StatusConflict, rec.Code, "hold-only client must be rejected") +} + +func TestPostCurrentSession_AttachedClientSucceeds(t *testing.T) { + t.Parallel() + c := newTestController() + ws := installSyntheticWorkspace(t, c) + + cid := uuid.New().String() + require.NoError(t, c.backend.AttachClient(ws.ID, cid)) + t.Cleanup(func() { c.backend.DetachClient(ws.ID, cid) }) + + rec := postCurrentSession(t, c, ws.ID, cid, "S1") + require.Equal(t, http.StatusOK, rec.Code) + + // Clearing also returns 200. + rec = postCurrentSession(t, c, ws.ID, cid, "") + require.Equal(t, http.StatusOK, rec.Code) +} diff --git a/internal/server/net_other.go b/internal/server/net_other.go new file mode 100644 index 0000000000000000000000000000000000000000..502d1436b4a091bd04a4b42dff82c7fcfafff7a8 --- /dev/null +++ b/internal/server/net_other.go @@ -0,0 +1,61 @@ +//go:build !windows + +package server + +import ( + "errors" + "io/fs" + "net" + "os" + "time" +) + +// staleSocketDialTimeout bounds the probe used to detect whether a Unix +// socket file on disk is backed by a live listener. +const staleSocketDialTimeout = 200 * time.Millisecond + +// listen binds a net.Listener on the given network and address. +// +// For unix sockets it self-heals from stale socket files: if the path +// already exists on disk, it first probes with a short net.DialTimeout. +// A successful dial means a live server owns the socket, so we proceed +// to net.Listen (which surfaces the usual "address already in use" +// error). A failed dial that isStaleSocketErr classifies as stale +// triggers an os.Remove of the path (ignoring fs.ErrNotExist) before +// the bind. +// +// The returned removedStale bool reports whether a stale socket file +// was removed prior to binding so callers can log it. The operation +// is idempotent: removing an absent file is a no-op, and a live +// socket is never removed. +func listen(network, address string) (net.Listener, bool, error) { + var removedStale bool + if network == "unix" && address != "" { + if _, err := os.Stat(address); err == nil { + conn, dialErr := net.DialTimeout(network, address, staleSocketDialTimeout) //nolint:noctx + if dialErr == nil { + // A live server owns the socket. Fall through to + // net.Listen so the caller sees the standard + // "address already in use" error. + conn.Close() + } else if isStaleSocketErr(dialErr) { + rmErr := os.Remove(address) + switch { + case rmErr == nil: + removedStale = true + case errors.Is(rmErr, fs.ErrNotExist): + // Another process removed it between our + // stat and remove; treat as a no-op. + default: + return nil, false, rmErr + } + } + } + } + //nolint:noctx + ln, err := net.Listen(network, address) + if err != nil { + return nil, removedStale, err + } + return ln, removedStale, nil +} diff --git a/internal/server/net_windows.go b/internal/server/net_windows.go new file mode 100644 index 0000000000000000000000000000000000000000..c5ae0ec106e7c1aa37e908be9fa24908bc25852c --- /dev/null +++ b/internal/server/net_windows.go @@ -0,0 +1,37 @@ +//go:build windows +// +build windows + +package server + +import ( + "net" + + "github.com/Microsoft/go-winio" +) + +// listen binds a net.Listener on the given network and address. +// +// On Windows there is no Unix-socket stale-file recovery to perform, +// so removedStale is always false. The signature matches the +// non-Windows implementation so callers can use a single code path. +func listen(network, address string) (net.Listener, bool, error) { + switch network { + case "npipe": + cfg := &winio.PipeConfig{ + MessageMode: true, + InputBufferSize: 65536, + OutputBufferSize: 65536, + } + ln, err := winio.ListenPipe(address, cfg) + if err != nil { + return nil, false, err + } + return ln, false, nil + default: + ln, err := net.Listen(network, address) //nolint:noctx + if err != nil { + return nil, false, err + } + return ln, false, nil + } +} diff --git a/internal/server/proto.go b/internal/server/proto.go new file mode 100644 index 0000000000000000000000000000000000000000..88a9e44a0fa71808120b0399242fa7414c6388ed --- /dev/null +++ b/internal/server/proto.go @@ -0,0 +1,781 @@ +package server + +import ( + "encoding/json" + "errors" + "fmt" + "net/http" + + "github.com/charmbracelet/crush/internal/backend" + "github.com/charmbracelet/crush/internal/proto" + "github.com/charmbracelet/crush/internal/session" + "github.com/google/uuid" +) + +type controllerV1 struct { + backend *backend.Backend + server *Server +} + +// handleGetHealth checks server health. +func (c *controllerV1) handleGetHealth(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) +} + +// handleGetVersion returns server version information. +func (c *controllerV1) handleGetVersion(w http.ResponseWriter, _ *http.Request) { + jsonEncode(w, c.backend.VersionInfo()) +} + +// handlePostControl sends a control command to the server. +func (c *controllerV1) handlePostControl(w http.ResponseWriter, r *http.Request) { + var req proto.ServerControl + 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 + } + + switch req.Command { + case proto.ServerControlShutdown, proto.ServerControlShutdownIfIdle: + // Both spellings are conditional. Only the backend can rule on + // idleness without racing a session that arrives between a + // client's own check and its request, and guarding the plain + // command too means clients predating the check cannot take live + // sessions down either. + if !c.backend.ShutdownIfIdle() { + c.handleError(w, r, backend.ErrServerNotIdle) + return + } + default: + c.server.logError(r, "Unknown command", "command", req.Command) + jsonError(w, http.StatusBadRequest, "unknown command") + return + } +} + +// handleGetConfig returns global server configuration. +func (c *controllerV1) handleGetConfig(w http.ResponseWriter, _ *http.Request) { + jsonEncode(w, c.backend.Config()) +} + +// handleGetWorkspaces lists all workspaces. +func (c *controllerV1) handleGetWorkspaces(w http.ResponseWriter, _ *http.Request) { + jsonEncode(w, c.backend.ListWorkspaces()) +} + +// handleGetWorkspace returns a single workspace by ID. +func (c *controllerV1) handleGetWorkspace(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + ws, err := c.backend.GetWorkspaceProto(id) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, ws) +} + +// handlePostWorkspaces creates a new workspace. +func (c *controllerV1) handlePostWorkspaces(w http.ResponseWriter, r *http.Request) { + var args proto.Workspace + if err := json.NewDecoder(r.Body).Decode(&args); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + _, result, err := c.backend.CreateWorkspace(args) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, result) +} + +// requireClientID reads the client_id query parameter and validates it +// as a UUID. On failure it writes a 400 and returns false. +func (c *controllerV1) requireClientID(w http.ResponseWriter, r *http.Request) (string, bool) { + cid := r.URL.Query().Get("client_id") + if cid == "" { + c.server.logError(r, "Missing client_id query parameter") + jsonError(w, http.StatusBadRequest, "client_id is required") + return "", false + } + if _, err := uuid.Parse(cid); err != nil { + c.server.logError(r, "Invalid client_id", "error", err) + jsonError(w, http.StatusBadRequest, "client_id is not a valid UUID") + return "", false + } + return cid, true +} + +// handlePostWorkspaceCurrentSession records the calling client's +// current session selection for the workspace. An empty session_id +// clears the entry (e.g. the client is on the landing screen). +func (c *controllerV1) handlePostWorkspaceCurrentSession(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + clientID, ok := c.requireClientID(w, r) + if !ok { + return + } + var req proto.CurrentSession + 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.SetCurrentSession(id, clientID, req.SessionID); err != nil { + c.handleError(w, r, err) + return + } +} + +// handleDeleteClient retires a client, releasing every claim it holds. +func (c *controllerV1) handleDeleteClient(w http.ResponseWriter, r *http.Request) { + if err := c.backend.RetireClient(r.PathValue("client_id")); err != nil { + c.handleError(w, r, err) + return + } +} + +// handleDeleteWorkspaces deletes a workspace. +func (c *controllerV1) handleDeleteWorkspaces(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + clientID, ok := c.requireClientID(w, r) + if !ok { + return + } + if err := c.backend.DeleteWorkspace(id, clientID); err != nil { + c.handleError(w, r, err) + return + } +} + +// handleGetWorkspaceConfig returns workspace configuration. +func (c *controllerV1) handleGetWorkspaceConfig(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + cfg, err := c.backend.GetWorkspaceConfig(id) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, cfg) +} + +// handleGetWorkspaceProviders lists available providers for a workspace. +func (c *controllerV1) handleGetWorkspaceProviders(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + providers, err := c.backend.GetWorkspaceProviders(id) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, providers) +} + +// handleGetWorkspaceEvents streams workspace events as Server-Sent Events. +func (c *controllerV1) handleGetWorkspaceEvents(w http.ResponseWriter, r *http.Request) { + flusher := http.NewResponseController(w) + id := r.PathValue("id") + clientID, ok := c.requireClientID(w, r) + if !ok { + return + } + // Subscribe to the event broker BEFORE attaching the client. + // AttachClient bumps the stream count that observers use to + // detect a live subscriber; subscribing first guarantees that + // once a client appears attached, any published event is + // delivered rather than dropped on a not-yet-registered stream. + events, err := c.backend.SubscribeEvents(r.Context(), id) + if err != nil { + c.handleError(w, r, err) + return + } + if err := c.backend.AttachClient(id, clientID); err != nil { + c.handleError(w, r, err) + return + } + defer c.backend.DetachClient(id, clientID) + + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + // Flush headers immediately so clients see the 200 response + // before any events arrive. Without this, a quiet workspace + // keeps the client's SubscribeEvents call blocked on the + // initial RoundTrip. + w.WriteHeader(http.StatusOK) + flusher.Flush() + + for { + select { + case <-r.Context().Done(): + c.server.logDebug(r, "Stopping event stream") + return + case ev, ok := <-events: + if !ok { + return + } + wrapped := wrapEvent(ev.Payload) + if wrapped == nil { + continue + } + data, err := json.Marshal(wrapped) + if err != nil { + c.server.logError(r, "Failed to marshal event", "error", err) + continue + } + + fmt.Fprintf(w, "data: %s\n\n", data) + flusher.Flush() + } + } +} + +// handleGetWorkspaceLSPs lists LSP clients for a workspace. +func (c *controllerV1) handleGetWorkspaceLSPs(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + states, err := c.backend.GetLSPStates(id) + if err != nil { + c.handleError(w, r, err) + return + } + result := make(map[string]proto.LSPClientInfo, len(states)) + for k, v := range states { + result[k] = proto.LSPClientInfo{ + Name: v.Name, + State: v.State, + Error: v.Error, + DiagnosticCount: v.DiagnosticCount, + ConnectedAt: v.ConnectedAt, + } + } + jsonEncode(w, result) +} + +// handleGetWorkspaceLSPDiagnostics returns diagnostics for an LSP client. +func (c *controllerV1) handleGetWorkspaceLSPDiagnostics(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + lspName := r.PathValue("lsp") + diagnostics, err := c.backend.GetLSPDiagnostics(id, lspName) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, diagnostics) +} + +// handleGetWorkspaceSessions lists sessions for a workspace. +func (c *controllerV1) handleGetWorkspaceSessions(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sessions, err := c.backend.ListSessions(r.Context(), id) + if err != nil { + c.handleError(w, r, err) + return + } + ws, _ := c.backend.GetWorkspace(id) + result := make([]proto.Session, len(sessions)) + for i, s := range sessions { + result[i] = sessionToProto(s) + result[i].IsBusy = isSessionBusy(ws, s.ID) + result[i].AttachedClients = attachedClients(ws, s.ID) + } + jsonEncode(w, result) +} + +// handlePostWorkspaceSessions creates a new session in a workspace. +func (c *controllerV1) handlePostWorkspaceSessions(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var args session.Session + if err := json.NewDecoder(r.Body).Decode(&args); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + sess, err := c.backend.CreateSession(r.Context(), id, args.Title) + if err != nil { + c.handleError(w, r, err) + return + } + ws, _ := c.backend.GetWorkspace(id) + out := sessionToProto(sess) + out.IsBusy = isSessionBusy(ws, sess.ID) + out.AttachedClients = attachedClients(ws, sess.ID) + jsonEncode(w, out) +} + +// handleGetWorkspaceSession returns a single session. +func (c *controllerV1) handleGetWorkspaceSession(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.PathValue("sid") + sess, err := c.backend.GetSession(r.Context(), id, sid) + if err != nil { + c.handleError(w, r, err) + return + } + ws, _ := c.backend.GetWorkspace(id) + out := sessionToProto(sess) + out.IsBusy = isSessionBusy(ws, sess.ID) + out.AttachedClients = attachedClients(ws, sess.ID) + jsonEncode(w, out) +} + +// handleGetWorkspaceSessionHistory returns the history for a session. +func (c *controllerV1) handleGetWorkspaceSessionHistory(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.PathValue("sid") + history, err := c.backend.ListSessionHistory(r.Context(), id, sid) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, history) +} + +// handleGetWorkspaceSessionMessages returns all messages for a session. +func (c *controllerV1) handleGetWorkspaceSessionMessages(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.PathValue("sid") + messages, err := c.backend.ListSessionMessages(r.Context(), id, sid) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, messagesToProto(messages)) +} + +// handlePutWorkspaceSession updates a session. +func (c *controllerV1) handlePutWorkspaceSession(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var sess session.Session + if err := json.NewDecoder(r.Body).Decode(&sess); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + saved, err := c.backend.SaveSession(r.Context(), id, sess) + if err != nil { + c.handleError(w, r, err) + return + } + ws, _ := c.backend.GetWorkspace(id) + out := sessionToProto(saved) + out.IsBusy = isSessionBusy(ws, saved.ID) + out.AttachedClients = attachedClients(ws, saved.ID) + jsonEncode(w, out) +} + +// handleDeleteWorkspaceSession deletes a session. +func (c *controllerV1) handleDeleteWorkspaceSession(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.PathValue("sid") + if err := c.backend.DeleteSession(r.Context(), id, sid); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handleGetWorkspaceSessionUserMessages returns user messages for a session. +func (c *controllerV1) handleGetWorkspaceSessionUserMessages(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.PathValue("sid") + messages, err := c.backend.ListUserMessages(r.Context(), id, sid) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, messagesToProto(messages)) +} + +// handleGetWorkspaceAllUserMessages returns all user messages across sessions. +func (c *controllerV1) handleGetWorkspaceAllUserMessages(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + messages, err := c.backend.ListAllUserMessages(r.Context(), id) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, messagesToProto(messages)) +} + +// handleGetWorkspaceSessionFileTrackerFiles lists files read in a session. +func (c *controllerV1) handleGetWorkspaceSessionFileTrackerFiles(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.PathValue("sid") + files, err := c.backend.FileTrackerListReadFiles(r.Context(), id, sid) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, files) +} + +// handlePostWorkspaceFileTrackerRead records a file read event. +func (c *controllerV1) handlePostWorkspaceFileTrackerRead(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.FileTrackerReadRequest + 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.FileTrackerRecordRead(r.Context(), id, req.SessionID, req.Path); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handleGetWorkspaceFileTrackerLastRead returns the last read time for a file. +func (c *controllerV1) handleGetWorkspaceFileTrackerLastRead(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.URL.Query().Get("session_id") + path := r.URL.Query().Get("path") + + t, err := c.backend.FileTrackerLastReadTime(r.Context(), id, sid, path) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, t) +} + +// handlePostWorkspaceLSPStart starts an LSP server for a path. +func (c *controllerV1) handlePostWorkspaceLSPStart(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.LSPStartRequest + 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.LSPStart(r.Context(), id, req.Path); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handlePostWorkspaceLSPStopAll stops all LSP servers. +func (c *controllerV1) handlePostWorkspaceLSPStopAll(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + if err := c.backend.LSPStopAll(r.Context(), id); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handleGetWorkspaceAgent returns agent info for a workspace. +func (c *controllerV1) handleGetWorkspaceAgent(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + info, err := c.backend.GetAgentInfo(id) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, info) +} + +// handlePostWorkspaceAgent sends a message to the agent. +func (c *controllerV1) handlePostWorkspaceAgent(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var msg proto.AgentMessage + if err := json.NewDecoder(r.Body).Decode(&msg); err != nil { + c.server.logError(r, "Failed to decode request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + + // The run's lifetime is detached from the prompting client's HTTP + // request: SendMessage validates and accepts the prompt, dispatches + // the run on a goroutine bound to the workspace context, and returns + // immediately. A dropping its TCP connection (network blip, TUI + // restart) or B canceling the session via the explicit cancel + // endpoint can no longer tear down a turn that other subscribed + // clients are still watching. Only the explicit cancel endpoint + // should be able to end a run. + if err := c.backend.SendMessage(id, msg); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusAccepted) +} + +// handlePostWorkspaceAgentInit initializes the agent for a workspace. +func (c *controllerV1) handlePostWorkspaceAgentInit(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.AgentInitRequest + if r.Body != nil && r.ContentLength > 0 { + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + c.server.logError(r, "Failed to decode agent init request", "error", err) + jsonError(w, http.StatusBadRequest, "failed to decode request") + return + } + } + + if err := c.backend.InitAgent(r.Context(), id, req.Interactive); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handlePostWorkspaceAgentUpdate updates the agent for a workspace. +func (c *controllerV1) handlePostWorkspaceAgentUpdate(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + if err := c.backend.UpdateAgent(r.Context(), id); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handleGetWorkspaceAgentSession returns a specific agent session. +func (c *controllerV1) handleGetWorkspaceAgentSession(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.PathValue("sid") + agentSession, err := c.backend.GetAgentSession(r.Context(), id, sid) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, agentSession) +} + +// handlePostWorkspaceAgentSessionCancel cancels a running agent session. +func (c *controllerV1) handlePostWorkspaceAgentSessionCancel(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.PathValue("sid") + if err := c.backend.CancelSession(id, sid); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handleGetWorkspaceAgentSessionPromptQueued returns whether a queued prompt exists. +func (c *controllerV1) handleGetWorkspaceAgentSessionPromptQueued(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.PathValue("sid") + queued, err := c.backend.QueuedPrompts(id, sid) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, queued) +} + +// handlePostWorkspaceAgentSessionPromptClear clears the prompt queue for a session. +func (c *controllerV1) handlePostWorkspaceAgentSessionPromptClear(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.PathValue("sid") + if err := c.backend.ClearQueue(id, sid); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handlePostWorkspaceAgentSessionSummarize summarizes a session. +func (c *controllerV1) handlePostWorkspaceAgentSessionSummarize(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.PathValue("sid") + if err := c.backend.SummarizeSession(r.Context(), id, sid); err != nil { + c.handleError(w, r, err) + return + } + w.WriteHeader(http.StatusOK) +} + +// handlePostWorkspaceAgentSessionShell runs a shell command in the workspace. +func (c *controllerV1) handlePostWorkspaceAgentSessionShell(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.PathValue("sid") + + var req proto.ShellCommandRequest + 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 + } + req.SessionID = sid + + resp, err := c.backend.RunShellCommand(r.Context(), id, req) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, resp) +} + +// handleGetWorkspaceAgentSessionPromptList returns the list of queued prompts. +func (c *controllerV1) handleGetWorkspaceAgentSessionPromptList(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + sid := r.PathValue("sid") + prompts, err := c.backend.QueuedPromptsList(id, sid) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, prompts) +} + +// handleGetWorkspaceAgentDefaultSmallModel returns the default small model for a provider. +func (c *controllerV1) handleGetWorkspaceAgentDefaultSmallModel(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + providerID := r.URL.Query().Get("provider_id") + model, err := c.backend.GetDefaultSmallModel(id, providerID) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, model) +} + +// handlePostWorkspacePermissionsGrant grants a permission request. +func (c *controllerV1) handlePostWorkspacePermissionsGrant(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.PermissionGrant + 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 + } + + resolved, err := c.backend.GrantPermission(id, req) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, proto.PermissionGrantResponse{Resolved: resolved}) +} + +// handlePostWorkspaceQuestionsAnswer submits answers for a batch question. +func (c *controllerV1) handlePostWorkspaceQuestionsAnswer(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.QuestionAnswer + 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 + } + + resolved, err := c.backend.AnswerQuestion(id, req) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, proto.QuestionAnswerResponse{Resolved: resolved}) +} + +// handlePostWorkspaceQuestionsCancel cancels the pending question +// batch for a workspace. +func (c *controllerV1) handlePostWorkspaceQuestionsCancel(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + cancelled, err := c.backend.CancelQuestion(id) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, proto.QuestionAnswerResponse{Resolved: cancelled}) +} + +// handlePostWorkspacePermissionsSkip sets whether to skip permission prompts. +func (c *controllerV1) handlePostWorkspacePermissionsSkip(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + + var req proto.PermissionSkipRequest + 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.SetPermissionsSkip(id, req.Skip); err != nil { + c.handleError(w, r, err) + return + } +} + +// handleGetWorkspacePermissionsSkip returns whether permission prompts are skipped. +func (c *controllerV1) handleGetWorkspacePermissionsSkip(w http.ResponseWriter, r *http.Request) { + id := r.PathValue("id") + skip, err := c.backend.GetPermissionsSkip(id) + if err != nil { + c.handleError(w, r, err) + return + } + jsonEncode(w, proto.PermissionSkipRequest{Skip: skip}) +} + +// handleError maps backend errors to HTTP status codes and writes the +// JSON error response. +// +// Runtime cancellation of an agent run no longer reaches here for the +// agent-prompt path: SendMessage is fire-and-forget (the handler returns +// 202 before the run starts) and Backend.runAgent swallows +// context.Canceled, surfacing the FinishReasonCanceled marker to SSE +// subscribers instead. The remaining callers pass synchronous backend +// errors, so context.Canceled gets no special case and would fall through +// to the default 500 like any other unexpected error. +func (c *controllerV1) handleError(w http.ResponseWriter, r *http.Request, err error) { + status := http.StatusInternalServerError + switch { + case errors.Is(err, backend.ErrWorkspaceNotFound): + status = http.StatusNotFound + case errors.Is(err, backend.ErrLSPClientNotFound): + status = http.StatusNotFound + case errors.Is(err, backend.ErrAgentNotInitialized): + status = http.StatusBadRequest + case errors.Is(err, backend.ErrPathRequired): + status = http.StatusBadRequest + case errors.Is(err, backend.ErrInvalidPermissionAction): + status = http.StatusBadRequest + case errors.Is(err, backend.ErrUnknownCommand): + status = http.StatusBadRequest + case errors.Is(err, backend.ErrInvalidClientID): + status = http.StatusBadRequest + case errors.Is(err, backend.ErrClientNotAttached): + // 409, not 404: the workspace exists, the caller just has no live + // stream yet. A 404 here is indistinguishable from "workspace + // gone" and would trip clients that treat 404 as the trigger for + // workspace recovery. + status = http.StatusConflict + case errors.Is(err, backend.ErrWorkspaceClosing), + errors.Is(err, backend.ErrServerNotIdle), + errors.Is(err, backend.ErrClientRetired): + status = http.StatusConflict + case errors.Is(err, backend.ErrServerShuttingDown): + // 503, not 409: the request is not wrong, this process is just + // leaving. Clients retry against its replacement. + status = http.StatusServiceUnavailable + } + c.server.logError(r, err.Error()) + jsonError(w, status, err.Error()) +} + +func jsonEncode(w http.ResponseWriter, v any) { + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(v) +} + +func jsonError(w http.ResponseWriter, status int, message string) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(proto.Error{Message: message}) +} diff --git a/internal/server/recover.go b/internal/server/recover.go new file mode 100644 index 0000000000000000000000000000000000000000..547dfb0c84fa7a8026ef0c91d761b196c9703de5 --- /dev/null +++ b/internal/server/recover.go @@ -0,0 +1,59 @@ +package server + +import ( + "log/slog" + "net/http" + "runtime/debug" +) + +// recoverHandler wraps the next handler in a panic-recovery middleware. +// If a handler panics, the panic is logged with a stack trace and a 500 +// JSON error is written to the client (when no response has been started +// yet). Without this, a panicking handler closes the connection silently +// and surfaces as an opaque EOF on the client side. +func (s *Server) recoverHandler(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + rrw := &recoverResponseWriter{ResponseWriter: w} + defer func() { + rec := recover() + if rec == nil { + return + } + // http.ErrAbortHandler is the documented way to abort a + // handler without logging; preserve that contract. + if rec == http.ErrAbortHandler { + panic(rec) + } + s.logError( + r, "Panic in handler", + slog.Any("panic", rec), + slog.String("stack", string(debug.Stack())), + ) + if !rrw.wroteHeader { + jsonError(rrw, http.StatusInternalServerError, "internal server error") + } + }() + next.ServeHTTP(rrw, r) + }) +} + +// recoverResponseWriter tracks whether the response has been started so +// the recovery middleware knows if it can still write a 500 error. +type recoverResponseWriter struct { + http.ResponseWriter + wroteHeader bool +} + +func (rrw *recoverResponseWriter) WriteHeader(code int) { + rrw.wroteHeader = true + rrw.ResponseWriter.WriteHeader(code) +} + +func (rrw *recoverResponseWriter) Write(b []byte) (int, error) { + rrw.wroteHeader = true + return rrw.ResponseWriter.Write(b) +} + +func (rrw *recoverResponseWriter) Unwrap() http.ResponseWriter { + return rrw.ResponseWriter +} diff --git a/internal/server/sessions_isbusy_test.go b/internal/server/sessions_isbusy_test.go new file mode 100644 index 0000000000000000000000000000000000000000..e5ba52744e14b3f811cea45dbb88e47f6c680646 --- /dev/null +++ b/internal/server/sessions_isbusy_test.go @@ -0,0 +1,334 @@ +package server + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "testing" + + "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/charmbracelet/crush/internal/session" + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +// stubCoordinator is a minimal agent.Coordinator that only reports +// per-session busy state. Every other method returns a zero value so +// the type satisfies the interface without dragging in the full +// coordinator dependency graph. +type stubCoordinator struct { + busy map[string]bool +} + +func (s *stubCoordinator) Run(ctx context.Context, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + return nil, nil +} + +func (s *stubCoordinator) RunAccepted(ctx context.Context, accept *agent.AcceptedRun, sessionID, prompt string, attachments ...message.Attachment) (*fantasy.AgentResult, error) { + return nil, nil +} + +func (s *stubCoordinator) BeginAccepted(sessionID string) *agent.AcceptedRun { + return nil +} +func (s *stubCoordinator) Cancel(string) {} +func (s *stubCoordinator) CancelAll() {} +func (s *stubCoordinator) IsBusy() bool { return false } +func (s *stubCoordinator) IsSessionBusy(id string) bool { + return s.busy[id] +} +func (s *stubCoordinator) QueuedPrompts(string) int { return 0 } +func (s *stubCoordinator) QueuedPromptsList(string) []string { return nil } +func (s *stubCoordinator) ClearQueue(string) {} +func (s *stubCoordinator) Summarize(context.Context, string) error { + return nil +} +func (s *stubCoordinator) Model() agent.Model { return agent.Model{} } +func (s *stubCoordinator) UpdateModels(context.Context) error { return nil } +func (s *stubCoordinator) GenerateTitle(context.Context, string, string) {} + +// stubSessions is a minimal session.Service that returns a fixed list +// (and supports Get by ID). All other methods return zero values; the +// IsBusy tests do not exercise them. +type stubSessions struct { + session.Service // embed nil to inherit the unexported broker methods + all []session.Session +} + +func (s *stubSessions) List(context.Context) ([]session.Session, error) { + return s.all, nil +} + +func (s *stubSessions) Get(_ context.Context, id string) (session.Session, error) { + for _, sess := range s.all { + if sess.ID == id { + return sess, nil + } + } + return session.Session{}, errors.New("not found") +} + +// buildBusyWorkspace returns a controller wired to a backend that owns +// a single workspace whose AgentCoordinator reports the named session +// as busy. +func buildBusyWorkspace(t *testing.T, sessionID string, busy bool) (*controllerV1, string) { + t.Helper() + + b := backend.New(context.Background(), nil, nil) + wsID := uuid.New().String() + coord := &stubCoordinator{busy: map[string]bool{sessionID: busy}} + a := &app.App{AgentCoordinator: coord} + a.Sessions = &stubSessions{all: []session.Session{{ID: sessionID, Title: "t"}}} + + ws := &backend.Workspace{ + ID: wsID, + Path: t.TempDir(), + App: a, + } + backend.InsertWorkspaceForTest(b, ws) + + s := &Server{backend: b} + return &controllerV1{backend: b, server: s}, wsID +} + +func TestSessionListIncludesIsBusy(t *testing.T) { + t.Parallel() + const sid = "s-busy" + c, wsID := buildBusyWorkspace(t, sid, true) + + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/v1/workspaces/"+wsID+"/sessions", nil) + req.SetPathValue("id", wsID) + rec := httptest.NewRecorder() + c.handleGetWorkspaceSessions(rec, req) + require.Equal(t, http.StatusOK, rec.Code) + + var got []proto.Session + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + require.Len(t, got, 1) + require.Equal(t, sid, got[0].ID) + require.True(t, got[0].IsBusy, "expected IsBusy=true for the busy session") +} + +func TestSessionListIdleSessionIsNotBusy(t *testing.T) { + t.Parallel() + const sid = "s-idle" + c, wsID := buildBusyWorkspace(t, sid, false) + + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/v1/workspaces/"+wsID+"/sessions", nil) + req.SetPathValue("id", wsID) + rec := httptest.NewRecorder() + c.handleGetWorkspaceSessions(rec, req) + require.Equal(t, http.StatusOK, rec.Code) + + var got []proto.Session + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + require.Len(t, got, 1) + require.False(t, got[0].IsBusy, "expected IsBusy=false for idle session") +} + +func TestSessionGetIncludesIsBusy(t *testing.T) { + t.Parallel() + const sid = "s-busy" + c, wsID := buildBusyWorkspace(t, sid, true) + + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/v1/workspaces/"+wsID+"/sessions/"+sid, nil) + req.SetPathValue("id", wsID) + req.SetPathValue("sid", sid) + rec := httptest.NewRecorder() + c.handleGetWorkspaceSession(rec, req) + require.Equal(t, http.StatusOK, rec.Code) + + var got proto.Session + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + require.Equal(t, sid, got.ID) + require.True(t, got.IsBusy) +} + +// TestIsSessionBusyNilSafe verifies the helper tolerates a missing +// workspace, app, or coordinator — phase A handlers rely on this so +// they can pass GetWorkspace's result through without an extra guard. +func TestIsSessionBusyNilSafe(t *testing.T) { + t.Parallel() + + require.False(t, isSessionBusy(nil, "x")) + require.False(t, isSessionBusy(&backend.Workspace{}, "x")) + require.False(t, isSessionBusy(&backend.Workspace{App: &app.App{}}, "x")) +} + +// TestProtoSessionIsBusyBackwardCompat verifies older consumers that +// unmarshal proto.Session without knowing about IsBusy still succeed +// and ignore the new field harmlessly. +func TestProtoSessionIsBusyBackwardCompat(t *testing.T) { + t.Parallel() + + wire := proto.Session{ID: "s1", Title: "t", IsBusy: true} + raw, err := json.Marshal(wire) + require.NoError(t, err) + + // Old client shape: same struct minus IsBusy. We model this by + // unmarshaling into a struct that doesn't declare the field. + type oldSession struct { + ID string `json:"id"` + Title string `json:"title"` + } + var old oldSession + require.NoError(t, json.Unmarshal(raw, &old)) + require.Equal(t, "s1", old.ID) + require.Equal(t, "t", old.Title) +} + +// buildMultiSessionWorkspace returns a controller wired to a backend +// that owns a workspace with the given session IDs. Used to exercise +// AttachedClients counts across more than one session. +func buildMultiSessionWorkspace(t *testing.T, sessionIDs ...string) (*controllerV1, *backend.Workspace) { + t.Helper() + + b := backend.New(context.Background(), nil, nil) + a := &app.App{AgentCoordinator: &stubCoordinator{}} + sessions := make([]session.Session, len(sessionIDs)) + for i, sid := range sessionIDs { + sessions[i] = session.Session{ID: sid, Title: sid} + } + a.Sessions = &stubSessions{all: sessions} + + ws := &backend.Workspace{ + ID: uuid.New().String(), + Path: t.TempDir(), + App: a, + } + backend.InsertWorkspaceForTest(b, ws) + // Synthetic workspaces have an incomplete App; bypass the + // default teardown to avoid panics when the last client detaches. + backend.SetWorkspaceShutdownFnForTest(ws, func() {}) + + s := &Server{backend: b} + return &controllerV1{backend: b, server: s}, ws +} + +// listSessions invokes handleGetWorkspaceSessions and returns the +// decoded response so tests can assert per-session counts. +func listSessions(t *testing.T, c *controllerV1, wsID string) []proto.Session { + t.Helper() + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/v1/workspaces/"+wsID+"/sessions", nil) + req.SetPathValue("id", wsID) + rec := httptest.NewRecorder() + c.handleGetWorkspaceSessions(rec, req) + require.Equal(t, http.StatusOK, rec.Code) + var got []proto.Session + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + return got +} + +func countsBySessionID(sessions []proto.Session) map[string]int { + out := make(map[string]int, len(sessions)) + for _, s := range sessions { + out[s.ID] = s.AttachedClients + } + return out +} + +// TestSessionListIncludesAttachedClients walks two sessions through +// the same lifecycle covered by TestAttachedClients_BasicLifecycle in +// the backend package, but observed at the handler boundary. +func TestSessionListIncludesAttachedClients(t *testing.T) { + t.Parallel() + c, ws := buildMultiSessionWorkspace(t, "S1", "S2") + + // No attached clients yet. + counts := countsBySessionID(listSessions(t, c, ws.ID)) + require.Equal(t, 0, counts["S1"]) + require.Equal(t, 0, counts["S2"]) + + // Attach A, set to S1: S1=1. + cidA := uuid.New().String() + require.NoError(t, c.backend.AttachClient(ws.ID, cidA)) + t.Cleanup(func() { c.backend.DetachClient(ws.ID, cidA) }) + require.NoError(t, c.backend.SetCurrentSession(ws.ID, cidA, "S1")) + counts = countsBySessionID(listSessions(t, c, ws.ID)) + require.Equal(t, 1, counts["S1"]) + require.Equal(t, 0, counts["S2"]) + + // Attach B, set to S1: S1=2. + cidB := uuid.New().String() + require.NoError(t, c.backend.AttachClient(ws.ID, cidB)) + require.NoError(t, c.backend.SetCurrentSession(ws.ID, cidB, "S1")) + counts = countsBySessionID(listSessions(t, c, ws.ID)) + require.Equal(t, 2, counts["S1"]) + require.Equal(t, 0, counts["S2"]) + + // B switches to S2: counts redistribute. + require.NoError(t, c.backend.SetCurrentSession(ws.ID, cidB, "S2")) + counts = countsBySessionID(listSessions(t, c, ws.ID)) + require.Equal(t, 1, counts["S1"]) + require.Equal(t, 1, counts["S2"]) + + // B detaches: S2 drops to 0. + c.backend.DetachClient(ws.ID, cidB) + counts = countsBySessionID(listSessions(t, c, ws.ID)) + require.Equal(t, 1, counts["S1"]) + require.Equal(t, 0, counts["S2"]) +} + +// TestSessionListExcludesHoldOnlyClient verifies that a registered +// client without an SSE stream (streams == 0) does not contribute to +// AttachedClients, even though it has an entry in the workspace's +// clients map. +func TestSessionListExcludesHoldOnlyClient(t *testing.T) { + t.Parallel() + c, ws := buildMultiSessionWorkspace(t, "S1") + + cid := uuid.New().String() + require.NoError(t, backend.RegisterClientForTesting(c.backend, ws, cid)) + t.Cleanup(func() { _ = c.backend.DeleteWorkspace(ws.ID, cid) }) + + counts := countsBySessionID(listSessions(t, c, ws.ID)) + require.Equal(t, 0, counts["S1"], "hold-only client must not be counted") +} + +// TestSessionListExcludesUnselectedAttachedClient verifies that a +// client with a live SSE stream but no current session +// (currentSessionID == "") does not show up under any session's count. +func TestSessionListExcludesUnselectedAttachedClient(t *testing.T) { + t.Parallel() + c, ws := buildMultiSessionWorkspace(t, "S1") + + cid := uuid.New().String() + require.NoError(t, c.backend.AttachClient(ws.ID, cid)) + t.Cleanup(func() { c.backend.DetachClient(ws.ID, cid) }) + // Intentionally do NOT call SetCurrentSession. + + counts := countsBySessionID(listSessions(t, c, ws.ID)) + require.Equal(t, 0, counts["S1"], + "attached client with no current session must not contribute to S1") +} + +// TestSessionGetIncludesAttachedClients verifies the single-session +// handler also populates AttachedClients. +func TestSessionGetIncludesAttachedClients(t *testing.T) { + t.Parallel() + c, ws := buildMultiSessionWorkspace(t, "S1") + + cid := uuid.New().String() + require.NoError(t, c.backend.AttachClient(ws.ID, cid)) + t.Cleanup(func() { c.backend.DetachClient(ws.ID, cid) }) + require.NoError(t, c.backend.SetCurrentSession(ws.ID, cid, "S1")) + + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, + "/v1/workspaces/"+ws.ID+"/sessions/S1", nil) + req.SetPathValue("id", ws.ID) + req.SetPathValue("sid", "S1") + rec := httptest.NewRecorder() + c.handleGetWorkspaceSession(rec, req) + require.Equal(t, http.StatusOK, rec.Code) + + var got proto.Session + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + require.Equal(t, 1, got.AttachedClients) +} diff --git a/internal/server/socket.go b/internal/server/socket.go new file mode 100644 index 0000000000000000000000000000000000000000..d691e16ddb1598a3adad0d9ec902c11500a3e00f --- /dev/null +++ b/internal/server/socket.go @@ -0,0 +1,10 @@ +//go:build !windows + +package server + +// isStaleSocketErr is the internal, non-Windows alias for the +// cross-platform IsStaleSocketErr. It is kept for the existing +// callers in net_other.go. +func isStaleSocketErr(err error) bool { + return IsStaleSocketErr(err) +} diff --git a/internal/server/socket_test.go b/internal/server/socket_test.go new file mode 100644 index 0000000000000000000000000000000000000000..38e3255866f09f86171f50d102dfd6f1cdd4ec8b --- /dev/null +++ b/internal/server/socket_test.go @@ -0,0 +1,193 @@ +//go:build !windows + +package server + +import ( + "errors" + "fmt" + "io/fs" + "net" + "os" + "path/filepath" + "strings" + "sync" + "syscall" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// fakeTimeoutErr is a minimal net.Error implementation whose Timeout() +// returns true. It is used to verify that IsStaleSocketErr never +// classifies a timeout as stale. +type fakeTimeoutErr struct{} + +func (fakeTimeoutErr) Error() string { return "fake timeout" } +func (fakeTimeoutErr) Timeout() bool { return true } +func (fakeTimeoutErr) Temporary() bool { return true } + +func TestIsStaleSocketErr(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + err error + want bool + }{ + {name: "nil", err: nil, want: false}, + {name: "ECONNREFUSED", err: syscall.ECONNREFUSED, want: true}, + { + name: "wrapped ECONNREFUSED", + err: fmt.Errorf("dial: %w", syscall.ECONNREFUSED), + want: true, + }, + {name: "fs.ErrNotExist", err: fs.ErrNotExist, want: true}, + { + name: "wrapped fs.ErrNotExist", + err: fmt.Errorf("stat: %w", fs.ErrNotExist), + want: true, + }, + {name: "timeout net.Error", err: fakeTimeoutErr{}, want: false}, + {name: "generic error", err: errors.New("boom"), want: false}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tc.want, IsStaleSocketErr(tc.err)) + }) + } +} + +func TestDefaultHost_XDGRuntimeDir(t *testing.T) { + dir := t.TempDir() + t.Setenv("XDG_RUNTIME_DIR", dir) + + host := DefaultHost() + + require.True(t, strings.HasPrefix(host, "unix://"), + "DefaultHost should return a unix:// URL, got %q", host) + path := strings.TrimPrefix(host, "unix://") + + // The composed path may exceed maxUnixSocketPathLen and fall back + // to /tmp; only assert containment when it did not. Recompose the + // path under dir (rather than checking the returned path length, + // which is short again after a /tmp fallback) to decide whether a + // fallback happened. The socket is named crush-.sock. + composed := filepath.Join(dir, filepath.Base(path)) + if len(composed) <= maxUnixSocketPathLen { + require.True(t, strings.HasPrefix(path, dir), + "socket path %q should live under %q", path, dir) + } + require.True(t, strings.HasSuffix(path, ".sock"), + "socket path %q should end in .sock", path) + require.Contains(t, filepath.Base(path), "crush", + "socket filename should contain 'crush'") +} + +func TestDefaultHost_FallbackTemp(t *testing.T) { + t.Setenv("XDG_RUNTIME_DIR", "") + + host := DefaultHost() + + require.True(t, strings.HasPrefix(host, "unix://"), + "DefaultHost should return a unix:// URL, got %q", host) + path := strings.TrimPrefix(host, "unix://") + require.NotEmpty(t, path, "fallback socket path must be non-empty") + require.True(t, strings.HasSuffix(path, ".sock"), + "socket path %q should end in .sock", path) + require.Contains(t, filepath.Base(path), "crush", + "socket filename should contain 'crush'") +} + +// staleSocketPath creates a deterministic stale unix socket file on +// disk: the socket node exists but no goroutine is accepting on it. +// It does so by binding a listener, disabling unlink-on-close, then +// closing the listener. The path is returned so the caller can probe +// it. A leftover file is best-effort removed via t.Cleanup. +func staleSocketPath(t *testing.T, path string) { + t.Helper() + ln, err := net.Listen("unix", path) //nolint:noctx + require.NoError(t, err) + ul, ok := ln.(*net.UnixListener) + require.True(t, ok, "expected *net.UnixListener, got %T", ln) + ul.SetUnlinkOnClose(false) + require.NoError(t, ul.Close()) + + // Verify it is actually stale: dialing should fail. + conn, dialErr := net.DialTimeout("unix", path, 200*time.Millisecond) //nolint:noctx + if dialErr == nil { + conn.Close() + t.Fatalf("expected stale socket at %q to refuse connections", path) + } + require.True(t, IsStaleSocketErr(dialErr), + "expected stale-socket dial error, got %v", dialErr) + + t.Cleanup(func() { + _ = os.Remove(path) + }) +} + +func TestListen_RemovesStaleSocket(t *testing.T) { + // t.TempDir() yields a path that may already be near the macOS + // sun_path limit; use a short filename to stay well under it. + dir := t.TempDir() + path := filepath.Join(dir, "s.sock") + + staleSocketPath(t, path) + + // Confirm the stale node is present before we call listen. + _, statErr := os.Stat(path) + require.NoError(t, statErr, "stale socket file should exist on disk") + + ln, removedStale, err := listen("unix", path) + require.NoError(t, err) + require.NotNil(t, ln) + require.True(t, removedStale, "listen should report removedStale=true") + t.Cleanup(func() { + _ = ln.Close() + }) +} + +func TestListen_LiveSocketNotRemoved(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "s.sock") + + ln1, err := net.Listen("unix", path) //nolint:noctx + require.NoError(t, err) + + // Drain accepts so the listener stays alive and responsive without + // blocking the test on a stray connection. + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + for { + c, err := ln1.Accept() + if err != nil { + return + } + _ = c.Close() + } + }() + t.Cleanup(func() { + _ = ln1.Close() + wg.Wait() + }) + + ln2, removedStale, err := listen("unix", path) + if ln2 != nil { + _ = ln2.Close() + } + require.Error(t, err, "listen on a live socket must fail") + require.False(t, removedStale, + "a live socket must never be removed (got removedStale=true)") + + // The live socket file must still be on disk and dialable. + _, statErr := os.Stat(path) + require.NoError(t, statErr, "live socket file should still exist") + conn, dialErr := net.DialTimeout("unix", path, 200*time.Millisecond) //nolint:noctx + require.NoError(t, dialErr, "live socket should still accept dials") + _ = conn.Close() +} diff --git a/internal/version/version.go b/internal/version/version.go new file mode 100644 index 0000000000000000000000000000000000000000..99b06243f176cf7e078a14ec3eae867738d48cbc --- /dev/null +++ b/internal/version/version.go @@ -0,0 +1,53 @@ +package version + +import ( + "os" + "runtime/debug" + "strconv" +) + +// Build-time parameters set via -ldflags. + +var ( + Version = "devel" + Commit = "unknown" + // BuildID is a unique identifier for this build. For release builds it + // equals Commit; for development builds (go run / go build without + // ldflags) it is derived from the executable's modification time, which + // changes on every recompilation. + BuildID = "" +) + +// A user may install crush using `go install github.com/charmbracelet/crush@latest`. +// without -ldflags, in which case the version above is unset. As a workaround +// we use the embedded build version that *is* set when using `go install` (and +// is only set for `go install` and not for `go build`). +func init() { + info, ok := debug.ReadBuildInfo() + if ok { + mainVersion := info.Main.Version + if mainVersion != "" && mainVersion != "(devel)" { + Version = mainVersion + } + } + + // Derive BuildID when not set via ldflags. + if BuildID == "" { + BuildID = deriveBuildID() + } +} + +// deriveBuildID uses the running executable's modification time as a unique +// build fingerprint. This changes on every recompilation (including `go run`), +// making it reliable for detecting stale servers during development. +func deriveBuildID() string { + exe, err := os.Executable() + if err != nil { + return "unknown" + } + fi, err := os.Stat(exe) + if err != nil { + return "unknown" + } + return strconv.FormatInt(fi.ModTime().UnixNano(), 36) +}