diff --git a/agents-api/agents_api/models/execution/create_execution_transition.py b/agents-api/agents_api/models/execution/create_execution_transition.py
index 02b2c22f0..0ed85969b 100644
--- a/agents-api/agents_api/models/execution/create_execution_transition.py
+++ b/agents-api/agents_api/models/execution/create_execution_transition.py
@@ -17,7 +17,6 @@
     wrap_in_class,
 )
 
-
 valid_transitions = {
     # Start state
     "init": ["wait", "error", "step", "cancelled"],
diff --git a/agents-api/agents_api/routers/sessions/session.py b/agents-api/agents_api/routers/sessions/session.py
index a9276a79e..5784f40af 100644
--- a/agents-api/agents_api/routers/sessions/session.py
+++ b/agents-api/agents_api/routers/sessions/session.py
@@ -11,7 +11,12 @@
 from openai.types.chat.chat_completion import ChatCompletion
 from pydantic import UUID4
 
-from ...autogen.openapi_model import DocIds, InputChatMLMessage, Tool
+from ...autogen.openapi_model import (
+    CreateEntryRequest,
+    DocIds,
+    InputChatMLMessage,
+    Tool,
+)
 from ...clients.embed import embed
 from ...clients.temporal import run_summarization_task, run_truncation_task
 from ...clients.worker.types import ChatML
@@ -38,7 +43,6 @@
 )
 from ...models.entry.create_entries import create_entries
 from ...models.session.get_cached_response import get_cached_response
-from ...models.session.prepare_chat_context import prepare_chat_context
 from ...models.session.prepare_session_data import prepare_session_data
 from ...models.session.set_cached_response import set_cached_response
 from .exceptions import InputTooBigError
@@ -208,7 +212,9 @@ async def run(
         # TODO: implement locking at some point
 
         # Get session data
-        session_data = get_session_data(self.developer_id, self.session_id)
+        session_data = prepare_session_data(
+            developer_id=self.developer_id, session_id=self.session_id
+        )
         if session_data is None:
             raise SessionNotFoundError(self.developer_id, self.session_id)
 
@@ -502,15 +508,20 @@ async def backward(
         entries: list[Entry] = []
         for m in new_input:
             entries.append(
-                Entry(
-                    session_id=self.session_id,
+                CreateEntryRequest(
                     role=m.role,
                     content=m.content,
                     name=m.name,
                 )
             )
 
-        entries.append(new_entry)
+        entries.append(
+            CreateEntryRequest(
+                role=new_entry.role,
+                content=new_entry.content,
+                name=new_entry.name,
+            )
+        )
         bg_task = None
 
         if (
@@ -524,7 +535,7 @@ async def backward(
             else:
                 raise PromptTooBigError(total_tokens, final_settings.token_budget)
 
-        create_entries_query(
+        create_entries(
             developer_id=self.developer_id, session_id=self.session_id, data=entries
         )
 
diff --git a/agents-api/agents_api/routers/tasks/routers.py b/agents-api/agents_api/routers/tasks/routers.py
index 69df9da08..84583721d 100644
--- a/agents-api/agents_api/routers/tasks/routers.py
+++ b/agents-api/agents_api/routers/tasks/routers.py
@@ -24,10 +24,14 @@
 from agents_api.clients.temporal import run_task_execution_workflow
 from agents_api.common.protocol.tasks import ExecutionInput
 from agents_api.dependencies.developer_id import get_developer_id
-from agents_api.models.execution.create_execution import create_execution as create_execution_query
-from agents_api.models.execution.get_execution import get_execution as get_execution_query
+from agents_api.models.execution.create_execution import (
+    create_execution as create_execution_query,
+)
+from agents_api.models.execution.get_execution import (
+    get_execution as get_execution_query,
+)
 from agents_api.models.execution.get_execution_transition import (
-    get_execution_transition_query,
+    get_execution_transition as get_execution_transition_query,
 )
 from agents_api.models.execution.list_execution_transitions import (
     list_execution_transitions as list_execution_transitions_query,