Persist compaction atomically

pkg/runtime/event.go

@@ -506,6 +506,11 @@ type SessionSummaryEvent struct {
 	Cost           float64     `json:"cost,omitempty"`
 	Model          string      `json:"model,omitempty"`
 	Usage          *chat.Usage `json:"usage,omitempty"`
+
+	// persisted is set when the compaction path wrote this summary directly.
+	// The persistence observer uses it to avoid writing the same summary again
+	// on normal observed RunStream executions. It is intentionally not serialized.
+	persisted bool
 }
 
 // SessionSummary builds the event announcing an applied compaction summary.

pkg/runtime/persistence_observer.go

@@ -119,6 +119,9 @@ func (p *PersistenceObserver) OnEvent(ctx context.Context, sess *session.Session
 		}
 
 	case *SessionSummaryEvent:
+		if e.persisted {
+			break
+		}
 		item := session.Item{Summary: e.Summary, FirstKeptEntry: e.FirstKeptEntry, Cost: e.Cost, Model: e.Model}
 		if e.Usage != nil {
 			// Copy so the persisted item doesn't alias the event's pointer.

pkg/runtime/remote_runtime.go

@@ -863,6 +863,10 @@ func (s *RemoteSessionStore) AddSubSession(context.Context, string, *session.Ses
 	return fmt.Errorf("add sub session: %w", ErrUnsupported)
 }
 
+func (s *RemoteSessionStore) PersistCompaction(context.Context, *session.Session, int64, int64, session.Item) error {
+	return fmt.Errorf("persist compaction: %w", ErrUnsupported)
+}
+
 func (s *RemoteSessionStore) AddSummary(context.Context, string, session.Item) error {
 	return fmt.Errorf("add summary: %w", ErrUnsupported)
 }

pkg/runtime/session_compaction.go

@@ -2,6 +2,7 @@ package runtime
 
 import (
 	"context"
+	"fmt"
 	"log/slog"
 
 	"github.com/docker/docker-agent/pkg/agent"
@@ -127,20 +128,33 @@ func (r *LocalRuntime) doCompact(ctx context.Context, sess *session.Session, a *
 	// the new summary's estimated size.
 	preInputTokens, preOutputTokens := sess.Usage()
 
-	// Apply the summary to the session. This is intrinsically
-	// runtime-private: it mutates session-internal state and persists
-	// through the runtime's session store.
-	sess.ApplyCompaction(result.InputTokens, 0, session.Item{
+	// Apply and persist the summary as one intrinsic successful-compaction step.
+	// Manual Summarize calls do not pass through the RunStream observer chain,
+	// while normal observed runs do; marking the emitted event as persisted
+	// keeps the observer from appending a duplicate row.
+	item := session.Item{
 		Summary:        result.Summary,
 		FirstKeptEntry: result.FirstKeptEntry,
 		Cost:           result.Cost,
 		Model:          result.Model,
 		Usage:          summaryUsage(result),
-	})
-	_ = r.sessionStore.UpdateSession(ctx, sess)
+	}
+	// Atomically persist the metadata and summary before mutating the live
+	// session. A failed write is a failed compaction: no success summary event
+	// is emitted and the in-memory continuation remains unchanged.
+	if err := r.sessionStore.PersistCompaction(ctx, sess, result.InputTokens, 0, item); err != nil {
+		slog.ErrorContext(ctx, "Failed to persist session compaction", "session_id", sess.ID, "error", err)
+		events.Emit(ErrorForSession(sess.ID, fmt.Sprintf("Failed to persist session compaction: %v", err)))
+		outcome = CompactionOutcomeFailed
+		return
+	}
 
 	slog.DebugContext(ctx, "Generated session summary", "session_id", sess.ID, "summary_length", len(result.Summary))
-	events.Emit(SessionSummary(sess.ID, result.Summary, a.Name(), result.FirstKeptEntry, result.Cost, result.Model, summaryUsage(result)))
+	summaryEvent := SessionSummary(sess.ID, result.Summary, a.Name(), result.FirstKeptEntry, result.Cost, result.Model, summaryUsage(result))
+	if e, ok := summaryEvent.(*SessionSummaryEvent); ok {
+		e.persisted = true
+	}
+	events.Emit(summaryEvent)
 
 	// after_compaction: observational. Fired only when a summary was
 	// actually applied to the session. The hook receives the

pkg/runtime/session_compaction_test.go

@@ -1,8 +1,12 @@
 package runtime
 
 import (
+	"context"
+	"errors"
 	"os"
+	"path/filepath"
 	"runtime"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -13,6 +17,7 @@ import (
 	"github.com/docker/docker-agent/pkg/config/latest"
 	"github.com/docker/docker-agent/pkg/modelsdev"
 	"github.com/docker/docker-agent/pkg/session"
+	"github.com/docker/docker-agent/pkg/session/sqlitestore"
 	"github.com/docker/docker-agent/pkg/team"
 )
 
@@ -301,6 +306,179 @@ func TestDoCompactAfterHookFires(t *testing.T) {
 		"after_compaction must receive the produced summary and the *pre-compaction* token counts")
 }
 
+type failingCompactionStore struct {
+	session.Store
+
+	err error
+}
+
+func (s failingCompactionStore) PersistCompaction(context.Context, *session.Session, int64, int64, session.Item) error {
+	return s.err
+}
+
+func TestDoCompactInMemoryStoreAppendsSummaryOnce(t *testing.T) {
+	store := session.NewInMemorySessionStore()
+	summaryStream := newStreamBuilder().AddContent("one summary").AddStopWithUsage(1, 1).Build()
+	prov := &queueProvider{id: "test/mock-model", streams: []chat.MessageStream{summaryStream}}
+	root := agent.New("root", "test", agent.WithModel(prov))
+	rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)),
+		WithSessionCompaction(false), WithSessionStore(store),
+		WithModelStore(mockModelStoreWithLimit{limit: 100_000}))
+	require.NoError(t, err)
+
+	sess := session.New(session.WithID("compact-memory"), session.WithMessages([]session.Item{
+		session.NewMessageItem(session.UserMessage("hi")),
+		session.NewMessageItem(&session.Message{Message: chat.Message{Role: chat.MessageRoleAssistant, Content: "hello"}}),
+	}))
+	require.NoError(t, store.AddSession(t.Context(), sess))
+
+	events := make(chan Event, 32)
+	rt.Summarize(t.Context(), sess, "", NewChannelSink(events))
+	close(events)
+	for range events {
+	}
+
+	reloaded, err := store.GetSession(t.Context(), sess.ID)
+	require.NoError(t, err)
+	var summaries int
+	for _, item := range reloaded.Messages {
+		if item.Summary != "" {
+			summaries++
+		}
+	}
+	assert.Equal(t, 1, summaries)
+	assert.Same(t, sess, reloaded, "aliasing store must mutate the live session exactly once")
+}
+
+func TestDoCompactPersistenceFailureReportsFailedWithoutApplying(t *testing.T) {
+	base := session.NewInMemorySessionStore()
+	persistErr := errors.New("disk full")
+	store := failingCompactionStore{Store: base, err: persistErr}
+	summaryStream := newStreamBuilder().AddContent("lost summary").AddStopWithUsage(1, 1).Build()
+	prov := &queueProvider{id: "test/mock-model", streams: []chat.MessageStream{summaryStream}}
+	root := agent.New("root", "test", agent.WithModel(prov))
+	rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)),
+		WithSessionCompaction(false), WithSessionStore(store),
+		WithModelStore(mockModelStoreWithLimit{limit: 100_000}))
+	require.NoError(t, err)
+
+	sess := session.New(session.WithID("compact-fails"), session.WithMessages([]session.Item{
+		session.NewMessageItem(session.UserMessage("hi")),
+		session.NewMessageItem(&session.Message{Message: chat.Message{Role: chat.MessageRoleAssistant, Content: "hello"}}),
+	}))
+	require.NoError(t, base.AddSession(t.Context(), sess))
+	before := len(sess.Messages)
+
+	events := make(chan Event, 32)
+	rt.Summarize(t.Context(), sess, "", NewChannelSink(events))
+	close(events)
+	var outcome string
+	var sawError, sawSummary bool
+	for ev := range events {
+		switch e := ev.(type) {
+		case *SessionCompactionEvent:
+			if e.Status == "completed" {
+				outcome = e.Outcome
+			}
+		case *ErrorEvent:
+			sawError = true
+		case *SessionSummaryEvent:
+			sawSummary = true
+		}
+	}
+	assert.Equal(t, CompactionOutcomeFailed, outcome)
+	assert.True(t, sawError)
+	assert.False(t, sawSummary, "an unpersisted summary must not be announced as successful")
+	assert.Len(t, sess.Messages, before, "failed persistence must leave live continuation unchanged")
+}
+
+func TestDoCompactPersistsSummaryForSQLiteReload(t *testing.T) {
+	dbPath := filepath.Join(t.TempDir(), "sessions.db")
+	store, err := sqlitestore.New(t.Context(), dbPath)
+	require.NoError(t, err)
+
+	summaryStream := newStreamBuilder().AddContent("persisted summary").AddStopWithUsage(12, 3).Build()
+	prov := &queueProvider{id: "test/mock-model", streams: []chat.MessageStream{summaryStream}}
+	root := agent.New("root", "test", agent.WithModel(prov))
+	rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)),
+		WithSessionCompaction(false),
+		WithSessionStore(store),
+		WithModelStore(mockModelStoreWithLimit{limit: 100_000}),
+	)
+	require.NoError(t, err)
+
+	sess := session.New(
+		session.WithID("compact-reload"),
+		session.WithMessages([]session.Item{
+			session.NewMessageItem(session.UserMessage("old question")),
+			session.NewMessageItem(&session.Message{Message: chat.Message{Role: chat.MessageRoleAssistant, Content: "old answer"}}),
+		}),
+	)
+	require.NoError(t, store.AddSession(t.Context(), sess))
+
+	events := make(chan Event, 32)
+	rt.Summarize(t.Context(), sess, "", NewChannelSink(events))
+	close(events)
+	for range events {
+	}
+	require.NoError(t, store.Close())
+
+	reopened, err := sqlitestore.New(t.Context(), dbPath)
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, reopened.Close()) })
+	reloaded, err := reopened.GetSession(t.Context(), sess.ID)
+	require.NoError(t, err)
+	require.Len(t, reloaded.Messages, 3)
+	assert.Equal(t, "persisted summary", reloaded.Messages[2].Summary)
+
+	messages := reloaded.GetMessages(root)
+	var continuation strings.Builder
+	for _, msg := range messages {
+		continuation.WriteString(msg.Content)
+	}
+	assert.Contains(t, continuation.String(), "persisted summary")
+	assert.NotContains(t, continuation.String(), "old question", "reload must continue from the compacted view")
+}
+
+func TestDoCompactObservedRunDoesNotDuplicateSummary(t *testing.T) {
+	dbPath := filepath.Join(t.TempDir(), "sessions.db")
+	store, err := sqlitestore.New(t.Context(), dbPath)
+	require.NoError(t, err)
+	summaryStream := newStreamBuilder().AddContent("one summary").AddStopWithUsage(1, 1).Build()
+	prov := &queueProvider{id: "test/mock-model", streams: []chat.MessageStream{summaryStream}}
+	root := agent.New("root", "test", agent.WithModel(prov))
+	rt, err := NewLocalRuntime(t.Context(), team.New(team.WithAgents(root)),
+		WithSessionCompaction(false),
+		WithSessionStore(store),
+		WithModelStore(mockModelStoreWithLimit{limit: 100_000}),
+	)
+	require.NoError(t, err)
+
+	sess := session.New(session.WithID("compact-dedup"), session.WithMessages([]session.Item{
+		session.NewMessageItem(session.UserMessage("hi")),
+		session.NewMessageItem(&session.Message{Message: chat.Message{Role: chat.MessageRoleAssistant, Content: "hello"}}),
+	}))
+	require.NoError(t, store.AddSession(t.Context(), sess))
+
+	inner := make(chan Event, 32)
+	observed := rt.observe(t.Context(), sess, inner)
+	rt.compactWithReason(t.Context(), sess, "", compactionReasonManual, NewChannelSink(inner))
+	close(inner)
+	for range observed {
+	}
+
+	reloaded, err := store.GetSession(t.Context(), sess.ID)
+	require.NoError(t, err)
+	t.Cleanup(func() { require.NoError(t, store.Close()) })
+	var summaries int
+	for _, item := range reloaded.Messages {
+		if item.Summary != "" {
+			summaries++
+		}
+	}
+	assert.Equal(t, 1, summaries, "intrinsic persistence and the observer must not both append the summary")
+}
+
 // TestDoCompactNoHooksMatchesPriorBehavior is a regression guard: with
 // no compaction-related hooks configured, compactWithReason must still
 // emit the same SessionCompaction started/completed pair that all

