diff --git a/main.py b/main.py
index 81c648c..b2c34ea 100644
--- a/main.py
+++ b/main.py
@@ -21,41 +21,41 @@ def interrupt_node(state: GraphState) -> GraphState:
     # Call interrupt; execution pauses until resumed
     interrupt(payload)
     # After resume, payload will contain 'answer'
-    # We expect the resumed payload to be the same dict with added 'answer'
     # The node receives the resumed payload as its return value
-    # So we just return the state with human_value set
-    return {"human_value": payload.get("answer", None), "foo": state.get("foo", "")}
+    # We store the answer in state
+    state["human_value"] = payload.get("answer")
+    return state
 
 # Build graph
 builder = StateGraph(GraphState)
 builder.add_node("interrupt_node", interrupt_node)
-builder.set_entry_point("interrupt_node")
 builder.add_edge(START, "interrupt_node")
-# No further nodes; graph ends after interrupt node
+builder.add_edge("interrupt_node", END)
+
+# Compile with checkpoint
 graph = builder.compile(checkpointer=InMemorySaver())
 
-# Run graph with interrupt handling
-if __name__ == "__main__":
-    config = {"configurable": {"thread_id": "hitl-demo"}}
-    # Initial state
-    state = {"human_value": None, "foo": "initial"}
-    # Stream execution
-    for chunk in graph.stream(state, config):
-        if "__interrupt__" in chunk:
-            # Extract interrupt payload
-            interrupt_payload = chunk["__interrupt__"][0].value
-            # Show question to user
-            answer = questionary.select(
-                interrupt_payload["question"],
-                choices=interrupt_payload["allow_responds"],
-            ).ask()
-            # Add answer to payload and resume
-            interrupt_payload["answer"] = answer
-            # Resume graph
-            for resume_chunk in graph.stream(Command(resume=interrupt_payload), config):
-                # Print final state when graph finishes
-                if "node" in resume_chunk:
-                    print("Final state:", resume_chunk["node"])  
-        else:
-            # Normal output (none expected here)
-            pass
+# Run graph with stream and handle interrupt
+config = {"configurable": {"thread_id": "demo"}}
+
+# First run: will pause at interrupt
+stream = graph.stream({}, config)
+for chunk in stream:
+    if "__interrupt__" in chunk:
+        # Extract payload
+        payload = chunk["__interrupt__"][0].value
+        # Show question and get answer
+        answer = questionary.select(
+            payload["question"],
+            choices=payload["allow_responds"],
+        ).ask()
+        # Add answer to payload
+        payload["answer"] = answer
+        # Resume graph
+        stream = graph.stream(Command(resume=payload), config)
+        continue
+    # Print normal output
+    print(chunk)
+
+# After stream ends, print final state
+print("Final state:", stream.final_state)