acp-sdk 0.10.0__py3-none-any.whl → 0.11.0__py3-none-any.whl

acp_sdk/client/client.py

@@ -1,8 +1,8 @@
+import asyncio
+import logging
 import ssl
 import typing
-import uuid
-from collections.abc import AsyncGenerator, AsyncIterator
-from contextlib import asynccontextmanager
+from collections.abc import AsyncIterator
 from types import TracebackType
 from typing import Self
 
@@ -15,12 +15,13 @@ from acp_sdk.client.utils import input_to_messages
 from acp_sdk.instrumentation import get_tracer
 from acp_sdk.models import (
     ACPError,
-    Agent,
+    AgentManifest,
     AgentName,
     AgentReadResponse,
     AgentsListResponse,
     AwaitResume,
     Error,
+    ErrorCode,
     ErrorEvent,
     Event,
     PingResponse,
@@ -33,16 +34,20 @@ from acp_sdk.models import (
     RunMode,
     RunResumeRequest,
     RunResumeResponse,
-    SessionId,
+    Session,
+    SessionReadResponse,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class Client:
     def __init__(
         self,
         *,
-        session_id: SessionId | None = None,
+        session: Session | None = None,
         client: httpx.AsyncClient | None = None,
+        manage_client: bool = True,
         auth: httpx._types.AuthTypes | None = None,
         params: httpx._types.QueryParamTypes | None = None,
         headers: httpx._types.HeaderTypes | None = None,
@@ -62,7 +67,10 @@ class Client:
         transport: httpx.AsyncBaseTransport | None = None,
         trust_env: bool = True,
     ) -> None:
-        self._session_id = session_id
+        self._session = session
+        self._session_last_refresh_base_url: httpx.URL | None = None
+        self._session_refresh_lock = asyncio.Lock()
+
         self._client = client or httpx.AsyncClient(
             auth=auth,
             params=params,
@@ -83,13 +91,24 @@ class Client:
             transport=transport,
             trust_env=trust_env,
         )
+        self._manage_client = manage_client
 
     @property
     def client(self) -> httpx.AsyncClient:
         return self._client
 
     async def __aenter__(self) -> Self:
-        await self._client.__aenter__()
+        if self._manage_client:
+            await self._client.__aenter__()
+        self._session_span_manager = (
+            (
+                get_tracer()
+                .start_as_current_span("session", attributes={"acp.session": str(self._session.id)})
+                .__enter__()
+            )
+            if self._session
+            else None
+        )
         return self
 
     async def __aexit__(
@@ -98,121 +117,152 @@ class Client:
         exc_value: BaseException | None = None,
         traceback: TracebackType | None = None,
     ) -> None:
-        await self._client.__aexit__(exc_type, exc_value, traceback)
+        if self._session_span_manager:
+            self._session_span_manager.__exit__(exc_type, exc_value, traceback)
+        if self._manage_client:
+            await self._client.__aexit__(exc_type, exc_value, traceback)
 
-    @asynccontextmanager
-    async def session(self, session_id: SessionId | None = None) -> AsyncGenerator[Self]:
-        session_id = session_id or uuid.uuid4()
-        with get_tracer().start_as_current_span("session", attributes={"acp.session": str(session_id)}):
-            yield Client(client=self._client, session_id=session_id)
+    def session(self, session: Session | None = None) -> Self:
+        return Client(client=self._client, manage_client=False, session=session or Session())
 
-    async def agents(self) -> AsyncIterator[Agent]:
-        response = await self._client.get("/agents")
+    async def agents(self, *, base_url: httpx.URL | str | None = None) -> AsyncIterator[AgentManifest]:
+        response = await self._client.get(self._create_url("/agents", base_url=base_url))
         self._raise_error(response)
         for agent in AgentsListResponse.model_validate(response.json()).agents:
             yield agent
 
-    async def agent(self, *, name: AgentName) -> Agent:
-        response = await self._client.get(f"/agents/{name}")
+    async def agent(self, *, name: AgentName, base_url: httpx.URL | str | None = None) -> AgentManifest:
+        response = await self._client.get(self._create_url(f"/agents/{name}", base_url=base_url))
         self._raise_error(response)
         response = AgentReadResponse.model_validate(response.json())
-        return Agent(**response.model_dump())
+        return AgentManifest(**response.model_dump())
 
-    async def ping(self) -> bool:
-        response = await self._client.get("/ping")
+    async def ping(self, *, base_url: httpx.URL | str | None = None) -> bool:
+        response = await self._client.get(self._create_url("/ping", base_url=base_url))
         self._raise_error(response)
         PingResponse.model_validate(response.json())
         return
 
-    async def run_sync(self, input: Input, *, agent: AgentName) -> Run:
+    async def run_sync(self, input: Input, *, agent: AgentName, base_url: httpx.URL | str | None = None) -> Run:
         response = await self._client.post(
-            "/runs",
+            self._create_url("/runs", base_url=base_url),
             content=RunCreateRequest(
                 agent_name=agent,
                 input=input_to_messages(input),
                 mode=RunMode.SYNC,
-                session_id=self._session_id,
+                **(await self._prepare_session_for_run(base_url=base_url)),
             ).model_dump_json(),
         )
         self._raise_error(response)
         response = RunCreateResponse.model_validate(response.json())
         return Run(**response.model_dump())
 
-    async def run_async(self, input: Input, *, agent: AgentName) -> Run:
+    async def run_async(self, input: Input, *, agent: AgentName, base_url: httpx.URL | str | None = None) -> Run:
         response = await self._client.post(
-            "/runs",
+            self._create_url("/runs", base_url=base_url),
             content=RunCreateRequest(
                 agent_name=agent,
                 input=input_to_messages(input),
                 mode=RunMode.ASYNC,
-                session_id=self._session_id,
+                **(await self._prepare_session_for_run(base_url=base_url)),
             ).model_dump_json(),
         )
         self._raise_error(response)
         response = RunCreateResponse.model_validate(response.json())
         return Run(**response.model_dump())
 
-    async def run_stream(self, input: Input, *, agent: AgentName) -> AsyncIterator[Event]:
+    async def run_stream(
+        self, input: Input, *, agent: AgentName, base_url: httpx.URL | str | None = None
+    ) -> AsyncIterator[Event]:
         async with aconnect_sse(
             self._client,
             "POST",
-            "/runs",
+            self._create_url("/runs", base_url=base_url),
             content=RunCreateRequest(
                 agent_name=agent,
                 input=input_to_messages(input),
                 mode=RunMode.STREAM,
-                session_id=self._session_id,
+                session=await self._prepare_session_for_run(base_url=base_url),
             ).model_dump_json(),
         ) as event_source:
             async for event in self._validate_stream(event_source):
                 yield event
 
-    async def run_status(self, *, run_id: RunId) -> Run:
-        response = await self._client.get(f"/runs/{run_id}")
+    async def run_status(self, *, run_id: RunId, base_url: httpx.URL | str | None = None) -> Run:
+        response = await self._client.get(self._create_url(f"/runs/{run_id}", base_url=base_url))
         self._raise_error(response)
         return Run.model_validate(response.json())
 
-    async def run_events(self, *, run_id: RunId) -> AsyncIterator[Event]:
-        response = await self._client.get(f"/runs/{run_id}/events")
+    async def run_events(self, *, run_id: RunId, base_url: httpx.URL | str | None = None) -> AsyncIterator[Event]:
+        response = await self._client.get(self._create_url(f"/runs/{run_id}/events", base_url=base_url))
         self._raise_error(response)
         response = RunEventsListResponse.model_validate(response.json())
         for event in response.events:
             yield event
 
-    async def run_cancel(self, *, run_id: RunId) -> Run:
-        response = await self._client.post(f"/runs/{run_id}/cancel")
+    async def run_cancel(self, *, run_id: RunId, base_url: httpx.URL | str | None = None) -> Run:
+        response = await self._client.post(self._create_url(f"/runs/{run_id}/cancel", base_url=base_url))
         self._raise_error(response)
         response = RunCancelResponse.model_validate(response.json())
         return Run(**response.model_dump())
 
-    async def run_resume_sync(self, await_resume: AwaitResume, *, run_id: RunId) -> Run:
+    async def run_resume_sync(
+        self, await_resume: AwaitResume, *, run_id: RunId, base_url: httpx.URL | str | None = None
+    ) -> Run:
         response = await self._client.post(
-            f"/runs/{run_id}",
+            self._create_url(f"/runs/{run_id}", base_url=base_url),
             content=RunResumeRequest(await_resume=await_resume, mode=RunMode.SYNC).model_dump_json(),
         )
         self._raise_error(response)
         response = RunResumeResponse.model_validate(response.json())
         return Run(**response.model_dump())
 
-    async def run_resume_async(self, await_resume: AwaitResume, *, run_id: RunId) -> Run:
+    async def run_resume_async(
+        self, await_resume: AwaitResume, *, run_id: RunId, base_url: httpx.URL | str | None = None
+    ) -> Run:
         response = await self._client.post(
-            f"/runs/{run_id}",
+            self._create_url(f"/runs/{run_id}", base_url=base_url),
             content=RunResumeRequest(await_resume=await_resume, mode=RunMode.ASYNC).model_dump_json(),
         )
         self._raise_error(response)
         response = RunResumeResponse.model_validate(response.json())
         return Run(**response.model_dump())
 
-    async def run_resume_stream(self, await_resume: AwaitResume, *, run_id: RunId) -> AsyncIterator[Event]:
+    async def run_resume_stream(
+        self, await_resume: AwaitResume, *, run_id: RunId, base_url: httpx.URL | str | None = None
+    ) -> AsyncIterator[Event]:
         async with aconnect_sse(
             self._client,
             "POST",
-            f"/runs/{run_id}",
+            self._create_url(f"/runs/{run_id}", base_url=base_url),
             content=RunResumeRequest(await_resume=await_resume, mode=RunMode.STREAM).model_dump_json(),
         ) as event_source:
             async for event in self._validate_stream(event_source):
                 yield event
 
+    async def refresh_session(
+        self, *, base_url: httpx.URL | str | None = None, timeout: httpx._types.TimeoutTypes = 5000
+    ) -> Session:
+        if not self._session:
+            raise RuntimeError("Client is not in a session")
+
+        async with self._session_refresh_lock:
+            url = self._create_url(
+                f"/sessions/{self._session.id}",
+                base_url=base_url or self._session_last_refresh_base_url,
+            )
+
+            try:
+                response = await self._client.get(url, timeout=timeout)
+                response = SessionReadResponse.model_validate(response.json())
+                self._session = Session(**response.model_dump())
+            except ACPError as e:
+                if e.error.code == ErrorCode.NOT_FOUND:
+                    pass
+                raise e
+
+            return self._session
+
     async def _validate_stream(
         self,
         event_source: EventSource,
@@ -231,3 +281,44 @@ class Client:
             response.raise_for_status()
         except httpx.HTTPError:
             raise ACPError(Error.model_validate(response.json()))
+
+    def _create_base_url(self, base_url: httpx.URL | str | None) -> httpx.URL:
+        base_url = httpx.URL(base_url or self._client.base_url)
+        if not base_url.raw_path.endswith(b"/"):
+            base_url = base_url.copy_with(raw_path=base_url.raw_path + b"/")
+        return base_url
+
+    def _create_url(self, endpoint: str, base_url: httpx.URL | str | None) -> httpx.URL:
+        merge_url = httpx.URL(endpoint)
+
+        if not merge_url.is_relative_url:
+            raise ValueError("Endpoint must be a relative URL")
+
+        base_url = self._create_base_url(base_url)
+        merge_raw_path = base_url.raw_path + merge_url.raw_path.lstrip(b"/")
+        return base_url.copy_with(raw_path=merge_raw_path)
+
+    async def _prepare_session_for_run(self, *, base_url: httpx.URL | str | None) -> dict:
+        if not self._session:
+            return {}
+
+        target_base_url = self._create_base_url(base_url=base_url)
+        try:
+            if not self._session_last_refresh_base_url:
+                return {"session": self._session}
+            if self._session_last_refresh_base_url == target_base_url:
+                # Same server, no need to forward session
+                return {"session_id": self._session.id}
+
+            session = await self.refresh_session()
+            return {"session": session}
+        except ACPError as e:
+            if e.error.code == ErrorCode.NOT_FOUND:
+                return {"session": self._session}
+            raise e
+        finally:
+            await self._update_session_refresh_url(target_base_url)
+
+    async def _update_session_refresh_url(self, url: httpx.URL) -> None:
+        async with self._session_refresh_lock:
+            self._session_last_refresh_base_url = url

acp_sdk/models/__init__.py

@@ -1,3 +1,4 @@
 from acp_sdk.models.errors import *  # noqa: F403
 from acp_sdk.models.models import *  # noqa: F403
 from acp_sdk.models.schemas import *  # noqa: F403
+from acp_sdk.models.types import *  # noqa: F403

acp_sdk/models/models.py

@@ -1,5 +1,6 @@
 import asyncio
 import uuid
+from collections.abc import AsyncIterator
 from datetime import datetime, timezone
 from enum import Enum
 from typing import Any, Literal, Optional, Union
@@ -7,6 +8,8 @@ from typing import Any, Literal, Optional, Union
 from pydantic import AnyUrl, BaseModel, ConfigDict, Field
 
 from acp_sdk.models.errors import ACPError, Error
+from acp_sdk.models.types import AgentName, ResourceId, ResourceUrl, RunId, SessionId
+from acp_sdk.shared import ResourceLoader, ResourceStore
 
 
 class AnyModel(BaseModel):
@@ -74,6 +77,40 @@ class Metadata(BaseModel):
     model_config = ConfigDict(extra="allow")
 
 
+class CitationMetadata(BaseModel):
+    """
+    Represents an inline citation, providing info about information source. This
+    is supposed to be rendered as an inline icon, optionally marking a text
+    range it belongs to.
+
+    If CitationMetadata is included together with content in the message part,
+    the citation belongs to that content and renders at the MessagePart position.
+    This way may be used for non-text content, like images and files.
+
+    Alternatively, `start_index` and `end_index` may define a text range,
+    counting characters in the current Message across all MessageParts with
+    content type `text/*`, where the citation will be rendered. If one of
+    `start_index` and `end_index` is missing or their values are equal, the
+    citation renders only as an inline icon at that position.
+
+    If both `start_index` and `end_index` are not present and MessagePart has no
+    content, the citation renders as inline icon only at the MessagePart position.
+
+    Properties:
+    - url: URL of the source document.
+    - title: Title of the source document.
+    - description: Accompanying text, which may be a general description of the
+                   source document, or a specific snippet.
+    """
+
+    kind: Literal["citation"] = "citation"
+    start_index: Optional[int]
+    end_index: Optional[int]
+    url: Optional[str]
+    title: Optional[str]
+    description: Optional[str]
+
+
 class MessagePart(BaseModel):
     name: Optional[str] = None
     content_type: Optional[str] = "text/plain"
@@ -83,9 +120,9 @@ class MessagePart(BaseModel):
 
     model_config = ConfigDict(extra="allow")
 
+    metadata: Optional[CitationMetadata] = Field(discriminator="kind", default=None)
+
     def model_post_init(self, __context: Any) -> None:
-        if self.content is None and self.content_url is None:
-            raise ValueError("Either content or content_url must be provided")
         if self.content is not None and self.content_url is not None:
             raise ValueError("Only one of content or content_url can be provided")
 
@@ -95,6 +132,7 @@ class Artifact(MessagePart):
 
 
 class Message(BaseModel):
+    role: Literal["user"] | Literal["agent"] | str = Field("user", pattern=r"^(user|agent(\/[a-zA-Z0-9_\-]+)?)$")
     parts: list[MessagePart]
     created_at: datetime | None = Field(default_factory=lambda: datetime.now(timezone.utc))
     completed_at: datetime | None = Field(default_factory=lambda: datetime.now(timezone.utc))
@@ -102,7 +140,10 @@ class Message(BaseModel):
     def __add__(self, other: "Message") -> "Message":
         if not isinstance(other, Message):
             raise TypeError(f"Cannot concatenate Message with {type(other).__name__}")
+        if self.role != other.role:
+            raise ValueError("Cannot concatenate messages with different roles")
         return Message(
+            role=self.role,
             parts=self.parts + other.parts,
             created_at=min(self.created_at, other.created_at) if self.created_at and other.created_at else None,
             completed_at=max(self.completed_at, other.completed_at)
@@ -146,11 +187,6 @@ class Message(BaseModel):
         return Message(parts=parts, created_at=self.created_at, completed_at=self.completed_at)
 
 
-AgentName = str
-SessionId = uuid.UUID
-RunId = uuid.UUID
-
-
 class RunMode(str, Enum):
     SYNC = "sync"
     ASYNC = "async"
@@ -280,11 +316,41 @@ Event = Union[
     RunCancelledEvent,
     RunFailedEvent,
     RunCompletedEvent,
-    MessagePartEvent,
 ]
 
 
-class Agent(BaseModel):
+class AgentManifest(BaseModel):
     name: str
     description: str | None = None
     metadata: Metadata = Metadata()
+
+
+class Session(BaseModel):
+    id: SessionId = Field(default_factory=uuid.uuid4)
+    history: list[ResourceUrl] = Field(default_factory=list)
+    state: ResourceUrl | None = None
+
+    loader: ResourceLoader | None = Field(None, exclude=True)
+    store: ResourceStore | None = Field(None, exclude=True)
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    async def load_history(self, *, loader: ResourceLoader | None = None) -> AsyncIterator[Message]:
+        loader = loader or self.loader or ResourceLoader()
+        for url in self.history:
+            data = await loader.load(url)
+            yield Message.model_validate_json(data)
+
+    async def load_state(self, *, loader: ResourceLoader | None = None) -> bytes:
+        loader = loader or self.loader or ResourceLoader()
+        data = await loader.load(self.state)
+        return data
+
+    async def store_state(self, data: bytes, *, store: ResourceStore | None = None) -> ResourceUrl:
+        store = store or self.store
+        if not store:
+            raise ValueError("Store must be specified")
+
+        id = ResourceId()
+        await store.store(id, data)
+        return await store.url(id)

acp_sdk/models/schemas.py

@@ -1,6 +1,16 @@
 from pydantic import BaseModel
 
-from acp_sdk.models.models import Agent, AgentName, AwaitResume, Event, Message, Run, RunMode, SessionId
+from acp_sdk.models.models import (
+    AgentManifest,
+    AgentName,
+    AwaitResume,
+    Event,
+    Message,
+    Run,
+    RunMode,
+    Session,
+    SessionId,
+)
 
 
 class PingResponse(BaseModel):
@@ -8,16 +18,17 @@ class PingResponse(BaseModel):
 
 
 class AgentsListResponse(BaseModel):
-    agents: list[Agent]
+    agents: list[AgentManifest]
 
 
-class AgentReadResponse(Agent):
+class AgentReadResponse(AgentManifest):
     pass
 
 
 class RunCreateRequest(BaseModel):
     agent_name: AgentName
     session_id: SessionId | None = None
+    session: Session | None = None
     input: list[Message]
     mode: RunMode = RunMode.SYNC
 
@@ -45,3 +56,7 @@ class RunCancelResponse(Run):
 
 class RunEventsListResponse(BaseModel):
     events: list[Event]
+
+
+class SessionReadResponse(Session):
+    pass

acp_sdk/models/types.py

@@ -0,0 +1,9 @@
+from uuid import UUID
+
+from pydantic import AnyHttpUrl
+
+AgentName = str
+SessionId = UUID
+RunId = UUID
+ResourceUrl = AnyHttpUrl
+ResourceId = UUID

acp_sdk/server/__init__.py

@@ -1,7 +1,11 @@
-from acp_sdk.server.agent import Agent as Agent
+from acp_sdk.server.agent import AgentManifest as AgentManifest
 from acp_sdk.server.agent import agent as agent
 from acp_sdk.server.app import create_app as create_app
 from acp_sdk.server.context import Context as Context
 from acp_sdk.server.server import Server as Server
+from acp_sdk.server.store import MemoryStore as MemoryStore
+from acp_sdk.server.store import PostgreSQLStore as PostgreSQLStore
+from acp_sdk.server.store import RedisStore as RedisStore
+from acp_sdk.server.store import Store as Store
 from acp_sdk.server.types import RunYield as RunYield
 from acp_sdk.server.types import RunYieldResume as RunYieldResume

acp_sdk/server/agent.py

@@ -1,23 +1,14 @@
 import abc
-import asyncio
 import inspect
 from collections.abc import AsyncGenerator, Coroutine, Generator
-from concurrent.futures import ThreadPoolExecutor
 from typing import Callable
 
-import janus
-
-from acp_sdk.models import (
-    AgentName,
-    Message,
-    SessionId,
-)
-from acp_sdk.models.models import Metadata
+from acp_sdk.models import AgentName, Message, Metadata
 from acp_sdk.server.context import Context
 from acp_sdk.server.types import RunYield, RunYieldResume
 
 
-class Agent(abc.ABC):
+class AgentManifest(abc.ABC):
     @property
     def name(self) -> AgentName:
         return self.__class__.__name__
@@ -38,75 +29,8 @@ class Agent(abc.ABC):
     ):
         pass
 
-    async def execute(
-        self, input: list[Message], session_id: SessionId | None, executor: ThreadPoolExecutor
-    ) -> AsyncGenerator[RunYield, RunYieldResume]:
-        yield_queue: janus.Queue[RunYield] = janus.Queue()
-        yield_resume_queue: janus.Queue[RunYieldResume] = janus.Queue()
-
-        context = Context(
-            session_id=session_id, executor=executor, yield_queue=yield_queue, yield_resume_queue=yield_resume_queue
-        )
-
-        if inspect.isasyncgenfunction(self.run):
-            run = asyncio.create_task(self._run_async_gen(input, context))
-        elif inspect.iscoroutinefunction(self.run):
-            run = asyncio.create_task(self._run_coro(input, context))
-        elif inspect.isgeneratorfunction(self.run):
-            run = asyncio.get_running_loop().run_in_executor(executor, self._run_gen, input, context)
-        else:
-            run = asyncio.get_running_loop().run_in_executor(executor, self._run_func, input, context)
-
-        try:
-            while not run.done() or yield_queue.async_q.qsize() > 0:
-                value = yield await yield_queue.async_q.get()
-                if isinstance(value, Exception):
-                    raise value
-                await yield_resume_queue.async_q.put(value)
-        except janus.AsyncQueueShutDown:
-            pass
-
-    async def _run_async_gen(self, input: list[Message], context: Context) -> None:
-        try:
-            gen: AsyncGenerator[RunYield, RunYieldResume] = self.run(input, context)
-            value = None
-            while True:
-                value = await context.yield_async(await gen.asend(value))
-        except StopAsyncIteration:
-            pass
-        except Exception as e:
-            await context.yield_async(e)
-        finally:
-            context.shutdown()
-
-    async def _run_coro(self, input: list[Message], context: Context) -> None:
-        try:
-            await context.yield_async(await self.run(input, context))
-        except Exception as e:
-            await context.yield_async(e)
-        finally:
-            context.shutdown()
-
-    def _run_gen(self, input: list[Message], context: Context) -> None:
-        try:
-            gen: Generator[RunYield, RunYieldResume] = self.run(input, context)
-            value = None
-            while True:
-                value = context.yield_sync(gen.send(value))
-        except StopIteration:
-            pass
-        except Exception as e:
-            context.yield_sync(e)
-        finally:
-            context.shutdown()
-
-    def _run_func(self, input: list[Message], context: Context) -> None:
-        try:
-            context.yield_sync(self.run(input, context))
-        except Exception as e:
-            context.yield_sync(e)
-        finally:
-            context.shutdown()
+
+Agent = AgentManifest
 
 
 def agent(
@@ -114,10 +38,10 @@ def agent(
     description: str | None = None,
     *,
     metadata: Metadata | None = None,
-) -> Callable[[Callable], Agent]:
+) -> Callable[[Callable], AgentManifest]:
     """Decorator to create an agent."""
 
-    def decorator(fn: Callable) -> Agent:
+    def decorator(fn: Callable) -> AgentManifest:
         signature = inspect.signature(fn)
         parameters = list(signature.parameters.values())
 
@@ -130,7 +54,7 @@ def agent(
 
         has_context_param = len(parameters) == 2
 
-        class DecoratorAgentBase(Agent):
+        class DecoratorAgentBase(AgentManifest):
             @property
             def name(self) -> str:
                 return name or fn.__name__
@@ -143,7 +67,7 @@ def agent(
             def metadata(self) -> Metadata:
                 return metadata or Metadata()
 
-        agent: Agent
+        agent: AgentManifest
         if inspect.isasyncgenfunction(fn):
 
             class AsyncGenDecoratorAgent(DecoratorAgentBase):