| 170 | } |
| 171 | |
| 172 | func (r *Runner) runConversation(ctx context.Context, chatID uuid.UUID, logger slog.Logger, events <-chan codersdk.ChatStreamEvent) error { |
| 173 | r.chatID = chatID |
| 174 | |
| 175 | for event := range events { |
| 176 | r.result.eventCount++ |
| 177 | |
| 178 | switch event.Type { |
| 179 | case codersdk.ChatStreamEventTypeStatus: |
| 180 | if event.Status == nil { |
| 181 | continue |
| 182 | } |
| 183 | done, err := r.handleStatusEvent(ctx, chatID, logger, event.Status.Status) |
| 184 | if err != nil { |
| 185 | return err |
| 186 | } |
| 187 | if done { |
| 188 | return nil |
| 189 | } |
| 190 | case codersdk.ChatStreamEventTypeMessagePart: |
| 191 | r.handleMessagePartEvent(ctx, logger) |
| 192 | case codersdk.ChatStreamEventTypeMessage: |
| 193 | // StreamChat replays persisted rows as message events, not |
| 194 | // message_part deltas, when a turn finished server-side before |
| 195 | // the stream attached. Route assistant rows through the same |
| 196 | // first-output path; skip user rows so persisted prompts do not |
| 197 | // count as model output. |
| 198 | if event.Message == nil || event.Message.Role != codersdk.ChatMessageRoleAssistant { |
| 199 | continue |
| 200 | } |
| 201 | r.handleMessagePartEvent(ctx, logger) |
| 202 | case codersdk.ChatStreamEventTypeRetry: |
| 203 | r.handleRetryEvent(ctx, logger, event.Retry) |
| 204 | case codersdk.ChatStreamEventTypeError: |
| 205 | r.handleErrorEvent(ctx, logger, event.Error) |
| 206 | } |
| 207 | } |
| 208 | |
| 209 | if ctx.Err() != nil { |
| 210 | return ctx.Err() |
| 211 | } |
| 212 | |
| 213 | r.result.failureStage = failureStageStreamEndedEarly |
| 214 | r.cfg.Metrics.ChatStageFailuresTotal.WithLabelValues(r.result.failureStage).Inc() |
| 215 | if r.lastStreamError != "" { |
| 216 | return xerrors.Errorf("chat %s stream ended before completing %d of %d turns: %s", chatID, r.result.turnsCompleted, r.cfg.Turns, r.lastStreamError) |
| 217 | } |
| 218 | return xerrors.Errorf("chat %s stream ended before completing %d of %d turns", chatID, r.result.turnsCompleted, r.cfg.Turns) |
| 219 | } |
| 220 | |
| 221 | func (r *Runner) handleStatusEvent(ctx context.Context, chatID uuid.UUID, logger slog.Logger, status codersdk.ChatStatus) (bool, error) { |
| 222 | if status == r.lastStatus { |