Fix handoff workflow context management and improve AG-UI demo (#5136)

python/samples/05-end-to-end/ag_ui_workflow_handoff/README.md

@@ -16,6 +16,10 @@ It includes:
 
 The backend uses Azure OpenAI responses and supports intent-driven, non-linear handoff routing.
 
+This demo keeps workflow state per `thread_id`. When the assistant ends a case with `Case complete.`, the UI blocks
+later top-level input on that same thread and asks the user to start a new case explicitly instead of resuming a
+terminated workflow.
+
 ## Folder Layout
 
 - `backend/server.py` - FastAPI + AG-UI endpoint + Handoff workflow
@@ -81,6 +85,28 @@ VITE_BACKEND_URL=http://127.0.0.1:8891 npm run dev
 7. When replacement is requested, wait for the `submit_replacement` reviewer interrupt and approve/reject it.
 8. If you asked for refund-only, the flow should close without replacement/shipping prompts.
 9. Confirm the case snapshot updates and workflow completion.
+10. After the case closes, another top-level message on the same thread is rejected with a notice.
+11. Click **Start New Case** to begin a fresh thread.
+
+## Important: `require_per_service_call_history_persistence`
+
+All agents participating in a handoff workflow **must** be constructed with
+`require_per_service_call_history_persistence=True`. The `HandoffBuilder` will
+raise a `ValueError` at build time if any participant is missing this flag.
+
+**Why this is required:** Handoff workflows use middleware that short-circuits
+tool calls via `MiddlewareTermination` when a handoff tool is invoked. Without
+per-service-call history persistence, local history providers would persist tool
+results that the service never received, causing call/result mismatches on
+subsequent turns.
+
+```python
+agent = Agent(
+    client=client,
+    name="my_agent",
+    require_per_service_call_history_persistence=True,  # Required for handoff
+)
+```
 
 ## What This Validates
 

python/samples/05-end-to-end/ag_ui_workflow_handoff/backend/server.py

@@ -17,12 +17,17 @@ import logging
 import logging.handlers
 import os
 import random
+from collections.abc import AsyncGenerator
+from typing import Any
 
 import uvicorn
 from agent_framework import (
     Agent,
     Message,
     Workflow,
+    WorkflowBuilder,
+    WorkflowContext,
+    executor,
     tool,
 )
 from agent_framework.ag_ui import AgentFrameworkWorkflow, add_agent_framework_fastapi_endpoint
@@ -101,6 +106,7 @@ def create_agents() -> tuple[Agent, Agent, Agent]:
             "4. If the issue is fully resolved, send a concise wrap-up that ends with exactly: Case complete."
         ),
         client=client,
+        require_per_service_call_history_persistence=True,
     )
 
     refund = Agent(
@@ -126,6 +132,7 @@ def create_agents() -> tuple[Agent, Agent, Agent]:
         ),
         client=client,
         tools=[lookup_order_details, submit_refund],
+        require_per_service_call_history_persistence=True,
     )
 
     order = Agent(
@@ -149,19 +156,25 @@ def create_agents() -> tuple[Agent, Agent, Agent]:
         ),
         client=client,
         tools=[lookup_order_details, submit_replacement],
+        require_per_service_call_history_persistence=True,
     )
 
     return triage, refund, order
 
 
+def is_case_complete_text(text: str) -> bool:
+    """Return True when a message ends with the explicit demo completion marker."""
+
+    return text.strip().lower().endswith("case complete.")
+
+
 def _termination_condition(conversation: list[Message]) -> bool:
     """Stop when any assistant emits an explicit completion marker."""
 
     for message in reversed(conversation):
         if message.role != "assistant":
             continue
-        text = (message.text or "").strip().lower()
-        if text.endswith("case complete."):
+        if is_case_complete_text(message.text or ""):
             return True
     return False
 
@@ -215,6 +228,71 @@ def create_handoff_workflow() -> Workflow:
     return builder.with_start_agent(triage).build()
 
 
