Improve summary to keep context (#159)

* improve summary to keep context

* improve loop

* remove debug msg

Author: kujtimiihoxha

internal/db/db.go

@@ -0,0 +1,9 @@
+-- +goose Up
+-- +goose StatementBegin
+ALTER TABLE sessions ADD COLUMN summary_message_id TEXT;
+-- +goose StatementEnd
+
+-- +goose Down
+-- +goose StatementBegin
+ALTER TABLE sessions DROP COLUMN summary_message_id;
+-- +goose StatementEnd

internal/db/files.sql.go

@@ -7,6 +7,7 @@ INSERT INTO sessions (
     prompt_tokens,
     completion_tokens,
     cost,
+    summary_message_id,
     updated_at,
     created_at
 ) VALUES (
@@ -17,6 +18,7 @@ INSERT INTO sessions (
     ?,
     ?,
     ?,
+    null,
     strftime('%s', 'now'),
     strftime('%s', 'now')
 ) RETURNING *;
@@ -38,6 +40,7 @@ SET
     title = ?,
     prompt_tokens = ?,
     completion_tokens = ?,
+    summary_message_id = ?,
     cost = ?
 WHERE id = ?
 RETURNING *;

internal/db/messages.sql.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/opencode-ai/opencode/internal/config"
 	"github.com/opencode-ai/opencode/internal/llm/models"
@@ -245,6 +246,23 @@ func (a *agent) processGeneration(ctx context.Context, sessionID, content string
 			}
 		}()
 	}
+	session, err := a.sessions.Get(ctx, sessionID)
+	if err != nil {
+		return a.err(fmt.Errorf("failed to get session: %w", err))
+	}
+	if session.SummaryMessageID != "" {
+		summaryMsgInex := -1
+		for i, msg := range msgs {
+			if msg.ID == session.SummaryMessageID {
+				summaryMsgInex = i
+				break
+			}
+		}
+		if summaryMsgInex != -1 {
+			msgs = msgs[summaryMsgInex:]
+			msgs[0].Role = message.User
+		}
+	}
 
 	userMsg, err := a.createUserMessage(ctx, sessionID, content, attachmentParts)
 	if err != nil {
@@ -614,37 +632,51 @@ func (a *agent) Summarize(ctx context.Context, sessionID string) error {
 			a.Publish(pubsub.CreatedEvent, event)
 			return
 		}
-		// Create a new session with the summary
-		newSession, err := a.sessions.Create(summarizeCtx, oldSession.Title+" - Continuation")
+		// Create a message in the new session with the summary
+		msg, err := a.messages.Create(summarizeCtx, oldSession.ID, message.CreateMessageParams{
+			Role: message.Assistant,
+			Parts: []message.ContentPart{
+				message.TextContent{Text: summary},
+				message.Finish{
+					Reason: message.FinishReasonEndTurn,
+					Time:   time.Now().Unix(),
+				},
+			},
+			Model: a.summarizeProvider.Model().ID,
+		})
 		if err != nil {
 			event = AgentEvent{
 				Type:  AgentEventTypeError,
-				Error: fmt.Errorf("failed to create new session: %w", err),
+				Error: fmt.Errorf("failed to create summary message: %w", err),
 				Done:  true,
 			}
+
 			a.Publish(pubsub.CreatedEvent, event)
 			return
 		}
-
-		// Create a message in the new session with the summary
-		_, err = a.messages.Create(summarizeCtx, newSession.ID, message.CreateMessageParams{
-			Role:  message.Assistant,
-			Parts: []message.ContentPart{message.TextContent{Text: summary}},
-			Model: a.summarizeProvider.Model().ID,
-		})
+		oldSession.SummaryMessageID = msg.ID
+		oldSession.CompletionTokens = response.Usage.OutputTokens
+		oldSession.PromptTokens = 0
+		model := a.summarizeProvider.Model()
+		usage := response.Usage
+		cost := model.CostPer1MInCached/1e6*float64(usage.CacheCreationTokens) +
+			model.CostPer1MOutCached/1e6*float64(usage.CacheReadTokens) +
+			model.CostPer1MIn/1e6*float64(usage.InputTokens) +
+			model.CostPer1MOut/1e6*float64(usage.OutputTokens)
+		oldSession.Cost += cost
+		_, err = a.sessions.Save(summarizeCtx, oldSession)
 		if err != nil {
 			event = AgentEvent{
 				Type:  AgentEventTypeError,
-				Error: fmt.Errorf("failed to create summary message: %w", err),
+				Error: fmt.Errorf("failed to save session: %w", err),
 				Done:  true,
 			}
-
 			a.Publish(pubsub.CreatedEvent, event)
-			return
 		}
+
 		event = AgentEvent{
 			Type:      AgentEventTypeSummarize,
-			SessionID: newSession.ID,
+			SessionID: oldSession.ID,
 			Progress:  "Summary complete",
 			Done:      true,
 		}

internal/db/migrations/20250515105448_add_summary_message_id.sql

@@ -16,6 +16,7 @@ type Session struct {
 	MessageCount     int64
 	PromptTokens     int64
 	CompletionTokens int64
+	SummaryMessageID string
 	Cost             float64
 	CreatedAt        int64
 	UpdatedAt        int64
@@ -105,7 +106,11 @@ func (s *service) Save(ctx context.Context, session Session) (Session, error) {
 		Title:            session.Title,
 		PromptTokens:     session.PromptTokens,
 		CompletionTokens: session.CompletionTokens,
-		Cost:             session.Cost,
+		SummaryMessageID: sql.NullString{
+			String: session.SummaryMessageID,
+			Valid:  session.SummaryMessageID != "",
+		},
+		Cost: session.Cost,
 	})
 	if err != nil {
 		return Session{}, err
@@ -135,6 +140,7 @@ func (s service) fromDBItem(item db.Session) Session {
 		MessageCount:     item.MessageCount,
 		PromptTokens:     item.PromptTokens,
 		CompletionTokens: item.CompletionTokens,
+		SummaryMessageID: item.SummaryMessageID.String,
 		Cost:             item.Cost,
 		CreatedAt:        item.CreatedAt,
 		UpdatedAt:        item.UpdatedAt,