diff --git a/src/oss/langchain/short-term-memory.mdx b/src/oss/langchain/short-term-memory.mdx
index 11b1853d9..f30c8723f 100644
--- a/src/oss/langchain/short-term-memory.mdx
+++ b/src/oss/langchain/short-term-memory.mdx
@@ -679,9 +679,9 @@ def update_user_info(
     runtime: ToolRuntime[CustomContext, CustomState],
 ) -> Command:
     """Look up and update user info."""
-    user_id = runtime.context.user_id  # [!code highlight]
+    user_id = runtime.context.user_id 
     name = "John Smith" if user_id == "user_123" else "Unknown user"
-    return Command(update={
+    return Command(update={  # [!code highlight]
         "user_name": name,
         # update the message history
         "messages": [
@@ -699,12 +699,11 @@ def greet(
     """Use this to greet the user once you found their info."""
     user_name = runtime.state["user_name"]
     return f"Hello {user_name}!"
-  # [!code highlight]
 agent = create_agent(
     model="openai:gpt-5-nano",
     tools=[update_user_info, greet],
-    state_schema=CustomState,
-    context_schema=CustomContext,  # [!code highlight]
+    state_schema=CustomState, # [!code highlight]
+    context_schema=CustomContext,  
 )
 
 agent.invoke(