Repository: microsoft/autogen
Pull request: fix: fix user input in m1
Closes #4975

Problem statement:
RuntimeError: Failed to get user input: Failed to get user input:
### What happened?

Error doesn't occur with the `--no-hil` option.

```console
(.venv-m1) gaganbansal@MacBook-Pro workspace % m1 "hi"
Enter your response: ---------- user ----------
hi
---------- MagenticOneOrchestrator ----------

We are working to address the following user request:

hi


To answer this request we have assembled the following team:

FileSurfer: An agent that can handle local files.
WebSurfer:
    A helpful assistant with access to a web browser.
    Ask them to perform web searches, open pages, and interact with content (e.g., clicking links, scrolling the viewport, etc., filling in form fields, etc.).
    It can also summarize the entire page, or answer questions based on the content of the page.
    It can also be asked to sleep and wait for pages to load, in cases where the pages seem to be taking a while to load.
Coder: A helpful and general-purpose AI assistant that has strong language skills, Python skills, and Linux command line skills.
Executor: A computer terminal that performs no other action than running Python scripts (provided to it quoted in ```python code blocks), or sh shell scripts (provided to it quoted in ```sh code blocks).
User: A huError processing publish message for User/8d19aae6-3888-4890-9cb2-dc14c5d22958
Traceback (most recent call last):
  File "/Users/gaganbansal/workspace/.venv-m1/lib/python3.13/site-packages/autogen_agentchat/agents/_user_proxy_agent.py", line 144, in _get_input
    return await async_func(prompt, cancellation_token)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/gaganbansal/workspace/.venv-m1/lib/python3.13/site-packages/autogen_agentchat/agents/_user_proxy_agent.py", line 24, in cancellable_input
    return await task
           ^^^^^^^^^^
  File "/Users/gaganbansal/workspace/.venv-m1/lib/python3.13/site-packages/aioconsole/stream.py", line 297, in ainput
    raise EOFError
EOFError

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/gaganbansal/workspace/.venv-m1/lib/python3.13/site-packages/autogen_agentchat/agents/_user_proxy_agent.py", line 167, in on_messages
    user_input = await self._get_inputman user


Here is an initial fact sheet to consider:

1. GIVEN OR VERIFIED FACTS
   - The specific request is simply "hi," which provides no specific facts or figures.

2. FACTS TO LOOK UP
   - None required, as the request does not provide any context or mention any factual information that would need verification or supplementation.

3. FACTS TO DERIVE
   - None, as there is no information given in the request to logically deduce further details or insights.

4. EDUCATED GUESSES
   - "hi" is a greeting, so the request could be interpreted as a casual initiation of conversation or inquiry.


Here is the plan to follow as best as possible:

Since the original request is simply "hi," which is a greeting, and does not specify any particular task or objective, here is a short bullet-point plan for addressing the request based on the available team members:

- User: Engage by responding to the greeting to establish further context or clarify any specific needs or questions.
  - "Hello! How can I assist you today?"

- WebSurfer: Remain on standby in case the User follows up with a request that involves needing information from the web.

- FileSurfer: Stand by if the User requires handling or accessing local files in any follow-up request.

- Coder: Be prepared to assist with any programming-related inquiries or tasks if they arise.

- Executor: Remain inactive unless directed by Coder to run a specific script.

This plan focuses on establishing communication and determining the next steps based on User input.

---------- MagenticOneOrchestrator ----------
Hello! How can I assist you today?
(prompt, cancellation_token)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/gaganbansal/workspace/.venv-m1/lib/python3.13/site-packages/autogen_agentchat/agents/_user_proxy_agent.py", line 154, in _get_input
    raise RuntimeError(f"Failed to get user input: {str(e)}") from e
RuntimeError: Failed to get user input:

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/gaganbansal/workspace/.venv-m1/lib/python3.13/site-packages/autogen_core/_single_threaded_agent_runtime.py", line 409, in _on_message
    return await agent.on_message(
           ^^^^^^^^^^^^^^^^^^^^^^^
    ...<2 lines>...
    )
    ^
  File "/Users/gaganbansal/workspace/.venv-m1/lib/python3.13/site-packages/autogen_core/_base_agent.py", line 113, in on_message
    return await self.on_message_impl(message, ctx)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/gaganbansal/workspace/.venv-m1/lib/python3.13/site-packages/autogen_agentchat/teams/_group_chat/_sequential_routed_agent.py", line 48, in on_message_impl
    return await super().on_message_impl(message, ctx)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/gaganbansal/workspace/.venv-m1/lib/python3.13/site-packages/autogen_core/_routed_agent.py", line 485, in on_message_impl
    return await h(self, message, ctx)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/gaganbansal/workspace/.venv-m1/lib/python3.13/site-packages/autogen_core/_routed_agent.py", line 268, in wrapper
    return_value = await func(self, message, ctx)  # type: ignore
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/gaganbansal/workspace/.venv-m1/lib/python3.13/site-packages/autogen_agentchat/teams/_group_chat/_chat_agent_container.py", line 53, in handle_request
    async for msg in self._agent.on_messages_stream(self._message_buffer, ctx.cancellation_token):
    ...<11 lines>...
            )
  File "/Users/gaganbansal/workspace/.venv-m1/lib/python3.13/site-packages/autogen_agentchat/agents/_base_chat_agent.py", line 98, in on_messages_stream
    response = await self.on_messages(messages, cancellation_token)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/gaganbansal/workspace/.venv-m1/lib/python3.13/site-packages/autogen_agentchat/agents/_user_proxy_agent.py", line 180, in on_messages
    raise RuntimeError(f"Failed to get user input: {str(e)}") from e
RuntimeError: Failed to get user input: Failed to get user input:
--- Logging error ---
```

### What did you expect to happen?

No errors

### How can we reproduce it (as minimally and precisely as possible)?

- I test on a Mac, iTerm
- install magentic-one-cli with uv/pip
- exported an openai key

### AutoGen version

v0.4

### Which package was this bug in

Core

### Model used

_No response_

### Python version

_No response_

### Operating system

_No response_

### Any additional info you think would be helpful for fixing this bug

_No response_

Hints from issue discussion:
@peterychang , any insights on this?

Commit message: undo changes

Patch to review:
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_user_proxy_agent.py b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_user_proxy_agent.py
index 2ad9a24682f0..de833df0fd3b 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_user_proxy_agent.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/agents/_user_proxy_agent.py
@@ -1,19 +1,17 @@
 import asyncio
+import uuid
+from contextlib import contextmanager
+from contextvars import ContextVar
 from inspect import iscoroutinefunction
-from typing import Awaitable, Callable, Optional, Sequence, Union, cast
+from typing import Any, AsyncGenerator, Awaitable, Callable, ClassVar, Generator, Optional, Sequence, Union, cast
 
 from aioconsole import ainput  # type: ignore
 from autogen_core import CancellationToken
 
 from ..base import Response
-from ..messages import ChatMessage, HandoffMessage, TextMessage
+from ..messages import AgentEvent, ChatMessage, HandoffMessage, TextMessage, UserInputRequestedEvent
 from ._base_chat_agent import BaseChatAgent
 
-# Define input function types more precisely
-SyncInputFunc = Callable[[str], str]
-AsyncInputFunc = Callable[[str, Optional[CancellationToken]], Awaitable[str]]
-InputFuncType = Union[SyncInputFunc, AsyncInputFunc]
-
 
 # TODO: ainput doesn't seem to play nicely with jupyter.
 #       No input window appears in this case.
@@ -109,6 +107,39 @@ async def cancellable_user_agent():
                     print(f"BaseException: {e}")
     """
 
+    # Define input function types more precisely
+    SyncInputFunc = Callable[[str], str]
+    AsyncInputFunc = Callable[[str, Optional[CancellationToken]], Awaitable[str]]
+    InputFuncType = Union[SyncInputFunc, AsyncInputFunc]
+    DEFAULT_INPUT_FUNC: ClassVar[InputFuncType] = cancellable_input
+
+    class InputRequestContext:
+        def __init__(self) -> None:
+            raise RuntimeError(
+                "InputRequestContext cannot be instantiated. It is a static class that provides context management for user input requests."
+            )
+
+        _INPUT_REQUEST_CONTEXT_VAR: ClassVar[ContextVar[str]] = ContextVar("_INPUT_REQUEST_CONTEXT_VAR")
+
+        @classmethod
+        @contextmanager
+        def populate_context(cls, ctx: str) -> Generator[None, Any, None]:
+            """:meta private:"""
+            token = UserProxyAgent.InputRequestContext._INPUT_REQUEST_CONTEXT_VAR.set(ctx)
+            try:
+                yield
+            finally:
+                UserProxyAgent.InputRequestContext._INPUT_REQUEST_CONTEXT_VAR.reset(token)
+
+        @classmethod
+        def request_id(cls) -> str:
+            try:
+                return cls._INPUT_REQUEST_CONTEXT_VAR.get()
+            except LookupError as e:
+                raise RuntimeError(
+                    "InputRequestContext.runtime() must be called within the input callback of a UserProxyAgent."
+                ) from e
+
     def __init__(
         self,
         name: str,
@@ -140,11 +171,11 @@ async def _get_input(self, prompt: str, cancellation_token: Optional[Cancellatio
         try:
             if self._is_async:
                 # Cast to AsyncInputFunc for proper typing
-                async_func = cast(AsyncInputFunc, self.input_func)
+                async_func = cast(UserProxyAgent.AsyncInputFunc, self.input_func)
                 return await async_func(prompt, cancellation_token)
             else:
                 # Cast to SyncInputFunc for proper typing
-                sync_func = cast(SyncInputFunc, self.input_func)
+                sync_func = cast(UserProxyAgent.SyncInputFunc, self.input_func)
                 loop = asyncio.get_event_loop()
                 return await loop.run_in_executor(None, sync_func, prompt)
 
@@ -153,9 +184,15 @@ async def _get_input(self, prompt: str, cancellation_token: Optional[Cancellatio
         except Exception as e:
             raise RuntimeError(f"Failed to get user input: {str(e)}") from e
 
-    async def on_messages(
-        self, messages: Sequence[ChatMessage], cancellation_token: Optional[CancellationToken] = None
-    ) -> Response:
+    async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
+        async for message in self.on_messages_stream(messages, cancellation_token):
+            if isinstance(message, Response):
+                return message
+        raise AssertionError("The stream should have returned the final result.")
+
+    async def on_messages_stream(
+        self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken
+    ) -> AsyncGenerator[AgentEvent | ChatMessage | Response, None]:
         """Handle incoming messages by requesting user input."""
         try:
             # Check for handoff first
@@ -164,15 +201,18 @@ async def on_messages(
                 f"Handoff received from {handoff.source}. Enter your response: " if handoff else "Enter your response: "
             )
 
-            user_input = await self._get_input(prompt, cancellation_token)
+            request_id = str(uuid.uuid4())
+
+            input_requested_event = UserInputRequestedEvent(request_id=request_id, source=self.name)
+            yield input_requested_event
+            with UserProxyAgent.InputRequestContext.populate_context(request_id):
+                user_input = await self._get_input(prompt, cancellation_token)
 
             # Return appropriate message type based on handoff presence
             if handoff:
-                return Response(
-                    chat_message=HandoffMessage(content=user_input, target=handoff.source, source=self.name)
-                )
+                yield Response(chat_message=HandoffMessage(content=user_input, target=handoff.source, source=self.name))
             else:
-                return Response(chat_message=TextMessage(content=user_input, source=self.name))
+                yield Response(chat_message=TextMessage(content=user_input, source=self.name))
 
         except asyncio.CancelledError:
             raise
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py b/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py
index 07fc3123eb4c..21fb32d9d584 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/messages.py
@@ -103,25 +103,40 @@ class ToolCallSummaryMessage(BaseChatMessage):
     type: Literal["ToolCallSummaryMessage"] = "ToolCallSummaryMessage"
 
 
+class UserInputRequestedEvent(BaseAgentEvent):
+    """An event signaling a that the user proxy has requested user input. Published prior to invoking the input callback."""
+
+    request_id: str
+    """Identifier for the user input request."""
+
+    content: Literal[""] = ""
+    """Empty content for compat with consumers expecting a content field."""
+
+    type: Literal["UserInputRequestedEvent"] = "UserInputRequestedEvent"
+
+
 ChatMessage = Annotated[
     TextMessage | MultiModalMessage | StopMessage | ToolCallSummaryMessage | HandoffMessage, Field(discriminator="type")
 ]
 """Messages for agent-to-agent communication only."""
 
 
-AgentEvent = Annotated[ToolCallRequestEvent | ToolCallExecutionEvent, Field(discriminator="type")]
+AgentEvent = Annotated[
+    ToolCallRequestEvent | ToolCallExecutionEvent | UserInputRequestedEvent, Field(discriminator="type")
+]
 """Events emitted by agents and teams when they work, not used for agent-to-agent communication."""
 
 
 __all__ = [
+    "AgentEvent",
     "BaseMessage",
-    "TextMessage",
+    "ChatMessage",
+    "HandoffMessage",
     "MultiModalMessage",
     "StopMessage",
-    "HandoffMessage",
-    "ToolCallRequestEvent",
+    "TextMessage",
     "ToolCallExecutionEvent",
+    "ToolCallRequestEvent",
     "ToolCallSummaryMessage",
-    "ChatMessage",
-    "AgentEvent",
+    "UserInputRequestedEvent",
 ]
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/ui/__init__.py b/python/packages/autogen-agentchat/src/autogen_agentchat/ui/__init__.py
index 65c4f1e07ad9..9cc0837c58c2 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/ui/__init__.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/ui/__init__.py
@@ -2,6 +2,6 @@
 This module implements utility classes for formatting/printing agent messages.
 """
 
-from ._console import Console
+from ._console import Console, UserInputManager
 
-__all__ = ["Console"]
+__all__ = ["Console", "UserInputManager"]
diff --git a/python/packages/autogen-agentchat/src/autogen_agentchat/ui/_console.py b/python/packages/autogen-agentchat/src/autogen_agentchat/ui/_console.py
index 79d39d6add7f..443ed1490cb0 100644
--- a/python/packages/autogen-agentchat/src/autogen_agentchat/ui/_console.py
+++ b/python/packages/autogen-agentchat/src/autogen_agentchat/ui/_console.py
@@ -1,14 +1,18 @@
+import asyncio
 import os
 import sys
 import time
-from typing import AsyncGenerator, List, Optional, TypeVar, cast
+from inspect import iscoroutinefunction
+from typing import AsyncGenerator, Dict, List, Optional, TypeVar, cast
 
 from aioconsole import aprint  # type: ignore
 from autogen_core import Image
+from autogen_core._cancellation_token import CancellationToken
 from autogen_core.models import RequestUsage
 
+from autogen_agentchat.agents._user_proxy_agent import UserProxyAgent
 from autogen_agentchat.base import Response, TaskResult
-from autogen_agentchat.messages import AgentEvent, ChatMessage, MultiModalMessage
+from autogen_agentchat.messages import AgentEvent, ChatMessage, MultiModalMessage, UserInputRequestedEvent
 
 
 def _is_running_in_iterm() -> bool:
@@ -22,11 +26,53 @@ def _is_output_a_tty() -> bool:
 T = TypeVar("T", bound=TaskResult | Response)
 
 
+class UserInputManager:
+    def __init__(self, callback: UserProxyAgent.InputFuncType = UserProxyAgent.DEFAULT_INPUT_FUNC):
+        self.input_events: Dict[str, asyncio.Event] = {}
+        self.callback = callback
+
+    def get_wrapped_callback(self) -> UserProxyAgent.AsyncInputFunc:
+        async def user_input_func_wrapper(prompt: str, cancellation_token: Optional[CancellationToken]) -> str:
+            # Lookup the event for the prompt, if it exists wait for it.
+            # If it doesn't exist, create it and store it.
+            # Get request ID:
+            request_id = UserProxyAgent.InputRequestContext.request_id()
+            if request_id in self.input_events:
+                event = self.input_events[request_id]
+            else:
+                event = asyncio.Event()
+                self.input_events[request_id] = event
+
+            await event.wait()
+
+            del self.input_events[request_id]
+
+            if iscoroutinefunction(self.callback):
+                # Cast to AsyncInputFunc for proper typing
+                async_func = cast(UserProxyAgent.AsyncInputFunc, self.callback)
+                return await async_func(prompt, cancellation_token)
+            else:
+                # Cast to SyncInputFunc for proper typing
+                sync_func = cast(UserProxyAgent.SyncInputFunc, self.callback)
+                loop = asyncio.get_event_loop()
+                return await loop.run_in_executor(None, sync_func, prompt)
+
+        return user_input_func_wrapper
+
+    def notify_event_received(self, request_id: str) -> None:
+        if request_id in self.input_events:
+            self.input_events[request_id].set()
+        else:
+            event = asyncio.Event()
+            self.input_events[request_id] = event
+
+
 async def Console(
     stream: AsyncGenerator[AgentEvent | ChatMessage | T, None],
     *,
     no_inline_images: bool = False,
     output_stats: bool = True,
+    user_input_manager: UserInputManager | None = None,
 ) -> T:
     """
     Consumes the message stream from :meth:`~autogen_agentchat.base.TaskRunner.run_stream`
@@ -62,6 +108,7 @@ async def Console(
                     f"Duration: {duration:.2f} seconds\n"
                 )
                 await aprint(output, end="")
+
             # mypy ignore
             last_processed = message  # type: ignore
 
@@ -91,9 +138,13 @@ async def Console(
                     f"Duration: {duration:.2f} seconds\n"
                 )
                 await aprint(output, end="")
+
             # mypy ignore
             last_processed = message  # type: ignore
-
+        # We don't want to print UserInputRequestedEvent messages, we just use them to signal the user input event.
+        elif isinstance(message, UserInputRequestedEvent):
+            if user_input_manager is not None:
+                user_input_manager.notify_event_received(message.request_id)
         else:
             # Cast required for mypy to be happy
             message = cast(AgentEvent | ChatMessage, message)  # type: ignore
diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one.py b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one.py
index 23aca97014c3..904c486426c0 100644
--- a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one.py
+++ b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one.py
@@ -116,7 +116,12 @@ async def example_usage_hil():
 
     """
 
-    def __init__(self, client: ChatCompletionClient, hil_mode: bool = False):
+    def __init__(
+        self,
+        client: ChatCompletionClient,
+        hil_mode: bool = False,
+        input_func: UserProxyAgent.InputFuncType | None = None,
+    ):
         self.client = client
         self._validate_client_capabilities(client)
 
@@ -126,7 +131,7 @@ def __init__(self, client: ChatCompletionClient, hil_mode: bool = False):
         executor = CodeExecutorAgent("Executor", code_executor=LocalCommandLineCodeExecutor())
         agents: List[ChatAgent] = [fs, ws, coder, executor]
         if hil_mode:
-            user_proxy = UserProxyAgent("User")
+            user_proxy = UserProxyAgent("User", input_func=input_func)
             agents.append(user_proxy)
         super().__init__(agents, model_client=client)
 
diff --git a/python/packages/magentic-one-cli/src/magentic_one_cli/_m1.py b/python/packages/magentic-one-cli/src/magentic_one_cli/_m1.py
index e5a07b164939..8efdf28ec268 100644
--- a/python/packages/magentic-one-cli/src/magentic_one_cli/_m1.py
+++ b/python/packages/magentic-one-cli/src/magentic_one_cli/_m1.py
@@ -3,6 +3,7 @@
 import warnings
 
 from autogen_agentchat.ui import Console
+from autogen_agentchat.ui._console import UserInputManager
 from autogen_ext.models.openai import OpenAIChatCompletionClient
 from autogen_ext.teams.magentic_one import MagenticOne
 
@@ -37,9 +38,10 @@ def main() -> None:
     args = parser.parse_args()
 
     async def run_task(task: str, hil_mode: bool) -> None:
+        input_manager = UserInputManager()
         client = OpenAIChatCompletionClient(model="gpt-4o")
-        m1 = MagenticOne(client=client, hil_mode=hil_mode)
-        await Console(m1.run_stream(task=task), output_stats=False)
+        m1 = MagenticOne(client=client, hil_mode=hil_mode, input_func=input_manager.get_wrapped_callback())
+        await Console(m1.run_stream(task=task), output_stats=False, user_input_manager=input_manager)
 
     task = args.task[0]
     asyncio.run(run_task(task, not args.no_hil))


Code review comments:
diff_hunk: @@ -3,6 +3,7 @@
 import warnings
 
 from autogen_agentchat.ui import Console
+from autogen_agentchat.ui._console import UserInputManager
path: python/packages/magentic-one-cli/src/magentic_one_cli/_m1.py
text: @user1:
```suggestion
from autogen_agentchat.ui import UserInputManager
```

diff_hunk: @@ -22,11 +26,53 @@ def _is_output_a_tty() -> bool:
 T = TypeVar("T", bound=TaskResult | Response)
 
 
+class UserInputManager:
+    def __init__(self, callback: UserProxyAgent.InputFuncType = UserProxyAgent.DEFAULT_INPUT_FUNC):
path: python/packages/autogen-agentchat/src/autogen_agentchat/ui/_console.py
text: @user1:
perhaps the `InputFuncType` should be outside of the `UserProxyAgent` in a common place, e.g., `autogen_agentchat.base` or something like that? It's a bit strange that the `ui`  submodule is depending on the `UserProxyAgent`.