Skip to content

test: accept context in chatd idle test hook #26509

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
hugodutka wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

test: accept context in chatd idle test hook #26509

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0c6fe1f
deferred review threads 26109 plan
hugodutka Jun 12, 2026
063179e
test: accept context in chatd idle test hook
hugodutka Jun 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion coderd/coderdtest/chat.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ func WaitForChatSettled(

server := api.ChatDaemonForTest()
require.NotNil(t, server)
chatd.WaitUntilIdleForTest(server)
require.NoError(t, chatd.WaitUntilIdleForTest(ctx, server))

chat, err := getChatByIDAsSystem(ctx, api.Database, chatID)
require.NoError(t, err)
Expand Down
24 changes: 12 additions & 12 deletions coderd/x/chatd/chatd_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2206,7 +2206,7 @@ func TestAutoPromoteQueuedMessagesPreservesPerTurnModelOrder(t *testing.T) {
testutil.TryReceive(ctx, t, secondRunStarted)
testutil.TryReceive(ctx, t, thirdRunStarted)
require.GreaterOrEqual(t, requestCount.Load(), int32(3))
chatd.WaitUntilIdleForTest(server)
require.NoError(t, chatd.WaitUntilIdleForTest(ctx, server))

queuedMessages, err := db.GetChatQueuedMessages(ctx, chat.ID)
require.NoError(t, err)
Expand Down Expand Up @@ -2397,7 +2397,7 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) {
require.GreaterOrEqual(t, requestCount.Load(), int32(3))

close(allowThirdRequestFinish)
chatd.WaitUntilIdleForTest(server)
require.NoError(t, chatd.WaitUntilIdleForTest(ctx, server))

queued, err := db.GetChatQueuedMessages(ctx, chat.ID)
require.NoError(t, err)
Expand Down Expand Up @@ -2646,7 +2646,7 @@ func TestEditMessageDebugCleanupDeletesPreEditRuns(t *testing.T) {
})
require.NoError(t, err)

chatd.WaitUntilIdleForTest(replica)
require.NoError(t, chatd.WaitUntilIdleForTest(ctx, replica))

// ErrNoRows on staleRun proves the fast-retry path DELETED the
// row: FinalizeStale (the only other debug-row writer on the
Expand Down Expand Up @@ -2737,7 +2737,7 @@ func TestEditMessageDebugCleanupPreservesRecentRuns(t *testing.T) {
})
require.NoError(t, err)

chatd.WaitUntilIdleForTest(replica)
require.NoError(t, chatd.WaitUntilIdleForTest(ctx, replica))

remaining, err := db.GetChatDebugRunByID(ctx, recentRun.ID)
require.NoError(t, err,
Expand Down Expand Up @@ -3323,7 +3323,7 @@ func TestRequiresActionChatPersistsWaitingStatusLabel(t *testing.T) {
got.LastTurnSummary.Valid &&
got.LastTurnSummary.String == "Waiting for user input"
}, testutil.IntervalFast)
chatd.WaitUntilIdleForTest(server)
require.NoError(t, chatd.WaitUntilIdleForTest(ctx, server))

require.Equal(t, database.ChatStatusRequiresAction, fromDB.Status,
"expected requires_action, got %s (last_error=%q)",
Expand Down Expand Up @@ -5329,7 +5329,7 @@ func waitForChatProcessed(
return c.Status != database.ChatStatusPending &&
c.Status != database.ChatStatusRunning
}, testutil.WaitShort, testutil.IntervalFast)
chatd.WaitUntilIdleForTest(server)
require.NoError(t, chatd.WaitUntilIdleForTest(ctx, server))
}

// newTestServer creates a passive server that never calls
Expand Down Expand Up @@ -8194,7 +8194,7 @@ func TestPassiveServerDoesNotProcess(t *testing.T) {
})
require.NoError(t, err)

chatd.WaitUntilIdleForTest(server)
require.NoError(t, chatd.WaitUntilIdleForTest(ctx, server))

// Re-read from DB to catch any unexpected processing.
stored, err := db.GetChatByID(ctx, chat.ID)
Expand Down Expand Up @@ -8806,7 +8806,7 @@ func TestInterruptChatDoesNotSendWebPushNotification(t *testing.T) {
}
return fromDB.Status == database.ChatStatusWaiting && !fromDB.WorkerID.Valid
}, testutil.IntervalFast)
chatd.WaitUntilIdleForTest(server)
require.NoError(t, chatd.WaitUntilIdleForTest(ctx, server))

fromDB, err := db.GetChatByID(ctx, chat.ID)
require.NoError(t, err)
Expand Down Expand Up @@ -9267,7 +9267,7 @@ func TestErroredChatClearsLastTurnSummaryAndSendsWebPush(t *testing.T) {
fromDB.Status == database.ChatStatusError &&
mockPush.dispatchCount.Load() >= 1
}, testutil.IntervalFast)
chatd.WaitUntilIdleForTest(server)
require.NoError(t, chatd.WaitUntilIdleForTest(ctx, server))

