feat: inference generates messages, simplify inference run events, fix encoding bug

backend/src/agent.py

@@ -144,13 +144,28 @@ class Agent:
             },
         )
 
-        async def on_output_event(stream):
+        async def on_inference(msg):
-            print("[agent] run output " + repr(stream))
+            kwargs = {'content': msg.content}
+
+            if msg.role == 'assistant':
+                # kind and subkind are default
+                pass
+            elif msg.role == 'reasoning':
+                kwargs['kind'] = "log"
+                kwargs['subkind'] = "reasoning"
+            elif msg.role == 'tool':
+                kwargs['kind'] = "log"
+                kwargs['subkind'] = "tool-call"
+                kwargs['content'] = repr(msg.content)
+            else:
+                # Unknown message
+                kwargs['kind'] = "log"
+
             new_message = await self.app.send_chat_message(
                 sender=self.id,
                 chat_id=message.chat_id,
                 thread_id=message.thread_id,
-                content=stream,
+                **kwargs
             )
             run.emit_event('chat-message', {
                 'messageId': new_message['id'],
@@ -158,41 +173,7 @@ class Agent:
                 'threadId': message.thread_id,
             })
 
-        async def on_reasoning_event(stream):
+        await run.on_async('inference', on_inference)
-            print("[agent] run output " + repr(stream))
-            new_message = await self.app.send_chat_message(
-                sender=self.id,
-                chat_id=message.chat_id,
-                thread_id=message.thread_id,
-                content=stream,
-                kind="log",
-                subkind="reasoning",
-            )
-            run.emit_event('chat-message', {
-                'messageId': new_message['id'],
-                'chatId': message.chat_id,
-                'threadId': message.thread_id,
-            })
-
-        async def on_toolcall_event(stream):
-            print("[agent] run output " + repr(stream))
-            new_message = await self.app.send_chat_message(
-                sender=self.id,
-                chat_id=message.chat_id,
-                thread_id=message.thread_id,
-                content=stream,
-                kind="log",
-                subkind="tool-call",
-            )
-            run.emit_event('chat-message', {
-                'messageId': new_message['id'],
-                'chatId': message.chat_id,
-                'threadId': message.thread_id,
-            })
-
-        await run.on_async('reasoning', on_reasoning_event)
-        await run.on_async('output', on_output_event)
-        await run.on_async('tool-call', on_toolcall_event)
 
         await self.run(run)
         return run

backend/src/agents/expert.py

@@ -20,47 +20,25 @@ async def chat_completion_loop(app, run, messages, model_id=None, tools=None):
     tools = list(tools or [])
     execute_tools = run.input.get('tool_policy', 'auto') == 'auto'
 
-    tool_calls = []
-    output_stream = None
-    reasoning_stream = None
-    inference_end = asyncio.Event()
-
-    def on_output(stream):
-        nonlocal output_stream
-        output_stream = stream
-
-    def on_reasoning(stream):
-        nonlocal reasoning_stream
-        reasoning_stream = stream
-
-    run.on('output', on_output)
-    run.on('reasoning', on_reasoning)
-    run.on('tool-call', tool_calls.append)
-    run.on('inference-end', lambda _: inference_end.set())
-
     while True:
