Fix formatting and highlight code sections (#1181)

## Overview
<!-- Brief description of what documentation is being added/updated -->

## Type of change
Improves code highlighting consistency in the documentation and ensures
that emphasis is placed on the short-term-memory-related logic rather
than unrelated context.
**Type:** Update existing documentation

## Checklist
<!-- Put an 'x' in all boxes that apply -->
- [x] I have read the [contributing guidelines](README.md)
- [ ] I have tested my changes locally using `docs dev`
- [ ] All code examples have been tested and work correctly
- [ ] I have used **root relative** paths for internal links
- [ ] I have updated navigation in `src/docs.json` if needed
- I have gotten approval from the relevant reviewers
- (Internal team members only / optional) I have created a preview
deployment using the [Create Preview Branch
workflow](https://github.com/langchain-ai/docs/actions/workflows/create-preview-branch.yml)

Author: donMichaelL

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="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(