fromDB, err := db.GetChatByID(ctx, chat.ID)
require.NoError(t, err)
Expand Down Expand Up @@ -11272,7 +11272,7 @@ func TestCreateChatImmediatelyProcessesNewChat(t *testing.T) {
// this receive would time out.
testutil.TryReceive(ctx, t, processed)

chatd.WaitUntilIdleForTest(server)
require.NoError(t, chatd.WaitUntilIdleForTest(ctx, server))

// Verify the chat was fully processed.
fromDB, err := db.GetChatByID(ctx, chat.ID)
Expand Down Expand Up @@ -11335,7 +11335,7 @@ func TestSendMessageImmediatelyProcessesWaitingChat(t *testing.T) {
// wait for the processing goroutine to finish so the chat
// transitions to "waiting" status.
testutil.TryReceive(ctx, t, firstProcessed)
chatd.WaitUntilIdleForTest(server)
require.NoError(t, chatd.WaitUntilIdleForTest(ctx, server))

// Now send a follow-up message, which should also be
// processed immediately without waiting for the acquire ticker.
Expand All @@ -11347,7 +11347,7 @@ func TestSendMessageImmediatelyProcessesWaitingChat(t *testing.T) {
require.NoError(t, err)

testutil.TryReceive(ctx, t, secondProcessed)
chatd.WaitUntilIdleForTest(server)
require.NoError(t, chatd.WaitUntilIdleForTest(ctx, server))

// Both turns processed. Verify the second request reached the LLM.
require.GreaterOrEqual(t, requestCount.Load(), int32(2),
Expand Down
6 changes: 3 additions & 3 deletions coderd/x/chatd/recording_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,7 +188,7 @@ func TestWaitAgentComputerUseRecording(t *testing.T) {

// Wait for background processing triggered by CreateChat to
// settle before setting up the mock agent connection.
WaitUntilIdleForTest(server)
require.NoError(t, WaitUntilIdleForTest(ctx, server))

// Now wire up the mock agent connection.
server.agentConnFn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
Expand Down Expand Up @@ -274,7 +274,7 @@ func TestWaitAgentComputerUseRecordingWithThumbnail(t *testing.T) {
"parent-recording-thumb", "computer-use-child-thumb",
)

WaitUntilIdleForTest(server)
require.NoError(t, WaitUntilIdleForTest(ctx, server))

server.agentConnFn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
require.Equal(t, agent.ID, agentID)
Expand Down Expand Up @@ -368,7 +368,7 @@ func TestWaitAgentNonComputerUseNoRecording(t *testing.T) {

// Wait for background processing triggered by CreateChat to
// settle before setting up the mock agent connection.
WaitUntilIdleForTest(server)
require.NoError(t, WaitUntilIdleForTest(ctx, server))

// Wire up the mock agent connection. The mock has zero
// expectations — gomock will fail if StartDesktopRecording
Expand Down
6 changes: 3 additions & 3 deletions coderd/x/chatd/subagent_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3315,7 +3315,7 @@ func TestWaitAgentDoesNotRelayRegularSubagentAttachments(t *testing.T) {
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})

parent, child := createParentChildChats(ctx, t, server, user, org, model)
WaitUntilIdleForTest(server)
require.NoError(t, WaitUntilIdleForTest(ctx, server))

insertedFile := insertLinkedChatFile(
ctx,
Expand Down Expand Up @@ -3503,7 +3503,7 @@ func TestAwaitSubagentCompletion(t *testing.T) {

// signalWake from CreateChat may trigger immediate processing.
// Wait for it to settle, then reset chats to the state we need.
WaitUntilIdleForTest(server)
require.NoError(t, WaitUntilIdleForTest(ctx, server))
setChatStatus(ctx, t, db, parent.ID, database.ChatStatusRunning, "")
setChatStatus(ctx, t, db, child.ID, database.ChatStatusRunning, "")

Expand Down Expand Up @@ -3582,7 +3582,7 @@ func TestAwaitSubagentCompletion(t *testing.T) {

// This case should return immediately, so use the shared
// real-clock server instead of a mock clock.
WaitUntilIdleForTest(server)
require.NoError(t, WaitUntilIdleForTest(ctx, server))
setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")

gotChat, report, err := server.awaitSubagentCompletion(
Expand Down
14 changes: 6 additions & 8 deletions coderd/x/chatd/testhooks.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,19 +1,17 @@
package chatd

import (
"context"
"time"
)
import "context"

// WaitUntilIdleForTest waits for background chat work tracked by the server to
// finish without shutting the server down. Tests use this to assert final
// database state only after asynchronous chat processing has completed.
// Close waits for the same tracked work, but also stops the server.
func WaitUntilIdleForTest(server *Server) {
func WaitUntilIdleForTest(ctx context.Context, server *Server) error {
if server.chatWorker != nil {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
_ = server.chatWorker.WaitIdle(ctx)
if err := server.chatWorker.WaitIdle(ctx); err != nil {
return err
}
}
server.drainInflight()
return nil
}
Loading
Loading