hordelab/hub
2
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: task cancellation (wip)

Browse Source
This commit is contained in:
Shivering Ocelot committed 2026-05-30 17:51:49 -05:00
1 parent 25b5fbd22a
commit 77addb889d
10 files changed
+319 -60
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

backend/src/agent.py
+35
View File
@@ -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
+10 -19
View File
@@ -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
+8 -1
View File
@@ -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
+52 -15
View File
@@ -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
+2 -22
View File
@@ -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))
backend/test_cancellation.py
+74
View File
@@ -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
+8
View File
@@ -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
+3 -3
View File
@@ -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
+78
View File
@@ -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
+49
View File
@@ -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())