diff --git a/autogpt/agent.py b/autogpt/agent.py
index 0f59e4b2..aee46576 100644
--- a/autogpt/agent.py
+++ b/autogpt/agent.py
@@ -1,5 +1,5 @@
 from agent_protocol import Agent, Step, Task
-
+import time 
 import autogpt.utils
 
 
@@ -11,12 +11,15 @@ class AutoGPT:
 
     async def task_handler(self, task: Task) -> None:
         print(f"task: {task.input}")
+        await Agent.db.create_step(task.task_id, task.input, is_last=True)
+        time.sleep(2)
         autogpt.utils.run(task.input)
-        await Agent.db.create_step(task.task_id, task.input)
+        
 
     async def step_handler(self, step: Step) -> Step:
         print(f"step: {step.input}")
-        step = Agent.db.get_step(step.step_id)
-        updated_step: Step = Agent.db.update_step(step.step_id, status="completed")
-        updated_step.output = step.input
+        agent_step = await Agent.db.get_step(step.task_id, step.step_id)
+        updated_step: Step = await Agent.db.update_step(agent_step.task_id, agent_step.step_id, status="completed")
+        updated_step.output = agent_step.input
+        print(f"Step completed: {updated_step}")
         return updated_step
diff --git a/autogpt/db.py b/autogpt/db.py
index 5aa2da9a..1597f264 100644
--- a/autogpt/db.py
+++ b/autogpt/db.py
@@ -9,7 +9,7 @@ import sqlite3
 from typing import Dict, List, Optional
 
 from agent_protocol import Artifact, Step, Task, TaskDB
-from agent_protocol.models import TaskInput
+from agent_protocol.models import TaskInput, Status
 
 
 class DataNotFoundError(Exception):
@@ -130,7 +130,15 @@ class AgentDB(TaskDB):
         cursor = self.conn.cursor()
         cursor.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,))
         if task := cursor.fetchone():
-            return Task(task_id=task[0], input=task[1], additional_input=task[2])
+            task = Task(task_id=task[0], input=task[1], additional_input=task[2])
+            cursor.execute("SELECT * FROM steps WHERE task_id=?", (task_id,))
+            steps = cursor.fetchall()
+            if steps:
+                for step in steps:
+                    status = Status.created if step[3] == "created" else Status.completed
+                    task.steps.append(Step(task_id=step[1], step_id=step[0], name=step[2], status=status, is_last=True if step[4] == 1 else False, additional_properties=step[5]))
+            print(f"Task: {task}")
+            return task
         else:
             raise DataNotFoundError("Task not found")
 
diff --git a/test.py b/test.py
new file mode 100644
index 00000000..7744e6b1
--- /dev/null
+++ b/test.py
@@ -0,0 +1,13 @@
+DATABASE_NAME="agent.db"
+import sqlite3
+# Read all data from database
+
+def read_all():
+    conn = sqlite3.connect(DATABASE_NAME)
+    cur = conn.cursor()
+    cur.execute("SELECT * FROM tasks")
+    rows = cur.fetchall()
+    conn.close()
+    return rows
+
+print(read_all())
\ No newline at end of file