fad58e2 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

feat: be able to continue non-interactive sessions (#2401)

fad58e2Daniil Sivak7 months ago

Taskfile.yaml+5

compare with main

unifiedsplit

@@ -184,3 +184,8 @@ tasks:

184 - go get charm.land/fantasy

185 - go get charm.land/catwalk

186 - go mod tidy

187+

188+ sqlc:

189+ desc: Generate code using SQLC

190+ cmds:

191+ - sqlc generate

internal/app/app.go+39−3

@@ -157,9 +157,40 @@ func (app *App) AgentNotifications() *pubsub.Broker[notify.Notification] {

157 return app.agentNotifications

158}

159

160+// resolveSession resolves which session to use for a non-interactive run

161+// If continueSessionID is set, it looks up that session by ID

162+// If useLast is set, it returns the most recently updated top-level session

163+// Otherwise, it creates a new session

164+func (app *App) resolveSession(ctx context.Context, continueSessionID string, useLast bool) (session.Session, error) {

165+ switch {

166+ case continueSessionID != "":

167+ if app.Sessions.IsAgentToolSession(continueSessionID) {

168+ return session.Session{}, fmt.Errorf("cannot continue an agent tool session: %s", continueSessionID)

169+ }

170+ sess, err := app.Sessions.Get(ctx, continueSessionID)

171+ if err != nil {

172+ return session.Session{}, fmt.Errorf("session not found: %s", continueSessionID)

173+ }

174+ if sess.ParentSessionID != "" {

175+ return session.Session{}, fmt.Errorf("cannot continue a child session: %s", continueSessionID)

176+ }

177+ return sess, nil

178+

179+ case useLast:

180+ sess, err := app.Sessions.GetLast(ctx)

181+ if err != nil {

182+ return session.Session{}, fmt.Errorf("no sessions found to continue")

183+ }

184+ return sess, nil

185+

186+ default:

187+ return app.Sessions.Create(ctx, agent.DefaultSessionName)

188+ }

189+}

190+

160// RunNonInteractive runs the application in non-interactive mode with the

191// RunNonInteractive runs the application in non-interactive mode with the

161// given prompt, printing to stdout.

192// given prompt, printing to stdout.

162−func (app *App) RunNonInteractive(ctx context.Context, output io.Writer, prompt, largeModel, smallModel string, hideSpinner bool) error {

193+func (app *App) RunNonInteractive(ctx context.Context, output io.Writer, prompt, largeModel, smallModel string, hideSpinner bool, continueSessionID string, useLast bool) error {

163 slog.Info("Running in non-interactive mode")

194 slog.Info("Running in non-interactive mode")

164

195

165 ctx, cancel := context.WithCancel(ctx)

196 ctx, cancel := context.WithCancel(ctx)

@@ -227,11 +258,16 @@ func (app *App) RunNonInteractive(ctx context.Context, output io.Writer, prompt,

227

258

228 defer stopSpinner()

259 defer stopSpinner()

229

260

230− sess, err := app.Sessions.Create(ctx, agent.DefaultSessionName)

261+ sess, err := app.resolveSession(ctx, continueSessionID, useLast)

231 if err != nil {

262 if err != nil {

232 return fmt.Errorf("failed to create session for non-interactive mode: %w", err)

263 return fmt.Errorf("failed to create session for non-interactive mode: %w", err)

233 }

264 }

265+

266+ if continueSessionID != "" || useLast {

267+ slog.Info("Continuing session for non-interactive run", "session_id", sess.ID)

268+ } else {

234− slog.Info("Created session for non-interactive run", "session_id", sess.ID)

269+ slog.Info("Created session for non-interactive run", "session_id", sess.ID)

270+ }

235

271

236 // Automatically approve all permission requests for this non-interactive

272 // Automatically approve all permission requests for this non-interactive

237 // session.

273 // session.

internal/app/resolve_session_test.gonew file+174

@@ -0,0 +1,174 @@

1+package app

3+import (

4+ "context"

5+ "database/sql"

6+ "fmt"

7+ "strings"

8+ "testing"

10+ "github.com/charmbracelet/crush/internal/pubsub"

11+ "github.com/charmbracelet/crush/internal/session"

12+ "github.com/stretchr/testify/require"

13+)

14+

15+// mockSessionService is a minimal mock of session.Service for testing resolveSession.

16+type mockSessionService struct {

17+ sessions []session.Session

18+ created []session.Session

19+}

20+

21+func (m *mockSessionService) Subscribe(context.Context) <-chan pubsub.Event[session.Session] {

22+ return make(chan pubsub.Event[session.Session])

23+}

24+

25+func (m *mockSessionService) Create(_ context.Context, title string) (session.Session, error) {

26+ s := session.Session{ID: "new-session-id", Title: title}

27+ m.created = append(m.created, s)

28+ return s, nil

29+}

30+

31+func (m *mockSessionService) CreateTitleSession(context.Context, string) (session.Session, error) {

32+ return session.Session{}, nil

33+}

34+

35+func (m *mockSessionService) CreateTaskSession(context.Context, string, string, string) (session.Session, error) {

36+ return session.Session{}, nil

37+}

38+

39+func (m *mockSessionService) Get(_ context.Context, id string) (session.Session, error) {

40+ for _, s := range m.sessions {

41+ if s.ID == id {

42+ return s, nil

43+ }

44+ }

45+ return session.Session{}, sql.ErrNoRows

46+}

47+

48+func (m *mockSessionService) GetLast(_ context.Context) (session.Session, error) {

49+ if len(m.sessions) > 0 {

50+ return m.sessions[0], nil

51+ }

52+ return session.Session{}, sql.ErrNoRows

53+}

54+

55+func (m *mockSessionService) List(context.Context) ([]session.Session, error) {

56+ return m.sessions, nil

57+}

58+

59+func (m *mockSessionService) Save(_ context.Context, s session.Session) (session.Session, error) {

60+ return s, nil

61+}

62+

63+func (m *mockSessionService) UpdateTitleAndUsage(context.Context, string, string, int64, int64, float64) error {

64+ return nil

65+}

66+

67+func (m *mockSessionService) Rename(context.Context, string, string) error {

68+ return nil

69+}

70+

71+func (m *mockSessionService) Delete(context.Context, string) error {

72+ return nil

73+}

74+

75+func (m *mockSessionService) CreateAgentToolSessionID(messageID, toolCallID string) string {

76+ return fmt.Sprintf("%s$$%s", messageID, toolCallID)

77+}

78+

79+func (m *mockSessionService) ParseAgentToolSessionID(sessionID string) (string, string, bool) {

80+ parts := strings.Split(sessionID, "$$")

81+ if len(parts) != 2 {

82+ return "", "", false

83+ }

84+ return parts[0], parts[1], true

85+}

86+

87+func (m *mockSessionService) IsAgentToolSession(sessionID string) bool {

88+ _, _, ok := m.ParseAgentToolSessionID(sessionID)

89+ return ok

90+}

91+

92+func newTestApp(sessions session.Service) *App {

93+ return &App{Sessions: sessions}

94+}

95+

96+func TestResolveSession_NewSession(t *testing.T) {

97+ mock := &mockSessionService{}

98+ app := newTestApp(mock)

99+

100+ sess, err := app.resolveSession(t.Context(), "", false)

101+ require.NoError(t, err)

102+ require.Equal(t, "new-session-id", sess.ID)

103+ require.Len(t, mock.created, 1)

104+}

105+

106+func TestResolveSession_ContinueByID(t *testing.T) {

107+ mock := &mockSessionService{

108+ sessions: []session.Session{

109+ {ID: "existing-id", Title: "Old session"},

110+ },

111+ }

112+ app := newTestApp(mock)

113+

114+ sess, err := app.resolveSession(t.Context(), "existing-id", false)

115+ require.NoError(t, err)

116+ require.Equal(t, "existing-id", sess.ID)

117+ require.Equal(t, "Old session", sess.Title)

118+ require.Empty(t, mock.created)

119+}

120+

121+func TestResolveSession_ContinueByID_NotFound(t *testing.T) {

122+ mock := &mockSessionService{}

123+ app := newTestApp(mock)

124+

125+ _, err := app.resolveSession(t.Context(), "nonexistent", false)

126+ require.Error(t, err)

127+ require.Contains(t, err.Error(), "session not found")

128+}

129+

130+func TestResolveSession_ContinueByID_ChildSession(t *testing.T) {

131+ mock := &mockSessionService{

132+ sessions: []session.Session{

133+ {ID: "child-id", ParentSessionID: "parent-id", Title: "Child session"},

134+ },

135+ }

136+ app := newTestApp(mock)

137+

138+ _, err := app.resolveSession(t.Context(), "child-id", false)

139+ require.Error(t, err)

140+ require.Contains(t, err.Error(), "cannot continue a child session")

141+}

142+

143+func TestResolveSession_ContinueByID_AgentToolSession(t *testing.T) {

144+ mock := &mockSessionService{}

145+ app := newTestApp(mock)

146+

147+ _, err := app.resolveSession(t.Context(), "msg123$$tool456", false)

148+ require.Error(t, err)

149+ require.Contains(t, err.Error(), "cannot continue an agent tool session")

150+}

151+

152+func TestResolveSession_Last(t *testing.T) {

153+ mock := &mockSessionService{

154+ sessions: []session.Session{

155+ {ID: "most-recent", Title: "Latest session"},

156+ {ID: "older", Title: "Older session"},

157+ },

158+ }

159+ app := newTestApp(mock)

160+

161+ sess, err := app.resolveSession(t.Context(), "", true)

162+ require.NoError(t, err)

163+ require.Equal(t, "most-recent", sess.ID)

164+ require.Empty(t, mock.created)

165+}

166+

167+func TestResolveSession_Last_NoSessions(t *testing.T) {

168+ mock := &mockSessionService{}

169+ app := newTestApp(mock)

170+

171+ _, err := app.resolveSession(t.Context(), "", true)

172+ require.Error(t, err)

173+ require.Contains(t, err.Error(), "no sessions found")

174+}

internal/cmd/run.go+34−5

@@ -36,12 +36,23 @@ crush run --quiet "Generate a README for this project"

37# Run in verbose mode (show logs)

38crush run --verbose "Generate a README for this project"

39+

40+# Continue a previous session

41+crush run --session {session-id} "Follow up on your last response"

42+

43+# Continue the most recent session

44+crush run --continue "Follow up on your last response"

45+

39 `,

46 `,

40 RunE: func(cmd *cobra.Command, args []string) error {

47 RunE: func(cmd *cobra.Command, args []string) error {

48+ var (

41− quiet, _ := cmd.Flags().GetBool("quiet")

49+ quiet, _ = cmd.Flags().GetBool("quiet")

42− verbose, _ := cmd.Flags().GetBool("verbose")

50+ verbose, _ = cmd.Flags().GetBool("verbose")

43− largeModel, _ := cmd.Flags().GetString("model")

51+ largeModel, _ = cmd.Flags().GetString("model")

44− smallModel, _ := cmd.Flags().GetString("small-model")

52+ smallModel, _ = cmd.Flags().GetString("small-model")

53+ sessionID, _ = cmd.Flags().GetString("session")

54+ useLast, _ = cmd.Flags().GetBool("continue")

55+ )

46 // Cancel on SIGINT or SIGTERM.

57 // Cancel on SIGINT or SIGTERM.

47 ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)

58 ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, os.Kill)

@@ -53,6 +64,14 @@ crush run --verbose "Generate a README for this project"

53 }

64 }

54 defer app.Shutdown()

65 defer app.Shutdown()

67+ if sessionID != "" {

68+ sess, err := resolveSessionID(ctx, app.Sessions, sessionID)

69+ if err != nil {

70+ return err

71+ }

72+ sessionID = sess.ID

73+ }

74+

56 if !app.Config().IsConfigured() {

75 if !app.Config().IsConfigured() {

57 return fmt.Errorf("no providers configured - please run 'crush' to set up a provider interactively")

76 return fmt.Errorf("no providers configured - please run 'crush' to set up a provider interactively")

58 }

77 }

@@ -76,7 +95,14 @@ crush run --verbose "Generate a README for this project"

76 event.SetNonInteractive(true)

95 event.SetNonInteractive(true)

77 event.AppInitialized()

96 event.AppInitialized()

98+ switch {

99+ case sessionID != "":

100+ event.SetContinueBySessionID(true)

101+ case useLast:

102+ event.SetContinueLastSession(true)

103+ }

104+

79− return app.RunNonInteractive(ctx, os.Stdout, prompt, largeModel, smallModel, quiet || verbose)

105+ return app.RunNonInteractive(ctx, os.Stdout, prompt, largeModel, smallModel, quiet || verbose, sessionID, useLast)

80 },

106 },

