feat: fix agent prompt data flow, fix tools in openai endpoint, account for completion usage info

backend/src/agent.py

@@ -77,38 +77,39 @@ class Agent:
                 **kwargs,
             },
         )
-        out_stream = Stream(self.app)
 
-        def on_output_event(stream):
+        done_event = asyncio.Event()
+        out_stream = Stream(self.app)
+        tool_calls = []
+        tool_names = None
+        if 'tools' in kwargs:
+            tool_names = [t['function']['name'] for t in kwargs['tools'] if isinstance(t, dict)]
+
         def on_chunk(chunk):
             out_stream.feed(chunk)
-            stream.listen(on_chunk, lambda: None)
+        def on_inference(msg):
-        def on_done(data):
+            kwargs = {'content': msg.content}
-            out_stream.close()
-        run.on('output', on_output_event)
-        run.on('done', on_done)
 
-        # Change the output shape based on the presence of tools
+            if msg.role == 'assistant':
-        if 'tools' in kwargs:
+                msg.content.listen(on_chunk, lambda: None)
-            tool_calls = []
+            elif msg.role == 'tool' and tool_names:
-            tool_names = [t['function']['name'] for t in kwargs['tools'] if isinstance(t, dict)]
+                if msg.content.tool_name in tool_names:
-            def on_tool_call(call):
+                    tool_calls.append(msg.content)
-                print('prompt call', call)
+        def on_done(ev):
-                if call.tool_name in tool_names:
+            done_event.set()
-                    tool_calls.append(call)
+            out_stream.close()
-            run.on('tool-call', on_tool_call)
+
+        run.on('inference', on_inference)
+        run.on('done', on_done)
         await self.run(run)
 
-            # Cannot stream tool calls out, wait for them
+        if tool_names:
-            await out_stream.content_wait()
+            await done_event.wait()
-            print('tool names' + repr(tool_names))
-            print('tool calls' + repr(tool_calls))
             return {
                 'output': out_stream,
                 'tool_calls': tool_calls,
             }
         else:
-            await self.run(run)
             return out_stream
 
 
@@ -157,9 +158,9 @@ class Agent:
                 kwargs['kind'] = "log"
                 kwargs['subkind'] = "tool-call"
                 kwargs['content'] = repr(msg.content)
-            else:
+
             # Unknown message
-                kwargs['kind'] = "log"
+            else: return
 
             new_message = await self.app.send_chat_message(
                 sender=self.id,

backend/src/agents/expert.py

@@ -47,8 +47,9 @@ 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)
-        elif turn_response:
+
-            break
+        # By default, end if there's no calls
+        else: break
 
     run.close()
 

backend/src/inference.py

@@ -11,8 +11,10 @@ from tool import Tool, ToolSpec, ToolCall
 #from llama_cpp import Llama
 from openrouter import OpenRouter
 
+# TODO: Collect usage data
+
 TEST_MOCK = False
-VALID_ROLES = ['user', 'assistant', 'system', 'tool', 'reasoning']
+VALID_ROLES = ['user', 'assistant', 'system', 'tool', 'reasoning', 'usage']
 CONTEXT_WINDOW = 10_000
 
 # Background inference tasks are kept referenced here so they aren't GC'd mid-flight.
@@ -60,14 +62,24 @@ def format_openai(messages):
             formatted.append(msg)
             return msg
 
+    # For debugging
+    _stream_count = 0
+    _badrole_count = 0
+    _badroles = set()
+
     for msg in messages:
+        content = msg.content
+        if isinstance(content, Stream):
+            _stream_count += 1
+            content = content.content
+
         if msg.role == 'tool':
             last = last_assistant_message()
             tool_calls = last.get('tool_calls')
             if tool_calls == None:
                 tool_calls = []
                 last['tool_calls'] = tool_calls
