Function TestInterruptChatPersistsPartialResponse

coderd/x/chatd/chatd_test.go:7210–7346

(t *testing.T)

}

func TestInterruptChatPersistsPartialResponse(t *testing.T) {
	t.Parallel()

	db, ps := dbtestutil.NewDB(t)
	ctx := testutil.Context(t, testutil.WaitLong)

	// Set up a mock OpenAI that streams a partial response and then
	// blocks until the request context is canceled (simulating an
	// interrupt mid-stream).
	chunksDelivered := make(chan struct{})
	openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
		if !req.Stream {
			return chattest.OpenAINonStreamingResponse("title")
		}
		chunks := make(chan chattest.OpenAIChunk, 1)
		go func() {
			defer close(chunks)
			// Send two partial text chunks so there is meaningful
			// content to persist.
			for _, c := range chattest.OpenAITextChunks("hello world") {
				chunks <- c
			}
			// Signal that chunks have been written to the HTTP response.
			select {
			case <-chunksDelivered:
			default:
				close(chunksDelivered)
			}
			// Block until interrupt cancels the context.
			<-req.Context().Done()
		}()
		return chattest.OpenAIResponse{StreamingChunks: chunks}
	})

	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
	server := chatd.New(chatd.Config{
		Logger:                     logger,
		Database:                   db,
		ReplicaID:                  uuid.New(),
		Pubsub:                     ps,
		PendingChatAcquireInterval: 10 * time.Millisecond,
		InFlightChatStaleAfter:     testutil.WaitSuperLong,
	})
	server.Start()
	t.Cleanup(func() {
		require.NoError(t, server.Close())
	})

	user, org, model := seedChatDependencies(t, db)
	setOpenAIProviderBaseURL(ctx, t, db, openAIURL)

	chat, err := server.CreateChat(ctx, chatd.CreateOptions{
		OrganizationID:     org.ID,
		OwnerID:            user.ID,
		Title:              "interrupt-persist-test",
		ModelConfigID:      model.ID,
		InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
	})