agnt5 0.2.8a6__cp310-abi3-macosx_11_0_arm64.whl → 0.2.8a8__cp310-abi3-macosx_11_0_arm64.whl

agnt5/agent.py

@@ -12,11 +12,12 @@ import logging
 import time
 from typing import Any, Callable, Dict, List, Optional, Union
 
-from .context import Context
+from .context import Context, get_current_context, set_current_context
 from . import lm
 from .lm import GenerateRequest, GenerateResponse, LanguageModel, Message, ModelConfig, ToolDefinition
 from .tool import Tool, ToolRegistry
 from ._telemetry import setup_module_logger
+from .exceptions import WaitingForUserInputException
 
 logger = setup_module_logger(__name__)
 
@@ -79,6 +80,7 @@ class AgentContext(Context):
 
         self._agent_name = agent_name
         self._session_id = session_id or run_id
+        self.parent_context = parent_context  # Store for context chain traversal
 
         # Determine state adapter based on parent context
         from .entity import EntityStateAdapter, _get_state_adapter
@@ -118,6 +120,16 @@ class AgentContext(Context):
         # Entity key for database persistence (without :messages suffix to match API expectations)
         self._entity_key = f"agent:{agent_name}:{self._session_id}"
 
+        # Determine storage mode: "workflow" if parent is WorkflowContext, else "standalone"
+        self._storage_mode = "standalone"  # Default mode
+        self._workflow_entity = None
+
+        if parent_context and hasattr(parent_context, '_workflow_entity'):
+            # Agent is running within a workflow - store conversation in workflow state
+            self._storage_mode = "workflow"
+            self._workflow_entity = parent_context._workflow_entity
+            logger.debug(f"Agent '{agent_name}' using workflow storage mode (workflow entity: {self._workflow_entity.key})")
+
     @property
     def state(self):
         """
@@ -151,10 +163,27 @@ class AgentContext(Context):
         Retrieve conversation history from state, loading from database if needed.
 
         Uses the EntityStateAdapter which delegates to Rust core for cache-first loading.
+        If running within a workflow, loads from workflow entity state instead.
 
         Returns:
             List of Message objects from conversation history
         """
+        if self._storage_mode == "workflow":
+            return await self._load_from_workflow_state()
+        else:
+            return await self._load_from_entity_storage()
+
+    async def _load_from_workflow_state(self) -> List[Message]:
+        """Load conversation history from workflow entity state."""
+        key = f"agent.{self._agent_name}"
+        agent_data = self._workflow_entity.state.get(key, {})
+        messages_data = agent_data.get("messages", [])
+
+        # Convert dict representations back to Message objects
+        return self._convert_dicts_to_messages(messages_data)
+
+    async def _load_from_entity_storage(self) -> List[Message]:
+        """Load conversation history from AgentSession entity (standalone mode)."""
         entity_type = "AgentSession"
         entity_key = self._entity_key
 
@@ -173,6 +202,10 @@ class AgentContext(Context):
             messages_data = []
 
         # Convert dict representations back to Message objects
+        return self._convert_dicts_to_messages(messages_data)
+
+    def _convert_dicts_to_messages(self, messages_data: list) -> List[Message]:
+        """Convert list of message dicts to Message objects."""
         messages = []
         for msg_dict in messages_data:
             if isinstance(msg_dict, dict):
@@ -199,10 +232,48 @@ class AgentContext(Context):
         Save conversation history to state and persist to database.
 
         Uses the EntityStateAdapter which delegates to Rust core for version-checked saves.
+        If running within a workflow, saves to workflow entity state instead.
 
         Args:
             messages: List of Message objects to persist
         """
+        if self._storage_mode == "workflow":
+            await self._save_to_workflow_state(messages)
+        else:
+            await self._save_to_entity_storage(messages)
+
+    async def _save_to_workflow_state(self, messages: List[Message]) -> None:
+        """Save conversation history to workflow entity state."""
+        # Convert Message objects to dict for JSON serialization
+        messages_data = []
+        for msg in messages:
+            messages_data.append({
+                "role": msg.role.value if hasattr(msg.role, 'value') else str(msg.role),
+                "content": msg.content,
+                "timestamp": time.time()
+            })
+
+        # Build agent data structure
+        key = f"agent.{self._agent_name}"
+        current_data = self._workflow_entity.state.get(key, {})
+        now = time.time()
+
+        agent_data = {
+            "session_id": self._session_id,
+            "agent_name": self._agent_name,
+            "created_at": current_data.get("created_at", now),
+            "last_message_time": now,
+            "message_count": len(messages_data),
+            "messages": messages_data,
+            "metadata": getattr(self, '_custom_metadata', {})
+        }
+
+        # Store in workflow state (WorkflowEntity handles persistence)
+        self._workflow_entity.state.set(key, agent_data)
+        logger.info(f"Saved conversation to workflow state: {key} ({len(messages_data)} messages)")
+
+    async def _save_to_entity_storage(self, messages: List[Message]) -> None:
+        """Save conversation history to AgentSession entity (standalone mode)."""
         # Convert Message objects to dict for JSON serialization
         messages_data = []
         for msg in messages:
@@ -271,6 +342,35 @@ class AgentContext(Context):
             print(f"User ID: {metadata['custom'].get('user_id')}")
             ```
         """
