From 496975ed947afed3b6f8532bd5c16e8a83d5a71f Mon Sep 17 00:00:00 2001
From: Rene Fichtmueller <rf@flexoptix.net>
Date: Fri, 17 Jul 2026 17:22:40 +0200
Subject: [PATCH] fix(gateway): keep reasoning trace split compatible

---
 .../gateway/src/modules/reasoning-trace.ts    | 62 ++++++++++++++++---
 1 file changed, 55 insertions(+), 7 deletions(-)

diff --git a/packages/gateway/src/modules/reasoning-trace.ts b/packages/gateway/src/modules/reasoning-trace.ts
index 2a4e06a..79207f9 100644
--- a/packages/gateway/src/modules/reasoning-trace.ts
+++ b/packages/gateway/src/modules/reasoning-trace.ts
@@ -1,18 +1,66 @@
 export interface ReasoningTraceSplit {
   visible: string;
-  trace: string | null;
+  finalAnswer: string;
+  trace: {
+    marker: 'thinking_tag' | 'think_tag' | 'markdown_header' | 'json_field';
+    content: string;
+    estimatedTokens: number;
+  } | null;
 }
 
-export function splitReasoningTrace(output: string): ReasoningTraceSplit {
-  const match = output.match(/<think>([\s\S]*?)<\/think>/i);
-  if (!match) return { visible: output, trace: null };
+function estimateTokens(text: string): number {
+  return Math.max(1, Math.ceil(text.length / 4));
+}
+
+function splitWithTrace(
+  output: string,
+  marker: NonNullable<ReasoningTraceSplit['trace']>['marker'],
+  traceContent: string,
+  finalAnswer: string
+): ReasoningTraceSplit {
+  const visible = finalAnswer.trim();
   return {
-    visible: output.replace(match[0], '').trim(),
-    trace: match[1]?.trim() ?? null,
+    visible,
+    finalAnswer: visible,
+    trace: {
+      marker,
+      content: traceContent.trim(),
+      estimatedTokens: estimateTokens(traceContent),
+    },
   };
 }
 
-export async function storeReasoningTrace(_callId: string, _trace: string | null): Promise<void> {
+export function splitReasoningTrace(output: string): ReasoningTraceSplit {
+  const trimmed = output.trim();
+
+  const thinkingMatch = trimmed.match(/<thinking>([\s\S]*?)<\/thinking>/i);
+  if (thinkingMatch) {
+    return splitWithTrace(trimmed, 'thinking_tag', thinkingMatch[1] ?? '', trimmed.replace(thinkingMatch[0], ''));
+  }
+
+  const thinkMatch = trimmed.match(/<think>([\s\S]*?)<\/think>/i);
+  if (thinkMatch) {
+    return splitWithTrace(trimmed, 'think_tag', thinkMatch[1] ?? '', trimmed.replace(thinkMatch[0], ''));
+  }
+
+  const markdownMatch = trimmed.match(/\*\*Reasoning:\*\*([\s\S]*?)\*\*Answer:\*\*([\s\S]*)/i);
+  if (markdownMatch) {
+    return splitWithTrace(trimmed, 'markdown_header', markdownMatch[1] ?? '', markdownMatch[2] ?? '');
+  }
+
+  try {
+    const parsed = JSON.parse(trimmed) as { reasoning?: unknown; answer?: unknown };
+    if (typeof parsed.reasoning === 'string' && typeof parsed.answer === 'string') {
+      return splitWithTrace(trimmed, 'json_field', parsed.reasoning, parsed.answer);
+    }
+  } catch {
+    // Plain model output, not a reasoning envelope.
+  }
+
+  return { visible: trimmed, finalAnswer: trimmed, trace: null };
+}
+
+export async function storeReasoningTrace(_callId: string, _trace: ReasoningTraceSplit['trace']): Promise<void> {
   // Reasoning trace persistence is optional. Keep the public route usable when
   // the storage table has not been deployed yet.
 }