[Arch proposal] ENVIRONMENT event source (All-Hands-AI#4584)

Co-authored-by: Xingyao Wang <[email protected]>

openhands/controller/agent_controller.py

@@ -156,7 +156,7 @@ async def report_error(self, message: str, exception: Exception | None = None):
         if exception is not None and isinstance(exception, litellm.AuthenticationError):
             detail = 'Please check your credentials. Is your API key correct?'
         self.event_stream.add_event(
-            ErrorObservation(f'{message}:{detail}'), EventSource.USER
+            ErrorObservation(f'{message}:{detail}'), EventSource.ENVIRONMENT
         )
 
     async def start_step_loop(self):
@@ -346,7 +346,8 @@ async def set_agent_state_to(self, new_state: AgentState):
 
         self.state.agent_state = new_state
         self.event_stream.add_event(
-            AgentStateChangedObservation('', self.state.agent_state), EventSource.AGENT
+            AgentStateChangedObservation('', self.state.agent_state),
+            EventSource.ENVIRONMENT,
         )
 
         if new_state == AgentState.INIT and self.state.resume_state:
@@ -423,7 +424,8 @@ async def _step(self) -> None:
         if self._is_stuck():
             # This need to go BEFORE report_error to sync metrics
             self.event_stream.add_event(
-                FatalErrorObservation('Agent got stuck in a loop'), EventSource.USER
+                FatalErrorObservation('Agent got stuck in a loop'),
+                EventSource.ENVIRONMENT,
             )
             return
 

openhands/core/cli.py

@@ -61,7 +61,7 @@ def display_event(event: Event):
         if hasattr(event, 'thought'):
             display_message(event.thought)
     if isinstance(event, MessageAction):
-        if event.source != EventSource.USER:
+        if event.source == EventSource.AGENT:
             display_message(event.content)
     if isinstance(event, CmdRunAction):
         display_command(event.command)
@@ -131,7 +131,7 @@ async def prompt_for_next_task():
         next_message = input('How can I help? >> ')
         if next_message == 'exit':
             event_stream.add_event(
-                ChangeAgentStateAction(AgentState.STOPPED), EventSource.USER
+                ChangeAgentStateAction(AgentState.STOPPED), EventSource.ENVIRONMENT
             )
             return
         action = MessageAction(content=next_message)

openhands/core/message.py

@@ -49,6 +49,8 @@ def serialize_model(self):
 
 
 class Message(BaseModel):
+    # NOTE: this is not the same as EventSource
+    # These are the roles in the LLM's APIs
     role: Literal['user', 'system', 'assistant', 'tool']
     content: list[TextContent | ImageContent] = Field(default_factory=list)
     cache_enabled: bool = False

openhands/events/event.py

@@ -9,6 +9,7 @@
 class EventSource(str, Enum):
     AGENT = 'agent'
     USER = 'user'
+    ENVIRONMENT = 'environment'
 
 
 @dataclass

openhands/runtime/base.py

@@ -136,6 +136,8 @@ async def on_event(self, event: Event) -> None:
             )
             observation._cause = event.id  # type: ignore[attr-defined]
             observation.tool_call_metadata = event.tool_call_metadata
+
+            # this might be unnecessary, since source should be set by the event stream when we're here
             source = event.source if event.source else EventSource.AGENT
             await self.event_stream.async_add_event(observation, source)  # type: ignore[arg-type]
 

openhands/security/invariant/analyzer.py

@@ -147,6 +147,7 @@ async def confirm(self, event: Event) -> None:
         new_event = action_from_dict(
             {'action': 'change_agent_state', 'args': {'agent_state': 'user_confirmed'}}
         )
+        # we should confirm only on agent actions
         event_source = event.source if event.source else EventSource.AGENT
         await call_sync_from_async(self.event_stream.add_event, new_event, event_source)
 

openhands/server/session/agent_session.py

@@ -118,7 +118,7 @@ async def _start(
             agent_configs=agent_configs,
         )
         self.event_stream.add_event(
-            ChangeAgentStateAction(AgentState.INIT), EventSource.USER
+            ChangeAgentStateAction(AgentState.INIT), EventSource.ENVIRONMENT
         )
         if self.controller:
             self.controller.agent_task = self.controller.start_step_loop()

openhands/server/session/session.py

@@ -73,10 +73,11 @@ async def loop_recv(self):
 
     async def _initialize_agent(self, data: dict):
         self.agent_session.event_stream.add_event(
-            ChangeAgentStateAction(AgentState.LOADING), EventSource.USER
+            ChangeAgentStateAction(AgentState.LOADING), EventSource.ENVIRONMENT
         )
         self.agent_session.event_stream.add_event(
-            AgentStateChangedObservation('', AgentState.LOADING), EventSource.AGENT
+            AgentStateChangedObservation('', AgentState.LOADING),
+            EventSource.ENVIRONMENT,
         )
         # Extract the agent-relevant arguments from the request
         args = {key: value for key, value in data.get('args', {}).items()}
@@ -138,12 +139,19 @@ async def on_event(self, event: Event):
             return
         if event.source == EventSource.AGENT:
             await self.send(event_to_dict(event))
-        elif event.source == EventSource.USER and isinstance(
+        # NOTE: ipython observations are not sent here currently
+        elif event.source == EventSource.ENVIRONMENT and isinstance(
             event, CmdOutputObservation
         ):
-            await self.send(event_to_dict(event))
+            # feedback from the environment to agent actions is understood as agent events by the UI
+            event_dict = event_to_dict(event)
+            event_dict['source'] = EventSource.AGENT
+            await self.send(event_dict)
         elif isinstance(event, ErrorObservation):
-            await self.send(event_to_dict(event))
+            # send error events as agent events to the UI
+            event_dict = event_to_dict(event)
+            event_dict['source'] = EventSource.AGENT
+            await self.send(event_dict)
 
     async def dispatch(self, data: dict):
         action = data.get('action', '')

tests/unit/test_agent_controller.py

@@ -207,7 +207,9 @@ async def on_event(event: Event):
                 'Non fatal error here to trigger loop'
             )
             non_fatal_error_obs._cause = event.id
-            await event_stream.async_add_event(non_fatal_error_obs, EventSource.USER)
+            await event_stream.async_add_event(
+                non_fatal_error_obs, EventSource.ENVIRONMENT
+            )
 
     event_stream.subscribe(EventStreamSubscriber.RUNTIME, on_event)
     runtime.event_stream = event_stream