Remove separate collection for aggregate results

apps/rowboat/app/actions/simulation_actions.ts

@@ -1,7 +1,7 @@
 'use server';
 
 import { ObjectId } from "mongodb";
-import { scenariosCollection, simulationRunsCollection, simulationResultsCollection, simulationAggregateResultsCollection } from "../lib/mongodb";
+import { scenariosCollection, simulationRunsCollection, simulationResultsCollection } from "../lib/mongodb";
 import { z } from 'zod';
 import { projectAuthCheck } from "./project_actions";
 import { type WithStringId } from "../lib/types/types";
@@ -220,13 +220,14 @@ export async function createAggregateResult(
 ): Promise<void> {
     await projectAuthCheck(projectId);
 
-    await simulationAggregateResultsCollection.insertOne({
-        projectId,
-        runId,
-        total,
-        pass,
-        fail,
-    });
+    await simulationRunsCollection.updateOne(
+        { _id: new ObjectId(runId), projectId },
+        {
+            $set: {
+                aggregateResults: { total, pass, fail }
+            }
+        }
+    );
 }
 
 export async function getAggregateResult(
@@ -235,19 +236,12 @@ export async function getAggregateResult(
 ): Promise<z.infer<typeof SimulationAggregateResult> | null> {
     await projectAuthCheck(projectId);
 
-    const result = await simulationAggregateResultsCollection.findOne({
+    const run = await simulationRunsCollection.findOne({
+        _id: new ObjectId(runId),
         projectId,
-        runId,
     });
 
-    if (!result) return null;
+    if (!run || !run.aggregateResults) return null;
 
-    // Only include the fields defined in the schema
-    return {
-        projectId: result.projectId,
-        runId: result.runId,
-        total: result.total,
-        pass: result.pass,
-        fail: result.fail
-    };
+    return run.aggregateResults;
 } 

apps/rowboat/app/lib/mongodb.ts

@@ -23,5 +23,4 @@ export const agentWorkflowsCollection = db.collection<z.infer<typeof Workflow>>(
 export const scenariosCollection = db.collection<z.infer<typeof Scenario>>("scenarios");
 export const apiKeysCollection = db.collection<z.infer<typeof ApiKey>>("api_keys");
 export const simulationRunsCollection = db.collection<z.infer<typeof SimulationRun>>("simulation_runs");
-export const simulationResultsCollection = db.collection<z.infer<typeof SimulationResult>>("simulation_results");
-export const simulationAggregateResultsCollection = db.collection<z.infer<typeof SimulationAggregateResult>>("simulation_aggregate_results");
+export const simulationResultsCollection = db.collection<z.infer<typeof SimulationResult>>("simulation_results");

apps/rowboat/app/lib/types/testing_types.ts

@@ -1,4 +1,5 @@
 import { z } from "zod";
+
 export const Scenario = z.object({
     projectId: z.string(),
     name: z.string().min(1, "Name cannot be empty"),
@@ -7,32 +8,38 @@ export const Scenario = z.object({
     context: z.string().default(''),
     createdAt: z.string().datetime(),
     lastUpdatedAt: z.string().datetime(),
-});export const SimulationArticleData = z.object({
+});
+
+export const SimulationArticleData = z.object({
     articleUrl: z.string(),
     articleTitle: z.string().default('').optional(),
     articleContent: z.string().default('').optional(),
 });
+
 export const SimulationScenarioData = z.object({
     scenario: z.string(),
     context: z.string().default(''),
 });
+
 export const SimulationChatMessagesData = z.object({
     chatMessages: z.string(),
 });
+
 export const SimulationData = z.union([SimulationArticleData, SimulationScenarioData, SimulationChatMessagesData]);
 
+export const SimulationAggregateResult = z.object({
+    total: z.number(),
+    pass: z.number(),
+    fail: z.number(),
+});
+
 export const SimulationRun = z.object({
     projectId: z.string(),
-    status: z.union([
-        z.literal('pending'),
-        z.literal('running'),
-        z.literal('completed'),
-        z.literal('cancelled'),
-        z.literal('failed')
-    ]),
+    status: z.enum(['pending', 'running', 'completed', 'cancelled', 'failed']),
     scenarioIds: z.array(z.string()),
-    startedAt: z.string().datetime(),
-    completedAt: z.string().datetime().optional(),
+    startedAt: z.string(),
+    completedAt: z.string().optional(),
+    aggregateResults: SimulationAggregateResult.optional(),
 });
 
 export const SimulationResult = z.object({
@@ -41,12 +48,4 @@ export const SimulationResult = z.object({
     scenarioId: z.string(),
     result: z.union([z.literal('pass'), z.literal('fail')]),
     details: z.string()
-});
-
-export const SimulationAggregateResult = z.object({
-    projectId: z.string(),
-    runId: z.string(),
-    total: z.number(),
-    pass: z.number(),
-    fail: z.number(),
 });

apps/rowboat/app/projects/[projectId]/simulation/app.tsx

@@ -227,12 +227,11 @@ export default function SimulationApp() {
           )
         );
 
-        // Calculate aggregate results
+        // Calculate and store aggregate results before marking as complete
         const total = scenarios.length;
         const pass = mockResults.filter(r => r.result === 'pass').length;
         const fail = mockResults.filter(r => r.result === 'fail').length;
 
-        // Store aggregate results
         await createAggregateResult(
           projectId as string,
           newRun._id,

apps/rowboat/app/projects/[projectId]/simulation/components/RunComponents.tsx

@@ -20,15 +20,11 @@ interface SimulationResultCardProps {
 export const SimulationResultCard = ({ run, results, scenarios }: SimulationResultCardProps) => {
   const [isExpanded, setIsExpanded] = useState(false);
   const [expandedScenarios, setExpandedScenarios] = useState<Set<string>>(new Set());
-  const [aggregateResult, setAggregateResult] = useState<z.infer<typeof SimulationAggregateResult> | null>(null);
 
-  useEffect(() => {
-    if (run.projectId && run._id) {
-      getAggregateResult(run.projectId, run._id)
-        .then(setAggregateResult)
-        .catch(console.error);
-    }
-  }, [run.projectId, run._id]);
+  // Replace the manual calculations with aggregate results from the run object
+  const totalScenarios = run.aggregateResults?.total ?? run.scenarioIds.length;
+  const passedScenarios = run.aggregateResults?.pass ?? 0;
+  const failedScenarios = run.aggregateResults?.fail ?? 0;
 
   const statusLabelClass = "px-3 py-1 rounded text-xs min-w-[60px] text-center uppercase font-semibold";
 
@@ -56,11 +52,6 @@ export const SimulationResultCard = ({ run, results, scenarios }: SimulationResu
     });
   };
 
-  // Replace the manual calculations with aggregate results
-  const totalScenarios = aggregateResult?.total ?? run.scenarioIds.length;
-  const passedScenarios = aggregateResult?.pass ?? 0;
-  const failedScenarios = aggregateResult?.fail ?? 0;
-
   const getDuration = () => {
     if (!run.completedAt) return 'In Progress';
     const start = new Date(run.startedAt);