6f0f519 · crush · git.iko.soy

git.iko.soy/crush

Fix shutdown (#1833)

commit6f0f5197e36c494d69ace2ee3d011696d5b656fc authorKujtim Hoxha <kujtimii.h@gmail.com> dateMon 12 Jan 2026, 15:11 (9 months ago) parentb2252f8

15 files changed+96−45

compare with main

unifiedsplit

internal/agent/agent.go+7−4 internal/agent/agentic_fetch_tool.go+4−2 internal/agent/tools/bash.go+4−1 internal/agent/tools/download.go+4−2 internal/agent/tools/edit.go+12−3 internal/agent/tools/fetch.go+4−2 internal/agent/tools/ls.go+4−2 internal/agent/tools/mcp-tools.go+4−1 internal/agent/tools/multiedit.go+9−3 internal/agent/tools/multiedit_test.go+2−2 internal/agent/tools/view.go+4−2 internal/agent/tools/write.go+4−1 internal/app/app.go+7−4 internal/permission/permission.go+13−8 internal/permission/permission_test.go+14−8

internal/agent/agent.go+7−4

@@ -445,7 +445,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*fantasy

445445 Content: content,

446446 IsError: true,

447447 }

448− _, createErr = a.messages.Create(context.Background(), currentAssistant.SessionID, message.CreateMessageParams{

448+ _, createErr = a.messages.Create(ctx, currentAssistant.SessionID, message.CreateMessageParams{

449449 Role: message.Tool,

450450 Parts: []message.ContentPart{

451451 toolResult,

@@ -876,14 +876,17 @@ func (a *sessionAgent) updateSessionUsage(model Model, session *session.Session,

876876}

877877

878878func (a *sessionAgent) Cancel(sessionID string) {

879− // Cancel regular requests.

880− if cancel, ok := a.activeRequests.Take(sessionID); ok && cancel != nil {

879+ // Cancel regular requests. Don't use Take() here - we need the entry to

880+ // remain in activeRequests so IsBusy() returns true until the goroutine

881+ // fully completes (including error handling that may access the DB).

882+ // The defer in processRequest will clean up the entry.

883+ if cancel, ok := a.activeRequests.Get(sessionID); ok && cancel != nil {

881884 slog.Info("Request cancellation initiated", "session_id", sessionID)

882885 cancel()

883886 }

884887

885888 // Also check for summarize requests.

886− if cancel, ok := a.activeRequests.Take(sessionID + "-summarize"); ok && cancel != nil {

889+ if cancel, ok := a.activeRequests.Get(sessionID + "-summarize"); ok && cancel != nil {

887890 slog.Info("Summarize cancellation initiated", "session_id", sessionID)

888891 cancel()

889892 }

internal/agent/agentic_fetch_tool.go+4−2

@@ -79,7 +79,7 @@ func (c *coordinator) agenticFetchTool(_ context.Context, client *http.Client) (

7979 description = "Search the web and analyze results"

8080 }

8181

82− p := c.permissions.Request(

82+ p, err := c.permissions.Request(ctx,

8383 permission.CreatePermissionRequest{

8484 SessionID: validationResult.SessionID,

8585 Path: c.cfg.WorkingDir(),

@@ -90,7 +90,9 @@ func (c *coordinator) agenticFetchTool(_ context.Context, client *http.Client) (

9090 Params: tools.AgenticFetchPermissionsParams(params),

9191 },

9292 )

93−

93+ if err != nil {

94+ return fantasy.ToolResponse{}, err

95+ }

9496 if !p {

9597 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

9698 }

internal/agent/tools/bash.go+4−1

@@ -215,7 +215,7 @@ func NewBashTool(permissions permission.Service, workingDir string, attribution

215215 return fantasy.ToolResponse{}, fmt.Errorf("session ID is required for executing shell command")

216216 }

217217 if !isSafeReadOnly {

218− p := permissions.Request(

218+ p, err := permissions.Request(ctx,

219219 permission.CreatePermissionRequest{

220220 SessionID: sessionID,

221221 Path: execWorkingDir,

@@ -226,6 +226,9 @@ func NewBashTool(permissions permission.Service, workingDir string, attribution

226226 Params: BashPermissionsParams(params),

227227 },

228228 )

229+ if err != nil {

230+ return fantasy.ToolResponse{}, err

231+ }

229232 if !p {

230233 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

231234 }

internal/agent/tools/download.go+4−2

@@ -70,7 +70,7 @@ func NewDownloadTool(permissions permission.Service, workingDir string, client *

7070 return fantasy.ToolResponse{}, fmt.Errorf("session ID is required for downloading files")

7171 }

7272

73− p := permissions.Request(

73+ p, err := permissions.Request(ctx,

7474 permission.CreatePermissionRequest{

7575 SessionID: sessionID,

7676 Path: filePath,

@@ -80,7 +80,9 @@ func NewDownloadTool(permissions permission.Service, workingDir string, client *

8080 Params: DownloadPermissionsParams(params),

8181 },

8282 )

83−

83+ if err != nil {

84+ return fantasy.ToolResponse{}, err

85+ }

8486 if !p {

8587 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

8688 }

internal/agent/tools/edit.go+12−3

@@ -122,7 +122,7 @@ func createNewFile(edit editContext, filePath, content string, call fantasy.Tool

122122 content,

123123 strings.TrimPrefix(filePath, edit.workingDir),

124124 )

125− p := edit.permissions.Request(

125+ p, err := edit.permissions.Request(edit.ctx,

126126 permission.CreatePermissionRequest{

127127 SessionID: sessionID,

128128 Path: fsext.PathOrPrefix(filePath, edit.workingDir),

@@ -137,6 +137,9 @@ func createNewFile(edit editContext, filePath, content string, call fantasy.Tool

137137 },

138138 },

139139 )

140+ if err != nil {

141+ return fantasy.ToolResponse{}, err

142+ }

140143 if !p {

141144 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

142145 }

@@ -243,7 +246,7 @@ func deleteContent(edit editContext, filePath, oldString string, replaceAll bool

243246 strings.TrimPrefix(filePath, edit.workingDir),

244247 )

245248

246− p := edit.permissions.Request(

249+ p, err := edit.permissions.Request(edit.ctx,

247250 permission.CreatePermissionRequest{

248251 SessionID: sessionID,

249252 Path: fsext.PathOrPrefix(filePath, edit.workingDir),

@@ -258,6 +261,9 @@ func deleteContent(edit editContext, filePath, oldString string, replaceAll bool

258261 },

259262 },

260263 )

264+ if err != nil {

265+ return fantasy.ToolResponse{}, err

266+ }

261267 if !p {

262268 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

263269 }

@@ -378,7 +384,7 @@ func replaceContent(edit editContext, filePath, oldString, newString string, rep

378384 strings.TrimPrefix(filePath, edit.workingDir),

379385 )

380386

381− p := edit.permissions.Request(

387+ p, err := edit.permissions.Request(edit.ctx,

382388 permission.CreatePermissionRequest{

383389 SessionID: sessionID,

384390 Path: fsext.PathOrPrefix(filePath, edit.workingDir),

@@ -393,6 +399,9 @@ func replaceContent(edit editContext, filePath, oldString, newString string, rep

393399 },

394400 },

395401 )

402+ if err != nil {

403+ return fantasy.ToolResponse{}, err

404+ }

396405 if !p {

397406 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

398407 }

internal/agent/tools/fetch.go+4−2

@@ -55,7 +55,7 @@ func NewFetchTool(permissions permission.Service, workingDir string, client *htt

5555 return fantasy.ToolResponse{}, fmt.Errorf("session ID is required for creating a new file")

5656 }

5757

58− p := permissions.Request(

58+ p, err := permissions.Request(ctx,

5959 permission.CreatePermissionRequest{

6060 SessionID: sessionID,

6161 Path: workingDir,

@@ -66,7 +66,9 @@ func NewFetchTool(permissions permission.Service, workingDir string, client *htt

6666 Params: FetchPermissionsParams(params),

6767 },

6868 )

69−

69+ if err != nil {

70+ return fantasy.ToolResponse{}, err

71+ }

7072 if !p {

7173 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

7274 }

internal/agent/tools/ls.go+4−2

@@ -79,7 +79,7 @@ func NewLsTool(permissions permission.Service, workingDir string, lsConfig confi

7979 return fantasy.ToolResponse{}, fmt.Errorf("session ID is required for accessing directories outside working directory")

8080 }

8181

82− granted := permissions.Request(

82+ granted, err := permissions.Request(ctx,

8383 permission.CreatePermissionRequest{

8484 SessionID: sessionID,

8585 Path: absSearchPath,

@@ -90,7 +90,9 @@ func NewLsTool(permissions permission.Service, workingDir string, lsConfig confi

9090 Params: LSPermissionsParams(params),

9191 },

9292 )

93−

93+ if err != nil {

94+ return fantasy.ToolResponse{}, err

95+ }

9496 if !granted {

9597 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

9698 }

internal/agent/tools/mcp-tools.go+4−1

@@ -89,7 +89,7 @@ func (m *Tool) Run(ctx context.Context, params fantasy.ToolCall) (fantasy.ToolRe

8989 return fantasy.ToolResponse{}, fmt.Errorf("session ID is required for creating a new file")

9090 }

9191 permissionDescription := fmt.Sprintf("execute %s with the following parameters:", m.Info().Name)

92− p := m.permissions.Request(

92+ p, err := m.permissions.Request(ctx,

9393 permission.CreatePermissionRequest{

9494 SessionID: sessionID,

9595 ToolCallID: params.ID,

@@ -100,6 +100,9 @@ func (m *Tool) Run(ctx context.Context, params fantasy.ToolCall) (fantasy.ToolRe

100100 Params: params.Input,

101101 },

102102 )

103+ if err != nil {

104+ return fantasy.ToolResponse{}, err

105+ }

103106 if !p {

104107 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

105108 }

internal/agent/tools/multiedit.go+9−3

@@ -173,7 +173,7 @@ func processMultiEditWithCreation(edit editContext, params MultiEditParams, call

173173 } else {

174174 description = fmt.Sprintf("Create file %s with %d edits", params.FilePath, editsApplied)

175175 }

176− p := edit.permissions.Request(permission.CreatePermissionRequest{

176+ p, err := edit.permissions.Request(edit.ctx, permission.CreatePermissionRequest{

177177 SessionID: sessionID,

178178 Path: fsext.PathOrPrefix(params.FilePath, edit.workingDir),

179179 ToolCallID: call.ID,

@@ -186,12 +186,15 @@ func processMultiEditWithCreation(edit editContext, params MultiEditParams, call

186186 NewContent: currentContent,

187187 },

188188 })

189+ if err != nil {

190+ return fantasy.ToolResponse{}, err

191+ }

189192 if !p {

190193 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

191194 }

192195

193196 // Write the file

194− err := os.WriteFile(params.FilePath, []byte(currentContent), 0o644)

197+ err = os.WriteFile(params.FilePath, []byte(currentContent), 0o644)

195198 if err != nil {

196199 return fantasy.ToolResponse{}, fmt.Errorf("failed to write file: %w", err)

197200 }

@@ -314,7 +317,7 @@ func processMultiEditExistingFile(edit editContext, params MultiEditParams, call

314317 } else {

315318 description = fmt.Sprintf("Apply %d edits to file %s", editsApplied, params.FilePath)

316319 }

317− p := edit.permissions.Request(permission.CreatePermissionRequest{

320+ p, err := edit.permissions.Request(edit.ctx, permission.CreatePermissionRequest{

318321 SessionID: sessionID,

319322 Path: fsext.PathOrPrefix(params.FilePath, edit.workingDir),

320323 ToolCallID: call.ID,

@@ -327,6 +330,9 @@ func processMultiEditExistingFile(edit editContext, params MultiEditParams, call

327330 NewContent: currentContent,

328331 },

329332 })

333+ if err != nil {

334+ return fantasy.ToolResponse{}, err

335+ }

330336 if !p {

331337 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

332338 }

internal/agent/tools/multiedit_test.go+2−2

@@ -19,8 +19,8 @@ type mockPermissionService struct {

1919 *pubsub.Broker[permission.PermissionRequest]

2020}

2121

22−func (m *mockPermissionService) Request(req permission.CreatePermissionRequest) bool {

23− return true

22+func (m *mockPermissionService) Request(ctx context.Context, req permission.CreatePermissionRequest) (bool, error) {

23+ return true, nil

2424}

2525

2626func (m *mockPermissionService) Grant(req permission.PermissionRequest) {}

internal/agent/tools/view.go+4−2

@@ -88,7 +88,7 @@ func NewViewTool(lspClients *csync.Map[string, *lsp.Client], permissions permiss

8888 return fantasy.ToolResponse{}, fmt.Errorf("session ID is required for accessing files outside working directory")

8989 }

9090

91− granted := permissions.Request(

91+ granted, err := permissions.Request(ctx,

9292 permission.CreatePermissionRequest{

9393 SessionID: sessionID,

9494 Path: absFilePath,

@@ -99,7 +99,9 @@ func NewViewTool(lspClients *csync.Map[string, *lsp.Client], permissions permiss

9999 Params: ViewPermissionsParams(params),

100100 },

101101 )

102−

102+ if err != nil {

103+ return fantasy.ToolResponse{}, err

104+ }

103105 if !granted {

104106 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

105107 }

internal/agent/tools/write.go+4−1

@@ -111,7 +111,7 @@ func NewWriteTool(lspClients *csync.Map[string, *lsp.Client], permissions permis

111111 strings.TrimPrefix(filePath, workingDir),

112112 )

113113

114− p := permissions.Request(

114+ p, err := permissions.Request(ctx,

115115 permission.CreatePermissionRequest{

116116 SessionID: sessionID,

117117 Path: fsext.PathOrPrefix(filePath, workingDir),

@@ -126,6 +126,9 @@ func NewWriteTool(lspClients *csync.Map[string, *lsp.Client], permissions permis

126126 },

127127 },

128128 )

129+ if err != nil {

130+ return fantasy.ToolResponse{}, err

131+ }

129132 if !p {

130133 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

131134 }

internal/app/app.go+7−4

@@ -392,13 +392,16 @@ func (app *App) Subscribe(program *tea.Program) {

392392func (app *App) Shutdown() {

393393 start := time.Now()

394394 defer func() { slog.Info("Shutdown took " + time.Since(start).String()) }()

395− var wg sync.WaitGroup

395+

396+ // First, cancel all agents and wait for them to finish. This must complete

397+ // before closing the DB so agents can finish writing their state.

396398 if app.AgentCoordinator != nil {

397− wg.Go(func() {

398− app.AgentCoordinator.CancelAll()

399− })

399+ app.AgentCoordinator.CancelAll()

400400 }

401401

402+ // Now run remaining cleanup tasks in parallel.

403+ var wg sync.WaitGroup

404+

402405 // Kill all background shells.

403406 wg.Go(func() {

404407 shell.GetBackgroundShellManager().KillAll()

internal/permission/permission.go+13−8

@@ -47,7 +47,7 @@ type Service interface {

4747 GrantPersistent(permission PermissionRequest)

4848 Grant(permission PermissionRequest)

4949 Deny(permission PermissionRequest)

50− Request(opts CreatePermissionRequest) bool

50+ Request(ctx context.Context, opts CreatePermissionRequest) (bool, error)

5151 AutoApproveSession(sessionID string)

5252 SetSkipRequests(skip bool)

5353 SkipRequests() bool

@@ -122,9 +122,9 @@ func (s *permissionService) Deny(permission PermissionRequest) {

122122 }

123123}

124124

125−func (s *permissionService) Request(opts CreatePermissionRequest) bool {

125+func (s *permissionService) Request(ctx context.Context, opts CreatePermissionRequest) (bool, error) {

126126 if s.skip {

127− return true

127+ return true, nil

128128 }

129129

130130 // tell the UI that a permission was requested

@@ -137,7 +137,7 @@ func (s *permissionService) Request(opts CreatePermissionRequest) bool {

137137 // Check if the tool/action combination is in the allowlist

138138 commandKey := opts.ToolName + ":" + opts.Action

139139 if slices.Contains(s.allowedTools, commandKey) || slices.Contains(s.allowedTools, opts.ToolName) {

140− return true

140+ return true, nil

141141 }

142142

143143 s.autoApproveSessionsMu.RLock()

@@ -145,7 +145,7 @@ func (s *permissionService) Request(opts CreatePermissionRequest) bool {

145145 s.autoApproveSessionsMu.RUnlock()

146146

147147 if autoApprove {

148− return true

148+ return true, nil

149149 }

150150

151151 fileInfo, err := os.Stat(opts.Path)

@@ -176,7 +176,7 @@ func (s *permissionService) Request(opts CreatePermissionRequest) bool {

176176 for _, p := range s.sessionPermissions {

177177 if p.ToolName == permission.ToolName && p.Action == permission.Action && p.SessionID == permission.SessionID && p.Path == permission.Path {

178178 s.sessionPermissionsMu.RUnlock()

179− return true

179+ return true, nil

180180 }

181181 }

182182 s.sessionPermissionsMu.RUnlock()

@@ -185,7 +185,7 @@ func (s *permissionService) Request(opts CreatePermissionRequest) bool {

185185 for _, p := range s.sessionPermissions {

186186 if p.ToolName == permission.ToolName && p.Action == permission.Action && p.SessionID == permission.SessionID && p.Path == permission.Path {

187187 s.sessionPermissionsMu.RUnlock()

188− return true

188+ return true, nil

189189 }

190190 }

191191 s.sessionPermissionsMu.RUnlock()

@@ -199,7 +199,12 @@ func (s *permissionService) Request(opts CreatePermissionRequest) bool {

199199 // Publish the request

200200 s.Publish(pubsub.CreatedEvent, permission)

201201

202− return <-respCh

202+ select {

203+ case <-ctx.Done():

204+ return false, ctx.Err()

205+ case granted := <-respCh:

206+ return granted, nil

207+ }

203208}

204209

205210func (s *permissionService) AutoApproveSession(sessionID string) {

internal/permission/permission_test.go+14−8

@@ -5,6 +5,7 @@ import (

55 "testing"

77 "github.com/stretchr/testify/assert"

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

89)

910

1011func TestPermissionService_AllowedCommands(t *testing.T) {

@@ -81,14 +82,16 @@ func TestPermissionService_AllowedCommands(t *testing.T) {

8182func TestPermissionService_SkipMode(t *testing.T) {

8283 service := NewPermissionService("/tmp", true, []string{})

8384

84− result := service.Request(CreatePermissionRequest{

85+ result, err := service.Request(t.Context(), CreatePermissionRequest{

8586 SessionID: "test-session",

8687 ToolName: "bash",

8788 Action: "execute",

8889 Description: "test command",

8990 Path: "/tmp",

9091 })

91−

92+ if err != nil {

93+ t.Errorf("unexpected error: %v", err)

94+ }

9295 if !result {

9396 t.Error("expected permission to be granted in skip mode")

9497 }

@@ -115,7 +118,7 @@ func TestPermissionService_SequentialProperties(t *testing.T) {

115118

116119 go func() {

117120 defer wg.Done()

118− result1 = service.Request(req1)

121+ result1, _ = service.Request(t.Context(), req1)

119122 }()

120123

121124 var permissionReq PermissionRequest

@@ -136,7 +139,8 @@ func TestPermissionService_SequentialProperties(t *testing.T) {

136139 Params: map[string]string{"file": "test.txt"},

137140 Path: "/tmp/test.txt",

138141 }

139− result2 := service.Request(req2)

142+ result2, err := service.Request(t.Context(), req2)

143+ require.NoError(t, err)

140144 assert.True(t, result2, "Second request should be auto-approved")

141145 })

142146 t.Run("Sequential requests with temporary grants", func(t *testing.T) {

@@ -156,7 +160,7 @@ func TestPermissionService_SequentialProperties(t *testing.T) {

156160 var wg sync.WaitGroup

157161

158162 wg.Go(func() {

159− result1 = service.Request(req)

163+ result1, _ = service.Request(t.Context(), req)

160164 })

161165

162166 var permissionReq PermissionRequest

@@ -170,7 +174,7 @@ func TestPermissionService_SequentialProperties(t *testing.T) {

170174 var result2 bool

171175

172176 wg.Go(func() {

173− result2 = service.Request(req)

177+ result2, _ = service.Request(t.Context(), req)

174178 })

175179

176180 event = <-events

@@ -215,7 +219,8 @@ func TestPermissionService_SequentialProperties(t *testing.T) {

215219 wg.Add(1)

216220 go func(index int, request CreatePermissionRequest) {

217221 defer wg.Done()

218− results = append(results, service.Request(request))

222+ result, _ := service.Request(t.Context(), request)

223+ results = append(results, result)

219224 }(i, req)

220225 }

221226

@@ -241,7 +246,8 @@ func TestPermissionService_SequentialProperties(t *testing.T) {

241246 assert.Equal(t, 2, grantedCount, "Should have 2 granted and 1 denied")

242247 secondReq := requests[1]

243248 secondReq.Description = "Repeat of second request"

244− result := service.Request(secondReq)

249+ result, err := service.Request(t.Context(), secondReq)

250+ require.NoError(t, err)

245251 assert.True(t, result, "Repeated request should be auto-approved due to persistent permission")

246252 })

247253}