Ocean/pipeline integration tests d552

server/src/index.ts

@@ -26,7 +26,7 @@
   callProviderStream,
   getModelName,
 } from './providers.js'
 import { callProviderStructured, StructuredOutputError } from './structuredOutput.js'
 
 import { checkUsage, recordUsage, getUsageStats } from './usage.js'
 import { dbHealthCheck, closePool } from './db.js'
@@ -219,6 +219,8 @@
   const availableProviders = providerOrder().filter(providerIsConfigured)
   let fullText = ''
   let usedProvider: ProviderName | null = null
+  let providerMeta: { provider: string; model: string } | null = null
+  let byokSucceeded = false
   const providerErrors: string[] = []
 
   if (byok?.apiKey && byok?.baseUrl) {
@@ -230,7 +232,17 @@
       } else {
         fullText = await chatWithOpenAiCompatible(byokParams, messages)
       }
-      usedProvider = 'openrouter' // label it generically
+      let byokHost = 'custom'
+      try {
+        byokHost = new URL(byok.baseUrl).hostname
+      } catch {
+        // keep custom
+      }
+      providerMeta = {
+        provider: byok.provider?.trim() || byokHost || 'byok',
+        model: byok.model?.trim() || byokHost,
+      }
+      byokSucceeded = true
     } catch (err) {
       const msg = err instanceof Error ? err.message : String(err)
       providerErrors.push(`byok(${byok.provider}): ${msg}`)
@@ -239,7 +251,7 @@
     }
   }
 