+def create_closed_case_notice_workflow() -> Workflow:
+    """Build a tiny workflow that explains why a completed case cannot continue."""
+
+    @executor(id="closed_case_notice")
+    async def closed_case_notice(message: Message | None, ctx: WorkflowContext[None, str]) -> None:
+        del message
+        await ctx.yield_output(
+            "Your case is complete, but you're trying to do something new. Please start a new thread."
+        )
+
+    return WorkflowBuilder(start_executor=closed_case_notice).build()
+
+
+class DemoHandoffWorkflow(AgentFrameworkWorkflow):
+    """Workflow wrapper that blocks new top-level input on completed demo threads."""
+
+    def __init__(self) -> None:
+        super().__init__(
+            workflow_factory=lambda _thread_id: create_handoff_workflow(),
+            name="ag_ui_handoff_workflow_demo",
+            description="Dynamic handoff workflow demo with tool approvals and request_info resumes.",
+        )
+        self._completed_threads: set[str] = set()
+        self._closed_case_notice_runner = AgentFrameworkWorkflow(workflow=create_closed_case_notice_workflow())
+
+    async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[Any]:
+        """Intercept completed threads and return a helpful notice instead of resuming them."""
+
+        thread_id = self._thread_id_from_input(input_data)
+        has_messages = isinstance(input_data.get("messages"), list) and len(input_data.get("messages", [])) > 0
+        has_resume = input_data.get("resume") is not None
+
+        if thread_id in self._completed_threads and has_messages and not has_resume:
+            async for event in self._closed_case_notice_runner.run(input_data):
+                yield event
+            return
+
+        message_text_by_id: dict[str, str] = {}
+        case_completed_this_run = False
+
+        async for event in super().run(input_data):
+            event_type = getattr(event, "type", None)
+            if event_type == "TEXT_MESSAGE_START":
+                message_id = getattr(event, "message_id", None)
+                if isinstance(message_id, str):
+                    message_text_by_id[message_id] = ""
+            elif event_type == "TEXT_MESSAGE_CONTENT":
+                message_id = getattr(event, "message_id", None)
+                delta = getattr(event, "delta", None)
+                if isinstance(message_id, str) and isinstance(delta, str):
+                    message_text_by_id[message_id] = f"{message_text_by_id.get(message_id, '')}{delta}"
+            elif event_type == "TEXT_MESSAGE_END":
+                message_id = getattr(event, "message_id", None)
+                if isinstance(message_id, str):
+                    final_text = message_text_by_id.pop(message_id, "")
+                    if is_case_complete_text(final_text):
+                        case_completed_this_run = True
+
+            yield event
+
+        if case_completed_this_run:
+            self._completed_threads.add(thread_id)
+            self.clear_thread_workflow(thread_id)
+
+
 def create_app() -> FastAPI:
     """Create and configure the FastAPI application."""
 
@@ -231,11 +309,7 @@ def create_app() -> FastAPI:
         allow_headers=["*"],
     )
 
