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

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):

acp_sdk/server/app.py

@@ -5,19 +5,24 @@ from contextlib import asynccontextmanager
 from datetime import timedelta
 from enum import Enum
 
-from fastapi import Depends, FastAPI, HTTPException, status
+import httpx
+import obstore.store
+from fastapi import Depends, FastAPI, HTTPException, Request, 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 obstore.exceptions import NotFoundError
 
 from acp_sdk.models import (
-    Agent as AgentModel,
-)
-from acp_sdk.models import (
+    ACPError,
     AgentName,
     AgentReadResponse,
     AgentsListResponse,
+    AwaitResume,
+    PingResponse,
+    ResourceId,
+    ResourceUrl,
     Run,
     RunCancelResponse,
     RunCreateRequest,
@@ -28,11 +33,15 @@ from acp_sdk.models import (
     RunReadResponse,
     RunResumeRequest,
     RunResumeResponse,
+    RunStatus,
+    Session,
+    SessionId,
+    SessionReadResponse,
+)
+from acp_sdk.models import (
+    AgentManifest as AgentModel,
 )
-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.agent import AgentManifest
 from acp_sdk.server.errors import (
     RequestValidationError,
     StarletteHTTPException,
@@ -42,9 +51,9 @@ from acp_sdk.server.errors import (
     validation_exception_handler,
 )
 from acp_sdk.server.executor import CancelData, Executor, RunData
-from acp_sdk.server.session import Session
 from acp_sdk.server.store import MemoryStore, Store
 from acp_sdk.server.utils import stream_sse, wait_util_stop
+from acp_sdk.shared import ResourceLoader, ResourceStore
 
 
 class Headers(str, Enum):
@@ -52,23 +61,34 @@ class Headers(str, Enum):
 
 
 def create_app(
-    *agents: Agent,
+    *agents: AgentManifest,
     store: Store | None = None,
+    resource_store: ResourceStore | None = None,
+    resource_loader: ResourceLoader | None = None,
+    forward_resources: bool = True,
     lifespan: Lifespan[AppType] | None = None,
     dependencies: list[Depends] | None = None,
 ) -> FastAPI:
+    if not forward_resources and (
+        resource_store is None
+        or isinstance(resource_store._store, (obstore.store.MemoryStore, obstore.store.LocalStore))
+    ):
+        raise ValueError("Resource forwarding must be enabled when resource store does not support HTTP URLs")
+
     executor: ThreadPoolExecutor
+    client = httpx.AsyncClient()
 
     @asynccontextmanager
     async def internal_lifespan(app: FastAPI) -> AsyncGenerator[None]:
         nonlocal executor
-        with ThreadPoolExecutor() as exec:
-            executor = exec
-            if not lifespan:
-                yield None
-            else:
-                async with lifespan(app) as state:
-                    yield state
+        async with client:
+            with ThreadPoolExecutor() as exec:
+                executor = exec
+                if not lifespan:
+                    yield None
+                else:
+                    async with lifespan(app) as state:
+                        yield state
 
     app = FastAPI(
         lifespan=internal_lifespan,
@@ -83,7 +103,7 @@ def create_app(
         allow_credentials=True,
     )
 
-    agents: dict[AgentName, Agent] = {agent.name: agent for agent in agents}
+    agents: dict[AgentName, AgentManifest] = {agent.name: agent for agent in agents}
 
     store = store or MemoryStore(limit=1000, ttl=timedelta(hours=1))
     run_store = store.as_store(model=RunData, prefix="run_")
@@ -91,6 +111,9 @@ def create_app(
     run_resume_store = store.as_store(model=AwaitResume, prefix="run_resume_")
     session_store = store.as_store(model=Session, prefix="session_")
 
+    resource_loader = resource_loader or ResourceLoader(client=client)
+    resource_store = resource_store or ResourceStore(store=obstore.store.MemoryStore())
+
     app.exception_handler(ACPError)(acp_error_handler)
     app.exception_handler(StarletteHTTPException)(http_exception_handler)
     app.exception_handler(RequestValidationError)(validation_exception_handler)
@@ -107,7 +130,7 @@ def create_app(
             run_data.run.status = RunStatus.CANCELLING
         return run_data
 
-    def find_agent(agent_name: AgentName) -> Agent:
+    def find_agent(agent_name: AgentName) -> AgentManifest:
         agent = agents.get(agent_name, None)
         if not agent:
             raise HTTPException(status_code=404, detail=f"Agent {agent_name} not found")
@@ -132,36 +155,54 @@ def create_app(
         return PingResponse()
 
     @app.post("/runs")
-    async def create_run(request: RunCreateRequest) -> RunCreateResponse:
+    async def create_run(request: RunCreateRequest, req: Request) -> RunCreateResponse:
         agent = find_agent(request.agent_name)
 
-        session = (
-            (await session_store.get(request.session_id)) or Session(id=request.session_id)
+        if request.session_id and request.session and request.session_id != request.session.id:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Session ID mismatch")
+
+        session = request.session or (
+            (
+                await session_store.get(request.session_id)
+                or Session(id=request.session_id, loader=resource_loader, store=resource_store)
+            )
             if request.session_id
-            else Session()
+            else Session(loader=resource_loader, store=resource_store)
         )
+
         nonlocal executor
         run_data = RunData(
-            run=Run(agent_name=agent.name, session_id=session.id),
-            input=request.input,
+            run=Run(
+                agent_name=agent.name,
+                session_id=session.id,
+            )
         )
         await run_store.set(run_data.key, run_data)
-
-        session.append(run_data.run.run_id)
         await session_store.set(session.id, session)
 
         headers = {Headers.RUN_ID: str(run_data.run.run_id)}
         ready = asyncio.Event()
 
+        async def create_resource_url(id: ResourceId) -> ResourceUrl:
+            if forward_resources:
+                return ResourceUrl(url=str(req.url_for("get_resource", resource_id=id)))
+            else:
+                return await resource_store.url(id)
+
         Executor(
             agent=agent,
             run_data=run_data,
-            history=await session.history(run_store),
+            session=session,
+            session_store=session_store,
             run_store=run_store,
             cancel_store=run_cancel_store,
             resume_store=run_resume_store,
             executor=executor,
-        ).execute(wait=ready)
+            request=req,
+            resource_store=resource_store,
+            resource_loader=resource_loader,
+            create_resource_url=create_resource_url,
+        ).execute(request.input, wait=ready)
 
         match request.mode:
             case RunMode.STREAM:
@@ -242,4 +283,21 @@ def create_app(
         run_data.run.status = RunStatus.CANCELLING
         return JSONResponse(status_code=status.HTTP_202_ACCEPTED, content=jsonable_encoder(run_data.run))
 
+    @app.get("/sessions/{session_id}")
+    async def read_session(session_id: SessionId) -> SessionReadResponse:
+        session = await session_store.get(session_id)
+        if not session:
+            raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
+        return session
+
+    if forward_resources:
+
+        @app.get("/resources/{resource_id}", name="get_resource")
+        async def read_resource(resource_id: ResourceId) -> StreamingResponse:
+            try:
+                result = await resource_store.load(resource_id)
+                return StreamingResponse(result)
+            except NotFoundError:
+                raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Resource {resource_id} not found")
+
     return app

acp_sdk/server/context.py

@@ -1,22 +1,30 @@
 from concurrent.futures import ThreadPoolExecutor
 
 import janus
+from fastapi import Request
 
-from acp_sdk.models import SessionId
+from acp_sdk.models import Session
 from acp_sdk.server.types import RunYield, RunYieldResume
+from acp_sdk.shared import ResourceLoader, ResourceStore
 
 
 class Context:
     def __init__(
         self,
         *,
-        session_id: SessionId | None = None,
+        session: Session,
+        store: ResourceStore,
+        loader: ResourceLoader,
         executor: ThreadPoolExecutor,
+        request: Request,
         yield_queue: janus.Queue[RunYield],
         yield_resume_queue: janus.Queue[RunYieldResume],
     ) -> None:
-        self.session_id = session_id
+        self.session = session
+        self.storage = store
+        self.loader = loader
         self.executor = executor
+        self.request = request
         self._yield_queue = yield_queue
         self._yield_resume_queue = yield_resume_queue
 

acp_sdk/server/executor.py

@@ -1,10 +1,14 @@
 import asyncio
+import inspect
 import logging
-from collections.abc import AsyncIterator
+import uuid
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Generator
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timezone
-from typing import Self
+from typing import Callable, Self
 
+import janus
+from fastapi import Request
 from pydantic import BaseModel, ValidationError
 
 from acp_sdk.instrumentation import get_tracer
@@ -22,6 +26,8 @@ from acp_sdk.models import (
     MessageCreatedEvent,
     MessagePart,
     MessagePartEvent,
+    ResourceId,
+    ResourceUrl,
     Run,
     RunAwaitingEvent,
     RunCancelledEvent,
@@ -30,15 +36,18 @@ from acp_sdk.models import (
     RunFailedEvent,
     RunInProgressEvent,
     RunStatus,
+    Session,
 )
-from acp_sdk.server.agent import Agent
+from acp_sdk.server.agent import AgentManifest
+from acp_sdk.server.context import Context
 from acp_sdk.server.logging import logger
 from acp_sdk.server.store import Store
+from acp_sdk.server.types import RunYield, RunYieldResume
+from acp_sdk.shared import ResourceLoader, ResourceStore
 
 
 class RunData(BaseModel):
     run: Run
-    input: list[Message]
     events: list[Event] = []
 
     @property
@@ -62,27 +71,38 @@ class Executor:
     def __init__(
         self,
         *,
-        agent: Agent,
+        agent: AgentManifest,
         run_data: RunData,
-        history: list[Message],
+        session: Session,
         executor: ThreadPoolExecutor,
+        request: Request,
         run_store: Store[RunData],
         cancel_store: Store[CancelData],
         resume_store: Store[AwaitResume],
+        session_store: Store[Session],
+        resource_store: ResourceStore,
+        resource_loader: ResourceLoader,
+        create_resource_url: Callable[[ResourceId], Awaitable[ResourceUrl]],
     ) -> None:
         self.agent = agent
-        self.history = history
+        self.session = session
         self.run_data = run_data
         self.executor = executor
+        self.request = request
 
         self.run_store = run_store
         self.cancel_store = cancel_store
         self.resume_store = resume_store
+        self.session_store = session_store
+        self.resource_store = resource_store
+        self.resource_loader = resource_loader
+
+        self.create_resource_url = create_resource_url
 
         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))
+    def execute(self, input: list[Message], *, wait: asyncio.Event) -> None:
+        self.task = asyncio.create_task(self._execute(input=input, executor=self.executor, wait=wait))
         self.watcher = asyncio.create_task(self._watch_for_cancellation())
 
     async def _push(self) -> None:
@@ -108,7 +128,16 @@ class Executor:
             except Exception:
                 logger.warning("Cancellation watcher failed, restarting")
 
-    async def _execute(self, run_data: RunData, *, executor: ThreadPoolExecutor, wait: asyncio.Event) -> None:
+    async def _record_session(self, history: list[Message]) -> None:
+        for message in history:
+            id = uuid.uuid4()
+            url = await self.create_resource_url(id)
+            await self.resource_store.store(id, message.model_dump_json().encode())
+            self.session.history.append(url)
+        await self.session_store.set(self.session.id, self.session)
+
+    async def _execute(self, input: list[Message], *, executor: ThreadPoolExecutor, wait: asyncio.Event) -> None:
+        run_data = self.run_data
         with get_tracer().start_as_current_span("run"):
             in_message = False
 
@@ -118,15 +147,22 @@ class Executor:
                     message = run_data.run.output[-1]
                     message.completed_at = datetime.now(timezone.utc)
                     await self._emit(MessageCompletedEvent(message=message))
+                    session_history.append(message)
                     in_message = False
 
+            session_history = input.copy()
             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
+                generator = self._execute_agent(
+                    input=input,
+                    session=self.session,
+                    storage=self.resource_store,
+                    loader=self.resource_loader,
+                    executor=executor,
+                    request=self.request,
                 )
                 self.logger.info("Run started")
 
@@ -141,18 +177,21 @@ class Executor:
                         if isinstance(next, str):
                             next = MessagePart(content=next)
                         if not in_message:
-                            run_data.run.output.append(Message(parts=[], completed_at=None))
+                            run_data.run.output.append(
+                                Message(role=f"agent/{self.agent.name}", 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)
+                        run_data.run.output.append(next.model_copy(update={"role": f"agent/{self.agent.name}"}))
                         await self._emit(MessageCreatedEvent(message=next))
                         for part in next.parts:
                             await self._emit(MessagePartEvent(part=part))
                         await self._emit(MessageCompletedEvent(message=next))
+                        session_history.append(next)
                     elif isinstance(next, AwaitRequest):
                         run_data.run.await_request = next
                         run_data.run.status = RunStatus.AWAITING
@@ -180,6 +219,10 @@ class Executor:
                 await flush_message()
                 run_data.run.status = RunStatus.COMPLETED
                 run_data.run.finished_at = datetime.now(timezone.utc)
+                try:
+                    await self._record_session(session_history)
+                except Exception as e:
+                    self.logger.warning(f"Failed to record session: {e}")
                 await self._emit(RunCompletedEvent(run=run_data.run))
                 self.logger.info("Run completed")
             except asyncio.CancelledError:
@@ -196,3 +239,85 @@ class Executor:
                 run_data.run.finished_at = datetime.now(timezone.utc)
                 await self._emit(RunFailedEvent(run=run_data.run))
                 self.logger.exception("Run failed")
+
+    async def _execute_agent(
+        self,
+        input: list[Message],
+        session: Session,
+        storage: ResourceStore,
+        loader: ResourceLoader,
+        executor: ThreadPoolExecutor,
+        request: Request,
+    ) -> AsyncGenerator[RunYield, RunYieldResume]:
+        yield_queue: janus.Queue[RunYield] = janus.Queue()
+        yield_resume_queue: janus.Queue[RunYieldResume] = janus.Queue()
+
+        context = Context(
+            session=session,
+            store=storage,
+            loader=loader,
+            executor=executor,
+            request=request,
+            yield_queue=yield_queue,
+            yield_resume_queue=yield_resume_queue,
+        )
+
+        if inspect.isasyncgenfunction(self.agent.run):
+            run = asyncio.create_task(self._run_async_gen(input, context))
+        elif inspect.iscoroutinefunction(self.agent.run):
+            run = asyncio.create_task(self._run_coro(input, context))
+        elif inspect.isgeneratorfunction(self.agent.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.agent.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.agent.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.agent.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.agent.run(input, context))
+        except Exception as e:
+            context.yield_sync(e)
+        finally:
+            context.shutdown()