feat: postgresql task store

Author: LaurentMnr95

.github/workflows/unit-tests.yml

@@ -15,6 +15,17 @@ jobs:
     runs-on: ubuntu-latest
 
     if: github.repository == 'google/a2a-python'
+    services:
+      postgres:
+        image: postgres:15-alpine
+        env:
+          POSTGRES_USER: postgres
+          POSTGRES_PASSWORD: postgres
+          POSTGRES_DB: a2a_test
+        ports:
+          - 5432:5432
+        volumes:
+          - ${{ github.workspace }}/docker/postgres/init.sql:/docker-entrypoint-initdb.d/init.sql
 
     strategy:
       matrix:
@@ -28,6 +39,9 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
+      - name: Set postgres test DSN
+        run: |
+          export POSTGRES_TEST_DSN="postgresql://postgres:postgres@localhost:5432/a2a_test"
 
       - name: Install uv
         run: |

pyproject.toml

@@ -8,6 +8,7 @@ authors = [{ name = "Google LLC", email = "[email protected]" }]
 requires-python = ">=3.10"
 keywords = ["A2A", "A2A SDK", "A2A Protocol", "Agent2Agent"]
 dependencies = [
+    "asyncpg>=0.30.0",
     "httpx>=0.28.1",
     "httpx-sse>=0.4.0",
     "opentelemetry-api>=1.33.0",
@@ -73,6 +74,7 @@ members = [
 
 [dependency-groups]
 dev = [
+    "asyncpg-stubs>=0.30.1",
     "datamodel-code-generator>=0.30.0",
     "mypy>=1.15.0",
     "pytest>=8.3.5",

src/a2a/server/tasks/__init__.py

@@ -2,6 +2,7 @@
 
 from a2a.server.tasks.inmemory_push_notifier import InMemoryPushNotifier
 from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
+from a2a.server.tasks.postgresql_task_store import PostgreSQLTaskStore
 from a2a.server.tasks.push_notifier import PushNotifier
 from a2a.server.tasks.result_aggregator import ResultAggregator
 from a2a.server.tasks.task_manager import TaskManager
@@ -12,6 +13,7 @@
 __all__ = [
     'InMemoryPushNotifier',
     'InMemoryTaskStore',
+    'PostgreSQLTaskStore',
     'PushNotifier',
     'ResultAggregator',
     'TaskManager',

src/a2a/types.py

@@ -731,7 +731,7 @@ class TaskResubscriptionRequest(BaseModel):
     """
 
 
-class TaskState(Enum):
+class TaskState(str, Enum):
     """
     Represents the possible states of a Task.
     """
@@ -1086,7 +1086,9 @@ class Artifact(BaseModel):
 
 
 class GetTaskPushNotificationConfigResponse(
-    RootModel[JSONRPCErrorResponse | GetTaskPushNotificationConfigSuccessResponse]
+    RootModel[
+        JSONRPCErrorResponse | GetTaskPushNotificationConfigSuccessResponse
+    ]
 ):
     root: JSONRPCErrorResponse | GetTaskPushNotificationConfigSuccessResponse
     """
@@ -1237,7 +1239,9 @@ class SendStreamingMessageRequest(BaseModel):
 
 
 class SetTaskPushNotificationConfigResponse(
-    RootModel[JSONRPCErrorResponse | SetTaskPushNotificationConfigSuccessResponse]
+    RootModel[
+        JSONRPCErrorResponse | SetTaskPushNotificationConfigSuccessResponse
+    ]
 ):
     root: JSONRPCErrorResponse | SetTaskPushNotificationConfigSuccessResponse
     """
@@ -1530,7 +1534,9 @@ class SendStreamingMessageSuccessResponse(BaseModel):
     """
 
 
-class CancelTaskResponse(RootModel[JSONRPCErrorResponse | CancelTaskSuccessResponse]):
+class CancelTaskResponse(
+    RootModel[JSONRPCErrorResponse | CancelTaskSuccessResponse]
+):
     root: JSONRPCErrorResponse | CancelTaskSuccessResponse
     """
     JSON-RPC response for the 'tasks/cancel' method.
@@ -1569,7 +1575,9 @@ class JSONRPCResponse(
     """
 
 
-class SendMessageResponse(RootModel[JSONRPCErrorResponse | SendMessageSuccessResponse]):
+class SendMessageResponse(
+    RootModel[JSONRPCErrorResponse | SendMessageSuccessResponse]
+):
     root: JSONRPCErrorResponse | SendMessageSuccessResponse
     """
     JSON-RPC response model for the 'message/send' method.