From 9373e57ebf5909dc73a5594f8944da72360b74f9 Mon Sep 17 00:00:00 2001 From: leookun Date: Mon, 10 Aug 2026 22:25:37 +0800 Subject: [PATCH] Enhance checkpoint handling and error management in forwarder - Added flushing of assistant text during provider completion to ensure no output is lost on transport failure. - Updated checkpoint blob synchronization tests to validate behavior under various conditions, including terminal and non-terminal states. - Introduced new functions for managing checkpoint terminal actions, improving clarity and maintainability of the code. - Implemented additional tests for imported blob handling and conversation state restoration, ensuring robustness in data integrity across operations. --- internal/backend/forwarder/actor.go | 3 + .../backend/forwarder/checkpoint_blobs.go | 103 ++++----- .../forwarder/checkpoint_blobs_test.go | 157 ++++++++++++- internal/backend/forwarder/compaction.go | 87 +++----- .../forwarder/compaction_append_only_test.go | 206 ++++++++++++++++++ internal/backend/forwarder/file_store.go | 14 +- internal/backend/forwarder/imported_blobs.go | 167 ++++++++++++++ .../forwarder/imported_history_test.go | 109 +++++++++ .../forwarder/interrupted_output_test.go | 48 ++++ internal/backend/forwarder/projector.go | 36 +-- internal/backend/forwarder/rewind.go | 21 ++ internal/backend/forwarder/runtime_summary.go | 3 +- internal/backend/forwarder/service.go | 39 +++- internal/backend/forwarder/token_usage.go | 57 +++-- internal/backend/forwarder/types.go | 24 +- 15 files changed, 884 insertions(+), 190 deletions(-) create mode 100644 internal/backend/forwarder/compaction_append_only_test.go create mode 100644 internal/backend/forwarder/imported_blobs.go create mode 100644 internal/backend/forwarder/imported_history_test.go diff --git a/internal/backend/forwarder/actor.go b/internal/backend/forwarder/actor.go index a712976..e30942b 100644 --- a/internal/backend/forwarder/actor.go +++ b/internal/backend/forwarder/actor.go @@ -730,6 +730,9 @@ func (service *Service) handleProviderDoneEvent(stream *ActiveStream, payload *s service.setTurnPhase(stream, TurnPhaseFailed) return service.closeStreamWithProviderError(stream, conversationID, turnSeq, requestID, accumulatedText, accumulatedReasoning, accumulatedReasoningSignature, accumulatedReasoningSignatureSource, accumulatedReasoningItemID, accumulatedReasoningStatus, accumulatedReasoningSummary, usage, providerErr, !hadToolInvocation) } + if err := service.flushAssistantText(stream, conversationID, turnSeq, requestID, accumulatedText, accumulatedReasoning, accumulatedReasoningSignature, accumulatedReasoningSignatureSource, accumulatedReasoningItemID, accumulatedReasoningStatus, accumulatedReasoningSummary, !hadToolInvocation); err != nil { + return service.failStream(stream, "unknown", fmt.Errorf("flush failed provider output: %w", err)) + } service.setTurnPhase(stream, TurnPhaseFailed) return service.failStream(stream, "unknown", payload.Err) } diff --git a/internal/backend/forwarder/checkpoint_blobs.go b/internal/backend/forwarder/checkpoint_blobs.go index edb65ed..9069962 100644 --- a/internal/backend/forwarder/checkpoint_blobs.go +++ b/internal/backend/forwarder/checkpoint_blobs.go @@ -19,15 +19,29 @@ type pendingCheckpointBlobWrite struct { blob CheckpointBlob } -func clonePendingTurnCompletion(completion *pendingTurnCompletion) *pendingTurnCompletion { +func successfulCheckpointTerminalAction(completion *pendingTurnCompletion) checkpointTerminalAction { if completion == nil { - return nil + return checkpointTerminalAction{} + } + return checkpointTerminalAction{ + Kind: checkpointTerminalActionComplete, + Completion: *completion, + } +} + +func failedCheckpointTerminalAction(errorCode string, errorMessage string) checkpointTerminalAction { + return checkpointTerminalAction{ + Kind: checkpointTerminalActionFail, + ErrorCode: strings.TrimSpace(errorCode), + ErrorMessage: strings.TrimSpace(errorMessage), } - cloned := *completion - return &cloned } func (service *Service) queueCheckpointProjection(stream *ActiveStream, projection *CheckpointProjection, completion *pendingTurnCompletion) error { + return service.queueCheckpointProjectionWithTerminal(stream, projection, successfulCheckpointTerminalAction(completion)) +} + +func (service *Service) queueCheckpointProjectionWithTerminal(stream *ActiveStream, projection *CheckpointProjection, terminal checkpointTerminalAction) error { if service == nil || stream == nil || projection == nil || projection.State == nil { return nil } @@ -43,8 +57,8 @@ func (service *Service) queueCheckpointProjection(stream *ActiveStream, projecti if stream.ConfirmedCheckpointBlobs == nil { stream.ConfirmedCheckpointBlobs = make(map[string]struct{}) } - if completion == nil && stream.PendingCheckpoint != nil { - completion = stream.PendingCheckpoint.Completion + if terminal.Kind == checkpointTerminalActionNone && stream.PendingCheckpoint != nil { + terminal = stream.PendingCheckpoint.Terminal } required := make(map[string]struct{}, len(projection.Blobs)) pendingKeys := make(map[string]struct{}, len(stream.PendingCheckpointBlobWrites)) @@ -74,11 +88,11 @@ func (service *Service) queueCheckpointProjection(stream *ActiveStream, projecti toWrite = append(toWrite, pendingCheckpointBlobWrite{requestID: requestID, blob: blob}) } stream.PendingCheckpoint = &pendingCheckpointPublish{ - State: state, - Required: required, - Completion: clonePendingTurnCompletion(completion), + State: state, + Required: required, + Terminal: terminal, } - if completion != nil { + if terminal.Kind != checkpointTerminalActionNone { stream.Phase = TurnPhaseCheckpointing } stream.UpdatedAt = time.Now().UTC() @@ -94,13 +108,8 @@ func (service *Service) queueCheckpointProjection(stream *ActiveStream, projecti if service.checkpointProjectionReady(stream) { return service.publishReadyCheckpoint(stream) } - // Keep the latest live UI state ahead of an immediate client abort. Blob writes are - // ordered before this snapshot; acknowledgements still gate terminal completion. - if completion == nil { - if err := service.publishPendingCheckpoint(stream); err != nil { - return service.finishAfterCheckpointSyncFailure(stream, fmt.Errorf("publish pending checkpoint: %w", err)) - } - } + // Checkpoints reference these Blob IDs, so the client must confirm every + // required Blob before the checkpoint becomes visible. service.scheduleStreamTimer( stream, providerTimerKey(streamTimerCheckpointBlobs, ""), @@ -113,31 +122,6 @@ func (service *Service) queueCheckpointProjection(stream *ActiveStream, projecti return nil } -func (service *Service) publishPendingCheckpoint(stream *ActiveStream) error { - if service == nil || stream == nil { - return nil - } - stream.mu.Lock() - pending := stream.PendingCheckpoint - if pending == nil || pending.Published { - stream.mu.Unlock() - return nil - } - pending.Published = true - state := pending.State - stream.UpdatedAt = time.Now().UTC() - stream.mu.Unlock() - if err := service.broker.Publish(stream.RequestID, StreamEvent{Message: buildCheckpointMessage(state)}); err != nil { - stream.mu.Lock() - if stream.PendingCheckpoint == pending { - pending.Published = false - } - stream.mu.Unlock() - return err - } - return nil -} - func (service *Service) checkpointProjectionReady(stream *ActiveStream) bool { if stream == nil { return false @@ -207,24 +191,18 @@ func (service *Service) publishReadyCheckpoint(stream *ActiveStream) error { } stream.PendingCheckpoint = nil state := pending.State - completion := clonePendingTurnCompletion(pending.Completion) - published := pending.Published + terminal := pending.Terminal stream.UpdatedAt = time.Now().UTC() stream.mu.Unlock() clearStreamTimer(stream, providerTimerKey(streamTimerCheckpointBlobs, "")) - if !published { - if err := service.broker.Publish(stream.RequestID, StreamEvent{Message: buildCheckpointMessage(state)}); err != nil { - if completion != nil { - log.Printf("forwarder checkpoint publish skipped before successful terminal request_id=%s err=%v", stream.RequestID, err) - return service.finishSuccessfulTurnAfterCheckpoint(stream, *completion) - } - return err + if err := service.broker.Publish(stream.RequestID, StreamEvent{Message: buildCheckpointMessage(state)}); err != nil { + if terminal.Kind != checkpointTerminalActionNone { + log.Printf("forwarder checkpoint publish skipped before terminal request_id=%s err=%v", stream.RequestID, err) + return service.finishCheckpointTerminalAction(stream, terminal) } + return err } - if completion != nil { - return service.finishSuccessfulTurnAfterCheckpoint(stream, *completion) - } - return nil + return service.finishCheckpointTerminalAction(stream, terminal) } func (service *Service) handleCheckpointBlobTimeout(stream *ActiveStream) error { @@ -251,12 +229,23 @@ func (service *Service) finishAfterCheckpointSyncFailure(stream *ActiveStream, c if cause != nil { log.Printf("forwarder checkpoint blob sync skipped request_id=%s conversation_id=%s err=%v", stream.RequestID, stream.ConversationID, cause) } - if pending != nil && pending.Completion != nil { - return service.finishSuccessfulTurnAfterCheckpoint(stream, *pending.Completion) + if pending != nil { + return service.finishCheckpointTerminalAction(stream, pending.Terminal) } return nil } +func (service *Service) finishCheckpointTerminalAction(stream *ActiveStream, terminal checkpointTerminalAction) error { + switch terminal.Kind { + case checkpointTerminalActionComplete: + return service.finishSuccessfulTurnAfterCheckpoint(stream, terminal.Completion) + case checkpointTerminalActionFail: + return service.finishFailedTurnAfterCheckpoint(stream, terminal.ErrorCode, terminal.ErrorMessage) + default: + return nil + } +} + func (service *Service) discardPendingCheckpoint(stream *ActiveStream, reason string) { if stream == nil { return diff --git a/internal/backend/forwarder/checkpoint_blobs_test.go b/internal/backend/forwarder/checkpoint_blobs_test.go index 7cf0f64..349aaec 100644 --- a/internal/backend/forwarder/checkpoint_blobs_test.go +++ b/internal/backend/forwarder/checkpoint_blobs_test.go @@ -8,22 +8,43 @@ import ( "cursor/gen/agentv1" ) -func TestCheckpointBlobSyncPublishesNonTerminalCheckpointBeforeAcknowledgements(t *testing.T) { +func TestCheckpointBlobSyncWaitsForAcknowledgementsBeforePublishingNonTerminalCheckpoint(t *testing.T) { service, stream, projection := testCheckpointBlobProjection(t) if err := service.queueCheckpointProjection(stream, projection, nil); err != nil { t.Fatalf("queueCheckpointProjection() error = %v", err) } events := readCheckpointTestEvents(t, service, stream) - if len(events) != len(projection.Blobs)+1 { - t.Fatalf("events before ACK = %d, want %d Blob writes and one checkpoint", len(events), len(projection.Blobs)) + if len(events) != len(projection.Blobs) { + t.Fatalf("events before ACK = %d, want %d Blob writes", len(events), len(projection.Blobs)) } - for _, event := range events[:len(projection.Blobs)] { + for _, event := range events { if event.Message.GetKvServerMessage().GetSetBlobArgs() == nil { t.Fatalf("event before ACK = %#v, want set_blob_args", event.Message) } } - if checkpoint := events[len(events)-1].Message.GetConversationCheckpointUpdate(); checkpoint == nil || len(checkpoint.GetTurns()) != 1 { - t.Fatalf("last event before ACK = %#v, want one Blob-backed turn", events[len(events)-1].Message) + + stream.mu.Lock() + var firstRequestID uint32 + for requestID := range stream.PendingCheckpointBlobWrites { + firstRequestID = requestID + break + } + stream.mu.Unlock() + if firstRequestID == 0 { + t.Fatal("checkpoint projection has no pending Blob writes") + } + if err := service.handleCheckpointBlobResult(stream, &agentv1.KvClientMessage{ + Id: firstRequestID, + Message: &agentv1.KvClientMessage_SetBlobResult{ + SetBlobResult: &agentv1.SetBlobResult{}, + }, + }); err != nil { + t.Fatalf("first Blob ACK error = %v", err) + } + for _, event := range readCheckpointTestEvents(t, service, stream) { + if event.Message.GetConversationCheckpointUpdate() != nil { + t.Fatal("checkpoint published after only a partial Blob acknowledgement") + } } acknowledgeCheckpointBlobs(t, service, stream) @@ -98,7 +119,123 @@ func TestCheckpointBlobTimeoutDoesNotFailSuccessfulTurn(t *testing.T) { } } -func TestCancellationKeepsPublishedCheckpointAndIgnoresLateAcknowledgements(t *testing.T) { +func TestCheckpointBlobSyncPublishesCheckpointBeforeFailedTerminal(t *testing.T) { + service, stream, _ := testCheckpointBlobProjection(t) + if err := service.failActiveStream( + stream, + stream.ConversationID, + stream.RequestID, + "model-call-1", + "provider_error", + "provider failed", + ); err != nil { + t.Fatalf("failActiveStream() error = %v", err) + } + + for _, event := range readCheckpointTestEvents(t, service, stream) { + if event.Message.GetConversationCheckpointUpdate() != nil || event.End { + t.Fatalf("event before ACK = %#v, want only Blob writes", event) + } + } + stream.mu.Lock() + phaseBeforeACK := stream.Phase + statusBeforeACK := stream.Status + stream.mu.Unlock() + if phaseBeforeACK != TurnPhaseCheckpointing || isTerminalStreamStatus(statusBeforeACK) { + t.Fatalf("before ACK phase=%s status=%s, want checkpointing and non-terminal", phaseBeforeACK, statusBeforeACK) + } + + acknowledgeCheckpointBlobs(t, service, stream) + events := readCheckpointTestEvents(t, service, stream) + checkpointIndex, endIndex := -1, -1 + for index, event := range events { + switch { + case event.Message.GetConversationCheckpointUpdate() != nil: + checkpointIndex = index + case event.End: + endIndex = index + if event.TerminalErrorCode != "provider_error" || event.TerminalErrorMessage != "provider failed" { + t.Fatalf("terminal event = %#v, want provider error", event) + } + } + } + if checkpointIndex < 0 || endIndex <= checkpointIndex { + t.Fatalf("terminal order checkpoint=%d end=%d", checkpointIndex, endIndex) + } + stream.mu.Lock() + phaseAfterACK := stream.Phase + statusAfterACK := stream.Status + stream.mu.Unlock() + if phaseAfterACK != TurnPhaseFailed || statusAfterACK != StreamStatusFailed { + t.Fatalf("after ACK phase=%s status=%s, want failed", phaseAfterACK, statusAfterACK) + } +} + +func TestCheckpointBlobTimeoutStillPublishesFailedTerminal(t *testing.T) { + service, stream, _ := testCheckpointBlobProjection(t) + if err := service.failActiveStream( + stream, + stream.ConversationID, + stream.RequestID, + "model-call-1", + "provider_error", + "provider failed", + ); err != nil { + t.Fatalf("failActiveStream() error = %v", err) + } + if err := service.handleCheckpointBlobTimeout(stream); err != nil { + t.Fatalf("handleCheckpointBlobTimeout() error = %v", err) + } + + events := readCheckpointTestEvents(t, service, stream) + var checkpoint, failedEnd bool + for _, event := range events { + checkpoint = checkpoint || event.Message.GetConversationCheckpointUpdate() != nil + failedEnd = failedEnd || event.End && event.TerminalErrorCode == "provider_error" && event.TerminalErrorMessage == "provider failed" + } + if checkpoint || !failedEnd { + t.Fatalf("timeout events checkpoint=%v failed_end=%v", checkpoint, failedEnd) + } +} + +func TestManualCompactionNoopWaitsForCheckpointBeforeTerminal(t *testing.T) { + service, stream, _ := testCheckpointBlobProjection(t) + conversation, _, _, err := service.snapshotCheckpointConversation(stream) + if err != nil { + t.Fatalf("snapshotCheckpointConversation() error = %v", err) + } + if _, err := service.store.SaveConversationWithEntries(stream.ConversationID, conversation, conversation.Entries); err != nil { + t.Fatalf("SaveConversationWithEntries() error = %v", err) + } + if err := service.finishManualCompactionNoop(stream); err != nil { + t.Fatalf("finishManualCompactionNoop() error = %v", err) + } + + for _, event := range readCheckpointTestEvents(t, service, stream) { + if event.Message.GetInteractionUpdate().GetTurnEnded() != nil || event.End { + t.Fatalf("terminal event before checkpoint Blob ACK = %#v", event) + } + } + acknowledgeCheckpointBlobs(t, service, stream) + + events := readCheckpointTestEvents(t, service, stream) + checkpointIndex, turnEndedIndex, endIndex := -1, -1, -1 + for index, event := range events { + switch { + case event.Message.GetConversationCheckpointUpdate() != nil: + checkpointIndex = index + case event.Message.GetInteractionUpdate().GetTurnEnded() != nil: + turnEndedIndex = index + case event.End: + endIndex = index + } + } + if checkpointIndex < 0 || turnEndedIndex <= checkpointIndex || endIndex <= turnEndedIndex { + t.Fatalf("terminal order checkpoint=%d turn_ended=%d end=%d", checkpointIndex, turnEndedIndex, endIndex) + } +} + +func TestCancellationDiscardsUnpublishedCheckpointAndIgnoresLateAcknowledgements(t *testing.T) { service, stream, projection := testCheckpointBlobProjection(t) if err := service.queueCheckpointProjection(stream, projection, nil); err != nil { t.Fatalf("queueCheckpointProjection() error = %v", err) @@ -110,8 +247,8 @@ func TestCancellationKeepsPublishedCheckpointAndIgnoresLateAcknowledgements(t *t checkpointBeforeCancel++ } } - if checkpointBeforeCancel != 1 { - t.Fatalf("checkpoints before cancel = %d, want 1", checkpointBeforeCancel) + if checkpointBeforeCancel != 0 { + t.Fatalf("checkpoints before cancel = %d, want 0", checkpointBeforeCancel) } stream.mu.Lock() requestIDs := make([]uint32, 0, len(stream.PendingCheckpointBlobWrites)) @@ -149,7 +286,7 @@ func TestCancellationKeepsPublishedCheckpointAndIgnoresLateAcknowledgements(t *t stream.mu.Lock() pending := stream.PendingCheckpoint stream.mu.Unlock() - if checkpointCount != 1 || !canceledEnd || pending != nil { + if checkpointCount != 0 || !canceledEnd || pending != nil { t.Fatalf("cancel events checkpoints=%d canceled_end=%v pending=%v", checkpointCount, canceledEnd, pending != nil) } } diff --git a/internal/backend/forwarder/compaction.go b/internal/backend/forwarder/compaction.go index f28220c..fc70caa 100644 --- a/internal/backend/forwarder/compaction.go +++ b/internal/backend/forwarder/compaction.go @@ -234,7 +234,7 @@ func (service *Service) buildLegacyCompactionPlan(base *compactionPlan, conversa if conversation == nil || base == nil { return nil, nil } - candidates := buildContextCompactionCandidates(checkpointProjectionEntries(conversation.Entries), base.CurrentTurnSeq, base.CurrentRequestID) + candidates := buildContextCompactionCandidates(replayablePromptProjectionEntries(conversation.Entries), base.CurrentTurnSeq, base.CurrentRequestID) if len(candidates) == 0 { return nil, nil } @@ -260,7 +260,7 @@ func (service *Service) buildAutoCompactionPlanFromHistory(base *compactionPlan, if err != nil { return nil, err } - currentCandidate, hasCurrentCandidate := buildCurrentTurnCompactionCandidate(checkpointProjectionEntries(conversation.Entries), base.CurrentTurnSeq, base.CurrentRequestID) + currentCandidate, hasCurrentCandidate := buildCurrentTurnCompactionCandidate(replayablePromptProjectionEntries(conversation.Entries), base.CurrentTurnSeq, base.CurrentRequestID) if !hasCurrentCandidate { return legacyPlan, nil } @@ -447,16 +447,8 @@ func (service *Service) handleCompactionEvent(stream *ActiveStream, payload *str if err := service.completeManualCompactionTurn(stream); err != nil { return service.failStream(stream, "unknown", err) } - if err := service.broker.Publish(stream.RequestID, StreamEvent{ - Message: buildTurnEndedMessage(0, 0, 0, 0), - }); err != nil { - return service.failStream(stream, "unknown", err) - } - if err := service.broker.Complete(stream.RequestID, "", ""); err != nil { - return service.failStream(stream, "unknown", err) - } - service.setTurnPhase(stream, TurnPhaseCompleted) - return nil + completion := manualCompactionTurnCompletion(stream) + return service.publishCheckpointWithCompletion(stream.RequestID, stream.ConversationID, &completion) } return service.requestProviderAction(stream, providerActionResume) } @@ -500,12 +492,8 @@ func (service *Service) finishManualCompactionNoop(stream *ActiveStream) error { if err := service.completeManualCompactionTurn(stream); err != nil { return err } - if err := service.broker.Publish(stream.RequestID, StreamEvent{ - Message: buildTurnEndedMessage(0, 0, 0, 0), - }); err != nil { - return err - } - return service.broker.Complete(stream.RequestID, "", "") + completion := manualCompactionTurnCompletion(stream) + return service.publishCheckpointWithCompletion(stream.RequestID, stream.ConversationID, &completion) } func (service *Service) completeManualCompactionTurn(stream *ActiveStream) error { @@ -530,10 +518,21 @@ func (service *Service) completeManualCompactionTurn(stream *ActiveStream) error if err := service.syncSummaryCarryForward(conversationID, requestID, modelCallID); err != nil { return err } - service.setTurnPhase(stream, TurnPhaseCompleted) return nil } +func manualCompactionTurnCompletion(stream *ActiveStream) pendingTurnCompletion { + if stream == nil { + return pendingTurnCompletion{} + } + return pendingTurnCompletion{ + ConversationID: strings.TrimSpace(stream.ConversationID), + RequestID: strings.TrimSpace(stream.RequestID), + TurnSeq: stream.TurnSeq, + ModelCallID: "turn:" + strings.TrimSpace(stream.RequestID), + } +} + func (service *Service) publishSummaryCompleted(stream *ActiveStream, hookMessage string) error { if service == nil || stream == nil { return nil @@ -568,6 +567,7 @@ func (service *Service) applyCompactionPlan(stream *ActiveStream, conversationID if err != nil { return err } + originalEntryCount := len(candidateConversation.Entries) if err := applyCompactionToConversation(candidateConversation, plan, summaryText); err != nil { return err } @@ -582,9 +582,9 @@ func (service *Service) applyCompactionPlan(stream *ActiveStream, conversationID if validationErr := validateCompactionCandidateBudget(recompiled, plan); validationErr != nil { return validationErr } - replacementEntries := append([]HistoryEntry(nil), candidateConversation.Entries...) + compactionEntries := append([]HistoryEntry(nil), candidateConversation.Entries[originalEntryCount:]...) if service.store != nil { - persisted, err := service.store.ReplaceEntries(conversationID, replacementEntries, func(item *ConversationFile) error { + persisted, _, err := service.store.AppendEntriesWithUpdate(conversationID, resetEntrySequences(compactionEntries), func(item *ConversationFile) error { if item == nil { return nil } @@ -605,10 +605,7 @@ func (service *Service) applyCompactionPlan(stream *ActiveStream, conversationID if item == nil { return nil } - item.Entries = nil - item.NextEntrySeq = 1 - item.NextTurnSeq = 1 - appendEntriesInPlace(item, resetEntrySequences(replacementEntries)) + appendEntriesInPlace(item, resetEntrySequences(compactionEntries)) item.TokenDetailsUsedTokens = 0 clearConversationAutoCompactionState(item) return nil @@ -643,14 +640,13 @@ func applyCompactionToConversation(conversation *ConversationFile, plan *Pending if conversation == nil || plan == nil { return nil } - replacementEntries, err := buildCompactedContextEntries(conversation, plan, summaryText) + compactionEntries, err := buildCompactedContextEntries(conversation, plan, summaryText) if err != nil { return err } - conversation.Entries = nil - conversation.NextEntrySeq = 1 - conversation.NextTurnSeq = 1 - appendEntriesInPlace(conversation, resetEntrySequences(replacementEntries)) + // Canonical history stays append-only. The prompt projector applies the + // latest summary marker when constructing model-visible replay. + appendEntriesInPlace(conversation, resetEntrySequences(compactionEntries)) conversation.TokenDetailsUsedTokens = 0 clearConversationAutoCompactionState(conversation) if conversation.TokenDetailsMaxTokens == 0 { @@ -671,40 +667,9 @@ func buildCompactedContextEntries(conversation *ConversationFile, plan *PendingC if ok { entries = append(entries, runtimeEntry) } - if conversation == nil || !plan.PreserveCurrentTurnInputs { - return entries, nil - } - entries = append(entries, buildAutoCompactionPreservedCurrentTurnEntries(conversation.Entries, plan)...) return entries, nil } -func buildAutoCompactionPreservedCurrentTurnEntries(entries []HistoryEntry, plan *PendingCompaction) []HistoryEntry { - if len(entries) == 0 || plan == nil || !plan.PreserveCurrentTurnInputs { - return nil - } - latestToolCallID := latestCompletedToolCallIDForTurn(entries, plan.CurrentTurnSeq, plan.CurrentRequestID) - preservedIndexes := autoCompactionPreservedEntryIndexes(entries, plan.CurrentTurnSeq, plan.CurrentRequestID, latestToolCallID) - if len(preservedIndexes) == 0 { - return nil - } - preserved := make([]HistoryEntry, 0, len(preservedIndexes)) - for index, entry := range entries { - if _, ok := preservedIndexes[index]; !ok { - continue - } - switch strings.TrimSpace(entry.Kind) { - case "compaction_summary", "compacted_summary", "compaction_request": - continue - case "tool_result": - if rewritten, ok := rewriteAutoCompactionToolResultEntry(entry, autoCompactionPreservedToolResultLimitBytes, false); ok { - entry = rewritten - } - } - preserved = append(preserved, entry) - } - return preserved -} - func newCompactionSummaryEntry(plan *PendingCompaction, summaryText string) HistoryEntry { payload, _ := json.Marshal(compactionSummaryEntryPayload{ Summary: strings.TrimSpace(summaryText), diff --git a/internal/backend/forwarder/compaction_append_only_test.go b/internal/backend/forwarder/compaction_append_only_test.go new file mode 100644 index 0000000..2edec32 --- /dev/null +++ b/internal/backend/forwarder/compaction_append_only_test.go @@ -0,0 +1,206 @@ +package forwarder + +import ( + "reflect" + "strings" + "testing" + + "google.golang.org/protobuf/encoding/protojson" + + "cursor/gen/agentv1" +) + +func TestApplyCompactionToConversationPreservesCanonicalHistory(t *testing.T) { + conversation := compactionAppendOnlyConversation(t) + originalEntries := append([]HistoryEntry(nil), conversation.Entries...) + plan := &PendingCompaction{ + Trigger: "manual", + CurrentTurnSeq: 2, + CurrentRequestID: "request-2", + } + + if err := applyCompactionToConversation(conversation, plan, "earlier context summary"); err != nil { + t.Fatalf("applyCompactionToConversation() error = %v", err) + } + if len(conversation.Entries) <= len(originalEntries) { + t.Fatalf("entries after compaction = %d, want the %d original entries plus a summary marker", len(conversation.Entries), len(originalEntries)) + } + if !reflect.DeepEqual(conversation.Entries[:len(originalEntries)], originalEntries) { + t.Fatal("compaction changed the canonical history prefix") + } + + projector := NewHistoryProjector() + projection, err := projector.ProjectCheckpointProjection(conversation) + if err != nil { + t.Fatalf("ProjectCheckpointProjection() error = %v", err) + } + if len(projection.State.GetTurns()) != 2 { + t.Fatalf("checkpoint turns after compaction = %d, want 2 visible turns", len(projection.State.GetTurns())) + } + replay, err := projector.ProjectPromptReplay(conversation) + if err != nil { + t.Fatalf("ProjectPromptReplay() error = %v", err) + } + if len(replay) != 1 || replay[0].Role != "user" || !strings.Contains(replay[0].Content, "earlier context summary") { + t.Fatalf("prompt replay after compaction = %#v, want only the compacted summary", replay) + } +} + +func TestCompactedPromptProjectionPlacesSummaryBeforePreservedCurrentTurn(t *testing.T) { + conversation := &ConversationFile{ + ConversationID: "conversation-1", + RootConversationID: "conversation-1", + Mode: "agent", + NextTurnSeq: 1, + NextEntrySeq: 1, + } + appendEntriesInPlace(conversation, []HistoryEntry{ + compactionTestUserEntry(t, 1, "request-1", "current question", "message-1"), + newToolCallEntry(1, "request-1", "call-1", "Read", "", "", checkpointTestReadToolCall(t, nil)), + newToolResultEntry(1, "request-1", "call-1", "Read", `{"path":"/tmp/example.txt"}`, "file contents", "", checkpointTestReadToolCall(t, nil)), + }) + plan := &PendingCompaction{ + Trigger: "auto", + CurrentTurnSeq: 1, + CurrentRequestID: "request-1", + PreserveCurrentTurnInputs: true, + } + if err := applyCompactionToConversation(conversation, plan, "current progress summary"); err != nil { + t.Fatalf("applyCompactionToConversation() error = %v", err) + } + + projected := compactedPromptProjectionEntries(conversation.Entries) + promptKinds := make([]string, 0, len(projected)) + for _, entry := range projected { + if isPromptReplayEntryKind(entry.Kind) { + promptKinds = append(promptKinds, entry.Kind) + } + } + want := []string{"compacted_summary", "user_message", "tool_call", "tool_result"} + if !reflect.DeepEqual(promptKinds, want) { + t.Fatalf("compacted prompt entry order = %#v, want %#v", promptKinds, want) + } +} + +func TestCompactionPlanningDoesNotRecompactArchivedHistory(t *testing.T) { + conversation := compactionAppendOnlyConversation(t) + if err := applyCompactionToConversation(conversation, &PendingCompaction{ + Trigger: "manual", + CurrentTurnSeq: 2, + CurrentRequestID: "request-2", + }, "archived history summary"); err != nil { + t.Fatalf("applyCompactionToConversation() error = %v", err) + } + appendEntriesInPlace(conversation, []HistoryEntry{ + compactionTestUserEntry(t, 3, "request-3", "new question", "message-3"), + }) + + plan, err := (&Service{}).buildLegacyCompactionPlan(&compactionPlan{ + CurrentTurnSeq: 3, + CurrentRequestID: "request-3", + }, conversation, false, 0) + if err != nil { + t.Fatalf("buildLegacyCompactionPlan() error = %v", err) + } + if plan != nil { + t.Fatalf("buildLegacyCompactionPlan() = %#v, want no already summarized candidates", plan) + } +} + +func TestApplyCompactionPlanPersistsHistoryAppendOnly(t *testing.T) { + store := NewConversationFileStore(t.TempDir()) + conversation := compactionAppendOnlyConversation(t) + if _, _, err := store.AppendEntries(conversation.ConversationID, resetEntrySequences(conversation.Entries)); err != nil { + t.Fatalf("AppendEntries() error = %v", err) + } + persisted, err := store.LoadConversation(conversation.ConversationID) + if err != nil { + t.Fatalf("initial LoadConversation() error = %v", err) + } + originalEntries := append([]HistoryEntry(nil), persisted.Entries...) + projector := NewHistoryProjector() + service := &Service{ + store: store, + projector: projector, + compiler: compactionProjectionCompiler{projector: projector}, + } + stream := &ActiveStream{ + RequestID: "request-2", + ConversationID: conversation.ConversationID, + TurnSeq: 2, + Mode: agentv1.AgentMode_AGENT_MODE_AGENT, + CheckpointConversation: persisted, + } + plan := &PendingCompaction{ + Trigger: "manual", + CurrentTurnSeq: 2, + CurrentRequestID: "request-2", + ContextWindowSize: 1_000_000, + } + if err := service.applyCompactionPlan(stream, conversation.ConversationID, plan, "persisted summary"); err != nil { + t.Fatalf("applyCompactionPlan() error = %v", err) + } + + loaded, err := store.LoadConversation(conversation.ConversationID) + if err != nil { + t.Fatalf("LoadConversation() error = %v", err) + } + if len(loaded.Entries) <= len(originalEntries) { + t.Fatalf("persisted entries after compaction = %d, want more than %d", len(loaded.Entries), len(originalEntries)) + } + for index := range originalEntries { + if !reflect.DeepEqual(loaded.Entries[index], originalEntries[index]) { + t.Fatalf("persisted history entry %d changed after compaction:\ngot %#v\nwant %#v", index, loaded.Entries[index], originalEntries[index]) + } + } +} + +type compactionProjectionCompiler struct { + projector *HistoryProjector +} + +func (compiler compactionProjectionCompiler) Compile(conversation *ConversationFile, _ agentv1.AgentMode, _ string, _ string) (CompiledConversation, error) { + messages, err := compiler.projector.ProjectPromptReplay(conversation) + return CompiledConversation{Messages: messages}, err +} + +func (compactionProjectionCompiler) DerivePromptContexts(*ConversationFile, agentv1.AgentMode, string) ([]PromptContextMessage, error) { + return nil, nil +} + +func compactionAppendOnlyConversation(t *testing.T) *ConversationFile { + t.Helper() + conversation := &ConversationFile{ + ConversationID: "conversation-1", + RootConversationID: "conversation-1", + Mode: "agent", + NextTurnSeq: 1, + NextEntrySeq: 1, + TokenDetailsUsedTokens: 42_000, + TokenDetailsMaxTokens: 50_000, + } + appendEntriesInPlace(conversation, []HistoryEntry{ + compactionTestUserEntry(t, 1, "request-1", "first question", "message-1"), + newAssistantTextEntry(1, "request-1", "first answer", "", ""), + compactionTestUserEntry(t, 2, "request-2", "second question", "message-2"), + newAssistantTextEntry(2, "request-2", "second answer", "", ""), + }) + return conversation +} + +func compactionTestUserEntry(t *testing.T, turnSeq int64, requestID string, text string, messageID string) HistoryEntry { + t.Helper() + payload, err := protojson.Marshal(&agentv1.UserMessage{Text: text, MessageId: messageID}) + if err != nil { + t.Fatalf("marshal user message: %v", err) + } + return HistoryEntry{ + TurnSeq: turnSeq, + RequestID: requestID, + Role: "user", + Kind: "user_message", + Payload: payload, + } +} + +var _ PromptCompiler = compactionProjectionCompiler{} diff --git a/internal/backend/forwarder/file_store.go b/internal/backend/forwarder/file_store.go index 785e8fd..dff749d 100644 --- a/internal/backend/forwarder/file_store.go +++ b/internal/backend/forwarder/file_store.go @@ -121,10 +121,15 @@ func (store *ConversationFileStore) LoadConversation(conversationID string) (*Co // AppendEntries 把已经发生的语义事件追加到 context.json，并同步 state.json。 func (store *ConversationFileStore) AppendEntries(conversationID string, entries []HistoryEntry) (*ConversationFile, []HistoryEntry, error) { + return store.AppendEntriesWithUpdate(conversationID, entries, nil) +} + +// AppendEntriesWithUpdate 原子追加 context entries，并在同一把会话锁内更新 state metadata。 +func (store *ConversationFileStore) AppendEntriesWithUpdate(conversationID string, entries []HistoryEntry, update func(*ConversationFile) error) (*ConversationFile, []HistoryEntry, error) { if store == nil { return nil, nil, fmt.Errorf("conversation file store is nil") } - if len(entries) == 0 { + if len(entries) == 0 && update == nil { conversation, err := store.LoadConversation(conversationID) return conversation, nil, err } @@ -162,6 +167,11 @@ func (store *ConversationFileStore) AppendEntries(conversationID string, entries conversation.Mode = alias } assigned := appendEntriesInPlace(conversation, entries) + if update != nil { + if err := update(conversation); err != nil { + return nil, nil, err + } + } deriveConversationLoopState(conversation) if err := store.writeConversationLocked(normalizedConversationID, conversation); err != nil { return nil, nil, err @@ -762,6 +772,7 @@ func mergeConversationMetadata(target *ConversationFile, source *ConversationFil target.CurrentPlanText = source.CurrentPlanText target.CurrentPlans = clonePlanRegistryEntries(source.CurrentPlans) target.CurrentTodos = cloneTodoItems(source.CurrentTodos) + target.ImportedTurnIDs = cloneByteSlices(source.ImportedTurnIDs) target.LatestRequestPrefix = cloneConversationRequestPrefix(source.LatestRequestPrefix) target.LastProviderCall = cloneConversationProviderCall(source.LastProviderCall) if !source.CreatedAt.IsZero() && (target.CreatedAt.IsZero() || source.CreatedAt.Before(target.CreatedAt)) { @@ -894,6 +905,7 @@ func cloneConversationFile(conversation *ConversationFile) *ConversationFile { cloned := *conversation cloned.CurrentPlans = clonePlanRegistryEntries(conversation.CurrentPlans) cloned.CurrentTodos = cloneTodoItems(conversation.CurrentTodos) + cloned.ImportedTurnIDs = cloneByteSlices(conversation.ImportedTurnIDs) cloned.LatestRequestPrefix = cloneConversationRequestPrefix(conversation.LatestRequestPrefix) cloned.LastProviderCall = cloneConversationProviderCall(conversation.LastProviderCall) cloned.Entries = append([]HistoryEntry(nil), conversation.Entries...) diff --git a/internal/backend/forwarder/imported_blobs.go b/internal/backend/forwarder/imported_blobs.go new file mode 100644 index 0000000..fbd17ed --- /dev/null +++ b/internal/backend/forwarder/imported_blobs.go @@ -0,0 +1,167 @@ +package forwarder + +import ( + "crypto/sha256" + "fmt" + + "google.golang.org/protobuf/proto" + + "cursor/gen/agentv1" + modeladapter "cursor/internal/backend/agent/model" + promptengine "cursor/internal/backend/agent/prompt" +) + +type importedBlobStore map[string][]byte + +func newImportedBlobStore(items []*agentv1.PreFetchedBlob) (importedBlobStore, error) { + if len(items) == 0 { + return nil, nil + } + store := make(importedBlobStore, len(items)) + for _, item := range items { + if item == nil || len(item.GetId()) == 0 { + continue + } + if len(item.GetId()) != sha256.Size { + return nil, fmt.Errorf("prefetched blob id length %d, want %d", len(item.GetId()), sha256.Size) + } + digest := sha256.Sum256(item.GetValue()) + if string(digest[:]) != string(item.GetId()) { + return nil, fmt.Errorf("prefetched blob %x failed SHA-256 validation", item.GetId()) + } + store[string(item.GetId())] = append([]byte(nil), item.GetValue()...) + } + return store, nil +} + +func (store importedBlobStore) resolve(id []byte) ([]byte, bool) { + if len(id) == 0 || len(store) == 0 { + return nil, false + } + value, ok := store[string(id)] + return append([]byte(nil), value...), ok +} + +func decodeImportedTurn(raw []byte, blobs importedBlobStore) (*agentv1.ConversationTurnStructure, []byte, error) { + if data, ok := blobs.resolve(raw); ok { + turn := &agentv1.ConversationTurnStructure{} + if err := proto.Unmarshal(data, turn); err != nil || turn.GetTurn() == nil { + return nil, nil, fmt.Errorf("decode imported turn blob %x: %w", raw, firstNonNilError(err, fmt.Errorf("turn payload is empty"))) + } + return turn, append([]byte(nil), raw...), nil + } + turn := &agentv1.ConversationTurnStructure{} + if err := proto.Unmarshal(raw, turn); err == nil && turn.GetTurn() != nil { + return turn, nil, nil + } + if len(raw) == sha256.Size { + return nil, append([]byte(nil), raw...), nil + } + return nil, nil, fmt.Errorf("decode imported inline turn") +} + +func decodeImportedUserMessage(raw []byte, blobs importedBlobStore) (*agentv1.UserMessage, error) { + data := raw + if resolved, ok := blobs.resolve(raw); ok { + data = resolved + } else if len(raw) == sha256.Size { + candidate := &agentv1.UserMessage{} + if err := proto.Unmarshal(raw, candidate); err != nil || !hasKnownUserMessageContent(candidate) { + return nil, fmt.Errorf("missing prefetched user message blob %x", raw) + } + return candidate, nil + } + message := &agentv1.UserMessage{} + if err := proto.Unmarshal(data, message); err != nil { + return nil, fmt.Errorf("decode imported turn user_message: %w", err) + } + return message, nil +} + +func decodeImportedStep(raw []byte, blobs importedBlobStore) (*agentv1.ConversationStep, error) { + data := raw + if resolved, ok := blobs.resolve(raw); ok { + data = resolved + } else if len(raw) == sha256.Size { + candidate := &agentv1.ConversationStep{} + if err := proto.Unmarshal(raw, candidate); err != nil || candidate.GetMessage() == nil { + return nil, fmt.Errorf("missing prefetched conversation step blob %x", raw) + } + return candidate, nil + } + step := &agentv1.ConversationStep{} + if err := proto.Unmarshal(data, step); err != nil { + return nil, fmt.Errorf("decode imported turn step: %w", err) + } + if step.GetMessage() == nil { + return nil, fmt.Errorf("decode imported turn step: payload is empty") + } + return step, nil +} + +func importedBlobTurnMessages(turn *agentv1.ConversationTurnStructure, blobs importedBlobStore) ([]modeladapter.Message, error) { + if turn == nil || turn.GetAgentConversationTurn() == nil { + return nil, nil + } + agentTurn := turn.GetAgentConversationTurn() + messages := make([]modeladapter.Message, 0, 1+len(agentTurn.GetSteps())) + if len(agentTurn.GetUserMessage()) > 0 { + userMessage, err := decodeImportedUserMessage(agentTurn.GetUserMessage(), blobs) + if err != nil { + return nil, err + } + if replay, ok := promptengine.BuildUserMessageReplayMessage(userMessage); ok { + messages = append(messages, toModelMessage(replay)) + } + } + for _, rawStep := range agentTurn.GetSteps() { + if len(rawStep) == 0 { + continue + } + step, err := decodeImportedStep(rawStep, blobs) + if err != nil { + return nil, err + } + for _, replay := range promptengine.BuildLegacyMessagesFromConversationStep(step) { + messages = append(messages, toModelMessage(replay)) + } + } + return messages, nil +} + +func importedTurnIDs(turns [][]byte, blobs importedBlobStore) ([][]byte, error) { + ids := make([][]byte, 0, len(turns)) + for _, raw := range turns { + if len(raw) == 0 { + continue + } + _, id, err := decodeImportedTurn(raw, blobs) + if err != nil { + return nil, err + } + if len(id) > 0 { + ids = append(ids, id) + } + } + return ids, nil +} + +func hasKnownUserMessageContent(message *agentv1.UserMessage) bool { + if message == nil { + return false + } + return message.GetText() != "" || + message.GetMessageId() != "" || + message.GetSelectedContext() != nil || + message.GetRichText() != "" || + len(message.GetConversationStateBlobId()) > 0 || + len(message.GetTextBlobId()) > 0 || + len(message.GetRichTextBlobId()) > 0 +} + +func firstNonNilError(err error, fallback error) error { + if err != nil { + return err + } + return fallback +} diff --git a/internal/backend/forwarder/imported_history_test.go b/internal/backend/forwarder/imported_history_test.go new file mode 100644 index 0000000..cbb7c60 --- /dev/null +++ b/internal/backend/forwarder/imported_history_test.go @@ -0,0 +1,109 @@ +package forwarder + +import ( + "crypto/sha256" + "testing" + + "google.golang.org/protobuf/proto" + + "cursor/gen/agentv1" +) + +func TestImportedConversationStateRestoresBlobOnlyForkAndCheckpointPrefix(t *testing.T) { + parent := compactionAppendOnlyConversation(t) + parent.Entries = parent.Entries[:2] + parent.NextEntrySeq = 3 + parent.NextTurnSeq = 2 + projection, err := NewHistoryProjector().ProjectCheckpointProjection(parent) + if err != nil { + t.Fatalf("ProjectCheckpointProjection() error = %v", err) + } + prefetched := make([]*agentv1.PreFetchedBlob, 0, len(projection.Blobs)) + for _, blob := range projection.Blobs { + prefetched = append(prefetched, &agentv1.PreFetchedBlob{Id: blob.ID, Value: blob.Data}) + } + state := proto.Clone(projection.State).(*agentv1.ConversationStateStructure) + state.RootPromptMessagesJson = nil + conversation, err := newRuntimeConversation("fork-conversation", agentv1.AgentMode_AGENT_MODE_AGENT) + if err != nil { + t.Fatalf("newRuntimeConversation() error = %v", err) + } + entries, err := (&Service{}).importConversationState(conversation, state, prefetched) + if err != nil { + t.Fatalf("importConversationState() error = %v", err) + } + if len(conversation.ImportedTurnIDs) != 1 || conversation.NextTurnSeq != 2 { + t.Fatalf("imported prefix turns=%d next_turn_seq=%d, want 1 and 2", len(conversation.ImportedTurnIDs), conversation.NextTurnSeq) + } + if len(entries) != 2 { + t.Fatalf("imported model entries = %d, want parent user and assistant", len(entries)) + } + appendEntriesInPlace(conversation, append(entries, + compactionTestUserEntry(t, 2, "request-2", "fork question", "message-2"), + )) + forkProjection, err := NewHistoryProjector().ProjectCheckpointProjection(conversation) + if err != nil { + t.Fatalf("fork ProjectCheckpointProjection() error = %v", err) + } + if len(forkProjection.State.GetTurns()) != 2 { + t.Fatalf("fork checkpoint turns = %d, want imported parent plus local fork turn", len(forkProjection.State.GetTurns())) + } + if string(forkProjection.State.GetTurns()[0]) != string(projection.State.GetTurns()[0]) { + t.Fatal("fork checkpoint did not preserve the imported parent turn ID as its prefix") + } +} + +func TestImportedConversationStateRejectsUnresolvedBlobTurn(t *testing.T) { + turnID := sha256.Sum256([]byte("missing imported turn")) + conversation, err := newRuntimeConversation("fork-conversation", agentv1.AgentMode_AGENT_MODE_AGENT) + if err != nil { + t.Fatalf("newRuntimeConversation() error = %v", err) + } + if _, err := (&Service{}).importConversationState(conversation, &agentv1.ConversationStateStructure{ + Turns: [][]byte{turnID[:]}, + }, nil); err == nil { + t.Fatal("importConversationState() accepted an unresolved Blob turn") + } +} + +func TestImportedTurnIDsPersistThroughConversationStore(t *testing.T) { + store := NewConversationFileStore(t.TempDir()) + turnID := sha256.Sum256([]byte("parent turn")) + conversation, err := newRuntimeConversation("fork-conversation", agentv1.AgentMode_AGENT_MODE_AGENT) + if err != nil { + t.Fatalf("newRuntimeConversation() error = %v", err) + } + conversation.ImportedTurnIDs = [][]byte{turnID[:]} + persisted, err := store.SaveConversationWithEntries(conversation.ConversationID, conversation, []HistoryEntry{ + compactionTestUserEntry(t, 2, "request-2", "fork question", "message-2"), + }) + if err != nil { + t.Fatalf("SaveConversationWithEntries() error = %v", err) + } + if len(persisted.ImportedTurnIDs) != 1 || string(persisted.ImportedTurnIDs[0]) != string(turnID[:]) { + t.Fatalf("persisted ImportedTurnIDs = %x, want %x", persisted.ImportedTurnIDs, turnID) + } + loaded, err := store.LoadConversation(conversation.ConversationID) + if err != nil { + t.Fatalf("LoadConversation() error = %v", err) + } + if len(loaded.ImportedTurnIDs) != 1 || string(loaded.ImportedTurnIDs[0]) != string(turnID[:]) { + t.Fatalf("loaded ImportedTurnIDs = %x, want %x", loaded.ImportedTurnIDs, turnID) + } +} + +func TestRewindImportedTurnPrefixUsesClientForkPoint(t *testing.T) { + ids := make([][]byte, 3) + for index := range ids { + digest := sha256.Sum256([]byte{byte(index + 1)}) + ids[index] = digest[:] + } + trimmed := rewindImportedTurnPrefix(ids, runRewindDecision{ + TargetTurnSeq: 4, + HasClientTurnCount: true, + ClientTurnCount: 1, + }) + if len(trimmed) != 1 || string(trimmed[0]) != string(ids[0]) { + t.Fatalf("rewindImportedTurnPrefix() = %x, want first imported turn only", trimmed) + } +} diff --git a/internal/backend/forwarder/interrupted_output_test.go b/internal/backend/forwarder/interrupted_output_test.go index 4097358..c747121 100644 --- a/internal/backend/forwarder/interrupted_output_test.go +++ b/internal/backend/forwarder/interrupted_output_test.go @@ -2,6 +2,7 @@ package forwarder import ( "encoding/json" + "errors" "strings" "testing" ) @@ -125,6 +126,53 @@ func TestCancelPersistsInterruptedProviderOutputIdempotently(t *testing.T) { } } +func TestGenericProviderFailurePersistsAccumulatedOutput(t *testing.T) { + service, stream, _ := testCheckpointBlobProjection(t) + conversation, _, _, err := service.snapshotCheckpointConversation(stream) + if err != nil { + t.Fatalf("snapshotCheckpointConversation() error = %v", err) + } + if _, err := service.store.SaveConversationWithEntries(stream.ConversationID, conversation, conversation.Entries); err != nil { + t.Fatalf("SaveConversationWithEntries() error = %v", err) + } + + stream.mu.Lock() + stream.CurrentModelCallID = "model-call-1" + stream.ProviderActive = true + stream.ProviderAccumulatedText = "partial answer before transport failure" + stream.Status = StreamStatusStreaming + stream.Phase = TurnPhaseProviderRunning + stream.mu.Unlock() + if err := service.handleProviderDoneEvent(stream, &streamProviderEvent{ + Done: true, + Err: errors.New("transport failed"), + }); err != nil { + t.Fatalf("handleProviderDoneEvent() error = %v", err) + } + + persisted, err := service.store.LoadConversation(stream.ConversationID) + if err != nil { + t.Fatalf("LoadConversation() error = %v", err) + } + foundPartialOutput := false + for _, entry := range persisted.Entries { + if entry.Kind != "assistant_text" { + continue + } + var payload assistantTextPayload + if err := json.Unmarshal(entry.Payload, &payload); err != nil { + t.Fatalf("decode assistant entry: %v", err) + } + if payload.Text == "partial answer before transport failure" { + foundPartialOutput = true + break + } + } + if !foundPartialOutput { + t.Fatal("generic provider failure discarded accumulated assistant output") + } +} + func TestCancelPreservesPersistedTurnActivityWithoutLiveAccumulator(t *testing.T) { service, stream, _ := testCheckpointBlobProjection(t) conversation, _, _, err := service.snapshotCheckpointConversation(stream) diff --git a/internal/backend/forwarder/projector.go b/internal/backend/forwarder/projector.go index 0b5980f..7f454b7 100644 --- a/internal/backend/forwarder/projector.go +++ b/internal/backend/forwarder/projector.go @@ -326,20 +326,24 @@ func compactedPromptProjectionEntries(entries []HistoryEntry) []HistoryEntry { latestToolCallID := latestCompletedToolCallIDForTurn(entries, compactionPayload.CurrentTurnSeq, compactionPayload.CurrentRequestID) preservedIndexes = autoCompactionPreservedEntryIndexes(entries, compactionPayload.CurrentTurnSeq, compactionPayload.CurrentRequestID, latestToolCallID) } - filtered := make([]HistoryEntry, 0, len(entries)-compactionIndex) - for index, entry := range entries { - if index < compactionIndex && isPromptReplayEntryKind(entry.Kind) { - if _, ok := preservedIndexes[index]; !ok { - continue - } + filtered := make([]HistoryEntry, 0, len(entries)-compactionIndex+len(preservedIndexes)) + for index := 0; index < compactionIndex; index++ { + if !isPromptReplayEntryKind(entries[index].Kind) { + filtered = append(filtered, entries[index]) } - if index < compactionIndex { - if rewritten, ok := compactedProjectionPreservedEntry(entry); ok { - entry = rewritten - } + } + filtered = append(filtered, entries[compactionIndex]) + for index := 0; index < compactionIndex; index++ { + if _, ok := preservedIndexes[index]; !ok || isCompactionSummaryKind(entries[index].Kind) { + continue + } + entry := entries[index] + if rewritten, ok := compactedProjectionPreservedEntry(entry); ok { + entry = rewritten } filtered = append(filtered, entry) } + filtered = append(filtered, entries[compactionIndex+1:]...) return filtered } @@ -575,7 +579,7 @@ func (projector *HistoryProjector) ProjectCheckpointProjection(conversation *Con if err != nil { return nil, err } - state.Turns = turnIDs + state.Turns = append(cloneByteSlices(conversation.ImportedTurnIDs), turnIDs...) replayMessages, err := projector.ProjectPromptReplay(conversation) if err != nil { return nil, err @@ -1292,7 +1296,7 @@ func filterCheckpointPersistentToolReplay(messages []promptengine.Message) []pro return filtered } -func restoreImportedReplayUserMessages(messages []promptengine.Message, importedTurns [][]byte) []promptengine.Message { +func restoreImportedReplayUserMessages(messages []promptengine.Message, importedTurns [][]byte, blobs importedBlobStore) []promptengine.Message { if len(messages) == 0 || len(importedTurns) == 0 { return messages } @@ -1301,16 +1305,16 @@ func restoreImportedReplayUserMessages(messages []promptengine.Message, imported if len(rawTurn) == 0 { continue } - turn := &agentv1.ConversationTurnStructure{} - if err := proto.Unmarshal(rawTurn, turn); err != nil { + turn, _, err := decodeImportedTurn(rawTurn, blobs) + if err != nil || turn == nil { continue } agentTurn := turn.GetAgentConversationTurn() if agentTurn == nil || len(agentTurn.GetUserMessage()) == 0 { continue } - userMessage := &agentv1.UserMessage{} - if err := proto.Unmarshal(agentTurn.GetUserMessage(), userMessage); err != nil { + userMessage, err := decodeImportedUserMessage(agentTurn.GetUserMessage(), blobs) + if err != nil { continue } replay, ok := promptengine.BuildUserMessageReplayMessage(userMessage) diff --git a/internal/backend/forwarder/rewind.go b/internal/backend/forwarder/rewind.go index ba7c851..a97401a 100644 --- a/internal/backend/forwarder/rewind.go +++ b/internal/backend/forwarder/rewind.go @@ -224,6 +224,7 @@ func (service *Service) applyRunRewindToConversation(conversation *ConversationF conversation.Entries = nil conversation.NextEntrySeq = 1 conversation.NextTurnSeq = 1 + conversation.ImportedTurnIDs = rewindImportedTurnPrefix(conversation.ImportedTurnIDs, decision) appendEntriesInPlace(conversation, appendReplacementRunEntries(decision.PrefixEntries, entries)) applyRunRewindConversationState(conversation, intent, turnSeq) deriveConversationLoopState(conversation) @@ -269,10 +270,30 @@ func applyRunRewindMetadata(conversation *ConversationFile, source *Conversation if source.TokenDetailsMaxTokens > 0 { conversation.TokenDetailsMaxTokens = source.TokenDetailsMaxTokens } + decision := runRewindDecision{TargetTurnSeq: turnSeq} + if intent.ConversationState != nil { + decision.HasClientTurnCount = true + decision.ClientTurnCount = len(intent.ConversationState.GetTurns()) + } + conversation.ImportedTurnIDs = rewindImportedTurnPrefix(source.ImportedTurnIDs, decision) } applyRunRewindConversationState(conversation, intent, turnSeq) } +func rewindImportedTurnPrefix(importedTurnIDs [][]byte, decision runRewindDecision) [][]byte { + keep := decision.TargetTurnSeq - 1 + if decision.HasClientTurnCount { + keep = int64(decision.ClientTurnCount) + } + if keep <= 0 || len(importedTurnIDs) == 0 { + return nil + } + if keep > int64(len(importedTurnIDs)) { + keep = int64(len(importedTurnIDs)) + } + return cloneByteSlices(importedTurnIDs[:keep]) +} + func (service *Service) logRunRewindDecision(requestID string, conversationID string, eventName string, decision runRewindDecision) { if service == nil || !decision.Evaluated { return diff --git a/internal/backend/forwarder/runtime_summary.go b/internal/backend/forwarder/runtime_summary.go index 9abbba3..b302c8a 100644 --- a/internal/backend/forwarder/runtime_summary.go +++ b/internal/backend/forwarder/runtime_summary.go @@ -50,7 +50,7 @@ func (service *Service) bootstrapRuntimeConversation(intent InboundIntent) (*Con } importedEntries := []HistoryEntry(nil) if len(conversation.Entries) == 0 && intent.ConversationState != nil { - importedEntries, err = service.importConversationState(conversation, intent.ConversationState) + importedEntries, err = service.importConversationState(conversation, intent.ConversationState, intent.PreFetchedBlobs) if err != nil { return nil, agentv1.AgentMode_AGENT_MODE_AGENT, 0, nil, err } @@ -138,6 +138,7 @@ func (service *Service) syncConversationRecord(conversationID string, conversati item.AutoCompactionReserveTokens = conversation.AutoCompactionReserveTokens item.AutoCompactionTriggeredAt = conversation.AutoCompactionTriggeredAt item.AutoCompactionSourceModelCallID = conversation.AutoCompactionSourceModelCallID + item.ImportedTurnIDs = cloneByteSlices(conversation.ImportedTurnIDs) item.LatestRequestPrefix = cloneConversationRequestPrefix(conversation.LatestRequestPrefix) item.LastProviderCall = cloneConversationProviderCall(conversation.LastProviderCall) item.CreatedAt = conversation.CreatedAt diff --git a/internal/backend/forwarder/service.go b/internal/backend/forwarder/service.go index 3f2e223..c97ace3 100644 --- a/internal/backend/forwarder/service.go +++ b/internal/backend/forwarder/service.go @@ -559,6 +559,7 @@ func (service *Service) decodeInboundIntent(requestID string, message *agentv1.A } intent.ConversationID = conversationID intent.ConversationState = runRequest.GetConversationState() + intent.PreFetchedBlobs = runRequest.GetPreFetchedBlobs() intent.UserMessage = extractUserMessage(message) intent.RequestContext = extractRequestContext(message) if service.shouldIgnoreEmptyResumeRunRequest(requestID, runRequest, intent.UserMessage, intent.RequestContext) { @@ -606,6 +607,7 @@ func (service *Service) decodeInboundIntent(requestID string, message *agentv1.A intent.ConversationID = conversationID intent.SubagentTypeName = strings.TrimSpace(prewarmRequest.GetSubagentTypeName()) intent.ConversationState = prewarmRequest.GetConversationState() + intent.PreFetchedBlobs = prewarmRequest.GetPreFetchedBlobs() intent.Mode, intent.ModeSource, intent.HasExplicitMode, err = extractPrewarmMode(prewarmRequest) if err != nil { return InboundIntent{}, err @@ -2246,6 +2248,15 @@ func (service *Service) finishSuccessfulTurnAfterCheckpoint(stream *ActiveStream return nil } +func (service *Service) finishFailedTurnAfterCheckpoint(stream *ActiveStream, terminalCode string, terminalMessage string) error { + if stream == nil { + return nil + } + err := service.broker.Fail(stream.RequestID, terminalCode, terminalMessage) + service.setTurnPhase(stream, TurnPhaseFailed) + return err +} + func (service *Service) failStreamIfNonTerminal(stream *ActiveStream, terminalCode string, cause error) error { if stream == nil || cause == nil { return nil @@ -2265,6 +2276,10 @@ func (service *Service) publishCheckpoint(requestID string, conversationID strin } func (service *Service) publishCheckpointWithCompletion(requestID string, _ string, completion *pendingTurnCompletion) error { + return service.publishCheckpointWithTerminalAction(requestID, successfulCheckpointTerminalAction(completion)) +} + +func (service *Service) publishCheckpointWithTerminalAction(requestID string, terminal checkpointTerminalAction) error { stream, ok := service.broker.Get(requestID) if !ok || stream == nil { return fmt.Errorf("request is not active: %s", requestID) @@ -2282,7 +2297,7 @@ func (service *Service) publishCheckpointWithCompletion(requestID string, _ stri } projection.State.PendingToolCalls = buildPendingToolCalls(pendingExecs, pendingInteractions) service.rewriteCheckpointTokenDetailsForClient(stream, conversation, projection.State) - return service.queueCheckpointProjection(stream, projection, completion) + return service.queueCheckpointProjectionWithTerminal(stream, projection, terminal) } func (service *Service) rewriteCheckpointTokenDetailsForClient(stream *ActiveStream, conversation *ConversationFile, state *agentv1.ConversationStateStructure) { @@ -2422,18 +2437,20 @@ func (service *Service) failActiveStream(stream *ActiveStream, conversationID st if cancel != nil { cancel() } - service.setTurnPhase(stream, TurnPhaseFailed) - var firstErr error - if err := service.syncSummaryCarryForward(conversationID, requestID, modelCallID); err != nil && firstErr == nil { - firstErr = err + if err := service.syncSummaryCarryForward(conversationID, requestID, modelCallID); err != nil { + log.Printf( + "forwarder summary sync before failed terminal skipped request_id=%s model_call_id=%s err=%v", + strings.TrimSpace(requestID), + strings.TrimSpace(modelCallID), + err, + ) } - if err := service.publishCheckpoint(requestID, conversationID); err != nil && firstErr == nil { - firstErr = err + terminal := failedCheckpointTerminalAction(terminalCode, terminalMessage) + if err := service.publishCheckpointWithTerminalAction(requestID, terminal); err != nil { + log.Printf("forwarder checkpoint queue before failed terminal skipped request_id=%s err=%v", strings.TrimSpace(requestID), err) + return service.finishFailedTurnAfterCheckpoint(stream, terminalCode, terminalMessage) } - if err := service.broker.Fail(requestID, terminalCode, terminalMessage); err != nil && firstErr == nil { - firstErr = err - } - return firstErr + return nil } // buildRunEntries 构造一次 run intent 需要写入 history 的首批 entry。 diff --git a/internal/backend/forwarder/token_usage.go b/internal/backend/forwarder/token_usage.go index ea3e894..069d69e 100644 --- a/internal/backend/forwarder/token_usage.go +++ b/internal/backend/forwarder/token_usage.go @@ -45,13 +45,25 @@ func (snapshot turnUsageSnapshot) requestTokensTotal() int64 { return snapshot.promptTokensTotal() + nonNegativeInt64(snapshot.OutputTokens) } -func (service *Service) importConversationState(item *ConversationFile, state *agentv1.ConversationStateStructure) ([]HistoryEntry, error) { +func (service *Service) importConversationState(item *ConversationFile, state *agentv1.ConversationStateStructure, prefetchedBlobs []*agentv1.PreFetchedBlob) ([]HistoryEntry, error) { if item == nil || state == nil { return nil, nil } + blobs, err := newImportedBlobStore(prefetchedBlobs) + if err != nil { + return nil, err + } + importedIDs, err := importedTurnIDs(state.GetTurns(), blobs) + if err != nil { + return nil, err + } item.TokenDetailsUsedTokens = state.GetTokenDetails().GetUsedTokens() + item.ImportedTurnIDs = importedIDs + if minimumNextTurnSeq := int64(len(item.ImportedTurnIDs)) + 1; item.NextTurnSeq < minimumNextTurnSeq { + item.NextTurnSeq = minimumNextTurnSeq + } entries := make([]HistoryEntry, 0, 2) - if messages, err := importedConversationStateModelMessages(state); err != nil { + if messages, err := importedConversationStateModelMessagesWithBlobs(state, blobs); err != nil { return nil, err } else { for _, message := range messages { @@ -105,6 +117,10 @@ func (service *Service) importConversationState(item *ConversationFile, state *a } func importedConversationStateModelMessages(state *agentv1.ConversationStateStructure) ([]modeladapter.Message, error) { + return importedConversationStateModelMessagesWithBlobs(state, nil) +} + +func importedConversationStateModelMessagesWithBlobs(state *agentv1.ConversationStateStructure, blobs importedBlobStore) ([]modeladapter.Message, error) { if state == nil { return nil, nil } @@ -113,7 +129,7 @@ func importedConversationStateModelMessages(state *agentv1.ConversationStateStru if err != nil { return nil, fmt.Errorf("decode imported replay messages: %w", err) } - decoded = restoreImportedReplayUserMessages(decoded, state.GetTurns()) + decoded = restoreImportedReplayUserMessages(decoded, state.GetTurns(), blobs) decoded = filterLegacyPlainWriteReplay(decoded) decoded = filterInternalPromptContextReplay(decoded) messages := make([]modeladapter.Message, 0, len(decoded)) @@ -133,35 +149,18 @@ func importedConversationStateModelMessages(state *agentv1.ConversationStateStru if len(rawTurn) == 0 { continue } - turn := &agentv1.ConversationTurnStructure{} - if err := proto.Unmarshal(rawTurn, turn); err != nil { - return nil, fmt.Errorf("decode imported turn: %w", err) + turn, turnID, err := decodeImportedTurn(rawTurn, blobs) + if err != nil { + return nil, err } - agentTurn := turn.GetAgentConversationTurn() - if agentTurn == nil { - continue + if turn == nil && len(turnID) > 0 { + return nil, fmt.Errorf("missing prefetched turn blob %x", turnID) } - if rawUser := agentTurn.GetUserMessage(); len(rawUser) > 0 { - userMessage := &agentv1.UserMessage{} - if err := proto.Unmarshal(rawUser, userMessage); err != nil { - return nil, fmt.Errorf("decode imported turn user_message: %w", err) - } - if replay, ok := promptengine.BuildUserMessageReplayMessage(userMessage); ok { - messages = append(messages, toModelMessage(replay)) - } - } - for _, rawStep := range agentTurn.GetSteps() { - if len(rawStep) == 0 { - continue - } - step := &agentv1.ConversationStep{} - if err := proto.Unmarshal(rawStep, step); err != nil { - return nil, fmt.Errorf("decode imported turn step: %w", err) - } - for _, replay := range promptengine.BuildLegacyMessagesFromConversationStep(step) { - messages = append(messages, toModelMessage(replay)) - } + turnMessages, err := importedBlobTurnMessages(turn, blobs) + if err != nil { + return nil, err } + messages = append(messages, turnMessages...) } return normalizeReplayMessageSequence(messages), nil } diff --git a/internal/backend/forwarder/types.go b/internal/backend/forwarder/types.go index 7eb7788..319062c 100644 --- a/internal/backend/forwarder/types.go +++ b/internal/backend/forwarder/types.go @@ -39,6 +39,7 @@ type ConversationFile struct { CurrentPlanText string `json:"current_plan_text,omitempty"` CurrentPlans map[string]*agentv1.PlanRegistryEntry `json:"current_plans,omitempty"` CurrentTodos []*agentv1.TodoItem `json:"current_todos,omitempty"` + ImportedTurnIDs [][]byte `json:"imported_turn_ids,omitempty"` LatestRequestPrefix *ConversationRequestPrefix `json:"latest_request_prefix,omitempty"` LastProviderCall *ConversationProviderCall `json:"last_provider_call,omitempty"` CreatedAt time.Time `json:"created_at"` @@ -224,11 +225,25 @@ type pendingTurnCompletion struct { Disposition pendingCompletionDisposition } +type checkpointTerminalActionKind uint8 + +const ( + checkpointTerminalActionNone checkpointTerminalActionKind = iota + checkpointTerminalActionComplete + checkpointTerminalActionFail +) + +type checkpointTerminalAction struct { + Kind checkpointTerminalActionKind + Completion pendingTurnCompletion + ErrorCode string + ErrorMessage string +} + type pendingCheckpointPublish struct { - State *agentv1.ConversationStateStructure - Required map[string]struct{} - Completion *pendingTurnCompletion - Published bool + State *agentv1.ConversationStateStructure + Required map[string]struct{} + Terminal checkpointTerminalAction } type PendingCompaction struct { @@ -430,6 +445,7 @@ type InboundIntent struct { SubagentTypeName string SubagentModelOverrides map[string]runtimecore.SubagentModelOverrideSelection ConversationState *agentv1.ConversationStateStructure + PreFetchedBlobs []*agentv1.PreFetchedBlob UserMessage *agentv1.UserMessage RequestContext *agentv1.RequestContext ClientMessage *agentv1.AgentClientMessage