+        if self._storage_mode == "workflow":
+            return await self._get_metadata_from_workflow()
+        else:
+            return await self._get_metadata_from_entity()
+
+    async def _get_metadata_from_workflow(self) -> Dict[str, Any]:
+        """Get metadata from workflow entity state."""
+        key = f"agent.{self._agent_name}"
+        agent_data = self._workflow_entity.state.get(key, {})
+
+        if not agent_data:
+            # No conversation exists yet - return defaults
+            return {
+                "created_at": None,
+                "last_activity": None,
+                "message_count": 0,
+                "custom": getattr(self, '_custom_metadata', {})
+            }
+
+        messages = agent_data.get("messages", [])
+        return {
+            "created_at": agent_data.get("created_at"),
+            "last_activity": agent_data.get("last_message_time"),
+            "message_count": len(messages),
+            "custom": agent_data.get("metadata", {})
+        }
+
+    async def _get_metadata_from_entity(self) -> Dict[str, Any]:
+        """Get metadata from AgentSession entity (standalone mode)."""
         entity_type = "AgentSession"
         entity_key = self._entity_key
 
@@ -633,6 +733,9 @@ class Agent:
             }
         }
 
+        # Auto-register agent for discovery by auto_register=True
+        AgentRegistry.register(self)
+
         # Store metadata
         self.metadata = {
             "description": instructions,
@@ -768,6 +871,47 @@ class Agent:
 
         return handoff_tool
 
+    def _detect_memory_scope(self, context: Optional[Context]) -> tuple[str, str]:
+        """
+        Auto-detect memory scope from context for agent conversation persistence.
+
+        Implements priority logic:
+        1. user_id → user-scoped memory (long-term)
+        2. session_id → session-scoped memory (multi-turn)
+        3. run_id → run-scoped memory (ephemeral)
+
+        Args:
+            context: WorkflowContext or other context with memory scoping fields
+
+        Returns:
+            Tuple of (entity_key, scope) where:
+            - entity_key: e.g., "user:user-456", "session:abc-123", "run:xyz-789"
+            - scope: "user", "session", or "run"
+
+        Example:
+            entity_key, scope = agent._detect_memory_scope(ctx)
+            # If ctx.user_id="user-123": ("user:user-123", "user")
+            # If ctx.session_id="sess-456": ("session:sess-456", "session")
+            # Otherwise: ("run:run-789", "run")
+        """
+        # Extract identifiers from context
+        user_id = getattr(context, 'user_id', None) if context else None
+        session_id = getattr(context, 'session_id', None) if context else None
+        run_id = getattr(context, 'run_id', None) if context else None
+
+        # Priority: user_id > session_id > run_id
+        if user_id:
+            return (f"user:{user_id}", "user")
+        elif session_id and session_id != run_id:  # Explicit session (not defaulting to run_id)
+            return (f"session:{session_id}", "session")
+        elif run_id:
+            return (f"run:{run_id}", "run")
+        else:
+            # Fallback: create ephemeral key
+            import uuid
+            fallback_run_id = f"agent-{self.name}-{uuid.uuid4().hex[:8]}"
+            return (f"run:{fallback_run_id}", "run")
+
     async def run(
         self,
         user_message: str,
@@ -777,7 +921,7 @@ class Agent:
 
         Args:
             user_message: User's input message
-            context: Optional context (auto-created if not provided)
+            context: Optional context (auto-created if not provided, or read from contextvar)
 
         Returns:
             AgentResult with output and execution details
@@ -789,6 +933,15 @@ class Agent:
             ```
         """
         # Create or adapt context
+        if context is None:
+            # Try to get context from task-local storage (set by workflow/function decorator)
+            context = get_current_context()
+
+        # IMPORTANT: Capture workflow context NOW before we replace it with AgentContext
+        # This allows LM calls inside the agent to emit workflow checkpoints
+        from .workflow import WorkflowContext
+        workflow_ctx = context if isinstance(context, WorkflowContext) else None
+
         if context is None:
             # Standalone execution - create AgentContext
             import uuid
@@ -802,13 +955,20 @@ class Agent:
             pass
         elif hasattr(context, '_workflow_entity'):
             # WorkflowContext - create AgentContext that inherits state
+            # Auto-detect memory scope based on user_id/session_id/run_id priority
+            entity_key, scope = self._detect_memory_scope(context)
+
             import uuid
             run_id = f"{context.run_id}:agent:{self.name}"
+            # Extract the ID from entity_key (e.g., "session:abc-123" → "abc-123")
+            detected_session_id = entity_key.split(":", 1)[1] if ":" in entity_key else context.run_id
+
             context = AgentContext(
                 run_id=run_id,
                 agent_name=self.name,
-                session_id=context.run_id,  # Share workflow's session
+                session_id=detected_session_id,  # Use auto-detected scope
                 parent_context=context,
+                runtime_context=getattr(context, '_runtime_context', None),  # Inherit trace context
             )
         else:
             # FunctionContext or other - create new AgentContext
@@ -817,198 +977,611 @@ class Agent:
             context = AgentContext(
                 run_id=run_id,
                 agent_name=self.name,
+                runtime_context=getattr(context, '_runtime_context', None),  # Inherit trace context
             )
 
-        # Load conversation history from state (if AgentContext)
-        if isinstance(context, AgentContext):
-            messages: List[Message] = await context.get_conversation_history()
-            # Add new user message
-            messages.append(Message.user(user_message))
-            # Save updated conversation
-            await context.save_conversation_history(messages)
-        else:
-            # Fallback for non-AgentContext (shouldn't happen with code above)
-            messages = [Message.user(user_message)]
+        # Emit checkpoint if called within a workflow context
+        if workflow_ctx is not None:
+            workflow_ctx._send_checkpoint("workflow.agent.started", {
+                "agent.name": self.name,
+                "agent.model": self.model_name,
+                "agent.tools": list(self.tools.keys()),
+                "agent.max_iterations": self.max_iterations,
+                "user_message": user_message,
+            })
 
-        # Create span for agent execution with trace linking
-        from ._core import create_span
+        # NEW: Check if this is a resume from HITL
+        if workflow_ctx and hasattr(workflow_ctx, "_agent_resume_info"):
+            resume_info = workflow_ctx._agent_resume_info
+            if resume_info["agent_name"] == self.name:
+                self.logger.info("Detected HITL resume, calling resume_from_hitl()")
 
-        with create_span(
-            self.name,
-            "agent",
-            context._runtime_context if hasattr(context, "_runtime_context") else None,
-            {
-                "agent.name": self.name,
-                "agent.model": self.model_name,  # Use model_name (always a string)
-                "agent.max_iterations": str(self.max_iterations),
-            },
-        ) as span:
-            all_tool_calls: List[Dict[str, Any]] = []
-
-            # Reasoning loop
-            for iteration in range(self.max_iterations):
-                # Build tool definitions for LLM
-                tool_defs = [
-                    ToolDefinition(
-                        name=tool.name,
-                        description=tool.description,
-                        parameters=tool.input_schema,
-                    )
-                    for tool in self.tools.values()
-                ]
-
-                # Convert messages to dict format for lm.generate()
-                messages_dict = []
-                for msg in messages:
-                    messages_dict.append({
-                        "role": msg.role.value,
-                        "content": msg.content
-                    })
+                # Clear resume info to avoid re-entry
+                delattr(workflow_ctx, "_agent_resume_info")
 
-                # Call LLM
-                # Check if we have a legacy LanguageModel instance or need to create one
-                if self._language_model is not None:
-                    # Legacy API: use provided LanguageModel instance
-                    request = GenerateRequest(
-                        model="mock-model",  # Not used by MockLanguageModel
-                        system_prompt=self.instructions,
-                        messages=messages,
-                        tools=tool_defs if tool_defs else [],
-                    )
-                    request.config.temperature = self.temperature
-                    if self.max_tokens:
-                        request.config.max_tokens = self.max_tokens
-                    if self.top_p:
-                        request.config.top_p = self.top_p
-                    response = await self._language_model.generate(request)
+                # Resume from checkpoint (context setup happens inside resume_from_hitl)
+                return await self.resume_from_hitl(
+                    context=workflow_ctx,
+                    agent_context=resume_info["agent_context"],
+                    user_response=resume_info["user_response"],
+                )
+
+        # Set context in task-local storage for automatic propagation to tools and LM calls
+        token = set_current_context(context)
+        try:
+            try:
+                # Load conversation history from state (if AgentContext)
+                if isinstance(context, AgentContext):
+                    messages: List[Message] = await context.get_conversation_history()
+                    # Add new user message
+                    messages.append(Message.user(user_message))
+                    # Save updated conversation
+                    await context.save_conversation_history(messages)
                 else:
-                    # New API: model is a string, create internal LM instance
-                    request = GenerateRequest(
-                        model=self.model,
-                        system_prompt=self.instructions,
-                        messages=messages,
-                        tools=tool_defs if tool_defs else [],
-                    )
-                    request.config.temperature = self.temperature
-                    if self.max_tokens:
-                        request.config.max_tokens = self.max_tokens
-                    if self.top_p:
-                        request.config.top_p = self.top_p
-
-                    # Create internal LM instance for generation
-                    # TODO: Use model_config when provided
-                    from .lm import _LanguageModel
-                    provider, model_name = self.model.split('/', 1)
-                    internal_lm = _LanguageModel(provider=provider.lower(), default_model=None)
-                    response = await internal_lm.generate(request)
-
-                # Add assistant response to messages
-                messages.append(Message.assistant(response.text))
-
-                # Check if LLM wants to use tools
-                if response.tool_calls:
-                    self.logger.debug(f"Agent calling {len(response.tool_calls)} tool(s)")
-
-                    # Store current conversation in context for potential handoffs
-                    # Use a simple dict attribute since we don't need full state persistence for this
-                    if not hasattr(context, '_agent_data'):
-                        context._agent_data = {}
-                    context._agent_data["_current_conversation"] = messages
-
-                    # Execute tool calls
-                    tool_results = []
-                    for tool_call in response.tool_calls:
-                        tool_name = tool_call["name"]
-                        tool_args_str = tool_call["arguments"]
-
-                        # Track tool call
-                        all_tool_calls.append(
-                            {
-                                "name": tool_name,
-                                "arguments": tool_args_str,
-                                "iteration": iteration + 1,
-                            }
-                        )
+                    # Fallback for non-AgentContext (shouldn't happen with code above)
+                    messages = [Message.user(user_message)]
+
+                # Create span for agent execution with trace linking
+                from ._core import create_span
+
+                with create_span(
+                    self.name,
+                    "agent",
+                    context._runtime_context if hasattr(context, "_runtime_context") else None,
+                    {
+                        "agent.name": self.name,
+                        "agent.model": self.model_name,  # Use model_name (always a string)
+                        "agent.max_iterations": str(self.max_iterations),
+                    },
+                ) as span:
+                    all_tool_calls: List[Dict[str, Any]] = []
+
+                    # Reasoning loop
+                    for iteration in range(self.max_iterations):
+                        # Build tool definitions for LLM
+                        tool_defs = [
+                            ToolDefinition(
+                                name=tool.name,
+                                description=tool.description,
+                                parameters=tool.input_schema,
+                            )
+                            for tool in self.tools.values()
+                        ]
 
-                        # Execute tool
-                        try:
-                            # Parse arguments
-                            tool_args = json.loads(tool_args_str)
+                        # Convert messages to dict format for lm.generate()
+                        messages_dict = []
+                        for msg in messages:
+                            messages_dict.append({
+                                "role": msg.role.value,
+                                "content": msg.content
+                            })
+
+                        # Call LLM
+                        # Check if we have a legacy LanguageModel instance or need to create one
+                        if self._language_model is not None:
+                            # Legacy API: use provided LanguageModel instance
+                            request = GenerateRequest(
+                                model="mock-model",  # Not used by MockLanguageModel
+                                system_prompt=self.instructions,
+                                messages=messages,
+                                tools=tool_defs if tool_defs else [],
+                            )
+                            request.config.temperature = self.temperature
+                            if self.max_tokens:
+                                request.config.max_tokens = self.max_tokens
+                            if self.top_p:
+                                request.config.top_p = self.top_p
+                            response = await self._language_model.generate(request)
+                        else:
+                            # New API: model is a string, create internal LM instance
+                            request = GenerateRequest(
+                                model=self.model,
+                                system_prompt=self.instructions,
+                                messages=messages,
+                                tools=tool_defs if tool_defs else [],
+                            )
+                            request.config.temperature = self.temperature
+                            if self.max_tokens:
+                                request.config.max_tokens = self.max_tokens
+                            if self.top_p:
+                                request.config.top_p = self.top_p
+
+                            # Create internal LM instance for generation
+                            # TODO: Use model_config when provided
+                            from .lm import _LanguageModel
+                            provider, model_name = self.model.split('/', 1)
+                            internal_lm = _LanguageModel(provider=provider.lower(), default_model=None)
+                            response = await internal_lm.generate(request)
+
+                        # Add assistant response to messages
+                        messages.append(Message.assistant(response.text))
+
+                        # Check if LLM wants to use tools
+                        if response.tool_calls:
+                            self.logger.debug(f"Agent calling {len(response.tool_calls)} tool(s)")
+
+                            # Store current conversation in context for potential handoffs
+                            # Use a simple dict attribute since we don't need full state persistence for this
+                            if not hasattr(context, '_agent_data'):
+                                context._agent_data = {}
+                            context._agent_data["_current_conversation"] = messages
+
+                            # Execute tool calls
+                            tool_results = []
+                            for tool_call in response.tool_calls:
+                                tool_name = tool_call["name"]
+                                tool_args_str = tool_call["arguments"]
+
+                                # Track tool call
+                                all_tool_calls.append(
+                                    {
+                                        "name": tool_name,
+                                        "arguments": tool_args_str,
+                                        "iteration": iteration + 1,
+                                    }
+                                )
 
-                            # Get tool
-                            tool = self.tools.get(tool_name)
-                            if not tool:
-                                result_text = f"Error: Tool '{tool_name}' not found"
-                            else:
                                 # Execute tool
-                                result = await tool.invoke(context, **tool_args)
-
-                                # Check if this was a handoff
-                                if isinstance(result, dict) and result.get("_handoff"):
-                                    self.logger.info(
-                                        f"Handoff detected to '{result['to_agent']}', "
-                                        f"terminating current agent"
-                                    )
-                                    # Save conversation before returning
-                                    if isinstance(context, AgentContext):
-                                        await context.save_conversation_history(messages)
-                                    # Return immediately with handoff result
-                                    return AgentResult(
-                                        output=result["output"],
-                                        tool_calls=all_tool_calls + result.get("tool_calls", []),
-                                        context=context,
-                                        handoff_to=result["to_agent"],
-                                        handoff_metadata=result,
+                                try:
+                                    # Parse arguments
+                                    tool_args = json.loads(tool_args_str)
+
+                                    # Get tool
+                                    tool = self.tools.get(tool_name)
+                                    if not tool:
+                                        result_text = f"Error: Tool '{tool_name}' not found"
+                                    else:
+                                        # Execute tool
+                                        result = await tool.invoke(context, **tool_args)
+
+                                        # Check if this was a handoff
+                                        if isinstance(result, dict) and result.get("_handoff"):
+                                            self.logger.info(
+                                                f"Handoff detected to '{result['to_agent']}', "
+                                                f"terminating current agent"
+                                            )
+                                            # Save conversation before returning
+                                            if isinstance(context, AgentContext):
+                                                await context.save_conversation_history(messages)
+                                            # Return immediately with handoff result
+                                            return AgentResult(
+                                                output=result["output"],
+                                                tool_calls=all_tool_calls + result.get("tool_calls", []),
+                                                context=context,
+                                                handoff_to=result["to_agent"],
+                                                handoff_metadata=result,
+                                            )
+
+                                        result_text = json.dumps(result) if result else "null"
+
+                                    tool_results.append(
+                                        {"tool": tool_name, "result": result_text, "error": None}
                                     )
 
-                                result_text = json.dumps(result) if result else "null"
+                                except WaitingForUserInputException as e:
+                                    # HITL PAUSE: Capture agent state and propagate exception
+                                    self.logger.info(f"Agent pausing for user input at iteration {iteration}")
+
+                                    # Serialize messages to dict format
+                                    messages_dict = [
+                                        {"role": msg.role.value, "content": msg.content}
+                                        for msg in messages
+                                    ]
+
+                                    # Enhance exception with agent execution context
+                                    raise WaitingForUserInputException(
+                                        question=e.question,
+                                        input_type=e.input_type,
+                                        options=e.options,
+                                        checkpoint_state=e.checkpoint_state,
+                                        agent_context={
+                                            "agent_name": self.name,
+                                            "iteration": iteration,
+                                            "messages": messages_dict,
+                                            "tool_results": tool_results,
+                                            "pending_tool_call": {
+                                                "name": tool_call["name"],
+                                                "arguments": tool_call["arguments"],
+                                                "tool_call_index": response.tool_calls.index(tool_call),
+                                            },
+                                            "all_tool_calls": all_tool_calls,
+                                            "model_config": {
+                                                "model": self.model,
+                                                "temperature": self.temperature,
+                                                "max_tokens": self.max_tokens,
+                                                "top_p": self.top_p,
+                                            },
+                                        },
+                                    ) from e
+
+                                except Exception as e:
+                                    # Regular tool errors - log and continue
+                                    self.logger.error(f"Tool execution error: {e}")
+                                    tool_results.append(
+                                        {"tool": tool_name, "result": None, "error": str(e)}
+                                    )
 
-                            tool_results.append(
-                                {"tool": tool_name, "result": result_text, "error": None}
+                            # Add tool results to conversation
+                            results_text = "\n".join(
+                                [
+                                    f"Tool: {tr['tool']}\nResult: {tr['result']}"
+                                    if tr["error"] is None
+                                    else f"Tool: {tr['tool']}\nError: {tr['error']}"
+                                    for tr in tool_results
+                                ]
                             )
-
-                        except Exception as e:
-                            self.logger.error(f"Tool execution error: {e}")
-                            tool_results.append(
-                                {"tool": tool_name, "result": None, "error": str(e)}
+                            messages.append(Message.user(f"Tool results:\n{results_text}\n\nPlease provide your final answer based on these results."))
+
+                            # Continue loop for agent to process results
+
+                        else:
+                            # No tool calls - agent is done
+                            self.logger.debug(f"Agent completed after {iteration + 1} iterations")
+                            # Save conversation before returning
+                            if isinstance(context, AgentContext):
+                                await context.save_conversation_history(messages)
+
+                            # Emit completion checkpoint
+                            if workflow_ctx:
+                                workflow_ctx._send_checkpoint("workflow.agent.completed", {
+                                    "agent.name": self.name,
+                                    "agent.iterations": iteration + 1,
+                                    "agent.tool_calls_count": len(all_tool_calls),
+                                    "output_length": len(response.text),
+                                })
+
+                            return AgentResult(
+                                output=response.text,
+                                tool_calls=all_tool_calls,
+                                context=context,
                             )
 
-                    # Add tool results to conversation
-                    results_text = "\n".join(
-                        [
-                            f"Tool: {tr['tool']}\nResult: {tr['result']}"
-                            if tr["error"] is None
-                            else f"Tool: {tr['tool']}\nError: {tr['error']}"
-                            for tr in tool_results
-                        ]
-                    )
-                    messages.append(Message.user(f"Tool results:\n{results_text}\n\nPlease provide your final answer based on these results."))
-
-                    # Continue loop for agent to process results
-
-                else:
-                    # No tool calls - agent is done
-                    self.logger.debug(f"Agent completed after {iteration + 1} iterations")
+                    # Max iterations reached
+                    self.logger.warning(f"Agent reached max iterations ({self.max_iterations})")
+                    final_output = messages[-1].content if messages else "No output generated"
                     # Save conversation before returning
                     if isinstance(context, AgentContext):
                         await context.save_conversation_history(messages)
+
+                    # Emit completion checkpoint with max iterations flag
+                    if workflow_ctx:
+                        workflow_ctx._send_checkpoint("workflow.agent.completed", {
+                            "agent.name": self.name,
+                            "agent.iterations": self.max_iterations,
+                            "agent.tool_calls_count": len(all_tool_calls),
+                            "agent.max_iterations_reached": True,
+                            "output_length": len(final_output),
+                        })
+
                     return AgentResult(
-                        output=response.text,
+                        output=final_output,
                         tool_calls=all_tool_calls,
                         context=context,
                     )
+            except Exception as e:
+                # Emit error checkpoint for observability
+                if workflow_ctx:
+                    workflow_ctx._send_checkpoint("workflow.agent.error", {
+                        "agent.name": self.name,
+                        "error": str(e),
+                        "error_type": type(e).__name__,
+                    })
+                raise
+        finally:
+            # Always reset context to prevent leakage between agent executions
+            from .context import _current_context
+            _current_context.reset(token)
+
+    async def resume_from_hitl(
+        self,
+        context: Context,
+        agent_context: Dict,
+        user_response: str,
+    ) -> AgentResult:
+        """
+        Resume agent execution after HITL pause.
+
+        This method reconstructs agent state from the checkpoint and injects
+        the user's response as the successful tool result, then continues
+        the conversation loop.
+
+        Args:
+            context: Current execution context (workflow or agent)
+            agent_context: Agent state from WaitingForUserInputException.agent_context
+            user_response: User's answer to the HITL question
+
+        Returns:
+            AgentResult with final output and tool calls
+        """
+        self.logger.info(f"Resuming agent '{self.name}' from HITL pause")
+
+        # 1. Restore conversation state
+        messages = [
+            Message(role=lm.MessageRole(msg["role"]), content=msg["content"])
+            for msg in agent_context["messages"]
+        ]
+        iteration = agent_context["iteration"]
+        all_tool_calls = agent_context["all_tool_calls"]
+
+        # 2. Restore partial tool results for current iteration
+        tool_results = agent_context["tool_results"]
+
+        # 3. Inject user response as successful tool result
+        pending_tool = agent_context["pending_tool_call"]
+        tool_results.append({
+            "tool": pending_tool["name"],
+            "result": json.dumps(user_response),
+            "error": None,
+        })
+
+        self.logger.debug(
+            f"Injected user response for tool '{pending_tool['name']}': {user_response}"
+        )
+
+        # 4. Add tool results to conversation
+        results_text = "\n".join([
+            f"Tool: {tr['tool']}\nResult: {tr['result']}"
+            if tr["error"] is None
+            else f"Tool: {tr['tool']}\nError: {tr['error']}"
+            for tr in tool_results
+        ])
+        messages.append(Message.user(
+            f"Tool results:\n{results_text}\n\n"
+            f"Please provide your final answer based on these results."
+        ))
+
+        # 5. Continue agent execution loop from next iteration
+        return await self._continue_execution_from_iteration(
+            context=context,
+            messages=messages,
+            iteration=iteration + 1,  # Next iteration
+            all_tool_calls=all_tool_calls,
+        )
+
+    async def _continue_execution_from_iteration(
+        self,
+        context: Context,
+        messages: List[Message],
+        iteration: int,
+        all_tool_calls: List[Dict],
+    ) -> AgentResult:
+        """
+        Continue agent execution from a specific iteration.
+
+        This is the core execution loop extracted to support both:
+        1. Normal execution (starting from iteration 0)
+        2. Resume after HITL (starting from iteration N)
 
-            # Max iterations reached
-            self.logger.warning(f"Agent reached max iterations ({self.max_iterations})")
-            final_output = messages[-1].content if messages else "No output generated"
-            # Save conversation before returning
-            if isinstance(context, AgentContext):
-                await context.save_conversation_history(messages)
-            return AgentResult(
-                output=final_output,
-                tool_calls=all_tool_calls,
-                context=context,
+        Args:
+            context: Execution context
+            messages: Conversation history
+            iteration: Starting iteration number
+            all_tool_calls: Accumulated tool calls
+
+        Returns:
+            AgentResult with output and tool calls
+        """
+        # Extract workflow context for checkpointing
+        workflow_ctx = None
+        if hasattr(context, "_workflow_entity"):
+            workflow_ctx = context
+        elif hasattr(context, "_agent_data") and "_workflow_ctx" in context._agent_data:
+            workflow_ctx = context._agent_data["_workflow_ctx"]
+
+        # Prepare tool definitions
+        tool_defs = [
+            ToolDefinition(
+                name=name,
+                description=tool.description or f"Tool: {name}",
+                parameters=tool.input_schema if hasattr(tool, "input_schema") else {},
             )
+            for name, tool in self.tools.items()
+        ]
+
+        # Main iteration loop (continue from specified iteration)
+        while iteration < self.max_iterations:
+            self.logger.debug(f"Agent iteration {iteration + 1}/{self.max_iterations}")
+
+            # Call LLM for next response
+            if self._language_model:
+                # Legacy API: model is a LanguageModel instance
+                request = GenerateRequest(
+                    system_prompt=self.instructions,
+                    messages=messages,
+                    tools=tool_defs if tool_defs else [],
+                )
+                request.config.temperature = self.temperature
+                if self.max_tokens:
+                    request.config.max_tokens = self.max_tokens
+                if self.top_p:
+                    request.config.top_p = self.top_p
+                response = await self._language_model.generate(request)
+            else:
+                # New API: model is a string, create internal LM instance
+                request = GenerateRequest(
+                    model=self.model,
+                    system_prompt=self.instructions,
+                    messages=messages,
+                    tools=tool_defs if tool_defs else [],
+                )
+                request.config.temperature = self.temperature
+                if self.max_tokens:
+                    request.config.max_tokens = self.max_tokens
+                if self.top_p:
+                    request.config.top_p = self.top_p
+
+                # Create internal LM instance for generation
+                from .lm import _LanguageModel
+                provider, model_name = self.model.split('/', 1)
+                internal_lm = _LanguageModel(provider=provider.lower(), default_model=None)
+                response = await internal_lm.generate(request)
+
+            # Add assistant response to messages
+            messages.append(Message.assistant(response.text))
+
+            # Check if LLM wants to use tools
+            if response.tool_calls:
+                self.logger.debug(f"Agent calling {len(response.tool_calls)} tool(s)")
+
+                # Store current conversation in context for potential handoffs
+                if not hasattr(context, '_agent_data'):
+                    context._agent_data = {}
+                context._agent_data["_current_conversation"] = messages
+
+                # Execute tool calls
+                tool_results = []
+                for tool_call in response.tool_calls:
+                    tool_name = tool_call["name"]
+                    tool_args_str = tool_call["arguments"]
+
+                    # Track tool call
+                    all_tool_calls.append({
+                        "name": tool_name,
+                        "arguments": tool_args_str,
+                        "iteration": iteration + 1,
+                    })
+
+                    # Execute tool
+                    try:
+                        # Parse arguments
+                        tool_args = json.loads(tool_args_str)
+
+                        # Get tool
+                        tool = self.tools.get(tool_name)
+                        if not tool:
+                            result_text = f"Error: Tool '{tool_name}' not found"
+                        else:
+                            # Execute tool
+                            result = await tool.invoke(context, **tool_args)
+
+                            # Check if this was a handoff
+                            if isinstance(result, dict) and result.get("_handoff"):
+                                self.logger.info(
+                                    f"Handoff detected to '{result['to_agent']}', "
+                                    f"terminating current agent"
+                                )
+                                # Save conversation before returning
+                                if isinstance(context, AgentContext):
+                                    await context.save_conversation_history(messages)
+                                # Return immediately with handoff result
+                                return AgentResult(
+                                    output=result["output"],
+                                    tool_calls=all_tool_calls + result.get("tool_calls", []),
+                                    context=context,
+                                    handoff_to=result["to_agent"],
+                                    handoff_metadata=result,
+                                )
+
+                            result_text = json.dumps(result) if result else "null"
+
+                        tool_results.append(
+                            {"tool": tool_name, "result": result_text, "error": None}
+                        )
+
+                    except WaitingForUserInputException as e:
+                        # HITL PAUSE: Capture agent state and propagate exception
+                        self.logger.info(f"Agent pausing for user input at iteration {iteration}")
+
+                        # Serialize messages to dict format
+                        messages_dict = [
+                            {"role": msg.role.value, "content": msg.content}
+                            for msg in messages
+                        ]
+
+                        # Enhance exception with agent execution context
+                        from .exceptions import WaitingForUserInputException
+                        raise WaitingForUserInputException(
+                            question=e.question,
+                            input_type=e.input_type,
+                            options=e.options,
+                            checkpoint_state=e.checkpoint_state,
+                            agent_context={
+                                "agent_name": self.name,
+                                "iteration": iteration,
+                                "messages": messages_dict,
+                                "tool_results": tool_results,
+                                "pending_tool_call": {
+                                    "name": tool_call["name"],
+                                    "arguments": tool_call["arguments"],
+                                    "tool_call_index": response.tool_calls.index(tool_call),
+                                },
+                                "all_tool_calls": all_tool_calls,
+                                "model_config": {
+                                    "model": self.model,
+                                    "temperature": self.temperature,
+                                    "max_tokens": self.max_tokens,
+                                    "top_p": self.top_p,
+                                },
+                            },
+                        ) from e
+
+                    except Exception as e:
+                        # Regular tool errors - log and continue
+                        self.logger.error(f"Tool execution error: {e}")
+                        tool_results.append(
+                            {"tool": tool_name, "result": None, "error": str(e)}
+                        )
+
+                # Add tool results to conversation
+                results_text = "\n".join([
+                    f"Tool: {tr['tool']}\nResult: {tr['result']}"
+                    if tr["error"] is None
+                    else f"Tool: {tr['tool']}\nError: {tr['error']}"
+                    for tr in tool_results
+                ])
+                messages.append(Message.user(
+                    f"Tool results:\n{results_text}\n\n"
+                    f"Please provide your final answer based on these results."
+                ))
+
+                # Continue loop for agent to process results
+
+            else:
+                # No tool calls - agent is done
+                self.logger.debug(f"Agent completed after {iteration + 1} iterations")
+                # Save conversation before returning
+                if isinstance(context, AgentContext):
+                    await context.save_conversation_history(messages)
+
+                # Emit completion checkpoint
+                if workflow_ctx:
+                    workflow_ctx._send_checkpoint("workflow.agent.completed", {
+                        "agent.name": self.name,
+                        "agent.iterations": iteration + 1,
+                        "agent.tool_calls_count": len(all_tool_calls),
+                        "output_length": len(response.text),
+                    })
+
+                return AgentResult(
+                    output=response.text,
+                    tool_calls=all_tool_calls,
+                    context=context,
+                )
+
+            iteration += 1
+
+        # Max iterations reached
+        self.logger.warning(f"Agent reached max iterations ({self.max_iterations})")
+        final_output = messages[-1].content if messages else "No output generated"
+        # Save conversation before returning
+        if isinstance(context, AgentContext):
+            await context.save_conversation_history(messages)
+
+        # Emit completion checkpoint with max iterations flag
+        if workflow_ctx:
+            workflow_ctx._send_checkpoint("workflow.agent.completed", {
+                "agent.name": self.name,
+                "agent.iterations": self.max_iterations,
+                "agent.tool_calls_count": len(all_tool_calls),
+                "agent.max_iterations_reached": True,
+                "output_length": len(final_output),
+            })
+
+        return AgentResult(
+            output=final_output,
+            tool_calls=all_tool_calls,
+            context=context,
+        )
 
 
 def agent(