Commit Graph

4 Commits

Author SHA1 Message Date
Mathias Fredriksson bdbcd3428b feat(coderd/chatd): unify chat storage on SDK parts and fix file-reference rendering (#22958)
File-reference parts in user messages were flattened to `TextContent` at
write time because fantasy has no file-reference content type. The
frontend never saw them as structured parts.

This moves all write paths (user, assistant, tool) from fantasy envelope
format to `codersdk.ChatMessagePart`. The streaming layer (`chatloop`)
is untouched, the conversion happens at the serialization boundary in
`persistStep`.

Old rows are still readable. `ParseContent` uses a structural heuristic
(`isFantasyEnvelopeFormat`) to distinguish legacy envelopes from SDK
parts. We chose this over try/fallback because fantasy envelopes
partially unmarshal into `ChatMessagePart` (the `type` field matches)
while silently losing content. A guard test enforces that no SDK part
can produce the envelope shape.

This is forward-only: new rows are unreadable by old code. Chat is
behind a feature flag so rollback risk is contained.

Also adds a typed `ChatMessageRole` to replace raw strings and
`fantasy.MessageRole*` casts at the persistence boundary. The type
covers `ChatMessage.Role`, `ChatStreamMessagePart.Role`, the
`PublishMessagePart` callback chain, and all DB write sites.
`fantasy.MessageRole*` remains only where we build `fantasy.Message`
structs for LLM dispatch.

Separately, `ProviderMetadata` was leaking to SSE clients via
`publishMessagePart`. `StripInternal` now runs on both the SSE and REST
paths, covering this.

Other cleanup:

- Old `db2sdk.contentBlockToPart` silently dropped metadata on
text/reasoning/tool-call content. New code preserves it.
- `providerMetadataToOptions` now logs warnings instead of silently
returning nil.
- `db2sdk` shrinks from ~250 lines of parallel conversion to ~15 lines
delegating to `chatprompt.ParseContent()`, removing the `fantasy` import
entirely.

Refs #22821
2026-03-13 17:53:26 +02:00
Kyle Carberry eecb7d0b66 fix: resolve bugs in chatd streaming system (#22720)
Split from #22693 per review feedback.

Fixes multiple bugs in coderd/chatd and sub-packages including race
conditions, transaction safety, stream buffer bounds, retry limits, and
enterprise relay improvements.

See commit message for full list.
2026-03-06 21:02:25 +00:00
Kyle Carberry 94a2e440a8 fix(chatd): extract session token from cookie for relay header (#22649)
## Problem

When a browser connects to the chat stream via WebSocket, it
authenticates using cookies only — the native WebSocket API cannot set
custom headers like `Coder-Session-Token`. The relay between replicas
copies the original request's `Cookie` header but did **not** set the
`Coder-Session-Token` header as a fallback.

This causes a **401 on the worker replica** when `EnableHostPrefix` is
enabled, because the `HTTPCookies.Middleware` strips bare
`coder_session_token` cookies (expecting the `__Host-` prefix). Without
a `Coder-Session-Token` header fallback, `apiKeyMiddleware` finds no
valid credentials.

### Root Cause

The data flow:
1. Browser → subscriber replica: `Cookie:
__Host-coder_session_token=xxx` (browser sends prefixed cookie)
2. Subscriber's `HTTPCookies.Middleware` normalizes: `Cookie:
coder_session_token=xxx` (strips prefix)
3. `relayHeaders()` copies `Cookie: coder_session_token=xxx` to relay
request
4. Worker replica's `HTTPCookies.Middleware` sees bare
`coder_session_token` → **strips it** (expects `__Host-` prefix)
5. `apiKeyMiddleware` → `APITokenFromRequest`: no cookie, no header →
**401**

## Fix

Modified `relayHeaders()` to extract the session token value from the
`Cookie` header and set it as the `Coder-Session-Token` header when no
explicit session token header is already present. The header is never
stripped by middleware, so the worker replica can always authenticate.

## Testing

- **`TestRelayHeaders`**: Unit tests for the updated `relayHeaders()`
function covering all scenarios (cookie-only, header+cookie, no auth,
nil source)
- **`TestExtractSessionTokenFromCookieHeader`**: Unit tests for the
helper function
- **`TestChatStreamRelay/RelayCookieOnlyAuth`**: Integration test with
plain HTTP, cookie-only WebSocket auth
- **`TestChatStreamRelay/RelayCookieOnlyAuthWithHostPrefix`**:
Integration test with `EnableHostPrefix=true`, confirming the 401 is
fixed
- **`cookieOnlySessionTokenProvider`**: Test helper that simulates
browser WebSocket behavior (sets Cookie header only on WebSocket dials,
no custom headers)

## Files Changed

- `enterprise/coderd/chatd/chatd.go` — `relayHeaders()` fix +
`extractSessionTokenFromCookieHeader()` helper
- `enterprise/coderd/chatd/relay_headers_internal_test.go` — unit tests
(new file)
- `enterprise/coderd/chats_test.go` — integration tests + test helper
type
2026-03-05 05:11:07 +00:00
Kyle Carberry 30d534b36b fix(chatd): fix relay race conditions, extract enterprise relay logic, move pubsub to OSS (#22589)
## Summary

Fixes a bug where interrupting a streaming chat and sending a new
message
left the relay connected to the wrong replica. Expanded into a broader
refactor that cleanly separates concerns:

- **OSS** owns pubsub subscription, message catch-up, queue updates,
  status forwarding, and local parts merging.
- **Enterprise** (`enterprise/coderd/chatd`) only manages relay dialing,
  reconnection, and stale-dial discarding for cross-replica streaming.

## Architecture

### OSS `coderd/chatd/chatd.go`

`Subscribe()` builds the initial snapshot then runs a single merge
goroutine that handles:

- Pubsub subscription for durable events (status, messages, queue,
errors)
- Message catch-up via `AfterMessageID`
- Local `message_part` forwarding
- Relay events from enterprise (when `SubscribeFn` is set)
- Sends `StatusNotification` to enterprise so it can manage relay
lifecycle

Key types:

- `SubscribeFn` — enterprise hook, returns relay-only events channel
- `SubscribeFnParams` — `ChatID`, `Chat`, `WorkerID`,
`StatusNotifications`, `RequestHeader`, `DB`, `Logger`
- `StatusNotification` — `Status` + `WorkerID`, sent to enterprise on
pubsub status changes

### Enterprise `enterprise/coderd/chatd/chatd.go`

`NewMultiReplicaSubscribeFn(cfg MultiReplicaSubscribeConfig)` returns a
`SubscribeFn` that:

- Opens an initial synchronous relay if the chat is running on a remote
worker
- Reads `StatusNotifications` from OSS to open/close relay connections
- Handles async dial, reconnect timers, stale-dial discarding
- Returns only relay `message_part` events

## Bug fixes

### Original bug: stale relay dial after interrupt

`openRelayAsync` goroutines used `mergedCtx` (subscription-level), not a
per-dial context. `closeRelay()` could not cancel in-flight dials. When
the user interrupts and a new replica picks up the chat, the old dial
goroutine could complete after the new one and deliver a stale
`relayResult`.

**Fix**: per-dial `dialCtx`/`dialCancel`, `expectedWorkerID` tracking,
`workerID` on `relayResult`. `closeRelay()` cancels the dial context and
drains `relayReadyCh`. Merge loop rejects mismatched worker IDs.

### Additional fixes

- `statusNotifications` send-on-closed-channel race — goroutine now owns
  `close()` via defer
- Enterprise spin-loop on `StatusNotifications` close — two-value
receive
  with nil-out
- `hasPubsub` set from `p.pubsub != nil` instead of subscription success
  — now tracks actual subscription result
- `lastMessageID` not initialized from `afterMessageID` — caused
  duplicate messages on catch-up
- `wrappedParts` goroutine leaked remote connection on `dialCtx` cancel
- `closeRelay()` did not drain `relayReadyCh`
- `setChatWaiting` race with `SendMessage(Interrupt)` — wrapped in
`InTx`
- `processChat` post-TX side effects fired when chat was taken by
another
  worker — added `errChatTakenByOtherWorker` sentinel
- Cancel closure data race on `reconnectTimer`
- Bare blocking send on pubsub error path
- `localParts` hot-spin after channel close
- No-pubsub branch dropped relay events and initial snapshot
- Failed relay dial caused permanent stall (no reconnect retry)
- DB error during reconnect timer caused permanent stall
- `time.NewTimer` replaced with `quartz.Clock` for testable timing

## Tests

9 enterprise tests covering:

- Relay reconnect on drop (mock clock)
- Async dial does not block merge loop
- Relay snapshot delivery
- Stale dial discarded after interrupt
- Cancel during in-flight dial
- Running-to-running worker switch
- Failed dial retries (mock clock)
- Local worker closes relay
- Multiple consecutive reconnects (mock clock)

All pass with `-race`.
2026-03-04 18:42:28 -05:00