a2a-lite 0.1.0__py3-none-any.whl → 0.2.0__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.0"
 
 __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
@@ -375,10 +396,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,
@@ -450,4 +495,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

@@ -116,10 +116,17 @@ 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)
@@ -131,11 +138,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]:

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.0",
 ]
 '''
     (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):
@@ -100,7 +104,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)
@@ -132,17 +136,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
@@ -172,13 +178,14 @@ class LiteAgentExecutor(AgentExecutor):
                 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
+            type_name = str(param_type)
             if "FilePart" in type_name:
                 from .parts import FilePart
                 if isinstance(value, dict):
@@ -234,7 +241,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/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 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):
@@ -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

a2a_lite/testing.py

@@ -3,21 +3,61 @@ Testing utilities for A2A Lite agents.
 
 Makes testing agents as simple as:
 
-    from a2a_lite.testing import TestClient
+    from a2a_lite.testing import AgentTestClient
 
     def test_my_agent():
-        client = TestClient(agent)
+        client = AgentTestClient(agent)
         result = client.call("greet", name="World")
         assert result == "Hello, World!"
 """
 from __future__ import annotations
 
 import json
+from dataclasses import dataclass
 from typing import Any, Dict, Optional
 from uuid import uuid4
 
 
-class TestClient:
+@dataclass
+class TestResult:
+    """
+    Structured result from a test client call.
+
+    Provides multiple ways to access the result:
+    - .data — parsed Python object (dict, list, int, str, etc.)
+    - .text — raw text string
+    - .json() — parse text as JSON (raises on invalid JSON)
+    - .raw_response — the full A2A response dict
+    """
+    _data: Any
+    _text: str
+    raw_response: Dict[str, Any]
+
+    @property
+    def data(self) -> Any:
+        """The parsed result value."""
+        return self._data
+
+    @property
+    def text(self) -> str:
+        """The raw text representation."""
+        return self._text
+
+    def json(self) -> Any:
+        """Parse the text as JSON."""
+        return json.loads(self._text)
+
+    def __eq__(self, other: Any) -> bool:
+        """Allow direct comparison with the data value for convenience."""
+        if isinstance(other, TestResult):
+            return self._data == other._data
+        return self._data == other
+
+    def __repr__(self) -> str:
+        return f"TestResult(data={self._data!r})"
+
+
+class AgentTestClient:
     """
     Simple test client for A2A Lite agents.
 
@@ -29,7 +69,7 @@ class TestClient:
             return a + b
 
         # In your test
-        client = TestClient(agent)
+        client = AgentTestClient(agent)
         assert client.call("add", a=2, b=3) == 5
     """
 
@@ -89,7 +129,7 @@ class TestClient:
         # Extract the actual result from A2A response
         return self._extract_result(data)
 
-    def _extract_result(self, response: Dict) -> Any:
+    def _extract_result(self, response: Dict) -> TestResult:
         """Extract the skill result from A2A response."""
         if "error" in response:
             raise TestClientError(response["error"])
@@ -103,11 +143,12 @@ class TestClient:
                 text = part.get("text", "")
                 # Try to parse as JSON
                 try:
-                    return json.loads(text)
+                    data = json.loads(text)
                 except json.JSONDecodeError:
-                    return text
+                    data = text
+                return TestResult(_data=data, _text=text, raw_response=response)
 
-        return result
+        return TestResult(_data=result, _text=json.dumps(result), raw_response=response)
 
     def get_agent_card(self) -> Dict[str, Any]:
         """
@@ -183,13 +224,13 @@ class TestClientError(Exception):
 
 
 # Async version for async tests
-class AsyncTestClient:
+class AsyncAgentTestClient:
     """
     Async test client for A2A Lite agents.
 
     Example:
         async def test_my_agent():
-            client = AsyncTestClient(agent)
+            client = AsyncAgentTestClient(agent)
             result = await client.call("greet", name="World")
             assert result == "Hello, World!"
     """
@@ -243,7 +284,7 @@ class AsyncTestClient:
 
         return self._extract_result(data)
 
-    def _extract_result(self, response: Dict) -> Any:
+    def _extract_result(self, response: Dict) -> TestResult:
         """Extract the skill result from A2A response."""
         if "error" in response:
             raise TestClientError(response["error"])
@@ -255,11 +296,12 @@ class AsyncTestClient:
             if part.get("kind") == "text" or part.get("type") == "text":
                 text = part.get("text", "")
                 try:
-                    return json.loads(text)
+                    data = json.loads(text)
                 except json.JSONDecodeError:
-                    return text
+                    data = text
+                return TestResult(_data=data, _text=text, raw_response=response)
 
-        return result
+        return TestResult(_data=result, _text=json.dumps(result), raw_response=response)
 
     async def close(self):
         """Close the client."""

a2a_lite/utils.py

@@ -5,6 +5,22 @@ from typing import Any, Dict, Type, get_origin, get_args, Union
 import inspect
 
 
+def _is_or_subclass(hint: Any, target_class: Type) -> bool:
+    """
+    Check if a type hint is, or is a subclass of, the target class.
+
+    Works with raw classes and string annotations.
+    """
+    try:
+        if hint is target_class:
+            return True
+        if isinstance(hint, type) and issubclass(hint, target_class):
+            return True
+    except TypeError:
+        pass
+    return False
+
+
 def type_to_json_schema(python_type: Type) -> Dict[str, Any]:
     """
     Convert Python type to JSON Schema.