-    demo_workflow = AgentFrameworkWorkflow(
-        workflow_factory=lambda _thread_id: create_handoff_workflow(),
-        name="ag_ui_handoff_workflow_demo",
-        description="Dynamic handoff workflow demo with tool approvals and request_info resumes.",
-    )
+    demo_workflow = DemoHandoffWorkflow()
 
     add_agent_framework_fastapi_endpoint(
         app=app,

python/samples/05-end-to-end/ag_ui_workflow_handoff/frontend/src/App.tsx

@@ -54,6 +54,16 @@ const STARTER_PROMPTS = [
   "Help me with a damaged-order refund and replacement.",
 ];
 
+const DEFAULT_CASE_SNAPSHOT: CaseSnapshot = {
+  orderId: "Not captured",
+  refundAmount: "Not captured",
+  refundApproved: "pending",
+  shippingPreference: "Not selected",
+};
+
+const CLOSED_CASE_NOTICE =
+  "This case is already complete. Start a new case to open a fresh thread for a new request.";
+
 function randomId(): string {
   if (typeof crypto !== "undefined" && typeof crypto.randomUUID === "function") {
     return crypto.randomUUID();
@@ -213,6 +223,10 @@ function normalizeTextForDedupe(text: string): string {
   return text.replace(/\s+/g, " ").trim();
 }
 
+function isCaseCompleteText(text: string): boolean {
+  return text.trim().toLowerCase().endsWith("case complete.");
+}
+
 function normalizeShippingPreference(text: string): string | null {
   const normalized = text.trim().toLowerCase();
   if (normalized.length === 0) {
@@ -263,24 +277,21 @@ export default function App(): JSX.Element {
   const assistantMessageIndexRef = useRef<Record<string, number>>({});
   const activeRunIdRef = useRef<string | null>(null);
   const pendingUsageRef = useRef<UsageDiagnostics | null>(null);
+  const caseClosedRef = useRef<boolean>(false);
 
   const [messages, setMessages] = useState<DisplayMessage[]>([]);
   const [requestInfoById, setRequestInfoById] = useState<Record<string, RequestInfoPayload>>({});
   const [pendingInterrupts, setPendingInterrupts] = useState<Interrupt[]>([]);
   const [activeAgent, setActiveAgent] = useState<AgentId>("triage_agent");
   const [visitedAgents, setVisitedAgents] = useState<Set<AgentId>>(new Set(["triage_agent"]));
-  const [caseSnapshot, setCaseSnapshot] = useState<CaseSnapshot>({
-    orderId: "Not captured",
-    refundAmount: "Not captured",
-    refundApproved: "pending",
-    shippingPreference: "Not selected",
-  });
+  const [caseSnapshot, setCaseSnapshot] = useState<CaseSnapshot>(DEFAULT_CASE_SNAPSHOT);
   const [statusText, setStatusText] = useState<string>("Ready");
   const [isRunning, setIsRunning] = useState<boolean>(false);
   const [inputText, setInputText] = useState<string>("");
   const [isApprovalModalOpen, setIsApprovalModalOpen] = useState<boolean>(false);
   const [latestUsage, setLatestUsage] = useState<UsageDiagnostics | null>(null);
   const [usageHistory, setUsageHistory] = useState<UsageDiagnostics[]>([]);
+  const [isCaseClosed, setIsCaseClosed] = useState<boolean>(false);
 
   const currentInterrupt = pendingInterrupts[0];
   const currentInterruptKind = currentInterrupt ? interruptKind(currentInterrupt) : "unknown";
@@ -288,6 +299,7 @@ export default function App(): JSX.Element {
   const interruptPrompt = currentInterrupt
     ? extractPromptFromInterrupt(currentInterrupt, currentRequestInfo)
     : "No pending interrupt.";
+  const canStartFreshCase = !currentInterrupt && isCaseClosed;
 
   const functionCall = currentInterrupt ? extractFunctionCallFromInterrupt(currentInterrupt) : null;
   const functionArguments = useMemo(() => parseFunctionArguments(functionCall), [functionCall]);
@@ -304,6 +316,34 @@ export default function App(): JSX.Element {
     setMessages((prev) => [...prev, message]);
   };
 
+  const pushSystemNotice = (text: string): void => {
+    setMessages((prev) => {
+      if (prev.length > 0 && prev[prev.length - 1]?.role === "system" && prev[prev.length - 1]?.text === text) {
+        return prev;
+      }
+      return [...prev, { id: randomId(), role: "system", text }];
+    });
+  };
+
+  const resetConversationState = (): void => {
+    threadIdRef.current = randomId();
+    assistantMessageIndexRef.current = {};
+    activeRunIdRef.current = null;
+    pendingUsageRef.current = null;
+    caseClosedRef.current = false;
+
+    setMessages([]);
+    setRequestInfoById({});
+    setPendingInterrupts([]);
+    setActiveAgent("triage_agent");
+    setVisitedAgents(new Set(["triage_agent"]));
+    setCaseSnapshot(DEFAULT_CASE_SNAPSHOT);
+    setStatusText("Ready");
+    setInputText("");
+    setIsApprovalModalOpen(false);
+    setIsCaseClosed(false);
+  };
+
   const rebuildAssistantMessageIndex = (items: DisplayMessage[]): void => {
     const next: Record<string, number> = {};
     items.forEach((item, index) => {
@@ -364,6 +404,10 @@ export default function App(): JSX.Element {
       }
       const candidate = prev[index];
       if (candidate.role === "user" || candidate.text.trim().length > 0) {
+        if (candidate.role === "assistant" && isCaseCompleteText(candidate.text)) {
+          caseClosedRef.current = true;
+          setIsCaseClosed(true);
+        }
         return prev;
       }
       const next = prev.filter((item) => item.id !== messageId);
@@ -565,7 +609,9 @@ export default function App(): JSX.Element {
         }
 
         setPendingInterrupts(interruptPayload);
-        setStatusText(interruptPayload.length > 0 ? "Waiting for input" : "Run complete");
+        setStatusText(
+          interruptPayload.length > 0 ? "Waiting for input" : caseClosedRef.current ? "Case complete" : "Run complete"
+        );
         setIsRunning(false);
         break;
       }
@@ -652,6 +698,12 @@ export default function App(): JSX.Element {
   };
 
   const startNewTurn = async (text: string): Promise<void> => {
+    if (caseClosedRef.current && pendingInterrupts.length === 0) {
+      pushSystemNotice(CLOSED_CASE_NOTICE);
+      setStatusText("Case complete");
+      return;
+    }
+
     pushMessage({ id: randomId(), role: "user", text });
 
     await runWithPayload({
@@ -873,7 +925,20 @@ export default function App(): JSX.Element {
 
           <article className="card interrupt-card">
             <h2>Pending Action</h2>
-            {!currentInterrupt && <p className="muted">No interrupt pending. Start with one of the prompts below.</p>}
+            {!currentInterrupt && (
+              <div className="pending-empty-state">
+                <p className="muted">
+                  {isCaseClosed
+                    ? "This case is closed. New top-level messages on this thread are blocked until you start a new case."
+                    : "No interrupt pending. Start with one of the prompts below."}
+                </p>
+                {canStartFreshCase && (
+                  <button type="button" className="case-reset" onClick={resetConversationState} disabled={isRunning}>
+                    Start New Case
+                  </button>
+                )}
+              </div>
+            )}
 
             {currentInterrupt && (
               <div className="interrupt-body">
@@ -907,7 +972,7 @@ export default function App(): JSX.Element {
               </div>
             )}
 
-            {!currentInterrupt && (
+            {!currentInterrupt && !isCaseClosed && (
               <div className="starter-prompts">
                 {STARTER_PROMPTS.map((prompt) => (
                   <button key={prompt} type="button" onClick={() => void startNewTurn(prompt)} disabled={isRunning}>
@@ -944,7 +1009,9 @@ export default function App(): JSX.Element {
                   ? "Waiting for reviewer approval..."
                   : currentInterruptKind === "handoff_input"
                     ? "Reply to continue..."
-                    : "Describe your issue..."
+                    : isCaseClosed
+                      ? "This case is complete. Click Start New Case to open a fresh thread..."
+                      : "Describe your issue..."
               }
               disabled={isRunning || currentInterruptKind === "approval"}
             />

python/samples/05-end-to-end/ag_ui_workflow_handoff/frontend/src/styles.css

@@ -297,6 +297,7 @@ body {
 }
 
 .approval-actions button,
+.case-reset,
 .starter-prompts button,
 .chat-input button {
   border: 0;
@@ -307,6 +308,7 @@ body {
 }
 
 .approval-actions button:disabled,
+.case-reset:disabled,
 .starter-prompts button:disabled,
 .chat-input button:disabled {
   opacity: 0.6;
@@ -391,6 +393,19 @@ body {
   gap: 10px;
 }
 
+.pending-empty-state {
+  display: grid;
+  gap: 10px;
+}
+
+.case-reset {
+  width: fit-content;
+  border: 1px solid #bdcfdc;
+  background: #ecf3f8;
+  color: #345267;
+  padding: 10px 14px;
+}
+
 .starter-prompts button {
   text-align: left;
   background: linear-gradient(125deg, #fff8ef 0%, #ffe7cf 100%);