-        # Reset per-iteration state
+        tool_calls = []
-        output_stream = None
+        turn_response = False
-        reasoning_stream = None
+        generator = inference.chat_completion(
-        inference_end.clear()
-
-        await inference.chat_completion(
             app,
             model_id=model_id,
             messages=messages,
             tools=tools,
-            agent_run=run,
+            opts=run.input,
         )
 
-        await inference_end.wait()
+        async for msg in generator:
+            messages.append(msg)
+            if msg.role == 'tool':
+                tool_calls.append(msg.content)
+            if msg.role == 'assistant':
+                turn_response = True
 
-        content = output_stream and await output_stream.content_wait()
+            run.emit_event('inference', msg)
-        reasoning = reasoning_stream and await reasoning_stream.content_wait()
-
-        if content:
-            messages.append(inference.Message('assistant', content))
-        if reasoning:
-            messages.append(inference.Message('reasoning', reasoning))
 
         if execute_tools and tool_calls:
             for tc in tool_calls:
@@ -69,9 +47,7 @@ async def chat_completion_loop(app, run, messages, model_id=None, tools=None):
                     print(f"[agent] unknown tool name: {tc.tool_name}")
                     continue
                 tc.result = await tool.execute_async(**tc.arguments)
-                messages.append(inference.Message('tool', tc))
+        elif turn_response:
-            tool_calls.clear()
-        elif content or not reasoning or not execute_tools:
             break
 
     run.close()

backend/src/inference.py

@@ -4,15 +4,28 @@ import json
 import os
 import traceback
 import random
+import inspect
 from stream import Stream
 from tool import Tool, ToolSpec, ToolCall
 
 #from llama_cpp import Llama
 from openrouter import OpenRouter
 
+TEST_MOCK = False
 VALID_ROLES = ['user', 'assistant', 'system', 'tool', 'reasoning']
 CONTEXT_WINDOW = 10_000
 
+# Background inference tasks are kept referenced here so they aren't GC'd mid-flight.
+_background_tasks = set()
+
+def spawn(coro):
+    task = asyncio.create_task(coro)
+    _background_tasks.add(task)
+    task.add_done_callback(_background_tasks.discard)
+    return task
+
+# These reflect "emitted messages" from any actor, individual turn parts, rather than entire turns.
+# role acts as a type and a sender simultaneously, and content is dynamic
 class Message:
     def __init__(self, role, content):
         if role not in VALID_ROLES:
@@ -21,11 +34,7 @@ class Message:
         self.content = content
 
     def to_dict(self):
-        """Serialize to the dict format expected by llama-cpp / ollama."""
+        return {"role": self.role, "content": self.content}
-        d = {"role": self.role, "content": self.content}
-        if self.reasoning:
-            d["reasoning"] = self.reasoning
-        return d
 
     def __repr__(self):
         return f"Message(role={self.role!r}, content={self.content!r})"
@@ -42,7 +51,6 @@ def format_tools(tools):
     return specs
 
 def format_openai(messages):
-    print('format openai ' + repr(messages))
     formatted = []
     def last_assistant_message():
         if formatted and formatted[-1]['role'] == 'assistant':
@@ -103,9 +111,12 @@ def safe_queue(callback):
     queue = asyncio.Queue()
     loop = asyncio.get_running_loop()
     async def read_loop():
+        nonlocal callback
         while True:
             args = await queue.get()
-            await callback(*args)
+            task = callback(*args)
+            if inspect.isawaitable(task):
+                await task
     loop.create_task(read_loop())
 
     def send(*args):
@@ -173,47 +184,38 @@ async def chat_completion(app,
     messages=None,
     tools=[],
     opts={},
-    agent_run=None
 ):
-    # TESTING OVERRIDE
-    # return await mock_completion(app, agent_run)
     db_model = (await app.get_model(model_id)).value
 
     model_engine = db_model['engine']
     model_name = db_model['model']
 
-    if model_engine == "llama-cpp":
+    print(f"[inference] chat completion: {model_engine} '{model_name}', {len(messages)} messages, {len(tools)} tools")
-        result = await local_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
-    elif model_engine == "openrouter":
-        result = await openrouter_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
-    else:
-        raise Exception("Unsupported model: " + repr(model))
 
-    if agent_run:
+    if TEST_MOCK:
-        return None
+        generator = mock_completion(app)
-    return result
+    elif model_engine == "llama-cpp":
+        generator = local_chat_completion(app, messages, model=model_name, tools=tools, opts=opts)
+    elif model_engine == "openrouter":
+        generator = openrouter_chat_completion(app, messages, model=model_name, tools=tools, opts=opts)
+    else:
+        raise Exception("Unsupported model engine: " + repr(model_engine))
+
+    async for message in generator:
+        yield message
+
+    print(f"[inference] chat completion end: {model_engine} '{model_name}'")
 
 
 # llama.cpp
-async def local_chat_completion(app, messages, model, tools=[], opts={}, agent_run=None):
+async def local_chat_completion(app, messages, model, tools=[], opts={}):
-    if not agent_run:
+    queue = asyncio.Queue()
-        raise Exception("agent_run required")
+    loop = asyncio.get_running_loop()
+    emit = safe_queue(queue.put_nowait)
 
-    content_stream = None
+    async def worker():
-    async def message_loop(name, value=None):
+        stream = None
-        nonlocal content_stream, agent_run
+        try:
-        if name == 'output':
-            if content_stream == None:
-                content_stream = Stream(app)
-                agent_run.emit_event('output', content_stream)
-            content_stream.feed(value)
-        elif name == 'close':
-            if content_stream:
-                content_stream.close()
-            agent_run.emit_event('inference-end', {})
-    send_message = safe_queue(message_loop)
-
-    def worker():
             llm = get_llama_cpp_model(app, model)
 
             kwargs = {}
@@ -233,50 +235,53 @@ async def local_chat_completion(app, messages, model, tools=[], opts={}, agent_r
                 kwargs['tool_choice'] = 'auto'
 
             # With stream=True, returns an iterator that infers tokens on demand (openai compatible)
-        llm_output = llm.create_chat_completion(
+            iterator = await loop.run_in_executor(None, lambda: llm.create_chat_completion(
                 messages=[msg.to_dict() for msg in messages],
                 stream=True,
                 **kwargs
-        )
+            ))
 
-        for chunk in llm_output:
+            sentinel = object()
-            # Check for cancellation
+            while True:
-            if agent_run and agent_run.cancelled:
+                chunk = await loop.run_in_executor(None, next, iterator, sentinel)
-                print(f"Cancelling local inference for run {agent_run.id}")
+                if chunk is sentinel:
                     break
 
                 delta = chunk['choices'][0]['delta']
                 if 'content' in delta:
-                send_message('output', delta['content'])
+                    if stream is None:
-            # Handle tool_call deltas if the model emits them
+                        stream = Stream(app)
+                        emit(Message('assistant', stream))
+                    stream.feed(delta['content'])
                 if 'tool_calls' in delta:
-                # Forward tool call chunks as JSON strings for the consumer
+                    for tc in delta['tool_calls']:
-                send_message("tool-call", {
+                        emit(Message('tool', tc))
-                    "id": tc.get("id"),
+        except Exception:
-                    "function": tc.get("function"),
+            print(traceback.format_exc())
-                })
+        finally:
+            if stream: stream.close()
+            emit(None)
 
-        send_message('close')
+    spawn(worker())
 
-    thread = threading.Thread(target=worker)
+    while True:
-    if agent_run:
+        message = await queue.get()
-        agent_run.thread = thread
+        if message is None:
-    thread.start()
+            break
-
+        yield message
-    if agent_run:
-        return None
-    return content_stream
 
 
-async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, agent_run=None):
+async def openrouter_chat_completion(app, messages, model, tools=[], opts={}):
-    if not agent_run:
+    if not app.openrouter_client:
-        raise Exception("agent_run required")
+        raise Exception("OPENROUTER_API_KEY environment variable is not set")
 
-    # TODO: Separate as InferenceRun class
+    queue = asyncio.Queue()
+
+    async def worker():
         content_stream = None
         reasoning_stream = None
-
         last_channel = None
+
         def open_channel(channel):
             nonlocal last_channel, content_stream, reasoning_stream
             if channel != last_channel:
@@ -286,33 +291,7 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, ag
                 reasoning_stream = None
             last_channel = channel
 
-
+        try:
-    async def message_loop(name, value=None):
-        nonlocal content_stream, reasoning_stream, agent_run
-        open_channel(name) # Hack: use the message as the channel name
-        if name == 'output':
-            if content_stream == None:
-                content_stream = Stream(app)
-                agent_run.emit_event('output', content_stream)
-            content_stream.feed(value)
-        if name == 'reasoning':
-            if reasoning_stream == None:
-                reasoning_stream = Stream(app)
-                agent_run.emit_event('reasoning', reasoning_stream)
-            reasoning_stream.feed(value)
-        if name == 'tool-call':
-            agent_run.emit_event('tool-call', value)
-        elif name == 'close':
-            if content_stream: content_stream.close()
-            if reasoning_stream: reasoning_stream.close()
-            agent_run.emit_event('inference-end', {})
-    emit = safe_queue(message_loop)
-
-    async def worker():
-        print("START OPENROUTER INFERENCE")
-        if not app.openrouter_client:
-            raise Exception("OPENROUTER_API_KEY environment variable is not set")
-
             kwargs = {
                 'model': model,
                 'messages': format_openai(messages),
@@ -326,13 +305,12 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, ag
             formatted_tools = format_tools(tools) if tools else None
             if formatted_tools:
                 kwargs['tools'] = formatted_tools
-        print(f"Tools: {tools} and {formatted_tools}")
 
-        temp = agent_run.input.get("temperature")
+            temp = opts.get("temperature")
             if temp is not None:
                 kwargs['temperature'] = temp
 
-        reasoning = agent_run.input.get("reasoning_effort")
+            reasoning = opts.get("reasoning_effort")
             if reasoning is not None:
                 kwargs['reasoning'] = {"effort": reasoning}
 
@@ -343,19 +321,22 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, ag
             tool_calls = []
 
             async for event in stream_response:
-            if agent_run and agent_run.cancelled:
-                stream_response.cancel()
-                break
-                
                 if event.choices and len(event.choices) > 0:
                     delta = event.choices[0].delta
 
-                if delta.content:
-                    emit('output', delta.content)
                     if delta.reasoning:
-                    emit('reasoning', delta.reasoning)
+                        open_channel('reasoning')
+                        if reasoning_stream is None:
+                            reasoning_stream = Stream(app)
+                            queue.put_nowait(Message('reasoning', reasoning_stream))
+                        reasoning_stream.feed(delta.reasoning)
+                    if delta.content:
+                        open_channel('output')
+                        if content_stream is None:
+                            content_stream = Stream(app)
+                            queue.put_nowait(Message('assistant', content_stream))
+                        content_stream.feed(delta.content)
                     if delta.tool_calls:
-                    print("[openrouter] tool call:" + repr(delta.tool_calls)) 
                         for tc in delta.tool_calls:
                             if tc.index >= len(tool_calls):
                                 tool_calls.append(ToolCall(tool_name="", arguments="", id=""))
@@ -372,33 +353,31 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, ag
                 else:
                     print("[openrouter] no event choices: " + repr(event))
 
+            if content_stream: content_stream.close()
+            if reasoning_stream: reasoning_stream.close()
+
             for tc in tool_calls:
                 try:
-                arguments = json.loads(tc.arguments)
+                    tc.arguments = json.loads(tc.arguments)
-                tc.arguments = arguments
+                    queue.put_nowait(Message('tool', tc))
-                emit('tool-call', tc)
                 except Exception:
                     print("[openrouter] Bad tool call arguments for " + repr(tc))
                     print(traceback.format_exc())
-
+        except Exception:
-        emit('close')
-        print("END OPENROUTER INFERENCE")
-
-    def handle_done(task):
-        try:
-            task.result()
-        except asyncio.CancelledError:
-            print("[openrouter] inference cancelled")
-        except Exception as e:
-            print(f"[openrouter] Worker error: {e}")
             print(traceback.format_exc())
+        finally:
+            queue.put_nowait(None)
 
-    task = asyncio.create_task(worker())
+    spawn(worker())
-    task.add_done_callback(handle_done)
+
-    agent_run.task = task
+    while True:
+        message = await queue.get()
+        if message is None:
+            break
+        yield message
 
 
-async def mock_completion(app, agent_run=None):
+async def mock_completion(app):
     VERBS = ['fix','scan','sync','deploy','parse','render','cache','trace',
         'merge','purge','index','probe','flush','bind','warm','drain']
     OBJECTS = ['nodes','packets','tokens','frames','shards','clusters',
@@ -406,20 +385,28 @@ async def mock_completion(app, agent_run=None):
        'handles','slots','cursors','segments']
     DELAY = 0.3
 
-    content_stream = Stream(app)
+    queue = asyncio.Queue()
-    if agent_run:
+    emit = safe_queue(queue.put_nowait)
-        agent_run.emit_event('output', content_stream)
 
-    async def content_task ():
+    async def worker():
+        stream = Stream(app)
+        emit(Message('assistant', stream))
+        try:
             await asyncio.sleep(DELAY)
-        content_stream.feed(random.choice(VERBS))
+            stream.feed(random.choice(VERBS))
             await asyncio.sleep(DELAY)
-        content_stream.feed(f" {random.randint(1, 99)} ")
+            stream.feed(f" {random.randint(1, 99)} ")
             await asyncio.sleep(DELAY)
-        content_stream.feed(random.choice(VERBS))
+            stream.feed(random.choice(VERBS))
             await asyncio.sleep(DELAY)
-        content_stream.close()
+        finally:
-        agent_run.emit_event('inference-end', {})
+            stream.close()
+            emit(None)
 
-    asyncio.create_task(content_task())
+    spawn(worker())
-    if not agent_run: return content_stream
+
+    while True:
+        message = await queue.get()
+        if message is None:
+            break
+        yield message

backend/src/server/websocket.py

@@ -1,4 +1,4 @@
-import json, asyncio
+import json, asyncio, traceback
 from quart import websocket
 from utils import encode_json
 from stream import Stream      # isinstance check in send-agent-prompt
@@ -40,10 +40,14 @@ def register(server, app, auth_manager):
 
         send_queue = asyncio.Queue()
         async def sender_impl():
+            # This thread MUST NOT crash
             nonlocal send_queue
             while True:
+                try:
                     msg = await send_queue.get()
                     await websocket.send(encode_json(msg))
+                except Exception as err:
+                    print(traceback.format_exc())
         async def send(msg):
             await send_queue.put(msg)
 
@@ -72,6 +76,7 @@ def register(server, app, auth_manager):
                 action = req['action']
                 body = req.get('body', {})
                 nonce = req.get('nonce', None)
+                print(f"[ws] action @{user_id} {action}")
 
                 async def agent_id_from_body (allow_none=False):
                     if 'agentId' in body:

backend/src/utils.py

@@ -9,7 +9,7 @@ def encode_value(value):
     if isinstance(value, ToolSpec):
         return value.to_openai_dict()
     elif isinstance(value, Message):
-        return {'type': value.type, 'content': encode_value(value.content)}
+        return {'role': value.role, 'content': encode_value(value.content)}
     elif isinstance(value, ToolCall):
         json = {
             'id': value.id,