git.iko.soy/crush
on main help sign in
Merge pull request #175 from charmbracelet/handle-permission-cancel
8122fecKujtim Hoxha
internal/llm/agent/agent.go+44−8
compare with main
unifiedsplit
@@ -501,12 +501,45 @@ func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msg
501 }
501 }
502 continue
502 continue
503 }
503 }
504+
505+ // Run tool in goroutine to allow cancellation
506+ type toolExecResult struct {
507+ response tools.ToolResponse
508+ err error
509+ }
510+ resultChan := make(chan toolExecResult, 1)
511+
512+ go func() {
504− toolResult, toolErr := tool.Run(ctx, tools.ToolCall{
513+ response, err := tool.Run(ctx, tools.ToolCall{
505− ID: toolCall.ID,
514+ ID: toolCall.ID,
506− Name: toolCall.Name,
515+ Name: toolCall.Name,
507− Input: toolCall.Input,
516+ Input: toolCall.Input,
517+ })
518+ resultChan <- toolExecResult{response: response, err: err}
508− })
519+ }()
520+
521+ var toolResponse tools.ToolResponse
522+ var toolErr error
523+
524+ select {
525+ case <-ctx.Done():
526+ a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")
527+ // Mark remaining tool calls as cancelled
528+ for j := i; j < len(toolCalls); j++ {
529+ toolResults[j] = message.ToolResult{
530+ ToolCallID: toolCalls[j].ID,
531+ Content: "Tool execution canceled by user",
532+ IsError: true,
533+ }
534+ }
535+ goto out
536+ case result := <-resultChan:
537+ toolResponse = result.response
538+ toolErr = result.err
539+ }
540+
509 if toolErr != nil {
541 if toolErr != nil {
542+ slog.Error("Tool execution error", "toolCall", toolCall.ID, "error", toolErr)
510 if errors.Is(toolErr, permission.ErrorPermissionDenied) {
543 if errors.Is(toolErr, permission.ErrorPermissionDenied) {
511 toolResults[i] = message.ToolResult{
544 toolResults[i] = message.ToolResult{
512 ToolCallID: toolCall.ID,
545 ToolCallID: toolCall.ID,
@@ -526,9 +559,9 @@ func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msg
526 }
559 }
527 toolResults[i] = message.ToolResult{
560 toolResults[i] = message.ToolResult{
528 ToolCallID: toolCall.ID,
561 ToolCallID: toolCall.ID,
529− Content: toolResult.Content,
562+ Content: toolResponse.Content,
530− Metadata: toolResult.Metadata,
563+ Metadata: toolResponse.Metadata,
531− IsError: toolResult.IsError,
564+ IsError: toolResponse.IsError,
532 }
565 }
533 }
566 }
534 }
567 }
@@ -796,6 +829,9 @@ func (a *agent) Summarize(ctx context.Context, sessionID string) error {
796}
829}
797
830
798func (a *agent) CancelAll() {
831func (a *agent) CancelAll() {
832+ if !a.IsBusy() {
833+ return
834+ }
799 a.activeRequests.Range(func(key, value any) bool {
835 a.activeRequests.Range(func(key, value any) bool {
800 a.Cancel(key.(string)) // key is sessionID
836 a.Cancel(key.(string)) // key is sessionID
801 return true
837 return true