-            tc = msg.content
+            tc = content
             if isinstance(tc, dict):
                 tc = ToolCall(
                     tool_name=tc['tool_name'],
@@ -94,10 +106,19 @@ def format_openai(messages):
         elif msg.role == 'reasoning':
             last = last_assistant_message()
             reasoning = last.get('reasoning', '')
-            last['reasoning'] = reasoning + msg.content
+            last['reasoning'] = reasoning + content
-        else:
+        elif msg.role in ['user', 'assistant', 'system']:
-            fmsg = {'role': msg.role, 'content': msg.content}
+            fmsg = {'role': msg.role, 'content': content}
             formatted.append(fmsg)
+        else:
+            _badrole_count += 1
+            _badroles.add(msg.role)
+
+    _warnings = []
+    if _stream_count: _warnings.append(f"{_stream_count} streams in message list")
+    if _badrole_count: _warnings.append(f"{_badrole_count} messages with bad roles {_badroles}")
+    if _warnings: print(f"[inference] WARNING: " + ", ".join(_warnings))
+
     return formatted
 
 
@@ -350,8 +371,14 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}):
                                     tool_call.tool_name += tc.function.name
                                 if tc.function.arguments:
                                     tool_call.arguments += tc.function.arguments
-                else:
+                if event.usage:
-                    print("[openrouter] no event choices: " + repr(event))
+                    usage = {}
+                    for key in ['completion_tokens', 'prompt_tokens', 'total_tokens', 'cost']:
+                        if hasattr(event.usage, key):
+                            value = getattr(event.usage, key)
+                            if value != None:
+                                usage[key] = value
+                    queue.put_nowait(Message('usage', usage))
 
             if content_stream: content_stream.close()
             if reasoning_stream: reasoning_stream.close()

backend/src/server/openai.py

@@ -1,7 +1,7 @@
 import json, time, traceback, asyncio, uuid
 from quart import Response, request
 from quart_auth import login_required, current_user
-from tool import ToolSpec, ToolParameter
+from tool import ToolSpec, ToolParameter, ToolCall
 from agent import AgentRun
 from utils import encode_json
 
@@ -72,43 +72,43 @@ def attach_run_handlers(run, tool_names):
     queue = asyncio.Queue()
     loop = asyncio.get_event_loop()
 
-    def on_output(stream):
+    async def on_content(chunk):
-        async def on_chunk(chunk):
         await queue.put(('content', chunk))
-        async def on_eof():
+    async def on_reasoning(chunk):
-            # on_done doesn't work?
-            await queue.put(('eof', None))
-        stream.listen(on_chunk, on_eof)
-
-    def on_reasoning(stream):
-        async def on_chunk(chunk):
         await queue.put(('reasoning', chunk))
     async def on_eof(): pass
-        stream.listen(on_chunk, on_eof)
 
-    def on_tool_call(tool_call):
+    def on_inference(msg):
-        deferred = route_tool_call(tool_call, tool_names)
+        # assistant and reasoning are assumed to have a Stream
+        # tool is assumed to have a ToolCall
+        if msg.role == 'assistant':
+            msg.content.listen(on_content, on_eof)
+        elif msg.role == 'reasoning':
+            msg.content.listen(on_reasoning, on_eof)
+        elif msg.role == 'tool':
+            deferred = route_tool_call(msg.content, tool_names)
             msg_type = 'deferred_call' if deferred else 'tool_call'
-        loop.create_task(queue.put((msg_type, tool_call)))
+            loop.create_task(queue.put((msg_type, msg.content)))
+        elif msg.role == 'usage':
+            loop.create_task(queue.put(('usage', msg.content)))
 
     def on_done(output):
         loop.create_task(queue.put(('eof', None)))
 
-    run.on('output', on_output)
+    run.on('inference', on_inference)
-    run.on('reasoning', on_reasoning)
-    run.on('tool-call', on_tool_call)
     run.on('done', on_done)
 
     return queue
 
 
-def format_chunk(agent_id, delta, finish_reason=None):
+def format_chunk(agent_id, delta, finish_reason=None, usage=None):
     """Build one OpenAI-format streaming chunk body."""
     return {
         "id": "chatcmpl-" + str(time.time()),
         "object": "chat.completion.chunk",
         "created": int(time.time()),
         "model": agent_id,
+        "usage": usage,
         "choices": [{"delta": delta, "index": 0, "finish_reason": finish_reason}]
     }
 
@@ -124,8 +124,8 @@ def register(server, app):
 
         # Extract required fields
         agent_id = data.get("model")
-        openai_messages = data.get("messages", [])
+        openai_messages = data.get("messages") or []
-        openai_tools = data.get("tools", [])
+        openai_tools = data.get("tools") or []
         stream_mode = data.get("stream", False)
 
         if not agent_id or not openai_messages:
