WIP Stash

backend/src/agent.py

@@ -2,6 +2,7 @@ import re
 import inspect
 import time
 import asyncio
+import traceback
 
 from temporary import generate_tmp_id
 from stream import Stream
@@ -66,7 +67,18 @@ class Agent:
             },
         )
         await self.run(run)
-        return run.get_output_stream()
+        out_stream = Stream(self.app)
+
+        def on_output_event(stream):
+            def on_chunk(chunk):
+                out_stream.feed(chunk)
+            stream.listen(on_chunk, lambda: None)
+        def on_done(data):
+            out_stream.close()
+        run.on('output', on_output_event)
+        run.on('done', on_done)
+
+        return out_stream
 
     async def received_message(self, message):
         chat_view = await self.app.get_chat_view(message.chat_id)
@@ -213,7 +225,8 @@ class AgentRun:
                 cb(data)
             except Exception as e:
                 print(f"ERROR: handling event '{event_name}' in run {self.id}\n" + str(e))
-        
+                print(traceback.format_exc())
+
         # Broadcast event to websocket listeners
         self.app.emit_event("run-event", {
             "runId": self.id,
@@ -221,24 +234,6 @@ class AgentRun:
             "data": data
         })
     
-    def get_output_stream(self):
-        if not self.output:
-            self.output = Stream(self.app)
-            self.emit_event('output', self.output)
-        return self.output
-    
-    def feed_output(self, chunk):
-        stream = self.get_output_stream()
-        if chunk:
-            self.output.feed(chunk)
-    
-    def feed_reasoning(self, chunk):
-        if not self.reasoning:
-            self.reasoning = Stream(self.app)
-            self.emit_event('reasoning', self.reasoning)
-        if chunk:
-            self.reasoning.feed(chunk)
-    
     def add_tools(self, tools):
         if "tools" in self.input:
             self.input["tools"] += tools
@@ -298,7 +293,7 @@ class AgentRun:
         self.close()
         
         # Emit cancellation event
-        self.emit_event("cancelled", {"runId": self.id, "agentId": self.agent_id})
+        self.emit_event("cancelled", {})
         
         # Broadcast cancellation event
         self.app.emit_event("run-cancelled", {
@@ -314,4 +309,4 @@ class AgentRun:
         self.end_time = int(time.time())
         if self.output: self.output.close()
         if self.reasoning: self.reasoning.close()
-        self.emit_event("done", {"output": self.output.content, "agent_id": self.agent_id})
+        self.emit_event("done", {})