fix(chatd): publish streaming message_part events during compaction (#22410)

## Problem

Context compaction in chatd persisted durable messages for the
`chat_summarized` tool call and result via `publishMessage`, but never
published `message_part` streaming events via `publishMessagePart`. This
meant connected clients had no streaming representation of the
compaction.

The client's `streamState` (built entirely from `message_part` events in
`streamState.ts`) never saw the compaction tool call, so:

- No **"Summarizing..."** running state was shown to the user during
summary generation (which can take up to 90s).
- The durable `message` events arrived after or interleaved with the
`status: waiting` event, causing the tool to appear as "Summarized" with
the chat appearing to just stop.

## Fix

### 1. `CompactionOptions.OnStart` callback (chatloop)

Added an `OnStart` callback to `CompactionOptions`, called in
`maybeCompact` right before `generateCompactionSummary` (the slow LLM
call). This gives `chatd` a hook to publish the tool-call `message_part`
immediately when compaction begins.

### 2. Tool-result streaming part (chatd)

`persistChatContextSummary` now publishes a tool-result `message_part`
before the durable `message` events, so clients transition from
"Summarizing..." to "Summarized" before the status change arrives.

### Event ordering is now:
1. `message_part` (tool call via `OnStart`) — client shows
"Summarizing..."
2. LLM generates summary (up to 90s)
3. `message_part` (tool result) — client shows "Summarized" in stream
state
4. `message` (assistant) — durable message persisted, stream state
resets
5. `message` (tool) — durable tool result persisted
6. `status: waiting` — chat transitions to idle

## Tests

