| 6132 | } |
| 6133 | |
| 6134 | func (p *Server) processChat(ctx context.Context, chat database.Chat) { |
| 6135 | logger := p.logger.With(slog.F("chat_id", chat.ID)) |
| 6136 | logger.Info(ctx, "processing chat request") |
| 6137 | |
| 6138 | p.metrics.Chats.WithLabelValues(chatloop.StateWaiting).Inc() |
| 6139 | defer p.metrics.Chats.WithLabelValues(chatloop.StateWaiting).Dec() |
| 6140 | |
| 6141 | chatCtx, cancel := context.WithCancelCause(ctx) |
| 6142 | defer cancel(nil) |
| 6143 | |
| 6144 | // Gate the control subscriber behind a channel that is closed |
| 6145 | // after we publish "running" status. This prevents stale |
| 6146 | // pubsub notifications (e.g. the "pending" notification from |
| 6147 | // SendMessage that triggered this processing) from |
| 6148 | // interrupting us before we start work. Due to async |
| 6149 | // PostgreSQL NOTIFY delivery, a notification published before |
| 6150 | // subscribeChatControl registers its queue can still arrive |
| 6151 | // after registration. |
| 6152 | controlArmed := make(chan struct{}) |
| 6153 | gatedCancel := func(cause error) { |
| 6154 | select { |
| 6155 | case <-controlArmed: |
| 6156 | cancel(cause) |
| 6157 | default: |
| 6158 | logger.Debug(ctx, "ignoring control notification before armed") |
| 6159 | } |
| 6160 | } |
| 6161 | |
| 6162 | controlCancel := p.subscribeChatControl(chatCtx, chat.ID, gatedCancel, logger) |
| 6163 | defer func() { |
| 6164 | if controlCancel != nil { |
| 6165 | controlCancel() |
| 6166 | } |
| 6167 | }() |
| 6168 | |
| 6169 | // Register with the centralized heartbeat loop instead of |
| 6170 | // running a per-chat goroutine. The loop issues a single batch |
| 6171 | // UPDATE for all chats on this worker and detects stolen chats |
| 6172 | // via set-difference. |
| 6173 | p.registerHeartbeat(&heartbeatEntry{ |
| 6174 | cancelWithCause: cancel, |
| 6175 | chatID: chat.ID, |
| 6176 | workspaceID: chat.WorkspaceID, |
| 6177 | logger: logger, |
| 6178 | }) |
| 6179 | defer p.unregisterHeartbeat(chat.ID) |
| 6180 | |
| 6181 | // Start buffering stream events BEFORE publishing the running |
| 6182 | // status. This closes a race where a subscriber sees |
| 6183 | // status=running but misses message_part events because |
| 6184 | // buffering hasn't started yet — the subscriber gets an empty |
| 6185 | // snapshot and publishToStream drops message_parts while |
| 6186 | // buffering is false. |
| 6187 | streamState := p.getOrCreateStreamState(chat.ID) |
| 6188 | streamState.mu.Lock() |
| 6189 | streamState.buffer = nil |
| 6190 | streamState.bufferRetainedAt = time.Time{} |
| 6191 | streamState.resetDropCounters() |