@@ -1,1133 +0,0 @@
1−package agent
2−
3−import (
4− "context"
5− "errors"
6− "fmt"
7− "log/slog"
8− "maps"
9− "slices"
10− "strings"
11− "time"
12−
13− "github.com/charmbracelet/catwalk/pkg/catwalk"
14− "github.com/charmbracelet/crush/internal/config"
15− "github.com/charmbracelet/crush/internal/csync"
16− "github.com/charmbracelet/crush/internal/event"
17− "github.com/charmbracelet/crush/internal/history"
18− "github.com/charmbracelet/crush/internal/llm/prompt"
19− "github.com/charmbracelet/crush/internal/llm/provider"
20− "github.com/charmbracelet/crush/internal/llm/tools"
21− "github.com/charmbracelet/crush/internal/log"
22− "github.com/charmbracelet/crush/internal/lsp"
23− "github.com/charmbracelet/crush/internal/message"
24− "github.com/charmbracelet/crush/internal/permission"
25− "github.com/charmbracelet/crush/internal/pubsub"
26− "github.com/charmbracelet/crush/internal/session"
27− "github.com/charmbracelet/crush/internal/shell"
28−)
29−
30−type AgentEventType string
31−
32−const (
33− AgentEventTypeError AgentEventType = "error"
34− AgentEventTypeResponse AgentEventType = "response"
35− AgentEventTypeSummarize AgentEventType = "summarize"
36−)
37−
38−type AgentEvent struct {
39− Type AgentEventType
40− Message message.Message
41− Error error
42−
43− // When summarizing
44− SessionID string
45− Progress string
46− Done bool
47−}
48−
49−type Service interface {
50− pubsub.Suscriber[AgentEvent]
51− Model() catwalk.Model
52− Run(ctx context.Context, sessionID string, content string, attachments ...message.Attachment) (<-chan AgentEvent, error)
53− Cancel(sessionID string)
54− CancelAll()
55− IsSessionBusy(sessionID string) bool
56− IsBusy() bool
57− Summarize(ctx context.Context, sessionID string) error
58− UpdateModel() error
59− QueuedPrompts(sessionID string) int
60− ClearQueue(sessionID string)
61−}
62−
63−type agent struct {
64− *pubsub.Broker[AgentEvent]
65− agentCfg config.Agent
66− sessions session.Service
67− messages message.Service
68− permissions permission.Service
69− baseTools *csync.Map[string, tools.BaseTool]
70− mcpTools *csync.Map[string, tools.BaseTool]
71− lspClients *csync.Map[string, *lsp.Client]
72−
73− // We need this to be able to update it when model changes
74− agentToolFn func() (tools.BaseTool, error)
75− cleanupFuncs []func()
76−
77− provider provider.Provider
78− providerID string
79−
80− titleProvider provider.Provider
81− summarizeProvider provider.Provider
82− summarizeProviderID string
83−
84− activeRequests *csync.Map[string, context.CancelFunc]
85− promptQueue *csync.Map[string, []string]
86−}
87−
88−var agentPromptMap = map[string]prompt.PromptID{
89− "coder": prompt.PromptCoder,
90− "task": prompt.PromptTask,
91−}
92−
93−func NewAgent(
94− ctx context.Context,
95− agentCfg config.Agent,
96− // These services are needed in the tools
97− permissions permission.Service,
98− sessions session.Service,
99− messages message.Service,
100− history history.Service,
101− lspClients *csync.Map[string, *lsp.Client],
102−) (Service, error) {
103− cfg := config.Get()
104−
105− var agentToolFn func() (tools.BaseTool, error)
106− if agentCfg.ID == "coder" && slices.Contains(agentCfg.AllowedTools, AgentToolName) {
107− agentToolFn = func() (tools.BaseTool, error) {
108− taskAgentCfg := config.Get().Agents["task"]
109− if taskAgentCfg.ID == "" {
110− return nil, fmt.Errorf("task agent not found in config")
111− }
112− taskAgent, err := NewAgent(ctx, taskAgentCfg, permissions, sessions, messages, history, lspClients)
113− if err != nil {
114− return nil, fmt.Errorf("failed to create task agent: %w", err)
115− }
116− return NewAgentTool(taskAgent, sessions, messages), nil
117− }
118− }
119−
120− providerCfg := config.Get().GetProviderForModel(agentCfg.Model)
121− if providerCfg == nil {
122− return nil, fmt.Errorf("provider for agent %s not found in config", agentCfg.Name)
123− }
124− model := config.Get().GetModelByType(agentCfg.Model)
125−
126− if model == nil {
127− return nil, fmt.Errorf("model not found for agent %s", agentCfg.Name)
128− }
129−
130− promptID := agentPromptMap[agentCfg.ID]
131− if promptID == "" {
132− promptID = prompt.PromptDefault
133− }
134− opts := []provider.ProviderClientOption{
135− provider.WithModel(agentCfg.Model),
136− provider.WithSystemMessage(prompt.GetPrompt(promptID, providerCfg.ID, config.Get().Options.ContextPaths...)),
137− }
138− agentProvider, err := provider.NewProvider(*providerCfg, opts...)
139− if err != nil {
140− return nil, err
141− }
142−
143− smallModelCfg := cfg.Models[config.SelectedModelTypeSmall]
144− var smallModelProviderCfg *config.ProviderConfig
145− if smallModelCfg.Provider == providerCfg.ID {
146− smallModelProviderCfg = providerCfg
147− } else {
148− smallModelProviderCfg = cfg.GetProviderForModel(config.SelectedModelTypeSmall)
149−
150− if smallModelProviderCfg.ID == "" {
151− return nil, fmt.Errorf("provider %s not found in config", smallModelCfg.Provider)
152− }
153− }
154− smallModel := cfg.GetModelByType(config.SelectedModelTypeSmall)
155− if smallModel.ID == "" {
156− return nil, fmt.Errorf("model %s not found in provider %s", smallModelCfg.Model, smallModelProviderCfg.ID)
157− }
158−
159− titleOpts := []provider.ProviderClientOption{
160− provider.WithModel(config.SelectedModelTypeSmall),
161− provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptTitle, smallModelProviderCfg.ID)),
162− }
163− titleProvider, err := provider.NewProvider(*smallModelProviderCfg, titleOpts...)
164− if err != nil {
165− return nil, err
166− }
167−
168− summarizeOpts := []provider.ProviderClientOption{
169− provider.WithModel(config.SelectedModelTypeLarge),
170− provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, providerCfg.ID)),
171− }
172− summarizeProvider, err := provider.NewProvider(*providerCfg, summarizeOpts...)
173− if err != nil {
174− return nil, err
175− }
176−
177− baseToolsFn := func() map[string]tools.BaseTool {
178− slog.Info("Initializing agent base tools", "agent", agentCfg.ID)
179− defer func() {
180− slog.Info("Initialized agent base tools", "agent", agentCfg.ID)
181− }()
182−
183− // Base tools available to all agents
184− cwd := cfg.WorkingDir()
185− result := make(map[string]tools.BaseTool)
186− for _, tool := range []tools.BaseTool{
187− tools.NewBashTool(permissions, cwd, cfg.Options.Attribution),
188− tools.NewDownloadTool(permissions, cwd),
189− tools.NewEditTool(lspClients, permissions, history, cwd),
190− tools.NewMultiEditTool(lspClients, permissions, history, cwd),
191− tools.NewFetchTool(permissions, cwd),
192− tools.NewGlobTool(cwd),
193− tools.NewGrepTool(cwd),
194− tools.NewLsTool(permissions, cwd),
195− tools.NewSourcegraphTool(),
196− tools.NewViewTool(lspClients, permissions, cwd),
197− tools.NewWriteTool(lspClients, permissions, history, cwd),
198− } {
199− result[tool.Name()] = tool
200− }
201− return result
202− }
203− mcpToolsFn := func() map[string]tools.BaseTool {
204− slog.Info("Initializing agent mcp tools", "agent", agentCfg.ID)
205− defer func() {
206− slog.Info("Initialized agent mcp tools", "agent", agentCfg.ID)
207− }()
208−
209− mcpToolsOnce.Do(func() {
210− doGetMCPTools(ctx, permissions, cfg)
211− })
212−
213− return maps.Collect(mcpTools.Seq2())
214− }
215−
216− a := &agent{
217− Broker: pubsub.NewBroker[AgentEvent](),
218− agentCfg: agentCfg,
219− provider: agentProvider,
220− providerID: string(providerCfg.ID),
221− messages: messages,
222− sessions: sessions,
223− titleProvider: titleProvider,
224− summarizeProvider: summarizeProvider,
225− summarizeProviderID: string(providerCfg.ID),
226− agentToolFn: agentToolFn,
227− activeRequests: csync.NewMap[string, context.CancelFunc](),
228− mcpTools: csync.NewLazyMap(mcpToolsFn),
229− baseTools: csync.NewLazyMap(baseToolsFn),
230− promptQueue: csync.NewMap[string, []string](),
231− permissions: permissions,
232− lspClients: lspClients,
233− }
234− a.setupEvents(ctx)
235− return a, nil
236−}
237−
238−func (a *agent) Model() catwalk.Model {
239− return *config.Get().GetModelByType(a.agentCfg.Model)
240−}
241−
242−func (a *agent) Cancel(sessionID string) {
243− // Cancel regular requests
244− if cancel, ok := a.activeRequests.Take(sessionID); ok && cancel != nil {
245− slog.Info("Request cancellation initiated", "session_id", sessionID)
246− cancel()
247− }
248−
249− // Also check for summarize requests
250− if cancel, ok := a.activeRequests.Take(sessionID + "-summarize"); ok && cancel != nil {
251− slog.Info("Summarize cancellation initiated", "session_id", sessionID)
252− cancel()
253− }
254−
255− if a.QueuedPrompts(sessionID) > 0 {
256− slog.Info("Clearing queued prompts", "session_id", sessionID)
257− a.promptQueue.Del(sessionID)
258− }
259−}
260−
261−func (a *agent) IsBusy() bool {
262− var busy bool
263− for cancelFunc := range a.activeRequests.Seq() {
264− if cancelFunc != nil {
265− busy = true
266− break
267− }
268− }
269− return busy
270−}
271−
272−func (a *agent) IsSessionBusy(sessionID string) bool {
273− _, busy := a.activeRequests.Get(sessionID)
274− return busy
275−}
276−
277−func (a *agent) QueuedPrompts(sessionID string) int {
278− l, ok := a.promptQueue.Get(sessionID)
279− if !ok {
280− return 0
281− }
282− return len(l)
283−}
284−
285−func (a *agent) generateTitle(ctx context.Context, sessionID string, content string) error {
286− if content == "" {
287− return nil
288− }
289− if a.titleProvider == nil {
290− return nil
291− }
292− session, err := a.sessions.Get(ctx, sessionID)
293− if err != nil {
294− return err
295− }
296− parts := []message.ContentPart{message.TextContent{
297− Text: fmt.Sprintf("Generate a concise title for the following content:\n\n%s", content),
298− }}
299−
300− // Use streaming approach like summarization
301− response := a.titleProvider.StreamResponse(
302− ctx,
303− []message.Message{
304− {
305− Role: message.User,
306− Parts: parts,
307− },
308− },
309− nil,
310− )
311−
312− var finalResponse *provider.ProviderResponse
313− for r := range response {
314− if r.Error != nil {
315− return r.Error
316− }
317− finalResponse = r.Response
318− }
319−
320− if finalResponse == nil {
321− return fmt.Errorf("no response received from title provider")
322− }
323−
324− title := strings.ReplaceAll(finalResponse.Content, "\n", " ")
325−
326− if idx := strings.Index(title, "</think>"); idx > 0 {
327− title = title[idx+len("</think>"):]
328− }
329−
330− title = strings.TrimSpace(title)
331− if title == "" {
332− return nil
333− }
334−
335− session.Title = title
336− _, err = a.sessions.Save(ctx, session)
337− return err
338−}
339−
340−func (a *agent) err(err error) AgentEvent {
341− return AgentEvent{
342− Type: AgentEventTypeError,
343− Error: err,
344− }
345−}
346−
347−func (a *agent) Run(ctx context.Context, sessionID string, content string, attachments ...message.Attachment) (<-chan AgentEvent, error) {
348− if !a.Model().SupportsImages && attachments != nil {
349− attachments = nil
350− }
351− events := make(chan AgentEvent, 1)
352− if a.IsSessionBusy(sessionID) {
353− existing, ok := a.promptQueue.Get(sessionID)
354− if !ok {
355− existing = []string{}
356− }
357− existing = append(existing, content)
358− a.promptQueue.Set(sessionID, existing)
359− return nil, nil
360− }
361−
362− genCtx, cancel := context.WithCancel(ctx)
363− a.activeRequests.Set(sessionID, cancel)
364− startTime := time.Now()
365−
366− go func() {
367− slog.Debug("Request started", "sessionID", sessionID)
368− defer log.RecoverPanic("agent.Run", func() {
369− events <- a.err(fmt.Errorf("panic while running the agent"))
370− })
371− var attachmentParts []message.ContentPart
372− for _, attachment := range attachments {
373− attachmentParts = append(attachmentParts, message.BinaryContent{Path: attachment.FilePath, MIMEType: attachment.MimeType, Data: attachment.Content})
374− }
375− result := a.processGeneration(genCtx, sessionID, content, attachmentParts)
376− if result.Error != nil {
377− if isCancelledErr(result.Error) {
378− slog.Error("Request canceled", "sessionID", sessionID)
379− } else {
380− slog.Error("Request errored", "sessionID", sessionID, "error", result.Error.Error())
381− event.Error(result.Error)
382− }
383− } else {
384− slog.Debug("Request completed", "sessionID", sessionID)
385− }
386− a.eventPromptResponded(sessionID, time.Since(startTime).Truncate(time.Second))
387− a.activeRequests.Del(sessionID)
388− cancel()
389− a.Publish(pubsub.CreatedEvent, result)
390− events <- result
391− close(events)
392− }()
393− a.eventPromptSent(sessionID)
394− return events, nil
395−}
396−
397−func (a *agent) processGeneration(ctx context.Context, sessionID, content string, attachmentParts []message.ContentPart) AgentEvent {
398− cfg := config.Get()
399− // List existing messages; if none, start title generation asynchronously.
400− msgs, err := a.messages.List(ctx, sessionID)
401− if err != nil {
402− return a.err(fmt.Errorf("failed to list messages: %w", err))
403− }
404− if len(msgs) == 0 {
405− go func() {
406− defer log.RecoverPanic("agent.Run", func() {
407− slog.Error("panic while generating title")
408− })
409− titleErr := a.generateTitle(ctx, sessionID, content)
410− if titleErr != nil && !errors.Is(titleErr, context.Canceled) && !errors.Is(titleErr, context.DeadlineExceeded) {
411− slog.Error("failed to generate title", "error", titleErr)
412− }
413− }()
414− }
415− session, err := a.sessions.Get(ctx, sessionID)
416− if err != nil {
417− return a.err(fmt.Errorf("failed to get session: %w", err))
418− }
419− if session.SummaryMessageID != "" {
420− summaryMsgInex := -1
421− for i, msg := range msgs {
422− if msg.ID == session.SummaryMessageID {
423− summaryMsgInex = i
424− break
425− }
426− }
427− if summaryMsgInex != -1 {
428− msgs = msgs[summaryMsgInex:]
429− msgs[0].Role = message.User
430− }
431− }
432−
433− userMsg, err := a.createUserMessage(ctx, sessionID, content, attachmentParts)
434− if err != nil {
435− return a.err(fmt.Errorf("failed to create user message: %w", err))
436− }
437− // Append the new user message to the conversation history.
438− msgHistory := append(msgs, userMsg)
439−
440− for {
441− // Check for cancellation before each iteration
442− select {
443− case <-ctx.Done():
444− return a.err(ctx.Err())
445− default:
446− // Continue processing
447− }
448− agentMessage, toolResults, err := a.streamAndHandleEvents(ctx, sessionID, msgHistory)
449− if err != nil {
450− if errors.Is(err, context.Canceled) {
451− agentMessage.AddFinish(message.FinishReasonCanceled, "Request cancelled", "")
452− a.messages.Update(context.Background(), agentMessage)
453− return a.err(ErrRequestCancelled)
454− }
455− return a.err(fmt.Errorf("failed to process events: %w", err))
456− }
457− if cfg.Options.Debug {
458− slog.Info("Result", "message", agentMessage.FinishReason(), "toolResults", toolResults)
459− }
460− if (agentMessage.FinishReason() == message.FinishReasonToolUse) && toolResults != nil {
461− // We are not done, we need to respond with the tool response
462− msgHistory = append(msgHistory, agentMessage, *toolResults)
463− // If there are queued prompts, process the next one
464− nextPrompt, ok := a.promptQueue.Take(sessionID)
465− if ok {
466− for _, prompt := range nextPrompt {
467− // Create a new user message for the queued prompt
468− userMsg, err := a.createUserMessage(ctx, sessionID, prompt, nil)
469− if err != nil {
470− return a.err(fmt.Errorf("failed to create user message for queued prompt: %w", err))
471− }
472− // Append the new user message to the conversation history
473− msgHistory = append(msgHistory, userMsg)
474− }
475− }
476−
477− continue
478− } else if agentMessage.FinishReason() == message.FinishReasonEndTurn {
479− queuePrompts, ok := a.promptQueue.Take(sessionID)
480− if ok {
481− for _, prompt := range queuePrompts {
482− if prompt == "" {
483− continue
484− }
485− userMsg, err := a.createUserMessage(ctx, sessionID, prompt, nil)
486− if err != nil {
487− return a.err(fmt.Errorf("failed to create user message for queued prompt: %w", err))
488− }
489− msgHistory = append(msgHistory, userMsg)
490− }
491− continue
492− }
493− }
494− if agentMessage.FinishReason() == "" {
495− // Kujtim: could not track down where this is happening but this means its cancelled
496− agentMessage.AddFinish(message.FinishReasonCanceled, "Request cancelled", "")
497− _ = a.messages.Update(context.Background(), agentMessage)
498− return a.err(ErrRequestCancelled)
499− }
500− return AgentEvent{
501− Type: AgentEventTypeResponse,
502− Message: agentMessage,
503− Done: true,
504− }
505− }
506−}
507−
508−func (a *agent) createUserMessage(ctx context.Context, sessionID, content string, attachmentParts []message.ContentPart) (message.Message, error) {
509− parts := []message.ContentPart{message.TextContent{Text: content}}
510− parts = append(parts, attachmentParts...)
511− return a.messages.Create(ctx, sessionID, message.CreateMessageParams{
512− Role: message.User,
513− Parts: parts,
514− })
515−}
516−
517−func (a *agent) getAllTools() ([]tools.BaseTool, error) {
518− var allTools []tools.BaseTool
519− for tool := range a.baseTools.Seq() {
520− if a.agentCfg.AllowedTools == nil || slices.Contains(a.agentCfg.AllowedTools, tool.Name()) {
521− allTools = append(allTools, tool)
522− }
523− }
524− if a.agentCfg.ID == "coder" {
525− allTools = slices.AppendSeq(allTools, a.mcpTools.Seq())
526− if a.lspClients.Len() > 0 {
527− allTools = append(allTools, tools.NewDiagnosticsTool(a.lspClients))
528− }
529− }
530− if a.agentToolFn != nil {
531− agentTool, agentToolErr := a.agentToolFn()
532− if agentToolErr != nil {
533− return nil, agentToolErr
534− }
535− allTools = append(allTools, agentTool)
536− }
537− return allTools, nil
538−}
539−
540−func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msgHistory []message.Message) (message.Message, *message.Message, error) {
541− ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID)
542−
543− // Create the assistant message first so the spinner shows immediately
544− assistantMsg, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{
545− Role: message.Assistant,
546− Parts: []message.ContentPart{},
547− Model: a.Model().ID,
548− Provider: a.providerID,
549− })
550− if err != nil {
551− return assistantMsg, nil, fmt.Errorf("failed to create assistant message: %w", err)
552− }
553−
554− allTools, toolsErr := a.getAllTools()
555− if toolsErr != nil {
556− return assistantMsg, nil, toolsErr
557− }
558− // Now collect tools (which may block on MCP initialization)
559− eventChan := a.provider.StreamResponse(ctx, msgHistory, allTools)
560−
561− // Add the session and message ID into the context if needed by tools.
562− ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
563−
564−loop:
565− for {
566− select {
567− case event, ok := <-eventChan:
568− if !ok {
569− break loop
570− }
571− if processErr := a.processEvent(ctx, sessionID, &assistantMsg, event); processErr != nil {
572− if errors.Is(processErr, context.Canceled) {
573− a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")
574− } else {
575− a.finishMessage(ctx, &assistantMsg, message.FinishReasonError, "API Error", processErr.Error())
576− }
577− return assistantMsg, nil, processErr
578− }
579− case <-ctx.Done():
580− a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")
581− return assistantMsg, nil, ctx.Err()
582− }
583− }
584−
585− toolResults := make([]message.ToolResult, len(assistantMsg.ToolCalls()))
586− toolCalls := assistantMsg.ToolCalls()
587− for i, toolCall := range toolCalls {
588− select {
589− case <-ctx.Done():
590− a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")
591− // Make all future tool calls cancelled
592− for j := i; j < len(toolCalls); j++ {
593− toolResults[j] = message.ToolResult{
594− ToolCallID: toolCalls[j].ID,
595− Content: "Tool execution canceled by user",
596− IsError: true,
597− }
598− }
599− goto out
600− default:
601− // Continue processing
602− var tool tools.BaseTool
603− allTools, _ = a.getAllTools()
604− for _, availableTool := range allTools {
605− if availableTool.Info().Name == toolCall.Name {
606− tool = availableTool
607− break
608− }
609− }
610−
611− // Tool not found
612− if tool == nil {
613− toolResults[i] = message.ToolResult{
614− ToolCallID: toolCall.ID,
615− Content: fmt.Sprintf("Tool not found: %s", toolCall.Name),
616− IsError: true,
617− }
618− continue
619− }
620−
621− // Run tool in goroutine to allow cancellation
622− type toolExecResult struct {
623− response tools.ToolResponse
624− err error
625− }
626− resultChan := make(chan toolExecResult, 1)
627−
628− go func() {
629− response, err := tool.Run(ctx, tools.ToolCall{
630− ID: toolCall.ID,
631− Name: toolCall.Name,
632− Input: toolCall.Input,
633− })
634− resultChan <- toolExecResult{response: response, err: err}
635− }()
636−
637− var toolResponse tools.ToolResponse
638− var toolErr error
639−
640− select {
641− case <-ctx.Done():
642− a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")
643− // Mark remaining tool calls as cancelled
644− for j := i; j < len(toolCalls); j++ {
645− toolResults[j] = message.ToolResult{
646− ToolCallID: toolCalls[j].ID,
647− Content: "Tool execution canceled by user",
648− IsError: true,
649− }
650− }
651− goto out
652− case result := <-resultChan:
653− toolResponse = result.response
654− toolErr = result.err
655− }
656−
657− if toolErr != nil {
658− slog.Error("Tool execution error", "toolCall", toolCall.ID, "error", toolErr)
659− if errors.Is(toolErr, permission.ErrorPermissionDenied) {
660− toolResults[i] = message.ToolResult{
661− ToolCallID: toolCall.ID,
662− Content: "Permission denied",
663− IsError: true,
664− }
665− for j := i + 1; j < len(toolCalls); j++ {
666− toolResults[j] = message.ToolResult{
667− ToolCallID: toolCalls[j].ID,
668− Content: "Tool execution canceled by user",
669− IsError: true,
670− }
671− }
672− a.finishMessage(ctx, &assistantMsg, message.FinishReasonPermissionDenied, "Permission denied", "")
673− break
674− }
675− }
676− toolResults[i] = message.ToolResult{
677− ToolCallID: toolCall.ID,
678− Content: toolResponse.Content,
679− Metadata: toolResponse.Metadata,
680− IsError: toolResponse.IsError,
681− }
682− }
683− }
684−out:
685− if len(toolResults) == 0 {
686− return assistantMsg, nil, nil
687− }
688− parts := make([]message.ContentPart, 0)
689− for _, tr := range toolResults {
690− parts = append(parts, tr)
691− }
692− msg, err := a.messages.Create(context.Background(), assistantMsg.SessionID, message.CreateMessageParams{
693− Role: message.Tool,
694− Parts: parts,
695− Provider: a.providerID,
696− })
697− if err != nil {
698− return assistantMsg, nil, fmt.Errorf("failed to create cancelled tool message: %w", err)
699− }
700−
701− return assistantMsg, &msg, err
702−}
703−
704−func (a *agent) finishMessage(ctx context.Context, msg *message.Message, finishReason message.FinishReason, message, details string) {
705− msg.AddFinish(finishReason, message, details)
706− _ = a.messages.Update(ctx, *msg)
707−}
708−
709−func (a *agent) processEvent(ctx context.Context, sessionID string, assistantMsg *message.Message, event provider.ProviderEvent) error {
710− select {
711− case <-ctx.Done():
712− return ctx.Err()
713− default:
714− // Continue processing.
715− }
716−
717− switch event.Type {
718− case provider.EventThinkingDelta:
719− assistantMsg.AppendReasoningContent(event.Thinking)
720− return a.messages.Update(ctx, *assistantMsg)
721− case provider.EventSignatureDelta:
722− assistantMsg.AppendReasoningSignature(event.Signature)
723− return a.messages.Update(ctx, *assistantMsg)
724− case provider.EventContentDelta:
725− assistantMsg.FinishThinking()
726− assistantMsg.AppendContent(event.Content)
727− return a.messages.Update(ctx, *assistantMsg)
728− case provider.EventToolUseStart:
729− assistantMsg.FinishThinking()
730− slog.Info("Tool call started", "toolCall", event.ToolCall)
731− assistantMsg.AddToolCall(*event.ToolCall)
732− return a.messages.Update(ctx, *assistantMsg)
733− case provider.EventToolUseDelta:
734− assistantMsg.AppendToolCallInput(event.ToolCall.ID, event.ToolCall.Input)
735− return a.messages.Update(ctx, *assistantMsg)
736− case provider.EventToolUseStop:
737− slog.Info("Finished tool call", "toolCall", event.ToolCall)
738− assistantMsg.FinishToolCall(event.ToolCall.ID)
739− return a.messages.Update(ctx, *assistantMsg)
740− case provider.EventError:
741− return event.Error
742− case provider.EventComplete:
743− assistantMsg.FinishThinking()
744− assistantMsg.SetToolCalls(event.Response.ToolCalls)
745− assistantMsg.AddFinish(event.Response.FinishReason, "", "")
746− if err := a.messages.Update(ctx, *assistantMsg); err != nil {
747− return fmt.Errorf("failed to update message: %w", err)
748− }
749− return a.trackUsage(ctx, sessionID, a.Model(), event.Response.Usage)
750− }
751−
752− return nil
753−}
754−
755−func (a *agent) trackUsage(ctx context.Context, sessionID string, model catwalk.Model, usage provider.TokenUsage) error {
756− sess, err := a.sessions.Get(ctx, sessionID)
757− if err != nil {
758− return fmt.Errorf("failed to get session: %w", err)
759− }
760−
761− cost := model.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) +
762− model.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) +
763− model.CostPer1MIn/1e6*float64(usage.InputTokens) +
764− model.CostPer1MOut/1e6*float64(usage.OutputTokens)
765−
766− a.eventTokensUsed(sessionID, usage, cost)
767−
768− sess.Cost += cost
769− sess.CompletionTokens = usage.OutputTokens + usage.CacheReadTokens
770− sess.PromptTokens = usage.InputTokens + usage.CacheCreationTokens
771−
772− _, err = a.sessions.Save(ctx, sess)
773− if err != nil {
774− return fmt.Errorf("failed to save session: %w", err)
775− }
776− return nil
777−}
778−
779−func (a *agent) Summarize(ctx context.Context, sessionID string) error {
780− if a.summarizeProvider == nil {
781− return fmt.Errorf("summarize provider not available")
782− }
783−
784− // Check if session is busy
785− if a.IsSessionBusy(sessionID) {
786− return ErrSessionBusy
787− }
788−
789− // Create a new context with cancellation
790− summarizeCtx, cancel := context.WithCancel(ctx)
791−
792− // Store the cancel function in activeRequests to allow cancellation
793− a.activeRequests.Set(sessionID+"-summarize", cancel)
794−
795− go func() {
796− defer a.activeRequests.Del(sessionID + "-summarize")
797− defer cancel()
798− event := AgentEvent{
799− Type: AgentEventTypeSummarize,
800− Progress: "Starting summarization...",
801− }
802−
803− a.Publish(pubsub.CreatedEvent, event)
804− // Get all messages from the session
805− msgs, err := a.messages.List(summarizeCtx, sessionID)
806− if err != nil {
807− event = AgentEvent{
808− Type: AgentEventTypeError,
809− Error: fmt.Errorf("failed to list messages: %w", err),
810− Done: true,
811− }
812− a.Publish(pubsub.CreatedEvent, event)
813− return
814− }
815− summarizeCtx = context.WithValue(summarizeCtx, tools.SessionIDContextKey, sessionID)
816−
817− if len(msgs) == 0 {
818− event = AgentEvent{
819− Type: AgentEventTypeError,
820− Error: fmt.Errorf("no messages to summarize"),
821− Done: true,
822− }
823− a.Publish(pubsub.CreatedEvent, event)
824− return
825− }
826−
827− event = AgentEvent{
828− Type: AgentEventTypeSummarize,
829− Progress: "Analyzing conversation...",
830− }
831− a.Publish(pubsub.CreatedEvent, event)
832−
833− // Add a system message to guide the summarization
834− summarizePrompt := "Provide a detailed but concise summary of our conversation above. Focus on information that would be helpful for continuing the conversation, including what we did, what we're doing, which files we're working on, and what we're going to do next."
835−
836− // Create a new message with the summarize prompt
837− promptMsg := message.Message{
838− Role: message.User,
839− Parts: []message.ContentPart{message.TextContent{Text: summarizePrompt}},
840− }
841−
842− // Append the prompt to the messages
843− msgsWithPrompt := append(msgs, promptMsg)
844−
845− event = AgentEvent{
846− Type: AgentEventTypeSummarize,
847− Progress: "Generating summary...",
848− }
849−
850− a.Publish(pubsub.CreatedEvent, event)
851−
852− // Send the messages to the summarize provider
853− response := a.summarizeProvider.StreamResponse(
854− summarizeCtx,
855− msgsWithPrompt,
856− nil,
857− )
858− var finalResponse *provider.ProviderResponse
859− for r := range response {
860− if r.Error != nil {
861− event = AgentEvent{
862− Type: AgentEventTypeError,
863− Error: fmt.Errorf("failed to summarize: %w", r.Error),
864− Done: true,
865− }
866− a.Publish(pubsub.CreatedEvent, event)
867− return
868− }
869− finalResponse = r.Response
870− }
871−
872− summary := strings.TrimSpace(finalResponse.Content)
873− if summary == "" {
874− event = AgentEvent{
875− Type: AgentEventTypeError,
876− Error: fmt.Errorf("empty summary returned"),
877− Done: true,
878− }
879− a.Publish(pubsub.CreatedEvent, event)
880− return
881− }
882− shell := shell.GetPersistentShell(config.Get().WorkingDir())
883− summary += "\n\n**Current working directory of the persistent shell**\n\n" + shell.GetWorkingDir()
884− event = AgentEvent{
885− Type: AgentEventTypeSummarize,
886− Progress: "Creating new session...",
887− }
888−
889− a.Publish(pubsub.CreatedEvent, event)
890− oldSession, err := a.sessions.Get(summarizeCtx, sessionID)
891− if err != nil {
892− event = AgentEvent{
893− Type: AgentEventTypeError,
894− Error: fmt.Errorf("failed to get session: %w", err),
895− Done: true,
896− }
897−
898− a.Publish(pubsub.CreatedEvent, event)
899− return
900− }
901− // Create a message in the new session with the summary
902− msg, err := a.messages.Create(summarizeCtx, oldSession.ID, message.CreateMessageParams{
903− Role: message.Assistant,
904− Parts: []message.ContentPart{
905− message.TextContent{Text: summary},
906− message.Finish{
907− Reason: message.FinishReasonEndTurn,
908− Time: time.Now().Unix(),
909− },
910− },
911− Model: a.summarizeProvider.Model().ID,
912− Provider: a.summarizeProviderID,
913− })
914− if err != nil {
915− event = AgentEvent{
916− Type: AgentEventTypeError,
917− Error: fmt.Errorf("failed to create summary message: %w", err),
918− Done: true,
919− }
920−
921− a.Publish(pubsub.CreatedEvent, event)
922− return
923− }
924− oldSession.SummaryMessageID = msg.ID
925− oldSession.CompletionTokens = finalResponse.Usage.OutputTokens
926− oldSession.PromptTokens = 0
927− model := a.summarizeProvider.Model()
928− usage := finalResponse.Usage
929− cost := model.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) +
930− model.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) +
931− model.CostPer1MIn/1e6*float64(usage.InputTokens) +
932− model.CostPer1MOut/1e6*float64(usage.OutputTokens)
933− oldSession.Cost += cost
934− _, err = a.sessions.Save(summarizeCtx, oldSession)
935− if err != nil {
936− event = AgentEvent{
937− Type: AgentEventTypeError,
938− Error: fmt.Errorf("failed to save session: %w", err),
939− Done: true,
940− }
941− a.Publish(pubsub.CreatedEvent, event)
942− }
943−
944− event = AgentEvent{
945− Type: AgentEventTypeSummarize,
946− SessionID: oldSession.ID,
947− Progress: "Summary complete",
948− Done: true,
949− }
950− a.Publish(pubsub.CreatedEvent, event)
951− // Send final success event with the new session ID
952− }()
953−
954− return nil
955−}
956−
957−func (a *agent) ClearQueue(sessionID string) {
958− if a.QueuedPrompts(sessionID) > 0 {
959− slog.Info("Clearing queued prompts", "session_id", sessionID)
960− a.promptQueue.Del(sessionID)
961− }
962−}
963−
964−func (a *agent) CancelAll() {
965− if !a.IsBusy() {
966− return
967− }
968− for key := range a.activeRequests.Seq2() {
969− a.Cancel(key) // key is sessionID
970− }
971−
972− for _, cleanup := range a.cleanupFuncs {
973− if cleanup != nil {
974− cleanup()
975− }
976− }
977−
978− timeout := time.After(5 * time.Second)
979− for a.IsBusy() {
980− select {
981− case <-timeout:
982− return
983− default:
984− time.Sleep(200 * time.Millisecond)
985− }
986− }
987−}
988−
989−func (a *agent) UpdateModel() error {
990− cfg := config.Get()
991−
992− // Get current provider configuration
993− currentProviderCfg := cfg.GetProviderForModel(a.agentCfg.Model)
994− if currentProviderCfg == nil || currentProviderCfg.ID == "" {
995− return fmt.Errorf("provider for agent %s not found in config", a.agentCfg.Name)
996− }
997−
998− // Check if provider has changed
999− if string(currentProviderCfg.ID) != a.providerID {
1000− // Provider changed, need to recreate the main provider
1001− model := cfg.GetModelByType(a.agentCfg.Model)
1002− if model.ID == "" {
1003− return fmt.Errorf("model not found for agent %s", a.agentCfg.Name)
1004− }
1005−
1006− promptID := agentPromptMap[a.agentCfg.ID]
1007− if promptID == "" {
1008− promptID = prompt.PromptDefault
1009− }
1010−
1011− opts := []provider.ProviderClientOption{
1012− provider.WithModel(a.agentCfg.Model),
1013− provider.WithSystemMessage(prompt.GetPrompt(promptID, currentProviderCfg.ID, cfg.Options.ContextPaths...)),
1014− }
1015−
1016− newProvider, err := provider.NewProvider(*currentProviderCfg, opts...)
1017− if err != nil {
1018− return fmt.Errorf("failed to create new provider: %w", err)
1019− }
1020−
1021− // Update the provider and provider ID
1022− a.provider = newProvider
1023− a.providerID = string(currentProviderCfg.ID)
1024− }
1025−
1026− // Check if providers have changed for title (small) and summarize (large)
1027− smallModelCfg := cfg.Models[config.SelectedModelTypeSmall]
1028− var smallModelProviderCfg config.ProviderConfig
1029− for p := range cfg.Providers.Seq() {
1030− if p.ID == smallModelCfg.Provider {
1031− smallModelProviderCfg = p
1032− break
1033− }
1034− }
1035− if smallModelProviderCfg.ID == "" {
1036− return fmt.Errorf("provider %s not found in config", smallModelCfg.Provider)
1037− }
1038−
1039− largeModelCfg := cfg.Models[config.SelectedModelTypeLarge]
1040− var largeModelProviderCfg config.ProviderConfig
1041− for p := range cfg.Providers.Seq() {
1042− if p.ID == largeModelCfg.Provider {
1043− largeModelProviderCfg = p
1044− break
1045− }
1046− }
1047− if largeModelProviderCfg.ID == "" {
1048− return fmt.Errorf("provider %s not found in config", largeModelCfg.Provider)
1049− }
1050−
1051− var maxTitleTokens int64 = 40
1052−
1053− // if the max output is too low for the gemini provider it won't return anything
1054− if smallModelCfg.Provider == "gemini" {
1055− maxTitleTokens = 1000
1056− }
1057− // Recreate title provider
1058− titleOpts := []provider.ProviderClientOption{
1059− provider.WithModel(config.SelectedModelTypeSmall),
1060− provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptTitle, smallModelProviderCfg.ID)),
1061− provider.WithMaxTokens(maxTitleTokens),
1062− }
1063− newTitleProvider, err := provider.NewProvider(smallModelProviderCfg, titleOpts...)
1064− if err != nil {
1065− return fmt.Errorf("failed to create new title provider: %w", err)
1066− }
1067− a.titleProvider = newTitleProvider
1068−
1069− // Recreate summarize provider if provider changed (now large model)
1070− if string(largeModelProviderCfg.ID) != a.summarizeProviderID {
1071− largeModel := cfg.GetModelByType(config.SelectedModelTypeLarge)
1072− if largeModel == nil {
1073− return fmt.Errorf("model %s not found in provider %s", largeModelCfg.Model, largeModelProviderCfg.ID)
1074− }
1075− summarizeOpts := []provider.ProviderClientOption{
1076− provider.WithModel(config.SelectedModelTypeLarge),
1077− provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, largeModelProviderCfg.ID)),
1078− }
1079− newSummarizeProvider, err := provider.NewProvider(largeModelProviderCfg, summarizeOpts...)
1080− if err != nil {
1081− return fmt.Errorf("failed to create new summarize provider: %w", err)
1082− }
1083− a.summarizeProvider = newSummarizeProvider
1084− a.summarizeProviderID = string(largeModelProviderCfg.ID)
1085− }
1086−
1087− return nil
1088−}
1089−
1090−func (a *agent) setupEvents(ctx context.Context) {
1091− ctx, cancel := context.WithCancel(ctx)
1092−
1093− go func() {
1094− subCh := SubscribeMCPEvents(ctx)
1095−
1096− for {
1097− select {
1098− case event, ok := <-subCh:
1099− if !ok {
1100− slog.Debug("MCPEvents subscription channel closed")
1101− return
1102− }
1103− switch event.Payload.Type {
1104− case MCPEventToolsListChanged:
1105− name := event.Payload.Name
1106− c, ok := mcpClients.Get(name)
1107− if !ok {
1108− slog.Warn("MCP client not found for tools update", "name", name)
1109− continue
1110− }
1111− cfg := config.Get()
1112− tools, err := getTools(ctx, name, a.permissions, c, cfg.WorkingDir())
1113− if err != nil {
1114− slog.Error("error listing tools", "error", err)
1115− updateMCPState(name, MCPStateError, err, nil, 0)
1116− _ = c.Close()
1117− continue
1118− }
1119− updateMcpTools(name, tools)
1120− a.mcpTools.Reset(maps.Collect(mcpTools.Seq2()))
1121− updateMCPState(name, MCPStateConnected, nil, c, a.mcpTools.Len())
1122− default:
1123− continue
1124− }
1125− case <-ctx.Done():
1126− slog.Debug("MCPEvents subscription cancelled")
1127− return
1128− }
1129− }
1130− }()
1131−
1132− a.cleanupFuncs = append(a.cleanupFuncs, cancel)
1133−}