feat(agentic-server): relay streaming chat completions and meter their usage

agentic/agentic-server/README.md

@@ -46,6 +46,10 @@ app.listen(3001, () => {
 | GET | `/v1/providers` | List configured providers |
 | GET | `/healthz` | Health check |
 
+## Streaming
+
+`POST /v1/chat/completions` with `stream: true` is relayed as a `text/event-stream`: the gateway forwards the provider's frames verbatim, so a client receives tokens as they are produced. It adds `stream_options: { include_usage: true }` unless the caller set `stream_options` itself, and meters the usage the final frame carries. Streaming is available for OpenAI-compatible providers; a provider whose stream would need translation (`ollama`, `anthropic`) is rejected `501` rather than answered with a non-streaming body.
+
 ## Identity & tenancy
 
 Requests carry tenant identity via headers: `X-Database-Id` (**required** — requests without it are rejected `400`), `X-Entity-Id`, and `X-Actor-Id`. Routing to a specific provider can be forced with `X-LLM-Provider`.

agentic/agentic-server/__tests__/streaming.test.ts

@@ -0,0 +1,187 @@
+/**
+ * Streaming tests for agentic-server.
+ *
+ * A harness that asks for `stream: true` needs its tokens as the provider emits
+ * them, so the gateway relays an OpenAI-compatible event stream verbatim rather
+ * than parsing the body as JSON. These tests assert the relayed frames, the
+ * usage the stream carries reaching the injected sink, the frame-splitting the
+ * scanner does across arbitrary chunk boundaries, and the loud rejection of a
+ * provider type whose stream the gateway cannot translate.
+ */
+
+import express from 'express';
+import type { Server } from 'http';
+
+import type { InferenceEntry } from '../src';
+import { createAgenticServer } from '../src';
+import { UsageScanner, withUsageStreamOptions } from '../src/streaming';
+
+const CHUNKS = [
+  'data: {"id":"c1","choices":[{"delta":{"content":"Hel"}}]}\n\n',
+  'data: {"id":"c1","choices":[{"delta":{"content":"lo"}}]}\n\n',
+  'data: {"id":"c1","choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":11,"completion_tokens":5,"total_tokens":16}}\n\n',
+  'data: [DONE]\n\n'
+];
+
+let mockLlmServer: Server;
+let mockLlmPort: number;
+let agenticServer: Server;
+let agenticPort: number;
+let ollamaServer: Server;
+let ollamaPort: number;
+
+let llmRequests: Array<{ body: any }>;
+let sinkEntries: InferenceEntry[];
+
+beforeAll(async () => {
+  const llmApp = express();
+  llmApp.use(express.json());
+
+  llmApp.post('/v1/chat/completions', (req: any, res: any) => {
+    llmRequests.push({ body: req.body });
+    res.setHeader('Content-Type', 'text/event-stream');
+    for (const chunk of CHUNKS) res.write(chunk);
+    res.end();
+  });
+
+  await new Promise<void>((resolve) => {
+    mockLlmServer = llmApp.listen(0, () => {
+      const addr = mockLlmServer.address();
+      mockLlmPort = typeof addr === 'object' && addr ? addr.port : 0;
+      resolve();
+    });
+  });
+
+  const app = createAgenticServer({
+    providerType: 'openai',
+    providerBaseUrl: `http://localhost:${mockLlmPort}`,
+    providerApiKey: 'test-key',
+    defaultModel: 'gpt-4o-mini',
+    inferenceSink: { logInference: (entry) => sinkEntries.push(entry) }
+  });
+
+  await new Promise<void>((resolve) => {
+    agenticServer = app.listen(0, () => {
+      const addr = agenticServer.address();
+      agenticPort = typeof addr === 'object' && addr ? addr.port : 0;
+      resolve();
+    });
+  });
+
+  const ollamaApp = createAgenticServer({
+    providerType: 'ollama',
+    providerBaseUrl: `http://localhost:${mockLlmPort}`,
+    defaultModel: 'llama3'
+  });
+
+  await new Promise<void>((resolve) => {
+    ollamaServer = ollamaApp.listen(0, () => {
+      const addr = ollamaServer.address();
+      ollamaPort = typeof addr === 'object' && addr ? addr.port : 0;
+      resolve();
+    });
+  });
+});
+
+afterAll(async () => {
+  await new Promise<void>((r, e) => agenticServer.close((err) => (err ? e(err) : r())));
+  await new Promise<void>((r, e) => ollamaServer.close((err) => (err ? e(err) : r())));
+  await new Promise<void>((r, e) => mockLlmServer.close((err) => (err ? e(err) : r())));
+});
+
+beforeEach(() => {
+  llmRequests = [];
+  sinkEntries = [];
+});
+
+// ─── SSE relay ────────────────────────────────────────────────────────────
+
+describe('POST /v1/chat/completions with stream: true', () => {
+  it('relays the provider event stream and meters the usage it carried', async () => {
+    const res = await fetch(`http://localhost:${agenticPort}/v1/chat/completions`, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/json',
+        'X-Database-Id': 'db-stream-1',
+        'X-Entity-Id': 'entity-stream-1'
+      },
+      body: JSON.stringify({
+        model: 'gpt-4o-mini',
+        stream: true,
+        messages: [{ role: 'user', content: 'Hello' }]
+      })
+    });
+
+    expect(res.status).toBe(200);
+    expect(res.headers.get('content-type')).toContain('text/event-stream');
+    expect(await res.text()).toBe(CHUNKS.join(''));
+
+    expect(llmRequests).toHaveLength(1);
+    expect(llmRequests[0].body.stream).toBe(true);
+    expect(llmRequests[0].body.stream_options).toEqual({ include_usage: true });
+
+    await new Promise((r) => setTimeout(r, 100));
+    expect(sinkEntries).toHaveLength(1);
+    expect(sinkEntries[0]).toMatchObject({
+      databaseId: 'db-stream-1',
+      entityId: 'entity-stream-1',
+      service: 'chat',
+      status: 'ok',
+      inputTokens: 11,
+      outputTokens: 5,
+      totalTokens: 16
+    });
+  });
+
+  it('rejects (501) a provider type whose stream the gateway cannot translate', async () => {
+    const res = await fetch(`http://localhost:${ollamaPort}/v1/chat/completions`, {
+      method: 'POST',
+      headers: { 'Content-Type': 'application/json', 'X-Database-Id': 'db-stream-2' },
+      body: JSON.stringify({ model: 'llama3', stream: true, messages: [] })
+    });
+
+    expect(res.status).toBe(501);
+    expect(llmRequests).toHaveLength(0);
+  });
+});
+
+// ─── Usage scanning ───────────────────────────────────────────────────────
+
+describe('UsageScanner', () => {
+  it('finds usage split across arbitrary chunk boundaries', () => {
+    const scanner = new UsageScanner();
+    const stream = CHUNKS.join('');
+    for (let i = 0; i < stream.length; i += 7) scanner.push(stream.slice(i, i + 7));
+
+    expect(scanner.result()).toEqual({ prompt_tokens: 11, completion_tokens: 5, total_tokens: 16 });
+  });
+
+  it('reports no usage for a stream that carried none', () => {
+    const scanner = new UsageScanner();
+    scanner.push('data: {"choices":[{"delta":{"content":"hi"}}]}\n\ndata: [DONE]\n\n');
+
+    expect(scanner.result()).toBeUndefined();
+  });
+
+  it('ignores a frame it cannot parse and keeps scanning', () => {
+    const scanner = new UsageScanner();
+    scanner.push(': keep-alive\n\ndata: not-json\n\n');
+    scanner.push('data: {"usage":{"prompt_tokens":3,"completion_tokens":4}}\n\n');
+
+    expect(scanner.result()).toEqual({ prompt_tokens: 3, completion_tokens: 4, total_tokens: 7 });
+  });
+});
+
+describe('withUsageStreamOptions', () => {
+  it('asks the provider for usage when the caller expressed no preference', () => {
+    expect(withUsageStreamOptions({ model: 'gpt-4o-mini' })).toEqual({
+      model: 'gpt-4o-mini',
+      stream_options: { include_usage: true }
+    });
+  });
+
+  it('keeps the stream options the caller chose', () => {
+    const body = { model: 'gpt-4o-mini', stream_options: { include_usage: false } };
+    expect(withUsageStreamOptions(body)).toBe(body);
+  });
+});

