From 77addb889d4b57fcc139137a5966d7a9c36deaee Mon Sep 17 00:00:00 2001 From: Shivering Ocelot Date: Sat, 30 May 2026 17:51:49 -0500 Subject: [PATCH] feat: task cancellation (wip) --- backend/src/agent.py | 35 +++++++++++++++ backend/src/agents/expert.py | 29 +++++------- backend/src/app.py | 9 +++- backend/src/inference.py | 67 +++++++++++++++++++++------- backend/src/server.py | 24 +--------- backend/test_cancellation.py | 74 +++++++++++++++++++++++++++++++ frontend/src/views/RunsView.vue | 8 ++++ readme.md | 6 +-- simple_test_cancellation.py | 78 +++++++++++++++++++++++++++++++++ test_cancellation.py | 49 +++++++++++++++++++++ 10 files changed, 319 insertions(+), 60 deletions(-) create mode 100644 backend/test_cancellation.py create mode 100644 simple_test_cancellation.py create mode 100644 test_cancellation.py diff --git a/backend/src/agent.py b/backend/src/agent.py index abb3387..def1a87 100644 --- a/backend/src/agent.py +++ b/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 diff --git a/backend/src/agents/expert.py b/backend/src/agents/expert.py index 0e7ac03..660c11f 100644 --- a/backend/src/agents/expert.py +++ b/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, ) diff --git a/backend/src/app.py b/backend/src/app.py index f452eb8..d56b433 100644 --- a/backend/src/app.py +++ b/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: diff --git a/backend/src/inference.py b/backend/src/inference.py index f1a7aa1..68f8d7b 100644 --- a/backend/src/inference.py +++ b/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") diff --git a/backend/src/server.py b/backend/src/server.py index e78188d..85d1e12 100644 --- a/backend/src/server.py +++ b/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)) diff --git a/backend/test_cancellation.py b/backend/test_cancellation.py new file mode 100644 index 0000000..3259190 --- /dev/null +++ b/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()) \ No newline at end of file diff --git a/frontend/src/views/RunsView.vue b/frontend/src/views/RunsView.vue index 847d040..b0f3670 100644 --- a/frontend/src/views/RunsView.vue +++ b/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 @@ {{ getStatusBadge(run) }} +
diff --git a/readme.md b/readme.md index daaf485..0a079cb 100644 --- a/readme.md +++ b/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. diff --git a/simple_test_cancellation.py b/simple_test_cancellation.py new file mode 100644 index 0000000..24d5c4d --- /dev/null +++ b/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()) \ No newline at end of file diff --git a/test_cancellation.py b/test_cancellation.py new file mode 100644 index 0000000..4b10107 --- /dev/null +++ b/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()) \ No newline at end of file