@@ -164,20 +164,36 @@ def register(server, app):
                         # message carries the output.
                         call.result = content
                     entry['resolved'] = True
-                messages.append({'role': 'tool', 'content': call.result, 'tool_call_id': call.id})
+                messages.append({'role': 'tool', 'content': call})
                 continue
 
             if role == 'assistant' and msg.get('tool_calls'):
-                # Rebuild the assistant's tool call turn from the cache so
+                # Rebuild the assistant's tool call turn from the request.
-                # unknown or expired calls are silently dropped and the agent
+                # Calls already in the cache (persistent conversation) are
-                # receives the original call objects, which carry their own
+                # reused as-is. Client-declared function calls carry their
-                # output once resolved.
+                # full payload on the wire, so they are reconstructed from
+                # the message when the cache is gone. Server-originated
+                # (deferred) calls only ever sent an id + type, so they can
+                # only be recovered from the cache and are dropped otherwise.
                 calls = []
                 for tc in msg['tool_calls']:
-                    entry = tool_call_cache.get(tc.get('id', ''))
+                    call_id = tc.get('id', '')
-                    if entry is None:
+                    entry = tool_call_cache.get(call_id)
-                        continue
+                    if entry is not None:
                         calls.append(entry['call'])
+                        continue
+
+                    func = tc.get('function')
+                    if tc.get('type') == 'function' and func:
+                        call = ToolCall(
+                            id=call_id,
+                            tool_name=func.get('name', ''),
+                            arguments=json.loads(func.get('arguments') or '{}'),
+                        )
+                        # Seeding the cache lets the following role='tool'
+                        # message resolve the call and attach its result.
+                        cache_tool_call(call, deferred=False)
+                        calls.append(call)
                 messages.append({
                     'role': 'assistant',
                     'content': content if type(content) == str else None,
@@ -239,17 +255,19 @@ def register(server, app):
                 # Never execute tools inside the run: emit them instead so
                 # the handler can send client tools back and defer native
                 # ones.
-                "defer_tool_execution": True,
+                "tool_policy": "ignore",
             },
         )
-        print("run input: " + repr(run.input))
 
         queue = attach_run_handlers(run, tool_names)
         await control.run(run)
 
+        usage = {'completion_tokens': 0, 'prompt_tokens': 0, 'total_tokens': 0, 'cost': 0}
+
         # --- Streaming ---
         if stream_mode:
             async def generate():
+                nonlocal usage
                 try:
                     tool_call_count = 0
                     while True:
@@ -281,6 +299,9 @@ def register(server, app):
                                 "type": "tool",
                             }]}
 
+                        elif msg_type == 'usage':
+                            usage = data
+
                         elif msg_type == 'eof':
                             # Determine finish_reason — was there a tool call?
                             finish_reason = "tool_calls" if any(
@@ -289,11 +310,15 @@ def register(server, app):
                             yield f"data: {json.dumps(format_chunk(agent_id, {}, finish_reason))}\n\n"
                             break
 
-                        if delta:
-                            yield f"data: {json.dumps(format_chunk(agent_id, delta))}\n\n"
                         else:
                             raise Exception(f"Unrecognized msg_type: {msg_type}. {repr(data)}")
 
+                        if delta:
+                            yield f"data: {json.dumps(format_chunk(agent_id, delta))}\n\n"
+
+                    # Usage
+                    yield f"data: {json.dumps(format_chunk(agent_id, {}, finish_reason, usage))}\n\n"
+
                 except Exception as e:
                     print("[openai] ERROR: " + str(e))
                     print(traceback.format_exc())
@@ -327,6 +352,8 @@ def register(server, app):
                         # Deferred calls stay opaque to the client: only the
                         # call id and type go on the wire.
                         tool_calls.append({"id": data.id, "type": "tool"})
+                    elif msg_type == 'usage':
+                        usage = data
                     else:
                         raise Exception(f"Unrecognized msg_type: {msg_type}. {repr(data)}")
             except Exception as e:
@@ -346,6 +373,7 @@ def register(server, app):
                 "object": "chat.completion",
                 "created": int(time.time()),
                 "model": agent_id,
+                "usage": usage,
                 "choices": [{
                     "message": message,
                     "finish_reason": finish_reason,