agentic/agentic-server/src/router.ts

@@ -6,13 +6,15 @@
  *   - Model-based routing: "anthropic/claude-3.5-sonnet" → anthropic provider
  *   - Header-based routing: X-LLM-Provider: ollama → ollama provider
  *   - Fire-and-forget inference metering via an injected InferenceSink
+ *   - Streaming chat completions relayed as server-sent events
  *   - /v1/usage reporting endpoint for external usage submission
  */
 
 import { Logger } from '@pgpmjs/logger';
 import { Router } from 'express';
 
 import { buildProviderHeaders, resolveProvider, resolveUpstreamUrl } from './providers';
+import { relayEventStream, withUsageStreamOptions } from './streaming';
 import {
   transformChatRequest,
   transformChatResponse,
@@ -50,15 +52,41 @@ export const createRouter = (options: AgenticServerOptions): Router => {
       model: req.body?.model
     });
 
+    // Only the OpenAI-compatible wire carries stream frames the gateway can
+    // relay; translating Ollama's or Anthropic's stream formats is separate
+    // work, and answering a stream request with a whole JSON body would break
+    // the client silently.
+    const streaming = req.body?.stream === true;
+    if (streaming && provider.type !== 'openai') {
+      res.status(501).json({
+        error: {
+          message: `streaming is not supported for provider type '${provider.type}'`
+        }
+      });
+      return;
+    }
+
     try {
       const upstreamUrl = resolveUpstreamUrl(provider, '/v1/chat/completions');
-      const body = transformChatRequest(provider, req.body || {});
+      const transformed = transformChatRequest(provider, req.body || {});
+      const body = streaming ? withUsageStreamOptions(transformed) : transformed;
       const headers = buildProviderHeaders(provider);
 
+      const abort = new AbortController();
+      // A client that hangs up mid-stream stops the upstream read; the request
+      // stream's own 'close' fires as soon as the body is consumed, so the
+      // response is what reports the disconnect.
+      if (streaming) {
+        res.on('close', () => {
+          if (!res.writableEnded) abort.abort();
+        });
+      }
+
       const upstream = await fetch(upstreamUrl, {
         method: 'POST',
         headers,
-        body: JSON.stringify(body)
+        body: JSON.stringify(body),
+        ...(streaming ? { signal: abort.signal } : {})
       });
 
       const latencyMs = Number(process.hrtime.bigint() - startTime) / 1e6;
@@ -87,6 +115,37 @@ export const createRouter = (options: AgenticServerOptions): Router => {
         return;
       }
 
+      if (streaming) {
+        const streamUsage = await relayEventStream(upstream, res);
+        const streamLatencyMs = Number(process.hrtime.bigint() - startTime) / 1e6;
+
+        log.info('inference complete', {
+          databaseId,
+          provider: provider.type,
+          model: req.body?.model,
+          promptTokens: streamUsage?.prompt_tokens,
+          completionTokens: streamUsage?.completion_tokens,
+          streamed: true
+        });
+
+        if (sink) {
+          sink.logInference({
+            databaseId, entityId, actorId,
+            model: String(req.body?.model || body.model || ''),
+            provider: provider.type,
+            service: 'chat',
+            operation: 'chat/completions',
+            inputTokens: streamUsage?.prompt_tokens || 0,
+            outputTokens: streamUsage?.completion_tokens || 0,
+            totalTokens: streamUsage?.total_tokens || 0,
+            latencyMs: streamLatencyMs,
+            status: 'ok',
+            ...(streamUsage ? { rawUsage: streamUsage } : {})
+          });
+        }
+        return;
+      }
+
       const data = await upstream.json() as Record<string, unknown>;
       const { body: responseBody, usage } = transformChatResponse(data, provider);
 
@@ -133,6 +192,13 @@ export const createRouter = (options: AgenticServerOptions): Router => {
         });
       }
 
+      // A stream that failed mid-relay has already sent its status and frames;
+      // the client sees the truncated stream rather than a JSON error.
+      if (res.headersSent) {
+        res.end();
+        return;
+      }
+
       res.status(502).json({
         error: { message: 'Failed to reach LLM provider', details: err.message }
       });