diff --git a/app/src/components/game/TreeChart.vue b/app/src/components/game/TreeChart.vue
index 687cbd1..b07408b 100644
--- a/app/src/components/game/TreeChart.vue
+++ b/app/src/components/game/TreeChart.vue
@@ -12,7 +12,7 @@ const sessionStore = useSession()
         <FileNode
           v-if="sessionStore.session"
           :node="sessionStore.session.graph"
-          :pwd="sessionStore.session.pwd"
+          pwd="/home/commander"
         />
       </li>
     </ul>
diff --git a/app/src/store/session.ts b/app/src/store/session.ts
index 141a1bc..b70d66c 100644
--- a/app/src/store/session.ts
+++ b/app/src/store/session.ts
@@ -2,6 +2,7 @@ import { defineStore } from 'pinia'
 import {
   FileGraph,
   type FileGraphUpdateEvent,
+  type FileObject,
 } from '../../../packages/file-graph'
 import type { IResponse, ITask } from '../../../packages/game'
 import type { SessionDetail } from '../../../server/src/routers/sessionRouter'
@@ -75,9 +76,7 @@ const useSession = defineStore('session', {
       this.session.responses = update.responses
       this.session.tasks = update.tasks
       this.session.status = update.status
-      if (update.graphUpdate) {
-        this.session.graph.handleEvent(update.graphUpdate)
-      }
+      console.log('update', update)
       // if (
       //   response.status === 'finished' &&
       //   this.session.status !== 'finished'
@@ -103,9 +102,10 @@ const useSession = defineStore('session', {
       term.onData((data: string) => {
         socket.emit('terminal', data)
       })
-      // socket.on('graph', (event: { discover: FileObject[]; pwd: string }) => {
-      //   this.updateGraph(event)
-      // })
+      socket.on('graph', (event: FileGraphUpdateEvent) => {
+        if (!this.session) return
+        this.session.graph.handleEvent(event)
+      })
       socket.on('update', (update: SessionUpdate) => {
         this.newUpdate(update)
       })
diff --git a/app/src/views/GamePage.vue b/app/src/views/GamePage.vue
index 209130f..ad7cc2b 100644
--- a/app/src/views/GamePage.vue
+++ b/app/src/views/GamePage.vue
@@ -84,7 +84,7 @@ onMounted(async () => {
       <div class="bg-border rounded-xl h-[90%] p-2">
         <div class="bg-bg rounded-lg h-full p-2">
           <section id="visualization" class="bg-bg-primary h-full rounded-md">
-            <!-- <VisualizationPart /> -->
+            <VisualizationPart />
           </section>
         </div>
       </div>
diff --git a/server/src/api/controllers/commandController.ts b/server/src/api/controllers/commandController.ts
index 4909ad2..f225d9a 100644
--- a/server/src/api/controllers/commandController.ts
+++ b/server/src/api/controllers/commandController.ts
@@ -58,16 +58,18 @@ export const newCommand = asyncHandler(async (req: Request, res: Response) => {
 
     if (gameSession.isFinished()) {
       await finishSession(session)
-    } else {
-      await session.save()
     }
     pushToSession(session.id, 'update', {
       responses: gameSession.getResponses(),
       tasks: gameSession.getTasks(),
       status: session.status,
-      graphUpdate: command.params,
     })
   }
+  if (command.params) {
+    session.graph = gameSession.getGraph()
+    pushToSession(session.id, 'graph', command.params)
+  }
+  await session.save()
 
   res.status(200).end()
 })