[Obs AI Assistant] Add tokenCount event to complete event stream in s…

…erverless tests (elastic#205397)

Closes elastic#203407
elastic#203408
elastic#203478

## Summary

### Problem
Some tests in serverless are failing after the migration to the
Inference endpoints. This is due to not emitting the tokenCount event,
which enables successfully closing the stream.

### Solution
Emit the tokenCount event in the serverless tests when the title or
conversation is simulated.

### Checklist

- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios
- [x] The PR description includes the appropriate Release Notes section,
and the correct `release_note:*` label is applied per the
[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)

x-pack/test_serverless/api_integration/test_suites/observability/ai_assistant/tests/complete/complete.spec.ts

@@ -67,8 +67,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
     },
   ];
 
-  // Failing: See https://github.com/elastic/kibana/issues/203407
-  describe.skip('/internal/observability_ai_assistant/chat/complete', function () {
+  describe('/internal/observability_ai_assistant/chat/complete', function () {
     // TODO: https://github.com/elastic/kibana/issues/192751
     this.tags(['skipMKI']);
     let proxy: LlmProxy;
@@ -107,6 +106,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
 
       await titleSimulator.status(200);
       await titleSimulator.next('My generated title');
+      await titleSimulator.tokenCount({ completion: 1, prompt: 1, total: 2 });
       await titleSimulator.complete();
 
       await conversationSimulator.status(200);

x-pack/test_serverless/api_integration/test_suites/observability/ai_assistant/tests/public_complete/public_complete.spec.ts

@@ -49,8 +49,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
     },
   ];
 
-  // Failing: See https://github.com/elastic/kibana/issues/203408
-  describe.skip('/api/observability_ai_assistant/chat/complete', function () {
+  describe('/api/observability_ai_assistant/chat/complete', function () {
     // TODO: https://github.com/elastic/kibana/issues/192751
     this.tags(['skipMKI']);
 
@@ -107,6 +106,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
 
       await titleSimulator.status(200);
       await titleSimulator.next('My generated title');
+      await titleSimulator.tokenCount({ completion: 1, prompt: 1, total: 2 });
       await titleSimulator.complete();
 
       await conversationSimulator.status(200);
@@ -202,6 +202,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
                 },
               ],
             });
+            await conversationSimulator.tokenCount({ completion: 0, prompt: 0, total: 0 });
             await conversationSimulator.complete();
           }
         );
@@ -259,6 +260,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
                 },
               ],
             });
+            await conversationSimulator.tokenCount({ completion: 0, prompt: 0, total: 0 });
             await conversationSimulator.complete();
           }
         );
@@ -275,6 +277,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
       before(async () => {
         responseBody = await getOpenAIResponse(async (conversationSimulator) => {
           await conversationSimulator.next('Hello');
+          await conversationSimulator.tokenCount({ completion: 1, prompt: 1, total: 2 });
           await conversationSimulator.complete();
         });
       });