pkg/session/session.go

@@ -914,15 +914,25 @@ func (s *Session) TitleSnapshot() string {
 	return s.Title
 }
 
-// ApplyCompaction atomically resets the session's cumulative token
-// counts and appends a summary item under s.mu so concurrent readers
-// (e.g. the persistence observer's UpdateSession snapshot) cannot
-// observe the new tokens without the matching summary item.
+// ApplyCompaction atomically resets the session's cumulative token counts,
+// derives scalar cost from canonical item history including the new summary,
+// and appends that summary under s.mu.
 func (s *Session) ApplyCompaction(inputTokens, outputTokens int64, item Item) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
+	s.applyCompactionLocked(inputTokens, outputTokens, s.totalCostLocked()+item.Cost, item)
+}
+
+func (s *Session) applyCompaction(inputTokens, outputTokens int64, resultingCost float64, item Item) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.applyCompactionLocked(inputTokens, outputTokens, resultingCost, item)
+}
+
+func (s *Session) applyCompactionLocked(inputTokens, outputTokens int64, resultingCost float64, item Item) {
 	s.InputTokens = inputTokens
 	s.OutputTokens = outputTokens
+	s.Cost = resultingCost
 	s.Messages = append(s.Messages, item)
 }
 
@@ -1666,7 +1676,10 @@ func (s *Session) MessagesSnapshot() []Item {
 func (s *Session) TotalCost() float64 {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
+	return s.totalCostLocked()
+}
 
+func (s *Session) totalCostLocked() float64 {
 	var cost float64
 	for _, item := range s.Messages {
 		switch {