diff --git a/examples/agent/healthcare-assistant/agent-with-instrumentation.py b/examples/agent/healthcare-assistant/agent-with-instrumentation.py
index fe7e2aa3..9d312003 100644
--- a/examples/agent/healthcare-assistant/agent-with-instrumentation.py
+++ b/examples/agent/healthcare-assistant/agent-with-instrumentation.py
@@ -128,7 +128,8 @@ async def invoke_chatbot(state):
     async def _process_query_async(self, messages: List[Dict[str, str]]) -> str:
         if not self.tools:
             self.load_tools()
-        self.graph = self._build_graph()
+        if self.graph is None:
+            self.graph = self._build_graph()
 
         langchain_messages: List[BaseMessage] = []
         for msg in messages:
@@ -144,7 +145,7 @@ async def _process_query_async(self, messages: List[Dict[str, str]]) -> str:
             splunk_ao_context.set_session(self.session_id)
 
             # One callback per request keeps each user turn in its own trace.
-            callback = SplunkAOAsyncCallback()
+            callback = SplunkAOAsyncCallback(flush_on_chain_end=True)
             run_config = {**self.langgraph_config, "callbacks": [callback]}
 
             result = await self.graph.ainvoke(
diff --git a/examples/agent/healthcare-assistant/app.py b/examples/agent/healthcare-assistant/app.py
index a0af9b44..3783fb8f 100644
--- a/examples/agent/healthcare-assistant/app.py
+++ b/examples/agent/healthcare-assistant/app.py
@@ -27,6 +27,7 @@ def _load_instrumented_agent():
 )
 from rag import get_rag_system
 from setup_env import setup_environment
+from splunk_ao import splunk_ao_context
 
 _APP_DIR = os.path.dirname(os.path.abspath(__file__))
 load_dotenv(os.path.join(_APP_DIR, ".env"))
@@ -104,7 +105,9 @@ def process_input(user_input: str | None):
                 elif isinstance(message, AIMessage):
                     conversation_messages.append({"role": "assistant", "content": message.content})
 
-        response = st.session_state.agent.process_query(conversation_messages)
+        # Pass only the latest user message to keep each trace clean (single input/output pair).
+        latest_user = [m for m in conversation_messages if m["role"] == "user"][-1:]
+        response = st.session_state.agent.process_query(latest_user)
         st.session_state.messages.append(
             {"message": AIMessage(content=response), "agent": "assistant"}
         )
@@ -147,14 +150,9 @@ def render_sidebar(app_config: dict) -> str:
             )
             if st.button("Log Hallucination", key="log_hallucination"):
                 with st.spinner("Logging hallucination to Splunk Agent Observability..."):
-                    existing_logger = (
-                        st.session_state.get("splunk_ao_logger")
-                        if st.session_state.get("splunk_ao_session_started", False)
-                        else None
-                    )
                     success = log_demo_hallucination(
                         config=app_config,
-                        existing_logger=existing_logger,
+                        existing_logger=splunk_ao_context,
                         session_id=st.session_state.get("session_id"),
                     )
                     if success:
diff --git a/examples/agent/healthcare-assistant/helpers/hallucination_helpers.py b/examples/agent/healthcare-assistant/helpers/hallucination_helpers.py
index f4e6e4d9..f417d506 100644
--- a/examples/agent/healthcare-assistant/helpers/hallucination_helpers.py
+++ b/examples/agent/healthcare-assistant/helpers/hallucination_helpers.py
@@ -47,6 +47,9 @@ def log_hallucination(
             logger.info("Creating new Splunk AO session for hallucination demo")
             splunk_ao_logger = SplunkAOLogger(project=project_name, agent_stream=agent_stream)
 
+        if external_session_id:
+            splunk_ao_logger.set_session(external_session_id)
+
         splunk_ao_logger.start_trace(
             input=question,
             name="Hallucination Demo",
diff --git a/examples/agent/healthcare-assistant/k8s.yaml b/examples/agent/healthcare-assistant/k8s.yaml
index 6f5181f2..b8283de6 100644
--- a/examples/agent/healthcare-assistant/k8s.yaml
+++ b/examples/agent/healthcare-assistant/k8s.yaml
@@ -16,7 +16,7 @@ spec:
     spec:
       containers:
         - name: healthcare-assistant
-          image: ghcr.io/splunk/healthcare-assistant:app-with-instrumentation
+          image: ghcr.io/splunk/healthcare-assistant-sao:0.3.0
           imagePullPolicy: Always
           ports:
             - containerPort: 8501
diff --git a/examples/agent/healthcare-assistant/requirements.txt b/examples/agent/healthcare-assistant/requirements.txt
index c4fc51e9..e28dc209 100644
--- a/examples/agent/healthcare-assistant/requirements.txt
+++ b/examples/agent/healthcare-assistant/requirements.txt
@@ -13,4 +13,4 @@ langchain-classic
 pyyaml
 toml
 pandas
-splunk-ao
+splunk-ao>=0.3.0