- **`OnStartFiresBeforePersist`**: Verifies callback ordering is
`on_start` → `generate` → `persist`.
- **`OnStartNotCalledBelowThreshold`**: Verifies `OnStart` is not called
when context usage is below the compaction threshold.
This commit is contained in:
Kyle Carberry
2026-02-27 16:33:39 -05:00
committed by GitHub
parent 8bb80b060e
commit 360df1d84f
3 changed files with 141 additions and 1 deletions
+29 -1
View File
@@ -1968,6 +1968,13 @@ func (p *Server) runChat(
streamCall.MaxOutputTokens = &maxOutputTokens
}
// Generate the tool call ID up front so that the OnStart
// streaming part and the Persist durable messages share
// the same identifier. Without this the client cannot
// correlate the "Summarizing..." tool call with the
// "Summarized" tool result.
compactionToolCallID := "chat_summarized_" + uuid.NewString()
compactionOptions := &chatloop.CompactionOptions{
ThresholdPercent: modelConfig.CompressionThreshold,
ContextLimit: modelConfig.ContextLimit,
@@ -1979,6 +1986,7 @@ func (p *Server) runChat(
persistCtx,
chat.ID,
modelConfig.ID,
compactionToolCallID,
result,
); err != nil {
return xerrors.Errorf("persist context summary: %w", err)
@@ -1992,6 +2000,16 @@ func (p *Server) runChat(
)
return nil
},
OnStart: func() {
// Publish a streaming tool-call part immediately so
// connected clients see "Summarizing..." while the
// LLM generates the summary.
p.publishMessagePart(chat.ID, string(fantasy.MessageRoleAssistant), codersdk.ChatMessagePart{
Type: codersdk.ChatMessagePartTypeToolCall,
ToolCallID: compactionToolCallID,
ToolName: "chat_summarized",
})
},
OnError: func(err error) {
logger.Warn(ctx, "failed to compact chat context", slog.Error(err))
},
@@ -2063,6 +2081,7 @@ func (p *Server) persistChatContextSummary(
ctx context.Context,
chatID uuid.UUID,
modelConfigID uuid.UUID,
toolCallID string,
result chatloop.CompactionResult,
) error {
if strings.TrimSpace(result.SystemSummary) == "" ||
@@ -2097,7 +2116,6 @@ func (p *Server) persistChatContextSummary(
return xerrors.Errorf("insert hidden summary message: %w", err)
}
toolCallID := "chat_summarized_" + uuid.NewString()
args, err := json.Marshal(map[string]any{
"source": "automatic",
"threshold_percent": result.ThresholdPercent,
@@ -2182,6 +2200,16 @@ func (p *Server) persistChatContextSummary(
return xerrors.Errorf("insert summary tool result message: %w", err)
}
// Publish a streaming tool-result part so connected clients
// transition from "Summarizing..." to "Summarized" before the
// durable messages and status change arrive.
p.publishMessagePart(chatID, string(fantasy.MessageRoleTool), codersdk.ChatMessagePart{
Type: codersdk.ChatMessagePartTypeToolResult,
ToolCallID: toolCallID,
ToolName: "chat_summarized",
Result: summaryResult,
})
p.publishMessage(chatID, assistantMessage)
p.publishMessage(chatID, toolMessage)
return nil
+5
View File
@@ -30,6 +30,7 @@ type CompactionOptions struct {
SystemSummaryPrefix string
Timeout time.Duration
Persist func(context.Context, CompactionResult) error
OnStart func()
OnError func(error)
}
@@ -134,6 +135,10 @@ func maybeCompact(
return nil
}
if config.OnStart != nil {
config.OnStart()
}
summary, err := generateCompactionSummary(
ctx,
runOpts.Model,
+107
View File
@@ -83,6 +83,113 @@ func TestRun_Compaction(t *testing.T) {
require.InDelta(t, 80.0, persistedCompaction.UsagePercent, 0.0001)
})
t.Run("OnStartFiresBeforePersist", func(t *testing.T) {
t.Parallel()
const summaryText = "compaction summary for ordering test"
// Track the order of callbacks to verify OnStart fires
// before the Generate call (summary generation) and
// before Persist.
var callOrder []string
model := &loopTestModel{
provider: "fake",
streamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) {
return streamFromParts([]fantasy.StreamPart{
{Type: fantasy.StreamPartTypeTextStart, ID: "text-1"},
{Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "done"},
{Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"},
{
Type: fantasy.StreamPartTypeFinish,
FinishReason: fantasy.FinishReasonStop,
Usage: fantasy.Usage{
InputTokens: 80,
TotalTokens: 85,
},
},
}), nil
},
generateFn: func(_ context.Context, _ fantasy.Call) (*fantasy.Response, error) {
callOrder = append(callOrder, "generate")
return &fantasy.Response{
Content: []fantasy.Content{
fantasy.TextContent{Text: summaryText},
},
}, nil
},
}
_, err := Run(context.Background(), RunOptions{
Model: model,
Messages: []fantasy.Message{
textMessage(fantasy.MessageRoleUser, "hello"),
},
MaxSteps: 1,
PersistStep: func(_ context.Context, _ PersistedStep) error {
return nil
},
ContextLimitFallback: 100,
Compaction: &CompactionOptions{
ThresholdPercent: 70,
SummaryPrompt: "summarize now",
OnStart: func() {
callOrder = append(callOrder, "on_start")
},
Persist: func(_ context.Context, _ CompactionResult) error {
callOrder = append(callOrder, "persist")
return nil
},
},
})
require.NoError(t, err)
require.Equal(t, []string{"on_start", "generate", "persist"}, callOrder)
})
t.Run("OnStartNotCalledBelowThreshold", func(t *testing.T) {
t.Parallel()
onStartCalled := false
model := &loopTestModel{
provider: "fake",
streamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) {
return streamFromParts([]fantasy.StreamPart{
{
Type: fantasy.StreamPartTypeFinish,
FinishReason: fantasy.FinishReasonStop,
Usage: fantasy.Usage{
InputTokens: 10,
},
},
}), nil
},
}
_, err := Run(context.Background(), RunOptions{
Model: model,
Messages: []fantasy.Message{
textMessage(fantasy.MessageRoleUser, "hello"),
},
MaxSteps: 1,
PersistStep: func(_ context.Context, _ PersistedStep) error {
return nil
},
ContextLimitFallback: 100,
Compaction: &CompactionOptions{
ThresholdPercent: 70,
OnStart: func() {
onStartCalled = true
},
Persist: func(_ context.Context, _ CompactionResult) error {
return nil
},
},
})
require.NoError(t, err)
require.False(t, onStartCalled, "OnStart should not fire when usage is below threshold")
})
t.Run("ErrorsAreReported", func(t *testing.T) {
t.Parallel()