Merge pull request #103 from DEEIX-AI/stream

feat: implement cancellation of pending message generations

Author: chenyme

backend/internal/application/conversation/generation_stream.go

@@ -53,7 +53,25 @@ func EnsureMessageGenerationRunID(raw string) string {
 
 // CancelMessageGeneration 取消用户显式停止的流式生成；浏览器刷新不会走这个路径。
 func (s *Service) CancelMessageGeneration(ctx context.Context, userID uint, runID string) bool {
-	return s.generationStreams.cancel(ctx, userID, normalizeRunID(runID))
+	normalizedRunID := normalizeRunID(runID)
+	canceled := s.generationStreams.cancel(ctx, userID, normalizedRunID)
+	if !canceled || s == nil || s.repo == nil {
+		return canceled
+	}
+	markCtx := ctx
+	var cancel context.CancelFunc
+	if markCtx == nil || markCtx.Err() != nil {
+		markCtx, cancel = context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+	}
+	_, _ = s.repo.CancelPendingGenerationMessagesByRunID(
+		markCtx,
+		userID,
+		normalizedRunID,
+		classifyRunErrorCode(ErrMessageGenerationCanceled),
+		ErrMessageGenerationCanceled.Error(),
+	)
+	return true
 }
 
 // PublishMessageGenerationEvent 发布生成流事件，并返回带 seq 的实际载荷。

backend/internal/infra/persistence/postgres/conversation/repository.go

@@ -889,6 +889,51 @@ func (r *Repo) UpdateMessageState(
 		Error)
 }
 
