adding better exporting

README.md

@@ -37,7 +37,7 @@ The global context limit defaults to 256k tokens and can be changed with `MAX_CO
 
 Inference requests are limited per client IP to 2 requests per second, 100 requests per five hours, and $10 of reported upstream cost per five hours. Railway's `X-Real-IP` header is used to identify clients.
 
-Set `DATABASE_URL` in `.env` (or the process environment) to a Postgres connection string. The proxy creates its `requests` table and index automatically on startup. Every incoming request and its complete response—including streaming responses and errors—is stored with headers, status, model, client IP, and timestamp. Successful inference requests also store a normalized training conversation containing the full chat history and generated assistant output, including reasoning, tool calls, and tool results.
+Set `DATABASE_URL` in `.env` (or the process environment) to a Postgres connection string. The proxy creates its `requests` table and index automatically on startup. Only requests to the OpenAI Chat Completions and Responses endpoints and the Anthropic Messages endpoints are stored. Their complete responses—including streaming responses and errors—are stored with headers, status, model, client IP, and timestamp. On startup, stored requests for all other endpoints are deleted. Successful chat completion requests also store a normalized training conversation containing the full chat history and generated assistant output, including reasoning, tool calls, and tool results.
 
 ```sh
 npm start
@@ -67,7 +67,9 @@ For Claude Code, either set `ANTHROPIC_BASE_URL=http://localhost:4005/ant` or us
 
 ## Exporting training data
 
-Exports use JSONL: each line is one inference request containing its complete input history followed by the generated assistant message. The normalized OpenAI-style `messages` preserve system/user/assistant roles, `reasoning_content`, assistant `tool_calls`, and `tool` results. Non-chat routes and failed inference requests are excluded.
+Each export creates two JSONL files. The regular `export-<timestamp>.jsonl` remains unchanged: each line is one inference request containing its complete input history followed by the generated assistant message. The additional `export-<timestamp>-traces.jsonl` removes cumulative intermediate snapshots and keeps each complete conversation branch. If a conversation is rewound and continued in multiple ways, every branch leaf is retained as its own full trace.
+
+Chat Completions, Responses, and Anthropic Messages requests are all exported in the same normalized OpenAI-style `messages` format. It preserves system/user/assistant roles, `reasoning_content`, assistant `tool_calls`, and `tool` results. Non-chat routes and failed inference requests are excluded.
 
 Each line also has `metadata` with the exact tool definitions supplied on that request, tool choice, model, API format, endpoint, streaming mode, generation parameters, caller-supplied request metadata, response ID/model, finish reason, usage, HTTP status, duration, request ID, and timestamp. Tool metadata remains in its original OpenAI, Responses, or Anthropic format so no provider-specific schema information is lost.
 
@@ -78,17 +80,17 @@ Each line also has `metadata` with the exact tool definitions supplied on that r
 Omit the limit to export every training request, newest first:
 
 ```sh
-npm run --silent export > requests.jsonl
+npm run --silent export
 ```
 
 Pass a positive limit to export that many of the most recent training requests:
 
 ```sh
-npm run --silent export -- 100 > recent-requests.jsonl
+npm run --silent export -- 100
 # Equivalent: npm run --silent export -- --limit 100
 ```
 
-The export command uses the same `DATABASE_URL` as the server. `--silent` suppresses npm's banner, and the script's progress message is written to standard error, so redirected JSONL remains valid. Rows collected before normalized training storage was added are normalized from their saved raw request and response during export.
+The export command uses the same `DATABASE_URL` as the server and saves both files in the project root. The limit applies to the regular per-request export; the traces file contains the complete branch leaves found among those requests. Rows collected before normalized training storage was added are normalized from their saved raw request and response during export.
 
 ## Tests
 

src/chat-completion-endpoints.js

@@ -0,0 +1,6 @@
+export const CHAT_COMPLETION_ENDPOINTS = Object.freeze([
+  '/oai/v1/chat/completions',
+  '/oai/v1/responses',
+  '/ant/v1/messages',
+  '/ant/v1/v1/messages',
+]);

