feat: task cancellation (wip)

backend/src/agent.py

@@ -137,6 +137,11 @@ class AgentRun:
         self.events = []
         self.event_listeners = {}
         self.output = None
+        
+        # Cancellation tracking
+        self.task = None  # For asyncio tasks
+        self.thread = None  # For threads
+        self.cancelled = False
 
         app.tmp.runs[self.id] = self
         
@@ -216,6 +221,36 @@ class AgentRun:
             "output": self.output,
         }
 
+    def cancel(self):
+        """Cancel the agent run and stop any ongoing inference."""
+        if self.cancelled or self.done:
+            return
+        
+        self.cancelled = True
+        print(f"Cancelling agent run {self.id}")
+        
+        # Cancel asyncio task if exists
+        if self.task and not self.task.done():
+            self.task.cancel()
+        
+        # Thread cannot be forcefully cancelled, but we can signal
+        if self.thread and self.thread.is_alive():
+            # The thread should be checking self.cancelled in its loop
+            print(f"Signalled thread to stop for run {self.id}")
+        
+        # Close the stream and mark as done
+        self.close()
+        
+        # Emit cancellation event
+        self.emit_event("cancelled", {"runId": self.id, "agentId": self.agent_id})
+        
+        # Broadcast cancellation event
+        self.app.emit_event("run-cancelled", {
+            "runId": self.id,
+            "agentId": self.agent_id,
+            "endTime": int(time.time()),
+        })
+
     def close(self):
         if self.done:
             return

backend/src/agents/expert.py

@@ -2,10 +2,8 @@ import inference
 from agent import Agent, AgentField
 
 CLASS_DESCRIPTION = """
-Expert AI in any field, powered by Llama 3 70b.
+Expert AI with customizable system prompt.
 Will get guidelines, documents, and other advanced features in the future.
-
-Available models are: [llama/qwen]-[big/small] (replicate), ollama:MODEL_NAME
 """.strip()
 
 class ExpertAgent(Agent):
@@ -22,7 +20,6 @@ class ExpertAgent(Agent):
         Agent.__init__(self, app, data)
 
     async def run(self, run):
-        inp = run.input or {}
         opts = {}
 
         if self.config.get("temperature") is not None:
@@ -30,26 +27,20 @@ class ExpertAgent(Agent):
         if self.config.get("repetition_penalty") is not None:
             opts["repetition_penalty"] = self.config["repetition_penalty"]
 
-        # 1. Build message list based on input mode
-        if "messages" in inp:
-            messages = inp["messages"]
-        else:
-            instruction = inp.get("instruction")
-            input_text = inp.get("input")
-            messages = inference.build_messages(
-                system_prompt=self.config.get("system_prompt", ""),
-                instruction=instruction,
-                input=input_text,
-                messages=inp.get("messages"),
-            )
+        # 1. Build message list based on input
+        messages = inference.build_messages(
+            system_prompt=self.config.get("system_prompt", ""),
+            instruction=run.input.get("instruction"),
+            input=run.input.get("input"),
+            messages=run.input.get("messages"),
+        )
 
-        # 2. Call inference, passing the run so it feeds run.stream
-        #    and can emit tool_call events through the run.
+        # 2. Call inference, passing the run to feed run.stream and tool calls
         await inference.chat_completion(
             self.app,
             model_id=self.config.get("model"),
             messages=messages,
-            tools=inp.get("tools", []),
+            tools=run.input.get("tools", []),
             opts=opts,
             agent_run=run,
         )

backend/src/app.py

@@ -36,7 +36,6 @@ class App(ChatMixin):
 
         self.replicate_client = replicate_token and replicate.Client(api_token=replicate_token) or None
         self.openrouter_client = openrouter_token and OpenRouter(api_key=openrouter_token) or None
-        
 
         self.agent_classes = [
             agents_mod.ExpertAgent,
@@ -332,6 +331,14 @@ class App(ChatMixin):
         )
         return run.id
 
+    async def cancel_agent_run(self, run_id):
+        """Cancel an ongoing agent run by ID."""
+        run = self.tmp.runs.get(run_id)
+        if run:
+            run.cancel()
+            return True
+        return False
+
     async def get_agent_runs(self, agent_id=None, user_id=None):
         runs = list(self.tmp.runs.values())
         if agent_id:

backend/src/inference.py

@@ -85,7 +85,7 @@ def build_messages(system_prompt="", instruction=None, input=None, messages=None
 
     # Add default messages to list
     if system_prompt:
-        messages.insert(0, LlamaMessage("system", system))
+        messages.insert(0, LlamaMessage("system", system_prompt))
     if prompt:
         messages.append(LlamaMessage("user", prompt))
 
@@ -176,6 +176,11 @@ async def local_chat_completion(app, messages, model, tools=[], opts={}, agent_r
         )
 
         for chunk in llm_output:
+            # Check for cancellation
+            if agent_run and agent_run.cancelled:
+                print(f"Cancelling local inference for run {agent_run.id}")
+                break
+                
             delta = chunk['choices'][0]['delta']
             if 'content' in delta:
                 chunk_text = delta['content']
@@ -195,6 +200,8 @@ async def local_chat_completion(app, messages, model, tools=[], opts={}, agent_r
         content_stream.close()
 
     thread = threading.Thread(target=worker)
+    if agent_run:
+        agent_run.thread = thread
     thread.start()
 
     if agent_run:
@@ -224,16 +231,28 @@ async def ollama_chat_completion(app, messages, model, tools=[], opts={}, agent_
         if formatted_tools:
             kwargs['tools'] = formatted_tools
 
-        stream_response = await ollama.AsyncClient().chat(**kwargs)
+        try:
+            stream_response = await ollama.AsyncClient().chat(**kwargs)
 
-        async for part in stream_response:
-            if part.message.content:
-                content_stream.feed(part.message.content)
+            async for part in stream_response:
+                if agent_run and agent_run.cancelled:
+                    print(f"Cancelling ollama inference for run {agent_run.id}")
+                    break
+                    
+                if part.message.content:
+                    content_stream.feed(part.message.content)
+        except asyncio.CancelledError:
+            print(f"Ollama task cancelled for run {agent_run.id}")
+            raise
 
         content_stream.close()
         print("END OLLAMA INFERENCE")
 
-    asyncio.create_task(worker())
+    task = asyncio.create_task(worker())
+    if agent_run:
+        agent_run.task = task
+    else:
+        task.start()
 
     if agent_run:
         return None
@@ -262,18 +281,29 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, ag
         if formatted_tools:
             kwargs['tools'] = formatted_tools
 
-        stream_response = await app.openrouter_client.chat.send_async(**kwargs)
+        try:
+            stream_response = await app.openrouter_client.chat.send_async(**kwargs)
 
-        async for event in stream_response:
-            if event.choices and len(event.choices) > 0:
-                delta = event.choices[0].delta
-                if delta.content:
-                    content_stream.feed(delta.content)
+            async for event in stream_response:
+                if agent_run and agent_run.cancelled:
+                    print(f"Cancelling openrouter inference for run {agent_run.id}")
+                    # TODO: Need to cancell?
+                    break
+                    
+                if event.choices and len(event.choices) > 0:
+                    delta = event.choices[0].delta
+                    if delta.content:
+                        content_stream.feed(delta.content)
+        except asyncio.CancelledError:
+            print(f"OpenRouter task cancelled for run {agent_run.id if agent_run else 'unknown'}")
+            raise
 
         content_stream.close()
         print("END OPENROUTER INFERENCE")
 
-    asyncio.create_task(worker())
+    task = asyncio.create_task(worker())
+    if agent_run:
+        agent_run.task = task
 
     if agent_run:
         return None
@@ -318,21 +348,28 @@ async def replicate_base_completion(app, model, input, agent_run=None):
             'replicate_client': app.replicate_client,
             'model': model,
             'content_stream': content_stream,
-            'input': input
+            'input': input,
+            'agent_run': agent_run
         }
     )
     thread.start()
     #response = await content_stream.content_wait()
     #return response.lstrip()
     if agent_run:
+        agent_run.thread = thread
         return None
     return content_stream
 
 
-def run_replicate_inference_loop(replicate_client, model, input, content_stream):
+def run_replicate_inference_loop(replicate_client, model, input, content_stream, agent_run=None):
     print("START INFERENCE")
     completion = replicate_client.run(model, input=input, stream=True)
     for event in completion:
+        if agent_run and agent_run.cancelled:
+            print(f"Cancelling replicate inference for run {agent_run.id}")
+            break
+            # TODO: does completion need explicit cancellation?
+            
         if event:
             content_stream.feed(str(event))
     print("END INFERENCE")

backend/src/server.py

@@ -224,22 +224,8 @@ async def start_server(app):
                         'nonce': nonce
                     })
 
-                elif action == "get-agent-run":
-                    await send({
-                        'run': await app.get_agent_run(body['runId']),
-                        'nonce': nonce
-                    })
-                    stream_id = body['streamId']
-                    print("Read Stream" + str(stream_id))
-                    stream = app.streams.get(stream_id)
-                    if stream is None:
-                        await send({'error': 'bad stream id: ' + str(stream_id), 'nonce': nonce})
-                    else:
-                        async def chunk_fn(chunk):
-                            await send({'data': chunk, 'eof': False, 'nonce': nonce})
-                        async def eof_fn():
-                            await send({'eof': True, 'fullData': stream.content, 'nonce': nonce})
-                        stream.listen(chunk_fn, eof_fn)
+                elif action == "cancel-agent-run":
+                    await app.cancel_agent_run(body['runId'])
 
                 elif action == "create-chat":
                     chat_id = await app.create_chat(
@@ -315,12 +301,6 @@ async def start_server(app):
                             response = await response.content_wait()
                     await send({'message': "prompt sent", 'response': response, 'nonce': nonce})
 
-                elif action == "get-agent-runs":
-                    await send({
-                        'runs': await app.get_agent_runs(agent_id=body.get('agentId'), user_id=user_id),
-                        'nonce': nonce
-                    })
-
                 elif action == "read-stream":
                     stream_id = body['streamId']
                     print("Read Stream" + str(stream_id))