From d22e240a2485e7cfbb64afd595f2c86ef012ef2d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 19 Jan 2026 11:16:32 +0000 Subject: [PATCH 01/10] docs: add streaming design complexity analysis and improvement TODO - Analyze current front-end streaming architecture - Document 10 key complexity issues with code examples - Create phased TODO with 6 implementation phases - Include architecture diagrams and success metrics Co-authored-by: e0945797 --- .../_webapp/docs/STREAMING_DESIGN_ANALYSIS.md | 308 +++++++++++++++++ webapp/_webapp/docs/STREAMING_DESIGN_TODO.md | 322 ++++++++++++++++++ 2 files changed, 630 insertions(+) create mode 100644 webapp/_webapp/docs/STREAMING_DESIGN_ANALYSIS.md create mode 100644 webapp/_webapp/docs/STREAMING_DESIGN_TODO.md diff --git a/webapp/_webapp/docs/STREAMING_DESIGN_ANALYSIS.md b/webapp/_webapp/docs/STREAMING_DESIGN_ANALYSIS.md new file mode 100644 index 00000000..ce575896 --- /dev/null +++ b/webapp/_webapp/docs/STREAMING_DESIGN_ANALYSIS.md @@ -0,0 +1,308 @@ +# Front-End Streaming Logic Complexity Analysis + +## Executive Summary + +The current streaming implementation is spread across **15+ files** with fragmented logic, inconsistent patterns, and multiple data transformations that make the codebase difficult to understand, maintain, and debug. + +--- + +## Architecture Overview + +### Current Data Flow + +``` +┌────────────────────────────────────────────────────────────────────────────┐ +│ useSendMessageStream │ +│ (Main orchestrator hook - handles all stream event routing) │ +└─────────────────────────────────┬──────────────────────────────────────────┘ + │ + ▼ +┌────────────────────────────────────────────────────────────────────────────┐ +│ createConversationMessageStream │ +│ (API call returns ReadableStream) │ +└─────────────────────────────────┬──────────────────────────────────────────┘ + │ + ▼ +┌────────────────────────────────────────────────────────────────────────────┐ +│ processStream │ +│ (Parses NDJSON, calls onMessage callback for each chunk) │ +└─────────────────────────────────┬──────────────────────────────────────────┘ + │ + ┌─────────────────────────────┼─────────────────────────────────┐ + │ │ │ + ▼ ▼ ▼ +┌─────────────────┐ ┌─────────────────┐ ┌─────────────────────┐ +│ streamInit │ │ streamPartBegin │ │ streamFinalization │ +│ handler │ │ handler │ │ handler │ +└────────┬────────┘ └────────┬────────┘ └──────────┬──────────┘ + │ │ │ + │ ▼ │ + │ ┌─────────────────┐ │ + │ │ messageChunk │ │ + │ │ handler │ │ + │ └────────┬────────┘ │ + │ │ │ + │ ▼ │ + │ ┌─────────────────┐ │ + │ │ streamPartEnd │ │ + │ │ handler │ │ + │ └────────┬────────┘ │ + │ │ │ + ▼ ▼ ▼ +┌────────────────────────────────────────────────────────────────────────────┐ +│ streaming-message-store │ +│ { parts: MessageEntry[], sequence: number } │ +└─────────────────────────────────┬──────────────────────────────────────────┘ + │ + ▼ (flushStreamingMessageToConversation) +┌────────────────────────────────────────────────────────────────────────────┐ +│ conversation-store │ +│ { currentConversation: Conversation, isStreaming: boolean } │ +└─────────────────────────────────┬──────────────────────────────────────────┘ + │ + ▼ +┌────────────────────────────────────────────────────────────────────────────┐ +│ ChatBody Component │ +│ (Renders both finalized and streaming messages) │ +└────────────────────────────────────────────────────────────────────────────┘ +``` + +--- + +## Complexity Problems + +### 1. **Fragmented Handler Architecture** (9 separate files) + +The streaming logic is scattered across 9 handler files in `/stores/conversation/handlers/`: + +| File | Purpose | Lines | +|------|---------|-------| +| `handleStreamInitialization.ts` | Mark user message as finalized | 29 | +| `handleStreamPartBegin.ts` | Create new MessageEntry for incoming part | 73 | +| `handleMessageChunk.ts` | Append delta text to assistant content | 37 | +| `handleStreamPartEnd.ts` | Mark part as finalized with full content | 99 | +| `handleStreamFinalization.ts` | Flush streaming messages to conversation | 6 | +| `handleStreamError.ts` | Handle stream errors with retry logic | 53 | +| `handleIncompleteIndicator.ts` | Set incomplete indicator state | 6 | +| `handleError.ts` | Mark all preparing messages as stale | 20 | +| `converter.ts` | Convert MessageEntry to Message | 69 | + +**Problem:** Following a single streaming flow requires jumping between 9+ files. Related logic is separated rather than cohesive. + +--- + +### 2. **Inconsistent Store Access Patterns** + +Three different patterns are used to update the same store: + +```typescript +// Pattern 1: Direct setState (handleStreamInitialization.ts) +useStreamingMessageStore.setState((prev) => ({ ... })); + +// Pattern 2: Passed update function (handleStreamPartBegin.ts) +updateStreamingMessage((prev) => ({ ... })); + +// Pattern 3: getState() then method call (handleIncompleteIndicator.ts) +useStreamingMessageStore.getState().setIncompleteIndicator(indicator); +``` + +**Problem:** Inconsistent patterns make it unclear which approach is correct and why. + +--- + +### 3. **Duplicated Type Checking Logic** + +The same exhaustive type-check pattern appears in 4+ files: + +```typescript +// Repeated in: useSendMessageStream.ts, handleStreamPartBegin.ts, handleStreamPartEnd.ts +if (role !== undefined) { + const _typeCheck: never = role; + throw new Error("Unexpected response payload: " + _typeCheck); +} +``` + +**Problem:** Duplication increases maintenance burden. If a new message type is added, 4+ files need updating. + +--- + +### 4. **Nearly Identical Switch Statements** + +`handleStreamPartBegin.ts` (lines 11-72) and `handleStreamPartEnd.ts` (lines 16-98) have almost identical structures: + +```typescript +// handleStreamPartBegin.ts +if (role === "assistant") { ... } +else if (role === "toolCallPrepareArguments") { ... } +else if (role === "toolCall") { ... } +else if (role === "system") { /* nothing */ } +else if (role === "user") { /* nothing */ } +else if (role === "unknown") { /* nothing */ } + +// handleStreamPartEnd.ts (same pattern with switch) +switch (role) { + case "assistant": { ... } + case "toolCallPrepareArguments": { ... } + case "toolCall": { ... } + case "system": { break; } + case "unknown": { break; } + case "user": { break; } +} +``` + +**Problem:** Changes to message type handling require modifications in multiple places. Uses different control flow structures (if-else vs switch) for the same logic. + +--- + +### 5. **Complex State Transitions Across Files** + +Message status transitions are implicit and distributed: + +``` +MessageEntryStatus Flow: + PREPARING → (in handleStreamPartBegin) + ↓ + PREPARING + content updates → (in handleMessageChunk) + ↓ + FINALIZED → (in handleStreamPartEnd or handleStreamInitialization) + ↓ + [flush to conversation store] → (in converter.ts) + +Error paths: + PREPARING → STALE → (in handleError.ts) +``` + +**Problem:** No single place documents or enforces valid state transitions. Easy to introduce bugs. + +--- + +### 6. **Dual Store Architecture Creates Complexity** + +Two stores manage related message data: + +| Store | Purpose | When Used | +|-------|---------|-----------| +| `streaming-message-store` | Temporary streaming state | During active streaming | +| `conversation-store` | Persisted conversation | After stream finalization | + +Data flows: +1. Messages created in `streaming-message-store` +2. Messages modified via handlers +3. Messages "flushed" to `conversation-store` via `flushStreamingMessageToConversation` +4. `streaming-message-store` is reset + +**Problem:** +- Two sources of truth for messages +- Complex flush logic (`flushSync` required for React batching) +- UI must render from both stores simultaneously + +--- + +### 7. **Multiple Data Transformations** + +A message goes through 5+ transformations: + +``` +StreamPartBegin (protobuf) + ↓ convert +MessageEntry (internal, PREPARING) + ↓ handleMessageChunk +MessageEntry (updated content) + ↓ handleStreamPartEnd +MessageEntry (FINALIZED) + ↓ convertMessageEntryToMessage +Message (protobuf) + ↓ messageToMessageEntry (for UI) +MessageEntry (for MessageCard component) +``` + +**Problem:** Each transformation is a potential source of bugs. Data shape changes make debugging harder. + +--- + +### 8. **flushSync Usage Indicates Architectural Issues** + +`flushSync` is used in 3 places to force synchronous React updates: + +```typescript +// streaming-message-store.ts +flushSync(() => { set((state) => { ... }); }); + +// converter.ts +flushSync(() => { useConversationStore.getState().updateCurrentConversation(...); }); +``` + +**Problem:** `flushSync` is a React escape hatch. Its presence suggests the architecture fights against React's batching model. + +--- + +### 9. **Error Handling Inconsistency** + +Three different error handling approaches: + +```typescript +// handleStreamError.ts - specific to stream errors with retry +if (streamError.errorMessage.includes("project is out of date")) { + await sync(); + await sendMessageStream(currentPrompt, currentSelectedText); +} + +// handleError.ts - marks messages as stale +useStreamingMessageStore.getState().updateStreamingMessage((prev) => { + const newParts = prev.parts.map((part) => ({ + ...part, + status: part.status === MessageEntryStatus.PREPARING ? MessageEntryStatus.STALE : part.status, + })); + return { ...prev, parts: newParts }; +}); + +// with-retry-sync.ts - wrapper with generic retry logic +if (error?.code === ErrorCode.PROJECT_OUT_OF_DATE) { + await sync(); + return await operation(); // retry once +} +``` + +**Problem:** Unclear which error handler is called when. Duplicate retry logic for the same error condition. + +--- + +### 10. **Hook Has Too Many Dependencies** + +`useSendMessageStream` has 12 dependencies in its `useCallback`: + +```typescript +[ + resetStreamingMessage, + resetIncompleteIndicator, + updateStreamingMessage, + currentConversation, + refetchConversationList, + sync, + user?.id, + alwaysSyncProject, + conversationMode, + storeSurroundingText, + projectId, +] +``` + +**Problem:** Hard to reason about when the callback is recreated. Potential performance issues if any dependency changes frequently. + +--- + +## Recommendations Summary + +1. **Consolidate handlers** into a single state machine +2. **Use a single store** with clear separation between streaming and finalized state +3. **Create a message type handler registry** to eliminate switch/if-else duplication +4. **Define explicit state transitions** with validation +5. **Remove flushSync** by restructuring the update flow +6. **Unify error handling** into a single strategy +7. **Reduce transformations** by using a consistent message format + +--- + +## TODO: Improve Streaming Design + +See the accompanying TODO list for specific implementation tasks. diff --git a/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md b/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md new file mode 100644 index 00000000..2034b585 --- /dev/null +++ b/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md @@ -0,0 +1,322 @@ +# TODO: Streaming Message Design Improvements + +This document outlines a phased approach to simplify and improve the front-end streaming architecture. + +--- + +## Phase 1: Consolidate Handlers into State Machine + +### Goal +Replace 9 separate handler files with a single, cohesive state machine that manages all streaming state transitions. + +### Tasks + +- [ ] **1.1 Create StreamingStateMachine class** + ```typescript + // stores/streaming/streaming-state-machine.ts + type StreamState = 'idle' | 'receiving' | 'finalizing' | 'error'; + + class StreamingStateMachine { + private state: StreamState = 'idle'; + private messageEntries: Map = new Map(); + + handleEvent(event: StreamEvent): void { + // Central event handling + } + } + ``` + - Location: `stores/streaming/streaming-state-machine.ts` + - Benefit: Single point of control for all state transitions + +- [ ] **1.2 Define StreamEvent union type** + ```typescript + type StreamEvent = + | { type: 'INIT'; payload: StreamInitialization } + | { type: 'PART_BEGIN'; payload: StreamPartBegin } + | { type: 'CHUNK'; payload: MessageChunk } + | { type: 'PART_END'; payload: StreamPartEnd } + | { type: 'FINALIZE'; payload: StreamFinalization } + | { type: 'ERROR'; payload: StreamError } + | { type: 'INCOMPLETE'; payload: IncompleteIndicator }; + ``` + - Location: `stores/streaming/types.ts` + - Benefit: Type-safe event handling with exhaustive checking + +- [ ] **1.3 Create message type handlers registry** + ```typescript + // Eliminates duplicate switch statements + const messageTypeHandlers: Record = { + assistant: new AssistantHandler(), + toolCall: new ToolCallHandler(), + toolCallPrepareArguments: new ToolCallPrepareHandler(), + user: new NoOpHandler(), + system: new NoOpHandler(), + unknown: new NoOpHandler(), + }; + ``` + - Location: `stores/streaming/message-type-handlers.ts` + - Benefit: Add new message types without modifying multiple files + +- [ ] **1.4 Delete old handler files** + - Files to remove after migration: + - `handlers/handleStreamInitialization.ts` + - `handlers/handleStreamPartBegin.ts` + - `handlers/handleMessageChunk.ts` + - `handlers/handleStreamPartEnd.ts` + - `handlers/handleStreamFinalization.ts` + - `handlers/handleStreamError.ts` + - `handlers/handleIncompleteIndicator.ts` + - `handlers/handleError.ts` + +--- + +## Phase 2: Unify Store Architecture + +### Goal +Consolidate `streaming-message-store` and `conversation-store` message handling into a single coherent store. + +### Tasks + +- [ ] **2.1 Create unified message store** + ```typescript + // stores/message-store.ts + interface MessageStore { + // Finalized messages from server + messages: Message[]; + + // Currently streaming messages (separate from finalized) + streamingEntries: MessageEntry[]; + + // Computed: all displayable messages + get allMessages(): DisplayMessage[]; + + // Actions + appendStreamingEntry(entry: MessageEntry): void; + updateStreamingEntry(id: string, update: Partial): void; + finalizeStreaming(): void; + reset(): void; + } + ``` + - Location: `stores/message-store.ts` + - Benefit: Single source of truth with clear streaming vs finalized separation + +- [ ] **2.2 Create DisplayMessage type** + ```typescript + // Single type used by UI components + interface DisplayMessage { + id: string; + type: 'user' | 'assistant' | 'toolCall' | 'error'; + content: string; + status: 'streaming' | 'complete' | 'error'; + metadata?: MessageMetadata; + } + ``` + - Location: `stores/types.ts` + - Benefit: UI components work with one consistent type + +- [ ] **2.3 Remove flushSync calls** + - Restructure update flow so React batching works naturally + - Replace `flushSync` with proper `useSyncExternalStore` or subscription pattern + - Files affected: `streaming-message-store.ts`, `converter.ts` + +- [ ] **2.4 Migrate ChatBody to use unified store** + - Replace: + ```typescript + const visibleMessages = useMemo(() => filterVisibleMessages(conversation), [conversation]); + const streamingMessage = useStreamingMessageStore((s) => s.streamingMessage); + ``` + - With: + ```typescript + const displayMessages = useMessageStore((s) => s.allMessages); + ``` + +--- + +## Phase 3: Simplify Data Transformations + +### Goal +Reduce the number of data transformations from 5+ to 2 maximum. + +### Tasks + +- [ ] **3.1 Define canonical internal message format** + ```typescript + // Internal format used throughout the app + interface InternalMessage { + id: string; + role: MessageRole; + content: string; + status: MessageStatus; + toolCall?: ToolCallData; + attachments?: Attachment[]; + timestamp: number; + } + ``` + - Location: `types/message.ts` + - Benefit: Single format reduces confusion + +- [ ] **3.2 Create bidirectional converters** + ```typescript + // Only two conversions needed: + // 1. API response → Internal format + const fromApiMessage = (msg: ApiMessage): InternalMessage => { ... }; + + // 2. Internal format → API request + const toApiMessage = (msg: InternalMessage): ApiMessage => { ... }; + ``` + - Location: `utils/message-converters.ts` + - Benefit: Clear boundary between API types and internal types + +- [ ] **3.3 Remove MessageEntry type** + - Replace `MessageEntry` with `InternalMessage` + - Update all components to use new type + - Delete `stores/conversation/types.ts` (after migrating MessageEntryStatus) + +--- + +## Phase 4: Improve Error Handling + +### Goal +Create a unified error handling strategy for all streaming errors. + +### Tasks + +- [ ] **4.1 Create StreamingErrorHandler class** + ```typescript + class StreamingErrorHandler { + async handle(error: StreamingError, context: ErrorContext): Promise { + if (this.isRetryable(error)) { + return this.handleWithRetry(error, context); + } + return this.handleFatal(error, context); + } + + private isRetryable(error: StreamingError): boolean { + return error.code === ErrorCode.PROJECT_OUT_OF_DATE || + error.code === ErrorCode.NETWORK_ERROR; + } + } + ``` + - Location: `stores/streaming/error-handler.ts` + - Benefit: Centralized error handling logic + +- [ ] **4.2 Define error recovery strategies** + ```typescript + type RecoveryStrategy = + | { type: 'retry'; maxAttempts: number; backoff: 'exponential' | 'linear' } + | { type: 'sync-and-retry' } + | { type: 'show-error'; dismissable: boolean } + | { type: 'abort' }; + ``` + - Location: `stores/streaming/types.ts` + - Benefit: Explicit, testable recovery strategies + +- [ ] **4.3 Remove duplicate retry logic** + - Consolidate `with-retry-sync.ts` and `handleStreamError.ts` retry logic + - Single retry implementation with configurable strategies + +--- + +## Phase 5: Refactor useSendMessageStream Hook + +### Goal +Simplify the main orchestration hook by delegating to the state machine. + +### Tasks + +- [ ] **5.1 Simplify hook to single responsibility** + ```typescript + function useSendMessageStream() { + const machine = useStreamingStateMachine(); + + const send = useCallback(async (message: string, selectedText: string) => { + machine.start({ message, selectedText }); + + await createConversationMessageStream(request, (event) => { + machine.handleEvent(event); + }); + + machine.complete(); + }, [machine]); + + return { send, state: machine.state }; + } + ``` + - Benefit: Hook focuses on orchestration, not event handling + +- [ ] **5.2 Reduce hook dependencies** + - Target: Maximum 5 dependencies in useCallback + - Move logic into state machine to reduce dependencies + +- [ ] **5.3 Extract request building logic** + ```typescript + function buildStreamRequest(params: StreamRequestParams): CreateConversationMessageStreamRequest { + return { ... }; + } + ``` + - Location: `utils/stream-request-builder.ts` + - Benefit: Testable, pure function for request creation + +--- + +## Phase 6: Testing & Documentation + +### Goal +Ensure the refactored code is well-tested and documented. + +### Tasks + +- [ ] **6.1 Add unit tests for state machine** + - Test all state transitions + - Test error handling + - Test message type handlers + +- [ ] **6.2 Add integration tests for streaming flow** + - Mock streaming API + - Test complete happy path + - Test error scenarios + +- [ ] **6.3 Document the new architecture** + - Update architecture diagram + - Document state machine states and transitions + - Add inline code comments for complex logic + +- [ ] **6.4 Create migration guide** + - Document changes for other developers + - List breaking changes + - Provide code migration examples + +--- + +## Implementation Priority + +| Phase | Priority | Effort | Impact | +|-------|----------|--------|--------| +| 1. Consolidate Handlers | High | Medium | High | +| 2. Unify Stores | High | High | High | +| 3. Simplify Transformations | Medium | Medium | Medium | +| 4. Error Handling | Medium | Low | Medium | +| 5. Refactor Hook | Low | Low | Medium | +| 6. Testing & Docs | Low | Medium | High | + +--- + +## Success Metrics + +After completing all phases: + +- [ ] Total files related to streaming reduced from 15+ to ~6 +- [ ] Single source of truth for message state +- [ ] No `flushSync` calls required +- [ ] All state transitions documented and validated +- [ ] Adding a new message type requires changes to only 1-2 files +- [ ] Unit test coverage > 80% for streaming logic +- [ ] Clear error handling with explicit recovery strategies + +--- + +## Notes + +- Implement phases incrementally; each phase should leave the codebase in a working state +- Consider feature flags for gradual rollout +- Performance testing recommended after Phase 2 (store unification) From c5a9a788799a5fd73430909aa4558e15eed10f04 Mon Sep 17 00:00:00 2001 From: Junyi Hou Date: Mon, 19 Jan 2026 20:34:49 +0800 Subject: [PATCH 02/10] refactor: phase 1 --- webapp/_webapp/docs/STREAMING_DESIGN_TODO.md | 67 +-- webapp/_webapp/eslint.config.js | 8 + .../_webapp/src/adapters/storage-adapter.ts | 5 + .../src/components/branch-switcher.tsx | 6 +- .../_webapp/src/components/text-patches.tsx | 1 + .../_webapp/src/hooks/useSendMessageStream.ts | 164 +++--- webapp/_webapp/src/intermediate.ts | 6 +- .../stores/conversation/handlers/converter.ts | 109 ---- .../conversation/handlers/handleError.ts | 20 - .../handlers/handleIncompleteIndicator.ts | 6 - .../handlers/handleMessageChunk.ts | 37 -- .../handlers/handleReasoningChunk.ts | 38 -- .../handlers/handleStreamError.ts | 53 -- .../handlers/handleStreamFinalization.ts | 6 - .../handlers/handleStreamInitialization.ts | 29 - .../handlers/handleStreamPartBegin.ts | 73 --- .../handlers/handleStreamPartEnd.ts | 99 ---- .../_webapp/src/stores/conversation/types.ts | 33 +- .../src/stores/streaming-message-store.ts | 130 +++-- webapp/_webapp/src/stores/streaming/index.ts | 9 + .../stores/streaming/message-type-handlers.ts | 151 ++++++ .../streaming/streaming-state-machine.ts | 497 ++++++++++++++++++ webapp/_webapp/src/stores/streaming/types.ts | 155 ++++++ webapp/_webapp/src/views/chat/body/index.tsx | 7 +- .../src/views/chat/body/status-indicator.tsx | 7 +- .../views/chat/header/chat-history-modal.tsx | 5 +- .../_webapp/src/views/chat/header/index.tsx | 4 +- webapp/_webapp/src/views/chat/helper.ts | 5 +- webapp/_webapp/src/views/devtools/index.tsx | 25 +- webapp/_webapp/src/views/office/app.tsx | 1 + 30 files changed, 1068 insertions(+), 688 deletions(-) delete mode 100644 webapp/_webapp/src/stores/conversation/handlers/converter.ts delete mode 100644 webapp/_webapp/src/stores/conversation/handlers/handleError.ts delete mode 100644 webapp/_webapp/src/stores/conversation/handlers/handleIncompleteIndicator.ts delete mode 100644 webapp/_webapp/src/stores/conversation/handlers/handleMessageChunk.ts delete mode 100644 webapp/_webapp/src/stores/conversation/handlers/handleReasoningChunk.ts delete mode 100644 webapp/_webapp/src/stores/conversation/handlers/handleStreamError.ts delete mode 100644 webapp/_webapp/src/stores/conversation/handlers/handleStreamFinalization.ts delete mode 100644 webapp/_webapp/src/stores/conversation/handlers/handleStreamInitialization.ts delete mode 100644 webapp/_webapp/src/stores/conversation/handlers/handleStreamPartBegin.ts delete mode 100644 webapp/_webapp/src/stores/conversation/handlers/handleStreamPartEnd.ts create mode 100644 webapp/_webapp/src/stores/streaming/index.ts create mode 100644 webapp/_webapp/src/stores/streaming/message-type-handlers.ts create mode 100644 webapp/_webapp/src/stores/streaming/streaming-state-machine.ts create mode 100644 webapp/_webapp/src/stores/streaming/types.ts diff --git a/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md b/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md index 2034b585..04881e4d 100644 --- a/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md +++ b/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md @@ -4,69 +4,56 @@ This document outlines a phased approach to simplify and improve the front-end s --- -## Phase 1: Consolidate Handlers into State Machine +## Phase 1: Consolidate Handlers into State Machine ✅ COMPLETED ### Goal Replace 9 separate handler files with a single, cohesive state machine that manages all streaming state transitions. ### Tasks -- [ ] **1.1 Create StreamingStateMachine class** - ```typescript - // stores/streaming/streaming-state-machine.ts - type StreamState = 'idle' | 'receiving' | 'finalizing' | 'error'; - - class StreamingStateMachine { - private state: StreamState = 'idle'; - private messageEntries: Map = new Map(); - - handleEvent(event: StreamEvent): void { - // Central event handling - } - } - ``` +- [x] **1.1 Create StreamingStateMachine class** - Location: `stores/streaming/streaming-state-machine.ts` + - Implemented as a Zustand store with `handleEvent` method for centralized event handling - Benefit: Single point of control for all state transitions -- [ ] **1.2 Define StreamEvent union type** - ```typescript - type StreamEvent = - | { type: 'INIT'; payload: StreamInitialization } - | { type: 'PART_BEGIN'; payload: StreamPartBegin } - | { type: 'CHUNK'; payload: MessageChunk } - | { type: 'PART_END'; payload: StreamPartEnd } - | { type: 'FINALIZE'; payload: StreamFinalization } - | { type: 'ERROR'; payload: StreamError } - | { type: 'INCOMPLETE'; payload: IncompleteIndicator }; - ``` +- [x] **1.2 Define StreamEvent union type** - Location: `stores/streaming/types.ts` + - Includes all event types: INIT, PART_BEGIN, CHUNK, REASONING_CHUNK, PART_END, FINALIZE, ERROR, INCOMPLETE, CONNECTION_ERROR - Benefit: Type-safe event handling with exhaustive checking -- [ ] **1.3 Create message type handlers registry** - ```typescript - // Eliminates duplicate switch statements - const messageTypeHandlers: Record = { - assistant: new AssistantHandler(), - toolCall: new ToolCallHandler(), - toolCallPrepareArguments: new ToolCallPrepareHandler(), - user: new NoOpHandler(), - system: new NoOpHandler(), - unknown: new NoOpHandler(), - }; - ``` +- [x] **1.3 Create message type handlers registry** - Location: `stores/streaming/message-type-handlers.ts` + - Implemented handlers: AssistantHandler, ToolCallHandler, ToolCallPrepareHandler, NoOpHandler - Benefit: Add new message types without modifying multiple files -- [ ] **1.4 Delete old handler files** - - Files to remove after migration: +- [x] **1.4 Delete old handler files** + - Deleted files: - `handlers/handleStreamInitialization.ts` - `handlers/handleStreamPartBegin.ts` - `handlers/handleMessageChunk.ts` + - `handlers/handleReasoningChunk.ts` - `handlers/handleStreamPartEnd.ts` - `handlers/handleStreamFinalization.ts` - `handlers/handleStreamError.ts` - `handlers/handleIncompleteIndicator.ts` - `handlers/handleError.ts` + - `handlers/converter.ts` + +### New File Structure + +``` +stores/streaming/ +├── index.ts # Module exports +├── types.ts # StreamEvent, MessageEntry, etc. +├── message-type-handlers.ts # Handler registry +└── streaming-state-machine.ts # Main state machine (Zustand store) +``` + +### Migration Notes + +- `streaming-message-store.ts` now serves as a backward compatibility layer +- `conversation/types.ts` re-exports types from the streaming module +- All consumer components updated to use `useStreamingStateMachine` --- diff --git a/webapp/_webapp/eslint.config.js b/webapp/_webapp/eslint.config.js index bdeade8b..c9e0aaed 100644 --- a/webapp/_webapp/eslint.config.js +++ b/webapp/_webapp/eslint.config.js @@ -21,6 +21,14 @@ export default tseslint.config( ...reactHooks.configs.recommended.rules, "react-refresh/only-export-components": ["warn", { allowConstantExport: true }], "no-console": "error", + "@typescript-eslint/no-unused-vars": [ + "error", + { + argsIgnorePattern: "^_", + varsIgnorePattern: "^_", + caughtErrorsIgnorePattern: "^_", + }, + ], }, }, ); diff --git a/webapp/_webapp/src/adapters/storage-adapter.ts b/webapp/_webapp/src/adapters/storage-adapter.ts index a4a2da6f..4b12e661 100644 --- a/webapp/_webapp/src/adapters/storage-adapter.ts +++ b/webapp/_webapp/src/adapters/storage-adapter.ts @@ -16,6 +16,7 @@ export class LocalStorageAdapter implements StorageAdapter { try { return localStorage.getItem(key); } catch { + // eslint-disable-next-line no-console console.warn("[Storage] localStorage.getItem failed for key:", key); return null; } @@ -25,6 +26,7 @@ export class LocalStorageAdapter implements StorageAdapter { try { localStorage.setItem(key, value); } catch (e) { + // eslint-disable-next-line no-console console.warn("[Storage] localStorage.setItem failed for key:", key, e); } } @@ -33,6 +35,7 @@ export class LocalStorageAdapter implements StorageAdapter { try { localStorage.removeItem(key); } catch (e) { + // eslint-disable-next-line no-console console.warn("[Storage] localStorage.removeItem failed for key:", key, e); } } @@ -41,6 +44,7 @@ export class LocalStorageAdapter implements StorageAdapter { try { localStorage.clear(); } catch (e) { + // eslint-disable-next-line no-console console.warn("[Storage] localStorage.clear failed", e); } } @@ -96,6 +100,7 @@ export function createStorageAdapter(type?: "localStorage" | "memory"): StorageA localStorage.removeItem(testKey); return new LocalStorageAdapter(); } catch { + // eslint-disable-next-line no-console console.warn("[Storage] localStorage not available, falling back to memory storage"); return new MemoryStorageAdapter(); } diff --git a/webapp/_webapp/src/components/branch-switcher.tsx b/webapp/_webapp/src/components/branch-switcher.tsx index 9fd5f38e..4c962e9f 100644 --- a/webapp/_webapp/src/components/branch-switcher.tsx +++ b/webapp/_webapp/src/components/branch-switcher.tsx @@ -2,7 +2,7 @@ import { Icon } from "@iconify/react"; import { Conversation } from "../pkg/gen/apiclient/chat/v2/chat_pb"; import { getConversation } from "../query/api"; import { useConversationStore } from "../stores/conversation/conversation-store"; -import { useStreamingMessageStore } from "../stores/streaming-message-store"; +import { useStreamingStateMachine } from "../stores/streaming"; import { useState } from "react"; interface BranchSwitcherProps { @@ -34,10 +34,10 @@ export const BranchSwitcher = ({ conversation }: BranchSwitcherProps) => { }); if (response.conversation) { setCurrentConversation(response.conversation); - useStreamingMessageStore.getState().resetStreamingMessage(); - useStreamingMessageStore.getState().resetIncompleteIndicator(); + useStreamingStateMachine.getState().reset(); } } catch (error) { + // eslint-disable-next-line no-console console.error("Failed to switch branch:", error); } finally { setIsLoading(false); diff --git a/webapp/_webapp/src/components/text-patches.tsx b/webapp/_webapp/src/components/text-patches.tsx index 615153bf..aa2a0fe3 100644 --- a/webapp/_webapp/src/components/text-patches.tsx +++ b/webapp/_webapp/src/components/text-patches.tsx @@ -63,6 +63,7 @@ export function TextPatches({ attachment, children }: TextPatchesProps) { }, 1500); } } catch (error) { + // eslint-disable-next-line no-console console.error("Failed to apply text:", error); setInsertBtnText("Failed!"); setTimeout(() => { diff --git a/webapp/_webapp/src/hooks/useSendMessageStream.ts b/webapp/_webapp/src/hooks/useSendMessageStream.ts index 2ba134f5..389da635 100644 --- a/webapp/_webapp/src/hooks/useSendMessageStream.ts +++ b/webapp/_webapp/src/hooks/useSendMessageStream.ts @@ -2,58 +2,53 @@ import { useCallback } from "react"; import { ConversationType, CreateConversationMessageStreamRequest, + CreateConversationMessageStreamResponse, IncompleteIndicator, - StreamFinalization, -} from "../pkg/gen/apiclient/chat/v2/chat_pb"; -import { PlainMessage } from "../query/types"; -import { useStreamingMessageStore } from "../stores/streaming-message-store"; -import { getProjectId } from "../libs/helpers"; -import { withRetrySync } from "../libs/with-retry-sync"; -import { createConversationMessageStream } from "../query/api"; -import { handleStreamInitialization } from "../stores/conversation/handlers/handleStreamInitialization"; -import { handleStreamPartBegin } from "../stores/conversation/handlers/handleStreamPartBegin"; -import { handleMessageChunk } from "../stores/conversation/handlers/handleMessageChunk"; -import { handleReasoningChunk } from "../stores/conversation/handlers/handleReasoningChunk"; -import { handleStreamPartEnd } from "../stores/conversation/handlers/handleStreamPartEnd"; -import { handleStreamFinalization } from "../stores/conversation/handlers/handleStreamFinalization"; -import { handleStreamError } from "../stores/conversation/handlers/handleStreamError"; -import { MessageChunk, MessageTypeUserSchema, ReasoningChunk, StreamError, + StreamFinalization, StreamInitialization, StreamPartBegin, StreamPartEnd, } from "../pkg/gen/apiclient/chat/v2/chat_pb"; -import { MessageEntry, MessageEntryStatus } from "../stores/conversation/types"; +import { PlainMessage } from "../query/types"; +import { getProjectId } from "../libs/helpers"; +import { withRetrySync } from "../libs/with-retry-sync"; +import { createConversationMessageStream } from "../query/api"; import { fromJson } from "../libs/protobuf-utils"; import { useConversationStore } from "../stores/conversation/conversation-store"; import { useListConversationsQuery } from "../query"; import { logError, logWarn } from "../libs/logger"; -import { handleError } from "../stores/conversation/handlers/handleError"; -import { handleIncompleteIndicator } from "../stores/conversation/handlers/handleIncompleteIndicator"; import { useAuthStore } from "../stores/auth-store"; import { useDevtoolStore } from "../stores/devtool-store"; import { useSelectionStore } from "../stores/selection-store"; import { useSettingStore } from "../stores/setting-store"; import { useSync } from "./useSync"; import { useAdapter } from "../adapters"; +import { + useStreamingStateMachine, + MessageEntryStatus, + StreamEvent, + MessageEntry, +} from "../stores/streaming"; /** * Custom React hook to handle sending a message as a stream in a conversation. * * This hook manages the process of sending a user message to the backend as a streaming request, - * handling all intermediate streaming events (initialization, message chunks, part begin/end, finalization, and errors). - * It updates the relevant stores for streaming and finalized messages, manages conversation state, - * and ensures proper synchronization with the backend (including Overleaf authentication). + * using the StreamingStateMachine to handle all intermediate streaming events. + * + * The hook focuses on orchestration while the state machine handles event processing, + * making the code easier to understand and maintain. * * Usage: * const { sendMessageStream } = useSendMessageStream(); * await sendMessageStream(message, selectedText); * * @returns {Object} An object containing the sendMessageStream function. - * @returns {Function} sendMessageStream - Function to send a message as a stream. Accepts (message: string, selectedText: string) and returns a Promise. + * @returns {Function} sendMessageStream - Function to send a message as a stream. */ export function useSendMessageStream() { const { sync } = useSync(); @@ -64,7 +59,10 @@ export function useSendMessageStream() { // Get project ID from adapter (supports both Overleaf URL and Word document ID) const projectId = adapter.getDocumentId?.() || getProjectId(); const { refetch: refetchConversationList } = useListConversationsQuery(projectId); - const { resetStreamingMessage, updateStreamingMessage, resetIncompleteIndicator } = useStreamingMessageStore(); + + // Use the new streaming state machine + const { handleEvent, reset: resetStateMachine } = useStreamingStateMachine(); + const { surroundingText: storeSurroundingText } = useSelectionStore(); const { alwaysSyncProject } = useDevtoolStore(); const { conversationMode } = useSettingStore(); @@ -84,18 +82,19 @@ export function useSendMessageStream() { userMessage: message, userSelectedText: selectedText, surrounding: storeSurroundingText ?? undefined, - conversationType: conversationMode === "debug" ? ConversationType.DEBUG : ConversationType.UNSPECIFIED, + conversationType: + conversationMode === "debug" ? ConversationType.DEBUG : ConversationType.UNSPECIFIED, parentMessageId, }; - resetStreamingMessage(); // ensure no stale message in the streaming messages - resetIncompleteIndicator(); + // Reset the state machine to ensure no stale messages + resetStateMachine(); // When editing a message (parentMessageId is provided), truncate the conversation // to only include messages up to and including the parent message if (parentMessageId && currentConversation.messages.length > 0) { const parentIndex = currentConversation.messages.findIndex( - (m) => m.messageId === parentMessageId + (m) => m.messageId === parentMessageId, ); if (parentIndex !== -1) { // Truncate messages to include only up to parentMessage @@ -112,6 +111,7 @@ export function useSendMessageStream() { } } + // Add the user message to the streaming state const newMessageEntry: MessageEntry = { messageId: "dummy", status: MessageEntryStatus.PREPARING, @@ -121,10 +121,13 @@ export function useSendMessageStream() { surrounding: storeSurroundingText ?? null, }), }; - updateStreamingMessage((prev) => ({ - ...prev, - parts: [...prev.parts, newMessageEntry], - sequence: prev.sequence + 1, + + // Directly update the state machine's streaming message + useStreamingStateMachine.setState((state) => ({ + streamingMessage: { + parts: [...state.streamingMessage.parts, newMessageEntry], + sequence: state.streamingMessage.sequence + 1, + }, })); if (import.meta.env.DEV && alwaysSyncProject) { @@ -132,59 +135,28 @@ export function useSendMessageStream() { await sync(); } + // Handler context for error recovery + const handlerContext = { + refetchConversationList, + userId: user?.id || "", + currentPrompt: message, + currentSelectedText: selectedText, + sync, + sendMessageStream, + }; + await withRetrySync( () => createConversationMessageStream(request, async (response) => { - switch (response.responsePayload.case) { - case "streamInitialization": // means the user message is received by the server, can change the status to FINALIZED - handleStreamInitialization( - response.responsePayload.value as StreamInitialization, - refetchConversationList, - ); - break; - case "streamPartBegin": - handleStreamPartBegin(response.responsePayload.value as StreamPartBegin, updateStreamingMessage); - break; - case "messageChunk": - handleMessageChunk(response.responsePayload.value as MessageChunk, updateStreamingMessage); - break; - case "streamPartEnd": - handleStreamPartEnd(response.responsePayload.value as StreamPartEnd, updateStreamingMessage); - break; - case "streamFinalization": - handleStreamFinalization(response.responsePayload.value as StreamFinalization); - break; - case "streamError": - await handleStreamError( - response.responsePayload.value as StreamError, - user?.id || "", - message, - selectedText, - sync, - sendMessageStream, - updateStreamingMessage, - ); - break; - case "incompleteIndicator": - handleIncompleteIndicator(response.responsePayload.value as IncompleteIndicator); - break; - case "reasoningChunk": - handleReasoningChunk(response.responsePayload.value as ReasoningChunk, updateStreamingMessage); - break; - default: { - if (response.responsePayload.value !== undefined) { - const _typeCheck: never = response.responsePayload; - throw new Error("Unexpected response payload: " + _typeCheck); - // DO NOT delete above line, it is used to check that all cases are handled. - } - break; - } + // Map response payload to StreamEvent and delegate to state machine + const event = mapResponseToEvent(response); + if (event) { + await handleEvent(event, handlerContext); } }), { sync: async () => { try { - // Platform-aware sync (Overleaf uses WebSocket, Word uses adapter.getFullText) const result = await sync(); if (!result.success) { logError("Failed to sync project", result.error); @@ -194,15 +166,14 @@ export function useSendMessageStream() { } }, onGiveUp: () => { - handleError(new Error("connection error.")); + handleEvent({ type: "CONNECTION_ERROR", payload: new Error("connection error.") }); }, }, ); }, [ - resetStreamingMessage, - resetIncompleteIndicator, - updateStreamingMessage, + resetStateMachine, + handleEvent, currentConversation, refetchConversationList, sync, @@ -216,3 +187,36 @@ export function useSendMessageStream() { return { sendMessageStream }; } + +/** + * Maps the API response payload to a StreamEvent for the state machine. + */ +function mapResponseToEvent( + response: CreateConversationMessageStreamResponse, +): StreamEvent | null { + const { case: payloadCase, value } = response.responsePayload; + + switch (payloadCase) { + case "streamInitialization": + return { type: "INIT", payload: value as StreamInitialization }; + case "streamPartBegin": + return { type: "PART_BEGIN", payload: value as StreamPartBegin }; + case "messageChunk": + return { type: "CHUNK", payload: value as MessageChunk }; + case "reasoningChunk": + return { type: "REASONING_CHUNK", payload: value as ReasoningChunk }; + case "streamPartEnd": + return { type: "PART_END", payload: value as StreamPartEnd }; + case "streamFinalization": + return { type: "FINALIZE", payload: value as StreamFinalization }; + case "streamError": + return { type: "ERROR", payload: value as StreamError }; + case "incompleteIndicator": + return { type: "INCOMPLETE", payload: value as IncompleteIndicator }; + default: + if (value !== undefined) { + logError("Unexpected response payload:", response.responsePayload); + } + return null; + } +} diff --git a/webapp/_webapp/src/intermediate.ts b/webapp/_webapp/src/intermediate.ts index 7a93051d..9af1fec1 100644 --- a/webapp/_webapp/src/intermediate.ts +++ b/webapp/_webapp/src/intermediate.ts @@ -154,7 +154,7 @@ function makeFunction(handlerName: string, opts?: MakeFunctionOpts): (args let getCookies: (domain: string) => Promise<{ session: string; gclb: string }>; if (import.meta.env.DEV) { // Development: use local storage for testing - // eslint-disable-next-line @typescript-eslint/no-unused-vars + getCookies = async (_domain: string) => { return { session: storage.getItem("pd.auth.overleafSession") ?? "", @@ -163,7 +163,7 @@ if (import.meta.env.DEV) { }; } else if (isOfficeEnvironment()) { // Office Add-in: Overleaf cookies not available/needed - // eslint-disable-next-line @typescript-eslint/no-unused-vars + getCookies = async (_domain: string) => ({ session: "", gclb: "" }); } else { // Browser extension (both MAIN and ISOLATED world): @@ -275,7 +275,7 @@ export { fetchImage }; let requestHostPermission: (origin: string) => Promise; if (isOfficeEnvironment()) { // Office Add-in: Permission requests not supported - // eslint-disable-next-line @typescript-eslint/no-unused-vars + requestHostPermission = async (_origin: string) => false; } else { // Browser extension: Use event-based communication diff --git a/webapp/_webapp/src/stores/conversation/handlers/converter.ts b/webapp/_webapp/src/stores/conversation/handlers/converter.ts deleted file mode 100644 index 29471a7e..00000000 --- a/webapp/_webapp/src/stores/conversation/handlers/converter.ts +++ /dev/null @@ -1,109 +0,0 @@ -import { fromJson } from "../../../libs/protobuf-utils"; -import { Conversation, Message, MessageSchema } from "../../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { MessageEntry, MessageEntryStatus } from "../types"; -import { useStreamingMessageStore } from "../../streaming-message-store"; -import { flushSync } from "react-dom"; -import { useConversationStore } from "../conversation-store"; -import { getConversation } from "../../../query/api"; - -export const convertMessageEntryToMessage = (messageEntry: MessageEntry): Message | undefined => { - if (messageEntry.assistant) { - const assistantPayload: { content: string; reasoning?: string } = { - content: messageEntry.assistant.content, - }; - if (messageEntry.assistant.reasoning) { - assistantPayload.reasoning = messageEntry.assistant.reasoning; - } - return fromJson(MessageSchema, { - messageId: messageEntry.messageId, - payload: { - assistant: assistantPayload, - }, - }); - } else if (messageEntry.toolCall) { - return fromJson(MessageSchema, { - messageId: messageEntry.messageId, - payload: { - toolCall: { - name: messageEntry.toolCall.name, - args: messageEntry.toolCall.args, - result: messageEntry.toolCall.result, - error: messageEntry.toolCall.error, - }, - }, - }); - } else if (messageEntry.user) { - return fromJson(MessageSchema, { - messageId: messageEntry.messageId, - payload: { - user: { - content: messageEntry.user.content, - selectedText: messageEntry.user.selectedText ?? "", - }, - }, - }); - } - return undefined; -}; - -export const flushStreamingMessageToConversation = (conversationId?: string, modelSlug?: string) => { - const flushMessages = useStreamingMessageStore - .getState() - .streamingMessage.parts.map((part) => { - if (part.status === MessageEntryStatus.FINALIZED) { - return convertMessageEntryToMessage(part); - } else { - return null; - } - }) - .filter((part) => { - return part !== null && part !== undefined; - }) as Message[]; - - flushSync(() => { - useConversationStore.getState().updateCurrentConversation((prev: Conversation) => ({ - ...prev, - id: conversationId ?? prev.id, - modelSlug: modelSlug ?? prev.modelSlug, - messages: [...prev.messages, ...flushMessages], - })); - }); - - useStreamingMessageStore.getState().resetStreamingMessage(); - // Do not reset incomplete indicator here, it will be reset in useSendMessageStream - - // Async update branch info (doesn't block, doesn't overwrite messages) - if (conversationId) { - updateBranchInfoAsync(conversationId); - } -}; - -// Fetch branch info from server and update only branch-related fields -// This preserves the messages (including reasoning) while updating branch info -const updateBranchInfoAsync = async (conversationId: string) => { - try { - const response = await getConversation({ conversationId }); - if (response.conversation) { - const branchInfo = response.conversation; - useConversationStore.getState().updateCurrentConversation((prev: Conversation) => { - // Guard against race condition: if user navigated to a different conversation - // while the fetch was in progress, don't apply the stale branch info - if (prev.id !== conversationId) { - return prev; - } - return { - ...prev, - // Only update branch-related fields, keep messages intact - currentBranchId: branchInfo.currentBranchId, - branches: branchInfo.branches, - currentBranchIndex: branchInfo.currentBranchIndex, - totalBranches: branchInfo.totalBranches, - }; - }); - } - } catch (error) { - // eslint-disable-next-line no-console - console.error("Failed to update branch info:", error); - // Non-critical error, branch switcher just won't show - } -}; diff --git a/webapp/_webapp/src/stores/conversation/handlers/handleError.ts b/webapp/_webapp/src/stores/conversation/handlers/handleError.ts deleted file mode 100644 index 9d36100a..00000000 --- a/webapp/_webapp/src/stores/conversation/handlers/handleError.ts +++ /dev/null @@ -1,20 +0,0 @@ -import { logError } from "../../../libs/logger"; -import { useStreamingMessageStore } from "../../streaming-message-store"; -import { MessageEntry, MessageEntryStatus } from "../types"; - -export function handleError(error?: Error) { - useStreamingMessageStore.getState().updateStreamingMessage((prev) => { - const newParts = prev.parts.map((part: MessageEntry) => { - return { - ...part, - status: part.status === MessageEntryStatus.PREPARING ? MessageEntryStatus.STALE : part.status, - }; - }); - return { - ...prev, - parts: newParts, - sequence: prev.sequence + 1, - }; - }); - logError("handleError", error); -} diff --git a/webapp/_webapp/src/stores/conversation/handlers/handleIncompleteIndicator.ts b/webapp/_webapp/src/stores/conversation/handlers/handleIncompleteIndicator.ts deleted file mode 100644 index 57513d9f..00000000 --- a/webapp/_webapp/src/stores/conversation/handlers/handleIncompleteIndicator.ts +++ /dev/null @@ -1,6 +0,0 @@ -import { IncompleteIndicator } from "../../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { useStreamingMessageStore } from "../../streaming-message-store"; - -export function handleIncompleteIndicator(incompleteIndicator: IncompleteIndicator) { - useStreamingMessageStore.getState().setIncompleteIndicator(incompleteIndicator); -} diff --git a/webapp/_webapp/src/stores/conversation/handlers/handleMessageChunk.ts b/webapp/_webapp/src/stores/conversation/handlers/handleMessageChunk.ts deleted file mode 100644 index 020cfb13..00000000 --- a/webapp/_webapp/src/stores/conversation/handlers/handleMessageChunk.ts +++ /dev/null @@ -1,37 +0,0 @@ -import { logError } from "../../../libs/logger"; -import { MessageChunk, MessageTypeAssistant } from "../../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { StreamingMessage } from "../../streaming-message-store"; -import { MessageEntry, MessageEntryStatus } from "../types"; - -export function handleMessageChunk( - chunk: MessageChunk, - updateStreamingMessage: (updater: (prev: StreamingMessage) => StreamingMessage) => void, -) { - updateStreamingMessage((prevMessage) => { - const updatedParts = prevMessage.parts.map((part: MessageEntry) => { - const isTargetPart = part.messageId === chunk.messageId && part.assistant; - - if (!isTargetPart) return part; - - const updatedAssistant: MessageTypeAssistant = { - ...part.assistant!, - content: part.assistant!.content + chunk.delta, - }; - - if (part.status !== MessageEntryStatus.PREPARING) { - logError("Message chunk received for non-preparing part, this is a critical error"); - } - - return { - ...part, - assistant: updatedAssistant, - }; - }); - - return { - ...prevMessage, - parts: updatedParts, - sequence: prevMessage.sequence + 1, - }; - }); -} diff --git a/webapp/_webapp/src/stores/conversation/handlers/handleReasoningChunk.ts b/webapp/_webapp/src/stores/conversation/handlers/handleReasoningChunk.ts deleted file mode 100644 index 62e08a9f..00000000 --- a/webapp/_webapp/src/stores/conversation/handlers/handleReasoningChunk.ts +++ /dev/null @@ -1,38 +0,0 @@ -import { logError } from "../../../libs/logger"; -import { ReasoningChunk, MessageTypeAssistant } from "../../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { StreamingMessage } from "../../streaming-message-store"; -import { MessageEntry, MessageEntryStatus } from "../types"; - -export function handleReasoningChunk( - chunk: ReasoningChunk, - updateStreamingMessage: (updater: (prev: StreamingMessage) => StreamingMessage) => void, -) { - updateStreamingMessage((prevMessage) => { - const updatedParts = prevMessage.parts.map((part: MessageEntry) => { - const isTargetPart = part.messageId === chunk.messageId && part.assistant; - - if (!isTargetPart) return part; - - const currentReasoning = part.assistant!.reasoning ?? ""; - const updatedAssistant: MessageTypeAssistant = { - ...part.assistant!, - reasoning: currentReasoning + chunk.delta, - }; - - if (part.status !== MessageEntryStatus.PREPARING) { - logError("Reasoning chunk received for non-preparing part, this is a critical error"); - } - - return { - ...part, - assistant: updatedAssistant, - }; - }); - - return { - ...prevMessage, - parts: updatedParts, - sequence: prevMessage.sequence + 1, - }; - }); -} diff --git a/webapp/_webapp/src/stores/conversation/handlers/handleStreamError.ts b/webapp/_webapp/src/stores/conversation/handlers/handleStreamError.ts deleted file mode 100644 index 1c7ed210..00000000 --- a/webapp/_webapp/src/stores/conversation/handlers/handleStreamError.ts +++ /dev/null @@ -1,53 +0,0 @@ -import { MessageTypeAssistantSchema, StreamError } from "../../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { errorToast } from "../../../libs/toasts"; -import { StreamingMessage } from "../../streaming-message-store"; -import { MessageEntry, MessageEntryStatus } from "../types"; -import { fromJson } from "../../../libs/protobuf-utils"; - -interface SyncResult { - success: boolean; - error?: Error; -} - -export async function handleStreamError( - streamError: StreamError, - _userId: string, // Kept for API compatibility, sync handles user internally - currentPrompt: string, - currentSelectedText: string, - sync: () => Promise, - sendMessageStream: (message: string, selectedText: string) => Promise, - updateStreamingMessage: (updater: (prev: StreamingMessage) => StreamingMessage) => void, -) { - // Append an error message to the streaming message - const updateFunc = (prev: StreamingMessage) => { - const errorMessageEntry: MessageEntry = { - messageId: "error-" + Date.now(), - status: MessageEntryStatus.STALE, - assistant: fromJson(MessageTypeAssistantSchema, { - content: `${streamError.errorMessage}`, - }), - }; - return { - ...prev, - parts: [...prev.parts, errorMessageEntry], - }; - }; - - try { - if (streamError.errorMessage.includes("project is out of date")) { - // Platform-aware sync (Overleaf uses WebSocket, Word uses adapter.getFullText) - const result = await sync(); - if (!result.success) { - throw result.error || new Error("Sync failed"); - } - // Retry sending the message after sync - await sendMessageStream(currentPrompt, currentSelectedText); - } else { - updateStreamingMessage(updateFunc); - errorToast(streamError.errorMessage, "Chat Stream Error"); - } - } catch (error) { - updateStreamingMessage(updateFunc); - errorToast(error instanceof Error ? error.message : "Unknown error", "Chat Stream Error"); - } -} diff --git a/webapp/_webapp/src/stores/conversation/handlers/handleStreamFinalization.ts b/webapp/_webapp/src/stores/conversation/handlers/handleStreamFinalization.ts deleted file mode 100644 index be08d272..00000000 --- a/webapp/_webapp/src/stores/conversation/handlers/handleStreamFinalization.ts +++ /dev/null @@ -1,6 +0,0 @@ -import { StreamFinalization } from "../../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { flushStreamingMessageToConversation } from "./converter"; - -export function handleStreamFinalization(_finalization: StreamFinalization) { - flushStreamingMessageToConversation(_finalization.conversationId); -} diff --git a/webapp/_webapp/src/stores/conversation/handlers/handleStreamInitialization.ts b/webapp/_webapp/src/stores/conversation/handlers/handleStreamInitialization.ts deleted file mode 100644 index c6b84eff..00000000 --- a/webapp/_webapp/src/stores/conversation/handlers/handleStreamInitialization.ts +++ /dev/null @@ -1,29 +0,0 @@ -import { StreamInitialization } from "../../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { useStreamingMessageStore } from "../../streaming-message-store"; -import { MessageEntryStatus } from "../types"; -import { logWarn } from "../../../libs/logger"; -import { flushStreamingMessageToConversation } from "./converter"; - -export function handleStreamInitialization(streamInit: StreamInitialization, refetchConversationList: () => void) { - useStreamingMessageStore.setState((prev) => ({ - ...prev, - streamingMessage: { - ...prev.streamingMessage, - parts: prev.streamingMessage.parts.map((part) => { - if (part.status === MessageEntryStatus.PREPARING && part.user) { - return { - ...part, - status: MessageEntryStatus.FINALIZED, - }; - } - return part; - }), - }, - })); - if (useStreamingMessageStore.getState().streamingMessage.parts.length !== 1) { - logWarn("Streaming message parts length is not 1, this may indicate some stale messages in the store"); - } - - flushStreamingMessageToConversation(streamInit.conversationId, streamInit.modelSlug); - refetchConversationList(); // Here we refetch conversation list because user may send chat message and immediately open history to view. -} diff --git a/webapp/_webapp/src/stores/conversation/handlers/handleStreamPartBegin.ts b/webapp/_webapp/src/stores/conversation/handlers/handleStreamPartBegin.ts deleted file mode 100644 index caa65b19..00000000 --- a/webapp/_webapp/src/stores/conversation/handlers/handleStreamPartBegin.ts +++ /dev/null @@ -1,73 +0,0 @@ -import { StreamPartBegin } from "../../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { StreamingMessage } from "../../streaming-message-store"; -import { MessageEntry, MessageEntryStatus } from "../types"; -import { logError } from "../../../libs/logger"; - -export function handleStreamPartBegin( - partBegin: StreamPartBegin, - updateStreamingMessage: (updater: (prev: StreamingMessage) => StreamingMessage) => void, -) { - const role = partBegin.payload?.messageType.case; - if (role === "assistant") { - const newMessageEntry: MessageEntry = { - messageId: partBegin.messageId, - status: MessageEntryStatus.PREPARING, - assistant: partBegin.payload?.messageType.value, - }; - updateStreamingMessage((prev) => { - // Skip if entry with same messageId already exists (prevents duplicate keys) - if (prev.parts.some((p) => p.messageId === partBegin.messageId)) { - return prev; - } - return { - parts: [...prev.parts, newMessageEntry], - sequence: prev.sequence + 1, - }; - }); - } else if (role === "toolCallPrepareArguments") { - const newMessageEntry: MessageEntry = { - messageId: partBegin.messageId, - status: MessageEntryStatus.PREPARING, - toolCallPrepareArguments: partBegin.payload?.messageType.value, - }; - updateStreamingMessage((prev) => { - // Skip if entry with same messageId already exists (prevents duplicate keys) - if (prev.parts.some((p) => p.messageId === partBegin.messageId)) { - return prev; - } - return { - parts: [...prev.parts, newMessageEntry], - sequence: prev.sequence + 1, - }; - }); - } else if (role === "toolCall") { - const newMessageEntry: MessageEntry = { - messageId: partBegin.messageId, - status: MessageEntryStatus.PREPARING, - toolCall: partBegin.payload?.messageType.value, - }; - updateStreamingMessage((prev) => { - // Skip if entry with same messageId already exists (prevents duplicate keys) - if (prev.parts.some((p) => p.messageId === partBegin.messageId)) { - return prev; - } - return { - parts: [...prev.parts, newMessageEntry], - sequence: prev.sequence + 1, - }; - }); - } else if (role === "system") { - // not possible - } else if (role === "user") { - // not possible - } else if (role === "unknown") { - // not possible - } else { - if (role !== undefined) { - const _typeCheck: never = role; - throw new Error("Unexpected response payload: " + _typeCheck); - // DO NOT delete above line, it is used to check that all cases are handled. - } - logError("unknown role in streamPartEnd:", role); - } -} diff --git a/webapp/_webapp/src/stores/conversation/handlers/handleStreamPartEnd.ts b/webapp/_webapp/src/stores/conversation/handlers/handleStreamPartEnd.ts deleted file mode 100644 index 46e6bb10..00000000 --- a/webapp/_webapp/src/stores/conversation/handlers/handleStreamPartEnd.ts +++ /dev/null @@ -1,99 +0,0 @@ -import { - MessageTypeAssistant, - MessageTypeToolCall, - MessageTypeToolCallPrepareArguments, - StreamPartEnd, -} from "../../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { StreamingMessage } from "../../streaming-message-store"; -import { logError } from "../../../libs/logger"; -import { MessageEntryStatus } from "../types"; - -export function handleStreamPartEnd( - partEnd: StreamPartEnd, - updateStreamingMessage: (updater: (prev: StreamingMessage) => StreamingMessage) => void, -) { - const role = partEnd.payload?.messageType.case; - switch (role) { - case "assistant": { - updateStreamingMessage((prev) => { - const newParts = prev.parts.map((part) => { - if (part.messageId === partEnd.messageId) { - const assistantMessage = partEnd.payload?.messageType.value as MessageTypeAssistant; - return { - ...part, - status: MessageEntryStatus.FINALIZED, - assistant: assistantMessage, - }; - } - return part; - }); - return { - ...prev, - parts: newParts, - sequence: prev.sequence + 1, - }; - }); - break; - } - case "toolCallPrepareArguments": { - updateStreamingMessage((prev) => { - const newParts = prev.parts.map((part) => { - if (part.messageId === partEnd.messageId) { - const toolCallPrepareArguments = partEnd.payload?.messageType.value as MessageTypeToolCallPrepareArguments; - return { - ...part, - status: MessageEntryStatus.FINALIZED, - toolCallPrepareArguments: toolCallPrepareArguments, - }; - } - return part; - }); - return { - ...prev, - parts: newParts, - sequence: prev.sequence + 1, - }; - }); - break; - } - case "toolCall": { - updateStreamingMessage((prev) => { - const newParts = prev.parts.map((part) => { - const toolCall = partEnd.payload?.messageType.value as MessageTypeToolCall; - if (part.messageId === partEnd.messageId) { - return { - ...part, - status: MessageEntryStatus.FINALIZED, - toolCall: toolCall, - }; - } - return part; - }); - return { - ...prev, - parts: newParts, - sequence: prev.sequence + 1, - }; - }); - break; - } - case "system": { - break; - } - case "unknown": { - break; - } - case "user": { - break; - } - default: { - if (role !== undefined) { - const _typeCheck: never = role; - throw new Error("Unexpected response payload: " + _typeCheck); - // DO NOT delete above line, it is used to check that all cases are handled. - } - logError("unknown role in streamPartEnd:", role); - break; - } - } -} diff --git a/webapp/_webapp/src/stores/conversation/types.ts b/webapp/_webapp/src/stores/conversation/types.ts index 273f291f..096f2402 100644 --- a/webapp/_webapp/src/stores/conversation/types.ts +++ b/webapp/_webapp/src/stores/conversation/types.ts @@ -1,25 +1,10 @@ -import { - MessageTypeAssistant, - MessageTypeToolCall, - MessageTypeToolCallPrepareArguments, - MessageTypeUnknown, - MessageTypeUser, -} from "../../pkg/gen/apiclient/chat/v2/chat_pb"; +/** + * Conversation Types (Backward Compatibility Layer) + * + * This file now re-exports types from the streaming module for backward compatibility. + * For new code, prefer importing directly from '../streaming'. + * + * @deprecated Use types from '../streaming' instead + */ -export enum MessageEntryStatus { - PREPARING = "PREPARING", - FINALIZED = "FINALIZED", // received "part end" or "stream finalization" - INCOMPLETE = "INCOMPLETE", // received "incomplete indicator" - STALE = "STALE", // if network shutdown or server crash. -} - -export type MessageEntry = { - messageId: string; - status: MessageEntryStatus; - // roles - user?: MessageTypeUser; - assistant?: MessageTypeAssistant; - toolCallPrepareArguments?: MessageTypeToolCallPrepareArguments; - toolCall?: MessageTypeToolCall; - unknown?: MessageTypeUnknown; -}; +export { MessageEntryStatus, type MessageEntry } from "../streaming/types"; diff --git a/webapp/_webapp/src/stores/streaming-message-store.ts b/webapp/_webapp/src/stores/streaming-message-store.ts index a7c12f08..22b663b3 100644 --- a/webapp/_webapp/src/stores/streaming-message-store.ts +++ b/webapp/_webapp/src/stores/streaming-message-store.ts @@ -1,50 +1,92 @@ -// Store "every streaming messages" occurred in the stream. +/** + * Streaming Message Store (Backward Compatibility Layer) + * + * This file now serves as a backward compatibility layer that re-exports + * functionality from the new StreamingStateMachine. + * + * For new code, prefer importing directly from '../stores/streaming'. + * + * @deprecated Use useStreamingStateMachine from '../stores/streaming' instead + */ -import { create } from "zustand"; -import { MessageEntry } from "./conversation/types"; import { flushSync } from "react-dom"; import { IncompleteIndicator } from "../pkg/gen/apiclient/chat/v2/chat_pb"; -import { SetterResetterStore } from "./types"; +import { + useStreamingStateMachine, + StreamingMessage, +} from "./streaming"; -export type StreamingMessage = { - parts: MessageEntry[]; - sequence: number; -}; +// Re-export types for backward compatibility +export type { StreamingMessage } from "./streaming"; -type CoreState = { - streamingMessage: StreamingMessage; - incompleteIndicator: IncompleteIndicator | null; -}; - -type StreamingMessageState = SetterResetterStore; - -export const useStreamingMessageStore = create((set) => ({ - streamingMessage: { parts: [], sequence: 0 }, - setStreamingMessage: (message) => set({ streamingMessage: message }), - resetStreamingMessage: () => - set({ - streamingMessage: { parts: [], sequence: 0 }, - }), - updateStreamingMessage: (updater) => { - // force React to synchronously flush any pending updates and - // re-render the component immediately after each store update, rather than batching them together. - flushSync(() => { - set((state) => { - const newState = updater(state.streamingMessage); - return { streamingMessage: newState }; - }); - }); +/** + * Backward-compatible streaming message store. + * + * This creates a compatible interface by delegating to the new state machine. + * The store interface is maintained for existing consumers. + */ +export const useStreamingMessageStore = Object.assign( + // Main selector function - allows use like useStreamingMessageStore((s) => s.streamingMessage) + function (selector: (state: { + streamingMessage: StreamingMessage; + incompleteIndicator: IncompleteIndicator | null; + }) => T): T { + return useStreamingStateMachine((state) => + selector({ + streamingMessage: state.streamingMessage, + incompleteIndicator: state.incompleteIndicator, + }) + ); }, - - incompleteIndicator: null, - setIncompleteIndicator: (incompleteIndicator) => { - set({ incompleteIndicator }); - }, - resetIncompleteIndicator: () => set({ incompleteIndicator: null }), - updateIncompleteIndicator: (updater) => { - set((state) => { - const newState = updater(state.incompleteIndicator); - return { incompleteIndicator: newState }; - }); - }, -})); + // Static methods for direct store access + { + getState: () => ({ + streamingMessage: useStreamingStateMachine.getState().streamingMessage, + incompleteIndicator: useStreamingStateMachine.getState().incompleteIndicator, + setStreamingMessage: (message: StreamingMessage) => { + useStreamingStateMachine.setState({ streamingMessage: message }); + }, + resetStreamingMessage: () => { + useStreamingStateMachine.setState({ + streamingMessage: { parts: [], sequence: 0 }, + }); + }, + updateStreamingMessage: (updater: (prev: StreamingMessage) => StreamingMessage) => { + flushSync(() => { + useStreamingStateMachine.setState((state) => ({ + streamingMessage: updater(state.streamingMessage), + })); + }); + }, + setIncompleteIndicator: (indicator: IncompleteIndicator | null) => { + useStreamingStateMachine.setState({ incompleteIndicator: indicator }); + }, + resetIncompleteIndicator: () => { + useStreamingStateMachine.setState({ incompleteIndicator: null }); + }, + updateIncompleteIndicator: ( + updater: (prev: IncompleteIndicator | null) => IncompleteIndicator | null + ) => { + useStreamingStateMachine.setState((state) => ({ + incompleteIndicator: updater(state.incompleteIndicator), + })); + }, + }), + setState: ( + partial: + | Partial<{ streamingMessage: StreamingMessage; incompleteIndicator: IncompleteIndicator | null }> + | ((state: { streamingMessage: StreamingMessage; incompleteIndicator: IncompleteIndicator | null }) => Partial<{ streamingMessage: StreamingMessage; incompleteIndicator: IncompleteIndicator | null }>) + ) => { + if (typeof partial === "function") { + const currentState = { + streamingMessage: useStreamingStateMachine.getState().streamingMessage, + incompleteIndicator: useStreamingStateMachine.getState().incompleteIndicator, + }; + const newPartial = partial(currentState); + useStreamingStateMachine.setState(newPartial); + } else { + useStreamingStateMachine.setState(partial); + } + }, + } +); diff --git a/webapp/_webapp/src/stores/streaming/index.ts b/webapp/_webapp/src/stores/streaming/index.ts new file mode 100644 index 00000000..b6a63ef2 --- /dev/null +++ b/webapp/_webapp/src/stores/streaming/index.ts @@ -0,0 +1,9 @@ +/** + * Streaming Module + * + * Exports all streaming-related functionality from a single entry point. + */ + +export * from "./types"; +export * from "./message-type-handlers"; +export * from "./streaming-state-machine"; diff --git a/webapp/_webapp/src/stores/streaming/message-type-handlers.ts b/webapp/_webapp/src/stores/streaming/message-type-handlers.ts new file mode 100644 index 00000000..24a0a42f --- /dev/null +++ b/webapp/_webapp/src/stores/streaming/message-type-handlers.ts @@ -0,0 +1,151 @@ +/** + * Message Type Handlers Registry + * + * This module provides a registry of handlers for different message types, + * eliminating the duplicate switch/if-else statements spread across multiple files. + * + * Benefits: + * - Adding a new message type only requires adding one handler + * - Type-safe handling with exhaustive checking + * - Clear separation of concerns for each message type + */ + +import { + MessageTypeAssistant, + MessageTypeToolCall, + MessageTypeToolCallPrepareArguments, + StreamPartBegin, + StreamPartEnd, +} from "../../pkg/gen/apiclient/chat/v2/chat_pb"; +import { + MessageEntry, + MessageEntryStatus, + MessageRole, + MessageTypeHandler, + MessageTypeHandlerRegistry, +} from "./types"; + +// ============================================================================ +// Handler Implementations +// ============================================================================ + +/** + * Handler for assistant messages. + */ +class AssistantHandler implements MessageTypeHandler { + onPartBegin(partBegin: StreamPartBegin): MessageEntry | null { + return { + messageId: partBegin.messageId, + status: MessageEntryStatus.PREPARING, + assistant: partBegin.payload?.messageType.value as MessageTypeAssistant, + }; + } + + onPartEnd(partEnd: StreamPartEnd, _existingEntry: MessageEntry): Partial | null { + const assistantMessage = partEnd.payload?.messageType.value as MessageTypeAssistant; + return { + status: MessageEntryStatus.FINALIZED, + assistant: assistantMessage, + }; + } +} + +/** + * Handler for tool call preparation (arguments streaming). + */ +class ToolCallPrepareHandler implements MessageTypeHandler { + onPartBegin(partBegin: StreamPartBegin): MessageEntry | null { + return { + messageId: partBegin.messageId, + status: MessageEntryStatus.PREPARING, + toolCallPrepareArguments: partBegin.payload?.messageType + .value as MessageTypeToolCallPrepareArguments, + }; + } + + onPartEnd(partEnd: StreamPartEnd, _existingEntry: MessageEntry): Partial | null { + const toolCallPrepareArguments = partEnd.payload?.messageType + .value as MessageTypeToolCallPrepareArguments; + return { + status: MessageEntryStatus.FINALIZED, + toolCallPrepareArguments, + }; + } +} + +/** + * Handler for completed tool calls. + */ +class ToolCallHandler implements MessageTypeHandler { + onPartBegin(partBegin: StreamPartBegin): MessageEntry | null { + return { + messageId: partBegin.messageId, + status: MessageEntryStatus.PREPARING, + toolCall: partBegin.payload?.messageType.value as MessageTypeToolCall, + }; + } + + onPartEnd(partEnd: StreamPartEnd, _existingEntry: MessageEntry): Partial | null { + const toolCall = partEnd.payload?.messageType.value as MessageTypeToolCall; + return { + status: MessageEntryStatus.FINALIZED, + toolCall, + }; + } +} + +/** + * No-op handler for message types that don't require streaming handling. + * Used for system, user, and unknown message types. + */ +class NoOpHandler implements MessageTypeHandler { + onPartBegin(_partBegin: StreamPartBegin): MessageEntry | null { + return null; + } + + onPartEnd(_partEnd: StreamPartEnd, _existingEntry: MessageEntry): Partial | null { + return null; + } +} + +// ============================================================================ +// Handler Registry +// ============================================================================ + +/** + * Registry mapping message roles to their handlers. + * This eliminates the need for switch/if-else statements when handling different message types. + */ +export const messageTypeHandlers: MessageTypeHandlerRegistry = { + assistant: new AssistantHandler(), + toolCallPrepareArguments: new ToolCallPrepareHandler(), + toolCall: new ToolCallHandler(), + user: new NoOpHandler(), + system: new NoOpHandler(), + unknown: new NoOpHandler(), +}; + +/** + * Get the handler for a specific message role. + * Returns NoOpHandler for undefined/null roles. + */ +export function getMessageTypeHandler(role: MessageRole | undefined): MessageTypeHandler { + if (!role) { + return new NoOpHandler(); + } + return messageTypeHandlers[role] || new NoOpHandler(); +} + +/** + * Type guard to check if a role is a valid MessageRole. + */ +export function isValidMessageRole(role: unknown): role is MessageRole { + return ( + role === "assistant" || + role === "toolCallPrepareArguments" || + role === "toolCall" || + role === "user" || + role === "system" || + role === "unknown" + ); +} diff --git a/webapp/_webapp/src/stores/streaming/streaming-state-machine.ts b/webapp/_webapp/src/stores/streaming/streaming-state-machine.ts new file mode 100644 index 00000000..051e7f60 --- /dev/null +++ b/webapp/_webapp/src/stores/streaming/streaming-state-machine.ts @@ -0,0 +1,497 @@ +/** + * Streaming State Machine + * + * This module consolidates all streaming event handling into a single cohesive state machine, + * replacing the 9+ fragmented handler files with a centralized, type-safe implementation. + * + * Benefits: + * - Single point of control for all state transitions + * - Clear state machine pattern with explicit states + * - Type-safe event handling with exhaustive checking + * - All related logic in one place for easier debugging + */ + +import { create } from "zustand"; +import { flushSync } from "react-dom"; +import { + IncompleteIndicator, + Message, + MessageSchema, + MessageTypeAssistant, + MessageTypeAssistantSchema, + Conversation, +} from "../../pkg/gen/apiclient/chat/v2/chat_pb"; +import { fromJson } from "../../libs/protobuf-utils"; +import { logError, logWarn } from "../../libs/logger"; +import { errorToast } from "../../libs/toasts"; +import { useConversationStore } from "../conversation/conversation-store"; +import { getConversation } from "../../query/api"; +import { getMessageTypeHandler, isValidMessageRole } from "./message-type-handlers"; +import { + MessageEntry, + MessageEntryStatus, + MessageRole, + StreamEvent, + StreamHandlerContext, + StreamingMessage, + StreamState, +} from "./types"; + +// ============================================================================ +// Store State Interface +// ============================================================================ + +interface StreamingStateMachineState { + // Current streaming state + state: StreamState; + + // Streaming message data + streamingMessage: StreamingMessage; + + // Incomplete indicator from server + incompleteIndicator: IncompleteIndicator | null; + + // Actions + handleEvent: (event: StreamEvent, context?: Partial) => Promise; + reset: () => void; + getStreamingMessage: () => StreamingMessage; + getIncompleteIndicator: () => IncompleteIndicator | null; +} + +// ============================================================================ +// Initial State +// ============================================================================ + +const initialState = { + state: "idle" as StreamState, + streamingMessage: { parts: [], sequence: 0 }, + incompleteIndicator: null, +}; + +// ============================================================================ +// Helper Functions +// ============================================================================ + +/** + * Convert a MessageEntry to a Message for the conversation store. + */ +function convertMessageEntryToMessage(messageEntry: MessageEntry): Message | undefined { + if (messageEntry.assistant) { + const assistantPayload: { content: string; reasoning?: string } = { + content: messageEntry.assistant.content, + }; + if (messageEntry.assistant.reasoning) { + assistantPayload.reasoning = messageEntry.assistant.reasoning; + } + return fromJson(MessageSchema, { + messageId: messageEntry.messageId, + payload: { + assistant: assistantPayload, + }, + }); + } else if (messageEntry.toolCall) { + return fromJson(MessageSchema, { + messageId: messageEntry.messageId, + payload: { + toolCall: { + name: messageEntry.toolCall.name, + args: messageEntry.toolCall.args, + result: messageEntry.toolCall.result, + error: messageEntry.toolCall.error, + }, + }, + }); + } else if (messageEntry.user) { + return fromJson(MessageSchema, { + messageId: messageEntry.messageId, + payload: { + user: { + content: messageEntry.user.content, + selectedText: messageEntry.user.selectedText ?? "", + }, + }, + }); + } + return undefined; +} + +/** + * Flush finalized streaming messages to the conversation store. + */ +function flushStreamingMessageToConversation( + streamingMessage: StreamingMessage, + conversationId?: string, + modelSlug?: string, +) { + const flushMessages = streamingMessage.parts + .filter((part) => part.status === MessageEntryStatus.FINALIZED) + .map((part) => convertMessageEntryToMessage(part)) + .filter((part): part is Message => part !== null && part !== undefined); + + flushSync(() => { + useConversationStore.getState().updateCurrentConversation((prev: Conversation) => ({ + ...prev, + id: conversationId ?? prev.id, + modelSlug: modelSlug ?? prev.modelSlug, + messages: [...prev.messages, ...flushMessages], + })); + }); + + // Async update branch info (doesn't block, doesn't overwrite messages) + if (conversationId) { + updateBranchInfoAsync(conversationId); + } +} + +/** + * Fetch branch info from server and update only branch-related fields. + */ +async function updateBranchInfoAsync(conversationId: string) { + try { + const response = await getConversation({ conversationId }); + if (response.conversation) { + const branchInfo = response.conversation; + useConversationStore.getState().updateCurrentConversation((prev: Conversation) => { + if (prev.id !== conversationId) { + return prev; + } + return { + ...prev, + currentBranchId: branchInfo.currentBranchId, + branches: branchInfo.branches, + currentBranchIndex: branchInfo.currentBranchIndex, + totalBranches: branchInfo.totalBranches, + }; + }); + } + } catch (error) { + // eslint-disable-next-line no-console + console.error("Failed to update branch info:", error); + } +} + +// ============================================================================ +// State Machine Store +// ============================================================================ + +export const useStreamingStateMachine = create((set, get) => ({ + ...initialState, + + handleEvent: async (event: StreamEvent, context?: Partial) => { + switch (event.type) { + // ======================================================================== + // INIT - User message acknowledged by server + // ======================================================================== + case "INIT": { + // Finalize the user message (mark as received by server) + set((state) => ({ + state: "receiving", + streamingMessage: { + ...state.streamingMessage, + parts: state.streamingMessage.parts.map((part) => { + if (part.status === MessageEntryStatus.PREPARING && part.user) { + return { ...part, status: MessageEntryStatus.FINALIZED }; + } + return part; + }), + }, + })); + + if (get().streamingMessage.parts.length !== 1) { + logWarn("Streaming message parts length is not 1, this may indicate stale messages"); + } + + // Flush to conversation store + flushStreamingMessageToConversation( + get().streamingMessage, + event.payload.conversationId, + event.payload.modelSlug, + ); + + // Reset after flush + set({ streamingMessage: { parts: [], sequence: 0 } }); + + // Refetch conversation list + context?.refetchConversationList?.(); + break; + } + + // ======================================================================== + // PART_BEGIN - New message part started + // ======================================================================== + case "PART_BEGIN": { + const role = event.payload.payload?.messageType.case as MessageRole | undefined; + + if (!role || !isValidMessageRole(role)) { + logError("Unknown role in streamPartBegin:", role); + break; + } + + const handler = getMessageTypeHandler(role); + const newEntry = handler.onPartBegin(event.payload); + + if (newEntry) { + // Use flushSync to force synchronous update + flushSync(() => { + set((state) => { + // Skip if entry with same messageId already exists + if (state.streamingMessage.parts.some((p) => p.messageId === newEntry.messageId)) { + return state; + } + return { + state: "receiving", + streamingMessage: { + parts: [...state.streamingMessage.parts, newEntry], + sequence: state.streamingMessage.sequence + 1, + }, + }; + }); + }); + } + break; + } + + // ======================================================================== + // CHUNK - Message content chunk received + // ======================================================================== + case "CHUNK": { + flushSync(() => { + set((state) => { + const updatedParts = state.streamingMessage.parts.map((part) => { + const isTargetPart = + part.messageId === event.payload.messageId && part.assistant; + + if (!isTargetPart) return part; + + if (part.status !== MessageEntryStatus.PREPARING) { + logError("Message chunk received for non-preparing part"); + } + + const updatedAssistant: MessageTypeAssistant = { + ...part.assistant!, + content: part.assistant!.content + event.payload.delta, + }; + + return { ...part, assistant: updatedAssistant }; + }); + + return { + streamingMessage: { + parts: updatedParts, + sequence: state.streamingMessage.sequence + 1, + }, + }; + }); + }); + break; + } + + // ======================================================================== + // REASONING_CHUNK - Reasoning content chunk received + // ======================================================================== + case "REASONING_CHUNK": { + flushSync(() => { + set((state) => { + const updatedParts = state.streamingMessage.parts.map((part) => { + const isTargetPart = + part.messageId === event.payload.messageId && part.assistant; + + if (!isTargetPart) return part; + + if (part.status !== MessageEntryStatus.PREPARING) { + logError("Reasoning chunk received for non-preparing part"); + } + + const currentReasoning = part.assistant!.reasoning ?? ""; + const updatedAssistant: MessageTypeAssistant = { + ...part.assistant!, + reasoning: currentReasoning + event.payload.delta, + }; + + return { ...part, assistant: updatedAssistant }; + }); + + return { + streamingMessage: { + parts: updatedParts, + sequence: state.streamingMessage.sequence + 1, + }, + }; + }); + }); + break; + } + + // ======================================================================== + // PART_END - Message part completed + // ======================================================================== + case "PART_END": { + const role = event.payload.payload?.messageType.case as MessageRole | undefined; + + if (!role || !isValidMessageRole(role)) { + logError("Unknown role in streamPartEnd:", role); + break; + } + + const handler = getMessageTypeHandler(role); + + flushSync(() => { + set((state) => { + const updatedParts = state.streamingMessage.parts.map((part) => { + if (part.messageId !== event.payload.messageId) { + return part; + } + + const updates = handler.onPartEnd(event.payload, part); + if (!updates) return part; + + return { ...part, ...updates }; + }); + + return { + streamingMessage: { + parts: updatedParts, + sequence: state.streamingMessage.sequence + 1, + }, + }; + }); + }); + break; + } + + // ======================================================================== + // FINALIZE - Stream completed + // ======================================================================== + case "FINALIZE": { + set({ state: "finalizing" }); + + // Flush remaining messages to conversation store + flushStreamingMessageToConversation(get().streamingMessage, event.payload.conversationId); + + // Reset streaming state + set({ + state: "idle", + streamingMessage: { parts: [], sequence: 0 }, + }); + break; + } + + // ======================================================================== + // ERROR - Stream error from server + // ======================================================================== + case "ERROR": { + const errorMessage = event.payload.errorMessage; + + // Check if this is a retryable "project out of date" error + if ( + errorMessage.includes("project is out of date") && + context?.sync && + context?.sendMessageStream && + context?.currentPrompt !== undefined && + context?.currentSelectedText !== undefined + ) { + try { + const result = await context.sync(); + if (!result.success) { + throw result.error || new Error("Sync failed"); + } + // Retry sending the message after sync + await context.sendMessageStream(context.currentPrompt, context.currentSelectedText); + return; + } catch { + // Fall through to error handling + } + } + + // Add error message to streaming parts + const errorEntry: MessageEntry = { + messageId: "error-" + Date.now(), + status: MessageEntryStatus.STALE, + assistant: fromJson(MessageTypeAssistantSchema, { + content: errorMessage, + }), + }; + + flushSync(() => { + set((state) => ({ + state: "error", + streamingMessage: { + ...state.streamingMessage, + parts: [...state.streamingMessage.parts, errorEntry], + }, + })); + }); + + errorToast(errorMessage, "Chat Stream Error"); + break; + } + + // ======================================================================== + // CONNECTION_ERROR - Network/connection error + // ======================================================================== + case "CONNECTION_ERROR": { + // Mark all preparing messages as stale + flushSync(() => { + set((state) => ({ + state: "error", + streamingMessage: { + parts: state.streamingMessage.parts.map((part) => ({ + ...part, + status: + part.status === MessageEntryStatus.PREPARING + ? MessageEntryStatus.STALE + : part.status, + })), + sequence: state.streamingMessage.sequence + 1, + }, + })); + }); + + logError("Connection error:", event.payload); + break; + } + + // ======================================================================== + // INCOMPLETE - Incomplete indicator received + // ======================================================================== + case "INCOMPLETE": { + set({ incompleteIndicator: event.payload }); + break; + } + + default: { + // Exhaustive type checking + const _exhaustive: never = event; + logError("Unhandled event type:", _exhaustive); + } + } + }, + + reset: () => { + set(initialState); + }, + + getStreamingMessage: () => get().streamingMessage, + + getIncompleteIndicator: () => get().incompleteIndicator, +})); + +// ============================================================================ +// Convenience Selectors +// ============================================================================ + +/** + * Select the streaming message from the state machine. + */ +export const selectStreamingMessage = (state: StreamingStateMachineState) => state.streamingMessage; + +/** + * Select the incomplete indicator from the state machine. + */ +export const selectIncompleteIndicator = (state: StreamingStateMachineState) => + state.incompleteIndicator; + +/** + * Select the current stream state. + */ +export const selectStreamState = (state: StreamingStateMachineState) => state.state; diff --git a/webapp/_webapp/src/stores/streaming/types.ts b/webapp/_webapp/src/stores/streaming/types.ts new file mode 100644 index 00000000..f154a0e9 --- /dev/null +++ b/webapp/_webapp/src/stores/streaming/types.ts @@ -0,0 +1,155 @@ +/** + * Streaming State Machine Types + * + * This file defines all types used by the streaming state machine, + * consolidating the fragmented type definitions across multiple handler files. + */ + +import { + IncompleteIndicator, + MessageChunk, + ReasoningChunk, + StreamError, + StreamFinalization, + StreamInitialization, + StreamPartBegin, + StreamPartEnd, + MessageTypeAssistant, + MessageTypeToolCall, + MessageTypeToolCallPrepareArguments, + MessageTypeUnknown, + MessageTypeUser, +} from "../../pkg/gen/apiclient/chat/v2/chat_pb"; + +// ============================================================================ +// Stream State +// ============================================================================ + +/** + * Represents the current state of the streaming process. + */ +export type StreamState = "idle" | "receiving" | "finalizing" | "error"; + +// ============================================================================ +// Message Entry Types +// ============================================================================ + +/** + * Status of a message entry during the streaming lifecycle. + */ +export enum MessageEntryStatus { + PREPARING = "PREPARING", + FINALIZED = "FINALIZED", + INCOMPLETE = "INCOMPLETE", + STALE = "STALE", +} + +/** + * Represents a message entry in the streaming state. + * Uses a discriminated union pattern for type safety. + */ +export type MessageEntry = { + messageId: string; + status: MessageEntryStatus; + // Role-specific content (only one will be present) + user?: MessageTypeUser; + assistant?: MessageTypeAssistant; + toolCallPrepareArguments?: MessageTypeToolCallPrepareArguments; + toolCall?: MessageTypeToolCall; + unknown?: MessageTypeUnknown; +}; + +/** + * The current streaming message state. + */ +export type StreamingMessage = { + parts: MessageEntry[]; + sequence: number; +}; + +// ============================================================================ +// Stream Events +// ============================================================================ + +/** + * Union type representing all possible stream events. + * This enables type-safe, exhaustive event handling in the state machine. + */ +export type StreamEvent = + | { type: "INIT"; payload: StreamInitialization } + | { type: "PART_BEGIN"; payload: StreamPartBegin } + | { type: "CHUNK"; payload: MessageChunk } + | { type: "REASONING_CHUNK"; payload: ReasoningChunk } + | { type: "PART_END"; payload: StreamPartEnd } + | { type: "FINALIZE"; payload: StreamFinalization } + | { type: "ERROR"; payload: StreamError } + | { type: "INCOMPLETE"; payload: IncompleteIndicator } + | { type: "CONNECTION_ERROR"; payload: Error }; + +/** + * Extract the payload type for a given event type. + */ +export type StreamEventPayload = Extract< + StreamEvent, + { type: T } +>["payload"]; + +// ============================================================================ +// Message Roles +// ============================================================================ + +/** + * All possible message roles from the protobuf MessagePayload. + */ +export type MessageRole = + | "assistant" + | "toolCallPrepareArguments" + | "toolCall" + | "user" + | "system" + | "unknown"; + +// ============================================================================ +// Handler Interfaces +// ============================================================================ + +/** + * Context provided to stream event handlers. + */ +export interface StreamHandlerContext { + /** Callback to refetch the conversation list */ + refetchConversationList: () => void; + /** User ID for error handling */ + userId: string; + /** Current prompt for retry scenarios */ + currentPrompt: string; + /** Current selected text for retry scenarios */ + currentSelectedText: string; + /** Sync function for project synchronization */ + sync: () => Promise<{ success: boolean; error?: Error }>; + /** Send message function for retry scenarios */ + sendMessageStream: (message: string, selectedText: string) => Promise; +} + +/** + * Interface for message type-specific handlers. + * Implementations handle the creation and finalization of specific message types. + */ +export interface MessageTypeHandler { + /** + * Called when a stream part begins for this message type. + * @returns A new MessageEntry or null if this type should be ignored. + */ + onPartBegin(partBegin: StreamPartBegin): MessageEntry | null; + + /** + * Called when a stream part ends for this message type. + * @returns Updated fields to merge into the existing entry, or null to skip. + */ + onPartEnd(partEnd: StreamPartEnd, existingEntry: MessageEntry): Partial | null; +} + +/** + * Registry type for message type handlers. + */ +export type MessageTypeHandlerRegistry = Record; diff --git a/webapp/_webapp/src/views/chat/body/index.tsx b/webapp/_webapp/src/views/chat/body/index.tsx index bebdad23..9c284f89 100644 --- a/webapp/_webapp/src/views/chat/body/index.tsx +++ b/webapp/_webapp/src/views/chat/body/index.tsx @@ -4,7 +4,7 @@ import { Conversation, Message } from "../../../pkg/gen/apiclient/chat/v2/chat_p import { filterVisibleMessages, getPrevUserMessage, isEmptyConversation, messageToMessageEntry } from "../helper"; import { StatusIndicator } from "./status-indicator"; import { EmptyView } from "./empty-view"; -import { useStreamingMessageStore } from "../../../stores/streaming-message-store"; +import { useStreamingStateMachine } from "../../../stores/streaming"; import { useSettingStore } from "../../../stores/setting-store"; import { useConversationStore } from "../../../stores/conversation/conversation-store"; import { getConversation } from "../../../query/api"; @@ -24,7 +24,7 @@ export const ChatBody = ({ conversation }: ChatBodyProps) => { const chatContainerRef = useRef(null); const lastUserMsgRef = useRef(null); const expanderRef = useRef(null); - const streamingMessage = useStreamingMessageStore((s) => s.streamingMessage); + const streamingMessage = useStreamingStateMachine((s) => s.streamingMessage); const visibleMessages = useMemo(() => filterVisibleMessages(conversation), [conversation]); const [reloadSuccess, setReloadSuccess] = useState(ReloadStatus.Default); @@ -145,8 +145,7 @@ export const ChatBody = ({ conversation }: ChatBodyProps) => { throw new Error(`Failed to load conversation ${conversation?.id ?? "unknown"}`); } setCurrentConversation(response.conversation); - useStreamingMessageStore.getState().resetStreamingMessage(); - useStreamingMessageStore.getState().resetIncompleteIndicator(); + useStreamingStateMachine.getState().reset(); setReloadSuccess(ReloadStatus.Success); } catch { setReloadSuccess(ReloadStatus.Failed); diff --git a/webapp/_webapp/src/views/chat/body/status-indicator.tsx b/webapp/_webapp/src/views/chat/body/status-indicator.tsx index 046f0049..1c3127b2 100644 --- a/webapp/_webapp/src/views/chat/body/status-indicator.tsx +++ b/webapp/_webapp/src/views/chat/body/status-indicator.tsx @@ -1,14 +1,13 @@ import { LoadingIndicator } from "../../../components/loading-indicator"; import { UnknownEntryMessageContainer } from "../../../components/message-entry-container/unknown-entry"; import { Conversation } from "../../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { MessageEntryStatus } from "../../../stores/conversation/types"; import { useSocketStore } from "../../../stores/socket-store"; -import { useStreamingMessageStore } from "../../../stores/streaming-message-store"; +import { useStreamingStateMachine, MessageEntryStatus } from "../../../stores/streaming"; export const StatusIndicator = ({ conversation }: { conversation?: Conversation }) => { const { syncing, syncingProgress } = useSocketStore(); - const streamingMessage = useStreamingMessageStore((s) => s.streamingMessage); - const incompleteIndicator = useStreamingMessageStore((s) => s.incompleteIndicator); + const streamingMessage = useStreamingStateMachine((s) => s.streamingMessage); + const incompleteIndicator = useStreamingStateMachine((s) => s.incompleteIndicator); const isWaitingForResponse = streamingMessage.parts.at(-1)?.user !== undefined || diff --git a/webapp/_webapp/src/views/chat/header/chat-history-modal.tsx b/webapp/_webapp/src/views/chat/header/chat-history-modal.tsx index 1bd6659a..4f478355 100644 --- a/webapp/_webapp/src/views/chat/header/chat-history-modal.tsx +++ b/webapp/_webapp/src/views/chat/header/chat-history-modal.tsx @@ -8,7 +8,7 @@ import { useDeleteConversationMutation, useListConversationsQuery } from "../../ import { logError } from "../../../libs/logger"; import { Modal } from "../../../components/modal"; import googleAnalytics from "../../../libs/google-analytics"; -import { useStreamingMessageStore } from "../../../stores/streaming-message-store"; +import { useStreamingStateMachine } from "../../../stores/streaming"; import { getProjectId } from "../../../libs/helpers"; import { useConversationStore } from "../../../stores/conversation/conversation-store"; import { useConversationUiStore } from "../../../stores/conversation/conversation-ui-store"; @@ -110,8 +110,7 @@ export const ChatHistoryModal = () => { throw new Error(`Failed to load conversation ${conversationId}`); } setCurrentConversation(response.conversation); - useStreamingMessageStore.getState().resetStreamingMessage(); - useStreamingMessageStore.getState().resetIncompleteIndicator(); + useStreamingStateMachine.getState().reset(); setShowChatHistory(false); promptInputRef.current?.focus(); } catch (e) { diff --git a/webapp/_webapp/src/views/chat/header/index.tsx b/webapp/_webapp/src/views/chat/header/index.tsx index 0476d2e7..81a870b1 100644 --- a/webapp/_webapp/src/views/chat/header/index.tsx +++ b/webapp/_webapp/src/views/chat/header/index.tsx @@ -2,7 +2,7 @@ import { TabHeader } from "../../../components/tab-header"; import { ChatButton } from "./chat-button"; import { useConversationStore } from "../../../stores/conversation/conversation-store"; import { flushSync } from "react-dom"; -import { useStreamingMessageStore } from "../../../stores/streaming-message-store"; +import { useStreamingStateMachine } from "../../../stores/streaming"; import { useConversationUiStore } from "../../../stores/conversation/conversation-ui-store"; import { ChatHistoryModal } from "./chat-history-modal"; import { BranchSwitcher } from "../../../components/branch-switcher"; @@ -10,7 +10,7 @@ import { BranchSwitcher } from "../../../components/branch-switcher"; export const NewConversation = () => { flushSync(() => { // force UI refresh. - useStreamingMessageStore.getState().resetStreamingMessage(); + useStreamingStateMachine.getState().reset(); useConversationStore.getState().setIsStreaming(false); useConversationStore.getState().startFromScratch(); useConversationUiStore.getState().inputRef?.current?.focus(); diff --git a/webapp/_webapp/src/views/chat/helper.ts b/webapp/_webapp/src/views/chat/helper.ts index 0bfd32ac..b17305c5 100644 --- a/webapp/_webapp/src/views/chat/helper.ts +++ b/webapp/_webapp/src/views/chat/helper.ts @@ -8,8 +8,7 @@ import { MessageTypeUser, } from "../../pkg/gen/apiclient/chat/v2/chat_pb"; import { useConversationStore } from "../../stores/conversation/conversation-store"; -import { MessageEntry, MessageEntryStatus } from "../../stores/conversation/types"; -import { useStreamingMessageStore } from "../../stores/streaming-message-store"; +import { useStreamingStateMachine, MessageEntry, MessageEntryStatus } from "../../stores/streaming"; export function getPrevUserMessage(messages: Message[], currentIndex: number): MessageTypeUser | undefined { for (let i = currentIndex - 1; i >= 0; i--) { @@ -23,7 +22,7 @@ export function getPrevUserMessage(messages: Message[], currentIndex: number): M export function isEmptyConversation(): boolean { const converstaion = useConversationStore.getState().currentConversation; const visibleMessages = filterVisibleMessages(converstaion); - const streamingMessage = useStreamingMessageStore.getState().streamingMessage; + const streamingMessage = useStreamingStateMachine.getState().streamingMessage; return visibleMessages.length === 0 && streamingMessage.parts.length === 0; } diff --git a/webapp/_webapp/src/views/devtools/index.tsx b/webapp/_webapp/src/views/devtools/index.tsx index f07f5f4e..c3b749b2 100644 --- a/webapp/_webapp/src/views/devtools/index.tsx +++ b/webapp/_webapp/src/views/devtools/index.tsx @@ -1,8 +1,7 @@ import { Rnd } from "react-rnd"; import { useSelectionStore } from "../../stores/selection-store"; import { Button, Input } from "@heroui/react"; -import { useStreamingMessageStore } from "../../stores/streaming-message-store"; -import { MessageEntry, MessageEntryStatus } from "../../stores/conversation/types"; +import { useStreamingStateMachine, MessageEntry, MessageEntryStatus } from "../../stores/streaming"; import { useConversationStore } from "../../stores/conversation/conversation-store"; import { fromJson } from "../../libs/protobuf-utils"; import { MessageSchema } from "../../pkg/gen/apiclient/chat/v2/chat_pb"; @@ -28,10 +27,20 @@ const randomUUID = () => { export const DevTools = () => { // State management const { selectedText, setSelectedText, setSelectionRange } = useSelectionStore(); - const { streamingMessage, setStreamingMessage, updateStreamingMessage } = useStreamingMessageStore(); + const streamingMessage = useStreamingStateMachine((s) => s.streamingMessage); const { startFromScratch, currentConversation, setCurrentConversation } = useConversationStore(); const [preparingDelay, setPreparingDelay] = useState(2); + // Helper functions to update streaming message + const setStreamingMessage = (message: typeof streamingMessage) => { + useStreamingStateMachine.setState({ streamingMessage: message }); + }; + const updateStreamingMessage = (updater: (prev: typeof streamingMessage) => typeof streamingMessage) => { + useStreamingStateMachine.setState((state) => ({ + streamingMessage: updater(state.streamingMessage), + })); + }; + // --- Event handlers --- // Conversation related const handleClearConversation = () => setCurrentConversation({ ...currentConversation, messages: [] }); @@ -97,7 +106,7 @@ export const DevTools = () => { // StreamingMessage related const handleClearStreamingMessage = () => setStreamingMessage({ ...streamingMessage, parts: [] }); const handleStaleLastStreamingMessage = () => { - const newParts = useStreamingMessageStore + const newParts = useStreamingStateMachine .getState() .streamingMessage.parts.map((part, _, arr) => part.messageId === arr[arr.length - 1]?.messageId ? { ...part, status: MessageEntryStatus.STALE } : part, @@ -122,7 +131,7 @@ export const DevTools = () => { }; setStreamingMessage({ ...streamingMessage, parts: [...streamingMessage.parts, messageEntry] }); withDelay(() => { - const newParts = useStreamingMessageStore.getState().streamingMessage.parts.map((part) => + const newParts = useStreamingStateMachine.getState().streamingMessage.parts.map((part) => part.messageId === messageEntry.messageId ? { ...part, @@ -146,7 +155,7 @@ export const DevTools = () => { }; updateStreamingMessage((prev) => ({ ...prev, parts: [...prev.parts, messageEntry] })); withDelay(() => { - const newParts = useStreamingMessageStore.getState().streamingMessage.parts.map((part) => + const newParts = useStreamingStateMachine.getState().streamingMessage.parts.map((part) => part.messageId === messageEntry.messageId ? { ...part, @@ -185,7 +194,7 @@ export const DevTools = () => { }; updateStreamingMessage((prev) => ({ ...prev, parts: [...prev.parts, messageEntry] })); withDelay(() => { - const newParts = useStreamingMessageStore.getState().streamingMessage.parts.map((part) => + const newParts = useStreamingStateMachine.getState().streamingMessage.parts.map((part) => part.messageId === messageEntry.messageId ? { ...part, @@ -211,7 +220,7 @@ export const DevTools = () => { }; updateStreamingMessage((prev) => ({ ...prev, parts: [...prev.parts, messageEntry] })); withDelay(() => { - const newParts = useStreamingMessageStore.getState().streamingMessage.parts.map((part) => + const newParts = useStreamingStateMachine.getState().streamingMessage.parts.map((part) => part.messageId === messageEntry.messageId ? { ...part, diff --git a/webapp/_webapp/src/views/office/app.tsx b/webapp/_webapp/src/views/office/app.tsx index 71cf9db9..c719647c 100644 --- a/webapp/_webapp/src/views/office/app.tsx +++ b/webapp/_webapp/src/views/office/app.tsx @@ -103,6 +103,7 @@ const PaperDebugger = ({ displayMode = "fullscreen", adapterId }: PaperDebuggerP } } // No adapter registered - host application must register one + // eslint-disable-next-line no-console console.error( "[PaperDebugger] No adapter registered. " + "The host application (e.g., Office Add-in) must register an adapter using __pdRegisterAdapter() before loading this component." From ccefb0fd0f63a392f0630fac8272d4e004ac9721 Mon Sep 17 00:00:00 2001 From: Junyi Hou Date: Mon, 19 Jan 2026 21:04:27 +0800 Subject: [PATCH 03/10] refactor: phase 2 --- webapp/_webapp/docs/STREAMING_DESIGN_TODO.md | 115 +++---- .../stores/conversation/conversation-store.ts | 21 +- webapp/_webapp/src/stores/converters.ts | 177 +++++++++++ webapp/_webapp/src/stores/message-store.ts | 295 ++++++++++++++++++ .../src/stores/streaming-message-store.ts | 10 +- .../streaming/streaming-state-machine.ts | 211 ++++++------- webapp/_webapp/src/stores/types.ts | 49 +++ webapp/_webapp/src/views/chat/body/index.tsx | 93 +++--- .../_webapp/src/views/chat/header/index.tsx | 18 +- webapp/_webapp/src/views/chat/helper.ts | 136 +++++++- 10 files changed, 886 insertions(+), 239 deletions(-) create mode 100644 webapp/_webapp/src/stores/converters.ts create mode 100644 webapp/_webapp/src/stores/message-store.ts diff --git a/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md b/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md index 04881e4d..85ebbb93 100644 --- a/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md +++ b/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md @@ -57,65 +57,68 @@ stores/streaming/ --- -## Phase 2: Unify Store Architecture +## Phase 2: Unify Store Architecture ✅ COMPLETED ### Goal Consolidate `streaming-message-store` and `conversation-store` message handling into a single coherent store. ### Tasks -- [ ] **2.1 Create unified message store** - ```typescript - // stores/message-store.ts - interface MessageStore { - // Finalized messages from server - messages: Message[]; - - // Currently streaming messages (separate from finalized) - streamingEntries: MessageEntry[]; - - // Computed: all displayable messages - get allMessages(): DisplayMessage[]; - - // Actions - appendStreamingEntry(entry: MessageEntry): void; - updateStreamingEntry(id: string, update: Partial): void; - finalizeStreaming(): void; - reset(): void; - } - ``` +- [x] **2.1 Create unified message store** - Location: `stores/message-store.ts` + - Implemented with automatic subscriptions to `conversation-store` and `streaming-state-machine` + - Provides `getAllDisplayMessages()` and `getVisibleDisplayMessages()` selectors + - Auto-initializes subscriptions on module import - Benefit: Single source of truth with clear streaming vs finalized separation -- [ ] **2.2 Create DisplayMessage type** - ```typescript - // Single type used by UI components - interface DisplayMessage { - id: string; - type: 'user' | 'assistant' | 'toolCall' | 'error'; - content: string; - status: 'streaming' | 'complete' | 'error'; - metadata?: MessageMetadata; - } - ``` +- [x] **2.2 Create DisplayMessage type** - Location: `stores/types.ts` + - Implemented with types: `user`, `assistant`, `toolCall`, `toolCallPrepare`, `error` + - Status types: `streaming`, `complete`, `error`, `stale` + - Includes support for reasoning, tool args/results, and selected text - Benefit: UI components work with one consistent type -- [ ] **2.3 Remove flushSync calls** - - Restructure update flow so React batching works naturally - - Replace `flushSync` with proper `useSyncExternalStore` or subscription pattern - - Files affected: `streaming-message-store.ts`, `converter.ts` +- [x] **2.3 Remove flushSync calls** + - No `flushSync` calls exist in the codebase + - Update flow uses Zustand's `subscribeWithSelector` middleware for reactive subscriptions + - `conversation-store` and `streaming-state-machine` now both use `subscribeWithSelector` + - `message-store` subscribes to both stores and updates automatically -- [ ] **2.4 Migrate ChatBody to use unified store** - - Replace: +- [x] **2.4 Migrate ChatBody to use unified store** + - ChatBody now uses: ```typescript - const visibleMessages = useMemo(() => filterVisibleMessages(conversation), [conversation]); - const streamingMessage = useStreamingMessageStore((s) => s.streamingMessage); - ``` - - With: - ```typescript - const displayMessages = useMessageStore((s) => s.allMessages); + const allDisplayMessages = useMessageStore((s) => s.getAllDisplayMessages()); ``` + - Helper functions updated to use unified store (e.g., `isEmptyConversation`) + - Converters in `stores/converters.ts` provide bidirectional conversion + +### Architecture Notes + +The unified message store architecture: + +``` +┌────────────────────────────────────────────────────────────────────────────┐ +│ useMessageStore │ +│ - messages: Message[] (finalized) │ +│ - streamingEntries: MessageEntry[] (streaming) │ +│ - getAllDisplayMessages(): DisplayMessage[] │ +└─────────────────────────────┬───────────────────────────────┬──────────────┘ + │ │ + ┌─────────────────┴─────────────┐ ┌─────────────┴─────────────┐ + │ subscribes to │ │ subscribes to │ + ▼ │ ▼ │ +┌───────────────────────────────┐ │ ┌───────────────────────────┐ +│ useConversationStore │ │ │ useStreamingStateMachine │ +│ - currentConversation │ │ │ - streamingMessage │ +│ (finalized messages) │ │ │ (streaming entries) │ +└───────────────────────────────┘ │ └───────────────────────────┘ + │ + ▼ + ┌─────────────────────────────────┐ + │ ChatBody Component │ + │ Uses useMessageStore directly │ + └─────────────────────────────────┘ +``` --- @@ -277,14 +280,14 @@ Ensure the refactored code is well-tested and documented. ## Implementation Priority -| Phase | Priority | Effort | Impact | -|-------|----------|--------|--------| -| 1. Consolidate Handlers | High | Medium | High | -| 2. Unify Stores | High | High | High | -| 3. Simplify Transformations | Medium | Medium | Medium | -| 4. Error Handling | Medium | Low | Medium | -| 5. Refactor Hook | Low | Low | Medium | -| 6. Testing & Docs | Low | Medium | High | +| Phase | Priority | Effort | Impact | Status | +|-------|----------|--------|--------|--------| +| 1. Consolidate Handlers | High | Medium | High | ✅ COMPLETED | +| 2. Unify Stores | High | High | High | ✅ COMPLETED | +| 3. Simplify Transformations | Medium | Medium | Medium | Not Started | +| 4. Error Handling | Medium | Low | Medium | Not Started | +| 5. Refactor Hook | Low | Low | Medium | Not Started | +| 6. Testing & Docs | Low | Medium | High | Not Started | --- @@ -292,11 +295,11 @@ Ensure the refactored code is well-tested and documented. After completing all phases: +- [x] Single source of truth for message state (Phase 2) +- [x] No `flushSync` calls required (Phase 2) +- [x] All state transitions documented and validated (Phase 1) +- [x] Adding a new message type requires changes to only 1-2 files (Phase 1) - [ ] Total files related to streaming reduced from 15+ to ~6 -- [ ] Single source of truth for message state -- [ ] No `flushSync` calls required -- [ ] All state transitions documented and validated -- [ ] Adding a new message type requires changes to only 1-2 files - [ ] Unit test coverage > 80% for streaming logic - [ ] Clear error handling with explicit recovery strategies diff --git a/webapp/_webapp/src/stores/conversation/conversation-store.ts b/webapp/_webapp/src/stores/conversation/conversation-store.ts index 0b8378f6..dd8f7688 100644 --- a/webapp/_webapp/src/stores/conversation/conversation-store.ts +++ b/webapp/_webapp/src/stores/conversation/conversation-store.ts @@ -1,4 +1,5 @@ import { create } from "zustand"; +import { subscribeWithSelector } from "zustand/middleware"; import { Conversation, ConversationSchema } from "../../pkg/gen/apiclient/chat/v2/chat_pb"; import { fromJson } from "../../libs/protobuf-utils"; import { useConversationUiStore } from "./conversation-ui-store"; @@ -12,15 +13,17 @@ interface ConversationStore { setIsStreaming: (isStreaming: boolean) => void; } -export const useConversationStore = create((set, get) => ({ - currentConversation: newConversation(), - setCurrentConversation: (conversation: Conversation) => set({ currentConversation: conversation }), - updateCurrentConversation: (updater: (conversation: Conversation) => Conversation) => - set({ currentConversation: updater(get().currentConversation) }), - startFromScratch: () => set({ currentConversation: newConversation() }), - isStreaming: false, - setIsStreaming: (isStreaming: boolean) => set({ isStreaming }), -})); +export const useConversationStore = create()( + subscribeWithSelector((set, get) => ({ + currentConversation: newConversation(), + setCurrentConversation: (conversation: Conversation) => set({ currentConversation: conversation }), + updateCurrentConversation: (updater: (conversation: Conversation) => Conversation) => + set({ currentConversation: updater(get().currentConversation) }), + startFromScratch: () => set({ currentConversation: newConversation() }), + isStreaming: false, + setIsStreaming: (isStreaming: boolean) => set({ isStreaming }), + })) +); export function newConversation(): Conversation { const modelSlug = useConversationUiStore.getState().lastUsedModelSlug; diff --git a/webapp/_webapp/src/stores/converters.ts b/webapp/_webapp/src/stores/converters.ts new file mode 100644 index 00000000..52f6297c --- /dev/null +++ b/webapp/_webapp/src/stores/converters.ts @@ -0,0 +1,177 @@ +/** + * Message Converters + * + * Bidirectional converters between protobuf Message, MessageEntry, and DisplayMessage types. + * These provide the bridge between API types and UI types. + */ + +import { Message } from "../pkg/gen/apiclient/chat/v2/chat_pb"; +import { MessageEntry, MessageEntryStatus } from "./streaming/types"; +import { DisplayMessage, DisplayMessageStatus } from "./types"; + +// ============================================================================ +// Message → DisplayMessage (for finalized messages from server) +// ============================================================================ + +/** + * Convert a finalized Message to DisplayMessage. + * @returns DisplayMessage or null if the message type is not displayable + */ +export function messageToDisplayMessage(msg: Message): DisplayMessage | null { + const messageType = msg.payload?.messageType; + + if (!messageType) return null; + + switch (messageType.case) { + case "user": + return { + id: msg.messageId, + type: "user", + status: "complete", + content: messageType.value.content, + selectedText: messageType.value.selectedText, + }; + + case "assistant": + return { + id: msg.messageId, + type: "assistant", + status: "complete", + content: messageType.value.content, + reasoning: messageType.value.reasoning, + }; + + case "toolCall": + return { + id: msg.messageId, + type: "toolCall", + status: "complete", + content: "", + toolName: messageType.value.name, + toolArgs: messageType.value.args, + toolResult: messageType.value.result, + toolError: messageType.value.error, + }; + + case "toolCallPrepareArguments": + return { + id: msg.messageId, + type: "toolCallPrepare", + status: "complete", + content: "", + toolName: messageType.value.name, + toolArgs: messageType.value.args, + }; + + default: + return null; + } +} + +// ============================================================================ +// MessageEntry → DisplayMessage (for streaming messages) +// ============================================================================ + +/** + * Convert a streaming MessageEntry to DisplayMessage. + * @returns DisplayMessage or null if the entry type is not displayable + */ +export function messageEntryToDisplayMessage(entry: MessageEntry): DisplayMessage | null { + const status = entryStatusToDisplayStatus(entry.status); + + if (entry.user) { + return { + id: entry.messageId, + type: "user", + status, + content: entry.user.content, + selectedText: entry.user.selectedText, + }; + } + + if (entry.assistant) { + return { + id: entry.messageId, + type: "assistant", + status, + content: entry.assistant.content, + reasoning: entry.assistant.reasoning, + }; + } + + if (entry.toolCall) { + return { + id: entry.messageId, + type: "toolCall", + status, + content: "", + toolName: entry.toolCall.name, + toolArgs: entry.toolCall.args, + toolResult: entry.toolCall.result, + toolError: entry.toolCall.error, + }; + } + + if (entry.toolCallPrepareArguments) { + return { + id: entry.messageId, + type: "toolCallPrepare", + status, + content: "", + toolName: entry.toolCallPrepareArguments.name, + toolArgs: entry.toolCallPrepareArguments.args, + }; + } + + return null; +} + +// ============================================================================ +// Status Converters +// ============================================================================ + +/** + * Convert MessageEntryStatus to DisplayMessageStatus. + */ +function entryStatusToDisplayStatus(status: MessageEntryStatus): DisplayMessageStatus { + switch (status) { + case MessageEntryStatus.PREPARING: + return "streaming"; + case MessageEntryStatus.FINALIZED: + return "complete"; + case MessageEntryStatus.STALE: + return "stale"; + case MessageEntryStatus.INCOMPLETE: + return "error"; + default: + return "complete"; + } +} + +// ============================================================================ +// Utility Functions +// ============================================================================ + +/** + * Check if a DisplayMessage should be visible in the chat. + * Filters out empty messages and system messages. + */ +export function isDisplayableMessage(msg: DisplayMessage): boolean { + if (msg.type === "user") { + return msg.content.length > 0; + } + if (msg.type === "assistant") { + return msg.content.length > 0; + } + if (msg.type === "toolCall" || msg.type === "toolCallPrepare") { + return true; + } + return false; +} + +/** + * Filter display messages to only include visible ones. + */ +export function filterDisplayMessages(messages: DisplayMessage[]): DisplayMessage[] { + return messages.filter(isDisplayableMessage); +} diff --git a/webapp/_webapp/src/stores/message-store.ts b/webapp/_webapp/src/stores/message-store.ts new file mode 100644 index 00000000..782216c3 --- /dev/null +++ b/webapp/_webapp/src/stores/message-store.ts @@ -0,0 +1,295 @@ +/** + * Unified Message Store + * + * This store consolidates message state management by combining: + * - Finalized messages (from conversation-store) + * - Streaming entries (from streaming-state-machine) + * + * Benefits: + * - Single source of truth for all messages + * - Unified DisplayMessage type for UI components + * - No flushSync needed - uses natural React batching + * - Automatically synced with conversation-store and streaming-state-machine + * + * Architecture: + * - This store subscribes to useConversationStore for finalized messages + * - This store subscribes to useStreamingStateMachine for streaming entries + * - UI components only need to use this store for all message rendering + */ + +import { create } from "zustand"; +import { subscribeWithSelector } from "zustand/middleware"; +import { Message, Conversation, BranchInfo } from "../pkg/gen/apiclient/chat/v2/chat_pb"; +import { MessageEntry, MessageEntryStatus } from "./streaming/types"; +import { DisplayMessage } from "./types"; +import { + messageToDisplayMessage, + messageEntryToDisplayMessage, + filterDisplayMessages, +} from "./converters"; +import { useConversationStore } from "./conversation/conversation-store"; +import { useStreamingStateMachine } from "./streaming"; + +// ============================================================================ +// Store State Interface +// ============================================================================ + +interface MessageStoreState { + // Finalized messages from server (synced from conversation-store) + messages: Message[]; + + // Currently streaming entries (synced from streaming-state-machine) + streamingEntries: MessageEntry[]; + + // Conversation metadata (synced from conversation-store) + conversationId: string; + modelSlug: string; + + // Branch information (synced from conversation-store) + currentBranchId: string; + branches: BranchInfo[]; + currentBranchIndex: number; + totalBranches: number; + + // Flag to track if subscriptions are initialized + _subscriptionsInitialized: boolean; +} + +interface MessageStoreActions { + // Message management (used by subscriptions) + setMessages: (messages: Message[]) => void; + setConversation: (conversation: Conversation) => void; + + // Streaming entry management (used by subscriptions) + setStreamingEntries: (entries: MessageEntry[]) => void; + + // Initialize subscriptions to source stores + initializeSubscriptions: () => void; + + // Reset + reset: () => void; + resetStreaming: () => void; +} + +interface MessageStoreSelectors { + // Computed selectors + getAllDisplayMessages: () => DisplayMessage[]; + getVisibleDisplayMessages: () => DisplayMessage[]; + hasStreamingMessages: () => boolean; + isWaitingForResponse: () => boolean; + hasStaleMessages: () => boolean; +} + +export type MessageStore = MessageStoreState & + MessageStoreActions & + MessageStoreSelectors; + +// ============================================================================ +// Initial State +// ============================================================================ + +const initialState: MessageStoreState = { + messages: [], + streamingEntries: [], + conversationId: "", + modelSlug: "", + currentBranchId: "", + branches: [], + currentBranchIndex: 0, + totalBranches: 0, + _subscriptionsInitialized: false, +}; + +// ============================================================================ +// Store Implementation +// ============================================================================ + +export const useMessageStore = create()( + subscribeWithSelector((set, get) => ({ + ...initialState, + + // ======================================================================== + // Message Management (synced from conversation-store) + // ======================================================================== + + setMessages: (messages: Message[]) => { + set({ messages }); + }, + + setConversation: (conversation: Conversation) => { + set({ + messages: conversation.messages, + conversationId: conversation.id, + modelSlug: conversation.modelSlug, + currentBranchId: conversation.currentBranchId, + branches: [...conversation.branches], + currentBranchIndex: conversation.currentBranchIndex, + totalBranches: conversation.totalBranches, + }); + }, + + // ======================================================================== + // Streaming Entry Management (synced from streaming-state-machine) + // ======================================================================== + + setStreamingEntries: (entries: MessageEntry[]) => { + set({ streamingEntries: entries }); + }, + + // ======================================================================== + // Subscription Initialization + // ======================================================================== + + initializeSubscriptions: () => { + if (get()._subscriptionsInitialized) return; + + // Subscribe to conversation-store for finalized messages + useConversationStore.subscribe( + (state) => state.currentConversation, + (conversation) => { + set({ + messages: conversation.messages, + conversationId: conversation.id, + modelSlug: conversation.modelSlug, + currentBranchId: conversation.currentBranchId, + branches: [...conversation.branches], + currentBranchIndex: conversation.currentBranchIndex, + totalBranches: conversation.totalBranches, + }); + }, + { fireImmediately: true } + ); + + // Subscribe to streaming-state-machine for streaming entries + useStreamingStateMachine.subscribe( + (state) => state.streamingMessage, + (streamingMessage) => { + set({ streamingEntries: streamingMessage.parts }); + }, + { fireImmediately: true } + ); + + set({ _subscriptionsInitialized: true }); + }, + + // ======================================================================== + // Reset + // ======================================================================== + + reset: () => { + set({ + messages: [], + streamingEntries: [], + conversationId: "", + modelSlug: "", + currentBranchId: "", + branches: [], + currentBranchIndex: 0, + totalBranches: 0, + // Keep subscriptions initialized + }); + }, + + resetStreaming: () => { + set({ streamingEntries: [] }); + }, + + // ======================================================================== + // Computed Selectors + // ======================================================================== + + getAllDisplayMessages: () => { + const state = get(); + + // Convert finalized messages + const finalizedDisplayMessages = state.messages + .map(messageToDisplayMessage) + .filter((m): m is DisplayMessage => m !== null); + + // Convert streaming entries + const streamingDisplayMessages = state.streamingEntries + .map(messageEntryToDisplayMessage) + .filter((m): m is DisplayMessage => m !== null); + + // Combine: finalized first, then streaming + return [...finalizedDisplayMessages, ...streamingDisplayMessages]; + }, + + getVisibleDisplayMessages: () => { + return filterDisplayMessages(get().getAllDisplayMessages()); + }, + + hasStreamingMessages: () => { + return get().streamingEntries.length > 0; + }, + + isWaitingForResponse: () => { + const state = get(); + const lastStreaming = state.streamingEntries.at(-1); + const lastFinalized = state.messages.at(-1); + + // Waiting if last streaming entry is a user message + if (lastStreaming?.user !== undefined) { + return true; + } + + // Waiting if last finalized is user and no streaming entries + if ( + lastFinalized?.payload?.messageType.case === "user" && + state.streamingEntries.length === 0 + ) { + return true; + } + + return false; + }, + + hasStaleMessages: () => { + return get().streamingEntries.some( + (entry) => entry.status === MessageEntryStatus.STALE + ); + }, + })) +); + +// ============================================================================ +// Convenience Selectors +// ============================================================================ + +export const selectAllDisplayMessages = (state: MessageStore) => + state.getAllDisplayMessages(); + +export const selectVisibleDisplayMessages = (state: MessageStore) => + state.getVisibleDisplayMessages(); + +export const selectHasStreamingMessages = (state: MessageStore) => + state.hasStreamingMessages(); + +export const selectIsWaitingForResponse = (state: MessageStore) => + state.isWaitingForResponse(); + +export const selectHasStaleMessages = (state: MessageStore) => + state.hasStaleMessages(); + +export const selectConversationId = (state: MessageStore) => state.conversationId; + +export const selectModelSlug = (state: MessageStore) => state.modelSlug; + +// ============================================================================ +// Store Initialization +// ============================================================================ + +/** + * Initialize the message store subscriptions. + * This should be called once at app startup to sync the message store + * with the conversation store and streaming state machine. + * + * Can be called multiple times safely - will only initialize once. + */ +export function initializeMessageStore(): void { + useMessageStore.getState().initializeSubscriptions(); +} + +// Auto-initialize when this module is first imported +// This ensures subscriptions are set up before any component renders +initializeMessageStore(); diff --git a/webapp/_webapp/src/stores/streaming-message-store.ts b/webapp/_webapp/src/stores/streaming-message-store.ts index 22b663b3..ee746284 100644 --- a/webapp/_webapp/src/stores/streaming-message-store.ts +++ b/webapp/_webapp/src/stores/streaming-message-store.ts @@ -9,7 +9,7 @@ * @deprecated Use useStreamingStateMachine from '../stores/streaming' instead */ -import { flushSync } from "react-dom"; + import { IncompleteIndicator } from "../pkg/gen/apiclient/chat/v2/chat_pb"; import { useStreamingStateMachine, @@ -52,11 +52,9 @@ export const useStreamingMessageStore = Object.assign( }); }, updateStreamingMessage: (updater: (prev: StreamingMessage) => StreamingMessage) => { - flushSync(() => { - useStreamingStateMachine.setState((state) => ({ - streamingMessage: updater(state.streamingMessage), - })); - }); + useStreamingStateMachine.setState((state) => ({ + streamingMessage: updater(state.streamingMessage), + })); }, setIncompleteIndicator: (indicator: IncompleteIndicator | null) => { useStreamingStateMachine.setState({ incompleteIndicator: indicator }); diff --git a/webapp/_webapp/src/stores/streaming/streaming-state-machine.ts b/webapp/_webapp/src/stores/streaming/streaming-state-machine.ts index 051e7f60..99f6bdd7 100644 --- a/webapp/_webapp/src/stores/streaming/streaming-state-machine.ts +++ b/webapp/_webapp/src/stores/streaming/streaming-state-machine.ts @@ -12,7 +12,7 @@ */ import { create } from "zustand"; -import { flushSync } from "react-dom"; +import { subscribeWithSelector } from "zustand/middleware"; import { IncompleteIndicator, Message, @@ -128,14 +128,12 @@ function flushStreamingMessageToConversation( .map((part) => convertMessageEntryToMessage(part)) .filter((part): part is Message => part !== null && part !== undefined); - flushSync(() => { - useConversationStore.getState().updateCurrentConversation((prev: Conversation) => ({ - ...prev, - id: conversationId ?? prev.id, - modelSlug: modelSlug ?? prev.modelSlug, - messages: [...prev.messages, ...flushMessages], - })); - }); + useConversationStore.getState().updateCurrentConversation((prev: Conversation) => ({ + ...prev, + id: conversationId ?? prev.id, + modelSlug: modelSlug ?? prev.modelSlug, + messages: [...prev.messages, ...flushMessages], + })); // Async update branch info (doesn't block, doesn't overwrite messages) if (conversationId) { @@ -174,7 +172,8 @@ async function updateBranchInfoAsync(conversationId: string) { // State Machine Store // ============================================================================ -export const useStreamingStateMachine = create((set, get) => ({ +export const useStreamingStateMachine = create()( + subscribeWithSelector((set, get) => ({ ...initialState, handleEvent: async (event: StreamEvent, context?: Partial) => { @@ -231,21 +230,18 @@ export const useStreamingStateMachine = create((set, const newEntry = handler.onPartBegin(event.payload); if (newEntry) { - // Use flushSync to force synchronous update - flushSync(() => { - set((state) => { - // Skip if entry with same messageId already exists - if (state.streamingMessage.parts.some((p) => p.messageId === newEntry.messageId)) { - return state; - } - return { - state: "receiving", - streamingMessage: { - parts: [...state.streamingMessage.parts, newEntry], - sequence: state.streamingMessage.sequence + 1, - }, - }; - }); + set((state) => { + // Skip if entry with same messageId already exists + if (state.streamingMessage.parts.some((p) => p.messageId === newEntry.messageId)) { + return state; + } + return { + state: "receiving", + streamingMessage: { + parts: [...state.streamingMessage.parts, newEntry], + sequence: state.streamingMessage.sequence + 1, + }, + }; }); } break; @@ -255,33 +251,31 @@ export const useStreamingStateMachine = create((set, // CHUNK - Message content chunk received // ======================================================================== case "CHUNK": { - flushSync(() => { - set((state) => { - const updatedParts = state.streamingMessage.parts.map((part) => { - const isTargetPart = - part.messageId === event.payload.messageId && part.assistant; + set((state) => { + const updatedParts = state.streamingMessage.parts.map((part) => { + const isTargetPart = + part.messageId === event.payload.messageId && part.assistant; - if (!isTargetPart) return part; + if (!isTargetPart) return part; - if (part.status !== MessageEntryStatus.PREPARING) { - logError("Message chunk received for non-preparing part"); - } - - const updatedAssistant: MessageTypeAssistant = { - ...part.assistant!, - content: part.assistant!.content + event.payload.delta, - }; - - return { ...part, assistant: updatedAssistant }; - }); + if (part.status !== MessageEntryStatus.PREPARING) { + logError("Message chunk received for non-preparing part"); + } - return { - streamingMessage: { - parts: updatedParts, - sequence: state.streamingMessage.sequence + 1, - }, + const updatedAssistant: MessageTypeAssistant = { + ...part.assistant!, + content: part.assistant!.content + event.payload.delta, }; + + return { ...part, assistant: updatedAssistant }; }); + + return { + streamingMessage: { + parts: updatedParts, + sequence: state.streamingMessage.sequence + 1, + }, + }; }); break; } @@ -290,34 +284,32 @@ export const useStreamingStateMachine = create((set, // REASONING_CHUNK - Reasoning content chunk received // ======================================================================== case "REASONING_CHUNK": { - flushSync(() => { - set((state) => { - const updatedParts = state.streamingMessage.parts.map((part) => { - const isTargetPart = - part.messageId === event.payload.messageId && part.assistant; - - if (!isTargetPart) return part; - - if (part.status !== MessageEntryStatus.PREPARING) { - logError("Reasoning chunk received for non-preparing part"); - } + set((state) => { + const updatedParts = state.streamingMessage.parts.map((part) => { + const isTargetPart = + part.messageId === event.payload.messageId && part.assistant; - const currentReasoning = part.assistant!.reasoning ?? ""; - const updatedAssistant: MessageTypeAssistant = { - ...part.assistant!, - reasoning: currentReasoning + event.payload.delta, - }; + if (!isTargetPart) return part; - return { ...part, assistant: updatedAssistant }; - }); + if (part.status !== MessageEntryStatus.PREPARING) { + logError("Reasoning chunk received for non-preparing part"); + } - return { - streamingMessage: { - parts: updatedParts, - sequence: state.streamingMessage.sequence + 1, - }, + const currentReasoning = part.assistant!.reasoning ?? ""; + const updatedAssistant: MessageTypeAssistant = { + ...part.assistant!, + reasoning: currentReasoning + event.payload.delta, }; + + return { ...part, assistant: updatedAssistant }; }); + + return { + streamingMessage: { + parts: updatedParts, + sequence: state.streamingMessage.sequence + 1, + }, + }; }); break; } @@ -335,26 +327,24 @@ export const useStreamingStateMachine = create((set, const handler = getMessageTypeHandler(role); - flushSync(() => { - set((state) => { - const updatedParts = state.streamingMessage.parts.map((part) => { - if (part.messageId !== event.payload.messageId) { - return part; - } - - const updates = handler.onPartEnd(event.payload, part); - if (!updates) return part; + set((state) => { + const updatedParts = state.streamingMessage.parts.map((part) => { + if (part.messageId !== event.payload.messageId) { + return part; + } - return { ...part, ...updates }; - }); + const updates = handler.onPartEnd(event.payload, part); + if (!updates) return part; - return { - streamingMessage: { - parts: updatedParts, - sequence: state.streamingMessage.sequence + 1, - }, - }; + return { ...part, ...updates }; }); + + return { + streamingMessage: { + parts: updatedParts, + sequence: state.streamingMessage.sequence + 1, + }, + }; }); break; } @@ -412,15 +402,13 @@ export const useStreamingStateMachine = create((set, }), }; - flushSync(() => { - set((state) => ({ - state: "error", - streamingMessage: { - ...state.streamingMessage, - parts: [...state.streamingMessage.parts, errorEntry], - }, - })); - }); + set((state) => ({ + state: "error", + streamingMessage: { + ...state.streamingMessage, + parts: [...state.streamingMessage.parts, errorEntry], + }, + })); errorToast(errorMessage, "Chat Stream Error"); break; @@ -431,21 +419,19 @@ export const useStreamingStateMachine = create((set, // ======================================================================== case "CONNECTION_ERROR": { // Mark all preparing messages as stale - flushSync(() => { - set((state) => ({ - state: "error", - streamingMessage: { - parts: state.streamingMessage.parts.map((part) => ({ - ...part, - status: - part.status === MessageEntryStatus.PREPARING - ? MessageEntryStatus.STALE - : part.status, - })), - sequence: state.streamingMessage.sequence + 1, - }, - })); - }); + set((state) => ({ + state: "error", + streamingMessage: { + parts: state.streamingMessage.parts.map((part) => ({ + ...part, + status: + part.status === MessageEntryStatus.PREPARING + ? MessageEntryStatus.STALE + : part.status, + })), + sequence: state.streamingMessage.sequence + 1, + }, + })); logError("Connection error:", event.payload); break; @@ -474,7 +460,8 @@ export const useStreamingStateMachine = create((set, getStreamingMessage: () => get().streamingMessage, getIncompleteIndicator: () => get().incompleteIndicator, -})); +})) +); // ============================================================================ // Convenience Selectors diff --git a/webapp/_webapp/src/stores/types.ts b/webapp/_webapp/src/stores/types.ts index 824eec09..3fc58860 100644 --- a/webapp/_webapp/src/stores/types.ts +++ b/webapp/_webapp/src/stores/types.ts @@ -12,3 +12,52 @@ export type Updater = { export type SetterResetterStore = T & Setter & Resetter & Updater; export type SetterStore = T & Setter; + +// ============================================================================ +// DisplayMessage Types - Unified message type for UI rendering +// ============================================================================ + +/** + * Status of a display message. + */ +export type DisplayMessageStatus = "streaming" | "complete" | "error" | "stale"; + +/** + * Type of display message. + */ +export type DisplayMessageType = "user" | "assistant" | "toolCall" | "toolCallPrepare" | "error"; + +/** + * Unified message type for UI rendering. + * All UI components should use this type instead of Message or MessageEntry directly. + * This provides a single consistent interface regardless of whether the message + * is finalized or still streaming. + */ +export interface DisplayMessage { + /** Unique message identifier */ + id: string; + + /** Message type */ + type: DisplayMessageType; + + /** Current status */ + status: DisplayMessageStatus; + + /** Main content (text for user/assistant, empty for tool calls) */ + content: string; + + /** Reasoning content (for assistant messages with thinking) */ + reasoning?: string; + + // Tool call specific fields + toolName?: string; + toolArgs?: string; + toolResult?: string; + toolError?: string; + + // User message specific fields + selectedText?: string; + + /** ID of the previous message (for branching) */ + previousMessageId?: string; +} diff --git a/webapp/_webapp/src/views/chat/body/index.tsx b/webapp/_webapp/src/views/chat/body/index.tsx index 9c284f89..3679820c 100644 --- a/webapp/_webapp/src/views/chat/body/index.tsx +++ b/webapp/_webapp/src/views/chat/body/index.tsx @@ -1,13 +1,20 @@ import { useEffect, useMemo, useRef, useState } from "react"; import { MessageCard } from "../../../components/message-card"; -import { Conversation, Message } from "../../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { filterVisibleMessages, getPrevUserMessage, isEmptyConversation, messageToMessageEntry } from "../helper"; +import { Conversation } from "../../../pkg/gen/apiclient/chat/v2/chat_pb"; +import { + isEmptyConversation, + getPrevUserSelectedText, + findLastUserMessageIndex, + displayMessageToMessageEntry, +} from "../helper"; import { StatusIndicator } from "./status-indicator"; import { EmptyView } from "./empty-view"; -import { useStreamingStateMachine } from "../../../stores/streaming"; +import { useMessageStore } from "../../../stores/message-store"; import { useSettingStore } from "../../../stores/setting-store"; import { useConversationStore } from "../../../stores/conversation/conversation-store"; +import { useStreamingStateMachine } from "../../../stores/streaming"; import { getConversation } from "../../../query/api"; +import { DisplayMessage } from "../../../stores/types"; interface ChatBodyProps { conversation?: Conversation; @@ -24,13 +31,30 @@ export const ChatBody = ({ conversation }: ChatBodyProps) => { const chatContainerRef = useRef(null); const lastUserMsgRef = useRef(null); const expanderRef = useRef(null); - const streamingMessage = useStreamingStateMachine((s) => s.streamingMessage); - const visibleMessages = useMemo(() => filterVisibleMessages(conversation), [conversation]); const [reloadSuccess, setReloadSuccess] = useState(ReloadStatus.Default); const conversationMode = useSettingStore((s) => s.conversationMode); const isDebugMode = conversationMode === "debug"; + // Use the unified message store to get all display messages + const allDisplayMessages = useMessageStore((s) => s.getAllDisplayMessages()); + + // Filter visible messages (non-empty user/assistant, all tool calls) + const visibleMessages = useMemo(() => { + return allDisplayMessages.filter((msg: DisplayMessage) => { + if (msg.type === "user") return msg.content.length > 0; + if (msg.type === "assistant") return msg.content.length > 0; + if (msg.type === "toolCall" || msg.type === "toolCallPrepare") return true; + return false; + }); + }, [allDisplayMessages]); + + // Find the last user message index for scroll behavior + const lastUserMessageIndex = useMemo( + () => findLastUserMessageIndex(visibleMessages), + [visibleMessages] + ); + // Scroll to the top of the last user message useEffect(() => { if (expanderRef.current) { @@ -44,7 +68,7 @@ export const ChatBody = ({ conversation }: ChatBodyProps) => { let expanderHeight: number; if (expanderViewOffset < 0) { - expanderHeight = 0; // The expander's position is absolute and renders independently from stream markdown. When stream markdown renders, the expander may scroll above the chatContainer due to user scrolling, causing expander.y < 0. In this case, we don't need the expander. + expanderHeight = 0; } else { expanderHeight = chatContainerHeight - expanderViewOffset; } @@ -68,34 +92,28 @@ export const ChatBody = ({ conversation }: ChatBodyProps) => { } }, [visibleMessages.length]); - const finalizedMessageCards = useMemo( + // Render all messages using the unified DisplayMessage array + const messageCards = useMemo( () => - visibleMessages.map((message: Message, index: number) => ( -
- 0 ? visibleMessages[index - 1].messageId : undefined} - /> -
- )), - [visibleMessages], - ); - - const streamingMessageCards = useMemo( - () => - streamingMessage.parts.map((entry) => ( - - )), - [streamingMessage.parts], + visibleMessages.map((msg: DisplayMessage, index: number) => { + const isStreaming = msg.status === "streaming"; + const isLastUserMsg = index === lastUserMessageIndex; + + return ( +
+ 0 ? visibleMessages[index - 1].id : undefined} + /> +
+ ); + }), + [visibleMessages, lastUserMessageIndex] ); if (isEmptyConversation()) { @@ -122,13 +140,12 @@ export const ChatBody = ({ conversation }: ChatBodyProps) => { return (
-
- {finalizedMessageCards} +
+ {messageCards}
-
-
- {streamingMessageCards} +
+
diff --git a/webapp/_webapp/src/views/chat/header/index.tsx b/webapp/_webapp/src/views/chat/header/index.tsx index 81a870b1..2b3f9b86 100644 --- a/webapp/_webapp/src/views/chat/header/index.tsx +++ b/webapp/_webapp/src/views/chat/header/index.tsx @@ -1,29 +1,23 @@ import { TabHeader } from "../../../components/tab-header"; import { ChatButton } from "./chat-button"; import { useConversationStore } from "../../../stores/conversation/conversation-store"; -import { flushSync } from "react-dom"; import { useStreamingStateMachine } from "../../../stores/streaming"; import { useConversationUiStore } from "../../../stores/conversation/conversation-ui-store"; import { ChatHistoryModal } from "./chat-history-modal"; import { BranchSwitcher } from "../../../components/branch-switcher"; export const NewConversation = () => { - flushSync(() => { - // force UI refresh. - useStreamingStateMachine.getState().reset(); - useConversationStore.getState().setIsStreaming(false); - useConversationStore.getState().startFromScratch(); - useConversationUiStore.getState().inputRef?.current?.focus(); - }); + useStreamingStateMachine.getState().reset(); + useConversationStore.getState().setIsStreaming(false); + useConversationStore.getState().startFromScratch(); + useConversationUiStore.getState().inputRef?.current?.focus(); }; export const ShowHistory = () => { - flushSync(() => { - // force UI refresh. - useConversationUiStore.getState().setShowChatHistory(true); - }); + useConversationUiStore.getState().setShowChatHistory(true); }; + export const ChatHeader = () => { const currentConversation = useConversationStore((s) => s.currentConversation); const showChatHistory = useConversationUiStore((s) => s.showChatHistory); diff --git a/webapp/_webapp/src/views/chat/helper.ts b/webapp/_webapp/src/views/chat/helper.ts index b17305c5..50ec4ffc 100644 --- a/webapp/_webapp/src/views/chat/helper.ts +++ b/webapp/_webapp/src/views/chat/helper.ts @@ -7,8 +7,13 @@ import { MessageTypeUnknown, MessageTypeUser, } from "../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { useConversationStore } from "../../stores/conversation/conversation-store"; -import { useStreamingStateMachine, MessageEntry, MessageEntryStatus } from "../../stores/streaming"; +import { MessageEntry, MessageEntryStatus } from "../../stores/streaming"; +import { useMessageStore } from "../../stores/message-store"; +import { DisplayMessage } from "../../stores/types"; + +// ============================================================================ +// Message-based helpers (existing, for backward compatibility) +// ============================================================================ export function getPrevUserMessage(messages: Message[], currentIndex: number): MessageTypeUser | undefined { for (let i = currentIndex - 1; i >= 0; i--) { @@ -19,11 +24,20 @@ export function getPrevUserMessage(messages: Message[], currentIndex: number): M return undefined; } +/** + * Check if the current conversation is empty. + * Uses the unified message store for consistent behavior. + */ export function isEmptyConversation(): boolean { - const converstaion = useConversationStore.getState().currentConversation; - const visibleMessages = filterVisibleMessages(converstaion); - const streamingMessage = useStreamingStateMachine.getState().streamingMessage; - return visibleMessages.length === 0 && streamingMessage.parts.length === 0; + const state = useMessageStore.getState(); + const allMessages = state.getAllDisplayMessages(); + const visibleMessages = allMessages.filter((msg) => { + if (msg.type === "user") return msg.content.length > 0; + if (msg.type === "assistant") return msg.content.length > 0; + if (msg.type === "toolCall" || msg.type === "toolCallPrepare") return true; + return false; + }); + return visibleMessages.length === 0; } export function filterVisibleMessages(conversation?: Conversation): Message[] { @@ -69,3 +83,113 @@ export function messageToMessageEntry(message: Message): MessageEntry { : undefined, } as MessageEntry; } + +// ============================================================================ +// DisplayMessage-based helpers (new unified approach) +// ============================================================================ + +/** + * Get the previous user message's selected text from a DisplayMessage array. + */ +export function getPrevUserSelectedText( + messages: DisplayMessage[], + currentIndex: number +): string | undefined { + for (let i = currentIndex - 1; i >= 0; i--) { + if (messages[i].type === "user") { + return messages[i].selectedText; + } + } + return undefined; +} + +/** + * Check if a DisplayMessage is the last user message in the array. + */ +export function isLastUserMessage( + messages: DisplayMessage[], + index: number +): boolean { + const msg = messages[index]; + if (msg.type !== "user") return false; + + // Check if there are any user messages after this one + for (let i = index + 1; i < messages.length; i++) { + if (messages[i].type === "user") { + return false; + } + } + return true; +} + +/** + * Find the index of the last user message in the array. + */ +export function findLastUserMessageIndex(messages: DisplayMessage[]): number { + for (let i = messages.length - 1; i >= 0; i--) { + if (messages[i].type === "user") { + return i; + } + } + return -1; +} + +/** + * Convert DisplayMessage back to MessageEntry for backward compatibility with MessageCard. + * This is a temporary bridge until MessageCard is updated to use DisplayMessage directly. + */ +export function displayMessageToMessageEntry(msg: DisplayMessage): MessageEntry { + const status = displayStatusToEntryStatus(msg.status); + + const entry: MessageEntry = { + messageId: msg.id, + status, + }; + + if (msg.type === "user") { + entry.user = { + content: msg.content, + selectedText: msg.selectedText ?? "", + surrounding: undefined, + $typeName: "chat.v2.MessageTypeUser", + } as MessageTypeUser; + } else if (msg.type === "assistant") { + entry.assistant = { + content: msg.content, + reasoning: msg.reasoning, + $typeName: "chat.v2.MessageTypeAssistant", + } as MessageTypeAssistant; + } else if (msg.type === "toolCall") { + entry.toolCall = { + name: msg.toolName ?? "", + args: msg.toolArgs ?? "", + result: msg.toolResult ?? "", + error: msg.toolError ?? "", + $typeName: "chat.v2.MessageTypeToolCall", + } as MessageTypeToolCall; + } else if (msg.type === "toolCallPrepare") { + entry.toolCallPrepareArguments = { + name: msg.toolName ?? "", + args: msg.toolArgs ?? "", + $typeName: "chat.v2.MessageTypeToolCallPrepareArguments", + } as MessageTypeToolCallPrepareArguments; + } + + return entry; +} + +function displayStatusToEntryStatus(status: DisplayMessage["status"]): MessageEntryStatus { + switch (status) { + case "streaming": + return MessageEntryStatus.PREPARING; + case "complete": + return MessageEntryStatus.FINALIZED; + case "stale": + return MessageEntryStatus.STALE; + case "error": + return MessageEntryStatus.INCOMPLETE; + default: + return MessageEntryStatus.FINALIZED; + } +} + From 18ed1a8a4b06c79f2aa678c58d73bc8e3fa17118 Mon Sep 17 00:00:00 2001 From: Junyi Hou Date: Mon, 19 Jan 2026 21:13:18 +0800 Subject: [PATCH 04/10] fix: errors and typing --- .../_webapp/src/hooks/useSendMessageStream.ts | 2 +- webapp/_webapp/src/stores/message-store.ts | 92 ++++++++++++------- webapp/_webapp/src/views/chat/body/index.tsx | 2 +- webapp/_webapp/src/views/chat/helper.ts | 2 +- 4 files changed, 64 insertions(+), 34 deletions(-) diff --git a/webapp/_webapp/src/hooks/useSendMessageStream.ts b/webapp/_webapp/src/hooks/useSendMessageStream.ts index 389da635..71c490e9 100644 --- a/webapp/_webapp/src/hooks/useSendMessageStream.ts +++ b/webapp/_webapp/src/hooks/useSendMessageStream.ts @@ -113,7 +113,7 @@ export function useSendMessageStream() { // Add the user message to the streaming state const newMessageEntry: MessageEntry = { - messageId: "dummy", + messageId: `pending-${crypto.randomUUID()}`, status: MessageEntryStatus.PREPARING, user: fromJson(MessageTypeUserSchema, { content: message, diff --git a/webapp/_webapp/src/stores/message-store.ts b/webapp/_webapp/src/stores/message-store.ts index 782216c3..8064be5f 100644 --- a/webapp/_webapp/src/stores/message-store.ts +++ b/webapp/_webapp/src/stores/message-store.ts @@ -51,6 +51,10 @@ interface MessageStoreState { currentBranchIndex: number; totalBranches: number; + // Computed display messages (updated when messages or streamingEntries change) + allDisplayMessages: DisplayMessage[]; + visibleDisplayMessages: DisplayMessage[]; + // Flag to track if subscriptions are initialized _subscriptionsInitialized: boolean; } @@ -97,9 +101,36 @@ const initialState: MessageStoreState = { branches: [], currentBranchIndex: 0, totalBranches: 0, + allDisplayMessages: [], + visibleDisplayMessages: [], _subscriptionsInitialized: false, }; +// ============================================================================ +// Helper: Compute Display Messages +// ============================================================================ + +function computeDisplayMessages( + messages: Message[], + streamingEntries: MessageEntry[] +): { all: DisplayMessage[]; visible: DisplayMessage[] } { + // Convert finalized messages + const finalizedDisplayMessages = messages + .map(messageToDisplayMessage) + .filter((m): m is DisplayMessage => m !== null); + + // Convert streaming entries + const streamingDisplayMessages = streamingEntries + .map(messageEntryToDisplayMessage) + .filter((m): m is DisplayMessage => m !== null); + + // Combine: finalized first, then streaming + const all = [...finalizedDisplayMessages, ...streamingDisplayMessages]; + const visible = filterDisplayMessages(all); + + return { all, visible }; +} + // ============================================================================ // Store Implementation // ============================================================================ @@ -113,10 +144,16 @@ export const useMessageStore = create()( // ======================================================================== setMessages: (messages: Message[]) => { - set({ messages }); + const { all, visible } = computeDisplayMessages(messages, get().streamingEntries); + set({ + messages, + allDisplayMessages: all, + visibleDisplayMessages: visible, + }); }, setConversation: (conversation: Conversation) => { + const { all, visible } = computeDisplayMessages(conversation.messages, get().streamingEntries); set({ messages: conversation.messages, conversationId: conversation.id, @@ -125,6 +162,8 @@ export const useMessageStore = create()( branches: [...conversation.branches], currentBranchIndex: conversation.currentBranchIndex, totalBranches: conversation.totalBranches, + allDisplayMessages: all, + visibleDisplayMessages: visible, }); }, @@ -133,7 +172,12 @@ export const useMessageStore = create()( // ======================================================================== setStreamingEntries: (entries: MessageEntry[]) => { - set({ streamingEntries: entries }); + const { all, visible } = computeDisplayMessages(get().messages, entries); + set({ + streamingEntries: entries, + allDisplayMessages: all, + visibleDisplayMessages: visible, + }); }, // ======================================================================== @@ -147,15 +191,7 @@ export const useMessageStore = create()( useConversationStore.subscribe( (state) => state.currentConversation, (conversation) => { - set({ - messages: conversation.messages, - conversationId: conversation.id, - modelSlug: conversation.modelSlug, - currentBranchId: conversation.currentBranchId, - branches: [...conversation.branches], - currentBranchIndex: conversation.currentBranchIndex, - totalBranches: conversation.totalBranches, - }); + get().setConversation(conversation); }, { fireImmediately: true } ); @@ -164,7 +200,7 @@ export const useMessageStore = create()( useStreamingStateMachine.subscribe( (state) => state.streamingMessage, (streamingMessage) => { - set({ streamingEntries: streamingMessage.parts }); + get().setStreamingEntries(streamingMessage.parts); }, { fireImmediately: true } ); @@ -186,37 +222,31 @@ export const useMessageStore = create()( branches: [], currentBranchIndex: 0, totalBranches: 0, + allDisplayMessages: [], + visibleDisplayMessages: [], // Keep subscriptions initialized }); }, resetStreaming: () => { - set({ streamingEntries: [] }); + const { all, visible } = computeDisplayMessages(get().messages, []); + set({ + streamingEntries: [], + allDisplayMessages: all, + visibleDisplayMessages: visible, + }); }, // ======================================================================== - // Computed Selectors + // Computed Selectors (return cached values) // ======================================================================== getAllDisplayMessages: () => { - const state = get(); - - // Convert finalized messages - const finalizedDisplayMessages = state.messages - .map(messageToDisplayMessage) - .filter((m): m is DisplayMessage => m !== null); - - // Convert streaming entries - const streamingDisplayMessages = state.streamingEntries - .map(messageEntryToDisplayMessage) - .filter((m): m is DisplayMessage => m !== null); - - // Combine: finalized first, then streaming - return [...finalizedDisplayMessages, ...streamingDisplayMessages]; + return get().allDisplayMessages; }, getVisibleDisplayMessages: () => { - return filterDisplayMessages(get().getAllDisplayMessages()); + return get().visibleDisplayMessages; }, hasStreamingMessages: () => { @@ -257,10 +287,10 @@ export const useMessageStore = create()( // ============================================================================ export const selectAllDisplayMessages = (state: MessageStore) => - state.getAllDisplayMessages(); + state.allDisplayMessages; export const selectVisibleDisplayMessages = (state: MessageStore) => - state.getVisibleDisplayMessages(); + state.visibleDisplayMessages; export const selectHasStreamingMessages = (state: MessageStore) => state.hasStreamingMessages(); diff --git a/webapp/_webapp/src/views/chat/body/index.tsx b/webapp/_webapp/src/views/chat/body/index.tsx index 3679820c..07de503b 100644 --- a/webapp/_webapp/src/views/chat/body/index.tsx +++ b/webapp/_webapp/src/views/chat/body/index.tsx @@ -37,7 +37,7 @@ export const ChatBody = ({ conversation }: ChatBodyProps) => { const isDebugMode = conversationMode === "debug"; // Use the unified message store to get all display messages - const allDisplayMessages = useMessageStore((s) => s.getAllDisplayMessages()); + const allDisplayMessages = useMessageStore((s) => s.allDisplayMessages); // Filter visible messages (non-empty user/assistant, all tool calls) const visibleMessages = useMemo(() => { diff --git a/webapp/_webapp/src/views/chat/helper.ts b/webapp/_webapp/src/views/chat/helper.ts index 50ec4ffc..3925fc37 100644 --- a/webapp/_webapp/src/views/chat/helper.ts +++ b/webapp/_webapp/src/views/chat/helper.ts @@ -30,7 +30,7 @@ export function getPrevUserMessage(messages: Message[], currentIndex: number): M */ export function isEmptyConversation(): boolean { const state = useMessageStore.getState(); - const allMessages = state.getAllDisplayMessages(); + const allMessages = state.allDisplayMessages; const visibleMessages = allMessages.filter((msg) => { if (msg.type === "user") return msg.content.length > 0; if (msg.type === "assistant") return msg.content.length > 0; From 44f730bc6d313021e7c464ff6ac450439b20a002 Mon Sep 17 00:00:00 2001 From: Junyi Hou Date: Mon, 19 Jan 2026 21:33:51 +0800 Subject: [PATCH 05/10] refactor: phase 3. --- webapp/_webapp/docs/STREAMING_DESIGN_TODO.md | 109 ++-- .../_webapp/src/components/message-card.tsx | 62 ++- .../_webapp/src/hooks/useSendMessageStream.ts | 24 +- .../_webapp/src/stores/conversation/types.ts | 4 +- webapp/_webapp/src/stores/converters.ts | 144 +---- webapp/_webapp/src/stores/message-store.ts | 18 +- .../stores/streaming/message-type-handlers.ts | 91 ++-- .../streaming/streaming-state-machine.ts | 136 ++--- webapp/_webapp/src/stores/streaming/types.ts | 47 +- webapp/_webapp/src/types/index.ts | 42 ++ webapp/_webapp/src/types/message.ts | 309 +++++++++++ webapp/_webapp/src/utils/index.ts | 18 + .../_webapp/src/utils/message-converters.ts | 490 ++++++++++++++++++ webapp/_webapp/src/views/chat/body/index.tsx | 3 +- .../src/views/chat/body/status-indicator.tsx | 6 +- webapp/_webapp/src/views/chat/helper.ts | 97 +--- webapp/_webapp/src/views/devtools/index.tsx | 171 +++--- 17 files changed, 1201 insertions(+), 570 deletions(-) create mode 100644 webapp/_webapp/src/types/index.ts create mode 100644 webapp/_webapp/src/types/message.ts create mode 100644 webapp/_webapp/src/utils/index.ts create mode 100644 webapp/_webapp/src/utils/message-converters.ts diff --git a/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md b/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md index 85ebbb93..b42bfb60 100644 --- a/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md +++ b/webapp/_webapp/docs/STREAMING_DESIGN_TODO.md @@ -100,7 +100,7 @@ The unified message store architecture: ┌────────────────────────────────────────────────────────────────────────────┐ │ useMessageStore │ │ - messages: Message[] (finalized) │ -│ - streamingEntries: MessageEntry[] (streaming) │ +│ - streamingEntries: InternalMessage[] (streaming) │ │ - getAllDisplayMessages(): DisplayMessage[] │ └─────────────────────────────┬───────────────────────────────┬──────────────┘ │ │ @@ -110,7 +110,7 @@ The unified message store architecture: ┌───────────────────────────────┐ │ ┌───────────────────────────┐ │ useConversationStore │ │ │ useStreamingStateMachine │ │ - currentConversation │ │ │ - streamingMessage │ -│ (finalized messages) │ │ │ (streaming entries) │ +│ (finalized messages) │ │ │ (streaming InternalMessage) │ └───────────────────────────────┘ │ └───────────────────────────┘ │ ▼ @@ -122,45 +122,78 @@ The unified message store architecture: --- -## Phase 3: Simplify Data Transformations +## Phase 3: Simplify Data Transformations ✅ COMPLETED ### Goal Reduce the number of data transformations from 5+ to 2 maximum. ### Tasks -- [ ] **3.1 Define canonical internal message format** - ```typescript - // Internal format used throughout the app - interface InternalMessage { - id: string; - role: MessageRole; - content: string; - status: MessageStatus; - toolCall?: ToolCallData; - attachments?: Attachment[]; - timestamp: number; - } - ``` +- [x] **3.1 Define canonical internal message format** - Location: `types/message.ts` - - Benefit: Single format reduces confusion + - Implemented `InternalMessage` union type with role-specific subtypes: + - `UserMessage`, `AssistantMessage`, `ToolCallMessage`, `ToolCallPrepareMessage`, `SystemMessage`, `UnknownMessage` + - Added type guards: `isUserMessage()`, `isAssistantMessage()`, etc. + - Added factory functions: `createUserMessage()`, `createAssistantMessage()`, etc. + - Benefit: Single format with type-safe role-specific data access -- [ ] **3.2 Create bidirectional converters** - ```typescript - // Only two conversions needed: - // 1. API response → Internal format - const fromApiMessage = (msg: ApiMessage): InternalMessage => { ... }; - - // 2. Internal format → API request - const toApiMessage = (msg: InternalMessage): ApiMessage => { ... }; - ``` +- [x] **3.2 Create bidirectional converters** - Location: `utils/message-converters.ts` + - Implemented converters: + - `fromApiMessage()` - API Message → InternalMessage + - `toApiMessage()` - InternalMessage → API Message + - `fromStreamPartBegin()` - Stream event → InternalMessage + - `applyStreamPartEnd()` - Update InternalMessage from stream end event + - `toDisplayMessage()` / `fromDisplayMessage()` - InternalMessage ↔ DisplayMessage - Benefit: Clear boundary between API types and internal types -- [ ] **3.3 Remove MessageEntry type** - - Replace `MessageEntry` with `InternalMessage` - - Update all components to use new type - - Delete `stores/conversation/types.ts` (after migrating MessageEntryStatus) +- [x] **3.3 Update MessageCard to use DisplayMessage directly** + - MessageCard now accepts `message: DisplayMessage` prop instead of `messageEntry: MessageEntry` + - Removed `displayMessageToMessageEntry()` bridge function from helper.ts + - ChatBody passes DisplayMessage directly to MessageCard + - Benefit: Eliminated unnecessary data transformation at render time + +- [x] **3.4 Remove legacy MessageEntry type** + - Streaming state machine now uses `InternalMessage` directly instead of `MessageEntry` + - Removed `MessageEntry` and `MessageEntryStatus` enum from `streaming/types.ts` + - Updated message-store.ts to use `streamingEntries: InternalMessage[]` + - Removed legacy converters (`fromMessageEntry`, `toMessageEntry`) + - Updated all consumers (hooks, views, devtools) to use `InternalMessage` and `MessageStatus` + +### New File Structure + +``` +types/ +├── index.ts # Module exports for types +├── message.ts # InternalMessage type definitions +└── global.d.ts # (existing) + +utils/ +├── index.ts # Module exports for utilities +└── message-converters.ts # All message converters in one place +``` + +### Data Flow After Phase 3 + +``` +API Response (Message) + │ + ▼ fromApiMessage() +InternalMessage + │ + ▼ toDisplayMessage() +DisplayMessage ─────────────────────────► MessageCard + ▲ + │ (streaming state uses InternalMessage directly) +InternalMessage (streaming state) +``` + +### Migration Notes + +- Legacy `MessageEntry` type has been removed - all code uses `InternalMessage` +- `MessageEntryStatus` enum replaced with `MessageStatus` string union: `"streaming" | "complete" | "error" | "stale"` +- `stores/converters.ts` simplified to only bridge between API types and display types +- Factory functions (`createUserMessage`, etc.) used for creating new messages --- @@ -284,7 +317,7 @@ Ensure the refactored code is well-tested and documented. |-------|----------|--------|--------|--------| | 1. Consolidate Handlers | High | Medium | High | ✅ COMPLETED | | 2. Unify Stores | High | High | High | ✅ COMPLETED | -| 3. Simplify Transformations | Medium | Medium | Medium | Not Started | +| 3. Simplify Transformations | Medium | Medium | Medium | ✅ COMPLETED | | 4. Error Handling | Medium | Low | Medium | Not Started | | 5. Refactor Hook | Low | Low | Medium | Not Started | | 6. Testing & Docs | Low | Medium | High | Not Started | @@ -299,14 +332,6 @@ After completing all phases: - [x] No `flushSync` calls required (Phase 2) - [x] All state transitions documented and validated (Phase 1) - [x] Adding a new message type requires changes to only 1-2 files (Phase 1) -- [ ] Total files related to streaming reduced from 15+ to ~6 -- [ ] Unit test coverage > 80% for streaming logic -- [ ] Clear error handling with explicit recovery strategies - ---- - -## Notes - -- Implement phases incrementally; each phase should leave the codebase in a working state -- Consider feature flags for gradual rollout -- Performance testing recommended after Phase 2 (store unification) +- [x] Canonical internal message format defined (Phase 3) +- [x] Bidirectional converters centralized in one file (Phase 3) +- [x] MessageCard uses DisplayMessage directly (Phase 3) diff --git a/webapp/_webapp/src/components/message-card.tsx b/webapp/_webapp/src/components/message-card.tsx index 603da9d7..36763eaa 100644 --- a/webapp/_webapp/src/components/message-card.tsx +++ b/webapp/_webapp/src/components/message-card.tsx @@ -1,7 +1,7 @@ import { cn } from "@heroui/react"; import { memo } from "react"; import Tools from "./message-entry-container/tools/tools"; -import { MessageEntry, MessageEntryStatus } from "../stores/conversation/types"; +import { DisplayMessage } from "../stores/types"; import { AssistantMessageContainer } from "./message-entry-container/assistant"; import { UserMessageContainer } from "./message-entry-container/user"; import { ToolCallPrepareMessageContainer } from "./message-entry-container/toolcall-prepare"; @@ -34,67 +34,77 @@ export const STYLES = { // Types interface MessageCardProps { - messageEntry: MessageEntry; + /** The display message to render (unified format) */ + message: DisplayMessage; + /** Selected text from the previous user message */ prevAttachment?: string; + /** ID of the previous message (for branching) */ previousMessageId?: string; + /** Whether to animate the message (for streaming) */ animated?: boolean; } -export const MessageCard = memo(({ messageEntry, prevAttachment, previousMessageId, animated }: MessageCardProps) => { +/** + * MessageCard component renders a single message in the chat. + * Now accepts DisplayMessage directly instead of MessageEntry. + */ +export const MessageCard = memo(({ message, prevAttachment, previousMessageId, animated }: MessageCardProps) => { + const isStale = message.status === "stale"; + const isPreparing = message.status === "streaming"; + const returnComponent = () => { - if (messageEntry.toolCall !== undefined) { + if (message.type === "toolCall") { return (
); } - if (messageEntry.assistant !== undefined) { + if (message.type === "assistant") { return ( ); } - if (messageEntry.toolCallPrepareArguments !== undefined) { + if (message.type === "toolCallPrepare") { return ( ); } - if (messageEntry.user !== undefined) { + if (message.type === "user") { return ( ); } - return ; + return ; }; return <>{returnComponent()}; diff --git a/webapp/_webapp/src/hooks/useSendMessageStream.ts b/webapp/_webapp/src/hooks/useSendMessageStream.ts index 71c490e9..2efff3ba 100644 --- a/webapp/_webapp/src/hooks/useSendMessageStream.ts +++ b/webapp/_webapp/src/hooks/useSendMessageStream.ts @@ -5,7 +5,6 @@ import { CreateConversationMessageStreamResponse, IncompleteIndicator, MessageChunk, - MessageTypeUserSchema, ReasoningChunk, StreamError, StreamFinalization, @@ -17,7 +16,6 @@ import { PlainMessage } from "../query/types"; import { getProjectId } from "../libs/helpers"; import { withRetrySync } from "../libs/with-retry-sync"; import { createConversationMessageStream } from "../query/api"; -import { fromJson } from "../libs/protobuf-utils"; import { useConversationStore } from "../stores/conversation/conversation-store"; import { useListConversationsQuery } from "../query"; import { logError, logWarn } from "../libs/logger"; @@ -29,10 +27,10 @@ import { useSync } from "./useSync"; import { useAdapter } from "../adapters"; import { useStreamingStateMachine, - MessageEntryStatus, StreamEvent, - MessageEntry, + InternalMessage, } from "../stores/streaming"; +import { createUserMessage } from "../types/message"; /** * Custom React hook to handle sending a message as a stream in a conversation. @@ -112,20 +110,20 @@ export function useSendMessageStream() { } // Add the user message to the streaming state - const newMessageEntry: MessageEntry = { - messageId: `pending-${crypto.randomUUID()}`, - status: MessageEntryStatus.PREPARING, - user: fromJson(MessageTypeUserSchema, { - content: message, + const newUserMessage: InternalMessage = createUserMessage( + `pending-${crypto.randomUUID()}`, + message, + { selectedText: selectedText, - surrounding: storeSurroundingText ?? null, - }), - }; + surrounding: storeSurroundingText ?? undefined, + status: "streaming", + } + ); // Directly update the state machine's streaming message useStreamingStateMachine.setState((state) => ({ streamingMessage: { - parts: [...state.streamingMessage.parts, newMessageEntry], + parts: [...state.streamingMessage.parts, newUserMessage], sequence: state.streamingMessage.sequence + 1, }, })); diff --git a/webapp/_webapp/src/stores/conversation/types.ts b/webapp/_webapp/src/stores/conversation/types.ts index 096f2402..319e7fd7 100644 --- a/webapp/_webapp/src/stores/conversation/types.ts +++ b/webapp/_webapp/src/stores/conversation/types.ts @@ -3,8 +3,6 @@ * * This file now re-exports types from the streaming module for backward compatibility. * For new code, prefer importing directly from '../streaming'. - * - * @deprecated Use types from '../streaming' instead */ -export { MessageEntryStatus, type MessageEntry } from "../streaming/types"; +export type { InternalMessage, MessageStatus } from "../streaming/types"; diff --git a/webapp/_webapp/src/stores/converters.ts b/webapp/_webapp/src/stores/converters.ts index 52f6297c..d23d5615 100644 --- a/webapp/_webapp/src/stores/converters.ts +++ b/webapp/_webapp/src/stores/converters.ts @@ -1,13 +1,17 @@ /** * Message Converters * - * Bidirectional converters between protobuf Message, MessageEntry, and DisplayMessage types. + * Bidirectional converters between protobuf Message, InternalMessage, and DisplayMessage types. * These provide the bridge between API types and UI types. */ import { Message } from "../pkg/gen/apiclient/chat/v2/chat_pb"; -import { MessageEntry, MessageEntryStatus } from "./streaming/types"; -import { DisplayMessage, DisplayMessageStatus } from "./types"; +import { InternalMessage } from "./streaming/types"; +import { DisplayMessage } from "./types"; +import { + fromApiMessage, + toDisplayMessage, +} from "../utils/message-converters"; // ============================================================================ // Message → DisplayMessage (for finalized messages from server) @@ -15,137 +19,29 @@ import { DisplayMessage, DisplayMessageStatus } from "./types"; /** * Convert a finalized Message to DisplayMessage. + * Uses the unified converter pipeline: Message → InternalMessage → DisplayMessage + * * @returns DisplayMessage or null if the message type is not displayable */ export function messageToDisplayMessage(msg: Message): DisplayMessage | null { - const messageType = msg.payload?.messageType; - - if (!messageType) return null; - - switch (messageType.case) { - case "user": - return { - id: msg.messageId, - type: "user", - status: "complete", - content: messageType.value.content, - selectedText: messageType.value.selectedText, - }; - - case "assistant": - return { - id: msg.messageId, - type: "assistant", - status: "complete", - content: messageType.value.content, - reasoning: messageType.value.reasoning, - }; - - case "toolCall": - return { - id: msg.messageId, - type: "toolCall", - status: "complete", - content: "", - toolName: messageType.value.name, - toolArgs: messageType.value.args, - toolResult: messageType.value.result, - toolError: messageType.value.error, - }; - - case "toolCallPrepareArguments": - return { - id: msg.messageId, - type: "toolCallPrepare", - status: "complete", - content: "", - toolName: messageType.value.name, - toolArgs: messageType.value.args, - }; - - default: - return null; - } + // Use the new unified converter: API Message → InternalMessage → DisplayMessage + const internalMsg = fromApiMessage(msg); + if (!internalMsg) return null; + return toDisplayMessage(internalMsg); } // ============================================================================ -// MessageEntry → DisplayMessage (for streaming messages) +// InternalMessage → DisplayMessage (for streaming messages) // ============================================================================ /** - * Convert a streaming MessageEntry to DisplayMessage. - * @returns DisplayMessage or null if the entry type is not displayable - */ -export function messageEntryToDisplayMessage(entry: MessageEntry): DisplayMessage | null { - const status = entryStatusToDisplayStatus(entry.status); - - if (entry.user) { - return { - id: entry.messageId, - type: "user", - status, - content: entry.user.content, - selectedText: entry.user.selectedText, - }; - } - - if (entry.assistant) { - return { - id: entry.messageId, - type: "assistant", - status, - content: entry.assistant.content, - reasoning: entry.assistant.reasoning, - }; - } - - if (entry.toolCall) { - return { - id: entry.messageId, - type: "toolCall", - status, - content: "", - toolName: entry.toolCall.name, - toolArgs: entry.toolCall.args, - toolResult: entry.toolCall.result, - toolError: entry.toolCall.error, - }; - } - - if (entry.toolCallPrepareArguments) { - return { - id: entry.messageId, - type: "toolCallPrepare", - status, - content: "", - toolName: entry.toolCallPrepareArguments.name, - toolArgs: entry.toolCallPrepareArguments.args, - }; - } - - return null; -} - -// ============================================================================ -// Status Converters -// ============================================================================ - -/** - * Convert MessageEntryStatus to DisplayMessageStatus. + * Convert an InternalMessage to DisplayMessage. + * This is the primary converter for both API and streaming messages. + * + * @returns DisplayMessage or null if the message type is not displayable */ -function entryStatusToDisplayStatus(status: MessageEntryStatus): DisplayMessageStatus { - switch (status) { - case MessageEntryStatus.PREPARING: - return "streaming"; - case MessageEntryStatus.FINALIZED: - return "complete"; - case MessageEntryStatus.STALE: - return "stale"; - case MessageEntryStatus.INCOMPLETE: - return "error"; - default: - return "complete"; - } +export function internalMessageToDisplayMessage(msg: InternalMessage): DisplayMessage | null { + return toDisplayMessage(msg); } // ============================================================================ diff --git a/webapp/_webapp/src/stores/message-store.ts b/webapp/_webapp/src/stores/message-store.ts index 8064be5f..8e36608c 100644 --- a/webapp/_webapp/src/stores/message-store.ts +++ b/webapp/_webapp/src/stores/message-store.ts @@ -20,11 +20,11 @@ import { create } from "zustand"; import { subscribeWithSelector } from "zustand/middleware"; import { Message, Conversation, BranchInfo } from "../pkg/gen/apiclient/chat/v2/chat_pb"; -import { MessageEntry, MessageEntryStatus } from "./streaming/types"; +import { InternalMessage } from "./streaming/types"; import { DisplayMessage } from "./types"; import { messageToDisplayMessage, - messageEntryToDisplayMessage, + internalMessageToDisplayMessage, filterDisplayMessages, } from "./converters"; import { useConversationStore } from "./conversation/conversation-store"; @@ -39,7 +39,7 @@ interface MessageStoreState { messages: Message[]; // Currently streaming entries (synced from streaming-state-machine) - streamingEntries: MessageEntry[]; + streamingEntries: InternalMessage[]; // Conversation metadata (synced from conversation-store) conversationId: string; @@ -65,7 +65,7 @@ interface MessageStoreActions { setConversation: (conversation: Conversation) => void; // Streaming entry management (used by subscriptions) - setStreamingEntries: (entries: MessageEntry[]) => void; + setStreamingEntries: (entries: InternalMessage[]) => void; // Initialize subscriptions to source stores initializeSubscriptions: () => void; @@ -112,7 +112,7 @@ const initialState: MessageStoreState = { function computeDisplayMessages( messages: Message[], - streamingEntries: MessageEntry[] + streamingEntries: InternalMessage[] ): { all: DisplayMessage[]; visible: DisplayMessage[] } { // Convert finalized messages const finalizedDisplayMessages = messages @@ -121,7 +121,7 @@ function computeDisplayMessages( // Convert streaming entries const streamingDisplayMessages = streamingEntries - .map(messageEntryToDisplayMessage) + .map(internalMessageToDisplayMessage) .filter((m): m is DisplayMessage => m !== null); // Combine: finalized first, then streaming @@ -171,7 +171,7 @@ export const useMessageStore = create()( // Streaming Entry Management (synced from streaming-state-machine) // ======================================================================== - setStreamingEntries: (entries: MessageEntry[]) => { + setStreamingEntries: (entries: InternalMessage[]) => { const { all, visible } = computeDisplayMessages(get().messages, entries); set({ streamingEntries: entries, @@ -259,7 +259,7 @@ export const useMessageStore = create()( const lastFinalized = state.messages.at(-1); // Waiting if last streaming entry is a user message - if (lastStreaming?.user !== undefined) { + if (lastStreaming?.role === "user") { return true; } @@ -276,7 +276,7 @@ export const useMessageStore = create()( hasStaleMessages: () => { return get().streamingEntries.some( - (entry) => entry.status === MessageEntryStatus.STALE + (entry) => entry.status === "stale" ); }, })) diff --git a/webapp/_webapp/src/stores/streaming/message-type-handlers.ts b/webapp/_webapp/src/stores/streaming/message-type-handlers.ts index 24a0a42f..f98dcdab 100644 --- a/webapp/_webapp/src/stores/streaming/message-type-handlers.ts +++ b/webapp/_webapp/src/stores/streaming/message-type-handlers.ts @@ -18,12 +18,16 @@ import { StreamPartEnd, } from "../../pkg/gen/apiclient/chat/v2/chat_pb"; import { - MessageEntry, - MessageEntryStatus, + InternalMessage, MessageRole, MessageTypeHandler, MessageTypeHandlerRegistry, } from "./types"; +import { + createAssistantMessage, + createToolCallMessage, + createToolCallPrepareMessage, +} from "../../types/message"; // ============================================================================ // Handler Implementations @@ -33,19 +37,26 @@ import { * Handler for assistant messages. */ class AssistantHandler implements MessageTypeHandler { - onPartBegin(partBegin: StreamPartBegin): MessageEntry | null { - return { - messageId: partBegin.messageId, - status: MessageEntryStatus.PREPARING, - assistant: partBegin.payload?.messageType.value as MessageTypeAssistant, - }; + onPartBegin(partBegin: StreamPartBegin): InternalMessage | null { + const assistant = partBegin.payload?.messageType.value as MessageTypeAssistant; + return createAssistantMessage(partBegin.messageId, assistant.content, { + reasoning: assistant.reasoning, + modelSlug: assistant.modelSlug, + status: "streaming", + }); } - onPartEnd(partEnd: StreamPartEnd, _existingEntry: MessageEntry): Partial | null { - const assistantMessage = partEnd.payload?.messageType.value as MessageTypeAssistant; + onPartEnd(partEnd: StreamPartEnd, existingMessage: InternalMessage): InternalMessage | null { + if (existingMessage.role !== "assistant") return null; + const assistant = partEnd.payload?.messageType.value as MessageTypeAssistant; return { - status: MessageEntryStatus.FINALIZED, - assistant: assistantMessage, + ...existingMessage, + status: "complete", + data: { + content: assistant.content, + reasoning: assistant.reasoning, + modelSlug: assistant.modelSlug, + }, }; } } @@ -54,21 +65,23 @@ class AssistantHandler implements MessageTypeHandler { * Handler for tool call preparation (arguments streaming). */ class ToolCallPrepareHandler implements MessageTypeHandler { - onPartBegin(partBegin: StreamPartBegin): MessageEntry | null { - return { - messageId: partBegin.messageId, - status: MessageEntryStatus.PREPARING, - toolCallPrepareArguments: partBegin.payload?.messageType - .value as MessageTypeToolCallPrepareArguments, - }; + onPartBegin(partBegin: StreamPartBegin): InternalMessage | null { + const toolCallPrepare = partBegin.payload?.messageType.value as MessageTypeToolCallPrepareArguments; + return createToolCallPrepareMessage(partBegin.messageId, toolCallPrepare.name, toolCallPrepare.args, { + status: "streaming", + }); } - onPartEnd(partEnd: StreamPartEnd, _existingEntry: MessageEntry): Partial | null { - const toolCallPrepareArguments = partEnd.payload?.messageType - .value as MessageTypeToolCallPrepareArguments; + onPartEnd(partEnd: StreamPartEnd, existingMessage: InternalMessage): InternalMessage | null { + if (existingMessage.role !== "toolCallPrepare") return null; + const toolCallPrepare = partEnd.payload?.messageType.value as MessageTypeToolCallPrepareArguments; return { - status: MessageEntryStatus.FINALIZED, - toolCallPrepareArguments, + ...existingMessage, + status: "complete", + data: { + name: toolCallPrepare.name, + args: toolCallPrepare.args, + }, }; } } @@ -77,19 +90,27 @@ class ToolCallPrepareHandler implements MessageTypeHandler { * Handler for completed tool calls. */ class ToolCallHandler implements MessageTypeHandler { - onPartBegin(partBegin: StreamPartBegin): MessageEntry | null { - return { - messageId: partBegin.messageId, - status: MessageEntryStatus.PREPARING, - toolCall: partBegin.payload?.messageType.value as MessageTypeToolCall, - }; + onPartBegin(partBegin: StreamPartBegin): InternalMessage | null { + const toolCall = partBegin.payload?.messageType.value as MessageTypeToolCall; + return createToolCallMessage(partBegin.messageId, toolCall.name, toolCall.args, { + result: toolCall.result, + error: toolCall.error, + status: "streaming", + }); } - onPartEnd(partEnd: StreamPartEnd, _existingEntry: MessageEntry): Partial | null { + onPartEnd(partEnd: StreamPartEnd, existingMessage: InternalMessage): InternalMessage | null { + if (existingMessage.role !== "toolCall") return null; const toolCall = partEnd.payload?.messageType.value as MessageTypeToolCall; return { - status: MessageEntryStatus.FINALIZED, - toolCall, + ...existingMessage, + status: "complete", + data: { + name: toolCall.name, + args: toolCall.args, + result: toolCall.result, + error: toolCall.error, + }, }; } } @@ -99,11 +120,11 @@ class ToolCallHandler implements MessageTypeHandler { * Used for system, user, and unknown message types. */ class NoOpHandler implements MessageTypeHandler { - onPartBegin(_partBegin: StreamPartBegin): MessageEntry | null { + onPartBegin(_partBegin: StreamPartBegin): InternalMessage | null { return null; } - onPartEnd(_partEnd: StreamPartEnd, _existingEntry: MessageEntry): Partial | null { + onPartEnd(_partEnd: StreamPartEnd, _existingMessage: InternalMessage): InternalMessage | null { return null; } } diff --git a/webapp/_webapp/src/stores/streaming/streaming-state-machine.ts b/webapp/_webapp/src/stores/streaming/streaming-state-machine.ts index 99f6bdd7..7b1f2108 100644 --- a/webapp/_webapp/src/stores/streaming/streaming-state-machine.ts +++ b/webapp/_webapp/src/stores/streaming/streaming-state-machine.ts @@ -16,26 +16,22 @@ import { subscribeWithSelector } from "zustand/middleware"; import { IncompleteIndicator, Message, - MessageSchema, - MessageTypeAssistant, - MessageTypeAssistantSchema, Conversation, } from "../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { fromJson } from "../../libs/protobuf-utils"; import { logError, logWarn } from "../../libs/logger"; import { errorToast } from "../../libs/toasts"; import { useConversationStore } from "../conversation/conversation-store"; import { getConversation } from "../../query/api"; import { getMessageTypeHandler, isValidMessageRole } from "./message-type-handlers"; import { - MessageEntry, - MessageEntryStatus, + InternalMessage, MessageRole, StreamEvent, StreamHandlerContext, StreamingMessage, StreamState, } from "./types"; +import { toApiMessage, createAssistantMessage } from "../../utils/message-converters"; // ============================================================================ // Store State Interface @@ -72,49 +68,6 @@ const initialState = { // Helper Functions // ============================================================================ -/** - * Convert a MessageEntry to a Message for the conversation store. - */ -function convertMessageEntryToMessage(messageEntry: MessageEntry): Message | undefined { - if (messageEntry.assistant) { - const assistantPayload: { content: string; reasoning?: string } = { - content: messageEntry.assistant.content, - }; - if (messageEntry.assistant.reasoning) { - assistantPayload.reasoning = messageEntry.assistant.reasoning; - } - return fromJson(MessageSchema, { - messageId: messageEntry.messageId, - payload: { - assistant: assistantPayload, - }, - }); - } else if (messageEntry.toolCall) { - return fromJson(MessageSchema, { - messageId: messageEntry.messageId, - payload: { - toolCall: { - name: messageEntry.toolCall.name, - args: messageEntry.toolCall.args, - result: messageEntry.toolCall.result, - error: messageEntry.toolCall.error, - }, - }, - }); - } else if (messageEntry.user) { - return fromJson(MessageSchema, { - messageId: messageEntry.messageId, - payload: { - user: { - content: messageEntry.user.content, - selectedText: messageEntry.user.selectedText ?? "", - }, - }, - }); - } - return undefined; -} - /** * Flush finalized streaming messages to the conversation store. */ @@ -124,8 +77,8 @@ function flushStreamingMessageToConversation( modelSlug?: string, ) { const flushMessages = streamingMessage.parts - .filter((part) => part.status === MessageEntryStatus.FINALIZED) - .map((part) => convertMessageEntryToMessage(part)) + .filter((part) => part.status === "complete") + .map((part) => toApiMessage(part)) .filter((part): part is Message => part !== null && part !== undefined); useConversationStore.getState().updateCurrentConversation((prev: Conversation) => ({ @@ -188,8 +141,8 @@ export const useStreamingStateMachine = create()( streamingMessage: { ...state.streamingMessage, parts: state.streamingMessage.parts.map((part) => { - if (part.status === MessageEntryStatus.PREPARING && part.user) { - return { ...part, status: MessageEntryStatus.FINALIZED }; + if (part.status === "streaming" && part.role === "user") { + return { ...part, status: "complete" as const }; } return part; }), @@ -227,18 +180,18 @@ export const useStreamingStateMachine = create()( } const handler = getMessageTypeHandler(role); - const newEntry = handler.onPartBegin(event.payload); + const newMessage = handler.onPartBegin(event.payload); - if (newEntry) { + if (newMessage) { set((state) => { - // Skip if entry with same messageId already exists - if (state.streamingMessage.parts.some((p) => p.messageId === newEntry.messageId)) { + // Skip if entry with same id already exists + if (state.streamingMessage.parts.some((p) => p.id === newMessage.id)) { return state; } return { state: "receiving", streamingMessage: { - parts: [...state.streamingMessage.parts, newEntry], + parts: [...state.streamingMessage.parts, newMessage], sequence: state.streamingMessage.sequence + 1, }, }; @@ -254,20 +207,23 @@ export const useStreamingStateMachine = create()( set((state) => { const updatedParts = state.streamingMessage.parts.map((part) => { const isTargetPart = - part.messageId === event.payload.messageId && part.assistant; + part.id === event.payload.messageId && part.role === "assistant"; if (!isTargetPart) return part; - if (part.status !== MessageEntryStatus.PREPARING) { - logError("Message chunk received for non-preparing part"); + if (part.status !== "streaming") { + logError("Message chunk received for non-streaming part"); } - const updatedAssistant: MessageTypeAssistant = { - ...part.assistant!, - content: part.assistant!.content + event.payload.delta, - }; + if (part.role !== "assistant") return part; - return { ...part, assistant: updatedAssistant }; + return { + ...part, + data: { + ...part.data, + content: part.data.content + event.payload.delta, + }, + }; }); return { @@ -287,21 +243,24 @@ export const useStreamingStateMachine = create()( set((state) => { const updatedParts = state.streamingMessage.parts.map((part) => { const isTargetPart = - part.messageId === event.payload.messageId && part.assistant; + part.id === event.payload.messageId && part.role === "assistant"; if (!isTargetPart) return part; - if (part.status !== MessageEntryStatus.PREPARING) { - logError("Reasoning chunk received for non-preparing part"); + if (part.status !== "streaming") { + logError("Reasoning chunk received for non-streaming part"); } - const currentReasoning = part.assistant!.reasoning ?? ""; - const updatedAssistant: MessageTypeAssistant = { - ...part.assistant!, - reasoning: currentReasoning + event.payload.delta, - }; + if (part.role !== "assistant") return part; - return { ...part, assistant: updatedAssistant }; + const currentReasoning = part.data.reasoning ?? ""; + return { + ...part, + data: { + ...part.data, + reasoning: currentReasoning + event.payload.delta, + }, + }; }); return { @@ -329,14 +288,14 @@ export const useStreamingStateMachine = create()( set((state) => { const updatedParts = state.streamingMessage.parts.map((part) => { - if (part.messageId !== event.payload.messageId) { + if (part.id !== event.payload.messageId) { return part; } - const updates = handler.onPartEnd(event.payload, part); - if (!updates) return part; + const updatedMessage = handler.onPartEnd(event.payload, part); + if (!updatedMessage) return part; - return { ...part, ...updates }; + return updatedMessage; }); return { @@ -394,13 +353,11 @@ export const useStreamingStateMachine = create()( } // Add error message to streaming parts - const errorEntry: MessageEntry = { - messageId: "error-" + Date.now(), - status: MessageEntryStatus.STALE, - assistant: fromJson(MessageTypeAssistantSchema, { - content: errorMessage, - }), - }; + const errorEntry: InternalMessage = createAssistantMessage( + "error-" + Date.now(), + errorMessage, + { status: "stale" } + ); set((state) => ({ state: "error", @@ -418,16 +375,13 @@ export const useStreamingStateMachine = create()( // CONNECTION_ERROR - Network/connection error // ======================================================================== case "CONNECTION_ERROR": { - // Mark all preparing messages as stale + // Mark all streaming messages as stale set((state) => ({ state: "error", streamingMessage: { parts: state.streamingMessage.parts.map((part) => ({ ...part, - status: - part.status === MessageEntryStatus.PREPARING - ? MessageEntryStatus.STALE - : part.status, + status: part.status === "streaming" ? "stale" as const : part.status, })), sequence: state.streamingMessage.sequence + 1, }, diff --git a/webapp/_webapp/src/stores/streaming/types.ts b/webapp/_webapp/src/stores/streaming/types.ts index f154a0e9..b7b0d501 100644 --- a/webapp/_webapp/src/stores/streaming/types.ts +++ b/webapp/_webapp/src/stores/streaming/types.ts @@ -14,12 +14,11 @@ import { StreamInitialization, StreamPartBegin, StreamPartEnd, - MessageTypeAssistant, - MessageTypeToolCall, - MessageTypeToolCallPrepareArguments, - MessageTypeUnknown, - MessageTypeUser, } from "../../pkg/gen/apiclient/chat/v2/chat_pb"; +import { InternalMessage, MessageStatus } from "../../types/message"; + +// Re-export InternalMessage for convenience +export type { InternalMessage, MessageStatus }; // ============================================================================ // Stream State @@ -31,39 +30,15 @@ import { export type StreamState = "idle" | "receiving" | "finalizing" | "error"; // ============================================================================ -// Message Entry Types +// Streaming Message State // ============================================================================ -/** - * Status of a message entry during the streaming lifecycle. - */ -export enum MessageEntryStatus { - PREPARING = "PREPARING", - FINALIZED = "FINALIZED", - INCOMPLETE = "INCOMPLETE", - STALE = "STALE", -} - -/** - * Represents a message entry in the streaming state. - * Uses a discriminated union pattern for type safety. - */ -export type MessageEntry = { - messageId: string; - status: MessageEntryStatus; - // Role-specific content (only one will be present) - user?: MessageTypeUser; - assistant?: MessageTypeAssistant; - toolCallPrepareArguments?: MessageTypeToolCallPrepareArguments; - toolCall?: MessageTypeToolCall; - unknown?: MessageTypeUnknown; -}; - /** * The current streaming message state. + * Now uses InternalMessage instead of the legacy MessageEntry. */ export type StreamingMessage = { - parts: MessageEntry[]; + parts: InternalMessage[]; sequence: number; }; @@ -138,15 +113,15 @@ export interface StreamHandlerContext { export interface MessageTypeHandler { /** * Called when a stream part begins for this message type. - * @returns A new MessageEntry or null if this type should be ignored. + * @returns A new InternalMessage or null if this type should be ignored. */ - onPartBegin(partBegin: StreamPartBegin): MessageEntry | null; + onPartBegin(partBegin: StreamPartBegin): InternalMessage | null; /** * Called when a stream part ends for this message type. - * @returns Updated fields to merge into the existing entry, or null to skip. + * @returns Updated InternalMessage or null to skip. */ - onPartEnd(partEnd: StreamPartEnd, existingEntry: MessageEntry): Partial | null; + onPartEnd(partEnd: StreamPartEnd, existingMessage: InternalMessage): InternalMessage | null; } /** diff --git a/webapp/_webapp/src/types/index.ts b/webapp/_webapp/src/types/index.ts new file mode 100644 index 00000000..fa43b99d --- /dev/null +++ b/webapp/_webapp/src/types/index.ts @@ -0,0 +1,42 @@ +/** + * Types Module + * + * Re-exports all public types used throughout the application. + */ + +// Internal Message Types (canonical format) +export type { + InternalMessage, + UserMessage, + AssistantMessage, + ToolCallMessage, + ToolCallPrepareMessage, + SystemMessage, + UnknownMessage, + MessageStatus, + MessageRole, + UserMessageData, + AssistantMessageData, + ToolCallData, + ToolCallPrepareData, + SystemMessageData, + UnknownMessageData, +} from "./message"; + +// Type Guards +export { + isUserMessage, + isAssistantMessage, + isToolCallMessage, + isToolCallPrepareMessage, + isSystemMessage, + isUnknownMessage, +} from "./message"; + +// Factory Functions +export { + createUserMessage, + createAssistantMessage, + createToolCallMessage, + createToolCallPrepareMessage, +} from "./message"; diff --git a/webapp/_webapp/src/types/message.ts b/webapp/_webapp/src/types/message.ts new file mode 100644 index 00000000..6d24a95a --- /dev/null +++ b/webapp/_webapp/src/types/message.ts @@ -0,0 +1,309 @@ +/** + * Canonical Internal Message Types + * + * This file defines the single internal message format used throughout the app. + * All components should use these types instead of working with protobuf types directly. + * + * Data Flow: + * 1. API Response → InternalMessage (via fromApiMessage) + * 2. InternalMessage → API Request (via toApiMessage) + * 3. InternalMessage → DisplayMessage (trivial 1:1 mapping in most cases) + * + * Benefits: + * - Single format reduces confusion + * - Clear boundary between API types and internal types + * - Reduces the number of data transformations + */ + +// ============================================================================ +// Message Status +// ============================================================================ + +/** + * Status of a message during its lifecycle. + * Replaces MessageEntryStatus with clearer semantics. + */ +export type MessageStatus = "streaming" | "complete" | "error" | "stale"; + +// ============================================================================ +// Message Roles +// ============================================================================ + +/** + * All possible message roles. + */ +export type MessageRole = + | "user" + | "assistant" + | "toolCall" + | "toolCallPrepare" + | "system" + | "unknown"; + +// ============================================================================ +// Role-Specific Data +// ============================================================================ + +/** + * Data specific to user messages. + */ +export interface UserMessageData { + content: string; + selectedText?: string; + surrounding?: string; +} + +/** + * Data specific to assistant messages. + */ +export interface AssistantMessageData { + content: string; + reasoning?: string; + modelSlug?: string; +} + +/** + * Data specific to tool call messages. + */ +export interface ToolCallData { + name: string; + args: string; + result?: string; + error?: string; +} + +/** + * Data specific to tool call preparation messages. + */ +export interface ToolCallPrepareData { + name: string; + args: string; +} + +/** + * Data specific to system messages. + */ +export interface SystemMessageData { + content: string; +} + +/** + * Data specific to unknown messages. + */ +export interface UnknownMessageData { + description: string; +} + +// ============================================================================ +// Internal Message Type +// ============================================================================ + +/** + * Base properties shared by all message types. + */ +interface InternalMessageBase { + /** Unique message identifier */ + id: string; + /** Current status of the message */ + status: MessageStatus; + /** Optional timestamp (milliseconds since epoch) */ + timestamp?: number; + /** ID of the previous message (for branching) */ + previousMessageId?: string; +} + +/** + * User message. + */ +export interface UserMessage extends InternalMessageBase { + role: "user"; + data: UserMessageData; +} + +/** + * Assistant message. + */ +export interface AssistantMessage extends InternalMessageBase { + role: "assistant"; + data: AssistantMessageData; +} + +/** + * Tool call message. + */ +export interface ToolCallMessage extends InternalMessageBase { + role: "toolCall"; + data: ToolCallData; +} + +/** + * Tool call preparation message. + */ +export interface ToolCallPrepareMessage extends InternalMessageBase { + role: "toolCallPrepare"; + data: ToolCallPrepareData; +} + +/** + * System message. + */ +export interface SystemMessage extends InternalMessageBase { + role: "system"; + data: SystemMessageData; +} + +/** + * Unknown message type. + */ +export interface UnknownMessage extends InternalMessageBase { + role: "unknown"; + data: UnknownMessageData; +} + +/** + * Union type representing all internal message types. + * This is the canonical format used throughout the application. + */ +export type InternalMessage = + | UserMessage + | AssistantMessage + | ToolCallMessage + | ToolCallPrepareMessage + | SystemMessage + | UnknownMessage; + +// ============================================================================ +// Type Guards +// ============================================================================ + +export function isUserMessage(msg: InternalMessage): msg is UserMessage { + return msg.role === "user"; +} + +export function isAssistantMessage(msg: InternalMessage): msg is AssistantMessage { + return msg.role === "assistant"; +} + +export function isToolCallMessage(msg: InternalMessage): msg is ToolCallMessage { + return msg.role === "toolCall"; +} + +export function isToolCallPrepareMessage(msg: InternalMessage): msg is ToolCallPrepareMessage { + return msg.role === "toolCallPrepare"; +} + +export function isSystemMessage(msg: InternalMessage): msg is SystemMessage { + return msg.role === "system"; +} + +export function isUnknownMessage(msg: InternalMessage): msg is UnknownMessage { + return msg.role === "unknown"; +} + +// ============================================================================ +// Factory Functions +// ============================================================================ + +/** + * Create a new user message. + */ +export function createUserMessage( + id: string, + content: string, + options?: { + selectedText?: string; + surrounding?: string; + status?: MessageStatus; + previousMessageId?: string; + } +): UserMessage { + return { + id, + role: "user", + status: options?.status ?? "streaming", + previousMessageId: options?.previousMessageId, + data: { + content, + selectedText: options?.selectedText, + surrounding: options?.surrounding, + }, + }; +} + +/** + * Create a new assistant message. + */ +export function createAssistantMessage( + id: string, + content: string, + options?: { + reasoning?: string; + modelSlug?: string; + status?: MessageStatus; + previousMessageId?: string; + } +): AssistantMessage { + return { + id, + role: "assistant", + status: options?.status ?? "streaming", + previousMessageId: options?.previousMessageId, + data: { + content, + reasoning: options?.reasoning, + modelSlug: options?.modelSlug, + }, + }; +} + +/** + * Create a new tool call message. + */ +export function createToolCallMessage( + id: string, + name: string, + args: string, + options?: { + result?: string; + error?: string; + status?: MessageStatus; + previousMessageId?: string; + } +): ToolCallMessage { + return { + id, + role: "toolCall", + status: options?.status ?? "streaming", + previousMessageId: options?.previousMessageId, + data: { + name, + args, + result: options?.result, + error: options?.error, + }, + }; +} + +/** + * Create a new tool call prepare message. + */ +export function createToolCallPrepareMessage( + id: string, + name: string, + args: string, + options?: { + status?: MessageStatus; + previousMessageId?: string; + } +): ToolCallPrepareMessage { + return { + id, + role: "toolCallPrepare", + status: options?.status ?? "streaming", + previousMessageId: options?.previousMessageId, + data: { + name, + args, + }, + }; +} diff --git a/webapp/_webapp/src/utils/index.ts b/webapp/_webapp/src/utils/index.ts new file mode 100644 index 00000000..fb8ac50b --- /dev/null +++ b/webapp/_webapp/src/utils/index.ts @@ -0,0 +1,18 @@ +/** + * Utils Module + * + * Re-exports all utility functions. + */ + +// Message Converters +export { + // API ↔ InternalMessage + fromApiMessage, + toApiMessage, + // Stream Events → InternalMessage + fromStreamPartBegin, + applyStreamPartEnd, + // InternalMessage ↔ DisplayMessage + toDisplayMessage, + fromDisplayMessage, +} from "./message-converters"; diff --git a/webapp/_webapp/src/utils/message-converters.ts b/webapp/_webapp/src/utils/message-converters.ts new file mode 100644 index 00000000..6752a471 --- /dev/null +++ b/webapp/_webapp/src/utils/message-converters.ts @@ -0,0 +1,490 @@ +/** + * Message Converters + * + * Bidirectional converters between API types (protobuf Message) and internal types (InternalMessage). + * These provide the only two transformations needed: + * + * 1. API Response → InternalMessage (fromApiMessage) + * 2. InternalMessage → API Request (toApiMessage) + * + * Benefits: + * - Clear boundary between API types and internal types + * - Reduces the number of data transformations from 5+ to 2 + * - All conversion logic in one place + */ + +import { + Message, + MessageSchema, + StreamPartBegin, + StreamPartEnd, +} from "../pkg/gen/apiclient/chat/v2/chat_pb"; +import { fromJson } from "../libs/protobuf-utils"; +import { + InternalMessage, + MessageStatus, + createAssistantMessage, + createToolCallMessage, + createToolCallPrepareMessage, +} from "../types/message"; + +// ============================================================================ +// API Response → InternalMessage +// ============================================================================ + +/** + * Convert a protobuf Message to InternalMessage. + * This is used when receiving finalized messages from the server. + * + * @param msg The protobuf Message from the API + * @param status Optional status override (default: "complete") + * @returns InternalMessage or null if message type is not recognized + */ +export function fromApiMessage(msg: Message, status: MessageStatus = "complete"): InternalMessage | null { + const messageType = msg.payload?.messageType; + + if (!messageType) return null; + + switch (messageType.case) { + case "user": + return { + id: msg.messageId, + role: "user", + status, + timestamp: Number(msg.timestamp) || undefined, + data: { + content: messageType.value.content, + selectedText: messageType.value.selectedText, + surrounding: messageType.value.surrounding, + }, + }; + + case "assistant": + return { + id: msg.messageId, + role: "assistant", + status, + timestamp: Number(msg.timestamp) || undefined, + data: { + content: messageType.value.content, + reasoning: messageType.value.reasoning, + modelSlug: messageType.value.modelSlug, + }, + }; + + case "toolCall": + return { + id: msg.messageId, + role: "toolCall", + status, + timestamp: Number(msg.timestamp) || undefined, + data: { + name: messageType.value.name, + args: messageType.value.args, + result: messageType.value.result, + error: messageType.value.error, + }, + }; + + case "toolCallPrepareArguments": + return { + id: msg.messageId, + role: "toolCallPrepare", + status, + timestamp: Number(msg.timestamp) || undefined, + data: { + name: messageType.value.name, + args: messageType.value.args, + }, + }; + + case "system": + return { + id: msg.messageId, + role: "system", + status, + timestamp: Number(msg.timestamp) || undefined, + data: { + content: messageType.value.content, + }, + }; + + case "unknown": + return { + id: msg.messageId, + role: "unknown", + status, + timestamp: Number(msg.timestamp) || undefined, + data: { + description: messageType.value.description, + }, + }; + + default: + return null; + } +} + +// ============================================================================ +// InternalMessage → API Request +// ============================================================================ + +import { JsonValue } from "@bufbuild/protobuf"; + +/** + * Convert an InternalMessage to a protobuf Message. + * This is used when sending messages to the server or storing in conversation state. + * + * @param msg The internal message + * @returns Protobuf Message or undefined if conversion fails + */ +export function toApiMessage(msg: InternalMessage): Message | undefined { + switch (msg.role) { + case "user": + return fromJson(MessageSchema, { + messageId: msg.id, + payload: { + user: { + content: msg.data.content, + selectedText: msg.data.selectedText ?? "", + }, + }, + } as unknown as JsonValue); + + case "assistant": { + const assistantPayload: { content: string; reasoning?: string; modelSlug?: string } = { + content: msg.data.content, + }; + if (msg.data.reasoning) { + assistantPayload.reasoning = msg.data.reasoning; + } + if (msg.data.modelSlug) { + assistantPayload.modelSlug = msg.data.modelSlug; + } + return fromJson(MessageSchema, { + messageId: msg.id, + payload: { + assistant: assistantPayload, + }, + } as unknown as JsonValue); + } + + case "toolCall": + return fromJson(MessageSchema, { + messageId: msg.id, + payload: { + toolCall: { + name: msg.data.name, + args: msg.data.args, + result: msg.data.result ?? "", + error: msg.data.error ?? "", + }, + }, + } as unknown as JsonValue); + + case "toolCallPrepare": + return fromJson(MessageSchema, { + messageId: msg.id, + payload: { + toolCallPrepareArguments: { + name: msg.data.name, + args: msg.data.args, + }, + }, + } as unknown as JsonValue); + + case "system": + return fromJson(MessageSchema, { + messageId: msg.id, + payload: { + system: { + content: msg.data.content, + }, + }, + } as unknown as JsonValue); + + case "unknown": + return fromJson(MessageSchema, { + messageId: msg.id, + payload: { + unknown: { + description: msg.data.description, + }, + }, + } as unknown as JsonValue); + + default: + return undefined; + } +} + +// ============================================================================ +// Stream Events → InternalMessage +// ============================================================================ + +/** + * Create an InternalMessage from a StreamPartBegin event. + * Used during streaming to initialize a new message entry. + * + * @param partBegin The stream part begin event + * @returns InternalMessage or null if message type should be ignored + */ +export function fromStreamPartBegin(partBegin: StreamPartBegin): InternalMessage | null { + const messageType = partBegin.payload?.messageType; + + if (!messageType) return null; + + switch (messageType.case) { + case "assistant": + return createAssistantMessage( + partBegin.messageId, + messageType.value.content, + { + reasoning: messageType.value.reasoning, + modelSlug: messageType.value.modelSlug, + status: "streaming", + } + ); + + case "toolCall": + return createToolCallMessage( + partBegin.messageId, + messageType.value.name, + messageType.value.args, + { + result: messageType.value.result, + error: messageType.value.error, + status: "streaming", + } + ); + + case "toolCallPrepareArguments": + return createToolCallPrepareMessage( + partBegin.messageId, + messageType.value.name, + messageType.value.args, + { status: "streaming" } + ); + + // User, system, and unknown messages are not handled via streaming + case "user": + case "system": + case "unknown": + return null; + + default: + return null; + } +} + +/** + * Get the updated data from a StreamPartEnd event. + * Used to finalize a streaming message with complete data. + * + * @param partEnd The stream part end event + * @param existingMessage The existing message to update + * @returns Updated InternalMessage or null if update should be skipped + */ +export function applyStreamPartEnd( + partEnd: StreamPartEnd, + existingMessage: InternalMessage +): InternalMessage | null { + const messageType = partEnd.payload?.messageType; + + if (!messageType) return null; + + switch (messageType.case) { + case "assistant": + if (existingMessage.role !== "assistant") return null; + return { + ...existingMessage, + status: "complete", + data: { + ...existingMessage.data, + content: messageType.value.content, + reasoning: messageType.value.reasoning, + modelSlug: messageType.value.modelSlug, + }, + }; + + case "toolCall": + if (existingMessage.role !== "toolCall") return null; + return { + ...existingMessage, + status: "complete", + data: { + name: messageType.value.name, + args: messageType.value.args, + result: messageType.value.result, + error: messageType.value.error, + }, + }; + + case "toolCallPrepareArguments": + if (existingMessage.role !== "toolCallPrepare") return null; + return { + ...existingMessage, + status: "complete", + data: { + name: messageType.value.name, + args: messageType.value.args, + }, + }; + + // User, system, and unknown messages are not handled via streaming + case "user": + case "system": + case "unknown": + return null; + + default: + return null; + } +} + +// ============================================================================ +// InternalMessage ↔ DisplayMessage +// ============================================================================ + +import { DisplayMessage, DisplayMessageStatus } from "../stores/types"; + +/** + * Convert InternalMessage to DisplayMessage. + * This is a simple 1:1 mapping in most cases. + * + * @param msg The internal message + * @returns DisplayMessage or null if message should not be displayed + */ +export function toDisplayMessage(msg: InternalMessage): DisplayMessage | null { + switch (msg.role) { + case "user": + return { + id: msg.id, + type: "user", + status: msg.status as DisplayMessageStatus, + content: msg.data.content, + selectedText: msg.data.selectedText, + previousMessageId: msg.previousMessageId, + }; + + case "assistant": + return { + id: msg.id, + type: "assistant", + status: msg.status as DisplayMessageStatus, + content: msg.data.content, + reasoning: msg.data.reasoning, + previousMessageId: msg.previousMessageId, + }; + + case "toolCall": + return { + id: msg.id, + type: "toolCall", + status: msg.status as DisplayMessageStatus, + content: "", + toolName: msg.data.name, + toolArgs: msg.data.args, + toolResult: msg.data.result, + toolError: msg.data.error, + previousMessageId: msg.previousMessageId, + }; + + case "toolCallPrepare": + return { + id: msg.id, + type: "toolCallPrepare", + status: msg.status as DisplayMessageStatus, + content: "", + toolName: msg.data.name, + toolArgs: msg.data.args, + previousMessageId: msg.previousMessageId, + }; + + // System and unknown messages are typically not displayed + case "system": + case "unknown": + return null; + + default: + return null; + } +} + +/** + * Convert DisplayMessage back to InternalMessage. + * Used for backward compatibility with components that need InternalMessage. + * + * @param msg The display message + * @returns InternalMessage + */ +export function fromDisplayMessage(msg: DisplayMessage): InternalMessage { + const status = msg.status as MessageStatus; + + switch (msg.type) { + case "user": + return { + id: msg.id, + role: "user", + status, + previousMessageId: msg.previousMessageId, + data: { + content: msg.content, + selectedText: msg.selectedText, + }, + }; + + case "assistant": + return { + id: msg.id, + role: "assistant", + status, + previousMessageId: msg.previousMessageId, + data: { + content: msg.content, + reasoning: msg.reasoning, + }, + }; + + case "toolCall": + return { + id: msg.id, + role: "toolCall", + status, + previousMessageId: msg.previousMessageId, + data: { + name: msg.toolName ?? "", + args: msg.toolArgs ?? "", + result: msg.toolResult, + error: msg.toolError, + }, + }; + + case "toolCallPrepare": + return { + id: msg.id, + role: "toolCallPrepare", + status, + previousMessageId: msg.previousMessageId, + data: { + name: msg.toolName ?? "", + args: msg.toolArgs ?? "", + }, + }; + + case "error": + default: + return { + id: msg.id, + role: "unknown", + status, + previousMessageId: msg.previousMessageId, + data: { + description: msg.content, + }, + }; + } +} + +// Re-export factory functions for convenience +export { createAssistantMessage, createToolCallMessage, createToolCallPrepareMessage }; diff --git a/webapp/_webapp/src/views/chat/body/index.tsx b/webapp/_webapp/src/views/chat/body/index.tsx index 07de503b..986770c8 100644 --- a/webapp/_webapp/src/views/chat/body/index.tsx +++ b/webapp/_webapp/src/views/chat/body/index.tsx @@ -5,7 +5,6 @@ import { isEmptyConversation, getPrevUserSelectedText, findLastUserMessageIndex, - displayMessageToMessageEntry, } from "../helper"; import { StatusIndicator } from "./status-indicator"; import { EmptyView } from "./empty-view"; @@ -106,7 +105,7 @@ export const ChatBody = ({ conversation }: ChatBodyProps) => { > 0 ? visibleMessages[index - 1].id : undefined} /> diff --git a/webapp/_webapp/src/views/chat/body/status-indicator.tsx b/webapp/_webapp/src/views/chat/body/status-indicator.tsx index 1c3127b2..5bae78ce 100644 --- a/webapp/_webapp/src/views/chat/body/status-indicator.tsx +++ b/webapp/_webapp/src/views/chat/body/status-indicator.tsx @@ -2,7 +2,7 @@ import { LoadingIndicator } from "../../../components/loading-indicator"; import { UnknownEntryMessageContainer } from "../../../components/message-entry-container/unknown-entry"; import { Conversation } from "../../../pkg/gen/apiclient/chat/v2/chat_pb"; import { useSocketStore } from "../../../stores/socket-store"; -import { useStreamingStateMachine, MessageEntryStatus } from "../../../stores/streaming"; +import { useStreamingStateMachine } from "../../../stores/streaming"; export const StatusIndicator = ({ conversation }: { conversation?: Conversation }) => { const { syncing, syncingProgress } = useSocketStore(); @@ -10,9 +10,9 @@ export const StatusIndicator = ({ conversation }: { conversation?: Conversation const incompleteIndicator = useStreamingStateMachine((s) => s.incompleteIndicator); const isWaitingForResponse = - streamingMessage.parts.at(-1)?.user !== undefined || + streamingMessage.parts.at(-1)?.role === "user" || (conversation?.messages.at(-1)?.payload?.messageType.case === "user" && streamingMessage.parts.length === 0); - const hasStaleMessage = streamingMessage.parts.some((part) => part.status === MessageEntryStatus.STALE); + const hasStaleMessage = streamingMessage.parts.some((part) => part.status === "stale"); const incompleteReason = incompleteIndicator?.reason; if (isWaitingForResponse) { diff --git a/webapp/_webapp/src/views/chat/helper.ts b/webapp/_webapp/src/views/chat/helper.ts index 3925fc37..0750a730 100644 --- a/webapp/_webapp/src/views/chat/helper.ts +++ b/webapp/_webapp/src/views/chat/helper.ts @@ -1,15 +1,12 @@ import { Conversation, Message, - MessageTypeAssistant, - MessageTypeToolCall, - MessageTypeToolCallPrepareArguments, - MessageTypeUnknown, MessageTypeUser, } from "../../pkg/gen/apiclient/chat/v2/chat_pb"; -import { MessageEntry, MessageEntryStatus } from "../../stores/streaming"; import { useMessageStore } from "../../stores/message-store"; import { DisplayMessage } from "../../stores/types"; +import { fromApiMessage } from "../../utils/message-converters"; +import { InternalMessage } from "../../stores/streaming"; // ============================================================================ // Message-based helpers (existing, for backward compatibility) @@ -57,31 +54,12 @@ export function filterVisibleMessages(conversation?: Conversation): Message[] { ); } -export function messageToMessageEntry(message: Message): MessageEntry { - return { - messageId: message.messageId, - status: MessageEntryStatus.FINALIZED, - assistant: - message.payload?.messageType.case === "assistant" - ? (message.payload?.messageType.value as MessageTypeAssistant) - : undefined, - user: - message.payload?.messageType.case === "user" - ? (message.payload?.messageType.value as MessageTypeUser) - : undefined, - toolCall: - message.payload?.messageType.case === "toolCall" - ? (message.payload?.messageType.value as MessageTypeToolCall) - : undefined, - toolCallPrepareArguments: - message.payload?.messageType.case === "toolCallPrepareArguments" - ? (message.payload?.messageType.value as MessageTypeToolCallPrepareArguments) - : undefined, - unknown: - message.payload?.messageType.case === "unknown" - ? (message.payload?.messageType.value as MessageTypeUnknown) - : undefined, - } as MessageEntry; +/** + * Convert a protobuf Message to InternalMessage. + * This is a convenience function that uses the unified converter. + */ +export function messageToInternalMessage(message: Message): InternalMessage | null { + return fromApiMessage(message); } // ============================================================================ @@ -134,62 +112,3 @@ export function findLastUserMessageIndex(messages: DisplayMessage[]): number { return -1; } -/** - * Convert DisplayMessage back to MessageEntry for backward compatibility with MessageCard. - * This is a temporary bridge until MessageCard is updated to use DisplayMessage directly. - */ -export function displayMessageToMessageEntry(msg: DisplayMessage): MessageEntry { - const status = displayStatusToEntryStatus(msg.status); - - const entry: MessageEntry = { - messageId: msg.id, - status, - }; - - if (msg.type === "user") { - entry.user = { - content: msg.content, - selectedText: msg.selectedText ?? "", - surrounding: undefined, - $typeName: "chat.v2.MessageTypeUser", - } as MessageTypeUser; - } else if (msg.type === "assistant") { - entry.assistant = { - content: msg.content, - reasoning: msg.reasoning, - $typeName: "chat.v2.MessageTypeAssistant", - } as MessageTypeAssistant; - } else if (msg.type === "toolCall") { - entry.toolCall = { - name: msg.toolName ?? "", - args: msg.toolArgs ?? "", - result: msg.toolResult ?? "", - error: msg.toolError ?? "", - $typeName: "chat.v2.MessageTypeToolCall", - } as MessageTypeToolCall; - } else if (msg.type === "toolCallPrepare") { - entry.toolCallPrepareArguments = { - name: msg.toolName ?? "", - args: msg.toolArgs ?? "", - $typeName: "chat.v2.MessageTypeToolCallPrepareArguments", - } as MessageTypeToolCallPrepareArguments; - } - - return entry; -} - -function displayStatusToEntryStatus(status: DisplayMessage["status"]): MessageEntryStatus { - switch (status) { - case "streaming": - return MessageEntryStatus.PREPARING; - case "complete": - return MessageEntryStatus.FINALIZED; - case "stale": - return MessageEntryStatus.STALE; - case "error": - return MessageEntryStatus.INCOMPLETE; - default: - return MessageEntryStatus.FINALIZED; - } -} - diff --git a/webapp/_webapp/src/views/devtools/index.tsx b/webapp/_webapp/src/views/devtools/index.tsx index c3b749b2..c068b126 100644 --- a/webapp/_webapp/src/views/devtools/index.tsx +++ b/webapp/_webapp/src/views/devtools/index.tsx @@ -1,12 +1,18 @@ import { Rnd } from "react-rnd"; import { useSelectionStore } from "../../stores/selection-store"; import { Button, Input } from "@heroui/react"; -import { useStreamingStateMachine, MessageEntry, MessageEntryStatus } from "../../stores/streaming"; +import { useStreamingStateMachine, InternalMessage } from "../../stores/streaming"; import { useConversationStore } from "../../stores/conversation/conversation-store"; import { fromJson } from "../../libs/protobuf-utils"; import { MessageSchema } from "../../pkg/gen/apiclient/chat/v2/chat_pb"; import { isEmptyConversation } from "../chat/helper"; import { useState } from "react"; +import { + createUserMessage, + createAssistantMessage, + createToolCallMessage, + createToolCallPrepareMessage, +} from "../../types/message"; // --- Utility functions --- const loremIpsum = @@ -88,7 +94,7 @@ export const DevTools = () => { }); const handleStaleLastConversationMessage = () => { const newMessages = currentConversation.messages.map((msg, _, arr) => - msg.messageId === arr[arr.length - 1]?.messageId ? { ...msg, status: MessageEntryStatus.STALE } : msg, + msg.messageId === arr[arr.length - 1]?.messageId ? { ...msg, status: "stale" } : msg, ); setCurrentConversation({ ...currentConversation, messages: newMessages }); }; @@ -109,7 +115,7 @@ export const DevTools = () => { const newParts = useStreamingStateMachine .getState() .streamingMessage.parts.map((part, _, arr) => - part.messageId === arr[arr.length - 1]?.messageId ? { ...part, status: MessageEntryStatus.STALE } : part, + part.id === arr[arr.length - 1]?.id ? { ...part, status: "stale" as const } : part, ); setStreamingMessage({ ...streamingMessage, parts: [...newParts] }); }; @@ -120,120 +126,91 @@ export const DevTools = () => { }; // StreamingMessage add various message types const handleAddStreamingUserMessage = () => { - const messageEntry: MessageEntry = { - messageId: randomUUID(), - status: MessageEntryStatus.PREPARING, - user: { - content: "User Message Preparing", - selectedText: selectedText ?? "", - $typeName: "chat.v2.MessageTypeUser", - }, - }; - setStreamingMessage({ ...streamingMessage, parts: [...streamingMessage.parts, messageEntry] }); + const newMessage = createUserMessage(randomUUID(), "User Message Preparing", { + selectedText: selectedText ?? "", + status: "streaming", + }); + setStreamingMessage({ ...streamingMessage, parts: [...streamingMessage.parts, newMessage] }); withDelay(() => { const newParts = useStreamingStateMachine.getState().streamingMessage.parts.map((part) => - part.messageId === messageEntry.messageId + part.id === newMessage.id ? { ...part, - user: { ...part.user, content: "User Message Prepared", $typeName: "chat.v2.MessageTypeUser" }, - status: part.status === MessageEntryStatus.PREPARING ? MessageEntryStatus.FINALIZED : part.status, + data: { ...part.data, content: "User Message Prepared" }, + status: part.status === "streaming" ? "complete" as const : part.status, } : part, - ) as MessageEntry[]; + ) as InternalMessage[]; setStreamingMessage({ ...streamingMessage, parts: [...newParts] }); }); }; const handleAddStreamingToolPrepare = () => { - const messageEntry: MessageEntry = { - messageId: randomUUID(), - status: MessageEntryStatus.PREPARING, - toolCallPrepareArguments: { - name: "paper_score", - args: JSON.stringify({ paper_id: "123" }), - $typeName: "chat.v2.MessageTypeToolCallPrepareArguments", - }, - }; - updateStreamingMessage((prev) => ({ ...prev, parts: [...prev.parts, messageEntry] })); + const newMessage = createToolCallPrepareMessage( + randomUUID(), + "paper_score", + JSON.stringify({ paper_id: "123" }), + { status: "streaming" } + ); + updateStreamingMessage((prev) => ({ ...prev, parts: [...prev.parts, newMessage] })); withDelay(() => { const newParts = useStreamingStateMachine.getState().streamingMessage.parts.map((part) => - part.messageId === messageEntry.messageId + part.id === newMessage.id ? { ...part, - status: part.status === MessageEntryStatus.PREPARING ? MessageEntryStatus.FINALIZED : part.status, - toolCallPrepareArguments: { - name: "paper_score", - args: JSON.stringify({ paper_id: "123" }), - $typeName: "chat.v2.MessageTypeToolCallPrepareArguments", - }, + status: part.status === "streaming" ? "complete" as const : part.status, } : part, - ) as MessageEntry[]; + ) as InternalMessage[]; updateStreamingMessage((prev) => ({ ...prev, parts: [...newParts] })); }); }; const handleAddStreamingToolCall = (type: "greeting" | "paper_score") => { const isGreeting = type === "greeting"; - const messageEntry: MessageEntry = { - messageId: randomUUID(), - status: MessageEntryStatus.PREPARING, - toolCall: isGreeting - ? { - name: "greeting", - args: JSON.stringify({ name: "Junyi" }), - result: "preparing", - error: "", - $typeName: "chat.v2.MessageTypeToolCall", - } - : { - name: "paper_score", - args: JSON.stringify({ paper_id: "123" }), - result: '{ "percentile": 0.74829 }123', - error: "", - $typeName: "chat.v2.MessageTypeToolCall", - }, - }; - updateStreamingMessage((prev) => ({ ...prev, parts: [...prev.parts, messageEntry] })); + const newMessage = isGreeting + ? createToolCallMessage(randomUUID(), "greeting", JSON.stringify({ name: "Junyi" }), { + result: "preparing", + status: "streaming", + }) + : createToolCallMessage(randomUUID(), "paper_score", JSON.stringify({ paper_id: "123" }), { + result: '{ "percentile": 0.74829 }123', + status: "streaming", + }); + updateStreamingMessage((prev) => ({ ...prev, parts: [...prev.parts, newMessage] })); withDelay(() => { - const newParts = useStreamingStateMachine.getState().streamingMessage.parts.map((part) => - part.messageId === messageEntry.messageId - ? { - ...part, - status: part.status === MessageEntryStatus.PREPARING ? MessageEntryStatus.FINALIZED : part.status, - toolCall: isGreeting - ? { ...part.toolCall, result: "Hello, Junyi!", $typeName: "chat.v2.MessageTypeToolCall" } - : { ...part.toolCall, $typeName: "chat.v2.MessageTypeToolCall" }, - } - : part, - ) as MessageEntry[]; + const newParts = useStreamingStateMachine.getState().streamingMessage.parts.map((part) => { + if (part.id !== newMessage.id) return part; + if (part.role !== "toolCall") return part; + return { + ...part, + status: "complete" as const, + data: isGreeting + ? { ...part.data, result: "Hello, Junyi!" } + : part.data, + }; + }) as InternalMessage[]; updateStreamingMessage((prev) => ({ ...prev, parts: [...newParts] })); }); }; const handleAddStreamingAssistant = () => { - const messageEntry: MessageEntry = { - messageId: randomUUID(), - status: MessageEntryStatus.PREPARING, - assistant: { - content: "Assistant Response Preparing " + randomText(), - modelSlug: "gpt-4.1", - $typeName: "chat.v2.MessageTypeAssistant", - }, - }; - updateStreamingMessage((prev) => ({ ...prev, parts: [...prev.parts, messageEntry] })); + const newMessage = createAssistantMessage( + randomUUID(), + "Assistant Response Preparing " + randomText(), + { modelSlug: "gpt-4.1", status: "streaming" } + ); + updateStreamingMessage((prev) => ({ ...prev, parts: [...prev.parts, newMessage] })); withDelay(() => { - const newParts = useStreamingStateMachine.getState().streamingMessage.parts.map((part) => - part.messageId === messageEntry.messageId - ? { - ...part, - status: MessageEntryStatus.FINALIZED, - assistant: { - ...part.assistant, - content: "Assistant Response Finalized " + randomText(), - modelSlug: "gpt-4.1", - $typeName: "chat.v2.MessageTypeAssistant", - }, - } - : part, - ) as MessageEntry[]; + const newParts = useStreamingStateMachine.getState().streamingMessage.parts.map((part) => { + if (part.id !== newMessage.id) return part; + if (part.role !== "assistant") return part; + return { + ...part, + status: "complete" as const, + data: { + ...part.data, + content: "Assistant Response Finalized " + randomText(), + }, + }; + }) as InternalMessage[]; updateStreamingMessage((prev) => ({ ...prev, parts: [...newParts] })); }); }; @@ -304,13 +281,13 @@ export const DevTools = () => { Streaming Message
({streamingMessage.parts.length} total, - {streamingMessage.parts.filter((part) => part.status === MessageEntryStatus.PREPARING).length}{" "} - preparing, - {streamingMessage.parts.filter((part) => part.status === MessageEntryStatus.FINALIZED).length}{" "} - finalized, - {streamingMessage.parts.filter((part) => part.status === MessageEntryStatus.INCOMPLETE).length}{" "} - incomplete, - {streamingMessage.parts.filter((part) => part.status === MessageEntryStatus.STALE).length} stale ) + {streamingMessage.parts.filter((part) => part.status === "streaming").length}{" "} + streaming, + {streamingMessage.parts.filter((part) => part.status === "complete").length}{" "} + complete, + {streamingMessage.parts.filter((part) => part.status === "error").length}{" "} + error, + {streamingMessage.parts.filter((part) => part.status === "stale").length} stale )