fix: tool invocation callback host

Author: skele2k

src/lightrace/client.py

@@ -18,6 +18,7 @@
     _current_trace_id,
     _get_tool_registry,
     _set_client_defaults,
+    _set_on_tool_registered,
     _set_otel_exporter,
     _tool_registry,
 )
@@ -58,6 +59,7 @@ def __init__(
         session_id: str | None = None,
         dev_server: bool = True,
         dev_server_port: int = 0,
+        dev_server_host: str | None = None,
     ):
         self._public_key = public_key or os.environ.get("LIGHTRACE_PUBLIC_KEY", "")
         self._secret_key = secret_key or os.environ.get("LIGHTRACE_SECRET_KEY", "")
@@ -69,6 +71,9 @@ def __init__(
         self._dev_server: DevServer | None = None
         self._dev_server_enabled = dev_server
         self._dev_server_port = dev_server_port
+        self._dev_server_host = dev_server_host or os.environ.get(
+            "LIGHTRACE_DEV_SERVER_HOST", "127.0.0.1"
+        )
 
         if not enabled:
             logger.info("Lightrace disabled — no events will be sent")
@@ -120,16 +125,30 @@ def _start_dev_server(self) -> None:
         self._dev_server = DevServer(
             port=self._dev_server_port,
             public_key=self._public_key,
+            callback_host=self._dev_server_host,
         )
+        self._pending_registration: threading.Timer | None = None
         try:
             port = self._dev_server.start()
             logger.info("Dev server listening on http://127.0.0.1:%d", port)
             self._register_tools_http()
+
+            # Re-register when new tools are added after init (debounced)
+            def _on_new_tool(name: str) -> None:
+                logger.debug("New tool registered: %s — scheduling re-registration", name)
+                if self._pending_registration is not None:
+                    self._pending_registration.cancel()
+                timer = threading.Timer(0.2, self._register_tools_http)
+                timer.daemon = True
+                timer.start()
+                self._pending_registration = timer
+
+            _set_on_tool_registered(_on_new_tool)
         except Exception as e:
             logger.error("Failed to start dev server: %s", e)
 
     def _register_tools_http(self) -> None:
-        """Register tool definitions with the Lightrace backend via HTTP (fire-and-forget)."""
+        """Register tool definitions with the Lightrace backend via HTTP with retry."""
         registry = _get_tool_registry()
         if not registry:
             return
@@ -149,19 +168,42 @@ def _register_tools_http(self) -> None:
 
         auth = base64.b64encode(f"{self._public_key}:{self._secret_key}".encode()).decode()
         host = self._host
+        max_retries = 3
 
         def _do_register() -> None:
-            try:
-                resp = httpx.post(
-                    f"{host}/api/public/tools/register",
-                    json={"callbackUrl": callback_url, "tools": tools},
-                    headers={"Authorization": f"Basic {auth}"},
-                    timeout=5.0,
-                )
-                if resp.status_code >= 400:
-                    logger.warning("Tool registration returned %d", resp.status_code)
-            except Exception as e:
-                logger.error("Failed to register tools: %s", e)
+            import time
+
+            for attempt in range(max_retries):
+                try:
+                    resp = httpx.post(
+                        f"{host}/api/public/tools/register",
+                        json={"callbackUrl": callback_url, "tools": tools},
+                        headers={"Authorization": f"Basic {auth}"},
+                        timeout=5.0,
+                    )
+                    if resp.status_code < 400:
+                        logger.info(
+                            "Registered %d tool(s): %s",
+                            len(tools),
+                            ", ".join(t["name"] for t in tools),
+                        )
+                        return
+                    logger.warning(
+                        "Tool registration returned %d (attempt %d/%d)",
+                        resp.status_code,
+                        attempt + 1,
+                        max_retries,
+                    )
+                except Exception as e:
+                    logger.warning(
+                        "Tool registration failed (attempt %d/%d): %s",
+                        attempt + 1,
+                        max_retries,
+                        e,
+                    )
+                if attempt < max_retries - 1:
+                    time.sleep(2**attempt)
+            logger.error("Tool registration failed after %d attempts", max_retries)
 
         threading.Thread(target=_do_register, daemon=True, name="lightrace-register").start()
 
@@ -225,6 +267,10 @@ def flush(self) -> None:
 
     def shutdown(self) -> None:
         """Flush and shut down the client."""
+        _set_on_tool_registered(None)
+        if hasattr(self, "_pending_registration") and self._pending_registration is not None:
+            self._pending_registration.cancel()
+            self._pending_registration = None
         if self._dev_server is not None:
             self._dev_server.stop()
             self._dev_server = None

src/lightrace/dev_server.py

@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 import asyncio
+import inspect
 import logging
 import threading
 import time
@@ -85,16 +86,32 @@ async def invoke(req: InvokeRequest, request: Request) -> JSONResponse:
         input_data = req.input
         start = time.monotonic()
 
+        # Smart dispatch: spread kwargs when input keys match function parameter names
+        spread = False
+        if isinstance(input_data, dict):
+            try:
+                sig = inspect.signature(func)
+                param_names = {
+                    p.name
+                    for p in sig.parameters.values()
+                    if p.kind
+                    not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
+                }
+                # Spread if the input dict keys are a subset of the function's param names
+                spread = len(param_names) > 0 and set(input_data.keys()).issubset(param_names)
+            except (ValueError, TypeError):
+                spread = False
+
         try:
             if asyncio.iscoroutinefunction(func):
-                if isinstance(input_data, dict):
+                if spread:
                     result = await func(**input_data)
                 elif input_data is not None:
                     result = await func(input_data)
                 else:
                     result = await func()
             else:
-                if isinstance(input_data, dict):
+                if spread:
                     result = await asyncio.to_thread(func, **input_data)
                 elif input_data is not None:
                     result = await asyncio.to_thread(func, input_data)
@@ -128,9 +145,10 @@ async def invoke(req: InvokeRequest, request: Request) -> JSONResponse:
 class DevServer:
     """Lightweight HTTP server that accepts tool invocation from the dashboard."""
 
-    def __init__(self, port: int = 0, public_key: str = ""):
+    def __init__(self, port: int = 0, public_key: str = "", callback_host: str = "127.0.0.1"):
         self._port = port
         self._public_key = public_key
+        self._callback_host = callback_host
         self._thread: threading.Thread | None = None
         self._assigned_port: int | None = None
         self._server: Any = None
@@ -205,4 +223,4 @@ def port(self) -> int | None:
     def callback_url(self) -> str | None:
         if self._assigned_port is None:
             return None
-        return f"http://127.0.0.1:{self._assigned_port}"
+        return f"http://{self._callback_host}:{self._assigned_port}"

src/lightrace/trace.py

@@ -28,6 +28,7 @@
 # Global references (set by Client on init)
 _otel_exporter: Any = None  # LightraceOtelExporter instance
 _tool_registry: dict[str, dict[str, Any]] = {}
+_on_tool_registered: Callable[[str], None] | None = None
 
 # Client defaults
 _client_defaults: dict[str, str | None] = {"user_id": None, "session_id": None}
@@ -47,6 +48,11 @@ def _get_tool_registry() -> dict[str, dict[str, Any]]:
     return _tool_registry
 
 
+def _set_on_tool_registered(callback: Callable[[str], None] | None) -> None:
+    global _on_tool_registered
+    _on_tool_registered = callback
+
+
 VALID_TYPES = {None, "span", "generation", "event", "tool", "chain"}
 
 
@@ -86,6 +92,8 @@ def decorator(func: F) -> F:
                 "input_schema": build_json_schema(func),
                 "description": None,
             }
+            if _on_tool_registered is not None:
+                _on_tool_registered(obs_name)
 
         if asyncio.iscoroutinefunction(func):