fix(agent): cancelation logic

Kujtim Hoxha created

Change summary

go.mod                         | 42 ++++++++++-------------
go.sum                         | 14 +++----
internal/agent/agent.go        | 62 +++++++++++++++++++++++++++--------
internal/tui/page/chat/chat.go |  6 +++
4 files changed, 78 insertions(+), 46 deletions(-)

Detailed changes

go.mod 🔗

@@ -15,6 +15,7 @@ require (
 	github.com/charmbracelet/bubbletea/v2 v2.0.0-beta.4.0.20250910155747-997384b0b35e
 	github.com/charmbracelet/catwalk v0.6.4-0.20251002104711-f8c6c1e5b4a5
 	github.com/charmbracelet/fang v0.4.2
+	github.com/charmbracelet/fantasy v0.0.0-20251003110041-a581de4d1d81
 	github.com/charmbracelet/glamour/v2 v2.0.0-20250811143442-a27abb32f018
 	github.com/charmbracelet/lipgloss/v2 v2.0.0-beta.3.0.20250917201909-41ff0bf215ea
 	github.com/charmbracelet/log/v2 v2.0.0-20250226163916-c379e29ff706
@@ -50,30 +51,8 @@ require (
 	cloud.google.com/go v0.116.0 // indirect
 	cloud.google.com/go/auth v0.9.3 // indirect
 	cloud.google.com/go/compute/metadata v0.5.0 // indirect
-	github.com/charmbracelet/x/json v0.2.0 // indirect
-	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
-	github.com/google/go-cmp v0.7.0 // indirect
-	github.com/google/s2a-go v0.1.8 // indirect
-	github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
-	github.com/gorilla/websocket v1.5.3 // indirect
-	github.com/openai/openai-go/v2 v2.3.0 // indirect
-	go.opencensus.io v0.24.0 // indirect
-	golang.org/x/crypto v0.41.0 // indirect
-	google.golang.org/genai v1.26.0 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1 // indirect
-	google.golang.org/grpc v1.66.2 // indirect
-	google.golang.org/protobuf v1.36.8 // indirect
-	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
-)
-
-require (
-	github.com/anthropics/anthropic-sdk-go v1.12.0 // indirect
-	github.com/charmbracelet/fantasy v0.0.0-20251003071236-5d39f0348e5d
-	github.com/go-viper/mapstructure/v2 v2.4.0 // indirect
-)
-
-require (
 	github.com/andybalholm/cascadia v1.3.3 // indirect
+	github.com/anthropics/anthropic-sdk-go v1.12.0 // indirect
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
 	github.com/aymerick/douceur v0.2.0 // indirect
 	github.com/bahlo/generic-list-go v0.2.0 // indirect
@@ -82,6 +61,7 @@ require (
 	github.com/charmbracelet/ultraviolet v0.0.0-20250915111650-81d4262876ef
 	github.com/charmbracelet/x/cellbuf v0.0.14-0.20250811133356-e0c5dbe5ea4a // indirect
 	github.com/charmbracelet/x/exp/slice v0.0.0-20250904123553-b4e2667e5ad5
+	github.com/charmbracelet/x/json v0.2.0 // indirect
 	github.com/charmbracelet/x/powernap v0.0.0-20250919153222-1038f7e6fef4
 	github.com/charmbracelet/x/term v0.2.1
 	github.com/charmbracelet/x/termios v0.1.1 // indirect
@@ -93,7 +73,13 @@ require (
 	github.com/dustin/go-humanize v1.0.1 // indirect
 	github.com/fsnotify/fsnotify v1.9.0 // indirect
 	github.com/go-logfmt/logfmt v0.6.0 // indirect
+	github.com/go-viper/mapstructure/v2 v2.4.0 // indirect
+	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
+	github.com/google/go-cmp v0.7.0 // indirect
+	github.com/google/s2a-go v0.1.8 // indirect
+	github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
 	github.com/gorilla/css v1.0.1 // indirect
+	github.com/gorilla/websocket v1.5.3 // indirect
 	github.com/hashicorp/golang-lru/v2 v2.0.7 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/klauspost/compress v1.18.0 // indirect
@@ -112,6 +98,7 @@ require (
 	github.com/muesli/mango-pflag v0.1.0 // indirect
 	github.com/muesli/roff v0.1.0 // indirect
 	github.com/ncruces/julianday v1.0.0 // indirect
+	github.com/openai/openai-go/v2 v2.3.0 // indirect
 	github.com/pierrec/lz4/v4 v4.1.22 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/posthog/posthog-go v1.6.10
@@ -132,14 +119,21 @@ require (
 	github.com/yosida95/uritemplate/v3 v3.0.2 // indirect
 	github.com/yuin/goldmark v1.7.8 // indirect
 	github.com/yuin/goldmark-emoji v1.0.5 // indirect
+	go.opencensus.io v0.24.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
+	golang.org/x/crypto v0.41.0 // indirect
 	golang.org/x/exp v0.0.0-20250620022241-b7579e27df2b // indirect
-	golang.org/x/image v0.26.0 // indirect
+	golang.org/x/image v0.27.0 // indirect
 	golang.org/x/net v0.43.0 // indirect
 	golang.org/x/sync v0.17.0 // indirect
 	golang.org/x/sys v0.36.0 // indirect
 	golang.org/x/term v0.34.0 // indirect
 	golang.org/x/text v0.29.0
+	google.golang.org/genai v1.26.0 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1 // indirect
+	google.golang.org/grpc v1.66.2 // indirect
+	google.golang.org/protobuf v1.36.8 // indirect
+	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 	gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	mvdan.cc/sh/moreinterp v0.0.0-20250902163504-3cf4fd5717a5

go.sum 🔗

@@ -51,14 +51,12 @@ github.com/charmbracelet/colorprofile v0.3.2 h1:9J27WdztfJQVAQKX2WOlSSRB+5gaKqqI
 github.com/charmbracelet/colorprofile v0.3.2/go.mod h1:mTD5XzNeWHj8oqHb+S1bssQb7vIHbepiebQ2kPKVKbI=
 github.com/charmbracelet/fang v0.4.2 h1:nWr7Tb82/TTNNGMGG35aTZ1X68loAOQmpb0qxkKXjas=
 github.com/charmbracelet/fang v0.4.2/go.mod h1:wHJKQYO5ReYsxx+yZl+skDtrlKO/4LLEQ6EXsdHhRhg=
-github.com/charmbracelet/fantasy v0.0.0-20251002051643-c96822199d77 h1:YHuUqaojkeu00YtQeXPqM/1RNJH/jqGNaQYFwa7JQTk=
-github.com/charmbracelet/fantasy v0.0.0-20251002051643-c96822199d77/go.mod h1:RZotHpq44tKZDe6Vf0kk1iDqnUgH7Scx+K/7uJ9Qwnw=
-github.com/charmbracelet/fantasy v0.0.0-20251003054629-33bf06cef92c h1:+gCA3Sv7g1jnZ96Em7j9u61H2O/1SkBAZ1LM9yd2bM8=
-github.com/charmbracelet/fantasy v0.0.0-20251003054629-33bf06cef92c/go.mod h1:RZotHpq44tKZDe6Vf0kk1iDqnUgH7Scx+K/7uJ9Qwnw=
-github.com/charmbracelet/fantasy v0.0.0-20251003055851-3196e9fa7380 h1:UYWO3cutUTV9ZLOH4hrXFy9hg4E1pVbLZsZLA1OA3+g=
-github.com/charmbracelet/fantasy v0.0.0-20251003055851-3196e9fa7380/go.mod h1:RZotHpq44tKZDe6Vf0kk1iDqnUgH7Scx+K/7uJ9Qwnw=
 github.com/charmbracelet/fantasy v0.0.0-20251003071236-5d39f0348e5d h1:HS9qu7CgQGPnPcNoHIZXGWGFS8SaxaBXYytLdakm+jY=
 github.com/charmbracelet/fantasy v0.0.0-20251003071236-5d39f0348e5d/go.mod h1:RZotHpq44tKZDe6Vf0kk1iDqnUgH7Scx+K/7uJ9Qwnw=
+github.com/charmbracelet/fantasy v0.0.0-20251003104828-c9bb5e0cf237 h1:OgqeFnDwgYOVwLoOZpvmBFEh+j7G4f5gOzkozHUMbqo=
+github.com/charmbracelet/fantasy v0.0.0-20251003104828-c9bb5e0cf237/go.mod h1:RZotHpq44tKZDe6Vf0kk1iDqnUgH7Scx+K/7uJ9Qwnw=
+github.com/charmbracelet/fantasy v0.0.0-20251003110041-a581de4d1d81 h1:TFQ9P46252Db/B/eDLHtL5Le9Fm0d6LtYDr59mquN7w=
+github.com/charmbracelet/fantasy v0.0.0-20251003110041-a581de4d1d81/go.mod h1:RZotHpq44tKZDe6Vf0kk1iDqnUgH7Scx+K/7uJ9Qwnw=
 github.com/charmbracelet/glamour/v2 v2.0.0-20250811143442-a27abb32f018 h1:PU4Zvpagsk5sgaDxn5W4sxHuLp9QRMBZB3bFSk40A4w=
 github.com/charmbracelet/glamour/v2 v2.0.0-20250811143442-a27abb32f018/go.mod h1:Z/GLmp9fzaqX4ze3nXG7StgWez5uBM5XtlLHK8V/qSk=
 github.com/charmbracelet/lipgloss/v2 v2.0.0-beta.3.0.20250917201909-41ff0bf215ea h1:g1HfUgSMvye8mgecMD1mPscpt+pzJoDEiSA+p2QXzdQ=
@@ -328,8 +326,8 @@ golang.org/x/crypto v0.41.0/go.mod h1:pO5AFd7FA68rFak7rOAGVuygIISepHftHnr8dr6+sU
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/exp v0.0.0-20250620022241-b7579e27df2b h1:M2rDM6z3Fhozi9O7NWsxAkg/yqS/lQJ6PmkyIV3YP+o=
 golang.org/x/exp v0.0.0-20250620022241-b7579e27df2b/go.mod h1:3//PLf8L/X+8b4vuAfHzxeRUl04Adcb341+IGKfnqS8=
-golang.org/x/image v0.26.0 h1:4XjIFEZWQmCZi6Wv8BoxsDhRU3RVnLX04dToTDAEPlY=
-golang.org/x/image v0.26.0/go.mod h1:lcxbMFAovzpnJxzXS3nyL83K27tmqtKzIJpctK8YO5c=
+golang.org/x/image v0.27.0 h1:C8gA4oWU/tKkdCfYT6T2u4faJu3MeNS5O8UPWlPF61w=
+golang.org/x/image v0.27.0/go.mod h1:xbdrClrAUway1MUTEZDq9mz/UpRwYAkFFNUslZtcB+g=
 golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
 golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
 golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc=

internal/agent/agent.go 🔗

@@ -3,6 +3,7 @@ package agent
 import (
 	"context"
 	_ "embed"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"log/slog"
@@ -177,7 +178,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
 		// Before each step create the new assistant message
 		PrepareStep: func(callContext context.Context, options ai.PrepareStepFunctionOptions) (_ context.Context, prepared ai.PrepareStepResult, err error) {
 			var assistantMsg message.Message
-			assistantMsg, err = a.messages.Create(genCtx, call.SessionID, message.CreateMessageParams{
+			assistantMsg, err = a.messages.Create(callContext, call.SessionID, message.CreateMessageParams{
 				Role:     message.Assistant,
 				Parts:    []message.ContentPart{},
 				Model:    a.largeModel.ModelCfg.Model,
@@ -187,7 +188,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
 				return callContext, prepared, err
 			}
 
-			callContext = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
+			callContext = context.WithValue(callContext, tools.MessageIDContextKey, assistantMsg.ID)
 
 			currentAssistant = &assistantMsg
 
@@ -200,7 +201,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
 			queuedCalls, _ := a.messageQueue.Get(call.SessionID)
 			a.messageQueue.Del(call.SessionID)
 			for _, queued := range queuedCalls {
-				userMessage, createErr := a.createUserMessage(genCtx, queued)
+				userMessage, createErr := a.createUserMessage(callContext, queued)
 				if createErr != nil {
 					return callContext, prepared, createErr
 				}
@@ -291,12 +292,15 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
 				IsError:    isError,
 				Metadata:   result.ClientMetadata,
 			}
-			a.messages.Create(context.Background(), currentAssistant.SessionID, message.CreateMessageParams{
+			_, err := a.messages.Create(context.Background(), currentAssistant.SessionID, message.CreateMessageParams{
 				Role: message.Tool,
 				Parts: []message.ContentPart{
 					toolResult,
 				},
 			})
+			if err != nil {
+				return err
+			}
 			return a.messages.Update(genCtx, *currentAssistant)
 		},
 		OnStepFinish: func(stepResult ai.StepResult) error {
@@ -328,16 +332,29 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
 		}
 		toolCalls := currentAssistant.ToolCalls()
 		toolResults := currentAssistant.ToolResults()
+		// INFO: we use the parent context here because the genCtx has been cancelled
+		msgs, createErr := a.messages.List(ctx, currentSession.ID)
+		if createErr != nil {
+			return nil, createErr
+		}
 		for _, tc := range toolCalls {
 			if !tc.Finished {
 				tc.Finished = true
 				tc.Input = "{}"
 			}
 			currentAssistant.AddToolCall(tc)
+
 			found := false
-			for _, tr := range toolResults {
-				if tr.ToolCallID == tc.ID {
-					found = true
+			for _, msg := range msgs {
+				if msg.Role == message.Tool {
+					for _, tr := range toolResults {
+						if tr.ToolCallID == tc.ID {
+							found = true
+							break
+						}
+					}
+				}
+				if found {
 					break
 				}
 			}
@@ -348,12 +365,21 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
 				} else if isPermissionErr {
 					content = "Permission denied"
 				}
-				currentAssistant.AddToolResult(message.ToolResult{
+				toolResult := message.ToolResult{
 					ToolCallID: tc.ID,
 					Name:       tc.Name,
 					Content:    content,
 					IsError:    true,
+				}
+				_, createErr = a.messages.Create(context.Background(), currentAssistant.SessionID, message.CreateMessageParams{
+					Role: message.Tool,
+					Parts: []message.ContentPart{
+						toolResult,
+					},
 				})
+				if createErr != nil {
+					return nil, createErr
+				}
 			}
 		}
 		if isCancelErr {
@@ -363,13 +389,11 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
 		} else {
 			currentAssistant.AddFinish(message.FinishReasonError, "API Error", err.Error())
 		}
-		// INFO: we use the parent context here because the genCtx might have been cancelled
+		// INFO: we use the parent context here because the genCtx has been cancelled
 		updateErr := a.messages.Update(ctx, *currentAssistant)
 		if updateErr != nil {
 			return nil, updateErr
 		}
-	}
-	if err != nil {
 		return nil, err
 	}
 	wg.Wait()
@@ -541,22 +565,32 @@ func (a *sessionAgent) generateTitle(ctx context.Context, session *session.Sessi
 		return
 	}
 
+	var maxOutput int64 = 40
+	if a.smallModel.CatwalkCfg.CanReason {
+		maxOutput = a.smallModel.CatwalkCfg.DefaultMaxTokens
+	}
+
 	agent := ai.NewAgent(a.smallModel.Model,
-		ai.WithSystemPrompt(string(titlePrompt)),
-		ai.WithMaxOutputTokens(40),
+		ai.WithSystemPrompt(string(titlePrompt)+"\n /no_think"),
+		ai.WithMaxOutputTokens(maxOutput),
 	)
 
 	resp, err := agent.Stream(ctx, ai.AgentStreamCall{
-		Prompt: fmt.Sprintf("Generate a concise title for the following content:\n\n%s", prompt),
+		Prompt: fmt.Sprintf("Generate a concise title for the following content:\n\n%s\n <think>\n\n</think>", prompt),
 	})
 	if err != nil {
 		slog.Error("error generating title", "err", err)
 		return
 	}
 
+	data, _ := json.Marshal(resp)
+
+	slog.Info("Title Response")
+	slog.Info(string(data))
 	title := resp.Response.Content.Text()
 
 	title = strings.ReplaceAll(title, "\n", " ")
+	slog.Info(title)
 
 	// remove thinking tags if present
 	if idx := strings.Index(title, "</think>"); idx > 0 {

internal/tui/page/chat/chat.go 🔗

@@ -2,6 +2,7 @@ package chat
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"time"
 
@@ -751,6 +752,11 @@ func (p *chatPage) sendMessage(text string, attachments []message.Attachment) te
 	cmds = append(cmds, func() tea.Msg {
 		_, err := p.app.AgentCoordinator.Run(context.Background(), session.ID, text, attachments...)
 		if err != nil {
+			isCancelErr := errors.Is(err, context.Canceled)
+			isPermissionErr := errors.Is(err, permission.ErrorPermissionDenied)
+			if isCancelErr || isPermissionErr {
+				return nil
+			}
 			return util.InfoMsg{
 				Type: util.InfoTypeError,
 				Msg:  err.Error(),