Skip to content

Commit 5613ce4

Browse files
committed
feat: postgresql task store
1 parent e696c8d commit 5613ce4

File tree

5 files changed

+618
-560
lines changed

5 files changed

+618
-560
lines changed

.github/workflows/unit-tests.yml

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,17 @@ jobs:
1515
runs-on: ubuntu-latest
1616

1717
if: github.repository == 'google/a2a-python'
18+
services:
19+
postgres:
20+
image: postgres:15-alpine
21+
env:
22+
POSTGRES_USER: postgres
23+
POSTGRES_PASSWORD: postgres
24+
POSTGRES_DB: a2a_test
25+
ports:
26+
- 5432:5432
27+
volumes:
28+
- ${{ github.workspace }}/docker/postgres/init.sql:/docker-entrypoint-initdb.d/init.sql
1829

1930
strategy:
2031
matrix:
@@ -28,6 +39,9 @@ jobs:
2839
uses: actions/setup-python@v5
2940
with:
3041
python-version: ${{ matrix.python-version }}
42+
- name: Set postgres test DSN
43+
run: |
44+
export POSTGRES_TEST_DSN="postgresql://postgres:postgres@localhost:5432/a2a_test"
3145
3246
- name: Install uv
3347
run: |

pyproject.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ authors = [{ name = "Google LLC", email = "[email protected]" }]
88
requires-python = ">=3.13"
99
keywords = ["A2A", "A2A SDK", "A2A Protocol", "Agent2Agent"]
1010
dependencies = [
11+
"asyncpg>=0.30.0",
1112
"httpx>=0.28.1",
1213
"httpx-sse>=0.4.0",
1314
"opentelemetry-api>=1.33.0",
@@ -70,6 +71,7 @@ members = [
7071

7172
[dependency-groups]
7273
dev = [
74+
"asyncpg-stubs>=0.30.1",
7375
"datamodel-code-generator>=0.30.0",
7476
"mypy>=1.15.0",
7577
"pytest>=8.3.5",

src/a2a/server/tasks/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
from a2a.server.tasks.inmemory_push_notifier import InMemoryPushNotifier
44
from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
5+
from a2a.server.tasks.postgresql_task_store import PostgreSQLTaskStore
56
from a2a.server.tasks.push_notifier import PushNotifier
67
from a2a.server.tasks.result_aggregator import ResultAggregator
78
from a2a.server.tasks.task_manager import TaskManager
@@ -12,6 +13,7 @@
1213
__all__ = [
1314
'InMemoryPushNotifier',
1415
'InMemoryTaskStore',
16+
'PostgreSQLTaskStore',
1517
'PushNotifier',
1618
'ResultAggregator',
1719
'TaskManager',

src/a2a/types.py

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -733,7 +733,7 @@ class TaskResubscriptionRequest(BaseModel):
733733
"""
734734

735735

736-
class TaskState(Enum):
736+
class TaskState(str, Enum):
737737
"""
738738
Represents the possible states of a Task.
739739
"""
@@ -1088,7 +1088,9 @@ class Artifact(BaseModel):
10881088

10891089

10901090
class GetTaskPushNotificationConfigResponse(
1091-
RootModel[JSONRPCErrorResponse | GetTaskPushNotificationConfigSuccessResponse]
1091+
RootModel[
1092+
JSONRPCErrorResponse | GetTaskPushNotificationConfigSuccessResponse
1093+
]
10921094
):
10931095
root: JSONRPCErrorResponse | GetTaskPushNotificationConfigSuccessResponse
10941096
"""
@@ -1239,7 +1241,9 @@ class SendStreamingMessageRequest(BaseModel):
12391241

12401242

12411243
class SetTaskPushNotificationConfigResponse(
1242-
RootModel[JSONRPCErrorResponse | SetTaskPushNotificationConfigSuccessResponse]
1244+
RootModel[
1245+
JSONRPCErrorResponse | SetTaskPushNotificationConfigSuccessResponse
1246+
]
12431247
):
12441248
root: JSONRPCErrorResponse | SetTaskPushNotificationConfigSuccessResponse
12451249
"""
@@ -1524,7 +1528,9 @@ class SendStreamingMessageSuccessResponse(BaseModel):
15241528
"""
15251529

15261530

1527-
class CancelTaskResponse(RootModel[JSONRPCErrorResponse | CancelTaskSuccessResponse]):
1531+
class CancelTaskResponse(
1532+
RootModel[JSONRPCErrorResponse | CancelTaskSuccessResponse]
1533+
):
15281534
root: JSONRPCErrorResponse | CancelTaskSuccessResponse
15291535
"""
15301536
JSON-RPC response for the 'tasks/cancel' method.
@@ -1563,7 +1569,9 @@ class JSONRPCResponse(
15631569
"""
15641570

15651571

1566-
class SendMessageResponse(RootModel[JSONRPCErrorResponse | SendMessageSuccessResponse]):
1572+
class SendMessageResponse(
1573+
RootModel[JSONRPCErrorResponse | SendMessageSuccessResponse]
1574+
):
15671575
root: JSONRPCErrorResponse | SendMessageSuccessResponse
15681576
"""
15691577
JSON-RPC response model for the 'message/send' method.

0 commit comments

Comments
 (0)