a2a-lite 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

a2a_lite/__init__.py

@@ -13,8 +13,8 @@ SIMPLE (8 lines):
     agent.run()
 
 TEST IT (3 lines):
-    from a2a_lite import TestClient
-    client = TestClient(agent)
+    from a2a_lite import AgentTestClient
+    client = AgentTestClient(agent)
     assert client.call("greet", name="World") == "Hello, World!"
 
 WITH PYDANTIC:
@@ -64,7 +64,7 @@ WITH FILES (opt-in):
 from .agent import Agent
 from .decorators import SkillDefinition
 from .discovery import AgentDiscovery, DiscoveredAgent
-from .testing import TestClient, AsyncTestClient
+from .testing import AgentTestClient, AsyncAgentTestClient, TestResult
 
 # Middleware - always available
 from .middleware import (
@@ -102,7 +102,7 @@ from .auth import (
     require_auth,
 )
 
-__version__ = "0.1.0"
+__version__ = "0.2.1"
 
 __all__ = [
     # Core
@@ -111,8 +111,9 @@ __all__ = [
     "AgentDiscovery",
     "DiscoveredAgent",
     # Testing
-    "TestClient",
-    "AsyncTestClient",
+    "AgentTestClient",
+    "AsyncAgentTestClient",
+    "TestResult",
     # Middleware
     "MiddlewareContext",
     "MiddlewareChain",

a2a_lite/agent.py

@@ -7,9 +7,12 @@ from __future__ import annotations
 
 import asyncio
 import inspect
+import logging
 from typing import Any, Callable, Optional, Dict, List, Type, Union, get_origin, get_args
 from dataclasses import dataclass, field
 
+logger = logging.getLogger(__name__)
+
 import uvicorn
 
 from a2a.server.apps import A2AStarletteApplication
@@ -23,7 +26,7 @@ from a2a.types import (
 
 from .executor import LiteAgentExecutor
 from .decorators import SkillDefinition
-from .utils import type_to_json_schema, extract_function_schemas
+from .utils import type_to_json_schema, extract_function_schemas, _is_or_subclass
 from .middleware import MiddlewareChain, MiddlewareContext
 from .streaming import is_generator_function
 from .webhooks import NotificationManager, WebhookClient
@@ -85,6 +88,8 @@ class Agent:
     # Optional enterprise features
     auth: Optional[Any] = None  # AuthProvider
     task_store: Optional[Any] = None  # TaskStore or "memory"
+    cors_origins: Optional[List[str]] = None
+    production: bool = False
 
     def __post_init__(self):
         # Internal state
@@ -161,16 +166,30 @@ class Agent:
             if is_streaming:
                 self._has_streaming = True
 
-            # Detect special parameter types
-            hints = getattr(func, '__annotations__', {})
-            needs_task_context = any(
-                str(h).endswith("TaskContext") or "TaskContext" in str(h)
-                for h in hints.values()
-            )
-            needs_interaction = any(
-                str(h).endswith("InteractionContext") or "InteractionContext" in str(h)
-                for h in hints.values()
-            )
+            # Detect special parameter types using proper type introspection
+            import typing
+            from .tasks import TaskContext as _TaskContext
+            from .human_loop import InteractionContext as _InteractionContext
+
+            needs_task_context = False
+            needs_interaction = False
+            task_context_param: str | None = None
+            interaction_param: str | None = None
+
+            try:
+                resolved_hints = typing.get_type_hints(func)
+            except Exception:
+                resolved_hints = getattr(func, '__annotations__', {})
+
+            for param_name, hint in resolved_hints.items():
+                if param_name == 'return':
+                    continue
+                if _is_or_subclass(hint, _TaskContext):
+                    needs_task_context = True
+                    task_context_param = param_name
+                elif _is_or_subclass(hint, _InteractionContext):
+                    needs_interaction = True
+                    interaction_param = param_name
 
             # Extract schemas
             input_schema, output_schema = extract_function_schemas(func)
@@ -186,6 +205,8 @@ class Agent:
                 is_streaming=is_streaming,
                 needs_task_context=needs_task_context,
                 needs_interaction=needs_interaction,
+                task_context_param=task_context_param,
+                interaction_param=interaction_param,
             )
 
             self._skills[skill_name] = skill_def
@@ -311,6 +332,10 @@ class Agent:
             task_store=self._task_store,
         )
 
+        # The SDK's InMemoryTaskStore handles protocol-level task lifecycle
+        # (task creation, state transitions per the A2A spec). This is separate
+        # from self._task_store which provides application-level tracking
+        # (progress updates, custom status) exposed via TaskContext to skills.
         request_handler = DefaultRequestHandler(
             agent_executor=executor,
             task_store=InMemoryTaskStore(),
@@ -375,10 +400,34 @@ class Agent:
             )
             console.print(f"[dim]mDNS discovery enabled for {self.name}[/]")
 
+        # Production mode warning
+        if self.production:
+            url_str = self.url or f"http://{display_host}:{port}"
+            if not url_str.startswith("https://"):
+                logger.warning(
+                    "Running in production mode over HTTP. "
+                    "Consider using HTTPS for secure communication."
+                )
+
+        # Build the ASGI app
+        app = app_builder.build()
+
+        # Add CORS middleware if configured
+        if self.cors_origins is not None:
+            from starlette.middleware.cors import CORSMiddleware
+            from starlette.middleware import Middleware as StarletteMiddleware
+            # Wrap the existing app with CORS
+            app.add_middleware(
+                CORSMiddleware,
+                allow_origins=self.cors_origins,
+                allow_methods=["*"],
+                allow_headers=["*"],
+            )
+
         # Start server
         try:
             uvicorn.run(
-                app_builder.build(),
+                app,
                 host=host,
                 port=port,
                 log_level=log_level,
@@ -440,6 +489,7 @@ class Agent:
             task_store=self._task_store,
         )
 
+        # SDK task store for protocol-level lifecycle (separate from app-level self._task_store)
         request_handler = DefaultRequestHandler(
             agent_executor=executor,
             task_store=InMemoryTaskStore(),
@@ -450,4 +500,16 @@ class Agent:
             http_handler=request_handler,
         )
 
-        return app_builder.build()
+        app = app_builder.build()
+
+        # Add CORS middleware if configured
+        if self.cors_origins is not None:
+            from starlette.middleware.cors import CORSMiddleware
+            app.add_middleware(
+                CORSMiddleware,
+                allow_origins=self.cors_origins,
+                allow_methods=["*"],
+                allow_headers=["*"],
+            )
+
+        return app

a2a_lite/auth.py

@@ -60,6 +60,17 @@ class AuthRequest:
     method: str = "POST"
     path: str = "/"
 
+    def get_header(self, name: str) -> Optional[str]:
+        """Get a header value (case-insensitive)."""
+        # Try exact match first, then case-insensitive
+        if name in self.headers:
+            return self.headers[name]
+        lower = name.lower()
+        for k, v in self.headers.items():
+            if k.lower() == lower:
+                return v
+        return None
+
 
 @dataclass
 class AuthResult:
@@ -116,13 +127,20 @@ class APIKeyAuth(AuthProvider):
         header: str = "X-API-Key",
         query_param: Optional[str] = None,
     ):
-        self.keys = set(keys)
+        # Store only hashes of keys for security
+        self._key_hashes = {
+            hashlib.sha256(k.encode()).hexdigest() for k in keys
+        }
         self.header = header
         self.query_param = query_param
 
+    def _hash_key(self, key: str) -> str:
+        """Hash a key using SHA-256."""
+        return hashlib.sha256(key.encode()).hexdigest()
+
     async def authenticate(self, request: AuthRequest) -> AuthResult:
-        # Check header
-        key = request.headers.get(self.header)
+        # Check header (case-insensitive)
+        key = request.get_header(self.header)
 
         # Check query param
         if not key and self.query_param:
@@ -131,11 +149,12 @@ class APIKeyAuth(AuthProvider):
         if not key:
             return AuthResult.failure("API key required")
 
-        if key not in self.keys:
+        key_hash = self._hash_key(key)
+        if key_hash not in self._key_hashes:
             return AuthResult.failure("Invalid API key")
 
-        # Use hash of key as user ID
-        user_id = hashlib.sha256(key.encode()).hexdigest()[:16]
+        # Use hash prefix as user ID
+        user_id = key_hash[:16]
         return AuthResult.success(user_id=user_id)
 
     def get_scheme(self) -> Dict[str, Any]:
@@ -171,7 +190,7 @@ class BearerAuth(AuthProvider):
         self.header = header
 
     async def authenticate(self, request: AuthRequest) -> AuthResult:
-        auth_header = request.headers.get(self.header, "")
+        auth_header = request.get_header(self.header) or ""
 
         if not auth_header.startswith("Bearer "):
             return AuthResult.failure("Bearer token required")
@@ -220,7 +239,7 @@ class OAuth2Auth(AuthProvider):
         self._jwks_client = None
 
     async def authenticate(self, request: AuthRequest) -> AuthResult:
-        auth_header = request.headers.get("Authorization", "")
+        auth_header = request.get_header("Authorization") or ""
 
         if not auth_header.startswith("Bearer "):
             return AuthResult.failure("Bearer token required")

a2a_lite/cli.py

@@ -75,7 +75,7 @@ version = "0.1.0"
 description = "A2A Agent: {name}"
 requires-python = ">=3.10"
 dependencies = [
-    "a2a-lite>=0.1.0",
+    "a2a-lite>=0.2.1",
 ]
 '''
     (project_path / "pyproject.toml").write_text(pyproject)

a2a_lite/decorators.py

@@ -2,7 +2,7 @@
 Decorator definitions and skill metadata.
 """
 from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict, List, Optional
 
 
 @dataclass
@@ -18,6 +18,8 @@ class SkillDefinition:
     is_streaming: bool = False
     needs_task_context: bool = False
     needs_interaction: bool = False
+    task_context_param: Optional[str] = None
+    interaction_param: Optional[str] = None
 
     def to_dict(self) -> Dict[str, Any]:
         """Convert to dictionary for serialization."""

a2a_lite/discovery.py

@@ -4,10 +4,13 @@ mDNS-based local agent discovery using Zeroconf.
 from __future__ import annotations
 
 import asyncio
+import logging
 import socket
 from typing import Dict, List, Optional
 from dataclasses import dataclass
 
+logger = logging.getLogger(__name__)
+
 from zeroconf import ServiceBrowser, ServiceListener, Zeroconf, ServiceInfo
 from zeroconf.asyncio import AsyncZeroconf, AsyncServiceBrowser
 
@@ -113,6 +116,7 @@ class AgentDiscovery:
             s.close()
             return ip
         except Exception:
+            logger.debug("Could not detect local IP, falling back to 127.0.0.1")
             return "127.0.0.1"
 
 

a2a_lite/executor.py

@@ -6,14 +6,18 @@ from __future__ import annotations
 import asyncio
 import inspect
 import json
+import logging
 from typing import Any, Callable, Dict, List, Optional
 
+logger = logging.getLogger(__name__)
+
 from a2a.server.agent_execution import AgentExecutor, RequestContext
 from a2a.server.events import EventQueue
 
 from .decorators import SkillDefinition
 from .middleware import MiddlewareChain, MiddlewareContext
 from .streaming import is_generator_function, stream_generator
+from .utils import _is_or_subclass
 
 
 class LiteAgentExecutor(AgentExecutor):
@@ -55,6 +59,22 @@ class LiteAgentExecutor(AgentExecutor):
         from a2a.utils import new_agent_text_message
 
         try:
+            # Authenticate the request
+            if self.auth_provider:
+                from .auth import AuthRequest, NoAuth
+                if not isinstance(self.auth_provider, NoAuth):
+                    headers = {}
+                    if context.call_context and context.call_context.state:
+                        headers = context.call_context.state.get('headers', {})
+                    auth_request = AuthRequest(headers=headers)
+                    auth_result = await self.auth_provider.authenticate(auth_request)
+                    if not auth_result.authenticated:
+                        error_msg = json.dumps({
+                            "error": auth_result.error or "Authentication failed",
+                        })
+                        await event_queue.enqueue_event(new_agent_text_message(error_msg))
+                        return
+
             # Extract message and parts
             message, parts = self._extract_message_and_parts(context)
 
@@ -100,7 +120,7 @@ class LiteAgentExecutor(AgentExecutor):
                     else:
                         hook(skill_name, result, ctx)
                 except Exception:
-                    pass
+                    logger.warning("Completion hook error for skill '%s'", skill_name, exc_info=True)
 
         except Exception as e:
             await self._handle_error(e, event_queue)
@@ -116,7 +136,14 @@ class LiteAgentExecutor(AgentExecutor):
         if skill_name is None:
             if not self.skills:
                 return {"error": "No skills registered"}
-            skill_name = list(self.skills.keys())[0]
+            # Only auto-select if there's exactly one skill
+            if len(self.skills) == 1:
+                skill_name = list(self.skills.keys())[0]
+            else:
+                return {
+                    "error": "No skill specified. Use {\"skill\": \"name\", \"params\": {...}} format.",
+                    "available_skills": list(self.skills.keys()),
+                }
 
         if skill_name not in self.skills:
             return {
@@ -132,17 +159,19 @@ class LiteAgentExecutor(AgentExecutor):
         # Inject special contexts if needed
         if skill_def.needs_task_context and self.task_store:
             from .tasks import TaskContext, Task, TaskStatus, TaskState
-            task = self.task_store.create(skill_name, params)
+            task = await self.task_store.create(skill_name, params)
             # Only pass event_queue for streaming skills (status updates go via SSE)
             eq = event_queue if skill_def.is_streaming else None
             task_ctx = TaskContext(task, eq)
-            params["task"] = task_ctx
+            param_name = skill_def.task_context_param or "task"
+            params[param_name] = task_ctx
 
         if skill_def.needs_interaction:
             from .human_loop import InteractionContext
             task_id = metadata.get("task_id", "unknown")
             interaction_ctx = InteractionContext(task_id, event_queue)
-            params["ctx"] = interaction_ctx
+            param_name = skill_def.interaction_param or "ctx"
+            params[param_name] = interaction_ctx
 
         # Call the handler
         handler = skill_def.handler
@@ -161,26 +190,33 @@ class LiteAgentExecutor(AgentExecutor):
         metadata: Dict[str, Any],
     ) -> Dict[str, Any]:
         """Convert parameters to Pydantic models and file parts if needed."""
+        import typing
         handler = skill_def.handler
-        hints = getattr(handler, '__annotations__', {})
+        try:
+            hints = typing.get_type_hints(handler)
+        except Exception:
+            hints = getattr(handler, '__annotations__', {})
+
+        from .parts import FilePart, DataPart
 
         converted = {}
         for param_name, value in params.items():
+            if param_name == 'return':
+                continue
             param_type = hints.get(param_name)
 
             if param_type is None:
                 converted[param_name] = value
                 continue
 
-            type_name = str(param_type)
-
             # Skip special context types
-            if "TaskContext" in type_name or "InteractionContext" in type_name:
+            from .tasks import TaskContext as _TaskContext
+            from .human_loop import InteractionContext as _InteractionContext
+            if _is_or_subclass(param_type, _TaskContext) or _is_or_subclass(param_type, _InteractionContext):
                 continue
 
             # Convert FilePart
-            if "FilePart" in type_name:
-                from .parts import FilePart
+            if _is_or_subclass(param_type, FilePart):
                 if isinstance(value, dict):
                     # Handle both A2A format and simple dict format
                     if "file" in value:
@@ -201,8 +237,7 @@ class LiteAgentExecutor(AgentExecutor):
                 continue
 
             # Convert DataPart
-            if "DataPart" in type_name:
-                from .parts import DataPart
+            if _is_or_subclass(param_type, DataPart):
                 if isinstance(value, dict):
                     # Handle both A2A format and simple dict format
                     if "type" in value and value.get("type") == "data":
@@ -234,7 +269,7 @@ class LiteAgentExecutor(AgentExecutor):
             if isinstance(data, dict) and 'skill' in data:
                 return data['skill'], data.get('params', {})
         except json.JSONDecodeError:
-            pass
+            logger.debug("Message is not JSON, treating as plain text")
 
         return None, {"message": message}
 

a2a_lite/middleware.py

@@ -159,7 +159,12 @@ def retry_middleware(max_retries: int = 3, delay: float = 1.0):
 
 def rate_limit_middleware(requests_per_minute: int = 60):
     """
-    Create a simple rate limiting middleware.
+    Create a simple in-process rate limiting middleware.
+
+    Note: This rate limiter is per-process. Under multi-worker uvicorn
+    (e.g., ``--workers 4``), each worker tracks limits independently.
+    For shared rate limiting across workers, use an external store
+    (Redis, etc.) and a custom middleware.
 
     Example:
         agent.add_middleware(rate_limit_middleware(requests_per_minute=100))

a2a_lite/tasks.py

@@ -27,12 +27,15 @@ Example (with task tracking - opt-in):
 """
 from __future__ import annotations
 
+import asyncio
+import logging
 from dataclasses import dataclass, field
-from datetime import datetime
+from datetime import datetime, timezone
 from enum import Enum
 from typing import Any, Callable, Dict, List, Optional
 from uuid import uuid4
-import asyncio
+
+logger = logging.getLogger(__name__)
 
 
 class TaskState(str, Enum):
@@ -52,7 +55,7 @@ class TaskStatus:
     state: TaskState
     message: Optional[str] = None
     progress: Optional[float] = None  # 0.0 to 1.0
-    timestamp: datetime = field(default_factory=datetime.utcnow)
+    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
 
     def to_dict(self) -> Dict[str, Any]:
         return {
@@ -74,8 +77,8 @@ class Task:
     error: Optional[str] = None
     artifacts: List[Any] = field(default_factory=list)
     history: List[TaskStatus] = field(default_factory=list)
-    created_at: datetime = field(default_factory=datetime.utcnow)
-    updated_at: datetime = field(default_factory=datetime.utcnow)
+    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    updated_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
 
     def update_status(
         self,
@@ -86,7 +89,7 @@ class Task:
         """Update task status."""
         self.history.append(self.status)
         self.status = TaskStatus(state=state, message=message, progress=progress)
-        self.updated_at = datetime.utcnow()
+        self.updated_at = datetime.now(timezone.utc)
 
 
 class TaskContext:
@@ -145,7 +148,7 @@ class TaskContext:
                 else:
                     callback(self._task.status)
             except Exception:
-                pass
+                logger.warning("Status callback error for task '%s'", self._task.id, exc_info=True)
 
         # Send SSE event if streaming
         if self._event_queue:
@@ -170,52 +173,58 @@ class TaskContext:
 
 class TaskStore:
     """
-    In-memory task store.
+    In-memory task store with async locking for thread safety.
 
     For production, extend this with Redis/DB backend.
     """
 
     def __init__(self):
         self._tasks: Dict[str, Task] = {}
+        self._lock = asyncio.Lock()
 
-    def create(self, skill: str, params: Dict[str, Any]) -> Task:
+    async def create(self, skill: str, params: Dict[str, Any]) -> Task:
         """Create a new task."""
-        task = Task(
-            id=uuid4().hex,
-            skill=skill,
-            params=params,
-            status=TaskStatus(state=TaskState.SUBMITTED),
-        )
-        self._tasks[task.id] = task
-        return task
-
-    def get(self, task_id: str) -> Optional[Task]:
+        async with self._lock:
+            task = Task(
+                id=uuid4().hex,
+                skill=skill,
+                params=params,
+                status=TaskStatus(state=TaskState.SUBMITTED),
+            )
+            self._tasks[task.id] = task
+            return task
+
+    async def get(self, task_id: str) -> Optional[Task]:
         """Get task by ID."""
-        return self._tasks.get(task_id)
+        async with self._lock:
+            return self._tasks.get(task_id)
 
-    def update(self, task: Task) -> None:
+    async def update(self, task: Task) -> None:
         """Update task in store."""
-        self._tasks[task.id] = task
+        async with self._lock:
+            self._tasks[task.id] = task
 
-    def list(
+    async def list(
         self,
         state: Optional[TaskState] = None,
         skill: Optional[str] = None,
         limit: int = 100,
     ) -> List[Task]:
         """List tasks with optional filters."""
-        tasks = list(self._tasks.values())
+        async with self._lock:
+            tasks = list(self._tasks.values())
 
-        if state:
-            tasks = [t for t in tasks if t.status.state == state]
-        if skill:
-            tasks = [t for t in tasks if t.skill == skill]
+            if state:
+                tasks = [t for t in tasks if t.status.state == state]
+            if skill:
+                tasks = [t for t in tasks if t.skill == skill]
 
-        return sorted(tasks, key=lambda t: t.created_at, reverse=True)[:limit]
+            return sorted(tasks, key=lambda t: t.created_at, reverse=True)[:limit]
 
-    def delete(self, task_id: str) -> bool:
+    async def delete(self, task_id: str) -> bool:
         """Delete a task."""
-        if task_id in self._tasks:
-            del self._tasks[task_id]
-            return True
-        return False
+        async with self._lock:
+            if task_id in self._tasks:
+                del self._tasks[task_id]
+                return True
+            return False