lionagi 0.14.4__py3-none-any.whl → 0.14.6__py3-none-any.whl

lionagi/libs/concurrency/primitives.py

@@ -1,28 +1,36 @@
-"""Resource management primitives for structured concurrency."""
+"""Resource management primitives for structured concurrency.
 
+Pure async primitives focused on correctness and simplicity.
+"""
+
+import math
 from types import TracebackType
-from typing import Optional
 
 import anyio
 
+from .resource_tracker import track_resource, untrack_resource
 
-class Lock:
-    """A mutex lock for controlling access to a shared resource.
 
-    This lock is reentrant, meaning the same task can acquire it multiple times
-    without deadlocking.
-    """
+class Lock:
+    """A mutex lock for controlling access to a shared resource."""
 
     def __init__(self):
         """Initialize a new lock."""
         self._lock = anyio.Lock()
+        self._acquired = False
+        track_resource(self, f"Lock-{id(self)}", "Lock")
 
-    async def __aenter__(self) -> None:
-        """Acquire the lock.
+    def __del__(self):
+        """Clean up resource tracking when lock is destroyed."""
+        try:
+            untrack_resource(self)
+        except Exception:
+            pass
 
-        If the lock is already held by another task, this will wait until it's released.
-        """
-        await self.acquire()
+    async def __aenter__(self) -> None:
+        """Acquire the lock."""
+        await self._lock.acquire()
+        self._acquired = True
 
     async def __aexit__(
         self,
@@ -31,44 +39,45 @@ class Lock:
         exc_tb: TracebackType | None,
     ) -> None:
         """Release the lock."""
-        self.release()
-
-    async def acquire(self) -> bool:
-        """Acquire the lock.
+        self._lock.release()
+        self._acquired = False
 
-        Returns:
-            True if the lock was acquired, False otherwise.
-        """
+    async def acquire(self) -> None:
+        """Acquire the lock directly."""
         await self._lock.acquire()
-        return True
+        self._acquired = True
 
     def release(self) -> None:
-        """Release the lock.
-
-        Raises:
-            RuntimeError: If the lock is not currently held by this task.
-        """
+        """Release the lock directly."""
+        if not self._acquired:
+            raise RuntimeError(
+                "Attempted to release lock that was not acquired by this task"
+            )
         self._lock.release()
+        self._acquired = False
 
 
 class Semaphore:
-    """A semaphore for limiting concurrent access to a resource."""
+    """A semaphore preventing excessive releases."""
 
     def __init__(self, initial_value: int):
-        """Initialize a new semaphore.
-
-        Args:
-            initial_value: The initial value of the semaphore (must be >= 0)
-        """
+        """Initialize a new semaphore."""
         if initial_value < 0:
             raise ValueError("The initial value must be >= 0")
+        self._initial_value = initial_value
+        self._current_acquisitions = 0
         self._semaphore = anyio.Semaphore(initial_value)
+        track_resource(self, f"Semaphore-{id(self)}", "Semaphore")
 
