Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 78 additions & 0 deletions src/smallestai/atoms/crew/tools/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,11 @@ async def get_weather(location: str):
def __init__(self):
"""Initialize empty registry."""
self._tools: Dict[str, FunctionToolInfo] = {}
# Set by discover() to the owning crew node. When present, tool
# executions emit tool_call_start/end/error events over the node's
# websocket so the platform surfaces them on the call's Events tab
# (the same tool-call events single-prompt agents already show).
self._owner: Any = None

def register(self, func_or_info: Union[Callable, FunctionToolInfo]):
"""
Expand Down Expand Up @@ -105,6 +110,7 @@ async def tool2(self): pass
agent = MyAgent()
registry.discover(agent) # Registers both tools
"""
self._owner = obj
tools = find_function_tools(obj)
for tool in tools:
self.register(tool)
Expand Down Expand Up @@ -184,15 +190,70 @@ async def _execute_sequential(self, tool_calls: List[ToolCall], context: Optiona
results.append(result)
return results

async def _emit_tool_event(
self,
name: str,
call: ToolCall,
arguments: Any,
*,
response: Any = None,
error: Optional[str] = None,
latency_ms: int = 0,
success: bool = True,
) -> None:
"""Emit a tool_call_start/end/error event over the owning node's websocket.

No-op when the registry isn't owned by a crew node (standalone use) or
the node can't send events. Observability must never break tool
execution, so all failures here are swallowed.
"""
send = getattr(self._owner, "send_event", None)
if send is None or not asyncio.iscoroutinefunction(send):
return

context: Dict[str, Any] = {"arguments": arguments}
if response is not None:
try:
json.dumps(response)
context["response"] = response
except (TypeError, ValueError):
context["response"] = str(response)

payload: Dict[str, Any] = {
"turn_id": "",
"tool_call_id": call.id,
"function_name": call.name,
"context": context,
}
if name == "tool_call_end":
payload["latency"] = latency_ms
payload["success"] = success
elif name == "tool_call_error":
payload["error"] = error
payload["success"] = False

try:
from smallestai.atoms.crew.events import SDKAgentLogEvent

await send(SDKAgentLogEvent(name=name, payload=payload))
except Exception as exc: # never let telemetry break a tool
logger.debug(f"tool-call event emit skipped ({name}): {exc}")

async def _execute_single(self, call: ToolCall, context: Optional[Any]) -> ToolResult:
"""Execute a single tool call."""
import time

started = time.monotonic()
arguments: Any = {}
try:
tool_info = self._tools.get(call.name)
if not tool_info:
raise ValueError(f"Unknown tool: {call.name}")

arguments = json.loads(call.arguments)

await self._emit_tool_event("tool_call_start", call, arguments)

func = tool_info.function
args, kwargs = self._prepare_arguments(func, arguments, context)

Expand All @@ -210,6 +271,15 @@ async def _execute_single(self, call: ToolCall, context: Optional[Any]) -> ToolR

logger.debug(f"Tool {call.name} completed successfully")

await self._emit_tool_event(
"tool_call_end",
call,
arguments,
response=result,
latency_ms=int((time.monotonic() - started) * 1000),
success=True,
)

return ToolResult(
tool_call_id=call.id,
name=call.name,
Expand All @@ -219,6 +289,14 @@ async def _execute_single(self, call: ToolCall, context: Optional[Any]) -> ToolR

except Exception as e:
logger.exception(f"Error executing tool {call.name}: {e}")
await self._emit_tool_event(
"tool_call_error",
call,
arguments,
error=str(e),
latency_ms=int((time.monotonic() - started) * 1000),
success=False,
)
return ToolResult(
tool_call_id=call.id,
name=call.name,
Expand Down
79 changes: 79 additions & 0 deletions tests/custom/test_crew_tool_call_events.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
"""Crew tool executions emit tool_call_start/end/error over the node's websocket,
so the platform surfaces them on the Events tab (parity with single-prompt agents).
"""

import asyncio

from smallestai.atoms.crew.clients.types import ToolCall
from smallestai.atoms.crew.events import SDKAgentLogEvent
from smallestai.atoms.crew.tools import ToolRegistry, function_tool


class _Node:
"""Minimal stand-in for a crew node: records events sent over the websocket."""

def __init__(self):
self.sent = []
self.registry = ToolRegistry()
self.registry.discover(self)

async def send_event(self, event):
self.sent.append(event)

@function_tool(name="do_thing")
async def do_thing(self, x: str) -> dict:
return {"ok": True, "echo": x}

@function_tool(name="boom")
async def boom(self) -> None:
raise RuntimeError("kaboom")


def _logs(node):
return [e for e in node.sent if isinstance(e, SDKAgentLogEvent)]


def test_success_emits_start_and_end():
node = _Node()
call = ToolCall(id="c1", name="do_thing", arguments='{"x": "hi"}')
asyncio.run(node.registry.execute([call], parallel=False))
logs = _logs(node)
names = [e.name for e in logs]
assert names == ["tool_call_start", "tool_call_end"]
start, end = logs
assert start.payload["function_name"] == "do_thing"
assert start.payload["tool_call_id"] == "c1"
assert start.payload["context"]["arguments"] == {"x": "hi"}
assert end.payload["success"] is True
assert end.payload["context"]["response"] == {"ok": True, "echo": "hi"}
assert "latency" in end.payload


def test_error_emits_start_and_error():
node = _Node()
call = ToolCall(id="c2", name="boom", arguments="{}")
asyncio.run(node.registry.execute([call], parallel=False))
names = [e.name for e in _logs(node)]
assert names == ["tool_call_start", "tool_call_error"]
err = _logs(node)[-1]
assert err.payload["success"] is False
assert "kaboom" in (err.payload.get("error") or "")


def test_standalone_registry_is_silent():
# No owner node -> no events, no crash.
reg = ToolRegistry()

@function_tool(name="plain")
async def plain(x: str) -> str:
return x

reg.register(plain)
results = asyncio.run(reg.execute([ToolCall(id="c3", name="plain", arguments='{"x": "y"}')]))
assert results[0].content


if __name__ == "__main__":
import pytest

pytest.main([__file__, "-v"])
Loading