From ff069f34debeb19928b3b3bed72bb4f68e4762e9 Mon Sep 17 00:00:00 2001
From: Pavel401 <pavelalam401@gmail.com>
Date: Sat, 23 May 2026 20:00:31 +0530
Subject: [PATCH] feat: add Mastra traces for debugging

---
 backend/package.json        |  2 ++
 backend/src/index.ts        | 11 +++++++----
 backend/src/mastra/index.ts | 34 +++++++++++++++++++++++++++++-----
 docker-compose.dev.yml      |  4 ++++
 4 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/backend/package.json b/backend/package.json
index 6903fbd..74314ea 100644
--- a/backend/package.json
+++ b/backend/package.json
@@ -13,6 +13,8 @@
     "@clerk/backend": "^3.4.11",
     "@fastify/cors": "^11.0.0",
     "@mastra/core": "^1.36.0",
+    "@mastra/observability": "^1.13.0",
+    "@mastra/pg": "^1.11.1",
     "@openrouter/ai-sdk-provider": "^2.9.0",
     "ai": "^6.0.0",
     "convex": "^1.39.1",
diff --git a/backend/src/index.ts b/backend/src/index.ts
index 330ade1..224cd30 100644
--- a/backend/src/index.ts
+++ b/backend/src/index.ts
@@ -5,8 +5,8 @@ import { env } from "./env.js";
 import clerkAuthPlugin, { requireAuth } from "./clerk-auth.js";
 import { inferSchema } from "./pipeline/schema-inference.js";
 import { datasetContextSchema } from "./pipeline/populate.js";
-import { populateWorkflow } from "./mastra/workflows/populate.js";
-import { convex, api } from "./convex.js";
+import { mastra } from "./mastra/index.js";
+import { convex, api, internal } from "./convex.js";
 
 const fastify = Fastify({ logger: true });
 
@@ -61,7 +61,7 @@ await fastify.register(async (instance) => {
     }
 
     try {
-      const dataset = await convex.query(api.datasets.get, { id: parsed.data.datasetId });
+      const dataset = await convex.query(internal.datasets.getInternal, { id: parsed.data.datasetId });
       if (!dataset) {
         return reply.code(404).send({ error: "Dataset not found" });
       }
@@ -69,8 +69,10 @@ await fastify.register(async (instance) => {
         return reply.code(403).send({ error: "Not authorized to populate this dataset" });
       }
 
-      const run = await populateWorkflow.createRun();
+      const workflow = mastra.getWorkflow("populateWorkflow");
+      const run = await workflow.createRun();
       const result = await run.start({ inputData: parsed.data });
+      await mastra.observability.getDefaultInstance()?.flush();
 
       req.log.info({ workflowStatus: result.status, steps: JSON.stringify(result.steps).slice(0, 2000) }, "Populate workflow completed");
 
@@ -80,6 +82,7 @@ await fastify.register(async (instance) => {
 
       return { success: true, result: result.result };
     } catch (err) {
+      await mastra.observability.getDefaultInstance()?.flush();
       const msg = err instanceof Error ? err.message : String(err);
       if (msg.includes("validator") || msg.includes("Invalid")) {
         return reply.code(400).send({ error: "Invalid datasetId" });
diff --git a/backend/src/mastra/index.ts b/backend/src/mastra/index.ts
index 9a7cae7..aad5ef6 100644
--- a/backend/src/mastra/index.ts
+++ b/backend/src/mastra/index.ts
@@ -1,9 +1,33 @@
-import { Mastra } from "@mastra/core/mastra";
-import { inferSchemaWorkflow } from "./workflows/infer-schema.js";
-import { populateWorkflow } from "./workflows/populate.js";
-import { populateAgent } from "./agents/populate.js";
+import { Mastra } from '@mastra/core/mastra'
+import { PostgresStore } from '@mastra/pg'
+import {
+  Observability,
+  MastraStorageExporter,
+  SensitiveDataFilter,
+} from '@mastra/observability'
+import { inferSchemaWorkflow } from "./workflows/infer-schema.js"
+import { populateWorkflow } from "./workflows/populate.js"
+import { populateAgent } from "./agents/populate.js"
+
+const storage = new PostgresStore({
+  id: 'mastra-pg-storage',
+  connectionString: process.env.MASTRA_DATABASE_URL!,
+})
 
 export const mastra = new Mastra({
   agents: { populateAgent },
   workflows: { inferSchemaWorkflow, populateWorkflow },
-});
+  storage,
+  observability: new Observability({
+    
+    configs: {
+      default: {
+        serviceName: 'bigset',
+        exporters: [new MastraStorageExporter({ strategy: 'auto' })],
+        spanOutputProcessors: [
+          new SensitiveDataFilter(),
+        ],
+      },
+    },
+  }),
+})
diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml
index 7a0eec1..1ab86f9 100644
--- a/docker-compose.dev.yml
+++ b/docker-compose.dev.yml
@@ -24,6 +24,7 @@ services:
       - "3501:3501"
     volumes:
       - ./backend/src:/app/src
+      - ./backend/.mastra:/app/.mastra
     environment:
       CLIENT_ORIGIN: http://localhost:3500
       CONVEX_URL: http://convex:3210
@@ -33,6 +34,7 @@ services:
       CLERK_PUBLISHABLE_KEY: ${NEXT_PUBLIC_CLERK_PUBLISHABLE_KEY:-}
       OPENROUTER_API_KEY: ${OPENROUTER_API_KEY:-}
       TINYFISH_API_KEY: ${TINYFISH_API_KEY:-}
+      MASTRA_DATABASE_URL: postgresql://bigset:bigset@db:5432/bigset_internal
     depends_on:
       convex:
         condition: service_healthy
@@ -45,6 +47,7 @@ services:
       - "4111:4111"
     volumes:
       - ./backend/src:/app/src
+      - ./backend/.mastra:/app/.mastra
     environment:
       HOST: 0.0.0.0
       PORT: 4111
@@ -52,6 +55,7 @@ services:
       CONVEX_URL: http://convex:3210
       CONVEX_SELF_HOSTED_ADMIN_KEY: ${CONVEX_SELF_HOSTED_ADMIN_KEY:-}
       TINYFISH_API_KEY: ${TINYFISH_API_KEY:-}
+      MASTRA_DATABASE_URL: postgresql://bigset:bigset@db:5432/bigset_internal
     depends_on:
       convex:
         condition: service_healthy