-    async def __aenter__(self) -> None:
-        """Acquire the semaphore.
+    def __del__(self):
+        """Clean up resource tracking when semaphore is destroyed."""
+        try:
+            untrack_resource(self)
+        except Exception:
+            pass
 
-        If the semaphore value is zero, this will wait until it's released.
-        """
+    async def __aenter__(self) -> None:
+        """Acquire the semaphore."""
         await self.acquire()
 
     async def __aexit__(
@@ -81,35 +90,60 @@ class Semaphore:
         self.release()
 
     async def acquire(self) -> None:
-        """Acquire the semaphore.
-
-        If the semaphore value is zero, this will wait until it's released.
-        """
+        """Acquire the semaphore."""
         await self._semaphore.acquire()
+        self._current_acquisitions += 1
 
     def release(self) -> None:
-        """Release the semaphore, incrementing its value."""
+        """Release the semaphore."""
+        if self._current_acquisitions <= 0:
+            raise RuntimeError(
+                "Cannot release semaphore: no outstanding acquisitions"
+            )
         self._semaphore.release()
+        self._current_acquisitions -= 1
+
+    @property
+    def current_acquisitions(self) -> int:
+        """Get the current number of outstanding acquisitions."""
+        return self._current_acquisitions
+
+    @property
+    def initial_value(self) -> int:
+        """Get the initial semaphore value."""
+        return self._initial_value
 
 
 class CapacityLimiter:
     """A context manager for limiting the number of concurrent operations."""
 
-    def __init__(self, total_tokens: float):
-        """Initialize a new capacity limiter.
-
-        Args:
-            total_tokens: The maximum number of tokens (>= 1)
-        """
-        if total_tokens < 1:
-            raise ValueError("The total number of tokens must be >= 1")
-        self._limiter = anyio.CapacityLimiter(total_tokens)
+    def __init__(self, total_tokens: int | float):
+        """Initialize a new capacity limiter."""
+        if total_tokens == math.inf:
+            processed_tokens = math.inf
+        elif isinstance(total_tokens, (int, float)) and total_tokens >= 1:
+            processed_tokens = (
+                int(total_tokens) if total_tokens != math.inf else math.inf
+            )
+        else:
+            raise ValueError(
+                "The total number of tokens must be >= 1 (int or math.inf)"
+            )
+
+        self._limiter = anyio.CapacityLimiter(processed_tokens)
+        self._borrower_counter = 0
+        self._active_borrowers = {}
+        track_resource(self, f"CapacityLimiter-{id(self)}", "CapacityLimiter")
+
+    def __del__(self):
+        """Clean up resource tracking when limiter is destroyed."""
+        try:
+            untrack_resource(self)
+        except Exception:
+            pass
 
     async def __aenter__(self) -> None:
-        """Acquire a token.
-
-        If no tokens are available, this will wait until one is released.
-        """
+        """Acquire a token."""
         await self.acquire()
 
     async def __aexit__(
@@ -122,35 +156,49 @@ class CapacityLimiter:
         self.release()
 
     async def acquire(self) -> None:
-        """Acquire a token.
-
-        If no tokens are available, this will wait until one is released.
-        """
-        await self._limiter.acquire()
+        """Acquire a token."""
+        # Create a unique borrower identity for each acquisition
+        self._borrower_counter += 1
+        borrower = f"borrower-{self._borrower_counter}"
+        await self._limiter.acquire_on_behalf_of(borrower)
+        self._active_borrowers[borrower] = True
 
     def release(self) -> None:
-        """Release a token.
+        """Release a token."""
+        # Find and release the first active borrower
+        if not self._active_borrowers:
+            raise RuntimeError("No tokens to release")
 
-        Raises:
-            RuntimeError: If the current task doesn't hold any tokens.
-        """
-        self._limiter.release()
+        borrower = next(iter(self._active_borrowers))
+        self._limiter.release_on_behalf_of(borrower)
+        del self._active_borrowers[borrower]
 
     @property
-    def total_tokens(self) -> float:
+    def total_tokens(self) -> int | float:
         """The total number of tokens."""
-        return self._limiter.total_tokens
+        return float(self._limiter.total_tokens)
 
     @total_tokens.setter
-    def total_tokens(self, value: float) -> None:
-        """Set the total number of tokens.
-
-        Args:
-            value: The new total number of tokens (>= 1)
-        """
-        if value < 1:
-            raise ValueError("The total number of tokens must be >= 1")
-        self._limiter.total_tokens = value
+    def total_tokens(self, value: int | float) -> None:
+        """Set the total number of tokens."""
+        if value == math.inf:
+            processed_value = math.inf
+        elif isinstance(value, (int, float)) and value >= 1:
+            processed_value = int(value) if value != math.inf else math.inf
+        else:
+            raise ValueError(
+                "The total number of tokens must be >= 1 (int or math.inf)"
+            )
+
+        current_borrowed = self._limiter.borrowed_tokens
+        if processed_value != math.inf and processed_value < current_borrowed:
+            raise ValueError(
+                f"Cannot set total_tokens to {processed_value}: {current_borrowed} tokens "
+                f"are currently borrowed. Wait for tokens to be released or "
+                f"set total_tokens to at least {current_borrowed}."
+            )
+
+        self._limiter.total_tokens = processed_value
 
     @property
     def borrowed_tokens(self) -> int:
@@ -158,28 +206,28 @@ class CapacityLimiter:
         return self._limiter.borrowed_tokens
 
     @property
-    def available_tokens(self) -> float:
+    def available_tokens(self) -> int | float:
         """The number of tokens currently available."""
         return self._limiter.available_tokens
 
 
 class Event:
-    """An event object for task synchronization.
-
-    An event can be in one of two states: set or unset. When set, tasks waiting
-    on the event are allowed to proceed.
-    """
+    """An event object for task synchronization."""
 
     def __init__(self):
         """Initialize a new event in the unset state."""
         self._event = anyio.Event()
+        track_resource(self, f"Event-{id(self)}", "Event")
 
-    def is_set(self) -> bool:
-        """Check if the event is set.
+    def __del__(self):
+        """Clean up resource tracking when event is destroyed."""
+        try:
+            untrack_resource(self)
+        except Exception:
+            pass
 
-        Returns:
-            True if the event is set, False otherwise.
-        """
+    def is_set(self) -> bool:
+        """Check if the event is set."""
         return self._event.is_set()
 
     def set(self) -> None:
@@ -195,21 +243,21 @@ class Condition:
     """A condition variable for task synchronization."""
 
     def __init__(self, lock: Lock | None = None):
-        """Initialize a new condition.
-
-        Args:
-            lock: The lock to use, or None to create a new one
-        """
+        """Initialize a new condition."""
         self._lock = lock or Lock()
         self._condition = anyio.Condition(self._lock._lock)
+        track_resource(self, f"Condition-{id(self)}", "Condition")
 
-    async def __aenter__(self) -> "Condition":
-        """Acquire the underlying lock.
+    def __del__(self):
+        """Clean up resource tracking when condition is destroyed."""
+        try:
+            untrack_resource(self)
+        except Exception:
+            pass
 
-        Returns:
-            The condition instance.
-        """
-        await self._lock.acquire()
+    async def __aenter__(self) -> "Condition":
+        """Acquire the underlying lock."""
+        await self._lock.__aenter__()
         return self
 
     async def __aexit__(
@@ -219,7 +267,7 @@ class Condition:
         exc_tb: TracebackType | None,
     ) -> None:
         """Release the underlying lock."""
-        self._lock.release()
+        await self._lock.__aexit__(exc_type, exc_val, exc_tb)
 
     async def wait(self) -> None:
         """Wait for a notification.

lionagi/libs/concurrency/resource_tracker.py

@@ -0,0 +1,182 @@
+"""Resource tracking utilities for concurrency primitives.
+
+This module provides lightweight resource leak detection and lifecycle tracking
+to address the security vulnerabilities identified in the hardening tests.
+"""
+
+import weakref
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any
+
+
+@dataclass
+class ResourceInfo:
+    """Information about a tracked resource."""
+
+    name: str
+    creation_time: datetime
+    resource_type: str
+
+
+class ResourceTracker:
+    """Lightweight resource lifecycle tracking for leak detection.
+
+    This addresses the over-engineering concerns by providing simple,
+    practical resource management without complex abstraction layers.
+    """
+
+    def __init__(self):
+        """Initialize a new resource tracker."""
+        self._active_resources: dict[int, ResourceInfo] = {}
+        self._weak_refs: weakref.WeakKeyDictionary = (
+            weakref.WeakKeyDictionary()
+        )
+
+    def track(
+        self, resource: Any, name: str, resource_type: str | None = None
+    ) -> None:
+        """Track a resource for leak detection.
+
+        Args:
+            resource: The resource to track
+            name: Human-readable name for the resource
+            resource_type: Optional type classification
+        """
+        if resource_type is None:
+            resource_type = type(resource).__name__
+
+        resource_info = ResourceInfo(
+            name=name,
+            creation_time=datetime.now(),
+            resource_type=resource_type,
+        )
+
+        # Use weak reference to avoid interfering with garbage collection
+        self._weak_refs[resource] = resource_info
+        self._active_resources[id(resource)] = resource_info
+
+    def untrack(self, resource: Any) -> None:
+        """Manually untrack a resource.
+
+        Args:
+            resource: The resource to stop tracking
+        """
+        resource_id = id(resource)
+        self._active_resources.pop(resource_id, None)
+        self._weak_refs.pop(resource, None)
+
+    def cleanup_check(self) -> list[ResourceInfo]:
+        """Check for potentially leaked resources.
+
+        Returns:
+            List of resource info for resources that may have leaked
+        """
+        # Clean up references to garbage collected objects
+        current_resources = []
+        for resource, info in list(self._weak_refs.items()):
+            current_resources.append(info)
+
+        return current_resources
+
+    def get_active_count(self) -> int:
+        """Get the number of currently tracked resources.
+
+        Returns:
+            Number of active tracked resources
+        """
+        return len(self._weak_refs)
+
+    def get_resource_summary(self) -> dict[str, int]:
+        """Get a summary of tracked resources by type.
+
+        Returns:
+            Dictionary mapping resource types to counts
+        """
+        summary = {}
+        for info in self._weak_refs.values():
+            resource_type = info.resource_type
+            summary[resource_type] = summary.get(resource_type, 0) + 1
+        return summary
+
+
+# Global tracker instance for convenience
+_global_tracker = ResourceTracker()
+
+
+def track_resource(
+    resource: Any, name: str, resource_type: str | None = None
+) -> None:
+    """Track a resource using the global tracker.
+
+    Args:
+        resource: The resource to track
+        name: Human-readable name for the resource
+        resource_type: Optional type classification
+    """
+    _global_tracker.track(resource, name, resource_type)
+
+
+def untrack_resource(resource: Any) -> None:
+    """Untrack a resource using the global tracker.
+
+    Args:
+        resource: The resource to stop tracking
+    """
+    _global_tracker.untrack(resource)
+
+
+def get_global_tracker() -> ResourceTracker:
+    """Get the global resource tracker instance.
+
+    Returns:
+        The global ResourceTracker instance
+    """
+    return _global_tracker
+
+
+def cleanup_check() -> list[ResourceInfo]:
+    """Check for potentially leaked resources using global tracker.
+
+    Returns:
+        List of resource info for resources that may have leaked
+    """
+    return _global_tracker.cleanup_check()
+
+
+class resource_leak_detector:
+    """Context manager for resource leak detection in tests and production.
+
+    Example:
+        async with resource_leak_detector() as tracker:
+            lock = Lock()
+            tracker.track(lock, "test_lock")
+            # ... use lock
+        # Automatically checks for leaks on exit
+    """
+
+    def __init__(self, raise_on_leak: bool = False):
+        """Initialize the leak detector.
+
+        Args:
+            raise_on_leak: Whether to raise an exception if leaks are detected
+        """
+        self.raise_on_leak = raise_on_leak
+        self.tracker = ResourceTracker()
+        self._initial_count = 0
+
+    async def __aenter__(self) -> ResourceTracker:
+        """Enter the context and return the tracker."""
+        self._initial_count = self.tracker.get_active_count()
+        return self.tracker
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the context and check for leaks."""
+        leaked_resources = self.tracker.cleanup_check()
+
+        if leaked_resources and self.raise_on_leak:
+            resource_summary = self.tracker.get_resource_summary()
+            raise RuntimeError(
+                f"Resource leak detected: {len(leaked_resources)} resources "
+                f"still active. Summary: {resource_summary}"
+            )

lionagi/libs/concurrency/task.py

@@ -1,8 +1,10 @@
 """Task group implementation for structured concurrency."""
 
+from __future__ import annotations
+
 from collections.abc import Awaitable, Callable
 from types import TracebackType
-from typing import Any, Optional, TypeVar
+from typing import Any, TypeVar
 
 import anyio
 
@@ -61,7 +63,7 @@ class TaskGroup:
             raise RuntimeError("Task group is not active")
         return await self._task_group.start(func, *args, name=name)
 
-    async def __aenter__(self) -> "TaskGroup":
+    async def __aenter__(self) -> TaskGroup:
         """Enter the task group context.
 
         Returns:

lionagi/operations/builder.py

@@ -14,6 +14,7 @@ from typing import Any
 from lionagi.operations.node import BranchOperations, Operation
 from lionagi.protocols.graph.edge import Edge
 from lionagi.protocols.graph.graph import Graph
+from lionagi.protocols.types import ID
 
 __all__ = (
     "OperationGraphBuilder",
@@ -76,6 +77,7 @@ class OperationGraphBuilder:
         node_id: str | None = None,
         depends_on: list[str] | None = None,
         inherit_context: bool = False,
+        branch=None,
         **parameters,
     ) -> str:
         """
@@ -108,6 +110,9 @@ class OperationGraphBuilder:
             # Add as metadata for easy lookup
             node.metadata["reference_id"] = node_id
 
+        if branch:
+            node.branch_id = ID.get_id(branch)
+
         # Handle dependencies
         if depends_on:
             for dep_id in depends_on:
@@ -227,6 +232,7 @@ class OperationGraphBuilder:
         source_node_ids: list[str] | None = None,
         inherit_context: bool = False,
         inherit_from_source: int = 0,
+        branch=None,
         **parameters,
     ) -> str:
         """
@@ -264,6 +270,9 @@ class OperationGraphBuilder:
         if node_id:
             node.metadata["reference_id"] = node_id
 
+        if branch:
+            node.branch_id = ID.get_id(branch)
+
         # Store context inheritance for aggregations
         if inherit_context and sources:
             node.metadata["inherit_context"] = True