diff --git a/src/app/(app)/problems/[id]/layout.tsx b/src/app/(app)/problems/[id]/layout.tsx
index 196c814..e7fbc5f 100644
--- a/src/app/(app)/problems/[id]/layout.tsx
+++ b/src/app/(app)/problems/[id]/layout.tsx
@@ -14,36 +14,31 @@ export default async function ProblemLayout({
 }: ProblemProps) {
   const { id } = await params;
 
-  const [problemData, editorLanguageConfigs, languageServerConfigs] =
-    await Promise.all([
-      prisma.problem.findUnique({
-        where: { id },
-        include: {
-          templates: true,
-          testcases: {
-            include: {
-              data: true,
-            },
+  const [problem, editorLanguageConfigs, languageServerConfigs] = await Promise.all([
+    prisma.problem.findUnique({
+      where: { id },
+      include: {
+        templates: true,
+        testcases: {
+          include: {
+            data: true,
           },
         },
-      }),
-      prisma.editorLanguageConfig.findMany(),
-      prisma.languageServerConfig.findMany(),
-    ]);
+      },
+    }),
+    prisma.editorLanguageConfig.findMany(),
+    prisma.languageServerConfig.findMany(),
+  ]);
 
-  if (!problemData) {
+  if (!problem) {
     return notFound();
   }
 
-  const { templates, testcases, ...problemWithoutTemplates } = problemData;
-
   return (
     <div className="flex flex-col h-screen">
       <ProblemStoreProvider
         problemId={id}
-        problem={problemWithoutTemplates}
-        templates={templates}
-        testcases={testcases}
+        problem={problem}
         editorLanguageConfigs={editorLanguageConfigs}
         languageServerConfigs={languageServerConfigs}
       >
diff --git a/src/hooks/use-problem.ts b/src/hooks/use-problem.ts
index ccd2cdc..4c97571 100644
--- a/src/hooks/use-problem.ts
+++ b/src/hooks/use-problem.ts
@@ -52,15 +52,14 @@ export const useProblem = () => {
     setCurrentValue,
     problemId,
     problem,
-    templates,
     editorLanguageConfigs,
     languageServerConfigs,
   } = useProblemStore((state) => state);
 
   // Get the default template for the current language from the templates list
   const currentTemplate = useMemo(
-    () => templates.find((t) => t.language === currentLang)?.template || "",
-    [templates, currentLang]
+    () => problem.templates.find((t) => t.language === currentLang)?.template || "",
+    [problem.templates, currentLang]
   );
 
   const currentEditorLanguageConfig = useMemo(
@@ -124,7 +123,6 @@ export const useProblem = () => {
     currentValue,
     problemId,
     problem,
-    templates,
     editorLanguageConfigs,
     languageServerConfigs,
     currentTemplate,
diff --git a/src/providers/problem-store-provider.tsx b/src/providers/problem-store-provider.tsx
index b0f2ce4..222da37 100644
--- a/src/providers/problem-store-provider.tsx
+++ b/src/providers/problem-store-provider.tsx
@@ -2,28 +2,22 @@
 
 import {
   EditorLanguage,
-  type Problem,
   type EditorLanguageConfig,
   type LanguageServerConfig,
-  type Template,
 } from "@/generated/client";
 import { useStore } from "zustand";
-import { TestcaseWithData } from "@/types/prisma";
+import { type ProblemWithDetails } from "@/types/prisma";
 import { type ReactNode, createContext, useRef, useContext } from "react";
 import { type ProblemStore, createProblemStore } from "@/stores/problem-store";
 
 export type ProblemStoreApi = ReturnType<typeof createProblemStore>;
 
-export const ProblemStoreContext = createContext<ProblemStoreApi | undefined>(
-  undefined
-);
+export const ProblemStoreContext = createContext<ProblemStoreApi | undefined>(undefined);
 
 export interface ProblemStoreProviderProps {
   children: ReactNode;
   problemId: string;
-  problem: Problem;
-  templates: Template[];
-  testcases: TestcaseWithData;
+  problem: ProblemWithDetails;
   editorLanguageConfigs: EditorLanguageConfig[];
   languageServerConfigs: LanguageServerConfig[];
 }
@@ -32,8 +26,6 @@ export const ProblemStoreProvider = ({
   children,
   problemId,
   problem,
-  templates,
-  testcases,
   editorLanguageConfigs,
   languageServerConfigs,
 }: ProblemStoreProviderProps) => {
@@ -50,8 +42,6 @@ export const ProblemStoreProvider = ({
       currentValue: "",
       problemId,
       problem,
-      templates,
-      testcases,
       editorLanguageConfigs,
       languageServerConfigs,
     });
diff --git a/src/stores/problem-store.ts b/src/stores/problem-store.ts
index a8f13c6..ca56545 100644
--- a/src/stores/problem-store.ts
+++ b/src/stores/problem-store.ts
@@ -2,12 +2,10 @@ import type {
   EditorLanguage,
   EditorLanguageConfig,
   LanguageServerConfig,
-  Problem,
-  Template,
 } from "@/generated/client";
 import type { editor } from "monaco-editor";
 import { createStore } from "zustand/vanilla";
-import { TestcaseWithData } from "@/types/prisma";
+import type { ProblemWithDetails } from "@/types/prisma";
 import { createJSONStorage, persist } from "zustand/middleware";
 
 export type ProblemState = {
@@ -19,9 +17,7 @@ export type ProblemState = {
   currentLang: EditorLanguage;
   currentValue: string;
   problemId: string;
-  problem: Problem;
-  templates: Template[];
-  testcases: TestcaseWithData;
+  problem: ProblemWithDetails;
   editorLanguageConfigs: EditorLanguageConfig[];
   languageServerConfigs: LanguageServerConfig[];
 };