opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-22 05:05:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(coderd/x/chatd): wake after auto-promoting queued message (#24714)

Browse Source 
`tryAutoPromoteQueuedMessage` in `processChat`'s deferred cleanup could
set a chat back to `pending` without waking the processor. The processor
only noticed on the next 10ms poll, so under load tests like
`TestAutoPromoteQueuedMessageFallsBackForInvalidQueuedModelConfigID`
could time out waiting for the second streaming request (#1500).

Call `p.signalWake()` after the promoted-message publishes when
`promotedMessage != nil`, matching the pattern used by `CreateChat`,
`SendMessage`, `EditMessage`, `PromoteQueued`, and `InterruptChat`. Make
the regression helper `testAutoPromoteQueuedMessageFallback`
deterministic by setting `PendingChatAcquireInterval = time.Hour` and
synchronizing on a `secondRunStarted` channel instead of polling
`requestCount`, so the test fails without the wake instead of relying on
the 10ms ticker.

Closes https://github.com/coder/internal/issues/1500

> Mux is acting on Mike's behalf.
This commit is contained in:
Michael Suchacz
2026-04-26 11:08:32 +02:00
committed by GitHub
parent aee85040f0
commit ed33e28b13
2 changed files with 80 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/x/chatd/chatd.go
+10 -4
View File
@@ -187,9 +187,9 @@ type Server struct {
// and workspace state for the centralized heartbeat loop.
heartbeatRegistry map[uuid.UUID]*heartbeatEntry
// wakeCh is signaled by SendMessage, EditMessage, CreateChat,
// and PromoteQueued so the run loop calls processOnce
// immediately instead of waiting for the next ticker.
// wakeCh is signaled whenever a chat transitions to
// pending so the run loop calls processOnce immediately
// instead of waiting for the next ticker.
wakeCh chan struct{}
}
@@ -5209,7 +5209,6 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) {
QueueUpdate: true,
})
}
if p.shouldPublishFinishedChatState(cleanupCtx, logger, finishResult.updatedChat) {
p.publishStatus(chat.ID, status, uuid.NullUUID{})
// Best-effort: use any generated title captured during
@@ -5222,6 +5221,13 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) {
p.publishChatPubsubEvent(finishResult.updatedChat, codersdk.ChatWatchEventKindStatusChange, nil)
}
if promotedMessage != nil {
// Wake the processor so it picks up the newly pending
// chat immediately instead of waiting for the next
// acquire-interval tick.
p.signalWake()
}
// When the chat is parked in requires_action,
// publish the stream event and global pubsub event
// after the DB status has committed. Publishing
coderd/x/chatd/chatd_test.go
+70 -43
View File
@@ -2687,6 +2687,8 @@ func TestAutoPromoteQueuedMessagesPreservesPerTurnModelOrder(t *testing.T) {
ctx := testutil.Context(t, testutil.WaitSuperLong)
firstRunStarted := make(chan struct{})
secondRunStarted := make(chan struct{}, 1)
thirdRunStarted := make(chan struct{}, 1)
allowFirstRunFinish := make(chan struct{})
var requestCount atomic.Int32
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
@@ -2709,15 +2711,27 @@ func TestAutoPromoteQueuedMessagesPreservesPerTurnModelOrder(t *testing.T) {
}()
return chattest.OpenAIResponse{StreamingChunks: chunks}
case 2:
select {
case secondRunStarted <- struct{}{}:
default:
}
return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("second run done")...)
case 3:
select {
case thirdRunStarted <- struct{}{}:
default:
}
return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("third run done")...)
default:
return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("extra run done")...)
}
})
server := newActiveTestServer(t, db, ps)
server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) {
// Disable periodic polling so chained promotions must be driven by
// signalWake.
cfg.PendingChatAcquireInterval = time.Hour
})
user, org, modelConfigA := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
modelConfigB := insertChatModelConfigWithCallConfig(
ctx,
@@ -2769,9 +2783,9 @@ func TestAutoPromoteQueuedMessagesPreservesPerTurnModelOrder(t *testing.T) {
close(allowFirstRunFinish)
require.Eventually(t, func() bool {
return requestCount.Load() >= 3
}, testutil.WaitSuperLong, testutil.IntervalFast)
testutil.TryReceive(ctx, t, secondRunStarted)
testutil.TryReceive(ctx, t, thirdRunStarted)
require.GreaterOrEqual(t, requestCount.Load(), int32(3))
chatd.WaitUntilIdleForTest(server)
queuedMessages, err := db.GetChatQueuedMessages(ctx, chat.ID)
@@ -2825,6 +2839,7 @@ func testAutoPromoteQueuedMessageFallback(t *testing.T, queuedModelConfigID uuid
ctx := testutil.Context(t, testutil.WaitSuperLong)
firstRunStarted := make(chan struct{})
secondRunStarted := make(chan struct{}, 1)
allowFirstRunFinish := make(chan struct{})
var requestCount atomic.Int32
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
@@ -2847,11 +2862,19 @@ func testAutoPromoteQueuedMessageFallback(t *testing.T, queuedModelConfigID uuid
}()
return chattest.OpenAIResponse{StreamingChunks: chunks}
default:
select {
case secondRunStarted <- struct{}{}:
default:
}
return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("fallback run done")...)
}
})
server := newActiveTestServer(t, db, ps)
server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) {
// Disable periodic polling so only signalWake can
// trigger the next processing run.
cfg.PendingChatAcquireInterval = time.Hour
})
user, org, modelConfig := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
OrganizationID: org.ID,
@@ -2875,9 +2898,8 @@ func testAutoPromoteQueuedMessageFallback(t *testing.T, queuedModelConfigID uuid
close(allowFirstRunFinish)
require.Eventually(t, func() bool {
return requestCount.Load() >= 2
}, testutil.WaitSuperLong, testutil.IntervalFast)
testutil.TryReceive(ctx, t, secondRunStarted)
require.GreaterOrEqual(t, requestCount.Load(), int32(2))
chatd.WaitUntilIdleForTest(server)
queuedMessages, err := db.GetChatQueuedMessages(ctx, chat.ID)
@@ -3000,8 +3022,6 @@ func TestPromoteQueuedMessageFallsBackForInvalidQueuedModelConfigID(t *testing.T
func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) {
t.Parallel()
const acquireInterval = 10 * time.Millisecond
db, ps := dbtestutil.NewDB(t)
ctx := testutil.Context(t, testutil.WaitLong)
@@ -3013,27 +3033,14 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) {
require.NoError(t, err)
clock := quartz.NewMock(t)
acquireTrap := clock.Trap().NewTicker("chatd", "acquire")
defer acquireTrap.Close()
assertPendingWithoutQueuedMessages := func(chatID uuid.UUID) {
t.Helper()
queued, dbErr := db.GetChatQueuedMessages(ctx, chatID)
require.NoError(t, dbErr)
require.Empty(t, queued)
fromDB, dbErr := db.GetChatByID(ctx, chatID)
require.NoError(t, dbErr)
require.Equal(t, database.ChatStatusPending, fromDB.Status)
require.False(t, fromDB.WorkerID.Valid)
}
streamStarted := make(chan struct{})
interrupted := make(chan struct{})
secondRequestStarted := make(chan struct{})
thirdRequestStarted := make(chan struct{})
secondRequestStarted := make(chan struct{}, 1)
thirdRequestStarted := make(chan struct{}, 1)
allowFinish := make(chan struct{})
allowSecondRequestFinish := make(chan struct{})
allowThirdRequestFinish := make(chan struct{})
var requestCount atomic.Int32
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
if !req.Stream {
@@ -3061,9 +3068,35 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) {
}()
return chattest.OpenAIResponse{StreamingChunks: chunks}
case 2:
close(secondRequestStarted)
select {
case secondRequestStarted <- struct{}{}:
default:
}
chunks := make(chan chattest.OpenAIChunk, 1)
go func() {
defer close(chunks)
chunks <- chattest.OpenAITextChunks("second run partial")[0]
select {
case <-allowSecondRequestFinish:
case <-req.Context().Done():
}
}()
return chattest.OpenAIResponse{StreamingChunks: chunks}
case 3:
close(thirdRequestStarted)
select {
case thirdRequestStarted <- struct{}{}:
default:
}
chunks := make(chan chattest.OpenAIChunk, 1)
go func() {
defer close(chunks)
chunks <- chattest.OpenAITextChunks("third run partial")[0]
select {
case <-allowThirdRequestFinish:
case <-req.Context().Done():
}
}()
return chattest.OpenAIResponse{StreamingChunks: chunks}
}
return chattest.OpenAIStreamingResponse(
@@ -3073,10 +3106,11 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) {
server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) {
cfg.Clock = clock
cfg.PendingChatAcquireInterval = acquireInterval
// Keep periodic polling frozen so request handoff is synchronized
// through explicit mock channels.
cfg.PendingChatAcquireInterval = time.Hour
cfg.InFlightChatStaleAfter = testutil.WaitSuperLong
})
acquireTrap.MustWait(ctx).MustRelease(ctx)
user, org, model := seedChatDependencies(ctx, t, db)
setOpenAIProviderBaseURL(ctx, t, db, openAIURL)
@@ -3090,7 +3124,6 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) {
})
require.NoError(t, err)
clock.Advance(acquireInterval).MustWait(ctx)
testutil.TryReceive(ctx, t, streamStarted)
queuedResult, err := server.SendMessage(ctx, chatd.SendMessageOptions{
@@ -3105,12 +3138,8 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) {
testutil.TryReceive(ctx, t, interrupted)
close(allowFinish)
chatd.WaitUntilIdleForTest(server)
assertPendingWithoutQueuedMessages(chat.ID)
testutil.TryReceive(ctx, t, secondRequestStarted)
// Keep the acquire loop frozen here so "queued" stays pending.
// That makes the later send queue because the chat is still busy,
// rather than because the scheduler happened to be slow.
laterQueuedResult, err := server.SendMessage(ctx, chatd.SendMessageOptions{
ChatID: chat.ID,
Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("later queued")},
@@ -3159,13 +3188,11 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) {
})
require.NoError(t, err)
clock.Advance(acquireInterval).MustWait(ctx)
testutil.TryReceive(ctx, t, secondRequestStarted)
chatd.WaitUntilIdleForTest(server)
assertPendingWithoutQueuedMessages(chat.ID)
clock.Advance(acquireInterval).MustWait(ctx)
close(allowSecondRequestFinish)
testutil.TryReceive(ctx, t, thirdRequestStarted)
require.GreaterOrEqual(t, requestCount.Load(), int32(3))
close(allowThirdRequestFinish)
chatd.WaitUntilIdleForTest(server)
queued, err := db.GetChatQueuedMessages(ctx, chat.ID)