From 5e698b2ec6b4d2e3251b6d3cfd1a47aae87d3527 Mon Sep 17 00:00:00 2001 From: Kujtim Hoxha Date: Fri, 3 Oct 2025 14:14:47 +0200 Subject: [PATCH] fix(agent): cancelation logic --- 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(-) diff --git a/go.mod b/go.mod index 6db70ed4cd386fa1f4cf216a7a366cd33476067e..ee67cdfd6858fcf9ba7c757de842bccb5af8ceaa 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 9817cc105873a0084fb76dd341fb7d59f76f8d27..bb0fdcc832e226346954f353c49574037fccd373 100644 --- a/go.sum +++ b/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= diff --git a/internal/agent/agent.go b/internal/agent/agent.go index ce197465fff9696241536b16cf3d0c3ce9508852..a62254675221680a20aef4089485db803eb289de 100644 --- a/internal/agent/agent.go +++ b/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 \n\n", 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, ""); idx > 0 { diff --git a/internal/tui/page/chat/chat.go b/internal/tui/page/chat/chat.go index 19f2a3ea20694e1a11f9ff586cd51562f82ba064..6e0a4de9112397db8d487d0ed6c0c394e177a6a8 100644 --- a/internal/tui/page/chat/chat.go +++ b/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(),