From 926bb30e18c5a159123a00d40c7f8561d30fa2b5 Mon Sep 17 00:00:00 2001
From: ariel <ariel.j.kwiatkowski@gmail.com>
Date: Fri, 22 Mar 2024 19:54:34 +0100
Subject: [PATCH] Add handling of ending events

---
 cogment_lab/core.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/cogment_lab/core.py b/cogment_lab/core.py
index f333723..619ed60 100644
--- a/cogment_lab/core.py
+++ b/cogment_lab/core.py
@@ -240,6 +240,10 @@ async def on_message(self, messages: list):
         """Handle received messages."""
         pass
 
+    async def on_ending(self, observation, rendered_frame):
+        """Handle trial ending."""
+        pass
+
     async def end(self):
         """Clean up when done."""
         pass
@@ -250,10 +254,17 @@ async def impl(self, actor_session: ActorSession):
         async for event in actor_session.all_events():
             event: RecvEvent
             self.current_event = event
-            if event.type != cogment.EventType.ACTIVE:
+            if event.type not in (cogment.EventType.ACTIVE, cogment.EventType.ENDING):
                 logging.info(f"Skipping event of type {event.type}")
                 continue
 
+            if event.type == cogment.EventType.ENDING:
+                observation = self.session_helper.get_observation(event)
+                await self.on_ending(observation.value, observation.rendered_frame)
+                continue
+
+            # type = ACTIVE
+
             if event.observation:
                 observation = self.session_helper.get_observation(event)
                 if observation is None: