feat: task cancellation (wip)

backend/src/agent.py

@@ -138,6 +138,11 @@ class AgentRun:
         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
         
         # Emit run start event
@@ -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
+        # 1. Build message list based on input
-        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,
+            instruction=run.input.get("instruction"),
-                input=input_text,
+            input=run.input.get("input"),
-                messages=inp.get("messages"),
+            messages=run.input.get("messages"),
         )
 
-        # 2. Call inference, passing the run so it feeds run.stream
+        # 2. Call inference, passing the run to feed run.stream and tool calls
-        #    and can emit tool_call events through the run.
         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

@@ -37,7 +37,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
 
+        try:
             stream_response = await ollama.AsyncClient().chat(**kwargs)
 
             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
 
+        try:
             stream_response = await app.openrouter_client.chat.send_async(**kwargs)
 
             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":
+                elif action == "cancel-agent-run":
-                    await send({
+                    await app.cancel_agent_run(body['runId'])
-                        '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 == "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))

backend/test_cancellation.py

@@ -0,0 +1,74 @@
+#!/usr/bin/env python3
+"""
+Simple test script to verify agent run cancellation functionality.
+"""
+
+import sys
+import os
+import time
+
+# Mock the dependencies we don't need
+class MockApp:
+    def __init__(self):
+        self.data_directory = "/tmp/horde_test"
+        self.tmp = MockTmpState()
+        self.event_listeners = []
+    
+    def emit_event(self, name, data=None):
+        pass
+
+class MockTmpState:
+    def __init__(self):
+        self.runs = {}
+
+# Mock the generate_tmp_id function
+def mock_generate_tmp_id(app):
+    return f"tmp-run-{int(time.time())}"
+
+# Monkey patch the temporary module
+import temporary
+temporary.generate_tmp_id = mock_generate_tmp_id
+
+from agent import AgentRun
+
+async def test_cancellation():
+    print("Testing agent run cancellation...")
+    
+    # Create a mock app
+    app = MockApp()
+    
+    # Create an agent run
+    run = AgentRun.create(app, agent_id="test-agent", user_id="test-user")
+    print(f"Created agent run: {run.id}")
+    
+    # Test that the run is not cancelled initially
+    assert not run.cancelled, "Run should not be cancelled initially"
+    print("✓ Run is not cancelled initially")
+    
+    # Test cancellation
+    run.cancel()
+    assert run.cancelled, "Run should be cancelled after calling cancel()"
+    assert run.done, "Run should be marked as done after cancellation"
+    print("✓ Run is cancelled and marked as done")
+    
+    # Test that calling cancel again doesn't cause issues
+    run.cancel()
+    print("✓ Multiple cancel calls work correctly")
+    
+    # Test that the run has the expected attributes
+    assert hasattr(run, 'task'), "Run should have a task attribute"
+    assert hasattr(run, 'thread'), "Run should have a thread attribute"
+    assert hasattr(run, 'cancelled'), "Run should have a cancelled attribute"
+    print("✓ Run has all required cancellation attributes")
+    
+    # Test serialization
+    serialized = run.serialize()
+    assert serialized['id'] == run.id, "Serialized run should have correct ID"
+    assert serialized['done'] == True, "Serialized run should be marked as done"
+    print("✓ Run serialization works correctly")
+    
+    print("All tests passed! Cancellation functionality is working correctly.")
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(test_cancellation())

frontend/src/views/RunsView.vue

@@ -64,6 +64,11 @@
     }
   };
 
