feat: record first token latency in runtime router

tkmind-proxy.mjs

@@ -219,6 +219,25 @@ function createRuntimeRouter({
       if (streamKey) multi.set(streamKey, status, { EX: 600 });
       await multi.exec().catch(() => null);
     },
+    async firstTokenObserved(target, latencyMs) {
+      const client = await getClient();
+      if (!client || !target) return;
+      const workerId = workerIdForTarget(target);
+      const sample = Math.max(0, Math.round(Number(latencyMs) || 0));
+      const ewmaKey = key('worker', workerId, 'ewma_first_token_ms');
+      const prev = readNumber(await client.get(ewmaKey).catch(() => null));
+      const next = prev > 0 ? Math.round(prev * 0.8 + sample * 0.2) : sample;
+      const now = String(Date.now());
+      await client
+        .multi()
+        .set(ewmaKey, String(next))
+        .set(key('worker', workerId, 'last_first_token_ms'), String(sample))
+        .set(key('worker', workerId, 'last_first_token_at'), now)
+        .incr(key('worker', workerId, 'first_token_count'))
+        .set(key('worker', workerId, 'heartbeat'), now, { EX: 30 })
+        .exec()
+        .catch(() => null);
+    },
     async getStatus() {
       const client = await getClient();
       const workers = [];
@@ -245,6 +264,9 @@ function createRuntimeRouter({
                 key('worker', workerId, 'container_pids'),
                 key('worker', workerId, 'container_health'),
                 key('worker', workerId, 'metrics_sampled_at'),
+                key('worker', workerId, 'last_first_token_ms'),
+                key('worker', workerId, 'last_first_token_at'),
+                key('worker', workerId, 'first_token_count'),
               ])
               .catch(() => [])
           : [];
@@ -269,6 +291,9 @@ function createRuntimeRouter({
           containerPids: readNumber(values?.[15]),
           containerHealth: values?.[16] ?? null,
           metricsSampledAt: values?.[17] ? Number(values[17]) : null,
+          lastFirstTokenMs: readNumber(values?.[18]),
+          lastFirstTokenAt: values?.[19] ? Number(values[19]) : null,
+          firstTokenCount: readNumber(values?.[20]),
           score: workerScoreFromValues(values),
         });
       }
@@ -909,6 +934,11 @@ export function createTkmindProxy({
     await runtimeRouter.streamEnded(sessionId, target, options).catch(() => null);
   }
 
+  function markFirstTokenObserved(target, latencyMs) {
+    if (!runtimeRouter || !target) return;
+    void runtimeRouter.firstTokenObserved(target, latencyMs).catch(() => null);
+  }
+
   async function getRuntimeStatus() {
     const targetStatuses = [];
     for (const target of targets) {
@@ -1503,6 +1533,7 @@ export function createTkmindProxy({
 
   const proxySessionEvents = async (req, res, sessionId, { onAfterFinish, onEvent } = {}) => {
     const upstreamAbort = new AbortController();
+    const streamRequestedAt = Date.now();
     let clientClosed = false;
     const abortUpstream = () => {
       clientClosed = true;
@@ -1564,6 +1595,16 @@ export function createTkmindProxy({
       });
 
       const source = Readable.fromWeb(upstream.body);
+      let firstChunkSeen = false;
+      const firstTokenProbe = new Transform({
+        transform(chunk, _encoding, callback) {
+          if (!firstChunkSeen && chunk?.length > 0) {
+            firstChunkSeen = true;
+            markFirstTokenObserved(sessionTarget, Date.now() - streamRequestedAt);
+          }
+          callback(null, chunk);
+        },
+      });
       const linkSanitizer = createSessionEventSanitizer(req.currentUser, { onEvent });
       const waitForDrain = () => new Promise((resolve) => res.once('drain', resolve));
       const writeClientChunk = async (chunk) => {
@@ -1595,7 +1636,7 @@ export function createTkmindProxy({
       let streamCloseStatus = 'closed';
       try {
         try {
-          await pipeline(source, linkSanitizer, billingTransform, clientSink);
+          await pipeline(source, firstTokenProbe, linkSanitizer, billingTransform, clientSink);
         } catch (err) {
           streamCloseStatus = clientClosed || upstreamAbort.signal.aborted ? 'aborted' : 'error';
           throw err;