-  if (!usedProvider) {
+  if (!byokSucceeded) {
     for (const provider of availableProviders) {
       try {
         const providerOpts = { jsonMode: !llmOnly, maxTokens: llmOnly ? undefined : 2048 }
@@ -251,6 +263,7 @@
           fullText = response.text
         }
         usedProvider = provider
+        providerMeta = { provider, model: getModelName(provider) }
         break
       } catch (err) {
         const msg = err instanceof Error ? err.message : String(err)
@@ -265,7 +278,7 @@
     }
   }
 
-  if (!usedProvider) {
+  if (!byokSucceeded && !usedProvider) {
     if (providerErrors.length) {
       console.warn('[llm] all providers failed:', providerErrors.join(' | '))
     }
@@ -291,7 +304,7 @@
       message: text || 'I could not generate a response. Try rephrasing your question.',
       actions: [],
       source: 'llm',
-      meta: usedProvider ? { provider: usedProvider, model: getModelName(usedProvider) } : undefined,
+      meta: providerMeta ?? undefined,
     }
   }
 
@@ -301,6 +314,7 @@
   // If the LLM returned text that doesn't parse to valid actions, retry once
   // with a correction hint. This catches the common case where the model
   // returns prose instead of JSON, or malformed JSON.
+  // Only retry against server providers (not BYOK) to avoid wrong credentials.
   if (!stream && parsed.actions.length === 0 && fullText.trim().length > 0 && usedProvider) {
     const retryHint: Array<{ role: 'system' | 'user' | 'assistant'; content: string }> = [
       ...messages,
@@ -329,7 +343,7 @@
     message: parsed.message,
     actions: parsed.actions,
     source: 'llm',
-    meta: usedProvider ? { provider: usedProvider, model: getModelName(usedProvider) } : undefined,
+    meta: providerMeta ?? undefined,
   }
 }
 

src/ai/agentClient.test.ts

@@ -0,0 +1,43 @@
+import { describe, expect, it } from 'vitest'
+import { parseCompleteSseEvent, serverResponseToChatMessage } from './agentClient'
+
+describe('parseCompleteSseEvent', () => {
+  it('retains provider meta from complete events', () => {
+    const raw = JSON.stringify({
+      type: 'complete',
+      message: 'Sorted column B',
+      actions: [],
+      source: 'llm',
+      meta: { provider: 'groq', model: 'llama-3.3-70b-versatile' },
+    })
+    const parsed = parseCompleteSseEvent(raw)
+    expect(parsed).not.toBeNull()
+    expect(parsed!.meta).toEqual({
+      provider: 'groq',
+      model: 'llama-3.3-70b-versatile',
+    })
+  })
+
+  it('returns null for token events', () => {
+    expect(parseCompleteSseEvent(JSON.stringify({ type: 'token', content: 'hi' }))).toBeNull()
+  })
+
+  it('returns null for malformed JSON', () => {
+    expect(parseCompleteSseEvent('{not-json')).toBeNull()
+  })
+})
+
+describe('serverResponseToChatMessage', () => {
+  it('copies provider meta onto ChatMessage', () => {
+    const msg = serverResponseToChatMessage({
+      message: 'Done',
+      actions: [],
+      source: 'llm',
+      meta: { provider: 'groq', model: 'llama-3.3-70b-versatile' },
+    })
+    expect(msg.providerMeta).toEqual({
+      provider: 'groq',
+      model: 'llama-3.3-70b-versatile',
+    })
+  })
+})

src/ai/agentClient.ts

@@ -1,5 +1,5 @@
 import { v4 as uuid } from 'uuid'
-import type { AgentAction, ChatMessage } from '@/types'
+import type { AgentAction, ChatMessage, ProviderMeta } from '@/types'
 import type { SpreadsheetContextPayload } from '@/ai/buildContext'
 import { getAuthHeaders } from '@/lib/cloudSync'
 import { getByokPayload } from '@/lib/userApiKey'
@@ -18,6 +18,33 @@ export interface ServerChatResponse {
   source: 'llm' | 'fallback' | 'template'
   reasoning?: string
   suggestions?: string[]
+  meta?: ProviderMeta
+}
+
+/** Parse an SSE `data:` JSON payload into a ServerChatResponse when type=complete. */
+export function parseCompleteSseEvent(jsonStr: string): ServerChatResponse | null {
+  try {
+    const event = JSON.parse(jsonStr) as {
+      type?: string
+      message?: string
+      actions?: ServerAgentAction[]
+      source?: string
+      reasoning?: string
+      suggestions?: string[]
+      meta?: ProviderMeta
+    }
+    if (event.type !== 'complete' || typeof event.message !== 'string') return null
+    return {
+      message: event.message,
+      actions: Array.isArray(event.actions) ? event.actions : [],
+      source: (event.source as ServerChatResponse['source']) ?? 'llm',
+      reasoning: event.reasoning,
+      suggestions: event.suggestions,
+      meta: event.meta,
+    }
+  } catch {
+    return null
+  }
 }
 
 export interface ServerHealth {
@@ -106,21 +133,13 @@ export async function chatWithAgentServerStream(
         if (!jsonStr) continue
 
         try {
-          const event = JSON.parse(jsonStr) as
-            | { type: 'token'; content: string }
-            | { type: 'complete'; message: string; actions: ServerAgentAction[]; source: string; reasoning?: string; suggestions?: string[] }
-
-          if (event.type === 'token') {
-            onToken(event.content)
-          } else if (event.type === 'complete') {
-            finalResponse = {
-              message: event.message,
-              actions: event.actions,
-              source: event.source as ServerChatResponse['source'],
-              reasoning: event.reasoning,
-              suggestions: event.suggestions,
-            }
+          const parsed = JSON.parse(jsonStr) as { type?: string; content?: string }
+          if (parsed.type === 'token' && typeof parsed.content === 'string') {
+            onToken(parsed.content)
+            continue
           }
+          const complete = parseCompleteSseEvent(jsonStr)
+          if (complete) finalResponse = complete
         } catch {
           // Skip malformed events
         }
@@ -153,5 +172,6 @@ export function serverResponseToChatMessage(
     timestamp,
     suggestions: response.suggestions,
     actions: actions.length > 0 ? actions : undefined,
+    providerMeta: response.meta,
   }
 }

src/ai/brain.ts

@@ -692,6 +692,7 @@ export async function processMessage(input: ProcessMessageInput): Promise<ToolRe
       message: combined || 'I looked at your sheet but didn\'t find enough to go on. Try selecting a range or asking a more specific question.',
       toolUsed: deterministic?.toolUsed ?? (finalLlmText ? 'llm' : 'insights'),
       reasoning: serverResult.reasoning,
+      providerMeta: serverResult.meta,
       suggestions: contextualSuggestions.length > 0
         ? contextualSuggestions
         : (deterministic?.suggestions ?? serverResult.suggestions),