src/db.js

@@ -1,4 +1,5 @@
 import pg from 'pg';
+import { CHAT_COMPLETION_ENDPOINTS } from './chat-completion-endpoints.js';
 
 const { Pool } = pg;
 
@@ -38,6 +39,14 @@ export async function ensureSchema(pool) {
   await pool.query('CREATE INDEX IF NOT EXISTS requests_created_at_idx ON requests (created_at DESC, id DESC)');
 }
 
+export async function removeNonChatCompletionRequests(pool) {
+  const result = await pool.query(`
+    DELETE FROM requests
+    WHERE split_part(endpoint, '?', 1) <> ALL($1::text[])
+  `, [CHAT_COMPLETION_ENDPOINTS]);
+  return result.rowCount ?? 0;
+}
+
 export async function insertRequest(pool, entry) {
   await pool.query(
     `INSERT INTO requests (

src/export.js

@@ -1,7 +1,27 @@
 import './env.js';
 import { once } from 'node:events';
+import { createWriteStream } from 'node:fs';
+import { join, dirname } from 'node:path';
+import { fileURLToPath } from 'node:url';
 import { createPool, ensureSchema } from './db.js';
 import { createTrainingExample } from './training-data.js';
+import { createFullTraceSelector } from './traces.js';
+import { CHAT_COMPLETION_ENDPOINTS } from './chat-completion-endpoints.js';
+
+const __dirname = dirname(fileURLToPath(import.meta.url));
+
+function timestamp() {
+  const d = new Date();
+  const pad = (n) => String(n).padStart(2, '0');
+  return (
+    `${d.getFullYear()}${pad(d.getMonth() + 1)}${pad(d.getDate())}` +
+    `-${pad(d.getHours())}${pad(d.getMinutes())}${pad(d.getSeconds())}`
+  );
+}
+
+const exportTimestamp = timestamp();
+const exportPath = join(__dirname, '..', `export-${exportTimestamp}.jsonl`);
+const tracesPath = join(__dirname, '..', `export-${exportTimestamp}-traces.jsonl`);
 
 function usage() {
   return 'Usage: npm run export -- [limit]\n       npm run export -- --limit <limit>';
@@ -28,26 +48,36 @@ try {
   const limit = parseLimit(process.argv.slice(2));
   pool = createPool();
   await ensureSchema(pool);
+  const stream = createWriteStream(exportPath, { flags: 'w' });
+  const tracesStream = createWriteStream(tracesPath, { flags: 'w' });
+  stream.on('error', (err) => {
+    console.error(`Error writing to ${exportPath}: ${err.message}`);
+    process.exitCode = 1;
+  });
+  tracesStream.on('error', (err) => {
+    console.error(`Error writing to ${tracesPath}: ${err.message}`);
+    process.exitCode = 1;
+  });
+  const isFullTrace = createFullTraceSelector();
   let exported = 0;
+  let tracesExported = 0;
   let cursor = null;
   while (limit === null || exported < limit) {
     const remaining = limit === null ? 1_000 : limit - exported;
     const batchSize = Math.min(1_000, Math.max(100, remaining * 2));
-    const params = cursor ? [cursor.cursor_created_at, cursor.id, batchSize] : [batchSize];
+    const params = cursor
+      ? [cursor.cursor_created_at, cursor.id, CHAT_COMPLETION_ENDPOINTS, batchSize]
+      : [CHAT_COMPLETION_ENDPOINTS, batchSize];
     const cursorClause = cursor ? 'AND (created_at, id) < ($1::timestamptz, $2::bigint)' : '';
-    const batchParameter = cursor ? '$3' : '$1';
+    const endpointParameter = cursor ? '$3' : '$1';
+    const batchParameter = cursor ? '$4' : '$2';
     const result = await pool.query(
       `SELECT id, created_at, created_at::text AS cursor_created_at, endpoint, request_body, response_body,
               response_status, duration_ms, stream, training_data
          FROM requests
         WHERE (method = 'POST' OR method IS NULL)
           AND response_status BETWEEN 200 AND 299
-          AND (
+          AND split_part(endpoint, '?', 1) = ANY(${endpointParameter}::text[])
-            endpoint LIKE '/oai/v1/chat/completions%'
-            OR endpoint LIKE '/oai/v1/responses%'
-            OR endpoint LIKE '/ant/v1/messages%'
-            OR endpoint LIKE '/ant/v1/v1/messages%'
-          )
           ${cursorClause}
         ORDER BY created_at DESC, id DESC
         LIMIT ${batchParameter}`,
@@ -74,14 +104,23 @@ try {
           ...(row.duration_ms == null ? {} : { duration_ms: row.duration_ms }),
         },
       };
-      if (!process.stdout.write(`${JSON.stringify(exportedExample)}\n`)) await once(process.stdout, 'drain');
+      if (!stream.write(`${JSON.stringify(exportedExample)}\n`)) await once(stream, 'drain');
+      if (isFullTrace(exportedExample)) {
+        if (!tracesStream.write(`${JSON.stringify(exportedExample)}\n`)) await once(tracesStream, 'drain');
+        tracesExported += 1;
+      }
       exported += 1;
       if (limit !== null && exported >= limit) break;
     }
     if (result.rows.length < batchSize || (limit !== null && exported >= limit)) break;
     cursor = result.rows.at(-1);
   }
-  console.error(`Exported ${exported} training request(s) as JSONL, newest first`);
+  const finished = Promise.all([once(stream, 'finish'), once(tracesStream, 'finish')]);
+  stream.end();
+  tracesStream.end();
+  await finished;
+  console.error(`Exported ${exported} training request(s) as JSONL to ${exportPath}, newest first`);
+  console.error(`Exported ${tracesExported} complete trace branch(es) as JSONL to ${tracesPath}, newest first`);
 } catch (error) {
   console.error(`Error: ${error.message}`);
   process.exitCode = 1;

src/proxy.js

@@ -15,6 +15,7 @@ import {
 } from './anthropic.js';
 import { translateResponsesRequest, translateResponsesResponse, translateResponsesChunk } from './responses.js';
 import { createTrainingExample } from './training-data.js';
+import { CHAT_COMPLETION_ENDPOINTS } from './chat-completion-endpoints.js';
 
 export const DEFAULT_UPSTREAM_BASE_URL = 'https://opencode.ai/zen/go/v1';
 
@@ -125,7 +126,7 @@ function withContextLimit(payload, maxContextTokens) {
 export function createProxyApp({ keys, fetchImpl = globalThis.fetch, upstreamBaseUrl = DEFAULT_UPSTREAM_BASE_URL, maxContextTokens = maxContextFromEnv(), rateLimiter = createRateLimiter(), pagesDirectory = defaultPagesDirectory, requestLogger } = {}) {
   const rotator = createKeyRotator(keys);
   const app = express();
-  if (requestLogger) app.use(collectRequestData(requestLogger));
+  if (requestLogger) app.post(CHAT_COMPLETION_ENDPOINTS, collectRequestData(requestLogger));
   app.use(express.json({ limit: '10mb' }));
 
   app.get('/', (_request, response) => {

src/server.js

@@ -1,12 +1,13 @@
 import path from 'node:path';
 import { createProxyApp } from './proxy.js';
 import { projectRoot } from './env.js';
-import { createPool, createRequestLogger, ensureSchema } from './db.js';
+import { createPool, createRequestLogger, ensureSchema, removeNonChatCompletionRequests } from './db.js';
 import { loadConfiguredKeys } from './key-store.js';
 
 const keys = await loadConfiguredKeys({ path: path.join(projectRoot, 'keys.txt') });
 const pool = createPool();
 await ensureSchema(pool);
+await removeNonChatCompletionRequests(pool);
 const app = createProxyApp({ keys, requestLogger: createRequestLogger(pool) });
 const port = Number(process.env.PORT || 4005);
 

src/traces.js

@@ -0,0 +1,54 @@
+function normalizedJson(value, parentKey = '') {
+  if (Array.isArray(value)) return value.map((item) => normalizedJson(item, parentKey));
+  if (!value || typeof value !== 'object') return value;
+
+  return Object.fromEntries(Object.keys(value).sort().flatMap((key) => {
+    const child = value[key];
+    if (child === undefined) return [];
+
+    // Tool arguments are JSON strings in normalized messages. Providers may
+    // reserialize them with different whitespace when they send the next turn.
+    if (parentKey === 'function' && key === 'arguments' && typeof child === 'string') {
+      try {
+        return [[key, normalizedJson(JSON.parse(child))]];
+      } catch {
+        // Non-JSON arguments still compare as their original string.
+      }
+    }
+    return [[key, normalizedJson(child, key)]];
+  }));
+}
+
+function messageKey(message) {
+  return JSON.stringify(normalizedJson(message));
+}
+
+/**
+ * Selects maximal conversations from examples ordered newest first.
+ *
+ * Each trie node represents a message prefix already seen in a newer request.
+ * A conversation whose final node already exists is an intermediate snapshot
+ * (or duplicate) of that newer conversation. Divergent children remain
+ * independent, so every leaf of a rewound/branched conversation is retained.
+ */
+export function createFullTraceSelector() {
+  const root = new Map();
+
+  return (example) => {
+    if (!Array.isArray(example?.messages) || example.messages.length === 0) return true;
+
+    let children = root;
+    let finalNodeExisted = true;
+    for (const message of example.messages) {
+      const key = messageKey(message);
+      let node = children.get(key);
+      if (!node) {
+        node = { children: new Map() };
+        children.set(key, node);
+        finalNodeExisted = false;
+      }
+      children = node.children;
+    }
+    return !finalNodeExisted;
+  };
+}

test/db.test.js

@@ -1,5 +1,5 @@
 import { describe, expect, it, vi } from 'vitest';
-import { createRequestLogger, ensureSchema, insertRequest } from '../src/db.js';
+import { createRequestLogger, ensureSchema, insertRequest, removeNonChatCompletionRequests } from '../src/db.js';
 
 describe('request database', () => {
   it('creates the requests table, compatible columns, and recent-request index', async () => {
@@ -37,6 +37,19 @@ describe('request database', () => {
     ]);
   });
 
+  it('removes requests that are not for the chat completions endpoint', async () => {
+    const pool = { query: vi.fn().mockResolvedValue({ rowCount: 4 }) };
+
+    await expect(removeNonChatCompletionRequests(pool)).resolves.toBe(4);
+
+    expect(pool.query).toHaveBeenCalledOnce();
+    expect(pool.query.mock.calls[0][0]).toContain("split_part(endpoint, '?', 1) <> ALL($1::text[])");
+    expect(pool.query.mock.calls[0][1]).toEqual([[
+      '/oai/v1/chat/completions', '/oai/v1/responses',
+      '/ant/v1/messages', '/ant/v1/v1/messages',
+    ]]);
+  });
+
   it('reports storage failures without rejecting request handling', async () => {
     const onError = vi.fn();
     const logger = createRequestLogger({ query: vi.fn().mockRejectedValue(new Error('database offline')) }, { onError });

test/proxy.test.js

@@ -126,17 +126,37 @@ describe('proxy', () => {
     });
   });
 
-  it('collects error and not-found responses too', async () => {
+  it('does not collect requests to non-chat-completion endpoints', async () => {
+    const fetchImpl = vi.fn().mockResolvedValue(response('{"object":"list","data":[]}'));
     const requestLogger = { log: vi.fn() };
-    const app = createProxyApp({ keys: ['key'], requestLogger });
+    const app = createProxyApp({ keys: ['key'], fetchImpl, requestLogger });
 
     await request(app).get('/missing').expect(404);
+    await request(app).get('/oai/v1/models').expect(200);
 
-    expect(requestLogger.log).toHaveBeenCalledOnce();
+    expect(requestLogger.log).not.toHaveBeenCalled();
-    expect(requestLogger.log.mock.calls[0][0]).toMatchObject({
+  });
-      method: 'GET', endpoint: '/missing', requestBody: null, responseStatus: 404,
+
-    });
+  it('collects OpenAI Responses and Anthropic Messages requests', async () => {
-    expect(requestLogger.log.mock.calls[0][0].responseBody).toContain('Cannot GET /missing');
+    const cases = [
+      { endpoint: '/oai/v1/responses', body: { model: 'm', input: 'Hi' } },
+      { endpoint: '/ant/v1/messages', body: { model: 'm', max_tokens: 10, messages: [{ role: 'user', content: 'Hi' }] } },
+      { endpoint: '/ant/v1/v1/messages', body: { model: 'm', max_tokens: 10, messages: [{ role: 'user', content: 'Hi' }] } },
+    ];
+
+    for (const { endpoint, body } of cases) {
+      const fetchImpl = vi.fn().mockResolvedValue(response(JSON.stringify({
+        id: 'chatcmpl-response', model: 'm',
+        choices: [{ message: { role: 'assistant', content: 'Hello' }, finish_reason: 'stop' }],
+      })));
+      const requestLogger = { log: vi.fn() };
+      const app = createProxyApp({ keys: ['key'], fetchImpl, requestLogger });
+
+      await request(app).post(endpoint).send(body).expect(200);
+
+      expect(requestLogger.log).toHaveBeenCalledOnce();
+      expect(requestLogger.log.mock.calls[0][0]).toMatchObject({ method: 'POST', endpoint, requestBody: body });
+    }
   });
 
   it('translates Responses requests and returns a Responses object', async () => {

test/traces.test.js

@@ -0,0 +1,55 @@
+import { describe, expect, it } from 'vitest';
+import { createFullTraceSelector } from '../src/traces.js';
+
+const message = (role, content, extra = {}) => ({ role, content, ...extra });
+const example = (...messages) => ({ messages });
+
+describe('full trace selection', () => {
+  it('keeps only the maximal snapshot of a growing conversation', () => {
+    const user = message('user', 'Start');
+    const first = message('assistant', 'First');
+    const next = message('user', 'Continue');
+    const last = message('assistant', 'Done');
+    const examples = [
+      example(user, first, next, last),
+      example(user, first),
+    ];
+
+    expect(examples.filter(createFullTraceSelector())).toEqual([examples[0]]);
+  });
+
+  it('keeps every leaf when a conversation is rewound and branched', () => {
+    const shared = [message('user', 'Start'), message('assistant', 'Shared')];
+    const branchC = example(...shared, message('user', 'Choose C'), message('assistant', 'C'));
+    const branchD = example(...shared, message('user', 'Choose D'), message('assistant', 'D'));
+    const intermediate = example(...shared);
+
+    const examples = [branchD, branchC, intermediate];
+    expect(examples.filter(createFullTraceSelector())).toEqual([branchD, branchC]);
+  });
+
+  it('matches JSON tool arguments despite provider whitespace changes', () => {
+    const requestToolCall = message('assistant', null, {
+      tool_calls: [{ id: 'call_1', type: 'function', function: { name: 'run', arguments: '{ "path": "a" }' } }],
+    });
+    const echoedToolCall = message('assistant', null, {
+      tool_calls: [{ type: 'function', function: { arguments: '{"path":"a"}', name: 'run' }, id: 'call_1' }],
+    });
+    const newest = example(
+      message('user', 'Run it'), echoedToolCall,
+      message('tool', 'ok', { tool_call_id: 'call_1' }), message('assistant', 'Done'),
+    );
+    const intermediate = example(message('user', 'Run it'), requestToolCall);
+
+    const examples = [newest, intermediate];
+    expect(examples.filter(createFullTraceSelector())).toEqual([newest]);
+  });
+
+  it('does not combine unrelated conversations', () => {
+    const examples = [
+      example(message('user', 'One'), message('assistant', '1')),
+      example(message('user', 'Two'), message('assistant', '2')),
+    ];
+    expect(examples.filter(createFullTraceSelector())).toEqual(examples);
+  });
+});