81}

107}

108

@@ -85,4 +111,7 @@ func init() {

85 runCmd.Flags().BoolP("verbose", "v", false, "Show logs")

111 runCmd.Flags().BoolP("verbose", "v", false, "Show logs")

86 runCmd.Flags().StringP("model", "m", "", "Model to use. Accepts 'model' or 'provider/model' to disambiguate models with the same name across providers")

112 runCmd.Flags().StringP("model", "m", "", "Model to use. Accepts 'model' or 'provider/model' to disambiguate models with the same name across providers")

87 runCmd.Flags().String("small-model", "", "Small model to use. If not provided, uses the default small model for the provider")

113 runCmd.Flags().String("small-model", "", "Small model to use. If not provided, uses the default small model for the provider")

114+ runCmd.Flags().StringP("session", "s", "", "Continue a previous session by ID")

115+ runCmd.Flags().BoolP("continue", "C", false, "Continue the most recent session")

116+ runCmd.MarkFlagsMutuallyExclusive("session", "continue")

88}

117}

internal/db/db.go+10

@@ -63,6 +63,9 @@ func Prepare(ctx context.Context, db DBTX) (*Queries, error) {

63 if q.getHourDayHeatmapStmt, err = db.PrepareContext(ctx, getHourDayHeatmap); err != nil {

64 return nil, fmt.Errorf("error preparing query GetHourDayHeatmap: %w", err)

65 }

66+ if q.getLastSessionStmt, err = db.PrepareContext(ctx, getLastSession); err != nil {

67+ return nil, fmt.Errorf("error preparing query GetLastSession: %w", err)

68+ }

66 if q.getMessageStmt, err = db.PrepareContext(ctx, getMessage); err != nil {

69 if q.getMessageStmt, err = db.PrepareContext(ctx, getMessage); err != nil {

67 return nil, fmt.Errorf("error preparing query GetMessage: %w", err)

70 return nil, fmt.Errorf("error preparing query GetMessage: %w", err)

68 }

71 }

@@ -202,6 +205,11 @@ func (q *Queries) Close() error {

202 err = fmt.Errorf("error closing getHourDayHeatmapStmt: %w", cerr)

205 err = fmt.Errorf("error closing getHourDayHeatmapStmt: %w", cerr)

203 }

206 }

