From 2510917b7344b69c7606fe1e029a43495847faea Mon Sep 17 00:00:00 2001
From: Kujtim Hoxha <kujtimii.h@gmail.com>
Date: Mon, 30 Jun 2025 15:00:05 +0200
Subject: [PATCH] chore: change how summarize sessions gets triggered

---
 internal/config/config.go   | 38 +++++++++++++++++++++++++++++++++++++
 internal/llm/agent/agent.go |  5 +++++
 internal/tui/tui.go         | 19 +++++++++++++++++--
 3 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/internal/config/config.go b/internal/config/config.go
index 0cccbdfffe5dc1c42e21b03bd9a20d7112005bee..e0e54ff9139fabf299c57aadab8b384335260f22 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -860,6 +860,44 @@ func GetAgentModel(agentID AgentID) Model {
 	return Model{}
 }
 
+// GetAgentEffectiveMaxTokens returns the effective max tokens for an agent,
+// considering any overrides from the preferred model configuration
+func GetAgentEffectiveMaxTokens(agentID AgentID) int64 {
+	cfg := Get()
+	agent, ok := cfg.Agents[agentID]
+	if !ok {
+		logging.Error("Agent not found", "agent_id", agentID)
+		return 0
+	}
+
+	var preferredModel PreferredModel
+	switch agent.Model {
+	case LargeModel:
+		preferredModel = cfg.Models.Large
+	case SmallModel:
+		preferredModel = cfg.Models.Small
+	default:
+		logging.Warn("Unknown model type for agent", "agent_id", agentID, "model_type", agent.Model)
+		preferredModel = cfg.Models.Large // Fallback to large model
+	}
+
+	// Get the base model configuration
+	baseModel := GetAgentModel(agentID)
+	if baseModel.ID == "" {
+		return 0
+	}
+
+	// Start with the default max tokens from the base model
+	maxTokens := baseModel.DefaultMaxTokens
+
+	// Override with preferred model max tokens if set
+	if preferredModel.MaxTokens > 0 {
+		maxTokens = preferredModel.MaxTokens
+	}
+
+	return maxTokens
+}
+
 func GetAgentProvider(agentID AgentID) ProviderConfig {
 	cfg := Get()
 	agent, ok := cfg.Agents[agentID]
diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go
index d165921f639ffee7127e4044c42d154f091a0dca..5568c8a619287619900e8d5e5d5d44c2e85de446 100644
--- a/internal/llm/agent/agent.go
+++ b/internal/llm/agent/agent.go
@@ -50,6 +50,7 @@ type AgentEvent struct {
 type Service interface {
 	pubsub.Suscriber[AgentEvent]
 	Model() config.Model
+	EffectiveMaxTokens() int64
 	Run(ctx context.Context, sessionID string, content string, attachments ...message.Attachment) (<-chan AgentEvent, error)
 	Cancel(sessionID string)
 	CancelAll()
@@ -229,6 +230,10 @@ func (a *agent) Model() config.Model {
 	return config.GetAgentModel(a.agentCfg.ID)
 }
 
+func (a *agent) EffectiveMaxTokens() int64 {
+	return config.GetAgentEffectiveMaxTokens(a.agentCfg.ID)
+}
+
 func (a *agent) Cancel(sessionID string) {
 	// Cancel regular requests
 	if cancelFunc, exists := a.activeRequests.LoadAndDelete(sessionID); exists {
diff --git a/internal/tui/tui.go b/internal/tui/tui.go
index 46a3f3c2a4ee2a6f4b2441a3d10e98f7e46eca2a..a37a55cfbbd05eba217b10691a15d0caed7c183a 100644
--- a/internal/tui/tui.go
+++ b/internal/tui/tui.go
@@ -228,8 +228,23 @@ func (a *appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 			if err == nil {
 				model := a.app.CoderAgent.Model()
 				contextWindow := model.ContextWindow
-				tokens := session.CompletionTokens + session.PromptTokens
-				if (tokens >= int64(float64(contextWindow)*0.95)) && !config.Get().Options.DisableAutoSummarize {
+				usedTokens := session.CompletionTokens + session.PromptTokens
+				remainingTokens := contextWindow - usedTokens
+				
+				// Get effective max tokens for this agent (considering overrides)
+				maxTokens := a.app.CoderAgent.EffectiveMaxTokens()
+				
+				// Apply 10% margin to max tokens
+				maxTokensWithMargin := int64(float64(maxTokens) * 1.1)
+				
+				// Trigger auto-summarize if remaining tokens < max tokens + 10% margin
+				// Also ensure we have a reasonable minimum threshold to avoid too-frequent summaries
+				minThreshold := int64(1000) // Minimum 1000 tokens remaining before triggering
+				if maxTokensWithMargin < minThreshold {
+					maxTokensWithMargin = minThreshold
+				}
+				
+				if remainingTokens < maxTokensWithMargin && !config.Get().Options.DisableAutoSummarize {
 					// Show compact confirmation dialog
 					cmds = append(cmds, util.CmdHandler(dialogs.OpenDialogMsg{
 						Model: compact.NewCompactDialogCmp(a.app.CoderAgent, a.selectedSessionID, false),