6f0f519 · crush · git.iko.soy

internal/agent/agent.go+7−4

compare with main

unifiedsplit

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

445 Content: content,

446 IsError: true,

447 }

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

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

449 Role: message.Tool,

450 Parts: []message.ContentPart{

451 toolResult,

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

876}

877

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

879− // Cancel regular requests.

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.

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

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

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

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

882 cancel()

885 cancel()

883 }

886 }

884

887

885 // Also check for summarize requests.

888 // 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 {

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

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

888 cancel()

891 cancel()

889 }

892 }

internal/agent/agentic_fetch_tool.go+4−2

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

79 description = "Search the web and analyze results"

80 }

81

82− p := c.permissions.Request(

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

83 permission.CreatePermissionRequest{

84 SessionID: validationResult.SessionID,

85 Path: c.cfg.WorkingDir(),

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

90 Params: tools.AgenticFetchPermissionsParams(params),

91 },

92 )

93−

93+ if err != nil {

94+ return fantasy.ToolResponse{}, err

95+ }

94 if !p {

96 if !p {

95 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

97 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

96 }

98 }

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

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

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

216 }

217 if !isSafeReadOnly {

218− p := permissions.Request(

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

219 permission.CreatePermissionRequest{

220 SessionID: sessionID,

221 Path: execWorkingDir,

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

226 Params: BashPermissionsParams(params),

227 },

228 )

229+ if err != nil {

230+ return fantasy.ToolResponse{}, err

231+ }

229 if !p {

232 if !p {

230 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

233 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

231 }

234 }

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

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

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

71 }

72

73− p := permissions.Request(

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

74 permission.CreatePermissionRequest{

75 SessionID: sessionID,

76 Path: filePath,

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

80 Params: DownloadPermissionsParams(params),

81 },

82 )

83−

83+ if err != nil {

84+ return fantasy.ToolResponse{}, err

85+ }

84 if !p {

86 if !p {

85 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

87 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

86 }

88 }

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

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

122 content,

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

124 )

125− p := edit.permissions.Request(

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

126 permission.CreatePermissionRequest{

127 SessionID: sessionID,

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

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

137 },

138 },

139 )

140+ if err != nil {

141+ return fantasy.ToolResponse{}, err

142+ }

140 if !p {

143 if !p {

141 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

144 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

142 }

145 }

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

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

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

244 )

247 )

245

248

246− p := edit.permissions.Request(

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

247 permission.CreatePermissionRequest{

250 permission.CreatePermissionRequest{

248 SessionID: sessionID,

251 SessionID: sessionID,

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

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

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

258 },

261 },

259 },

262 },

260 )

263 )

264+ if err != nil {

265+ return fantasy.ToolResponse{}, err

266+ }

261 if !p {

267 if !p {

262 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

268 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

263 }

269 }

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

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

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

379 )

385 )

380

386

381− p := edit.permissions.Request(

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

382 permission.CreatePermissionRequest{

388 permission.CreatePermissionRequest{

383 SessionID: sessionID,

389 SessionID: sessionID,

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

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

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

393 },

399 },

394 },

400 },

395 )

401 )

402+ if err != nil {

403+ return fantasy.ToolResponse{}, err

404+ }

396 if !p {

405 if !p {

397 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

406 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

398 }

407 }

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

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

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

56 }

57

58− p := permissions.Request(

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

59 permission.CreatePermissionRequest{

60 SessionID: sessionID,

61 Path: workingDir,

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

66 Params: FetchPermissionsParams(params),

67 },

68 )

69−

69+ if err != nil {

70+ return fantasy.ToolResponse{}, err

71+ }

70 if !p {

72 if !p {

71 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

73 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

72 }

74 }

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

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

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

80 }

81

82− granted := permissions.Request(

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

83 permission.CreatePermissionRequest{

84 SessionID: sessionID,

85 Path: absSearchPath,

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

90 Params: LSPermissionsParams(params),

91 },

92 )

93−

93+ if err != nil {

94+ return fantasy.ToolResponse{}, err

95+ }

94 if !granted {

96 if !granted {

95 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

97 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

96 }

98 }

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

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

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

90 }

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

92− p := m.permissions.Request(

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

93 permission.CreatePermissionRequest{

94 SessionID: sessionID,

95 ToolCallID: params.ID,

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

100 Params: params.Input,

101 },

102 )

103+ if err != nil {

104+ return fantasy.ToolResponse{}, err

105+ }

103 if !p {

106 if !p {

104 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

107 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

105 }

108 }

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

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

173 } else {

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

175 }

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

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

177 SessionID: sessionID,

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

179 ToolCallID: call.ID,

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

186 NewContent: currentContent,

187 },

188 })

189+ if err != nil {

190+ return fantasy.ToolResponse{}, err

191+ }

189 if !p {

192 if !p {

190 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

193 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

191 }

194 }

192

195

193 // Write the file

196 // Write the file

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

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

195 if err != nil {

198 if err != nil {

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

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

197 }

200 }

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

314 } else {

317 } else {

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

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

316 }

319 }

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

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

318 SessionID: sessionID,

321 SessionID: sessionID,

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

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

320 ToolCallID: call.ID,

323 ToolCallID: call.ID,

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

327 NewContent: currentContent,

330 NewContent: currentContent,

328 },

331 },

329 })

332 })

333+ if err != nil {

334+ return fantasy.ToolResponse{}, err

335+ }

330 if !p {

336 if !p {

331 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

337 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

332 }

338 }

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

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

19 *pubsub.Broker[permission.PermissionRequest]

