@@ -62,36 +62,41 @@ type Model struct {
6262}
6363
6464type sessionAgent struct {
65− largeModel Model
66− smallModel Model
67− systemPrompt string
68− tools []ai.AgentTool
69− sessions session.Service
70− messages message.Service
65+ largeModel Model
66+ smallModel Model
67+ systemPrompt string
68+ tools []ai.AgentTool
69+ sessions session.Service
70+ messages message.Service
71+ disableAutoSummarize bool
7172
7273 messageQueue *csync.Map[string, []SessionAgentCall]
7374 activeRequests *csync.Map[string, context.CancelFunc]
7475}
7576
76−type SessionAgentOption func(*sessionAgent)
77+type SessionAgentOptions struct {
78+ LargeModel Model
79+ SmallModel Model
80+ SystemPrompt string
81+ DisableAutoSummarize bool
82+ Sessions session.Service
83+ Messages message.Service
84+ Tools []ai.AgentTool
85+}
7786
7887func NewSessionAgent(
79− largeModel Model,
80− smallModel Model,
81− systemPrompt string,
82− sessions session.Service,
83− messages message.Service,
84− tools ...ai.AgentTool,
88+ opts SessionAgentOptions,
8589) SessionAgent {
8690 return &sessionAgent{
87− largeModel: largeModel,
88− smallModel: smallModel,
89− systemPrompt: systemPrompt,
90− sessions: sessions,
91− messages: messages,
92− tools: tools,
93− messageQueue: csync.NewMap[string, []SessionAgentCall](),
94− activeRequests: csync.NewMap[string, context.CancelFunc](),
91+ largeModel: opts.LargeModel,
92+ smallModel: opts.SmallModel,
93+ systemPrompt: opts.SystemPrompt,
94+ sessions: opts.Sessions,
95+ messages: opts.Messages,
96+ disableAutoSummarize: opts.DisableAutoSummarize,
97+ tools: opts.Tools,
98+ messageQueue: csync.NewMap[string, []SessionAgentCall](),
99+ activeRequests: csync.NewMap[string, context.CancelFunc](),
95100 }
96101}
97102
@@ -164,6 +169,7 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
164169 history, files := a.preparePrompt(msgs, call.Attachments...)
165170
166171 var currentAssistant *message.Message
172+ var shouldSummarize bool
167173 result, err := agent.Stream(genCtx, ai.AgentStreamCall{
168174 Prompt: call.Prompt,
169175 Files: files,
@@ -292,16 +298,16 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
292298 IsError: isError,
293299 Metadata: result.ClientMetadata,
294300 }
295− _, err := a.messages.Create(context.Background(), currentAssistant.SessionID, message.CreateMessageParams{
301+ _, createMsgErr := a.messages.Create(genCtx, currentAssistant.SessionID, message.CreateMessageParams{
296302 Role: message.Tool,
297303 Parts: []message.ContentPart{
298304 toolResult,
299305 },
300306 })
301− if err != nil {
302− return err
307+ if createMsgErr != nil {
308+ return createMsgErr
303309 }
304− return a.messages.Update(genCtx, *currentAssistant)
310+ return nil
305311 },
306312 OnStepFinish: func(stepResult ai.StepResult) error {
307313 finishReason := message.FinishReasonUnknown
@@ -323,6 +329,18 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
323329 }
324330 return a.messages.Update(genCtx, *currentAssistant)
325331 },
332+ StopWhen: []ai.StopCondition{
333+ func(_ []ai.StepResult) bool {
334+ contextWindow := a.largeModel.CatwalkCfg.ContextWindow
335+ tokens := currentSession.CompletionTokens + currentSession.PromptTokens
336+ percentage := (float64(tokens) / float64(contextWindow)) * 100
337+ if (percentage > 80) && !a.disableAutoSummarize {
338+ shouldSummarize = true
339+ return true
340+ }
341+ return false
342+ },
343+ },
326344 })
327345 if err != nil {
328346 isCancelErr := errors.Is(err, context.Canceled)
@@ -358,28 +376,29 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
358376 break
359377 }
360378 }
361− if !found {
362− content := "There was an error while executing the tool"
363− if isCancelErr {
364− content = "Tool execution canceled by user"
365− } else if isPermissionErr {
366− content = "Permission denied"
367− }
368− toolResult := message.ToolResult{
369− ToolCallID: tc.ID,
370− Name: tc.Name,
371− Content: content,
372− IsError: true,
373− }
374− _, createErr = a.messages.Create(context.Background(), currentAssistant.SessionID, message.CreateMessageParams{
375− Role: message.Tool,
376− Parts: []message.ContentPart{
377− toolResult,
378− },
379− })
380− if createErr != nil {
381− return nil, createErr
382− }
379+ if found {
380+ continue
381+ }
382+ content := "There was an error while executing the tool"
383+ if isCancelErr {
384+ content = "Tool execution canceled by user"
385+ } else if isPermissionErr {
386+ content = "Permission denied"
387+ }
388+ toolResult := message.ToolResult{
389+ ToolCallID: tc.ID,
390+ Name: tc.Name,
391+ Content: content,
392+ IsError: true,
393+ }
394+ _, createErr = a.messages.Create(context.Background(), currentAssistant.SessionID, message.CreateMessageParams{
395+ Role: message.Tool,
396+ Parts: []message.ContentPart{
397+ toolResult,
398+ },
399+ })
400+ if createErr != nil {
401+ return nil, createErr
383402 }
384403 }
385404 if isCancelErr {
@@ -398,6 +417,13 @@ func (a *sessionAgent) Run(ctx context.Context, call SessionAgentCall) (*ai.Agen
398417 }
399418 wg.Wait()
400419
420+ if shouldSummarize {
421+ a.activeRequests.Del(call.SessionID)
422+ if summarizeErr := a.Summarize(genCtx, call.SessionID); summarizeErr != nil {
423+ return nil, summarizeErr
424+ }
425+ }
426+
401427 queuedMessages, ok := a.messageQueue.Get(call.SessionID)
402428 if !ok || len(queuedMessages) == 0 {
403429 return result, err
@@ -437,20 +463,21 @@ func (a *sessionAgent) Summarize(ctx context.Context, sessionID string) error {
437463 ai.WithSystemPrompt(string(summaryPrompt)),
438464 )
439465 summaryMessage, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{
440− Role: message.Assistant,
441− Model: a.largeModel.Model.Model(),
442− Provider: a.largeModel.Model.Provider(),
466+ Role: message.Assistant,
467+ Model: a.largeModel.Model.Model(),
468+ Provider: a.largeModel.Model.Provider(),
469+ IsSummaryMessage: true,
443470 })
444471 if err != nil {
445472 return err
446473 }
447474
448− resp, err := agent.Stream(ctx, ai.AgentStreamCall{
475+ resp, err := agent.Stream(genCtx, ai.AgentStreamCall{
449476 Prompt: "Provide a detailed summary of our conversation above.",
450477 Messages: aiMsgs,
451478 OnReasoningDelta: func(id string, text string) error {
452479 summaryMessage.AppendReasoningContent(text)
453− return a.messages.Update(ctx, summaryMessage)
480+ return a.messages.Update(genCtx, summaryMessage)
454481 },
455482 OnReasoningEnd: func(id string, reasoning ai.ReasoningContent) error {
456483 // handle anthropic signature
@@ -460,14 +487,20 @@ func (a *sessionAgent) Summarize(ctx context.Context, sessionID string) error {
460487 }
461488 }
462489 summaryMessage.FinishThinking()
463− return a.messages.Update(ctx, summaryMessage)
490+ return a.messages.Update(genCtx, summaryMessage)
464491 },
465492 OnTextDelta: func(id, text string) error {
466493 summaryMessage.AppendContent(text)
467− return a.messages.Update(ctx, summaryMessage)
494+ return a.messages.Update(genCtx, summaryMessage)
468495 },
469496 })
470497 if err != nil {
498+ isCancelErr := errors.Is(err, context.Canceled)
499+ if isCancelErr {
500+ // User cancelled summarize we need to remove the summary message
501+ deleteErr := a.messages.Delete(ctx, summaryMessage.ID)
502+ return deleteErr
503+ }
471504 return err
472505 }
473506