@aws/nx-plugin 0.112.1 → 0.113.0

package/src/utils/agent-connection/files/py-core-mcp/agentcore_mcp_client.py.template

@@ -6,10 +6,32 @@ from mcp.client.streamable_http import streamablehttp_client
 from strands.tools.mcp.mcp_client import MCPClient
 
 from .auth import SigV4HTTPXAuth
+from .session_context import get_current_session_id
+
+SESSION_HEADER = "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"
+
+
+class _SessionHeaderAuth(httpx.Auth):
+    """Stamps the AgentCore session header from the current async context."""
+
+    requires_request_body = True
+    requires_response_body = True
+
+    def __init__(self, inner: httpx.Auth | None = None):
+        self._inner = inner
+
+    def auth_flow(self, request: httpx.Request):  # type: ignore[override]
+        sid = get_current_session_id()
+        if sid:
+            request.headers[SESSION_HEADER] = sid
+        if self._inner is None:
+            yield request
+            return
+        yield from self._inner.auth_flow(request)
 
 
 class AgentCoreMCPClient:
-    """Factory for clients to call MCP servers hosted on Bedrock AgentCore Runtime."""
+    """Factory for MCP clients that forward the current async-context session."""
 
     @staticmethod
     def _build_url(agent_runtime_arn: str) -> tuple[str, str]:
@@ -19,64 +41,49 @@ class AgentCoreMCPClient:
         """
         region = agent_runtime_arn.split(":")[3]
         encoded_arn = agent_runtime_arn.replace(":", "%3A").replace("/", "%2F")
-        url = f"https://bedrock-agentcore.{region}.amazonaws.com/runtimes/{encoded_arn}/invocations?qualifier=DEFAULT"
+        url = (
+            f"https://bedrock-agentcore.{region}.amazonaws.com/runtimes/"
+            f"{encoded_arn}/invocations?qualifier=DEFAULT"
+        )
         return region, url
 
     @staticmethod
-    def _create(
-        agent_runtime_arn: str,
-        session_id: str,
-        headers: dict | None = None,
-        auth_handler: httpx.Auth | None = None,
-    ) -> MCPClient:
-        _, url = AgentCoreMCPClient._build_url(agent_runtime_arn)
+    def _create(url: str, auth: httpx.Auth | None = None) -> MCPClient:
         return MCPClient(
             lambda: streamablehttp_client(
                 url,
-                auth=auth_handler,
+                auth=_SessionHeaderAuth(auth),
                 timeout=120,
                 terminate_on_close=False,
-                headers={
-                    "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": session_id,
-                    **(headers if headers is not None else {}),
-                },
             )
         )
 
     @staticmethod
-    def with_iam_auth(agent_runtime_arn: str, session_id: str) -> MCPClient:
-        """Create an MCP client with IAM (SigV4) authentication.
-
-        Region and credentials are derived automatically from the ARN
-        and the default credential provider chain.
-        """
-        region, _ = AgentCoreMCPClient._build_url(agent_runtime_arn)
-        boto_session = boto3.Session(region_name=region)
-        credentials = boto_session.get_credentials()
+    def with_iam_auth(agent_runtime_arn: str) -> MCPClient:
+        """SigV4-authenticated client for a Bedrock AgentCore runtime."""
+        region, url = AgentCoreMCPClient._build_url(agent_runtime_arn)
+        credentials = boto3.Session(region_name=region).get_credentials()
         return AgentCoreMCPClient._create(
-            agent_runtime_arn=agent_runtime_arn,
-            session_id=session_id,
-            auth_handler=SigV4HTTPXAuth(credentials, "bedrock-agentcore", region),
+            url, SigV4HTTPXAuth(credentials, "bedrock-agentcore", region)
         )
 
     @staticmethod
     def with_jwt_auth(
-        agent_runtime_arn: str,
-        access_token_provider: Callable[[], str],
-        session_id: str,
+        agent_runtime_arn: str, access_token_provider: Callable[[], str]
     ) -> MCPClient:
-        """Create an MCP client with JWT authentication.
-
-        Args:
-            agent_runtime_arn: The ARN of the AgentCore runtime.
-            access_token_provider: A callable that returns the JWT access token.
-            session_id: Session ID for the connection.
-        """
-        access_token = access_token_provider()
-        return AgentCoreMCPClient._create(
-            agent_runtime_arn=agent_runtime_arn,
-            session_id=session_id,
-            headers={
-                "Authorization": f"Bearer {access_token}",
-            },
+        """Bearer-authenticated client for a Bedrock AgentCore runtime."""
+        _, url = AgentCoreMCPClient._build_url(agent_runtime_arn)
+        return MCPClient(
+            lambda: streamablehttp_client(
+                url,
+                auth=_SessionHeaderAuth(),
+                headers={"Authorization": f"Bearer {access_token_provider()}"},
+                timeout=120,
+                terminate_on_close=False,
+            )
         )
+
+    @staticmethod
+    def without_auth(url: str) -> MCPClient:
+        """Plain-HTTP client — for local dev."""
+        return AgentCoreMCPClient._create(url)

package/src/utils/agent-connection/files/py-core-runtime-config/session_context.py.template

@@ -0,0 +1,22 @@
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+
+_session_id_var: ContextVar[str | None] = ContextVar(
+    "agentcore_session_id", default=None
+)
+
+
+def get_current_session_id() -> str | None:
+    """The session ID for the current async scope, or ``None``."""
+    return _session_id_var.get()
+
+
+@contextmanager
+def session_id_context(session_id: str) -> Iterator[None]:
+    """Bind *session_id* as the current session for the scope of the block."""
+    token = _session_id_var.set(session_id)
+    try:
+        yield
+    finally:
+        _session_id_var.reset(token)

package/src/utils/agent-connection/files/py-core-runtime-config/with_session_id.py.template

@@ -0,0 +1,44 @@
+from collections.abc import Callable, Iterator
+from contextlib import AbstractContextManager, ExitStack, contextmanager
+from typing import Any
+
+from .session_context import get_current_session_id
+
+
+@contextmanager
+def with_session_id(
+    agent_factory: Callable[[], AbstractContextManager[Any]],
+    *,
+    name: str,
+    description: str,
+) -> Iterator[Any]:
+    """Wrap an agent factory so each session gets its own cached Agent."""
+    stack = ExitStack()
+    agents: dict[str, Any] = {}
+
+    def _for_session() -> Any:
+        sid = get_current_session_id() or "default"
+        if sid not in agents:
+            agents[sid] = stack.enter_context(agent_factory())
+        return agents[sid]
+
+    proxy_name = name
+    proxy_description = description
+
+    class _SessionRoutingAgent:
+        name: str = proxy_name
+        description: str = proxy_description
+
+        def __getattr__(self, attr: str) -> Any:
+            return getattr(_for_session(), attr)
+
+        def stream_async(self, *args, **kwargs):
+            return _for_session().stream_async(*args, **kwargs)
+
+        def invoke_async(self, *args, **kwargs):
+            return _for_session().invoke_async(*args, **kwargs)
+
+    try:
+        yield _SessionRoutingAgent()
+    finally:
+        stack.close()

package/src/utils/bundle/bundle.js

@@ -16,6 +16,7 @@ const nx_1 = require("../nx");
 const createPythonBundleTarget = ({ packageName, pythonPlatform, bundleTargetName, }) => {
     return {
         cache: true,
+        inputs: ['default', '^production'],
         executor: 'nx:run-commands',
         outputs: [`{workspaceRoot}/dist/{projectRoot}/${bundleTargetName}`],
         options: {

package/src/utils/bundle/bundle.js.map

@@ -1 +1 @@
-{"version":3,"file":"bundle.js","sourceRoot":"","sources":["../../../../../../packages/nx-plugin/src/utils/bundle/bundle.ts"],"names":[],"mappings":";AAAA;;;GAGG;;;AAEH,uCAQoB;AACpB,0CAA2C;AAC3C,gCAAqC;AACrC,oCAAiD;AACjD,8BAA0D;AAsB1D;;GAEG;AACH,MAAM,wBAAwB,GAAG,CAAC,EAChC,WAAW,EACX,cAAc,EACd,gBAAgB,GACgB,EAAuB,EAAE;IACzD,OAAO;QACL,KAAK,EAAE,IAAI;QACX,QAAQ,EAAE,iBAAiB;QAC3B,OAAO,EAAE,CAAC,sCAAsC,gBAAgB,EAAE,CAAC;QACnE,OAAO,EAAE;YACP,QAAQ,EAAE;gBACR,+EAA+E,WAAW,0BAA0B,gBAAgB,mBAAmB;gBACvJ,+FAA+F,cAAc,gCAAgC,gBAAgB,0BAA0B,gBAAgB,mBAAmB;aAC3N;YACD,QAAQ,EAAE,KAAK;SAChB;KACF,CAAC;AACJ,CAAC,CAAC;AAEF;;GAEG;AACI,MAAM,qBAAqB,GAAG,CACnC,OAA6B,EAC7B,IAAmC,EACnC,EAAE;IACF,IAAI,CAAC,OAAO,CAAC,OAAO,EAAE,CAAC;QACrB,OAAO,CAAC,OAAO,GAAG,EAAE,CAAC;IACvB,CAAC;IAED,MAAM,cAAc,GAAG,IAAI,EAAE,cAAc,IAAI,uBAAuB,CAAC;IACvE,MAAM,gBAAgB,GACpB,cAAc,KAAK,wBAAwB,CAAC,CAAC,CAAC,YAAY,CAAC,CAAC,CAAC,YAAY,CAAC;IAE5E,IAAI,CAAC,OAAO,CAAC,OAAO,EAAE,CAAC,gBAAgB,CAAC,EAAE,CAAC;QACzC,OAAO,CAAC,OAAO,CAAC,gBAAgB,CAAC,GAAG;YAClC,GAAG,wBAAwB,CAAC;gBAC1B,WAAW,EAAE,OAAO,CAAC,IAAI;gBACzB,cAAc;gBACd,gBAAgB;aACjB,CAAC;YACF,SAAS,EAAE,CAAC,SAAS,CAAC;SACvB,CAAC;IACJ,CAAC;IAED,mFAAmF;IACnF,IAAA,sCAAiC,EAAC,OAAO,EAAE,QAAQ,EAAE,gBAAgB,CAAC,CAAC;IACvE,IAAA,sCAAiC,EAAC,OAAO,EAAE,OAAO,EAAE,QAAQ,CAAC,CAAC;IAE9D,OAAO;QACL,gBAAgB;QAChB,eAAe,EAAE,IAAA,0BAAiB,EAAC,MAAM,EAAE,OAAO,CAAC,IAAI,EAAE,gBAAgB,CAAC;KAC3E,CAAC;AACJ,CAAC,CAAC;AA/BW,QAAA,qBAAqB,yBA+BhC;AA0BF;;GAEG;AACI,MAAM,yBAAyB,GAAG,KAAK,EAC5C,IAAU,EACV,OAA6B,EAC7B,IAAsC,EACtC,EAAE;IACF,OAAO,CAAC,OAAO,KAAK,EAAE,CAAC;IAEvB,qDAAqD;IACrD,IAAA,sBAAa,EACX,IAAI,EACJ,IAAA,0BAAiB,EAAC,SAAS,EAAE,OAAO,EAAE,IAAI,CAAC,EAC3C,OAAO,CAAC,IAAI,EACZ,EAAE,EACF,EAAE,iBAAiB,EAAE,0BAAiB,CAAC,YAAY,EAAE,CACtD,CAAC;IAEF,wBAAwB;IACxB,IAAI,CAAC,OAAO,CAAC,OAAO,CAAC,MAAM,EAAE,CAAC;QAC5B,OAAO,CAAC,OAAO,CAAC,MAAM,GAAG;YACvB,KAAK,EAAE,IAAI;YACX,OAAO,EAAE,CAAC,2CAA2C,CAAC;YACtD,QAAQ,EAAE,iBAAiB;YAC3B,OAAO,EAAE;gBACP,OAAO,EAAE,gCAAgC;gBACzC,GAAG,EAAE,eAAe;aACrB;YACD,SAAS,EAAE,CAAC,SAAS,CAAC;SACvB,CAAC;IACJ,CAAC;IAED,iCAAiC;IACjC,IAAA,sCAAiC,EAAC,OAAO,EAAE,OAAO,EAAE,QAAQ,CAAC,CAAC;IAE9D,MAAM,kBAAkB,GAAG,IAAA,0BAAiB,EAC1C,OAAO,CAAC,IAAI,EACZ,oBAAoB,CACrB,CAAC;IAEF,MAAM,UAAU,GAAG,IAAA,0BAAiB,EAClC,IAAA,6BAAqB,EAAC,IAAI,EAAE,OAAO,CAAC,IAAI,CAAC,EACzC,MAAM,EACN,OAAO,CAAC,IAAI,EACZ,QAAQ,EACR,IAAI,CAAC,eAAe,IAAI,GAAG,EAC3B,UAAU,CACX,CAAC;IAEF,MAAM,QAAQ,GAAG,IAAI,CAAC,QAAQ,EAAE,MAAM;QACpC,CAAC,CAAC,IAAI,CAAC,QAAQ;aACV,GAAG,CAAC,CAAC,GAAG,EAAE,EAAE,CACX,OAAO,GAAG,KAAK,QAAQ,CAAC,CAAC,CAAC,IAAI,GAAG,GAAG,CAAC,CAAC,CAAC,IAAI,GAAG,CAAC,MAAM,GAAG,CACzD;aACA,IAAI,CAAC,IAAI,CAAC;QACf,CAAC,CAAC,EAAE,CAAC;IAEP,MAAM,KAAK,GAAG;;cAEF,IAAI,CAAC,cAAc;uBACV,UAAU;iBAChB,IAAI,CAAC,QAAQ,IAAI,MAAM;MAClC,QAAQ,CAAC,CAAC,CAAC,cAAc,QAAQ,IAAI,CAAC,CAAC,CAAC,EAAE;IAC5C,CAAC;IAEH,mEAAmE;IACnE,yEAAyE;IACzE,+EAA+E;IAC/E,MAAM,QAAQ,GAAG,MAAM,IAAA,iBAAW,EAChC,IAAI,EACJ,kBAAkB,EAClB,gEAAgE,IAAI,CAAC,cAAc,qBAAqB,KAAK,KAAK,CACnH,CAAC;IAEF,IAAI,CAAC,QAAQ,EAAE,CAAC;QACd,wEAAwE;QACxE,uEAAuE;QACvE,kEAAkE;QAClE,8DAA8D;QAC9D,MAAM,QAAQ,GAAG,MAAM,IAAA,iBAAW,EAChC,IAAI,EACJ,kBAAkB,EAClB,6DAA6D,CAC9D,CAAC;QACF,IAAI,QAAQ,EAAE,CAAC;YACb,MAAM,IAAA,iBAAW,EACf,IAAI,EACJ,kBAAkB,EAClB,qFAAqF,KAAK,KAAK,CAChG,CAAC;QACJ,CAAC;IACH,CAAC;IAED,IAAA,qCAA4B,EAAC,IAAI,EAAE,EAAE,EAAE,IAAA,uBAAY,EAAC,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC;AACrE,CAAC,CAAC;AA5FW,QAAA,yBAAyB,6BA4FpC"}
+{"version":3,"file":"bundle.js","sourceRoot":"","sources":["../../../../../../packages/nx-plugin/src/utils/bundle/bundle.ts"],"names":[],"mappings":";AAAA;;;GAGG;;;AAEH,uCAQoB;AACpB,0CAA2C;AAC3C,gCAAqC;AACrC,oCAAiD;AACjD,8BAA0D;AAsB1D;;GAEG;AACH,MAAM,wBAAwB,GAAG,CAAC,EAChC,WAAW,EACX,cAAc,EACd,gBAAgB,GACgB,EAAuB,EAAE;IACzD,OAAO;QACL,KAAK,EAAE,IAAI;QACX,MAAM,EAAE,CAAC,SAAS,EAAE,aAAa,CAAC;QAClC,QAAQ,EAAE,iBAAiB;QAC3B,OAAO,EAAE,CAAC,sCAAsC,gBAAgB,EAAE,CAAC;QACnE,OAAO,EAAE;YACP,QAAQ,EAAE;gBACR,+EAA+E,WAAW,0BAA0B,gBAAgB,mBAAmB;gBACvJ,+FAA+F,cAAc,gCAAgC,gBAAgB,0BAA0B,gBAAgB,mBAAmB;aAC3N;YACD,QAAQ,EAAE,KAAK;SAChB;KACF,CAAC;AACJ,CAAC,CAAC;AAEF;;GAEG;AACI,MAAM,qBAAqB,GAAG,CACnC,OAA6B,EAC7B,IAAmC,EACnC,EAAE;IACF,IAAI,CAAC,OAAO,CAAC,OAAO,EAAE,CAAC;QACrB,OAAO,CAAC,OAAO,GAAG,EAAE,CAAC;IACvB,CAAC;IAED,MAAM,cAAc,GAAG,IAAI,EAAE,cAAc,IAAI,uBAAuB,CAAC;IACvE,MAAM,gBAAgB,GACpB,cAAc,KAAK,wBAAwB,CAAC,CAAC,CAAC,YAAY,CAAC,CAAC,CAAC,YAAY,CAAC;IAE5E,IAAI,CAAC,OAAO,CAAC,OAAO,EAAE,CAAC,gBAAgB,CAAC,EAAE,CAAC;QACzC,OAAO,CAAC,OAAO,CAAC,gBAAgB,CAAC,GAAG;YAClC,GAAG,wBAAwB,CAAC;gBAC1B,WAAW,EAAE,OAAO,CAAC,IAAI;gBACzB,cAAc;gBACd,gBAAgB;aACjB,CAAC;YACF,SAAS,EAAE,CAAC,SAAS,CAAC;SACvB,CAAC;IACJ,CAAC;IAED,mFAAmF;IACnF,IAAA,sCAAiC,EAAC,OAAO,EAAE,QAAQ,EAAE,gBAAgB,CAAC,CAAC;IACvE,IAAA,sCAAiC,EAAC,OAAO,EAAE,OAAO,EAAE,QAAQ,CAAC,CAAC;IAE9D,OAAO;QACL,gBAAgB;QAChB,eAAe,EAAE,IAAA,0BAAiB,EAAC,MAAM,EAAE,OAAO,CAAC,IAAI,EAAE,gBAAgB,CAAC;KAC3E,CAAC;AACJ,CAAC,CAAC;AA/BW,QAAA,qBAAqB,yBA+BhC;AA0BF;;GAEG;AACI,MAAM,yBAAyB,GAAG,KAAK,EAC5C,IAAU,EACV,OAA6B,EAC7B,IAAsC,EACtC,EAAE;IACF,OAAO,CAAC,OAAO,KAAK,EAAE,CAAC;IAEvB,qDAAqD;IACrD,IAAA,sBAAa,EACX,IAAI,EACJ,IAAA,0BAAiB,EAAC,SAAS,EAAE,OAAO,EAAE,IAAI,CAAC,EAC3C,OAAO,CAAC,IAAI,EACZ,EAAE,EACF,EAAE,iBAAiB,EAAE,0BAAiB,CAAC,YAAY,EAAE,CACtD,CAAC;IAEF,wBAAwB;IACxB,IAAI,CAAC,OAAO,CAAC,OAAO,CAAC,MAAM,EAAE,CAAC;QAC5B,OAAO,CAAC,OAAO,CAAC,MAAM,GAAG;YACvB,KAAK,EAAE,IAAI;YACX,OAAO,EAAE,CAAC,2CAA2C,CAAC;YACtD,QAAQ,EAAE,iBAAiB;YAC3B,OAAO,EAAE;gBACP,OAAO,EAAE,gCAAgC;gBACzC,GAAG,EAAE,eAAe;aACrB;YACD,SAAS,EAAE,CAAC,SAAS,CAAC;SACvB,CAAC;IACJ,CAAC;IAED,iCAAiC;IACjC,IAAA,sCAAiC,EAAC,OAAO,EAAE,OAAO,EAAE,QAAQ,CAAC,CAAC;IAE9D,MAAM,kBAAkB,GAAG,IAAA,0BAAiB,EAC1C,OAAO,CAAC,IAAI,EACZ,oBAAoB,CACrB,CAAC;IAEF,MAAM,UAAU,GAAG,IAAA,0BAAiB,EAClC,IAAA,6BAAqB,EAAC,IAAI,EAAE,OAAO,CAAC,IAAI,CAAC,EACzC,MAAM,EACN,OAAO,CAAC,IAAI,EACZ,QAAQ,EACR,IAAI,CAAC,eAAe,IAAI,GAAG,EAC3B,UAAU,CACX,CAAC;IAEF,MAAM,QAAQ,GAAG,IAAI,CAAC,QAAQ,EAAE,MAAM;QACpC,CAAC,CAAC,IAAI,CAAC,QAAQ;aACV,GAAG,CAAC,CAAC,GAAG,EAAE,EAAE,CACX,OAAO,GAAG,KAAK,QAAQ,CAAC,CAAC,CAAC,IAAI,GAAG,GAAG,CAAC,CAAC,CAAC,IAAI,GAAG,CAAC,MAAM,GAAG,CACzD;aACA,IAAI,CAAC,IAAI,CAAC;QACf,CAAC,CAAC,EAAE,CAAC;IAEP,MAAM,KAAK,GAAG;;cAEF,IAAI,CAAC,cAAc;uBACV,UAAU;iBAChB,IAAI,CAAC,QAAQ,IAAI,MAAM;MAClC,QAAQ,CAAC,CAAC,CAAC,cAAc,QAAQ,IAAI,CAAC,CAAC,CAAC,EAAE;IAC5C,CAAC;IAEH,mEAAmE;IACnE,yEAAyE;IACzE,+EAA+E;IAC/E,MAAM,QAAQ,GAAG,MAAM,IAAA,iBAAW,EAChC,IAAI,EACJ,kBAAkB,EAClB,gEAAgE,IAAI,CAAC,cAAc,qBAAqB,KAAK,KAAK,CACnH,CAAC;IAEF,IAAI,CAAC,QAAQ,EAAE,CAAC;QACd,wEAAwE;QACxE,uEAAuE;QACvE,kEAAkE;QAClE,8DAA8D;QAC9D,MAAM,QAAQ,GAAG,MAAM,IAAA,iBAAW,EAChC,IAAI,EACJ,kBAAkB,EAClB,6DAA6D,CAC9D,CAAC;QACF,IAAI,QAAQ,EAAE,CAAC;YACb,MAAM,IAAA,iBAAW,EACf,IAAI,EACJ,kBAAkB,EAClB,qFAAqF,KAAK,KAAK,CAChG,CAAC;QACJ,CAAC;IACH,CAAC;IAED,IAAA,qCAA4B,EAAC,IAAI,EAAE,EAAE,EAAE,IAAA,uBAAY,EAAC,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC;AACrE,CAAC,CAAC;AA5FW,QAAA,yBAAyB,6BA4FpC"}