20}

21

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

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

23− return true

23+ return true, nil

24}

25

26func (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

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

89 }

90

91− granted := permissions.Request(

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

92 permission.CreatePermissionRequest{

93 SessionID: sessionID,

94 Path: absFilePath,

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

99 Params: ViewPermissionsParams(params),

100 },

101 )

102−

102+ if err != nil {

103+ return fantasy.ToolResponse{}, err

104+ }

103 if !granted {

105 if !granted {

104 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

106 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

105 }

107 }

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

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

111 strings.TrimPrefix(filePath, workingDir),

112 )

113

114− p := permissions.Request(

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

115 permission.CreatePermissionRequest{

116 SessionID: sessionID,

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

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

126 },

127 },

128 )

129+ if err != nil {

130+ return fantasy.ToolResponse{}, err

131+ }

129 if !p {

132 if !p {

130 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

133 return fantasy.ToolResponse{}, permission.ErrorPermissionDenied

131 }

134 }

internal/app/app.go+7−4

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

392func (app *App) Shutdown() {

393 start := time.Now()

394 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.

396 if app.AgentCoordinator != nil {

398 if app.AgentCoordinator != nil {

397− wg.Go(func() {

398− app.AgentCoordinator.CancelAll()

399+ app.AgentCoordinator.CancelAll()

399− })

400 }

401

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

403+ var wg sync.WaitGroup

404+

402 // Kill all background shells.

405 // Kill all background shells.

403 wg.Go(func() {

406 wg.Go(func() {

404 shell.GetBackgroundShellManager().KillAll()

407 shell.GetBackgroundShellManager().KillAll()

internal/permission/permission.go+13−8

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

47 GrantPersistent(permission PermissionRequest)

48 Grant(permission PermissionRequest)

49 Deny(permission PermissionRequest)

50− Request(opts CreatePermissionRequest) bool

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

51 AutoApproveSession(sessionID string)

52 SetSkipRequests(skip bool)

53 SkipRequests() bool

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

122 }

123}

124

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

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

126 if s.skip {

127− return true

127+ return true, nil

128 }

129

130 // tell the UI that a permission was requested

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

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

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

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

140− return true

140+ return true, nil

141 }

142

143 s.autoApproveSessionsMu.RLock()

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

145 s.autoApproveSessionsMu.RUnlock()

146

147 if autoApprove {

148− return true

148+ return true, nil

149 }

150

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

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

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

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

178 s.sessionPermissionsMu.RUnlock()

179− return true

179+ return true, nil

180 }

181 }

182 s.sessionPermissionsMu.RUnlock()

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

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

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

187 s.sessionPermissionsMu.RUnlock()

188− return true

188+ return true, nil

189 }

190 }

191 s.sessionPermissionsMu.RUnlock()

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

199 // Publish the request

200 s.Publish(pubsub.CreatedEvent, permission)

201

202+ select {

203+ case <-ctx.Done():

202− return <-respCh

204+ return false, ctx.Err()

205+ case granted := <-respCh:

206+ return granted, nil

207+ }

203}

208}

204

209

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

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

internal/permission/permission_test.go+14−8

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

5 "testing"

6

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

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

8)

9)

9

10

10func TestPermissionService_AllowedCommands(t *testing.T) {

11func TestPermissionService_AllowedCommands(t *testing.T) {

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

81func TestPermissionService_SkipMode(t *testing.T) {

82func TestPermissionService_SkipMode(t *testing.T) {

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

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

83

84

84− result := service.Request(CreatePermissionRequest{

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

85 SessionID: "test-session",

86 SessionID: "test-session",

86 ToolName: "bash",

87 ToolName: "bash",

87 Action: "execute",

88 Action: "execute",

88 Description: "test command",

89 Description: "test command",

89 Path: "/tmp",

90 Path: "/tmp",

90 })

91 })

91−

92+ if err != nil {

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

94+ }

92 if !result {

95 if !result {

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

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

94 }

97 }

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

115

118

116 go func() {

119 go func() {

117 defer wg.Done()

120 defer wg.Done()

118− result1 = service.Request(req1)

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

119 }()

122 }()

120

123

121 var permissionReq PermissionRequest

124 var permissionReq PermissionRequest

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

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

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

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

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

138 }

141 }

139− result2 := service.Request(req2)

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

143+ require.NoError(t, err)

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

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

141 })

145 })

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

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

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

156 var wg sync.WaitGroup

160 var wg sync.WaitGroup

157

161

158 wg.Go(func() {

162 wg.Go(func() {

159− result1 = service.Request(req)

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

160 })

164 })

161

165

162 var permissionReq PermissionRequest

166 var permissionReq PermissionRequest

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

170 var result2 bool

174 var result2 bool

171

175

172 wg.Go(func() {

176 wg.Go(func() {

173− result2 = service.Request(req)

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

174 })

178 })

175

179

176 event = <-events

180 event = <-events

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

215 wg.Add(1)

219 wg.Add(1)

216 go func(index int, request CreatePermissionRequest) {

220 go func(index int, request CreatePermissionRequest) {

217 defer wg.Done()

221 defer wg.Done()

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

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

223+ results = append(results, result)

219 }(i, req)

224 }(i, req)

220 }

225 }

221

226

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

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

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

242 secondReq := requests[1]

247 secondReq := requests[1]

243 secondReq.Description = "Repeat of second request"

248 secondReq.Description = "Repeat of second request"

244− result := service.Request(secondReq)

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

250+ require.NoError(t, err)

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

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

246 })

252 })

247}

253}