204 }

207 }

208+ if q.getLastSessionStmt != nil {

209+ if cerr := q.getLastSessionStmt.Close(); cerr != nil {

210+ err = fmt.Errorf("error closing getLastSessionStmt: %w", cerr)

211+ }

212+ }

205 if q.getMessageStmt != nil {

213 if q.getMessageStmt != nil {

206 if cerr := q.getMessageStmt.Close(); cerr != nil {

214 if cerr := q.getMessageStmt.Close(); cerr != nil {

207 err = fmt.Errorf("error closing getMessageStmt: %w", cerr)

215 err = fmt.Errorf("error closing getMessageStmt: %w", cerr)

@@ -369,6 +377,7 @@ type Queries struct {

369 getFileByPathAndSessionStmt *sql.Stmt

377 getFileByPathAndSessionStmt *sql.Stmt

370 getFileReadStmt *sql.Stmt

378 getFileReadStmt *sql.Stmt

371 getHourDayHeatmapStmt *sql.Stmt

379 getHourDayHeatmapStmt *sql.Stmt

380+ getLastSessionStmt *sql.Stmt

372 getMessageStmt *sql.Stmt

381 getMessageStmt *sql.Stmt

373 getRecentActivityStmt *sql.Stmt

382 getRecentActivityStmt *sql.Stmt

374 getSessionByIDStmt *sql.Stmt

383 getSessionByIDStmt *sql.Stmt

@@ -411,6 +420,7 @@ func (q *Queries) WithTx(tx *sql.Tx) *Queries {

411 getFileByPathAndSessionStmt: q.getFileByPathAndSessionStmt,

420 getFileByPathAndSessionStmt: q.getFileByPathAndSessionStmt,

412 getFileReadStmt: q.getFileReadStmt,

421 getFileReadStmt: q.getFileReadStmt,

413 getHourDayHeatmapStmt: q.getHourDayHeatmapStmt,

422 getHourDayHeatmapStmt: q.getHourDayHeatmapStmt,

423+ getLastSessionStmt: q.getLastSessionStmt,

414 getMessageStmt: q.getMessageStmt,

424 getMessageStmt: q.getMessageStmt,

415 getRecentActivityStmt: q.getRecentActivityStmt,

425 getRecentActivityStmt: q.getRecentActivityStmt,

416 getSessionByIDStmt: q.getSessionByIDStmt,

426 getSessionByIDStmt: q.getSessionByIDStmt,

internal/db/querier.go+1

@@ -22,6 +22,7 @@ type Querier interface {

22 GetFileByPathAndSession(ctx context.Context, arg GetFileByPathAndSessionParams) (File, error)

23 GetFileRead(ctx context.Context, arg GetFileReadParams) (ReadFile, error)

24 GetHourDayHeatmap(ctx context.Context) ([]GetHourDayHeatmapRow, error)

25+ GetLastSession(ctx context.Context) (Session, error)

25 GetMessage(ctx context.Context, id string) (Message, error)

26 GetMessage(ctx context.Context, id string) (Message, error)

26 GetRecentActivity(ctx context.Context) ([]GetRecentActivityRow, error)

27 GetRecentActivity(ctx context.Context) ([]GetRecentActivityRow, error)

27 GetSessionByID(ctx context.Context, id string) (Session, error)

28 GetSessionByID(ctx context.Context, id string) (Session, error)

internal/db/sessions.sql.go+26

@@ -83,6 +83,32 @@ func (q *Queries) DeleteSession(ctx context.Context, id string) error {

83 return err

84}

86+const getLastSession = `-- name: GetLastSession :one

87+SELECT id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at, summary_message_id, todos

88+FROM sessions

89+ORDER BY updated_at DESC

90+LIMIT 1

91+`

92+

93+func (q *Queries) GetLastSession(ctx context.Context) (Session, error) {

94+ row := q.queryRow(ctx, q.getLastSessionStmt, getLastSession)

95+ var i Session

96+ err := row.Scan(

97+ &i.ID,

98+ &i.ParentSessionID,

99+ &i.Title,

100+ &i.MessageCount,

101+ &i.PromptTokens,

102+ &i.CompletionTokens,

103+ &i.Cost,

104+ &i.UpdatedAt,

105+ &i.CreatedAt,

106+ &i.SummaryMessageID,

107+ &i.Todos,

108+ )

109+ return i, err

110+}

111+

86const getSessionByID = `-- name: GetSessionByID :one

112const getSessionByID = `-- name: GetSessionByID :one

87SELECT id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at, summary_message_id, todos

113SELECT id, parent_session_id, title, message_count, prompt_tokens, completion_tokens, cost, updated_at, created_at, summary_message_id, todos

88FROM sessions

114FROM sessions

internal/db/sql/sessions.sql+6

@@ -28,6 +28,12 @@ SELECT *

28FROM sessions

29WHERE id = ? LIMIT 1;

31+-- name: GetLastSession :one

32+SELECT *

33+FROM sessions

34+ORDER BY updated_at DESC

35+LIMIT 1;

36+

31-- name: ListSessions :many

37-- name: ListSessions :many

32SELECT *

38SELECT *

33FROM sessions

39FROM sessions

internal/event/event.go+13−3

@@ -17,7 +17,9 @@ const (

17 endpoint = "https://data.charm.land"

18 key = "phc_4zt4VgDWLqbYnJYEwLRxFoaTL2noNrQij0C6E8k3I0V"

20− nonInteractiveEventName = "NonInteractive"

20+ nonInteractiveAttrName = "NonInteractive"

21+ continueSessionByIDAttrName = "ContinueSessionByID"

22+ continueLastSessionAttrName = "ContinueLastSession"

21)

23)

23var (

25var (

@@ -30,11 +32,19 @@ var (

30 Set("SHELL", filepath.Base(os.Getenv("SHELL"))).

32 Set("SHELL", filepath.Base(os.Getenv("SHELL"))).

31 Set("Version", version.Version).

33 Set("Version", version.Version).

32 Set("GoVersion", runtime.Version()).

34 Set("GoVersion", runtime.Version()).

33− Set(nonInteractiveEventName, false)

35+ Set(nonInteractiveAttrName, false)

34)

36)

36func SetNonInteractive(nonInteractive bool) {

38func SetNonInteractive(nonInteractive bool) {

37− baseProps = baseProps.Set(nonInteractiveEventName, nonInteractive)

39+ baseProps = baseProps.Set(nonInteractiveAttrName, nonInteractive)

40+}

41+

42+func SetContinueBySessionID(continueBySessionID bool) {

43+ baseProps = baseProps.Set(continueSessionByIDAttrName, continueBySessionID)

44+}

45+

46+func SetContinueLastSession(continueLastSession bool) {

47+ baseProps = baseProps.Set(continueLastSessionAttrName, continueLastSession)

38}

48}

40func Init() {

50func Init() {

internal/session/session.go+9

@@ -66,6 +66,7 @@ type Service interface {

66 CreateTitleSession(ctx context.Context, parentSessionID string) (Session, error)

67 CreateTaskSession(ctx context.Context, toolCallID, parentSessionID, title string) (Session, error)

68 Get(ctx context.Context, id string) (Session, error)

69+ GetLast(ctx context.Context) (Session, error)

69 List(ctx context.Context) ([]Session, error)

70 List(ctx context.Context) ([]Session, error)

70 Save(ctx context.Context, session Session) (Session, error)

71 Save(ctx context.Context, session Session) (Session, error)

71 UpdateTitleAndUsage(ctx context.Context, sessionID, title string, promptTokens, completionTokens int64, cost float64) error

72 UpdateTitleAndUsage(ctx context.Context, sessionID, title string, promptTokens, completionTokens int64, cost float64) error

@@ -166,6 +167,14 @@ func (s *service) Get(ctx context.Context, id string) (Session, error) {

166 return s.fromDBItem(dbSession), nil

167 return s.fromDBItem(dbSession), nil

167}

168}

168

169

170+func (s *service) GetLast(ctx context.Context) (Session, error) {

171+ dbSession, err := s.q.GetLastSession(ctx)

172+ if err != nil {

173+ return Session{}, err

174+ }

175+ return s.fromDBItem(dbSession), nil

176+}

177+

169func (s *service) Save(ctx context.Context, session Session) (Session, error) {

178func (s *service) Save(ctx context.Context, session Session) (Session, error) {

170 todosJSON, err := marshalTodos(session.Todos)

179 todosJSON, err := marshalTodos(session.Todos)

171 if err != nil {

180 if err != nil {