Function TestPersistToolResultWithBinaryData

coderd/x/chatd/chatd_test.go:4021–4197

(t *testing.T)

}

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

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

	const binaryOutputBase64 = "SEVBREVSAAAAc29tZSBkYXRhAABtb3JlIGRhdGEARU5E"
	binaryOutput, err := io.ReadAll(base64.NewDecoder(
		base64.StdEncoding,
		strings.NewReader(binaryOutputBase64),
	))
	require.NoError(t, err)

	var streamedCallCount atomic.Int32
	var streamedCallsMu sync.Mutex
	streamedCalls := make([][]chattest.OpenAIMessage, 0, 2)

	openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
		if !req.Stream {
			return chattest.OpenAINonStreamingResponse("Binary tool result test")
		}

		streamedCallsMu.Lock()
		streamedCalls = append(streamedCalls, append([]chattest.OpenAIMessage(nil), req.Messages...))
		streamedCallsMu.Unlock()

		if streamedCallCount.Add(1) == 1 {
			return chattest.OpenAIStreamingResponse(
				chattest.OpenAIToolCallChunk(
					"execute",
					`{"command":"cat /home/coder/binary_file.bin"}`,
				),
			)
		}
		// Include literal \u0000 in the response text, which is
		// what a real LLM writes when explaining binary output.
		// json.Marshal encodes the backslash as \\, producing
		// \\u0000 in the JSON bytes. The sanitizer must not
		// corrupt this into invalid JSON.
		return chattest.OpenAIStreamingResponse(
			chattest.OpenAITextChunks("The file contains \\u0000 null bytes.")...,
		)
	})

	// Use "openai-compat" provider so the chatd framework uses the
	// /chat/completions endpoint, where the mock server supports
	// streaming tool calls. The default "openai" provider routes to
	// /responses which only handles text deltas in the mock.
	user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL)
	ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID)

	ctrl := gomock.NewController(t)
	mockConn := agentconnmock.NewMockAgentConn(ctrl)
	mockConn.EXPECT().
		SetExtraHeaders(gomock.Any()).
		AnyTimes()
	mockConn.EXPECT().
		ContextConfig(gomock.Any()).