From 598d4e16aaac3f98506c4d5c2c261c24733f42cd Mon Sep 17 00:00:00 2001
From: dal <dallin@buster.so>
Date: Tue, 8 Jul 2025 12:01:11 -0600
Subject: [PATCH] log braintrust

---
 .../analyst-agent-task/analyst-agent-task.ts  | 13 ++++++++-
 .../message-post-processing.ts                | 28 ++++++++++++++-----
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/apps/trigger/src/tasks/analyst-agent-task/analyst-agent-task.ts b/apps/trigger/src/tasks/analyst-agent-task/analyst-agent-task.ts
index 5e46232ec..870af214c 100644
--- a/apps/trigger/src/tasks/analyst-agent-task/analyst-agent-task.ts
+++ b/apps/trigger/src/tasks/analyst-agent-task/analyst-agent-task.ts
@@ -1,5 +1,5 @@
 import { logger, schemaTask, tasks } from '@trigger.dev/sdk';
-import { initLogger, wrapTraced } from 'braintrust';
+import { currentSpan, initLogger, wrapTraced } from 'braintrust';
 import { AnalystAgentTaskInputSchema, type AnalystAgentTaskOutput } from './types';
 
 // Task 2 & 4: Database helpers (IMPLEMENTED)
@@ -441,6 +441,17 @@ export const analystAgentTask: ReturnType<
       const workflowStartMethodStart = Date.now();
       const tracedWorkflow = wrapTraced(
         async () => {
+          currentSpan().log({
+            metadata: {
+              userName: braintrustMetadata.userName || 'Unknown',
+              userId: braintrustMetadata.userId,
+              organizationName: braintrustMetadata.organizationName || 'Unknown',
+              organizationId: braintrustMetadata.organizationId,
+              messageId: braintrustMetadata.messageId,
+              chatId: braintrustMetadata.chatId,
+            },
+          });
+
           return await run.start({
             inputData: workflowInput,
             runtimeContext,
diff --git a/apps/trigger/src/tasks/message-post-processing/message-post-processing.ts b/apps/trigger/src/tasks/message-post-processing/message-post-processing.ts
index 4cdba6573..7c63d0a80 100644
--- a/apps/trigger/src/tasks/message-post-processing/message-post-processing.ts
+++ b/apps/trigger/src/tasks/message-post-processing/message-post-processing.ts
@@ -1,9 +1,9 @@
 import postProcessingWorkflow, {
   type PostProcessingWorkflowOutput,
 } from '@buster/ai/workflows/post-processing-workflow';
-import { eq, getDb, messages } from '@buster/database';
+import { eq, getBraintrustMetadata, getDb, messages } from '@buster/database';
 import { logger, schemaTask } from '@trigger.dev/sdk/v3';
-import { initLogger, wrapTraced } from 'braintrust';
+import { currentSpan, initLogger, wrapTraced } from 'braintrust';
 import { z } from 'zod/v4';
 import {
   buildWorkflowInput,
@@ -155,17 +155,20 @@ export const messagePostProcessingTask: ReturnType<
       });
 
       // Step 2: Fetch all required data concurrently
-      const [conversationMessages, previousPostProcessingResults, datasets] = await Promise.all([
-        fetchConversationHistory(messageContext.chatId),
-        fetchPreviousPostProcessingMessages(messageContext.chatId, messageContext.createdAt),
-        fetchUserDatasets(messageContext.createdBy),
-      ]);
+      const [conversationMessages, previousPostProcessingResults, datasets, braintrustMetadata] =
+        await Promise.all([
+          fetchConversationHistory(messageContext.chatId),
+          fetchPreviousPostProcessingMessages(messageContext.chatId, messageContext.createdAt),
+          fetchUserDatasets(messageContext.createdBy),
+          getBraintrustMetadata({ messageId: payload.messageId }),
+        ]);
 
       logger.log('Fetched required data', {
         messageId: payload.messageId,
         conversationMessagesCount: conversationMessages.length,
         previousPostProcessingCount: previousPostProcessingResults.length,
         datasetsCount: datasets.length,
+        braintrustMetadata, // Log the metadata to verify it's working
       });
 
       // Step 3: Build workflow input
@@ -191,6 +194,17 @@ export const messagePostProcessingTask: ReturnType<
 
       const tracedWorkflow = wrapTraced(
         async () => {
+          currentSpan().log({
+            metadata: {
+              userName: braintrustMetadata.userName || 'Unknown',
+              userId: braintrustMetadata.userId,
+              organizationName: braintrustMetadata.organizationName || 'Unknown',
+              organizationId: braintrustMetadata.organizationId,
+              messageId: braintrustMetadata.messageId,
+              chatId: braintrustMetadata.chatId,
+            },
+          });
+
           const run = postProcessingWorkflow.createRun();
           return await run.start({
             inputData: workflowInput,