+// CancelPendingGenerationMessagesByRunID 将用户显式取消的 pending 回合更新为稳定终态。
+func (r *Repo) CancelPendingGenerationMessagesByRunID(
+	ctx context.Context,
+	userID uint,
+	runID string,
+	errorCode string,
+	errorMessage string,
+) (bool, error) {
+	normalizedRunID := strings.TrimSpace(runID)
+	if userID == 0 || normalizedRunID == "" {
+		return false, repository.ErrInvalidInput
+	}
+	normalizedErrorCode := strings.TrimSpace(errorCode)
+	normalizedErrorMessage := truncateText(strings.TrimSpace(errorMessage), 255)
+	returnedRows := int64(0)
+	err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		userResult := tx.Model(&models.Message{}).
+			Where("user_id = ? AND run_id = ? AND role = ? AND status = ?", userID, normalizedRunID, "user", "pending").
+			Updates(map[string]interface{}{
+				"status":        "success",
+				"error_code":    "",
+				"error_message": "",
+			})
+		if userResult.Error != nil {
+			return userResult.Error
+		}
+		assistantResult := tx.Model(&models.Message{}).
+			Where("user_id = ? AND run_id = ? AND role = ? AND status = ?", userID, normalizedRunID, "assistant", "pending").
+			Updates(map[string]interface{}{
+				"status":        "canceled",
+				"error_code":    normalizedErrorCode,
+				"error_message": normalizedErrorMessage,
+			})
+		if assistantResult.Error != nil {
+			return assistantResult.Error
+		}
+		returnedRows = userResult.RowsAffected + assistantResult.RowsAffected
+		return nil
+	})
+	if err != nil {
+		return false, translateError(err)
+	}
+	return returnedRows > 0, nil
+}
+
 // InterruptPendingAssistantMessageByRunID 将失去活跃生成流的 pending assistant 标记为错误。
 func (r *Repo) InterruptPendingAssistantMessageByRunID(
 	ctx context.Context,

backend/internal/repository/conversation_core.go

@@ -71,6 +71,7 @@ type MessageRepository interface {
 	GetMessageByPublicIDForUser(ctx context.Context, userID uint, publicID string) (*domainconversation.Message, error)
 	UpdateMessageUsage(ctx context.Context, messageID uint, inputTokens int64, outputTokens int64, cacheReadTokens int64, cacheWriteTokens int64, reasoningTokens int64) error
 	UpdateMessageState(ctx context.Context, messageID uint, status string, errorCode string, errorMessage string) error
+	CancelPendingGenerationMessagesByRunID(ctx context.Context, userID uint, runID string, errorCode string, errorMessage string) (bool, error)
 	InterruptPendingAssistantMessageByRunID(ctx context.Context, userID uint, runID string, errorCode string, errorMessage string) (bool, error)
 	UpdateAssistantMessageCompletion(ctx context.Context, messageID uint, content string, outputTokens int64, reasoningTokens int64, latencyMS int64, status string, errorCode string, errorMessage string) error
 	UpdateMessageBilling(ctx context.Context, messageID uint, billedCurrency string, billedNanousd int64, pricingSnapshot string) error

backend/internal/transport/http/server.go

@@ -302,13 +302,26 @@ func isRegularFile(filePath string) bool {
 }
 
 func applyFrontendCacheHeaders(c *gin.Context, requestPath string) {
-	if strings.HasPrefix(requestPath, "/_next/static/") || strings.HasPrefix(requestPath, "/fonts/") {
+	if isImmutableFrontendAsset(requestPath) {
 		c.Header("Cache-Control", "public, max-age=31536000, immutable")
 		return
 	}
+	if isNextExportDataAsset(requestPath) {
+		c.Header("Cache-Control", "public, max-age=86400, stale-while-revalidate=604800")
+		return
+	}
 	c.Header("Cache-Control", "public, max-age=3600")
 }
 
+func isImmutableFrontendAsset(requestPath string) bool {
+	return strings.HasPrefix(requestPath, "/_next/static/") || strings.HasPrefix(requestPath, "/fonts/")
+}
+
+func isNextExportDataAsset(requestPath string) bool {
+	fileName := path.Base(requestPath)
+	return strings.HasPrefix(fileName, "__next.") && strings.EqualFold(path.Ext(fileName), ".txt")
+}
+
 func readyzHandler(hc HealthChecker) gin.HandlerFunc {
 	return func(c *gin.Context) {
 		ctx, cancel := context.WithTimeout(c.Request.Context(), 5*time.Second)

backend/internal/transport/http/server_test.go

@@ -34,6 +34,57 @@ func TestFrontendStaticFallbackServesExportedPage(t *testing.T) {
 	if strings.TrimSpace(recorder.Body.String()) != "chat page" {
 		t.Fatalf("expected chat page, got %q", recorder.Body.String())
 	}
+	if got := recorder.Header().Get("Cache-Control"); got != "no-cache" {
+		t.Fatalf("expected exported page no-cache, got %q", got)
+	}
+}
+
+func TestFrontendStaticCachesNextExportData(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	root := t.TempDir()
+	if err := os.WriteFile(filepath.Join(root, "__next._tree.txt"), []byte("tree"), 0o644); err != nil {
+		t.Fatalf("write next data: %v", err)
+	}
+
+	engine := gin.New()
+	registerFrontendStatic(engine, root, nil)
+
+	recorder := httptest.NewRecorder()
+	request := httptest.NewRequest(http.MethodGet, "/__next._tree.txt?conversation_id=demo&_rsc=abc", nil)
+	engine.ServeHTTP(recorder, request)
+
+	if recorder.Code != http.StatusOK {
+		t.Fatalf("expected status 200, got %d", recorder.Code)
+	}
+	if got := recorder.Header().Get("Cache-Control"); got != "public, max-age=86400, stale-while-revalidate=604800" {
+		t.Fatalf("expected next export data cache header, got %q", got)
+	}
+}
+
+func TestFrontendStaticCachesImmutableBuildAssets(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	root := t.TempDir()
+	chunkDir := filepath.Join(root, "_next", "static", "chunks")
+	if err := os.MkdirAll(chunkDir, 0o755); err != nil {
+		t.Fatalf("create chunk dir: %v", err)
+	}
+	if err := os.WriteFile(filepath.Join(chunkDir, "app.js"), []byte("chunk"), 0o644); err != nil {
+		t.Fatalf("write chunk: %v", err)
+	}
+
+	engine := gin.New()
+	registerFrontendStatic(engine, root, nil)
+
+	recorder := httptest.NewRecorder()
+	request := httptest.NewRequest(http.MethodGet, "/_next/static/chunks/app.js", nil)
+	engine.ServeHTTP(recorder, request)
+
+	if recorder.Code != http.StatusOK {
+		t.Fatalf("expected status 200, got %d", recorder.Code)
+	}
+	if got := recorder.Header().Get("Cache-Control"); got != "public, max-age=31536000, immutable" {
+		t.Fatalf("expected immutable cache header, got %q", got)
+	}
 }
 
 func TestFrontendStaticFallbackSkipsAPIPaths(t *testing.T) {

frontend/features/chat/components/app-chat-area.tsx

@@ -370,6 +370,7 @@ export function AppChatArea() {
     setAttachments,
     releaseAttachments,
     activeGenerationRunsRef,
+    resumingRunID,
   });
   const generating = sending || Boolean(resumingRunID);
   const uploadDropDisabled = generating || loading || uploading;

frontend/features/chat/hooks/use-chat-branch-state.ts

@@ -165,11 +165,13 @@ export function useChatBranchState({
   resetToken,
   messages,
   pendingExchange,
+  liveRunIDs,
 }: {
   conversationID: string | null;
   resetToken: number;
   messages: MessageDTO[];
   pendingExchange: PendingExchange | null;
+  liveRunIDs?: ReadonlySet<string>;
 }) {
   const t = useTranslations("chat.messages");
   const [branchSelections, setBranchSelections] = React.useState<Record<string, string>>({});
@@ -181,13 +183,17 @@ export function useChatBranchState({
   const serverTreeMessages = React.useMemo(
     () =>
       messages.map((item) =>
-        mapServerMessage(item, {
-          generationInterrupted: t("generationInterrupted"),
-          streamInterrupted: t("streamInterrupted"),
-          imageRunning: t("imageRunning"),
-        }),
+        mapServerMessage(
+          item,
+          {
+            generationInterrupted: t("generationInterrupted"),
+            streamInterrupted: t("streamInterrupted"),
+            imageRunning: t("imageRunning"),
+          },
+          { liveRunIDs },
+        ),
       ),
-    [messages, t],
+    [liveRunIDs, messages, t],
   );
   const serverMessagePublicIDs = React.useMemo(
     () => new Set(serverTreeMessages.map((item) => item.publicID).filter(Boolean)),
@@ -234,7 +240,7 @@ export function useChatBranchState({
           item.role === "assistant" &&
           ((pendingExchange.runID && item.runID === pendingExchange.runID) ||
             item.publicID === (pendingExchange.assistantPublicID || pendingExchange.tempAssistantPublicID)) &&
-          item.isPending,
+          item.isStreaming,
       ),
   );
 

frontend/features/chat/hooks/use-chat-runtime.ts

@@ -35,6 +35,7 @@ export function useChatRuntime({
   setAttachments,
   releaseAttachments,
   activeGenerationRunsRef,
+  resumingRunID = "",
 }: {
   conversationID: string | null;
   resetToken: number;
@@ -58,16 +59,22 @@ export function useChatRuntime({
   setAttachments: React.Dispatch<React.SetStateAction<PendingAttachment[]>>;
   releaseAttachments: (items: PendingAttachment[]) => void;
   activeGenerationRunsRef?: React.RefObject<Set<string>>;
+  resumingRunID?: string;
 }) {
   const [showConversationLayout, setShowConversationLayout] = React.useState(false);
   const [pendingExchange, setPendingExchange] = React.useState<PendingExchange | null>(null);
   const previousResetTokenRef = React.useRef(resetToken);
+  const liveServerRunIDs = React.useMemo(() => {
+    const normalized = resumingRunID.trim();
+    return normalized ? new Set([normalized]) : undefined;
+  }, [resumingRunID]);
 
   const branchState = useChatBranchState({
     conversationID,
     resetToken,
     messages,
     pendingExchange,
+    liveRunIDs: liveServerRunIDs,
   });
 
   const submitState = useChatSubmitStream({

frontend/features/chat/model/chat-thread.ts

@@ -160,6 +160,7 @@ export function mapServerMessage(
   labels: { generationInterrupted: string; streamInterrupted?: string; imageRunning?: string } = {
     generationInterrupted: "Generation interrupted",
   },
+  options: { liveRunIDs?: ReadonlySet<string> } = {},
 ): ChatAreaMessage {
   const publicID = item.publicID.trim();
   const msg: ChatAreaMessage = {
@@ -210,9 +211,11 @@ export function mapServerMessage(
       };
     }
     if (item.status === "pending") {
-      msg.isPending = true;
-      msg.isStreaming = true;
-      msg.activityLabel = item.contentType === "image" ? labels.imageRunning : undefined;
+      const liveRunID = item.runID?.trim() || "";
+      const live = Boolean(liveRunID && options.liveRunIDs?.has(liveRunID));
+      msg.isPending = live;
+      msg.isStreaming = live;
+      msg.activityLabel = live && item.contentType === "image" ? labels.imageRunning : undefined;
     }
   }
   return msg;