acp-sdk 0.9.1__py3-none-any.whl → 0.10.1__py3-none-any.whl

acp_sdk/server/app.py

@@ -1,13 +1,14 @@
+import asyncio
 from collections.abc import AsyncGenerator
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import asynccontextmanager
-from datetime import datetime, timedelta
+from datetime import timedelta
 from enum import Enum
 
-from cachetools import TTLCache
 from fastapi import Depends, FastAPI, HTTPException, status
 from fastapi.applications import AppType, Lifespan
 from fastapi.encoders import jsonable_encoder
+from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, StreamingResponse
 
 from acp_sdk.models import (
@@ -27,12 +28,11 @@ from acp_sdk.models import (
     RunReadResponse,
     RunResumeRequest,
     RunResumeResponse,
-    SessionId,
 )
 from acp_sdk.models.errors import ACPError
+from acp_sdk.models.models import AwaitResume, RunStatus
 from acp_sdk.models.schemas import PingResponse
 from acp_sdk.server.agent import Agent
-from acp_sdk.server.bundle import RunBundle
 from acp_sdk.server.errors import (
     RequestValidationError,
     StarletteHTTPException,
@@ -41,8 +41,10 @@ from acp_sdk.server.errors import (
     http_exception_handler,
     validation_exception_handler,
 )
+from acp_sdk.server.executor import CancelData, Executor, RunData
 from acp_sdk.server.session import Session
-from acp_sdk.server.utils import stream_sse
+from acp_sdk.server.store import MemoryStore, Store
+from acp_sdk.server.utils import stream_sse, wait_util_stop
 
 
 class Headers(str, Enum):
@@ -51,8 +53,7 @@ class Headers(str, Enum):
 
 def create_app(
     *agents: Agent,
-    run_limit: int = 1000,
-    run_ttl: timedelta = timedelta(hours=1),
+    store: Store | None = None,
     lifespan: Lifespan[AppType] | None = None,
     dependencies: list[Depends] | None = None,
 ) -> FastAPI:
@@ -74,20 +75,37 @@ def create_app(
         dependencies=dependencies,
     )
 
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["https://agentcommunicationprotocol.dev"],
+        allow_methods=["*"],
+        allow_headers=["*"],
+        allow_credentials=True,
+    )
+
     agents: dict[AgentName, Agent] = {agent.name: agent for agent in agents}
-    runs: TTLCache[RunId, RunBundle] = TTLCache(maxsize=run_limit, ttl=run_ttl, timer=datetime.now)
-    sessions: TTLCache[SessionId, Session] = TTLCache(maxsize=run_limit, ttl=run_ttl, timer=datetime.now)
+
+    store = store or MemoryStore(limit=1000, ttl=timedelta(hours=1))
+    run_store = store.as_store(model=RunData, prefix="run_")
+    run_cancel_store = store.as_store(model=CancelData, prefix="run_cancel_")
+    run_resume_store = store.as_store(model=AwaitResume, prefix="run_resume_")
+    session_store = store.as_store(model=Session, prefix="session_")
 
     app.exception_handler(ACPError)(acp_error_handler)
     app.exception_handler(StarletteHTTPException)(http_exception_handler)
     app.exception_handler(RequestValidationError)(validation_exception_handler)
     app.exception_handler(Exception)(catch_all_exception_handler)
 
-    def find_run_bundle(run_id: RunId) -> RunBundle:
-        bundle = runs.get(run_id)
-        if not bundle:
+    async def find_run_data(run_id: RunId) -> RunData:
+        run_data = await run_store.get(run_id)
+        if not run_data:
             raise HTTPException(status_code=404, detail=f"Run {run_id} not found")
-        return bundle
+        if run_data.run.status.is_terminal:
+            return run_data
+        cancel_data = await run_cancel_store.get(run_data.key)
+        if cancel_data is not None:
+            run_data.run.status = RunStatus.CANCELLING
+        return run_data
 
     def find_agent(agent_name: AgentName) -> Agent:
         agent = agents.get(agent_name, None)
@@ -117,94 +135,111 @@ def create_app(
     async def create_run(request: RunCreateRequest) -> RunCreateResponse:
         agent = find_agent(request.agent_name)
 
-        session = sessions.get(request.session_id, Session(id=request.session_id)) if request.session_id else Session()
+        session = (
+            (await session_store.get(request.session_id)) or Session(id=request.session_id)
+            if request.session_id
+            else Session()
+        )
         nonlocal executor
-        bundle = RunBundle(
-            agent=agent,
+        run_data = RunData(
             run=Run(agent_name=agent.name, session_id=session.id),
             input=request.input,
-            history=list(session.history()),
-            executor=executor,
         )
-        session.append(bundle)
+        await run_store.set(run_data.key, run_data)
 
-        runs[bundle.run.run_id] = bundle
-        sessions[session.id] = session
+        session.append(run_data.run.run_id)
+        await session_store.set(session.id, session)
 
-        headers = {Headers.RUN_ID: str(bundle.run.run_id)}
+        headers = {Headers.RUN_ID: str(run_data.run.run_id)}
+        ready = asyncio.Event()
+
+        Executor(
+            agent=agent,
+            run_data=run_data,
+            history=await session.history(run_store),
+            run_store=run_store,
+            cancel_store=run_cancel_store,
+            resume_store=run_resume_store,
+            executor=executor,
+        ).execute(wait=ready)
 
         match request.mode:
             case RunMode.STREAM:
                 return StreamingResponse(
-                    stream_sse(bundle),
+                    stream_sse(run_data, run_store, 0, ready=ready),
                     headers=headers,
                     media_type="text/event-stream",
                 )
             case RunMode.SYNC:
-                await bundle.join()
+                await wait_util_stop(run_data, run_store, ready=ready)
                 return JSONResponse(
                     headers=headers,
-                    content=jsonable_encoder(bundle.run),
+                    content=jsonable_encoder(run_data.run),
                 )
             case RunMode.ASYNC:
+                ready.set()
                 return JSONResponse(
                     status_code=status.HTTP_202_ACCEPTED,
                     headers=headers,
-                    content=jsonable_encoder(bundle.run),
+                    content=jsonable_encoder(run_data.run),
                 )
             case _:
                 raise NotImplementedError()
 
     @app.get("/runs/{run_id}")
     async def read_run(run_id: RunId) -> RunReadResponse:
-        bundle = find_run_bundle(run_id)
+        bundle = await find_run_data(run_id)
         return bundle.run
 
     @app.get("/runs/{run_id}/events")
     async def list_run_events(run_id: RunId) -> RunEventsListResponse:
-        bundle = find_run_bundle(run_id)
+        bundle = await find_run_data(run_id)
         return RunEventsListResponse(events=bundle.events)
 
     @app.post("/runs/{run_id}")
     async def resume_run(run_id: RunId, request: RunResumeRequest) -> RunResumeResponse:
-        bundle = find_run_bundle(run_id)
+        run_data = await find_run_data(run_id)
 
-        if bundle.run.await_request is None:
+        if run_data.run.await_request is None:
             raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=f"Run {run_id} has no await request")
 
-        if bundle.run.await_request.type != request.await_resume.type:
+        if run_data.run.await_request.type != request.await_resume.type:
             raise HTTPException(
                 status_code=status.HTTP_403_FORBIDDEN,
-                detail=f"Run {run_id} is expecting resume of type {bundle.run.await_request.type}",
+                detail=f"Run {run_id} is expecting resume of type {run_data.run.await_request.type}",
             )
 
-        await bundle.resume(request.await_resume)
+        run_data.run.status = RunStatus.IN_PROGRESS
+        await run_store.set(run_data.key, run_data)
+        await run_resume_store.set(run_data.key, request.await_resume)
+
         match request.mode:
             case RunMode.STREAM:
                 return StreamingResponse(
-                    stream_sse(bundle),
+                    stream_sse(run_data, run_store, len(run_data.events)),
                     media_type="text/event-stream",
                 )
             case RunMode.SYNC:
-                await bundle.join()
-                return bundle.run
+                run_data = await wait_util_stop(run_data, run_store)
+                return run_data.run
             case RunMode.ASYNC:
                 return JSONResponse(
                     status_code=status.HTTP_202_ACCEPTED,
-                    content=jsonable_encoder(bundle.run),
+                    content=jsonable_encoder(run_data.run),
                 )
             case _:
                 raise NotImplementedError()
 
     @app.post("/runs/{run_id}/cancel")
     async def cancel_run(run_id: RunId) -> RunCancelResponse:
-        bundle = find_run_bundle(run_id)
-        if bundle.run.status.is_terminal:
+        run_data = await find_run_data(run_id)
+        if run_data.run.status.is_terminal:
             raise HTTPException(
                 status_code=status.HTTP_403_FORBIDDEN,
-                detail=f"Run in terminal status {bundle.run.status} can't be cancelled",
+                detail=f"Run in terminal status {run_data.run.status} can't be cancelled",
             )
-        await bundle.cancel()
-        return JSONResponse(status_code=status.HTTP_202_ACCEPTED, content=jsonable_encoder(bundle.run))
+        await run_cancel_store.set(run_data.key, CancelData())
+        run_data.run.status = RunStatus.CANCELLING
+        return JSONResponse(status_code=status.HTTP_202_ACCEPTED, content=jsonable_encoder(run_data.run))
 
     return app

acp_sdk/server/executor.py

@@ -0,0 +1,198 @@
+import asyncio
+import logging
+from collections.abc import AsyncIterator
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime, timezone
+from typing import Self
+
+from pydantic import BaseModel, ValidationError
+
+from acp_sdk.instrumentation import get_tracer
+from acp_sdk.models import (
+    ACPError,
+    AnyModel,
+    AwaitRequest,
+    AwaitResume,
+    Error,
+    ErrorCode,
+    Event,
+    GenericEvent,
+    Message,
+    MessageCompletedEvent,
+    MessageCreatedEvent,
+    MessagePart,
+    MessagePartEvent,
+    Run,
+    RunAwaitingEvent,
+    RunCancelledEvent,
+    RunCompletedEvent,
+    RunCreatedEvent,
+    RunFailedEvent,
+    RunInProgressEvent,
+    RunStatus,
+)
+from acp_sdk.server.agent import Agent
+from acp_sdk.server.logging import logger
+from acp_sdk.server.store import Store
+
+
+class RunData(BaseModel):
+    run: Run
+    input: list[Message]
+    events: list[Event] = []
+
+    @property
+    def key(self) -> str:
+        return str(self.run.run_id)
+
+    async def watch(self, store: Store[Self], *, ready: asyncio.Event | None = None) -> AsyncIterator[Self]:
+        async for data in store.watch(self.key, ready=ready):
+            if data is None:
+                raise RuntimeError("Missing data")
+            yield data
+            if data.run.status.is_terminal:
+                break
+
+
+class CancelData(BaseModel):
+    pass
+
+
+class Executor:
+    def __init__(
+        self,
+        *,
+        agent: Agent,
+        run_data: RunData,
+        history: list[Message],
+        executor: ThreadPoolExecutor,
+        run_store: Store[RunData],
+        cancel_store: Store[CancelData],
+        resume_store: Store[AwaitResume],
+    ) -> None:
+        self.agent = agent
+        self.history = history
+        self.run_data = run_data
+        self.executor = executor
+
+        self.run_store = run_store
+        self.cancel_store = cancel_store
+        self.resume_store = resume_store
+
+        self.logger = logging.LoggerAdapter(logger, {"run_id": str(run_data.run.run_id)})
+
+    def execute(self, *, wait: asyncio.Event) -> None:
+        self.task = asyncio.create_task(self._execute(self.run_data, executor=self.executor, wait=wait))
+        self.watcher = asyncio.create_task(self._watch_for_cancellation())
+
+    async def _push(self) -> None:
+        await self.run_store.set(self.run_data.run.run_id, self.run_data)
+
+    async def _emit(self, event: Event) -> None:
+        freeze = event.model_copy(deep=True)
+        self.run_data.events.append(freeze)
+        await self._push()
+
+    async def _await(self) -> AwaitResume:
+        async for resume in self.resume_store.watch(self.run_data.key):
+            if resume is not None:
+                await self.resume_store.set(self.run_data.key, None)
+                return resume
+
+    async def _watch_for_cancellation(self) -> None:
+        while not self.task.done():
+            try:
+                async for data in self.cancel_store.watch(self.run_data.key):
+                    if data is not None:
+                        self.task.cancel()
+            except Exception:
+                logger.warning("Cancellation watcher failed, restarting")
+
+    async def _execute(self, run_data: RunData, *, executor: ThreadPoolExecutor, wait: asyncio.Event) -> None:
+        with get_tracer().start_as_current_span("run"):
+            in_message = False
+
+            async def flush_message() -> None:
+                nonlocal in_message
+                if in_message:
+                    message = run_data.run.output[-1]
+                    message.completed_at = datetime.now(timezone.utc)
+                    await self._emit(MessageCompletedEvent(message=message))
+                    in_message = False
+
+            try:
+                await wait.wait()
+
+                await self._emit(RunCreatedEvent(run=run_data.run))
+
+                generator = self.agent.execute(
+                    input=self.history + run_data.input, session_id=run_data.run.session_id, executor=executor
+                )
+                self.logger.info("Run started")
+
+                run_data.run.status = RunStatus.IN_PROGRESS
+                await self._emit(RunInProgressEvent(run=run_data.run))
+
+                await_resume = None
+                while True:
+                    next = await generator.asend(await_resume)
+
+                    if isinstance(next, (MessagePart, str)):
+                        if isinstance(next, str):
+                            next = MessagePart(content=next)
+                        if not in_message:
+                            run_data.run.output.append(Message(parts=[], completed_at=None))
+                            in_message = True
+                            await self._emit(MessageCreatedEvent(message=run_data.run.output[-1]))
+                        run_data.run.output[-1].parts.append(next)
+                        await self._emit(MessagePartEvent(part=next))
+                    elif isinstance(next, Message):
+                        await flush_message()
+                        run_data.run.output.append(next)
+                        await self._emit(MessageCreatedEvent(message=next))
+                        for part in next.parts:
+                            await self._emit(MessagePartEvent(part=part))
+                        await self._emit(MessageCompletedEvent(message=next))
+                    elif isinstance(next, AwaitRequest):
+                        run_data.run.await_request = next
+                        run_data.run.status = RunStatus.AWAITING
+                        await self._emit(RunAwaitingEvent(run=run_data.run))
+                        self.logger.info("Run awaited")
+                        await_resume = await self._await()
+                        run_data.run.status = RunStatus.IN_PROGRESS
+                        await self._emit(RunInProgressEvent(run=run_data.run))
+                        self.logger.info("Run resumed")
+                    elif isinstance(next, Error):
+                        raise ACPError(error=next)
+                    elif isinstance(next, BaseException):
+                        raise next
+                    elif next is None:
+                        await flush_message()
+                    elif isinstance(next, BaseModel):
+                        await self._emit(GenericEvent(generic=AnyModel(**next.model_dump())))
+                    else:
+                        try:
+                            generic = AnyModel.model_validate(next)
+                            await self._emit(GenericEvent(generic=generic))
+                        except ValidationError:
+                            raise TypeError("Invalid yield")
+            except StopAsyncIteration:
+                await flush_message()
+                run_data.run.status = RunStatus.COMPLETED
+                run_data.run.finished_at = datetime.now(timezone.utc)
+                await self._emit(RunCompletedEvent(run=run_data.run))
+                self.logger.info("Run completed")
+            except asyncio.CancelledError:
+                run_data.run.status = RunStatus.CANCELLED
+                run_data.run.finished_at = datetime.now(timezone.utc)
+                await self._emit(RunCancelledEvent(run=run_data.run))
+                self.logger.info("Run cancelled")
+            except Exception as e:
+                if isinstance(e, ACPError):
+                    run_data.run.error = e.error
+                else:
+                    run_data.run.error = Error(code=ErrorCode.SERVER_ERROR, message=str(e))
+                run_data.run.status = RunStatus.FAILED
+                run_data.run.finished_at = datetime.now(timezone.utc)
+                await self._emit(RunFailedEvent(run=run_data.run))
+                self.logger.exception("Run failed")

acp_sdk/server/server.py

@@ -1,8 +1,8 @@
 import asyncio
 import os
+import re
 from collections.abc import AsyncGenerator, Awaitable
 from contextlib import asynccontextmanager
-from datetime import timedelta
 from typing import Any, Callable
 
 import requests
@@ -16,6 +16,7 @@ from acp_sdk.server.agent import agent as agent_decorator
 from acp_sdk.server.app import create_app
 from acp_sdk.server.logging import configure_logger as configure_logger_func
 from acp_sdk.server.logging import logger
+from acp_sdk.server.store import Store
 from acp_sdk.server.telemetry import configure_telemetry as configure_telemetry_func
 from acp_sdk.server.utils import async_request_with_retry
 
@@ -54,8 +55,7 @@ class Server:
         configure_logger: bool = True,
         configure_telemetry: bool = False,
         self_registration: bool = True,
-        run_limit: int = 1000,
-        run_ttl: timedelta = timedelta(hours=1),
+        store: Store | None = None,
         host: str = "127.0.0.1",
         port: int = 8000,
         uds: str | None = None,
@@ -118,7 +118,7 @@ class Server:
 
         import uvicorn
 
-        app = create_app(*self.agents, lifespan=self.lifespan, run_limit=run_limit, run_ttl=run_ttl)
+        app = create_app(*self.agents, lifespan=self.lifespan, store=store)
 
         if configure_logger:
             configure_logger_func()
@@ -184,8 +184,7 @@ class Server:
         configure_logger: bool = True,
         configure_telemetry: bool = False,
         self_registration: bool = True,
-        run_limit: int = 1000,
-        run_ttl: timedelta = timedelta(hours=1),
+        store: Store | None = None,
         host: str = "127.0.0.1",
         port: int = 8000,
         uds: str | None = None,
@@ -243,8 +242,7 @@ class Server:
                 configure_logger=configure_logger,
                 configure_telemetry=configure_telemetry,
                 self_registration=self_registration,
-                run_limit=run_limit,
-                run_ttl=run_ttl,
+                store=store,
                 host=host,
                 port=port,
                 uds=uds,
@@ -311,27 +309,38 @@ class Server:
 
     async def _register_agent(self) -> None:
         """If not in PRODUCTION mode, register agent to the beeai platform and provide missing env variables"""
-        if os.getenv("PRODUCTION_MODE", False):
+        if os.getenv("PRODUCTION_MODE", "").lower() in ["true", "1"]:
             logger.debug("Agent is not automatically registered in the production mode.")
             return
 
         url = os.getenv("PLATFORM_URL", "http://127.0.0.1:8333")
+        host = re.sub(r"localhost|127\.0\.0\.1", "host.docker.internal", self.server.config.host)
         request_data = {
-            "location": f"http://{self.server.config.host}:{self.server.config.port}",
+            "location": f"http://{host}:{self.server.config.port}",
         }
+        await async_request_with_retry(lambda client, data=request_data: client.get(f"{url}/api/v1/providers"))
         try:
             await async_request_with_retry(
-                lambda client, data=request_data: client.post(f"{url}/api/v1/providers/register/unmanaged", json=data)
+                lambda client, data=request_data: client.post(
+                    f"{url}/api/v1/providers", json=data, params={"auto_remove": True}
+                )
             )
             logger.info("Agent registered to the beeai server.")
 
             # check missing env keyes
             envs_request = await async_request_with_retry(lambda client: client.get(f"{url}/api/v1/variables"))
             envs = envs_request.get("env")
+            os.environ["LLM_MODEL"] = "dummy"
+            os.environ["LLM_API_KEY"] = "dummy"
+            os.environ["LLM_API_BASE"] = f"{url.rstrip('/')}/api/v1/llm"
+
             for agent in self.agents:
                 # register all available envs
                 missing_keyes = []
                 for env in agent.metadata.model_dump().get("env", []):
+                    # Those envs are set to use LLM gateway from platform server
+                    if env["name"] in {"LLM_MODEL", "LLM_API_KEY", "LLM_API_BASE"}:
+                        continue
                     server_env = envs.get(env.get("name"))
                     if server_env:
                         logger.debug(f"Env variable {env['name']} = '{server_env}' added dynamically")

acp_sdk/server/session.py

@@ -1,21 +1,24 @@
 import uuid
-from collections.abc import Iterator
 
-from acp_sdk.models import Message, SessionId
-from acp_sdk.models.models import RunStatus
-from acp_sdk.server.bundle import RunBundle
+from pydantic import BaseModel, Field
 
+from acp_sdk.models import Message, RunId, RunStatus, SessionId
+from acp_sdk.server.executor import RunData
+from acp_sdk.server.store import Store
 
-class Session:
-    def __init__(self, id: SessionId | None = None) -> None:
-        self.id: SessionId = id or uuid.uuid4()
-        self.bundles: list[RunBundle] = []
 
-    def append(self, bundle: RunBundle) -> None:
-        self.bundles.append(bundle)
+class Session(BaseModel):
+    id: SessionId = Field(default_factory=uuid.uuid4)
+    runs: list[RunId] = []
 
-    def history(self) -> Iterator[Message]:
-        for bundle in self.bundles:
-            if bundle.run.status == RunStatus.COMPLETED:
-                yield from bundle.input
-                yield from bundle.run.output
+    def append(self, run_id: RunId) -> None:
+        self.runs.append(run_id)
+
+    async def history(self, store: Store[RunData]) -> list[Message]:
+        history = []
+        for run_id in self.runs:
+            run_data = await store.get(run_id)
+            if run_data is not None and run_data.run.status == RunStatus.COMPLETED:
+                history.extend(run_data.input)
+                history.extend(run_data.run.output)
+        return history

acp_sdk/server/store/__init__.py

@@ -0,0 +1,4 @@
+from acp_sdk.server.store.memory_store import MemoryStore as MemoryStore
+from acp_sdk.server.store.postgresql_store import PostgreSQLStore as PostgreSQLStore
+from acp_sdk.server.store.redis_store import RedisStore as RedisStore
+from acp_sdk.server.store.store import Store as Store

acp_sdk/server/store/memory_store.py

@@ -0,0 +1,35 @@
+import asyncio
+from collections.abc import AsyncIterator
+from datetime import datetime
+from typing import Generic
+
+from cachetools import TTLCache
+
+from acp_sdk.server.store.store import Store, T
+from acp_sdk.server.store.utils import Stringable
+
+
+class MemoryStore(Store[T], Generic[T]):
+    def __init__(self, *, limit: int, ttl: int | None = None) -> None:
+        super().__init__()
+        self._cache: TTLCache[str, T] = TTLCache(maxsize=limit, ttl=ttl, timer=datetime.now)
+        self._event = asyncio.Event()
+
+    async def get(self, key: Stringable) -> T | None:
+        value = self._cache.get(str(key))
+        return value.model_copy(deep=True) if value else value
+
+    async def set(self, key: Stringable, value: T | None) -> None:
+        if value is None:
+            del self._cache[str(key)]
+        else:
+            self._cache[str(key)] = value.model_copy(deep=True)
+        self._event.set()
+
+    async def watch(self, key: Stringable, *, ready: asyncio.Event | None = None) -> AsyncIterator[T | None]:
+        if ready:
+            ready.set()
+        while True:
+            await self._event.wait()
+            self._event.clear()
+            yield await self.get(key)

acp_sdk/server/store/postgresql_store.py

@@ -0,0 +1,69 @@
+import asyncio
+from collections.abc import AsyncIterator
+from typing import Generic
+
+from psycopg import AsyncConnection
+from psycopg.rows import dict_row
+
+from acp_sdk.server.store.store import Store, StoreModel, T
+from acp_sdk.server.store.utils import Stringable
+
+
+class PostgreSQLStore(Store[T], Generic[T]):
+    def __init__(self, *, aconn: AsyncConnection, table: str = "acp_store", channel: str = "acp_update") -> None:
+        super().__init__()
+        self._aconn = aconn
+        self._table = table
+        self._channel = channel
+
+    async def get(self, key: Stringable) -> T | None:
+        await self._ensure_table()
+        async with self._aconn.cursor(row_factory=dict_row) as cur:
+            await cur.execute(f"SELECT value FROM {self._table} WHERE key = %s", (str(key),))
+            result = await cur.fetchone()
+            if result is None:
+                return None
+            return StoreModel.model_validate(result["value"])
+
+    async def set(self, key: Stringable, value: T | None) -> None:
+        await self._ensure_table()
+        async with self._aconn.cursor() as cur:
+            if value is None:
+                await cur.execute(
+                    f"DELETE FROM {self._table} WHERE key = %s",
+                    (str(key),),
+                )
+            else:
+                await cur.execute(
+                    f"""
+                    INSERT INTO {self._table} (key, value)
+                    VALUES (%s, %s)
+                    ON CONFLICT (key)
+                    DO UPDATE SET value = EXCLUDED.value
+                    """,
+                    (str(key), value.model_dump_json()),
+                )
+            await cur.execute(f"NOTIFY {self._channel}, '{key!s}'")  # NOTIFY appears not to accept params
+            await self._aconn.commit()
+
+    async def watch(self, key: Stringable, *, ready: asyncio.Event | None = None) -> AsyncIterator[T | None]:
+        notify_conn = await AsyncConnection.connect(
+            conninfo=f"{self._aconn.info.dsn} password={self._aconn.info.password}", autocommit=True
+        )
+        async with notify_conn:
+            await notify_conn.execute(f"LISTEN {self._channel}")
+            if ready:
+                ready.set()
+            async for notify in notify_conn.notifies():
+                if notify.payload == str(key):
+                    yield await self.get(key)
+
+    async def _ensure_table(self) -> None:
+        async with self._aconn.cursor() as cur:
+            await cur.execute(f"""
+                CREATE TABLE IF NOT EXISTS {self._table} (
+                    key TEXT PRIMARY KEY,
+                    value JSONB NOT NULL
+                )
+            """)
+            await self._aconn.commit()

acp_sdk/server/store/redis_store.py

@@ -0,0 +1,40 @@
+import asyncio
+from collections.abc import AsyncIterator
+from typing import Generic
+
+from redis.asyncio import Redis
+
+from acp_sdk.server.store.store import Store, StoreModel, T
+from acp_sdk.server.store.utils import Stringable
+
+
+class RedisStore(Store[T], Generic[T]):
+    def __init__(self, *, redis: Redis) -> None:
+        super().__init__()
+        self._redis = redis
+
+    async def get(self, key: Stringable) -> T | None:
+        value = await self._redis.get(str(key))
+        return StoreModel.model_validate_json(value) if value else value
+
+    async def set(self, key: Stringable, value: T | None) -> None:
+        if value is None:
+            await self._redis.delete(str(key))
+        else:
+            await self._redis.set(name=str(key), value=value.model_dump_json())
+
+    async def watch(self, key: Stringable, *, ready: asyncio.Event | None = None) -> AsyncIterator[T]:
+        await self._redis.config_set("notify-keyspace-events", "KEA")
+
+        pubsub = self._redis.pubsub()
+        channel = f"__keyspace@0__:{key!s}"
+        await pubsub.subscribe(channel)
+        if ready:
+            ready.set()
+        try:
+            async for message in pubsub.listen():
+                if message["type"] == "message":
+                    yield await self.get(key)
+        finally:
+            await pubsub.unsubscribe(channel)
+            await pubsub.close()

acp_sdk/server/store/store.py

@@ -0,0 +1,55 @@
+import asyncio
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator
+from typing import Generic, TypeVar
+
+from pydantic import BaseModel, ConfigDict
+
+from acp_sdk.server.store.utils import Stringable
+
+
+class StoreModel(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+
+T = TypeVar("T", bound=BaseModel)
+U = TypeVar("U", bound=BaseModel)
+
+
+class Store(Generic[T], ABC):
+    @abstractmethod
+    async def get(self, key: Stringable) -> T | None:
+        pass
+
+    @abstractmethod
+    async def set(self, key: Stringable, value: T | None) -> None:
+        pass
+
+    @abstractmethod
+    def watch(self, key: Stringable, *, ready: asyncio.Event | None = None) -> AsyncIterator[T | None]:
+        pass
+
+    def as_store(self, model: type[U], prefix: Stringable = "") -> "Store[U]":
+        return StoreView(model=model, store=self, prefix=prefix)
+
+
+class StoreView(Store[U], Generic[U]):
+    def __init__(self, *, model: type[U], store: Store[T], prefix: Stringable = "") -> None:
+        super().__init__()
+        self._model = model
+        self._store = store
+        self._prefix = prefix
+
+    async def get(self, key: Stringable) -> U | None:
+        value = await self._store.get(self._get_key(key))
+        return self._model.model_validate(value.model_dump()) if value else value
+
+    async def set(self, key: Stringable, value: U | None) -> None:
+        await self._store.set(self._get_key(key), value)
+
+    async def watch(self, key: Stringable, *, ready: asyncio.Event | None = None) -> AsyncIterator[U | None]:
+        async for value in self._store.watch(self._get_key(key), ready=ready):
+            yield self._model.model_validate(value.model_dump()) if value else value
+
+    def _get_key(self, key: Stringable) -> str:
+        return f"{self._prefix!s}{key!s}"

acp_sdk/server/store/utils.py

@@ -0,0 +1,5 @@
+from typing import Protocol
+
+
+class Stringable(Protocol):
+    def __str__(self) -> str: ...

acp_sdk/server/utils.py

@@ -6,17 +6,41 @@ import httpx
 import requests
 from pydantic import BaseModel
 
-from acp_sdk.server.bundle import RunBundle
+from acp_sdk.models import RunStatus
+from acp_sdk.server.executor import RunData
 from acp_sdk.server.logging import logger
+from acp_sdk.server.store.store import Store
 
 
 def encode_sse(model: BaseModel) -> str:
     return f"data: {model.model_dump_json()}\n\n"
 
 
-async def stream_sse(bundle: RunBundle) -> AsyncGenerator[str]:
-    async for event in bundle.stream():
-        yield encode_sse(event)
+async def watch_util_stop(
+    run_data: RunData, store: Store[RunData], *, ready: asyncio.Event | None = None
+) -> AsyncGenerator[RunData]:
+    async for data in run_data.watch(store, ready=ready):
+        yield data
+        if data.run.status == RunStatus.AWAITING:
+            break
+
+
+async def wait_util_stop(run_data: RunData, store: Store[RunData], *, ready: asyncio.Event | None = None) -> RunData:
+    data = run_data
+    async for latest_data in watch_util_stop(run_data, store, ready=ready):
+        data = latest_data
+    return data
+
+
+async def stream_sse(
+    run_data: RunData, store: Store[RunData], idx: int, *, ready: asyncio.Event | None = None
+) -> AsyncGenerator[str]:
+    next_event_idx = idx
+    async for data in watch_util_stop(run_data, store, ready=ready):
+        new_events = data.events[next_event_idx:]
+        next_event_idx = len(data.events)
+        for event in new_events:
+            yield encode_sse(event)
 
 
 async def async_request_with_retry(

{acp_sdk-0.9.1.dist-info → acp_sdk-0.10.1.dist-info}/METADATA

@@ -1,27 +1,29 @@
 Metadata-Version: 2.4
 Name: acp-sdk
-Version: 0.9.1
+Version: 0.10.1
 Summary: Agent Communication Protocol SDK
 Author: IBM Corp.
 Maintainer-email: Tomas Pilar <thomas7pilar@gmail.com>
 License-Expression: Apache-2.0
 Requires-Python: <4.0,>=3.11
-Requires-Dist: cachetools>=5.5.2
-Requires-Dist: fastapi[standard]>=0.115.8
-Requires-Dist: httpx-sse>=0.4.0
-Requires-Dist: httpx>=0.26.0
-Requires-Dist: janus>=2.0.0
-Requires-Dist: opentelemetry-api>=1.31.1
-Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.31.1
+Requires-Dist: cachetools>=5.5
+Requires-Dist: fastapi[standard]>=0.115
+Requires-Dist: httpx-sse>=0.4
+Requires-Dist: httpx>=0.26
+Requires-Dist: janus>=2.0
+Requires-Dist: opentelemetry-api>=1.31
+Requires-Dist: opentelemetry-exporter-otlp-proto-http>=1.31
 Requires-Dist: opentelemetry-instrumentation-fastapi>=0.52b1
 Requires-Dist: opentelemetry-instrumentation-httpx>=0.52b1
-Requires-Dist: opentelemetry-sdk>=1.31.1
-Requires-Dist: pydantic>=2.0.0
+Requires-Dist: opentelemetry-sdk>=1.31
+Requires-Dist: psycopg[binary]>=3.2
+Requires-Dist: pydantic>=2.0
+Requires-Dist: redis>=6.1
 Description-Content-Type: text/markdown
 
 # Agent Communication Protocol SDK for Python
 
-Agent Communication Protocol SDK for Python provides allows developers to serve and consume agents over the Agent Communication Protocol.
+Agent Communication Protocol SDK for Python helps developers to serve and consume agents over the Agent Communication Protocol.
 
 ## Prerequisites
 

{acp_sdk-0.9.1.dist-info → acp_sdk-0.10.1.dist-info}/RECORD

@@ -12,16 +12,22 @@ acp_sdk/models/models.py,sha256=So5D9VjJxX2jdOF8BGc3f9D2ssbAyFEO2Rgr1KykI-M,7730
 acp_sdk/models/schemas.py,sha256=_ah7_zHsQJGxDXvnzsBvASdRsQHVphFQ7Sum6A04iRw,759
 acp_sdk/server/__init__.py,sha256=mxBBBFaZuMEUENRMLwp1XZkuLeT9QghcFmNvjnqvAAU,377
 acp_sdk/server/agent.py,sha256=6VBKn_qVXqUl79G8T7grwhnuLMwr67d4UGagMGX1hMs,6586
-acp_sdk/server/app.py,sha256=ViumwzMngn63wtpFllu3nnIOt9JL9PQYrTcI8v-_uY4,7257
-acp_sdk/server/bundle.py,sha256=umD2GgDp17lUddu0adpp1zUcm1JJvDrDpIZ0uR-6VeY,7204
+acp_sdk/server/app.py,sha256=XeMn8hfU2J2O2QxSvsH5q6lowxCfwS1-DTWz-gUVotA,8733
 acp_sdk/server/context.py,sha256=MgnLV6qcDIhc_0BjW7r4Jj1tHts4ZuwpdTGIBnz2Mgo,1036
 acp_sdk/server/errors.py,sha256=GSO8yYIqEeX8Y4Lz86ks35dMTHiQiXuOrLYYx0eXsbI,2110
+acp_sdk/server/executor.py,sha256=YL0J9cVY1QZtdTeqwjJaKDpB_T6_sByHlHc52kgNAJo,7742
 acp_sdk/server/logging.py,sha256=Oc8yZigCsuDnHHPsarRzu0RX3NKaLEgpELM2yovGKDI,411
-acp_sdk/server/server.py,sha256=0thIxmqLV-hqwSqP_TUSZAl9pog06lROAQgnSIIabME,13621
-acp_sdk/server/session.py,sha256=ekz1o6Sy1tQZlpaoS_VgbvFuUQh2qpiHG71mvBdvhgc,662
+acp_sdk/server/server.py,sha256=slePybTpocXD8k9HxOvxCiyvigiDUdZmkz3kiI-dCFU,14113
+acp_sdk/server/session.py,sha256=vGUVpKzUGefI1c7LeK08Bvd8zvJIRfsdJEt2KhYoEg0,764
 acp_sdk/server/telemetry.py,sha256=lbB2ppijUcqbHUOn0e-15LGcVvT_qrMguq8qBokICac,2016
 acp_sdk/server/types.py,sha256=gLb5wCkMYhmu2laj_ymK-TPfN9LSjRgKOP1H_893UzA,304
-acp_sdk/server/utils.py,sha256=BhZKBNaLgczX6aYjxYva-6VI1bKmHtYQ5YDA5LrwF50,1831
-acp_sdk-0.9.1.dist-info/METADATA,sha256=fdUTU2kHk-OakxDoaTCLUagieJuM8DjsKZiFsOzDEK8,1650
-acp_sdk-0.9.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-acp_sdk-0.9.1.dist-info/RECORD,,
+acp_sdk/server/utils.py,sha256=BYSn4Bd95Bn-oEH1W1yE_pWpYUOdtYPh-vMnou4nsdk,2721
+acp_sdk/server/store/__init__.py,sha256=zzKic0byQTM86cyC2whwZeNP4prfy_HZrbSriTaV5j8,282
+acp_sdk/server/store/memory_store.py,sha256=9hOoJiHdaDz-0hNHWhO_qZpp-I_pnTqYBIhQnlWilu4,1194
+acp_sdk/server/store/postgresql_store.py,sha256=bHyAgf4vSn_07wGXXq6juFwm3JldYNOjU9RARcDSYQo,2717
+acp_sdk/server/store/redis_store.py,sha256=IKXvDseOFMcoGjVYPOkOBhPnJAchy_RyeMayKLoVCGA,1378
+acp_sdk/server/store/store.py,sha256=jGmYy9oiuVjhYYJY8QRo4g2J2Qyt1HLTmq_eHy4aI7c,1806
+acp_sdk/server/store/utils.py,sha256=JumEOMs1h1uGlnHnUGeguee-srGzT7_Y2NVEYt01QuY,92
+acp_sdk-0.10.1.dist-info/METADATA,sha256=2CkcWAS4e-6bIeoFsv2eyxLcbeTSgy15gq3VCP1dxbE,1685
+acp_sdk-0.10.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+acp_sdk-0.10.1.dist-info/RECORD,,

acp_sdk/server/bundle.py

@@ -1,182 +0,0 @@
-import asyncio
-import logging
-from collections.abc import AsyncGenerator
-from concurrent.futures import ThreadPoolExecutor
-from datetime import datetime, timezone
-
-from pydantic import BaseModel, ValidationError
-
-from acp_sdk.instrumentation import get_tracer
-from acp_sdk.models import (
-    ACPError,
-    AnyModel,
-    AwaitRequest,
-    AwaitResume,
-    Error,
-    ErrorCode,
-    Event,
-    GenericEvent,
-    Message,
-    MessageCompletedEvent,
-    MessageCreatedEvent,
-    MessagePart,
-    MessagePartEvent,
-    Run,
-    RunAwaitingEvent,
-    RunCancelledEvent,
-    RunCompletedEvent,
-    RunCreatedEvent,
-    RunFailedEvent,
-    RunInProgressEvent,
-    RunStatus,
-)
-from acp_sdk.server.agent import Agent
-from acp_sdk.server.logging import logger
-
-
-class RunBundle:
-    def __init__(
-        self, *, agent: Agent, run: Run, input: list[Message], history: list[Message], executor: ThreadPoolExecutor
-    ) -> None:
-        self.agent = agent
-        self.run = run
-        self.input = input
-        self.history = history
-
-        self.stream_queue: asyncio.Queue[Event] = asyncio.Queue()
-        self.events: list[Event] = []
-
-        self.await_queue: asyncio.Queue[AwaitResume] = asyncio.Queue(maxsize=1)
-        self.await_or_terminate_event = asyncio.Event()
-
-        self.task = asyncio.create_task(self._execute(input, executor=executor))
-
-    async def stream(self) -> AsyncGenerator[Event]:
-        while True:
-            event = await self.stream_queue.get()
-            if event is None:
-                break
-            yield event
-            self.stream_queue.task_done()
-
-    async def emit(self, event: Event) -> None:
-        freeze = event.model_copy(deep=True)
-        self.events.append(freeze)
-        await self.stream_queue.put(freeze)
-
-    async def await_(self) -> AwaitResume:
-        await self.stream_queue.put(None)
-        self.await_queue.empty()
-        self.await_or_terminate_event.set()
-        self.await_or_terminate_event.clear()
-        resume = await self.await_queue.get()
-        self.await_queue.task_done()
-        return resume
-
-    async def resume(self, resume: AwaitResume) -> None:
-        self.stream_queue = asyncio.Queue()
-        await self.await_queue.put(resume)
-        self.run.status = RunStatus.IN_PROGRESS
-        self.run.await_request = None
-
-    async def cancel(self) -> None:
-        self.task.cancel()
-        self.run.status = RunStatus.CANCELLING
-        self.run.await_request = None
-
-    async def join(self) -> None:
-        await self.await_or_terminate_event.wait()
-
-    async def _execute(self, input: list[Message], *, executor: ThreadPoolExecutor) -> None:
-        with get_tracer().start_as_current_span("run"):
-            run_logger = logging.LoggerAdapter(logger, {"run_id": str(self.run.run_id)})
-
-            in_message = False
-
-            async def flush_message() -> None:
-                nonlocal in_message
-                if in_message:
-                    message = self.run.output[-1]
-                    message.completed_at = datetime.now(timezone.utc)
-                    await self.emit(MessageCompletedEvent(message=message))
-                    in_message = False
-
-            try:
-                await self.emit(RunCreatedEvent(run=self.run))
-
-                generator = self.agent.execute(
-                    input=self.history + input, session_id=self.run.session_id, executor=executor
-                )
-                run_logger.info("Run started")
-
-                self.run.status = RunStatus.IN_PROGRESS
-                await self.emit(RunInProgressEvent(run=self.run))
-
-                await_resume = None
-                while True:
-                    next = await generator.asend(await_resume)
-
-                    if isinstance(next, (MessagePart, str)):
-                        if isinstance(next, str):
-                            next = MessagePart(content=next)
-                        if not in_message:
-                            self.run.output.append(Message(parts=[], completed_at=None))
-                            in_message = True
-                            await self.emit(MessageCreatedEvent(message=self.run.output[-1]))
-                        self.run.output[-1].parts.append(next)
-                        await self.emit(MessagePartEvent(part=next))
-                    elif isinstance(next, Message):
-                        await flush_message()
-                        self.run.output.append(next)
-                        await self.emit(MessageCreatedEvent(message=next))
-                        for part in next.parts:
-                            await self.emit(MessagePartEvent(part=part))
-                        await self.emit(MessageCompletedEvent(message=next))
-                    elif isinstance(next, AwaitRequest):
-                        self.run.await_request = next
-                        self.run.status = RunStatus.AWAITING
-                        await self.emit(RunAwaitingEvent(run=self.run))
-                        run_logger.info("Run awaited")
-                        await_resume = await self.await_()
-                        await self.emit(RunInProgressEvent(run=self.run))
-                        run_logger.info("Run resumed")
-                    elif isinstance(next, Error):
-                        raise ACPError(error=next)
-                    elif isinstance(next, ACPError):
-                        raise next
-                    elif next is None:
-                        await flush_message()
-                    elif isinstance(next, BaseModel):
-                        await self.emit(GenericEvent(generic=AnyModel(**next.model_dump())))
-                    else:
-                        try:
-                            generic = AnyModel.model_validate(next)
-                            await self.emit(GenericEvent(generic=generic))
-                        except ValidationError:
-                            raise TypeError("Invalid yield")
-            except StopAsyncIteration:
-                await flush_message()
-                self.run.status = RunStatus.COMPLETED
-                self.run.finished_at = datetime.now(timezone.utc)
-                await self.emit(RunCompletedEvent(run=self.run))
-                run_logger.info("Run completed")
-            except asyncio.CancelledError:
-                self.run.status = RunStatus.CANCELLED
-                self.run.finished_at = datetime.now(timezone.utc)
-                await self.emit(RunCancelledEvent(run=self.run))
-                run_logger.info("Run cancelled")
-            except Exception as e:
-                if isinstance(e, ACPError):
-                    self.run.error = e.error
-                else:
-                    self.run.error = Error(code=ErrorCode.SERVER_ERROR, message=str(e))
-                self.run.status = RunStatus.FAILED
-                self.run.finished_at = datetime.now(timezone.utc)
-                await self.emit(RunFailedEvent(run=self.run))
-                run_logger.exception("Run failed")
-                raise
-            finally:
-                self.await_or_terminate_event.set()
-                await self.stream_queue.put(None)
-                if not self.task.done():
-                    self.task.cancel()