chore: change how summarize sessions gets triggered

Kujtim Hoxha created 4 months ago

Change summary

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

Detailed changes

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]

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 {

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),