@@ -1,1133 +0,0 @@
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"
30−type AgentEventType string
33− AgentEventTypeError AgentEventType = "error"
34− AgentEventTypeResponse AgentEventType = "response"
35− AgentEventTypeSummarize AgentEventType = "summarize"
38−type AgentEvent struct {
40− Message message.Message
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)
55− IsSessionBusy(sessionID string) bool
57− Summarize(ctx context.Context, sessionID string) error
59− QueuedPrompts(sessionID string) int
60− ClearQueue(sessionID string)
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]
73− // We need this to be able to update it when model changes
74− agentToolFn func() (tools.BaseTool, error)
75− cleanupFuncs []func()
77− provider provider.Provider
80− titleProvider provider.Provider
81− summarizeProvider provider.Provider
82− summarizeProviderID string
84− activeRequests *csync.Map[string, context.CancelFunc]
85− promptQueue *csync.Map[string, []string]
88−var agentPromptMap = map[string]prompt.PromptID{
89− "coder": prompt.PromptCoder,
90− "task": prompt.PromptTask,
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],
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")
112− taskAgent, err := NewAgent(ctx, taskAgentCfg, permissions, sessions, messages, history, lspClients)
114− return nil, fmt.Errorf("failed to create task agent: %w", err)
116− return NewAgentTool(taskAgent, sessions, messages), nil
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)
124− model := config.Get().GetModelByType(agentCfg.Model)
127− return nil, fmt.Errorf("model not found for agent %s", agentCfg.Name)
130− promptID := agentPromptMap[agentCfg.ID]
132− promptID = prompt.PromptDefault
134− opts := []provider.ProviderClientOption{
135− provider.WithModel(agentCfg.Model),
136− provider.WithSystemMessage(prompt.GetPrompt(promptID, providerCfg.ID, config.Get().Options.ContextPaths...)),
138− agentProvider, err := provider.NewProvider(*providerCfg, opts...)
143− smallModelCfg := cfg.Models[config.SelectedModelTypeSmall]
144− var smallModelProviderCfg *config.ProviderConfig
145− if smallModelCfg.Provider == providerCfg.ID {
146− smallModelProviderCfg = providerCfg
148− smallModelProviderCfg = cfg.GetProviderForModel(config.SelectedModelTypeSmall)
150− if smallModelProviderCfg.ID == "" {
151− return nil, fmt.Errorf("provider %s not found in config", smallModelCfg.Provider)
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)
159− titleOpts := []provider.ProviderClientOption{
160− provider.WithModel(config.SelectedModelTypeSmall),
161− provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptTitle, smallModelProviderCfg.ID)),
163− titleProvider, err := provider.NewProvider(*smallModelProviderCfg, titleOpts...)
168− summarizeOpts := []provider.ProviderClientOption{
169− provider.WithModel(config.SelectedModelTypeLarge),
170− provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, providerCfg.ID)),
172− summarizeProvider, err := provider.NewProvider(*providerCfg, summarizeOpts...)
177− baseToolsFn := func() map[string]tools.BaseTool {
178− slog.Info("Initializing agent base tools", "agent", agentCfg.ID)
180− slog.Info("Initialized agent base tools", "agent", agentCfg.ID)
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),
199− result[tool.Name()] = tool
203− mcpToolsFn := func() map[string]tools.BaseTool {
204− slog.Info("Initializing agent mcp tools", "agent", agentCfg.ID)
206− slog.Info("Initialized agent mcp tools", "agent", agentCfg.ID)
209− mcpToolsOnce.Do(func() {
210− doGetMCPTools(ctx, permissions, cfg)
213− return maps.Collect(mcpTools.Seq2())
217− Broker: pubsub.NewBroker[AgentEvent](),
219− provider: agentProvider,
220− providerID: string(providerCfg.ID),
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,
238−func (a *agent) Model() catwalk.Model {
239− return *config.Get().GetModelByType(a.agentCfg.Model)
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)
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)
255− if a.QueuedPrompts(sessionID) > 0 {
256− slog.Info("Clearing queued prompts", "session_id", sessionID)
257− a.promptQueue.Del(sessionID)
261−func (a *agent) IsBusy() bool {
263− for cancelFunc := range a.activeRequests.Seq() {
264− if cancelFunc != nil {
272−func (a *agent) IsSessionBusy(sessionID string) bool {
273− _, busy := a.activeRequests.Get(sessionID)
277−func (a *agent) QueuedPrompts(sessionID string) int {
278− l, ok := a.promptQueue.Get(sessionID)
285−func (a *agent) generateTitle(ctx context.Context, sessionID string, content string) error {
289− if a.titleProvider == nil {
292− session, err := a.sessions.Get(ctx, sessionID)
296− parts := []message.ContentPart{message.TextContent{
297− Text: fmt.Sprintf("Generate a concise title for the following content:\n\n%s", content),
300− // Use streaming approach like summarization
301− response := a.titleProvider.StreamResponse(
312− var finalResponse *provider.ProviderResponse
313− for r := range response {
317− finalResponse = r.Response
320− if finalResponse == nil {
321− return fmt.Errorf("no response received from title provider")
324− title := strings.ReplaceAll(finalResponse.Content, "\n", " ")
326− if idx := strings.Index(title, "</think>"); idx > 0 {
327− title = title[idx+len("</think>"):]
330− title = strings.TrimSpace(title)
335− session.Title = title
336− _, err = a.sessions.Save(ctx, session)
340−func (a *agent) err(err error) AgentEvent {
342− Type: AgentEventTypeError,
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 {
351− events := make(chan AgentEvent, 1)
352− if a.IsSessionBusy(sessionID) {
353− existing, ok := a.promptQueue.Get(sessionID)
355− existing = []string{}
357− existing = append(existing, content)
358− a.promptQueue.Set(sessionID, existing)
362− genCtx, cancel := context.WithCancel(ctx)
363− a.activeRequests.Set(sessionID, cancel)
364− startTime := time.Now()
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"))
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})
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)
380− slog.Error("Request errored", "sessionID", sessionID, "error", result.Error.Error())
381− event.Error(result.Error)
384− slog.Debug("Request completed", "sessionID", sessionID)
386− a.eventPromptResponded(sessionID, time.Since(startTime).Truncate(time.Second))
387− a.activeRequests.Del(sessionID)
389− a.Publish(pubsub.CreatedEvent, result)
393− a.eventPromptSent(sessionID)
397−func (a *agent) processGeneration(ctx context.Context, sessionID, content string, attachmentParts []message.ContentPart) AgentEvent {
399− // List existing messages; if none, start title generation asynchronously.
400− msgs, err := a.messages.List(ctx, sessionID)
402− return a.err(fmt.Errorf("failed to list messages: %w", err))
406− defer log.RecoverPanic("agent.Run", func() {
407− slog.Error("panic while generating title")
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)
415− session, err := a.sessions.Get(ctx, sessionID)
417− return a.err(fmt.Errorf("failed to get session: %w", err))
419− if session.SummaryMessageID != "" {
420− summaryMsgInex := -1
421− for i, msg := range msgs {
422− if msg.ID == session.SummaryMessageID {
427− if summaryMsgInex != -1 {
428− msgs = msgs[summaryMsgInex:]
429− msgs[0].Role = message.User
433− userMsg, err := a.createUserMessage(ctx, sessionID, content, attachmentParts)
435− return a.err(fmt.Errorf("failed to create user message: %w", err))
437− // Append the new user message to the conversation history.
438− msgHistory := append(msgs, userMsg)
441− // Check for cancellation before each iteration
444− return a.err(ctx.Err())
446− // Continue processing
448− agentMessage, toolResults, err := a.streamAndHandleEvents(ctx, sessionID, msgHistory)
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)
455− return a.err(fmt.Errorf("failed to process events: %w", err))
457− if cfg.Options.Debug {
458− slog.Info("Result", "message", agentMessage.FinishReason(), "toolResults", toolResults)
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)
466− for _, prompt := range nextPrompt {
467− // Create a new user message for the queued prompt
468− userMsg, err := a.createUserMessage(ctx, sessionID, prompt, nil)
470− return a.err(fmt.Errorf("failed to create user message for queued prompt: %w", err))
472− // Append the new user message to the conversation history
473− msgHistory = append(msgHistory, userMsg)
478− } else if agentMessage.FinishReason() == message.FinishReasonEndTurn {
479− queuePrompts, ok := a.promptQueue.Take(sessionID)
481− for _, prompt := range queuePrompts {
485− userMsg, err := a.createUserMessage(ctx, sessionID, prompt, nil)
487− return a.err(fmt.Errorf("failed to create user message for queued prompt: %w", err))
489− msgHistory = append(msgHistory, userMsg)
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)
501− Type: AgentEventTypeResponse,
502− Message: agentMessage,
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{
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)
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))
530− if a.agentToolFn != nil {
531− agentTool, agentToolErr := a.agentToolFn()
532− if agentToolErr != nil {
533− return nil, agentToolErr
535− allTools = append(allTools, agentTool)
537− return allTools, nil
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)
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,
551− return assistantMsg, nil, fmt.Errorf("failed to create assistant message: %w", err)
554− allTools, toolsErr := a.getAllTools()
555− if toolsErr != nil {
556− return assistantMsg, nil, toolsErr
558− // Now collect tools (which may block on MCP initialization)
559− eventChan := a.provider.StreamResponse(ctx, msgHistory, allTools)
561− // Add the session and message ID into the context if needed by tools.
562− ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)
567− case event, ok := <-eventChan:
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", "")
575− a.finishMessage(ctx, &assistantMsg, message.FinishReasonError, "API Error", processErr.Error())
577− return assistantMsg, nil, processErr
580− a.finishMessage(context.Background(), &assistantMsg, message.FinishReasonCanceled, "Request cancelled", "")
581− return assistantMsg, nil, ctx.Err()
585− toolResults := make([]message.ToolResult, len(assistantMsg.ToolCalls()))
586− toolCalls := assistantMsg.ToolCalls()
587− for i, toolCall := range toolCalls {
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",
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
613− toolResults[i] = message.ToolResult{
614− ToolCallID: toolCall.ID,
615− Content: fmt.Sprintf("Tool not found: %s", toolCall.Name),
621− // Run tool in goroutine to allow cancellation
622− type toolExecResult struct {
623− response tools.ToolResponse
626− resultChan := make(chan toolExecResult, 1)
629− response, err := tool.Run(ctx, tools.ToolCall{
631− Name: toolCall.Name,
632− Input: toolCall.Input,
634− resultChan <- toolExecResult{response: response, err: err}
637− var toolResponse tools.ToolResponse
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",
652− case result := <-resultChan:
653− toolResponse = result.response
654− toolErr = result.err
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",
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",
672− a.finishMessage(ctx, &assistantMsg, message.FinishReasonPermissionDenied, "Permission denied", "")
676− toolResults[i] = message.ToolResult{
677− ToolCallID: toolCall.ID,
678− Content: toolResponse.Content,
679− Metadata: toolResponse.Metadata,
680− IsError: toolResponse.IsError,
685− if len(toolResults) == 0 {
686− return assistantMsg, nil, nil
688− parts := make([]message.ContentPart, 0)
689− for _, tr := range toolResults {
690− parts = append(parts, tr)
692− msg, err := a.messages.Create(context.Background(), assistantMsg.SessionID, message.CreateMessageParams{
695− Provider: a.providerID,
698− return assistantMsg, nil, fmt.Errorf("failed to create cancelled tool message: %w", err)
701− return assistantMsg, &msg, err
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)
709−func (a *agent) processEvent(ctx context.Context, sessionID string, assistantMsg *message.Message, event provider.ProviderEvent) error {
714− // Continue processing.
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:
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)
749− return a.trackUsage(ctx, sessionID, a.Model(), event.Response.Usage)
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)
758− return fmt.Errorf("failed to get session: %w", err)
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)
766− a.eventTokensUsed(sessionID, usage, cost)
769− sess.CompletionTokens = usage.OutputTokens + usage.CacheReadTokens
770− sess.PromptTokens = usage.InputTokens + usage.CacheCreationTokens
772− _, err = a.sessions.Save(ctx, sess)
774− return fmt.Errorf("failed to save session: %w", err)
779−func (a *agent) Summarize(ctx context.Context, sessionID string) error {
780− if a.summarizeProvider == nil {
781− return fmt.Errorf("summarize provider not available")
784− // Check if session is busy
785− if a.IsSessionBusy(sessionID) {
786− return ErrSessionBusy
789− // Create a new context with cancellation
790− summarizeCtx, cancel := context.WithCancel(ctx)
792− // Store the cancel function in activeRequests to allow cancellation
793− a.activeRequests.Set(sessionID+"-summarize", cancel)
796− defer a.activeRequests.Del(sessionID + "-summarize")
798− event := AgentEvent{
799− Type: AgentEventTypeSummarize,
800− Progress: "Starting summarization...",
803− a.Publish(pubsub.CreatedEvent, event)
804− // Get all messages from the session
805− msgs, err := a.messages.List(summarizeCtx, sessionID)
808− Type: AgentEventTypeError,
809− Error: fmt.Errorf("failed to list messages: %w", err),
812− a.Publish(pubsub.CreatedEvent, event)
815− summarizeCtx = context.WithValue(summarizeCtx, tools.SessionIDContextKey, sessionID)
819− Type: AgentEventTypeError,
820− Error: fmt.Errorf("no messages to summarize"),
823− a.Publish(pubsub.CreatedEvent, event)
828− Type: AgentEventTypeSummarize,
829− Progress: "Analyzing conversation...",
831− a.Publish(pubsub.CreatedEvent, event)
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."
836− // Create a new message with the summarize prompt
837− promptMsg := message.Message{
839− Parts: []message.ContentPart{message.TextContent{Text: summarizePrompt}},
842− // Append the prompt to the messages
843− msgsWithPrompt := append(msgs, promptMsg)
846− Type: AgentEventTypeSummarize,
847− Progress: "Generating summary...",
850− a.Publish(pubsub.CreatedEvent, event)
852− // Send the messages to the summarize provider
853− response := a.summarizeProvider.StreamResponse(
858− var finalResponse *provider.ProviderResponse
859− for r := range response {
862− Type: AgentEventTypeError,
863− Error: fmt.Errorf("failed to summarize: %w", r.Error),
866− a.Publish(pubsub.CreatedEvent, event)
869− finalResponse = r.Response
872− summary := strings.TrimSpace(finalResponse.Content)
875− Type: AgentEventTypeError,
876− Error: fmt.Errorf("empty summary returned"),
879− a.Publish(pubsub.CreatedEvent, event)
882− shell := shell.GetPersistentShell(config.Get().WorkingDir())
883− summary += "\n\n**Current working directory of the persistent shell**\n\n" + shell.GetWorkingDir()
885− Type: AgentEventTypeSummarize,
886− Progress: "Creating new session...",
889− a.Publish(pubsub.CreatedEvent, event)
890− oldSession, err := a.sessions.Get(summarizeCtx, sessionID)
893− Type: AgentEventTypeError,
894− Error: fmt.Errorf("failed to get session: %w", err),
898− a.Publish(pubsub.CreatedEvent, event)
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},
907− Reason: message.FinishReasonEndTurn,
908− Time: time.Now().Unix(),
911− Model: a.summarizeProvider.Model().ID,
912− Provider: a.summarizeProviderID,
916− Type: AgentEventTypeError,
917− Error: fmt.Errorf("failed to create summary message: %w", err),
921− a.Publish(pubsub.CreatedEvent, event)
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)
937− Type: AgentEventTypeError,
938− Error: fmt.Errorf("failed to save session: %w", err),
941− a.Publish(pubsub.CreatedEvent, event)
945− Type: AgentEventTypeSummarize,
946− SessionID: oldSession.ID,
947− Progress: "Summary complete",
950− a.Publish(pubsub.CreatedEvent, event)
951− // Send final success event with the new session ID
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)
964−func (a *agent) CancelAll() {
968− for key := range a.activeRequests.Seq2() {
969− a.Cancel(key) // key is sessionID
972− for _, cleanup := range a.cleanupFuncs {
978− timeout := time.After(5 * time.Second)
984− time.Sleep(200 * time.Millisecond)
989−func (a *agent) UpdateModel() error {
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)
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)
1006− promptID := agentPromptMap[a.agentCfg.ID]
1007− if promptID == "" {
1008− promptID = prompt.PromptDefault
1011− opts := []provider.ProviderClientOption{
1012− provider.WithModel(a.agentCfg.Model),
1013− provider.WithSystemMessage(prompt.GetPrompt(promptID, currentProviderCfg.ID, cfg.Options.ContextPaths...)),
1016− newProvider, err := provider.NewProvider(*currentProviderCfg, opts...)
1018− return fmt.Errorf("failed to create new provider: %w", err)
1021− // Update the provider and provider ID
1022− a.provider = newProvider
1023− a.providerID = string(currentProviderCfg.ID)
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
1035− if smallModelProviderCfg.ID == "" {
1036− return fmt.Errorf("provider %s not found in config", smallModelCfg.Provider)
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
1047− if largeModelProviderCfg.ID == "" {
1048− return fmt.Errorf("provider %s not found in config", largeModelCfg.Provider)
1051− var maxTitleTokens int64 = 40
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
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),
1063− newTitleProvider, err := provider.NewProvider(smallModelProviderCfg, titleOpts...)
1065− return fmt.Errorf("failed to create new title provider: %w", err)
1067− a.titleProvider = newTitleProvider
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)
1075− summarizeOpts := []provider.ProviderClientOption{
1076− provider.WithModel(config.SelectedModelTypeLarge),
1077− provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, largeModelProviderCfg.ID)),
1079− newSummarizeProvider, err := provider.NewProvider(largeModelProviderCfg, summarizeOpts...)
1081− return fmt.Errorf("failed to create new summarize provider: %w", err)
1083− a.summarizeProvider = newSummarizeProvider
1084− a.summarizeProviderID = string(largeModelProviderCfg.ID)
1090−func (a *agent) setupEvents(ctx context.Context) {
1091− ctx, cancel := context.WithCancel(ctx)
1094− subCh := SubscribeMCPEvents(ctx)
1098− case event, ok := <-subCh:
1100− slog.Debug("MCPEvents subscription channel closed")
1103− switch event.Payload.Type {
1104− case MCPEventToolsListChanged:
1105− name := event.Payload.Name
1106− c, ok := mcpClients.Get(name)
1108− slog.Warn("MCP client not found for tools update", "name", name)
1111− cfg := config.Get()
1112− tools, err := getTools(ctx, name, a.permissions, c, cfg.WorkingDir())
1114− slog.Error("error listing tools", "error", err)
1115− updateMCPState(name, MCPStateError, err, nil, 0)
1119− updateMcpTools(name, tools)
1120− a.mcpTools.Reset(maps.Collect(mcpTools.Seq2()))
1121− updateMCPState(name, MCPStateConnected, nil, c, a.mcpTools.Len())
1126− slog.Debug("MCPEvents subscription cancelled")
1132− a.cleanupFuncs = append(a.cleanupFuncs, cancel)