Associating token usage of stream to user for billing purposes

[nikas-belogolov]

I'm implementing server-side usage tracking in a TanStack Start app. I want to track each user's token usage, but currently there doesnt seem to be a way to do this without intercepting the text stream and checking if chunk is of type "done" to get token usage:

export const Route = createFileRoute("/api/chat")({
  server: {
		handlers: {
			POST: async ({ request }) => {

        if (!process.env.GEMINI_API_KEY) {
          return new Response(
            JSON.stringify({
              error: "API key not configured",
            }),
            {
              status: 500,
              headers: { "Content-Type": "application/json" },
            },
          );
        }

        const { messages, data: { conversationId } } = await request.json();

        console.log("conversationId", conversationId)

        try {

          const stream = chat({
            adapter: geminiText("gemini-2.5-flash"),
            messages,
            conversationId,
          });

          // Intercept the stream to count tokens
          const interceptedStream = (async function* () {
            for await (const chunk of stream) {
              if (chunk.type === "done") {
                console.log("Server-side Token Usage:", chunk.usage);
                // You can save usage to DB here
                
              } 
              yield chunk;
            }
          })();

          // Convert stream to HTTP response
          return toServerSentEventsResponse(interceptedStream);

        } catch (error) {
          console.error(error);
          return new Response(
            JSON.stringify({
              error:
                error instanceof Error ? error.message : "An error occurred",
            }),
            {
              status: 500,
              headers: { "Content-Type": "application/json" },
            },
          );
        }
			},
		},
	},
});

What is the recommended way of associating token usage of a stream to a certain user for billing? I think it's a pretty basic thing for a SaSS apps but I haven't found a good way to do this yet.

Maybe add callbacks to the chat function similar to AI SDK? Or fix aiEventClient to include also other metadata