diff --git a/apps/agentstack-sdk-py/examples/agent.py b/apps/agentstack-sdk-py/examples/agent.py index 7539092b9a..370589c10f 100644 --- a/apps/agentstack-sdk-py/examples/agent.py +++ b/apps/agentstack-sdk-py/examples/agent.py @@ -93,7 +93,7 @@ async def execute( agent = beeai_framework.agents.react.ReActAgent( llm=beeai_framework.adapters.openai.backend.chat.OpenAIChatModel( model_id=self.context_llm[context.context_id]["default"].api_model, - api_key=self.context_llm[context.context_id]["default"].api_key, + api_key=self.context_llm[context.context_id]["default"].api_key.get_secret_value(), base_url=self.context_llm[context.context_id]["default"].api_base, ), tools=[ diff --git a/apps/agentstack-sdk-py/examples/secrets_agent.py b/apps/agentstack-sdk-py/examples/secrets_agent.py index 1e88a2d96c..99526f1c27 100644 --- a/apps/agentstack-sdk-py/examples/secrets_agent.py +++ b/apps/agentstack-sdk-py/examples/secrets_agent.py @@ -43,7 +43,7 @@ async def secrets_agent( ): """Agent that uses request a secret that can be provided during runtime""" if secrets and secrets.data and secrets.data.secret_fulfillments: - yield f"IBM Cloud API key: {secrets.data.secret_fulfillments['ibm_cloud'].secret}" + yield f"IBM Cloud API key: {secrets.data.secret_fulfillments['ibm_cloud'].secret.get_secret_value()}" else: runtime_provided_secrets = await secrets.request_secrets( params=SecretsServiceExtensionParams( @@ -51,7 +51,7 @@ async def secrets_agent( ) ) if runtime_provided_secrets and runtime_provided_secrets.secret_fulfillments: - yield f"IBM Cloud API key: {runtime_provided_secrets.secret_fulfillments['ibm_cloud'].secret}" + yield f"IBM Cloud API key: {runtime_provided_secrets.secret_fulfillments['ibm_cloud'].secret.get_secret_value()}" else: yield "No IBM Cloud API key provided" diff --git a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/auth/oauth/oauth.py b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/auth/oauth/oauth.py index 0df594d580..e8df72e385 100644 --- a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/auth/oauth/oauth.py +++ b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/auth/oauth/oauth.py @@ -19,6 +19,7 @@ from agentstack_sdk.a2a.extensions.auth.oauth.storage import MemoryTokenStorageFactory, TokenStorageFactory from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec from agentstack_sdk.a2a.types import AgentMessage, AuthRequired, RunYieldResume +from agentstack_sdk.util.pydantic import REVEAL_SECRETS, SecureBaseModel if TYPE_CHECKING: from agentstack_sdk.server.context import RunContext @@ -26,15 +27,15 @@ _DEFAULT_DEMAND_NAME = "default" -class AuthRequest(pydantic.BaseModel): +class AuthRequest(SecureBaseModel): authorization_endpoint_url: pydantic.AnyUrl -class AuthResponse(pydantic.BaseModel): +class AuthResponse(SecureBaseModel): redirect_uri: pydantic.AnyUrl -class OAuthFulfillment(pydantic.BaseModel): +class OAuthFulfillment(SecureBaseModel): redirect_uri: pydantic.AnyUrl @@ -122,7 +123,10 @@ async def handle_callback() -> tuple[str, str | None]: def create_auth_request(self, *, authorization_endpoint_url: pydantic.AnyUrl): data = AuthRequest(authorization_endpoint_url=authorization_endpoint_url) - return AgentMessage(text="Authorization required", metadata={self.spec.URI: data.model_dump(mode="json")}) + return AgentMessage( + text="Authorization required", + metadata={self.spec.URI: data.model_dump(mode="json", context={REVEAL_SECRETS: True})}, + ) def parse_auth_response(self, *, message: A2AMessage): if not message or not message.metadata or not (data := message.metadata.get(self.spec.URI)): @@ -147,5 +151,5 @@ def create_auth_response(self, *, task_id: str, redirect_uri: pydantic.AnyUrl): role=Role.user, parts=[TextPart(text="Authorization completed")], task_id=task_id, - metadata={self.spec.URI: data.model_dump(mode="json")}, + metadata={self.spec.URI: data.model_dump(mode="json", context={REVEAL_SECRETS: True})}, ) diff --git a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/auth/secrets/secrets.py b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/auth/secrets/secrets.py index ad4a50b76a..87b6b148e7 100644 --- a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/auth/secrets/secrets.py +++ b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/auth/secrets/secrets.py @@ -8,22 +8,28 @@ import pydantic from a2a.server.agent_execution.context import RequestContext from a2a.types import Message as A2AMessage +from opentelemetry import trace from typing_extensions import override from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec from agentstack_sdk.a2a.types import AgentMessage, AuthRequired +from agentstack_sdk.util.pydantic import REDACT_SECRETS, REVEAL_SECRETS, SecureBaseModel +from agentstack_sdk.util.telemetry import flatten_dict if TYPE_CHECKING: from agentstack_sdk.server.context import RunContext +A2A_EXTENSION_SECRETS_REQUESTED = "a2a_extension.secrets.requested" +A2A_EXTENSION_SECRETS_RESOLVED = "a2a_extension.secrets.resolved" + class SecretDemand(pydantic.BaseModel): name: str description: str | None = None -class SecretFulfillment(pydantic.BaseModel): - secret: str +class SecretFulfillment(SecureBaseModel): + secret: pydantic.SecretStr class SecretsServiceExtensionParams(pydantic.BaseModel): @@ -61,15 +67,25 @@ def parse_secret_response(self, message: A2AMessage) -> SecretsServiceExtensionM return SecretsServiceExtensionMetadata.model_validate(data) async def request_secrets(self, params: SecretsServiceExtensionParams) -> SecretsServiceExtensionMetadata: + span = trace.get_current_span() + span.add_event( + A2A_EXTENSION_SECRETS_REQUESTED, + attributes=flatten_dict(params.model_dump(context={REDACT_SECRETS: True})), + ) resume = await self.context.yield_async( AuthRequired( message=AgentMessage( - metadata={self.spec.URI: params.model_dump(mode="json")}, + metadata={self.spec.URI: params.model_dump(mode="json", context={REVEAL_SECRETS: True})}, ) ) ) if isinstance(resume, A2AMessage): - return self.parse_secret_response(message=resume) + response = self.parse_secret_response(message=resume) + span.add_event( + A2A_EXTENSION_SECRETS_RESOLVED, + attributes=flatten_dict(response.model_dump(context={REDACT_SECRETS: True})), + ) + return response else: raise ValueError("Secrets has not been provided in response.") diff --git a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/base.py b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/base.py index 27804ab2aa..dea0d15246 100644 --- a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/base.py +++ b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/base.py @@ -13,8 +13,17 @@ from a2a.server.agent_execution.context import RequestContext from a2a.types import AgentCard, AgentExtension from a2a.types import Message as A2AMessage +from opentelemetry import trace +from opentelemetry.trace import SpanKind +from pydantic import BaseModel from typing_extensions import override +from agentstack_sdk.util.pydantic import REDACT_SECRETS +from agentstack_sdk.util.telemetry import ( + flatten_dict, + trace_class, +) + ParamsT = typing.TypeVar("ParamsT") MetadataFromClientT = typing.TypeVar("MetadataFromClientT") MetadataFromServerT = typing.TypeVar("MetadataFromServerT") @@ -25,6 +34,10 @@ from agentstack_sdk.server.dependencies import Dependency +A2A_EXTENSION_URI = "a2a_extension.uri" +A2A_EXTENSION_METADATA_RECEIVED_EVENT = "a2a_extension.metadata.received" + + def _get_generic_args(cls: type, base_class: type) -> tuple[typing.Any, ...]: for base in getattr(cls, "__orig_bases__", ()): if typing.get_origin(base) is base_class and (args := typing.get_args(base)): @@ -121,7 +134,14 @@ class BaseExtensionServer(abc.ABC, typing.Generic[ExtensionSpecT, MetadataFromCl def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) - cls.MetadataFromClient = _get_generic_args(cls, BaseExtensionServer)[1] + + generic_args = _get_generic_args(cls, BaseExtensionServer) + trace_class( + kind=SpanKind.SERVER, + exclude_list=["lifespan", "_fork"], + attributes={A2A_EXTENSION_URI: generic_args[0].URI}, + )(cls) + cls.MetadataFromClient = generic_args[1] _metadata_from_client: MetadataFromClientT | None = None _dependencies: dict[str, Dependency] = {} # noqa: RUF012 @@ -151,6 +171,11 @@ def parse_client_metadata(self, message: A2AMessage) -> MetadataFromClientT | No def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, request_context: RequestContext): if self._metadata_from_client is None: self._metadata_from_client = self.parse_client_metadata(message) + if isinstance(self._metadata_from_client, BaseModel): + trace.get_current_span().add_event( + A2A_EXTENSION_METADATA_RECEIVED_EVENT, + attributes=flatten_dict(self._metadata_from_client.model_dump(context={REDACT_SECRETS: True})), + ) def _fork(self) -> typing.Self: """Creates a clone of this instance with the same arguments as the original""" @@ -182,7 +207,10 @@ class BaseExtensionClient(abc.ABC, typing.Generic[ExtensionSpecT, MetadataFromSe def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) - cls.MetadataFromServer = _get_generic_args(cls, BaseExtensionClient)[1] + + generic_args = _get_generic_args(cls, BaseExtensionClient) + trace_class(kind=SpanKind.CLIENT, attributes={A2A_EXTENSION_URI: generic_args[0].URI})(cls) + cls.MetadataFromServer = generic_args[1] def __init__(self, spec: ExtensionSpecT) -> None: self.spec = spec diff --git a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/interactions/approval.py b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/interactions/approval.py index efccf3557b..050231f340 100644 --- a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/interactions/approval.py +++ b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/interactions/approval.py @@ -9,33 +9,40 @@ import a2a.types from mcp import Implementation, Tool +from opentelemetry import trace from pydantic import BaseModel, Discriminator, Field, TypeAdapter from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec from agentstack_sdk.a2a.types import AgentMessage, InputRequired +from agentstack_sdk.util.pydantic import REDACT_SECRETS, REVEAL_SECRETS, SecureBaseModel +from agentstack_sdk.util.telemetry import flatten_dict if TYPE_CHECKING: from agentstack_sdk.server.context import RunContext +A2A_EXTENSION_APPROVAL_REQUESTED = "a2a_extension.approval.requested" +A2A_EXTENSION_APPROVAL_RESOLVED = "a2a_extension.approval.resolved" + + class ApprovalRejectionError(RuntimeError): pass -class GenericApprovalRequest(BaseModel): +class GenericApprovalRequest(SecureBaseModel): action: Literal["generic"] = "generic" title: str | None = Field(None, description="A human-readable title for the action being approved.") description: str | None = Field(None, description="A human-readable description of the action being approved.") -class ToolCallServer(BaseModel): +class ToolCallServer(SecureBaseModel): name: str = Field(description="The programmatic name of the server.") title: str | None = Field(description="A human-readable title for the server.") version: str = Field(description="The version of the server.") -class ToolCallApprovalRequest(BaseModel): +class ToolCallApprovalRequest(SecureBaseModel): action: Literal["tool-call"] = "tool-call" title: str | None = Field(None, description="A human-readable title of the tool.") @@ -60,7 +67,7 @@ def from_mcp_tool( ApprovalRequest = Annotated[GenericApprovalRequest | ToolCallApprovalRequest, Discriminator("action")] -class ApprovalResponse(BaseModel): +class ApprovalResponse(SecureBaseModel): decision: Literal["approve", "reject"] @property @@ -86,7 +93,10 @@ class ApprovalExtensionMetadata(BaseModel): class ApprovalExtensionServer(BaseExtensionServer[ApprovalExtensionSpec, ApprovalExtensionMetadata]): def create_request_message(self, *, request: ApprovalRequest): - return AgentMessage(text="Approval requested", metadata={self.spec.URI: request.model_dump(mode="json")}) + return AgentMessage( + text="Approval requested", + metadata={self.spec.URI: request.model_dump(mode="json", context={REVEAL_SECRETS: True})}, + ) def parse_response(self, *, message: a2a.types.Message): if not message.metadata or not (data := message.metadata.get(self.spec.URI)): @@ -99,11 +109,21 @@ async def request_approval( *, context: RunContext, ) -> ApprovalResponse: + span = trace.get_current_span() + span.add_event( + A2A_EXTENSION_APPROVAL_REQUESTED, + attributes=flatten_dict(request.model_dump(context={REDACT_SECRETS: True})), + ) message = self.create_request_message(request=request) message = await context.yield_async(InputRequired(message=message)) if not message: raise RuntimeError("Yield did not return a message") - return self.parse_response(message=message) + response = self.parse_response(message=message) + span.add_event( + A2A_EXTENSION_APPROVAL_RESOLVED, + attributes=flatten_dict(response.model_dump(context={REDACT_SECRETS: True})), + ) + return response class ApprovalExtensionClient(BaseExtensionClient[ApprovalExtensionSpec, NoneType]): @@ -113,7 +133,7 @@ def create_response_message(self, *, response: ApprovalResponse, task_id: str | role=a2a.types.Role.user, parts=[], task_id=task_id, - metadata={self.spec.URI: response.model_dump(mode="json")}, + metadata={self.spec.URI: response.model_dump(mode="json", context={REVEAL_SECRETS: True})}, ) def parse_request(self, *, message: a2a.types.Message): @@ -122,4 +142,4 @@ def parse_request(self, *, message: a2a.types.Message): return TypeAdapter(ApprovalRequest).validate_python(data) def metadata(self) -> dict[str, Any]: - return {self.spec.URI: ApprovalExtensionMetadata().model_dump(mode="json")} + return {self.spec.URI: ApprovalExtensionMetadata().model_dump(mode="json", context={REVEAL_SECRETS: True})} diff --git a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/embedding.py b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/embedding.py index 4a69c5af77..96a44baa64 100644 --- a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/embedding.py +++ b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/embedding.py @@ -13,12 +13,13 @@ from typing_extensions import override from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec +from agentstack_sdk.util.pydantic import REVEAL_SECRETS, SecureBaseModel if TYPE_CHECKING: from agentstack_sdk.server.context import RunContext -class EmbeddingFulfillment(pydantic.BaseModel): +class EmbeddingFulfillment(SecureBaseModel): identifier: str | None = None """ Name of the model for identification and optimization purposes. Usually corresponds to LiteLLM identifiers. @@ -31,7 +32,7 @@ class EmbeddingFulfillment(pydantic.BaseModel): Base URL for an OpenAI-compatible API. It should provide at least /v1/chat/completions """ - api_key: str + api_key: pydantic.SecretStr """ API key to attach as a `Authorization: Bearer $api_key` header. """ @@ -101,6 +102,6 @@ class EmbeddingServiceExtensionClient(BaseExtensionClient[EmbeddingServiceExtens def fulfillment_metadata(self, *, embedding_fulfillments: dict[str, EmbeddingFulfillment]) -> dict[str, Any]: return { self.spec.URI: EmbeddingServiceExtensionMetadata(embedding_fulfillments=embedding_fulfillments).model_dump( - mode="json" + mode="json", context={REVEAL_SECRETS: True} ) } diff --git a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/llm.py b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/llm.py index ba302f396d..2f925cae47 100644 --- a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/llm.py +++ b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/llm.py @@ -13,12 +13,13 @@ from typing_extensions import override from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec +from agentstack_sdk.util.pydantic import REVEAL_SECRETS, SecureBaseModel if TYPE_CHECKING: from agentstack_sdk.server.context import RunContext -class LLMFulfillment(pydantic.BaseModel): +class LLMFulfillment(SecureBaseModel): identifier: str | None = None """ Name of the model for identification and optimization purposes. Usually corresponds to LiteLLM identifiers. @@ -31,7 +32,7 @@ class LLMFulfillment(pydantic.BaseModel): Base URL for an OpenAI-compatible API. It should provide at least /v1/chat/completions """ - api_key: str + api_key: pydantic.SecretStr """ API key to attach as a `Authorization: Bearer $api_key` header. """ @@ -97,4 +98,8 @@ def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, class LLMServiceExtensionClient(BaseExtensionClient[LLMServiceExtensionSpec, NoneType]): def fulfillment_metadata(self, *, llm_fulfillments: dict[str, LLMFulfillment]) -> dict[str, Any]: - return {self.spec.URI: LLMServiceExtensionMetadata(llm_fulfillments=llm_fulfillments).model_dump(mode="json")} + return { + self.spec.URI: LLMServiceExtensionMetadata(llm_fulfillments=llm_fulfillments).model_dump( + mode="json", context={REVEAL_SECRETS: True} + ) + } diff --git a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/mcp.py b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/mcp.py index 0b516f14f6..fc64e8ff77 100644 --- a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/mcp.py +++ b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/mcp.py @@ -13,6 +13,7 @@ from a2a.types import Message as A2AMessage from mcp.client.stdio import StdioServerParameters, stdio_client from mcp.client.streamable_http import streamablehttp_client # pyrefly: ignore [deprecated] -- TODO: upgrade +from pydantic import AnyUrl from typing_extensions import override from agentstack_sdk.a2a.extensions.auth.oauth.oauth import OAuthExtensionServer @@ -20,6 +21,7 @@ from agentstack_sdk.a2a.extensions.services.platform import PlatformApiExtensionServer from agentstack_sdk.platform.client import get_platform_client from agentstack_sdk.util.logging import logger +from agentstack_sdk.util.pydantic import REVEAL_SECRETS, SecureBaseModel, redact_dict, redact_str if TYPE_CHECKING: from agentstack_sdk.server.context import RunContext @@ -30,20 +32,32 @@ _DEFAULT_ALLOWED_TRANSPORTS: list[_TRANSPORT_TYPES] = ["streamable_http"] -class StdioTransport(pydantic.BaseModel): +class StdioTransport(SecureBaseModel): type: Literal["stdio"] = "stdio" command: str args: list[str] env: dict[str, str] | None = None + @pydantic.field_serializer("args") + def _redact_args(self, v: list[str], info) -> list[str]: + return [redact_str(arg, info) for arg in v] -class StreamableHTTPTransport(pydantic.BaseModel): + @pydantic.field_serializer("env") + def _redact_env(self, v: dict[str, str] | None, info) -> dict[str, str] | None: + return redact_dict(v, info) if v is not None else None + + +class StreamableHTTPTransport(SecureBaseModel): type: Literal["streamable_http"] = "streamable_http" - url: str + url: AnyUrl headers: dict[str, str] | None = None + @pydantic.field_serializer("headers") + def _redact_headers(self, v: dict[str, str] | None, info) -> dict[str, str] | None: + return redact_dict(v, info) if v is not None else None + MCPTransport = Annotated[StdioTransport | StreamableHTTPTransport, pydantic.Field(discriminator="type")] @@ -114,7 +128,9 @@ def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, for fullfilment in self.data.mcp_fulfillments.values(): if fullfilment.transport.type == "streamable_http": try: - fullfilment.transport.url = re.sub("^{platform_url}", platform_url, str(fullfilment.transport.url)) + fullfilment.transport.url = AnyUrl( + re.sub("^{platform_url}", platform_url, str(fullfilment.transport.url)) + ) except Exception: logger.warning("Platform URL substitution failed", exc_info=True) @@ -162,7 +178,7 @@ async def create_client(self, demand: str = _DEFAULT_DEMAND_NAME): elif isinstance(transport, StreamableHTTPTransport): # pyrefly: ignore [deprecated] -- TODO: upgrade async with streamablehttp_client( - url=transport.url, + url=str(transport.url), headers=transport.headers, auth=await self._create_auth(transport), ) as ( @@ -180,7 +196,7 @@ async def _create_auth(self, transport: StreamableHTTPTransport): platform and platform.data and platform.data.base_url - and transport.url.startswith(str(platform.data.base_url)) + and str(transport.url).startswith(str(platform.data.base_url)) ): return await platform.create_httpx_auth() oauth = self._get_oauth_server() @@ -191,4 +207,8 @@ async def _create_auth(self, transport: StreamableHTTPTransport): class MCPServiceExtensionClient(BaseExtensionClient[MCPServiceExtensionSpec, NoneType]): def fulfillment_metadata(self, *, mcp_fulfillments: dict[str, MCPFulfillment]) -> dict[str, Any]: - return {self.spec.URI: MCPServiceExtensionMetadata(mcp_fulfillments=mcp_fulfillments).model_dump(mode="json")} + return { + self.spec.URI: MCPServiceExtensionMetadata(mcp_fulfillments=mcp_fulfillments).model_dump( + mode="json", context={REVEAL_SECRETS: True} + ) + } diff --git a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/platform.py b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/platform.py index 7e1acfeffe..a46dfcff46 100644 --- a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/platform.py +++ b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/platform.py @@ -26,14 +26,15 @@ from agentstack_sdk.platform.client import PlatformClient from agentstack_sdk.server.middleware.platform_auth_backend import PlatformAuthenticatedUser from agentstack_sdk.util.httpx import BearerAuth +from agentstack_sdk.util.pydantic import REVEAL_SECRETS, SecureBaseModel if TYPE_CHECKING: from agentstack_sdk.server.context import RunContext -class PlatformApiExtensionMetadata(pydantic.BaseModel): +class PlatformApiExtensionMetadata(SecureBaseModel): base_url: HttpUrl | None = None - auth_token: pydantic.Secret[str] | None = None + auth_token: pydantic.SecretStr | None = None expires_at: pydantic.AwareDatetime | None = None @@ -87,7 +88,8 @@ def handle_incoming_message(self, message: A2AMessage, run_context: RunContext, self._metadata_from_client = self._metadata_from_client or PlatformApiExtensionMetadata() data = self._metadata_from_client data.base_url = data.base_url or HttpUrl(os.getenv("PLATFORM_URL", "http://127.0.0.1:8333")) - data.auth_token = data.auth_token or self._get_header_token(request_context) + auth_token = data.auth_token or self._get_header_token(request_context) + data.auth_token = pydantic.SecretStr(auth_token.get_secret_value()) if auth_token else None if not data.auth_token: raise ExtensionError(self.spec, "Platform extension metadata was not provided") @@ -118,14 +120,11 @@ def api_auth_metadata( base_url: HttpUrl | None = None, ) -> dict[str, dict[str, str]]: return { - self.spec.URI: { - **PlatformApiExtensionMetadata( - base_url=base_url, - auth_token=pydantic.Secret("replaced below"), - expires_at=expires_at, - ).model_dump(mode="json"), - "auth_token": auth_token if isinstance(auth_token, str) else auth_token.get_secret_value(), - } + self.spec.URI: PlatformApiExtensionMetadata( + base_url=base_url, + auth_token=auth_token, + expires_at=expires_at, + ).model_dump(mode="json", context={REVEAL_SECRETS: True}) } diff --git a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/tools/call.py b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/tools/call.py index 916d38dfd5..d3a252c6ad 100644 --- a/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/tools/call.py +++ b/apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/tools/call.py @@ -15,6 +15,7 @@ from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec from agentstack_sdk.a2a.extensions.tools.exceptions import ToolCallRejectionError from agentstack_sdk.a2a.types import AgentMessage, InputRequired +from agentstack_sdk.util.pydantic import REVEAL_SECRETS if TYPE_CHECKING: from agentstack_sdk.server.context import RunContext @@ -67,7 +68,8 @@ class ToolCallExtensionMetadata(BaseModel): class ToolCallExtensionServer(BaseExtensionServer[ToolCallExtensionSpec, ToolCallExtensionMetadata]): def create_request_message(self, *, request: ToolCallRequest): return AgentMessage( - text="Tool call approval requested", metadata={self.spec.URI: request.model_dump(mode="json")} + text="Tool call approval requested", + metadata={self.spec.URI: request.model_dump(mode="json", context={REVEAL_SECRETS: True})}, ) def parse_response(self, *, message: a2a.types.Message): @@ -102,7 +104,7 @@ def create_response_message(self, *, response: ToolCallResponse, task_id: str | role=a2a.types.Role.user, parts=[], task_id=task_id, - metadata={self.spec.URI: response.model_dump(mode="json")}, + metadata={self.spec.URI: response.model_dump(mode="json", context={REVEAL_SECRETS: True})}, ) def parse_request(self, *, message: a2a.types.Message): @@ -111,4 +113,4 @@ def parse_request(self, *, message: a2a.types.Message): return ToolCallRequest.model_validate(data) def metadata(self) -> dict[str, Any]: - return {self.spec.URI: ToolCallExtensionMetadata().model_dump(mode="json")} + return {self.spec.URI: ToolCallExtensionMetadata().model_dump(mode="json", context={REVEAL_SECRETS: True})} diff --git a/apps/agentstack-sdk-py/src/agentstack_sdk/util/pydantic.py b/apps/agentstack-sdk-py/src/agentstack_sdk/util/pydantic.py new file mode 100644 index 0000000000..d36e95016a --- /dev/null +++ b/apps/agentstack-sdk-py/src/agentstack_sdk/util/pydantic.py @@ -0,0 +1,72 @@ +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +from typing import Any + +from pydantic import AnyUrl, BaseModel, Secret, SecretBytes, SecretStr, field_serializer +from pydantic_core.core_schema import SerializationInfo + +REVEAL_SECRETS = "reveal_secrets" +REDACT_SECRETS = "redact_secrets" + +_REDACTED = "***redacted***" + + +class SecureBaseModel(BaseModel): + """ + Base class that automatically handles SecretStr redaction unless explicitly revealed via context. + Inherit from this instead of BaseModel for models that may contain secrets. + """ + + @field_serializer("*") + @classmethod + def _redact_secrets(cls, v: Any, info: SerializationInfo) -> Any: + if isinstance(v, (Secret, SecretStr, SecretBytes)): + return redact_secret(v, info) + elif isinstance(v, AnyUrl): + return redact_url(v, info) + return v + + +def should_reveal(info: SerializationInfo) -> bool: + return bool(info.context and info.context.get(REVEAL_SECRETS, False)) + + +def should_redact(info: SerializationInfo) -> bool: + return bool(info.context and info.context.get(REDACT_SECRETS, False)) + + +def redact_secret( + v: Secret | SecretStr | SecretBytes, info: SerializationInfo +) -> Secret | SecretStr | SecretBytes | str | bytes: + if should_redact(info): + return _REDACTED + elif should_reveal(info): + return v.get_secret_value() + else: + return v + + +def redact_str(v: str, info: SerializationInfo) -> str: + return _REDACTED if should_redact(info) else v + + +def redact_url(v: AnyUrl, info: SerializationInfo) -> AnyUrl: + return ( + AnyUrl.build( + scheme=v.scheme, + username=_REDACTED if v.username else None, + password=_REDACTED if v.password else None, + host=v.host or "", + path=v.path.lstrip("/") if v.path else None, + port=v.port, + query=_REDACTED if v.query else None, + fragment=_REDACTED if v.fragment else None, + ) + if should_redact(info) + else v + ) + + +def redact_dict(v: dict[str, str], info: SerializationInfo) -> dict[str, str]: + return {k: redact_str(val, info) for k, val in v.items()} if should_redact(info) else v diff --git a/apps/agentstack-sdk-py/src/agentstack_sdk/util/telemetry.py b/apps/agentstack-sdk-py/src/agentstack_sdk/util/telemetry.py new file mode 100644 index 0000000000..c3f86820cb --- /dev/null +++ b/apps/agentstack-sdk-py/src/agentstack_sdk/util/telemetry.py @@ -0,0 +1,278 @@ +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +# Modified version of https://github.com/a2aproject/a2a-python/blob/2acd838796d44ab9bfe6ba8c8b4ea0c2571a59dc/src/a2a/utils/telemetry.py + +import asyncio +import functools +import inspect +import logging +from collections.abc import Callable, Iterable, Mapping +from typing import Any + +from opentelemetry import trace +from opentelemetry.trace import SpanKind, StatusCode + +from agentstack_sdk import __version__ + +INSTRUMENTING_MODULE_NAME = "agentstack-python-sdk" +INSTRUMENTING_MODULE_VERSION = __version__ + +logger = logging.getLogger(__name__) +tracer = trace.get_tracer(INSTRUMENTING_MODULE_NAME, INSTRUMENTING_MODULE_VERSION) + + +def trace_function( + func: Callable | None = None, + *, + span_name: str | None = None, + kind: SpanKind = SpanKind.INTERNAL, + attributes: dict[str, Any] | None = None, + attribute_extractor: Callable | None = None, +) -> Callable: + """A decorator to automatically trace a function call with OpenTelemetry. + + This decorator can be used to wrap both sync and async functions. + When applied, it creates a new span for each call to the decorated function. + The span will record the execution time, status (OK or ERROR), and any + exceptions that occur. + + It can be used in two ways: + + 1. As a direct decorator: `@trace_function` + 2. As a decorator factory to provide arguments: `@trace_function(span_name="custom.name")` + + Args: + func (callable, optional): The function to be decorated. If None, + the decorator returns a partial function, allowing it to be called + with arguments. Defaults to None. + span_name (str, optional): Custom name for the span. If None, + it defaults to ``f'{func.__module__}.{func.__name__}'``. + Defaults to None. + kind (SpanKind, optional): The ``opentelemetry.trace.SpanKind`` for the + created span. Defaults to ``SpanKind.INTERNAL``. + attributes (dict, optional): A dictionary of static attributes to be + set on the span. Keys are attribute names (str) and values are + the corresponding attribute values. Defaults to None. + attribute_extractor (callable, optional): A function that can be used + to dynamically extract and set attributes on the span. + It is called within a ``finally`` block, ensuring it runs even if + the decorated function raises an exception. + The function signature should be: + ``attribute_extractor(span, args, kwargs, result, exception)`` + where: + - ``span`` : the OpenTelemetry ``Span`` object. + - ``args`` : a tuple of positional arguments passed + - ``kwargs`` : a dictionary of keyword arguments passed + - ``result`` : return value (None if an exception occurred) + - ``exception`` : exception object if raised (None otherwise). + Any exception raised by the ``attribute_extractor`` itself will be + caught and logged. Defaults to None. + + Returns: + callable: The wrapped function that includes tracing, or a partial + decorator if ``func`` is None. + """ + if func is None: + return functools.partial( + trace_function, + span_name=span_name, + kind=kind, + attributes=attributes, + attribute_extractor=attribute_extractor, + ) + + actual_span_name = span_name or f"{func.__module__}.{func.__name__}" + + is_async_func = inspect.iscoroutinefunction(func) + + @functools.wraps(func) + async def async_wrapper(*args, **kwargs) -> Any: + """Async Wrapper for the decorator.""" + with tracer.start_as_current_span(actual_span_name, kind=kind) as span: + if attributes: + for k, v in attributes.items(): + span.set_attribute(k, v) + + result = None + exception = None + + try: + # Async wrapper, await for the function call to complete. + result = await func(*args, **kwargs) # pyrefly: ignore[not-callable] + span.set_status(StatusCode.OK) + # asyncio.CancelledError extends from BaseException + except asyncio.CancelledError as ce: + exception = None + span.record_exception(ce) + raise + except Exception as e: + exception = e + span.record_exception(e) + span.set_status(StatusCode.ERROR, description=str(e)) + raise + finally: + if attribute_extractor: + try: + attribute_extractor(span, args, kwargs, result, exception) + except Exception: + logger.exception( + "attribute_extractor error in span %s", + actual_span_name, + ) + return result + + @functools.wraps(func) + def sync_wrapper(*args, **kwargs) -> Any: + """Sync Wrapper for the decorator.""" + with tracer.start_as_current_span(actual_span_name, kind=kind) as span: + if attributes: + for k, v in attributes.items(): + span.set_attribute(k, v) + + result = None + exception = None + + try: + # Sync wrapper, execute the function call. + result = func(*args, **kwargs) # pyrefly: ignore[not-callable] + span.set_status(StatusCode.OK) + + except Exception as e: + exception = e + span.record_exception(e) + span.set_status(StatusCode.ERROR, description=str(e)) + raise + finally: + if attribute_extractor: + try: + attribute_extractor(span, args, kwargs, result, exception) + except Exception: + logger.exception( + "attribute_extractor error in span %s", + actual_span_name, + ) + return result + + return async_wrapper if is_async_func else sync_wrapper + + +def trace_class( + include_list: list[str] | None = None, + exclude_list: list[str] | None = None, + kind: SpanKind = SpanKind.INTERNAL, + attributes: dict[str, Any] | None = None, +) -> Callable: + """A class decorator to automatically trace specified methods of a class. + + This decorator iterates over the methods of a class and applies the + `trace_function` decorator to them, based on the `include_list` and + `exclude_list` criteria. Methods starting or ending with double underscores + (dunder methods, e.g., `__init__`, `__call__`) are always excluded by default. + + Args: + include_list (list[str], optional): A list of method names to + explicitly include for tracing. If provided, only methods in this + list (that are not dunder methods) will be traced. + Defaults to None (trace all non-dunder methods). + exclude_list (list[str], optional): A list of method names to exclude + from tracing. This is only considered if `include_list` is not + provided. Dunder methods are implicitly excluded. + Defaults to an empty list. + kind (SpanKind, optional): The `opentelemetry.trace.SpanKind` for the + created spans on the methods. Defaults to `SpanKind.INTERNAL`. + + Returns: + callable: A decorator function that, when applied to a class, + modifies the class to wrap its specified methods with tracing. + + Example: + To trace all methods except 'internal_method': + ```python + @trace_class(exclude_list=['internal_method']) + class MyService: + def public_api(self): + pass + + def internal_method(self): + pass + ``` + + To trace only 'method_one' and 'method_two': + ```python + @trace_class(include_list=['method_one', 'method_two']) + class AnotherService: + def method_one(self): + pass + + def method_two(self): + pass + + def not_traced_method(self): + pass + ``` + """ + exclude_list = exclude_list or [] + + def decorator(cls: Any) -> Any: + for name, method in inspect.getmembers(cls, inspect.isfunction): + if name.startswith("__") and name.endswith("__"): + continue + if include_list and name not in include_list: + continue + if not include_list and name in exclude_list: + continue + + span_name = f"{cls.__module__}.{cls.__name__}.{name}" + setattr( + cls, + name, + trace_function(span_name=span_name, kind=kind, attributes=attributes)(method), + ) + return cls + + return decorator + + +def flatten_dict( + d: Mapping[str, Any], + parent_key: str = "", + sep: str = ".", + skip_none: bool = True, + max_depth: int | None = None, +) -> dict[str, Any]: + """ + Flatten nested dict → OpenTelemetry-compatible flat dict. + + - Lists are usually kept as-is (OTel supports primitive arrays) + - None values are skipped by default + """ + items: list[tuple[str, Any]] = [] + + current_depth = parent_key.count(sep) + 1 if parent_key else 0 + if max_depth is not None and current_depth > max_depth: + # Optional: prevent runaway depth + items.append((parent_key, "[max depth exceeded]")) + return dict(items) + + for k, v in d.items(): + new_key = f"{parent_key}{sep}{k}" if parent_key else k + + if skip_none and v is None: + continue + + if isinstance(v, Mapping): + items.extend(flatten_dict(v, new_key, sep, skip_none, max_depth).items()) + + elif isinstance(v, (str, bool, int, float, type(None))): + items.append((new_key, v)) + + elif isinstance(v, Iterable) and not isinstance(v, (str, bytes)): + # Keep lists/arrays as-is — OpenTelemetry supports them + items.append((new_key, list(v))) # make sure it's a list + + else: + # Fallback: convert unknown types to string + items.append((new_key, str(v))) + + return dict(items) diff --git a/apps/agentstack-sdk-py/tests/e2e/test_extensions.py b/apps/agentstack-sdk-py/tests/e2e/test_extensions.py index 7c4385621b..69957a6ea6 100644 --- a/apps/agentstack-sdk-py/tests/e2e/test_extensions.py +++ b/apps/agentstack-sdk-py/tests/e2e/test_extensions.py @@ -44,7 +44,7 @@ async def chunked_artifact_producer( # Agent producing chunked artifacts await asyncio.sleep(random() * 0.5) - api_key = next(iter(llm_ext.data.llm_fulfillments.values())).api_key + api_key = next(iter(llm_ext.data.llm_fulfillments.values())).api_key.get_secret_value() yield api_key async with create_server_with_agent(chunked_artifact_producer) as (server, test_client): diff --git a/docs/development/agent-integration/rag.mdx b/docs/development/agent-integration/rag.mdx index fc501d338b..bf8e8e1395 100644 --- a/docs/development/agent-integration/rag.mdx +++ b/docs/development/agent-integration/rag.mdx @@ -243,7 +243,9 @@ def get_embedding_client( if not embedding_config: raise ValueError("Default embedding configuration not found") - embedding_client = AsyncOpenAI(api_key=embedding_config.api_key, base_url=embedding_config.api_base) + embedding_client = AsyncOpenAI( + api_key=embedding_config.api_key.get_secret_value(), base_url=embedding_config.api_base + ) embedding_model = embedding_config.api_model return embedding_client, embedding_model diff --git a/examples/agent-integration/rag/simple-rag-agent/src/simple_rag_agent/embedding/client.py b/examples/agent-integration/rag/simple-rag-agent/src/simple_rag_agent/embedding/client.py index ff6bbc2a7e..1f5562a693 100644 --- a/examples/agent-integration/rag/simple-rag-agent/src/simple_rag_agent/embedding/client.py +++ b/examples/agent-integration/rag/simple-rag-agent/src/simple_rag_agent/embedding/client.py @@ -15,6 +15,8 @@ def get_embedding_client( if not embedding_config: raise ValueError("Default embedding configuration not found") - embedding_client = AsyncOpenAI(api_key=embedding_config.api_key, base_url=embedding_config.api_base) + embedding_client = AsyncOpenAI( + api_key=embedding_config.api_key.get_secret_value(), base_url=embedding_config.api_base + ) embedding_model = embedding_config.api_model return embedding_client, embedding_model