chore: misc fixes and improvements for connectors (#1484)

Signed-off-by: Tomas Pilar <[email protected]>

Author: pilartomas

apps/agentstack-sdk-py/src/agentstack_sdk/a2a/extensions/services/mcp.py

@@ -3,9 +3,10 @@
 
 from __future__ import annotations
 
+import re
 from contextlib import asynccontextmanager
 from types import NoneType
-from typing import Annotated, Any, Literal, Self
+from typing import TYPE_CHECKING, Annotated, Any, Literal, Self
 
 import a2a.types
 import pydantic
@@ -14,6 +15,10 @@
 
 from agentstack_sdk.a2a.extensions.auth.oauth.oauth import OAuthExtensionServer
 from agentstack_sdk.a2a.extensions.base import BaseExtensionClient, BaseExtensionServer, BaseExtensionSpec
+from agentstack_sdk.util.logging import logger
+
+if TYPE_CHECKING:
+    from agentstack_sdk.server.context import RunContext
 
 _TRANSPORT_TYPES = Literal["streamable_http", "stdio"]
 
@@ -95,6 +100,23 @@ class MCPServiceExtensionMetadata(pydantic.BaseModel):
 
 
 class MCPServiceExtensionServer(BaseExtensionServer[MCPServiceExtensionSpec, MCPServiceExtensionMetadata]):
+    def handle_incoming_message(self, message: a2a.types.Message, context: RunContext):
+        from agentstack_sdk.platform import get_platform_client
+
+        super().handle_incoming_message(message, context)
+        if not self.data:
+            return
+
+        platform_url = str(get_platform_client().base_url)
+        for fullfilment in self.data.mcp_fulfillments.values():
+            if fullfilment.transport.type == "streamable_http":
+                try:
+                    fullfilment.transport.url = pydantic.AnyHttpUrl(
+                        re.sub(r"^http[s]?://{platform_url}", platform_url, str(fullfilment.transport.url))
+                    )
+                except Exception:
+                    logger.warning("Platform URL substitution failed", exc_info=True)
+
     def parse_client_metadata(self, message: a2a.types.Message) -> MCPServiceExtensionMetadata | None:
         metadata = super().parse_client_metadata(message)
         if metadata:

apps/agentstack-sdk-py/src/agentstack_sdk/server/agent.py

@@ -41,9 +41,9 @@
 from agentstack_sdk.server.constants import _IMPLICIT_DEPENDENCY_PREFIX
 from agentstack_sdk.server.context import RunContext
 from agentstack_sdk.server.dependencies import extract_dependencies
-from agentstack_sdk.server.logging import logger
 from agentstack_sdk.server.store.context_store import ContextStore
 from agentstack_sdk.server.utils import cancel_task, close_queue
+from agentstack_sdk.util.logging import logger
 from agentstack_sdk.util.utils import extract_messages
 
 AgentFunction: TypeAlias = Callable[[], AsyncGenerator[RunYield, RunYieldResume]]

apps/agentstack-sdk-py/src/agentstack_sdk/server/server.py

@@ -36,12 +36,12 @@
 from agentstack_sdk.platform.provider import Provider
 from agentstack_sdk.server.agent import Agent, AgentFactory
 from agentstack_sdk.server.agent import agent as agent_decorator
-from agentstack_sdk.server.logging import configure_logger as configure_logger_func
-from agentstack_sdk.server.logging import logger
 from agentstack_sdk.server.store.context_store import ContextStore
 from agentstack_sdk.server.store.memory_context_store import InMemoryContextStore
 from agentstack_sdk.server.telemetry import configure_telemetry as configure_telemetry_func
 from agentstack_sdk.server.utils import cancel_task
+from agentstack_sdk.util.logging import configure_logger as configure_logger_func
+from agentstack_sdk.util.logging import logger
 
 
 class Server:

…-py/src/agentstack_sdk/server/logging.py …dk-py/src/agentstack_sdk/util/logging.pyapps/agentstack-sdk-py/src/agentstack_sdk/server/logging.py renamed to apps/agentstack-sdk-py/src/agentstack_sdk/util/logging.py apps/agentstack-sdk-py/src/agentstack_sdk/server/logging.py renamed to apps/agentstack-sdk-py/src/agentstack_sdk/util/logging.py

@@ -27,7 +27,7 @@
 
 
 def setup_database_engine(config: Configuration) -> AsyncEngine:
-    return config.persistence.create_async_engine(isolation_level="READ COMMITTED")
+    return config.persistence.create_async_engine(isolation_level="READ COMMITTED", hide_parameters=True)
 
 
 async def setup_kubernetes_client(namespace: str | None = None, kubeconfig: pathlib.Path | str | dict | None = None):

apps/agentstack-server/src/agentstack_server/bootstrap.py

@@ -5,7 +5,7 @@
 from typing import Annotated, Literal
 from uuid import UUID, uuid4
 
-from pydantic import AnyUrl, AwareDatetime, BaseModel, Field
+from pydantic import AnyUrl, AwareDatetime, BaseModel, ConfigDict, Field
 
 from agentstack_server.domain.models.common import Metadata
 from agentstack_server.utils.utils import utc_now
@@ -27,12 +27,15 @@ class Token(BaseModel):
     refresh_token: str | None = None
     token_type: Literal["bearer"]
 
+    model_config = ConfigDict(extra="allow")
+
 
 class Authorization(BaseModel):
     client_id: str | None = None
     client_secret: str | None = None
     flow: AuthFlow | None = None
     token: Token | None = None
+    token_endpoint: AnyUrl | None = None
 
 
 class ConnectorState(StrEnum):

apps/agentstack-server/src/agentstack_server/domain/models/connector.py

@@ -14,7 +14,7 @@
 blueprint = Blueprint()
 
 
-@blueprint.periodic(cron="* * * * * */15")
+@blueprint.periodic(cron="* * * * * */30")
 @blueprint.task(queueing_lock="refresh_connectors", queue=str(Queues.CRON_CONNECTOR))
 @inject
 async def refresh_connectors(timestamp: int, service: ConnectorService):

apps/agentstack-server/src/agentstack_server/jobs/crons/connector.py

@@ -9,6 +9,7 @@
 
 from agentstack_server.configuration import Configuration
 from agentstack_server.jobs.crons.cleanup import blueprint as cleanup_crons
+from agentstack_server.jobs.crons.connector import blueprint as connector_crons
 from agentstack_server.jobs.crons.provider import blueprint as provider_crons
 from agentstack_server.jobs.tasks.context import blueprint as context_tasks
 from agentstack_server.jobs.tasks.file import blueprint as file_tasks
@@ -52,4 +53,5 @@ def exit_app_on_db_error(*_args, **_kwargs):
     app.add_tasks_from(blueprint=provider_build_tasks, namespace="provider_build_tasks")
     app.add_tasks_from(blueprint=provider_crons, namespace="cron_provider")
     app.add_tasks_from(blueprint=cleanup_crons, namespace="cron_cleanup")
+    app.add_tasks_from(blueprint=connector_crons, namespace="cron_connector")
     return app

apps/agentstack-server/src/agentstack_server/jobs/procrastinate.py

@@ -141,7 +141,10 @@ async def disconnect_connector(self, *, connector_id: UUID, user: User | None =
 
         await self._revoke_auth_token(connector=connector)
 
+        if connector.auth:
+            connector.auth.flow = None
         connector.state = ConnectorState.disconnected
+        connector.disconnect_reason = "Client request"
 
         async with self._uow() as uow:
             await uow.connectors.update(connector=connector)
@@ -176,12 +179,16 @@ async def oauth_callback(self, *, callback_url: str, state: str, error: str | No
                 auth_metadata = await self._discover_auth_metadata(connector=connector)
                 if not auth_metadata:
                     raise RuntimeError("Authorization server no longer contains necessary metadata")
+                token_endpoint = auth_metadata.get("token_endpoint")
+                if not token_endpoint:
+                    raise RuntimeError("Authorization server has no token endpoint in metadata")
                 token = await client.fetch_token(
-                    auth_metadata.get("token_endpoint"),
+                    token_endpoint,
                     authorization_response=callback_url,
                     code_verifier=connector.auth.flow.code_verifier,
                 )
                 connector.auth.token = Token.model_validate(token)
+                connector.auth.token_endpoint = AnyUrl(str(token_endpoint))
                 connector.auth.flow = None
             try:
                 await self.probe_connector(connector=connector)
@@ -238,6 +245,9 @@ async def refresh_connector(self, *, connector_id: UUID, user: User | None = Non
         try:
             await self.probe_connector(connector=connector)
         except Exception as err:
+            await self._revoke_auth_token(connector=connector)
+            if connector.auth:
+                connector.auth.flow = None
             connector.state = ConnectorState.disconnected
             connector.disconnect_reason = str(err)
         finally:
@@ -295,6 +305,7 @@ async def _revoke_auth_token(self, *, connector: Connector) -> None:
                 logger.warning("Token revocation failed", exc_info=True)
 
             connector.auth.token = None
+            connector.auth.token_endpoint = None
             async with self._uow() as uow:
                 await uow.connectors.update(connector=connector)
                 await uow.commit()
@@ -329,6 +340,8 @@ async def update_token(token, refresh_token=None, access_token=None):
             code_challenge_method="S256",
             headers=headers,
             timeout=timeout,
+            leeway=60,  # A job probes connectors every 30 seconds, ensuring the token is valid roughly for at least 30 seconds per request.
+            token_endpoint=str(connector.auth.token_endpoint),
         )
 
     async def _discover_auth_metadata(self, *, connector: Connector) -> AuthorizationServerMetadata | None: