synth-ai 0.2.8.dev4__py3-none-any.whl → 0.2.23.dev3__py3-none-any.whl

synth_ai/tracing_v3/config.py

@@ -1,7 +1,162 @@
-"""Configuration for tracing v3 with Turso/sqld."""
+"""Configuration helpers for tracing v3.
+
+This module centralises the logic for discovering which datastore the tracer
+should use. Historically the project defaulted to a local SQLite file which
+breaks under parallel load. The new resolver inspects environment variables
+and defaults to Turso/libSQL whenever credentials are supplied, while keeping a
+SQLite fallback for contributors without remote access.
+"""
+
+from __future__ import annotations
 
 import os
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
+
+from synth_ai.tracing_v3.constants import canonical_trace_db_path
+
+# STARTUP DIAGNOSTIC - Commented out to reduce noise
+# print(f"[TRACING_V3_CONFIG_LOADED] Python={sys.version_info.major}.{sys.version_info.minor} MODAL_IS_REMOTE={os.getenv('MODAL_IS_REMOTE')}", flush=True)
+
+# ---------------------------------------------------------------------------
+# DSN resolution helpers
+# ---------------------------------------------------------------------------
+
+_CANONICAL_DB_PATH = canonical_trace_db_path()
+_DEFAULT_TRACE_DIR = Path(os.getenv("SYNTH_TRACES_DIR", _CANONICAL_DB_PATH.parent))
+
+
+def _normalise_path(path: Path) -> Path:
+    """Resolve relative paths and expand user/home markers."""
+    path = path.expanduser()
+    if not path.is_absolute():
+        path = (Path.cwd() / path).resolve()
+    return path
+
+
+def _is_modal_environment() -> bool:
+    """Detect if running in Modal container.
+    
+    Modal automatically sets MODAL_IS_REMOTE=1 in all deployed containers.
+    We check this first, then fall back to other Modal env vars.
+    """
+    # Modal sets this in all deployed containers
+    if os.getenv("MODAL_IS_REMOTE") == "1":
+        return True
+    
+    # Additional Modal env vars as fallback
+    return bool(
+        os.getenv("MODAL_TASK_ID")
+        or os.getenv("MODAL_ENVIRONMENT")
+        or os.getenv("SERVICE", "").upper() == "MODAL"
+    )
+
+
+def _split_auth_from_url(url: str) -> tuple[str, str | None]:
+    """Strip any auth_token query parameter from a DSN."""
+    parsed = urlparse(url)
+    if not parsed.query:
+        return url, None
+
+    params = dict(parse_qsl(parsed.query, keep_blank_values=True))
+    token = params.pop("auth_token", None)
+    query = urlencode(params, doseq=True)
+    # urlunparse will omit the '?' automatically when query is empty
+    sanitised = urlunparse(parsed._replace(query=query))
+    return sanitised, token
+
+
+def _default_sqlite_url(*, ensure_dir: bool = True) -> tuple[str, str | None]:
+    """Generate a SQLite URL from SYNTH_TRACES_DIR if set, otherwise raise."""
+    traces_dir = os.getenv("SYNTH_TRACES_DIR")
+    if traces_dir:
+        dir_path = _normalise_path(Path(traces_dir))
+        if ensure_dir:
+            dir_path.mkdir(parents=True, exist_ok=True)
+        db_path = dir_path / "synth_traces.db"
+        sqlite_url = f"sqlite+aiosqlite:///{db_path}"
+        return sqlite_url, None
+    raise RuntimeError("SQLite fallback is disabled; configure LIBSQL_URL or run sqld locally.")
+
+
+def resolve_trace_db_settings(*, ensure_dir: bool = True) -> tuple[str, str | None]:
+    """Resolve the tracing database URL and optional auth token.
+
+    Resolution order:
+      1. `SYNTH_TRACES_DB` (explicit DSN override)
+      2. `LIBSQL_URL` / `TURSO_DATABASE_URL` (remote libSQL endpoints)
+      3. `TURSO_LOCAL_DB_URL` (legacy env for local sqld)
+      4. Modal environment: plain SQLite file (no sqld, no auth)
+      5. Local dev: sqld default
+    """
+    import logging
+    logger = logging.getLogger(__name__)
+
+    explicit = os.getenv("SYNTH_TRACES_DB")
+    if explicit:
+        logger.info(f"[TRACE_CONFIG] Using explicit SYNTH_TRACES_DB: {explicit}")
+        return _split_auth_from_url(explicit)
+
+    remote = os.getenv("LIBSQL_URL") or os.getenv("TURSO_DATABASE_URL")
+    if remote:
+        logger.info(f"[TRACE_CONFIG] Using remote Turso: {remote}")
+        url, token = _split_auth_from_url(remote)
+        if token:
+            return url, token
+        env_token = os.getenv("LIBSQL_AUTH_TOKEN") or os.getenv("TURSO_AUTH_TOKEN")
+        return url, env_token
+
+    local_override = os.getenv("TURSO_LOCAL_DB_URL")
+    if local_override:
+        logger.info(f"[TRACE_CONFIG] Using TURSO_LOCAL_DB_URL: {local_override}")
+        url, token = _split_auth_from_url(local_override)
+        if token:
+            return url, token
+        env_token = os.getenv("LIBSQL_AUTH_TOKEN") or os.getenv("TURSO_AUTH_TOKEN")
+        return url, env_token
+
+    # Check for SYNTH_TRACES_DIR to generate SQLite URL
+    traces_dir = os.getenv("SYNTH_TRACES_DIR")
+    if traces_dir:
+        try:
+            sqlite_url, _ = _default_sqlite_url(ensure_dir=ensure_dir)
+            logger.info(f"[TRACE_CONFIG] Using SQLite from SYNTH_TRACES_DIR: {sqlite_url}")
+            return sqlite_url, None
+        except RuntimeError:
+            pass  # Fall through to other options
+
+    # Modal environment: use plain SQLite file (no sqld daemon, no auth required)
+    is_modal = _is_modal_environment()
+    logger.info(f"[TRACE_CONFIG] Modal detection: {is_modal} (MODAL_IS_REMOTE={os.getenv('MODAL_IS_REMOTE')})")
+    if is_modal:
+        logger.info("[TRACE_CONFIG] Using Modal SQLite: file:/tmp/synth_traces.db")
+        return "file:/tmp/synth_traces.db", None
+
+    # Local dev: default to sqld HTTP API
+    default_url = os.getenv("LIBSQL_DEFAULT_URL", "http://127.0.0.1:8081")
+    logger.info(f"[TRACE_CONFIG] Using local sqld: {default_url}")
+    return default_url, None
+
+
+def resolve_trace_db_url(*, ensure_dir: bool = True) -> str:
+    """Return just the DSN, discarding any auth token."""
+    url, _ = resolve_trace_db_settings(ensure_dir=ensure_dir)
+    return url
+
+
+def resolve_trace_db_auth_token() -> str | None:
+    """Return the resolved auth token for the tracing datastore."""
+    _, token = resolve_trace_db_settings()
+    return token
+
+
+# ---------------------------------------------------------------------------
+# Config dataclasses
+# ---------------------------------------------------------------------------
+
+DEFAULT_DB_FILE = str(_normalise_path(_DEFAULT_TRACE_DIR) / _CANONICAL_DB_PATH.name)
 
 
 @dataclass
@@ -9,25 +164,15 @@ class TursoConfig:
     """Configuration for Turso/sqld connection."""
 
     # Default values matching serve.sh
-    DEFAULT_DB_FILE = "traces/v3/synth_ai.db"
+    DEFAULT_DB_FILE = DEFAULT_DB_FILE
     DEFAULT_HTTP_PORT = 8080
 
-    # Local embedded database for async SQLAlchemy
-    # Resolve to the actual SQLite file used by sqld if the base path is a directory
-    def _resolve_sqlite_db_url() -> str:  # type: ignore[no-redef]
-        base_path = os.path.abspath(os.getenv("SQLD_DB_PATH", "traces/v3/synth_ai.db"))
-        # If sqld is managing this DB, the real SQLite file lives under dbs/default/data
-        candidate = os.path.join(base_path, "dbs", "default", "data")
-        if os.path.isdir(base_path) and os.path.exists(candidate):
-            return f"sqlite+aiosqlite:///{candidate}"
-        return f"sqlite+aiosqlite:///{base_path}"
-
-    # Use env override if provided; otherwise resolve based on SQLD layout
-    db_url: str = os.getenv("TURSO_LOCAL_DB_URL", _resolve_sqlite_db_url())
+    # Resolve DB URL and auth token from environment (libSQL preferred)
+    db_url: str = field(default_factory=resolve_trace_db_url)
 
     # Remote database sync configuration
-    sync_url: str = os.getenv("TURSO_DATABASE_URL", "")
-    auth_token: str = os.getenv("TURSO_AUTH_TOKEN", "")
+    sync_url: str = os.getenv("LIBSQL_SYNC_URL") or os.getenv("TURSO_SYNC_URL", "")
+    auth_token: str = resolve_trace_db_auth_token() or ""
     sync_interval: int = int(
         os.getenv("TURSO_SYNC_SECONDS", "2")
     )  # 2 seconds for responsive local development
@@ -48,20 +193,20 @@ class TursoConfig:
 
     # Daemon settings (for local sqld) - match serve.sh defaults
     sqld_binary: str = os.getenv("SQLD_BINARY", "sqld")
-    sqld_db_path: str = os.getenv("SQLD_DB_PATH", "traces/v3/synth_ai.db")
+    sqld_db_path: str = os.getenv("SQLD_DB_PATH", DEFAULT_DB_FILE)
     sqld_http_port: int = int(os.getenv("SQLD_HTTP_PORT", "8080"))
     sqld_idle_shutdown: int = int(os.getenv("SQLD_IDLE_SHUTDOWN", "0"))  # 0 = no idle shutdown
 
-    def get_connect_args(self) -> dict:
+    def get_connect_args(self) -> dict[str, str]:
         """Get SQLAlchemy connection arguments."""
-        args = {}
+        args: dict[str, str] = {}
         if self.auth_token:
             args["auth_token"] = self.auth_token
         return args
 
-    def get_engine_kwargs(self) -> dict:
+    def get_engine_kwargs(self) -> dict[str, Any]:
         """Get SQLAlchemy engine creation kwargs."""
-        kwargs = {
+        kwargs: dict[str, Any] = {
             "echo": self.echo_sql,
             "future": True,
         }

synth_ai/tracing_v3/constants.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from datetime import datetime
+from pathlib import Path
+
+TRACE_DB_DIR = Path("traces")
+TRACE_DB_BASENAME = "turso_task_app_traces"
+
+
+def canonical_trace_db_name(*, timestamp: datetime | None = None) -> str:
+    """Return the canonical trace database filename (with optional timestamp suffix)."""
+
+    if timestamp is None:
+        return f"{TRACE_DB_BASENAME}.db"
+    return f"{TRACE_DB_BASENAME}_{timestamp.strftime('%Y-%m-%d_%H-%M-%S')}.db"
+
+
+def canonical_trace_db_path(*, timestamp: datetime | None = None) -> Path:
+    """Return the canonical trace database path within the default trace directory."""
+
+    return TRACE_DB_DIR / canonical_trace_db_name(timestamp=timestamp)

synth_ai/tracing_v3/db_config.py

@@ -4,8 +4,11 @@ Centralized database configuration for v3 tracing.
 
 import logging
 import os
+import shutil
 from typing import TYPE_CHECKING, Optional
 
+from synth_ai.tracing_v3.constants import canonical_trace_db_path
+
 if TYPE_CHECKING:
     from .turso.daemon import SqldDaemon
 
@@ -16,7 +19,7 @@ class DatabaseConfig:
     """Centralized database configuration management."""
 
     # Default values from serve.sh
-    DEFAULT_DB_FILE = "traces/v3/synth_ai.db"
+    DEFAULT_DB_FILE = str(canonical_trace_db_path())
     DEFAULT_HTTP_PORT = 8080
 
     def __init__(
@@ -27,11 +30,12 @@ class DatabaseConfig:
 
         Args:
             db_path: Path to database file. If None, uses DEFAULT_DB_FILE from serve.sh.
-            http_port: HTTP port for sqld daemon. If None, uses DEFAULT_HTTP_PORT from serve.sh.
+            http_port: Hrana WebSocket port for sqld daemon (env: SQLD_HTTP_PORT). If None, uses DEFAULT_HTTP_PORT.
             use_sqld: Whether to use sqld daemon or direct SQLite.
         """
-        self.use_sqld = use_sqld
-        self.http_port = http_port or int(os.getenv("SQLD_HTTP_PORT", self.DEFAULT_HTTP_PORT))
+        self.use_sqld = use_sqld and self._sqld_binary_available()
+        # Note: SQLD_HTTP_PORT is actually the hrana port (8080), not the HTTP API port
+        self.hrana_port = http_port or int(os.getenv("SQLD_HTTP_PORT", self.DEFAULT_HTTP_PORT))
         self._daemon: SqldDaemon | None = None
 
         # Set up database path to match serve.sh configuration
@@ -54,21 +58,34 @@ class DatabaseConfig:
         abs_path = os.path.abspath(self.db_file)
         sqld_data_path = os.path.join(abs_path, "dbs", "default", "data")
 
-        if os.path.exists(sqld_data_path):
-            # sqld is managing the database
-            logger.debug(f"✅ Using sqld-managed database at: {sqld_data_path}")
-            actual_db_path = sqld_data_path
-        else:
-            # Direct SQLite file
-            if not os.path.exists(abs_path):
-                logger.debug(f"⚠️  Database file not found at: {abs_path}")
-                logger.debug("🔧 Make sure to run './serve.sh' to start the turso/sqld service")
-            else:
-                logger.debug(f"📁 Using direct SQLite file at: {abs_path}")
-            actual_db_path = abs_path
+        if not os.path.exists(sqld_data_path) and not os.path.exists(abs_path):
+            raise RuntimeError(
+                "sqld data directory not found. Run `sqld --db-path <path>` before using the tracing database."
+            )
 
-        # SQLite URLs need 3 slashes for absolute paths
-        return f"sqlite+aiosqlite:///{actual_db_path}"
+        # Use http:// for local sqld HTTP API port
+        # sqld has two ports: hrana_port (Hrana WebSocket) and hrana_port+1 (HTTP API)
+        # Python libsql client uses HTTP API with http:// URLs
+        http_api_port = self.hrana_port + 1
+        return f"http://127.0.0.1:{http_api_port}"
+
+    def _sqld_binary_available(self) -> bool:
+        """Check if the sqld (Turso) binary is available on PATH."""
+        # Respect explicit SQLD_BINARY override when present
+        binary_override = os.getenv("SQLD_BINARY")
+        candidates = [binary_override, "sqld", "libsql-server"]
+
+        for candidate in candidates:
+            if candidate and shutil.which(candidate):
+                return True
+
+        if binary_override:
+            raise RuntimeError(
+                f"Configured SQLD_BINARY='{binary_override}' but the executable was not found on PATH."
+            )
+        raise RuntimeError(
+            "sqld binary not detected; install Turso's sqld or set SQLD_BINARY so that libSQL can be used."
+        )
 
     def start_daemon(self, wait_time: float = 2.0):
         """
@@ -87,7 +104,7 @@ class DatabaseConfig:
             # Import here to avoid circular dependency
             from .turso.daemon import SqldDaemon
 
-            self._daemon = SqldDaemon(db_path=self.db_base_path, http_port=self.http_port)
+            self._daemon = SqldDaemon(db_path=self.db_base_path, hrana_port=self.hrana_port)
 
         self._daemon.start()
 
@@ -133,11 +150,13 @@ def get_default_db_config() -> DatabaseConfig:
         # Check if sqld is already running (started by serve.sh)
         import subprocess
 
-        sqld_port = int(os.getenv("SQLD_HTTP_PORT", DatabaseConfig.DEFAULT_HTTP_PORT))
+        sqld_hrana_port = int(os.getenv("SQLD_HTTP_PORT", DatabaseConfig.DEFAULT_HTTP_PORT))
+        sqld_http_port = sqld_hrana_port + 1
         sqld_running = False
         try:
+            # Check for either hrana or http port in the process command line
             result = subprocess.run(
-                ["pgrep", "-f", f"sqld.*--http-listen-addr.*:{sqld_port}"],
+                ["pgrep", "-f", f"sqld.*(--hrana-listen-addr.*:{sqld_hrana_port}|--http-listen-addr.*:{sqld_http_port})"],
                 capture_output=True,
                 text=True,
             )
@@ -145,18 +164,12 @@ def get_default_db_config() -> DatabaseConfig:
                 # sqld is already running, don't start a new one
                 sqld_running = True
                 use_sqld = False
-                logger.debug(f"✅ Detected sqld already running on port {sqld_port}")
+                logger.debug(f"✅ Detected sqld already running on ports {sqld_hrana_port} (hrana) and {sqld_http_port} (http)")
         except Exception as e:
             logger.debug(f"Could not check for sqld process: {e}")
 
         if not sqld_running and use_sqld:
-            logger.warning("⚠️  sqld service not detected!")
-            logger.warning("🔧 Please start the turso/sqld service by running:")
-            logger.warning("   ./serve.sh")
-            logger.warning("")
-            logger.warning("This will start:")
-            logger.warning("  - sqld daemon (SQLite server) on port 8080")
-            logger.warning("  - Environment service on port 8901")
+            logger.warning("sqld service not detected. Start the Turso daemon (./serve.sh) before running tracing workloads.")
 
         _default_config = DatabaseConfig(db_path=db_path, use_sqld=use_sqld)
 

synth_ai/tracing_v3/decorators.py

@@ -1,4 +1,3 @@
-from __future__ import annotations
 """Async-aware decorators for tracing v3.
 
 This module provides decorators and context management utilities for the tracing
@@ -23,12 +22,15 @@ The decorators support both sync and async functions where appropriate,
 though async is preferred for consistency with the rest of the system.
 """
 
+from __future__ import annotations
+
 import asyncio
 import contextvars
 import functools
 import time
-from collections.abc import Callable
-from typing import Any, TypeVar
+from collections.abc import Awaitable, Callable, Mapping
+from contextvars import Token
+from typing import Any, TypeVar, cast, overload
 
 from .abstractions import LMCAISEvent, TimeRecord
 from .utils import calculate_cost, detect_provider
@@ -37,13 +39,13 @@ from .utils import calculate_cost, detect_provider
 # These variables automatically propagate across async call boundaries,
 # allowing deeply nested code to access tracing context without explicit passing
 _session_id_ctx: contextvars.ContextVar[str | None] = contextvars.ContextVar(
-    "session_id", default=None
+    "session_id"
 )
 _turn_number_ctx: contextvars.ContextVar[int | None] = contextvars.ContextVar(
-    "turn_number", default=None
+    "turn_number"
 )
 _session_tracer_ctx: contextvars.ContextVar[Any | None] = contextvars.ContextVar(
-    "session_tracer", default=None
+    "session_tracer"
 )
 
 
@@ -91,6 +93,16 @@ def get_session_tracer() -> Any:
 T = TypeVar("T")
 
 
+@overload
+def with_session(require: bool = True) -> Callable[[Callable[..., Awaitable[T]]], Callable[..., Awaitable[T]]]:
+    ...
+
+
+@overload
+def with_session(require: bool = True) -> Callable[[Callable[..., T]], Callable[..., T]]:
+    ...
+
+
 def with_session(require: bool = True):
     """Decorator that ensures a session is active.
 
@@ -112,25 +124,31 @@ def with_session(require: bool = True):
         ```
     """
 
-    def decorator(fn: Callable[..., T]) -> Callable[..., T]:
+    def decorator(fn: Callable[..., Awaitable[T]] | Callable[..., T]) -> Callable[..., Awaitable[T]] | Callable[..., T]:
         if asyncio.iscoroutinefunction(fn):
 
             @functools.wraps(fn)
-            async def async_wrapper(*args, **kwargs):
+            async def async_wrapper(*args: Any, **kwargs: Any) -> T:
                 session_id = get_session_id()
                 if require and session_id is None:
-                    raise RuntimeError(f"No active session for {fn.__name__}")
-                return await fn(*args, **kwargs)
+                    raise RuntimeError(
+                        f"No active session for {getattr(fn, '__name__', 'unknown')}"
+                    )
+                async_fn = cast(Callable[..., Awaitable[T]], fn)
+                return await async_fn(*args, **kwargs)
 
             return async_wrapper
         else:
 
             @functools.wraps(fn)
-            def sync_wrapper(*args, **kwargs):
+            def sync_wrapper(*args: Any, **kwargs: Any) -> T:
                 session_id = get_session_id()
                 if require and session_id is None:
-                    raise RuntimeError(f"No active session for {fn.__name__}")
-                return fn(*args, **kwargs)
+                    raise RuntimeError(
+                        f"No active session for {getattr(fn, '__name__', 'unknown')}"
+                    )
+                sync_fn = cast(Callable[..., T], fn)
+                return sync_fn(*args, **kwargs)
 
             return sync_wrapper
 
@@ -138,7 +156,7 @@ def with_session(require: bool = True):
 
 
 def trace_llm_call(
-    model_name: str = None,
+    model_name: str | None = None,
     system_id: str = "llm",
     extract_tokens: bool = True,
     extract_cost: bool = True,
@@ -171,31 +189,36 @@ def trace_llm_call(
         ```
     """
 
-    def decorator(fn: Callable[..., T]) -> Callable[..., T]:
+    def decorator(fn: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
         if asyncio.iscoroutinefunction(fn):
+            async_fn: Callable[..., Awaitable[T]] = fn
 
             @functools.wraps(fn)
-            async def async_wrapper(*args, **kwargs):
+            async def async_wrapper(*args: Any, **kwargs: Any) -> T:
                 tracer = get_session_tracer()
                 if not tracer:
-                    return await fn(*args, **kwargs)
+                    return await async_fn(*args, **kwargs)
 
                 start_time = time.time()
                 system_state_before = kwargs.get("state_before", {})
 
                 try:
-                    result = await fn(*args, **kwargs)
+                    result = await async_fn(*args, **kwargs)
 
                     # Extract metrics from result - this assumes the result follows
                     # common LLM API response formats (OpenAI, Anthropic, etc.)
-                    if extract_tokens and isinstance(result, dict):
-                        input_tokens = result.get("usage", {}).get("prompt_tokens")
-                        output_tokens = result.get("usage", {}).get("completion_tokens")
-                        total_tokens = result.get("usage", {}).get("total_tokens")
-                        actual_model = result.get("model", model_name)
-                    else:
-                        input_tokens = output_tokens = total_tokens = None
-                        actual_model = model_name
+                    input_tokens = output_tokens = total_tokens = None
+                    actual_model = model_name
+                    if extract_tokens and isinstance(result, Mapping):
+                        result_mapping = cast(Mapping[str, Any], result)
+                        usage = result_mapping.get("usage")
+                        if isinstance(usage, Mapping):
+                            input_tokens = usage.get("prompt_tokens")
+                            output_tokens = usage.get("completion_tokens")
+                            total_tokens = usage.get("total_tokens")
+                        value = result_mapping.get("model")
+                        if isinstance(value, str):
+                            actual_model = value
 
                     latency_ms = int((time.time() - start_time) * 1000)
 
@@ -208,14 +231,16 @@ def trace_llm_call(
                         input_tokens=input_tokens,
                         output_tokens=output_tokens,
                         total_tokens=total_tokens,
-                        cost_usd=calculate_cost(actual_model, input_tokens or 0, output_tokens or 0)
+                        cost_usd=calculate_cost(
+                            actual_model or "unknown", input_tokens or 0, output_tokens or 0
+                        )
                         if extract_cost
                         else None,
                         latency_ms=latency_ms,
                         system_state_before=system_state_before,
                         system_state_after=kwargs.get("state_after", {}),
                         metadata={
-                            "function": fn.__name__,
+                            "function": getattr(fn, "__name__", "unknown"),
                             "step_id": kwargs.get("step_id"),
                         },
                     )
@@ -234,7 +259,7 @@ def trace_llm_call(
                             provider=detect_provider(model_name),
                             latency_ms=int((time.time() - start_time) * 1000),
                             metadata={
-                                "function": fn.__name__,
+                                "function": getattr(fn, "__name__", "unknown"),
                                 "error": str(e),
                                 "error_type": type(e).__name__,
                             },
@@ -249,7 +274,7 @@ def trace_llm_call(
     return decorator
 
 
-def trace_method(event_type: str = "runtime", system_id: str = None):
+def trace_method(event_type: str = "runtime", system_id: str | None = None):
     """Generic method tracing decorator.
 
     Traces any method call by recording it as a RuntimeEvent. Supports both
@@ -269,43 +294,51 @@ def trace_method(event_type: str = "runtime", system_id: str = None):
         ```
     """
 
-    def decorator(fn: Callable[..., T]) -> Callable[..., T]:
+    def decorator(
+        fn: Callable[..., Awaitable[T]] | Callable[..., T]
+    ) -> Callable[..., Awaitable[T]] | Callable[..., T]:
         if asyncio.iscoroutinefunction(fn):
+            async_fn = cast(Callable[..., Awaitable[T]], fn)
 
             @functools.wraps(fn)
-            async def async_wrapper(self, *args, **kwargs):
+            async def async_wrapper(*args: Any, **kwargs: Any) -> T:
                 tracer = get_session_tracer()
                 if not tracer:
-                    return await fn(self, *args, **kwargs)
+                    return await async_fn(*args, **kwargs)
 
                 from .abstractions import RuntimeEvent
 
                 # Use class name as system_id if not provided
-                actual_system_id = system_id or self.__class__.__name__
+                self_obj = args[0] if args else None
+                inferred_system_id = (
+                    self_obj.__class__.__name__ if self_obj is not None else "unknown"
+                )
+                actual_system_id = system_id or inferred_system_id
 
                 event = RuntimeEvent(
                     system_instance_id=actual_system_id,
                     time_record=TimeRecord(event_time=time.time()),
                     actions=[],  # Can be overridden in metadata
                     metadata={
-                        "method": fn.__name__,
+                        "method": getattr(fn, "__name__", "unknown"),
                         "args": str(args)[:100],  # Truncate for safety
                         "step_id": kwargs.get("step_id"),
                     },
                 )
 
                 await tracer.record_event(event)
-                return await fn(self, *args, **kwargs)
+                return await async_fn(*args, **kwargs)
 
             return async_wrapper
         else:
 
             @functools.wraps(fn)
-            def sync_wrapper(self, *args, **kwargs):
+            def sync_wrapper(*args: Any, **kwargs: Any) -> T:
                 # For sync methods, we can't easily trace without blocking
                 # the event loop. This is a limitation of the async-first design.
                 # Consider converting to async or using a different approach
-                return fn(self, *args, **kwargs)
+                sync_fn = cast(Callable[..., T], fn)
+                return sync_fn(*args, **kwargs)
 
             return sync_wrapper
 
@@ -335,11 +368,11 @@ class SessionContext:
         ```
     """
 
-    def __init__(self, session_id: str, tracer=None):
+    def __init__(self, session_id: str, tracer: Any | None = None):
         self.session_id = session_id
         self.tracer = tracer
-        self._token = None
-        self._tracer_token = None
+        self._token: Token[str | None] | None = None
+        self._tracer_token: Token[Any] | None = None
 
     def __enter__(self):
         # Store tokens to restore previous context on exit
@@ -350,8 +383,9 @@ class SessionContext:
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         # Restore previous context - this is crucial for proper isolation
-        _session_id_ctx.reset(self._token)
-        if self._tracer_token:
+        if self._token is not None:
+            _session_id_ctx.reset(self._token)
+        if self._tracer_token is not None:
             _session_tracer_ctx.reset(self._tracer_token)
 
     async def __aenter__(self):
@@ -361,6 +395,7 @@ class SessionContext:
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        _session_id_ctx.reset(self._token)
-        if self._tracer_token:
+        if self._token is not None:
+            _session_id_ctx.reset(self._token)
+        if self._tracer_token is not None:
             _session_tracer_ctx.reset(self._tracer_token)