+  async function cancelRun (runId) {
+    const res = await backend.sendAction("cancel-agent-run", {runId});
+    console.log(`Cancelled run ${runId}`, res);
+  };
+
   onMounted(async () => {
     try {
       const res = await backend.sendAction("get-agent-runs");
@@ -104,6 +109,9 @@
         <span :class="['run-status', run.done ? 'status-done' : 'status-running']">
           {{ getStatusBadge(run) }}
         </span>
+        <button v-if="!run.done" @click="cancelRun(run.id)">
+          Cancel
+        </button>
       </div>
       <div v-if="expandedRunId === run.id" class="run-details">
         <div class="detail-row">

readme.md

@@ -15,8 +15,8 @@ Now, visit `http://127.0.0.1:5000` in your web browser to see the application in
 
 # Frontend
 
-Frontend uses flutter. For Ubuntu you can use [this guideline](https://ubuntu.com/blog/getting-started-with-flutter-on-ubuntu) to install it.
+Frontend is a Vue 3 Vite application, compiled ahead of time into [/frontend/dist] via `make web`.
 
-## Configuration
+# Data Directory
 
-You can set the backend URL for Android and native builds using the environment variable `BACKEND_URL`. Make sure to set this variable before running the application. Defaults to `null` if nothing is set.
+By default [/local-data] is used as the data directory, where database and user files are stored.

simple_test_cancellation.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+"""
+Simple test script to verify agent run cancellation functionality.
+"""
+
+import sys
+import os
+import time
+import json
+
+# Add the backend directory to the path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'backend'))
+
+# Mock the dependencies we don't need
+class MockApp:
+    def __init__(self):
+        self.data_directory = "/tmp/horde_test"
+        self.tmp = MockTmpState()
+        self.event_listeners = []
+    
+    def emit_event(self, name, data=None):
+        pass
+
+class MockTmpState:
+    def __init__(self):
+        self.runs = {}
+
+# Mock the generate_tmp_id function
+def mock_generate_tmp_id(app):
+    return f"tmp-run-{int(time.time())}"
+
+# Monkey patch the temporary module
+import backend.src.temporary as temp_module
+temp_module.generate_tmp_id = mock_generate_tmp_id
+
+from backend.src.agent import AgentRun
+
+async def test_cancellation():
+    print("Testing agent run cancellation...")
+    
+    # Create a mock app
+    app = MockApp()
+    
+    # Create an agent run
+    run = AgentRun.create(app, agent_id="test-agent", user_id="test-user")
+    print(f"Created agent run: {run.id}")
+    
+    # Test that the run is not cancelled initially
+    assert not run.cancelled, "Run should not be cancelled initially"
+    print("✓ Run is not cancelled initially")
+    
+    # Test cancellation
+    run.cancel()
+    assert run.cancelled, "Run should be cancelled after calling cancel()"
+    assert run.done, "Run should be marked as done after cancellation"
+    print("✓ Run is cancelled and marked as done")
+    
+    # Test that calling cancel again doesn't cause issues
+    run.cancel()
+    print("✓ Multiple cancel calls work correctly")
+    
+    # Test that the run has the expected attributes
+    assert hasattr(run, 'task'), "Run should have a task attribute"
+    assert hasattr(run, 'thread'), "Run should have a thread attribute"
+    assert hasattr(run, 'cancelled'), "Run should have a cancelled attribute"
+    print("✓ Run has all required cancellation attributes")
+    
+    # Test serialization
+    serialized = run.serialize()
+    assert serialized['id'] == run.id, "Serialized run should have correct ID"
+    assert serialized['done'] == True, "Serialized run should be marked as done"
+    print("✓ Run serialization works correctly")
+    
+    print("All tests passed! Cancellation functionality is working correctly.")
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(test_cancellation())

test_cancellation.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+"""
+Simple test script to verify agent run cancellation functionality.
+"""
+
+import asyncio
+import sys
+import os
+
+# Add the backend directory to the path
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'backend'))
+
+from src.app import App
+from src.agent import AgentRun
+
+async def test_cancellation():
+    print("Testing agent run cancellation...")
+    
+    # Create a mock app
+    app = App(data_directory="/tmp/horde_test")
+    
+    # Create an agent run
+    run = AgentRun.create(app, agent_id="test-agent", user_id="test-user")
+    print(f"Created agent run: {run.id}")
+    
+    # Test that the run is not cancelled initially
+    assert not run.cancelled, "Run should not be cancelled initially"
+    print("✓ Run is not cancelled initially")
+    
+    # Test cancellation
+    run.cancel()
+    assert run.cancelled, "Run should be cancelled after calling cancel()"
+    assert run.done, "Run should be marked as done after cancellation"
+    print("✓ Run is cancelled and marked as done")
+    
+    # Test that calling cancel again doesn't cause issues
+    run.cancel()
+    print("✓ Multiple cancel calls work correctly")
+    
+    # Test that the run has the expected attributes
+    assert hasattr(run, 'task'), "Run should have a task attribute"
+    assert hasattr(run, 'thread'), "Run should have a thread attribute"
+    assert hasattr(run, 'cancelled'), "Run should have a cancelled attribute"
+    print("✓ Run has all required cancellation attributes")
+    
+    print("All tests passed! Cancellation functionality is working correctly.")
+
+if __name__ == "__main__":
+    asyncio.run(test_cancellation())