diff --git a/api/models.py b/api/models.py
index 110751a..f9d88a3 100644
--- a/api/models.py
+++ b/api/models.py
@@ -100,8 +100,17 @@ class SM:
         type: Literal[SM.Type.GRACEFUL_DISCONNECT]
 
 
+AgentEventUnion = Annotated[
+    Union[
+        SM.EventTextChunk,
+        SM.EventEnd,
+    ],
+    Field(discriminator="subtype")
+]
+
+
 ServerMessage = Annotated[
-    Union[SM.Status, SM.AgentEvent, SM.Error, SM.GracefulDisconnect],
+    Union[SM.Status, AgentEventUnion, SM.Error, SM.GracefulDisconnect],
     Field(discriminator="type")
 ]
 """