Skip to content

datarobot_genai.dragent.frontends.session

session

DRAgentUserManager

Bases: UserManager

Add DataRobot signed auth-context resolution to NAT's standard identity extractors.

NAT 1.6 replaced the extensible context-based user_id resolution with UserManager.extract_user_from_connection() (#1775), which only supports standard auth (Bearer JWT, cookies, API key). DataRobot passes user identity via X-DataRobot-Authorization-Context (signed app-context JWT), which the vanilla extractor does not understand. This subclass handles that header first, then delegates to super().extract_user_from_connection() so standard auth still works.

Source code in datarobot_genai/dragent/frontends/session.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
class DRAgentUserManager(UserManager):
    """Add DataRobot signed auth-context resolution to NAT's standard identity extractors.

    NAT 1.6 replaced the extensible context-based user_id resolution with
    ``UserManager.extract_user_from_connection()`` (#1775), which only supports
    standard auth (Bearer JWT, cookies, API key). DataRobot passes user identity
    via ``X-DataRobot-Authorization-Context`` (signed app-context JWT), which the
    vanilla extractor does not understand. This subclass handles that header
    first, then delegates to ``super().extract_user_from_connection()`` so
    standard auth still works.
    """

    @classmethod
    def extract_user_from_connection(cls, connection: Request | WebSocket) -> UserInfo | None:
        if isinstance(connection, Request):
            auth_ctx = _auth_handler.get_context(dict(connection.headers))
            if auth_ctx is not None:
                logger.debug(
                    "Resolved user_id from X-DataRobot-Authorization-Context: %s", auth_ctx.user.id
                )
                return UserInfo._from_session_cookie(auth_ctx.user.id)
        return super().extract_user_from_connection(connection)

DRAgentAGUISessionManager

Bases: SessionManager

Source code in datarobot_genai/dragent/frontends/session.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
class DRAgentAGUISessionManager(SessionManager):
    @asynccontextmanager
    async def session(
        self,
        user_id: str | None = None,
        http_connection: HTTPConnection | None = None,
        user_message_id: str | None = None,
        conversation_id: str | None = None,
        user_input_callback: Callable[[InteractionPrompt], Awaitable[HumanResponse]] | None = None,
        user_authentication_callback: Callable[
            [AuthProviderBaseConfig, AuthFlowType], Awaitable[AuthenticatedContext | None]
        ]
        | None = None,
    ) -> AsyncIterator[Session]:
        """Bridge A2A preset, resolve DR headers, default for per-user, and inject A2A headers.

        NAT 1.6+ assigns ``self._context_state.user_id`` from the explicit ``user_id``
        argument only. The A2A adapter calls ``session()`` with no arguments; our
        executor sets ``context_id`` on the context var first, but the parent
        ``session()`` would overwrite it with ``None``. Copy any preset non-empty
        value into ``user_id`` before delegating so per-user workflows work without
        a Bearer JWT (local dev / message-only A2A).

        Additionally, A2A HTTP request headers stored in the module-level
        ``_a2a_headers`` ContextVar by :class:`_PerUserCompatibleAgentExecutor` are
        injected into ``ContextState._metadata``.
        """
        if user_id is None:
            preset = self._context_state.user_id.get()
            if preset:
                user_id = preset

        # Resolve identity via our UserManager (knows DR signed auth-context in
        # addition to NAT's standard extractors). Done explicitly rather than
        # via a module rebind so the default-user fallback below stays scoped
        # to this session() — DRAgentUserManager remains a pure identity
        # resolver that other callers can safely use.
        if user_id is None and isinstance(http_connection, (Request, WebSocket)):
            user_info = DRAgentUserManager.extract_user_from_connection(http_connection)
            if user_info is not None:
                user_id = user_info.get_user_id()

        # Per-user workflows need *a* key in the per-user-builder dict.
        # Fall back to a constant so they do not crash when no identity is
        # available (e.g. locust against a deployed agent). This is a
        # workflow-keying concern, not an identity claim — all such requests
        # share one builder/state.
        if user_id is None and self.is_workflow_per_user:
            logger.debug(
                "No identity resolved for per-user workflow — using %s as key",
                DEFAULT_DR_AGENT_USER_ID,
            )
            user_id = DEFAULT_DR_AGENT_USER_ID

        # Inject A2A request headers BEFORE super().session() so they are
        # available during per-user builder creation (agent card discovery
        # reads Context.get().metadata.headers). This is safe because NAT's
        # session() does NOT call set_metadata_from_http_request() when
        # http_connection is None (the A2A case).
        token_metadata = None
        preset_headers = _a2a_headers.get()
        if preset_headers:
            attrs = _build_metadata_from_headers(preset_headers)
            token_metadata = self._context_state._metadata.set(attrs)

        # Wrap the entire super().session() in try/finally so _metadata is
        # always reset — even if super().session().__aenter__() raises
        # (e.g. per-user builder creation failure).
        try:
            async with super().session(
                user_id=user_id,
                http_connection=http_connection,
                user_message_id=user_message_id,
                conversation_id=conversation_id,
                user_input_callback=user_input_callback,
                user_authentication_callback=user_authentication_callback,
            ) as sess:
                yield sess
        finally:
            if token_metadata is not None:
                self._context_state._metadata.reset(token_metadata)

    def get_workflow_input_schema(self) -> type[BaseModel]:
        """Get workflow input schema for OpenAPI documentation."""
        return DRAgentRunAgentInput

    def get_workflow_streaming_output_schema(self) -> type[BaseModel]:
        """Get workflow streaming output schema for OpenAPI documentation."""
        return DRAgentEventResponse

session async 

session(user_id: str | None = None, http_connection: HTTPConnection | None = None, user_message_id: str | None = None, conversation_id: str | None = None, user_input_callback: Callable[[InteractionPrompt], Awaitable[HumanResponse]] | None = None, user_authentication_callback: Callable[[AuthProviderBaseConfig, AuthFlowType], Awaitable[AuthenticatedContext | None]] | None = None) -> AsyncIterator[Session]

Bridge A2A preset, resolve DR headers, default for per-user, and inject A2A headers.

NAT 1.6+ assigns self._context_state.user_id from the explicit user_id argument only. The A2A adapter calls session() with no arguments; our executor sets context_id on the context var first, but the parent session() would overwrite it with None. Copy any preset non-empty value into user_id before delegating so per-user workflows work without a Bearer JWT (local dev / message-only A2A).

Additionally, A2A HTTP request headers stored in the module-level _a2a_headers ContextVar by :class:_PerUserCompatibleAgentExecutor are injected into ContextState._metadata.

Source code in datarobot_genai/dragent/frontends/session.py 
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
@asynccontextmanager
async def session(
    self,
    user_id: str | None = None,
    http_connection: HTTPConnection | None = None,
    user_message_id: str | None = None,
    conversation_id: str | None = None,
    user_input_callback: Callable[[InteractionPrompt], Awaitable[HumanResponse]] | None = None,
    user_authentication_callback: Callable[
        [AuthProviderBaseConfig, AuthFlowType], Awaitable[AuthenticatedContext | None]
    ]
    | None = None,
) -> AsyncIterator[Session]:
    """Bridge A2A preset, resolve DR headers, default for per-user, and inject A2A headers.

    NAT 1.6+ assigns ``self._context_state.user_id`` from the explicit ``user_id``
    argument only. The A2A adapter calls ``session()`` with no arguments; our
    executor sets ``context_id`` on the context var first, but the parent
    ``session()`` would overwrite it with ``None``. Copy any preset non-empty
    value into ``user_id`` before delegating so per-user workflows work without
    a Bearer JWT (local dev / message-only A2A).

    Additionally, A2A HTTP request headers stored in the module-level
    ``_a2a_headers`` ContextVar by :class:`_PerUserCompatibleAgentExecutor` are
    injected into ``ContextState._metadata``.
    """
    if user_id is None:
        preset = self._context_state.user_id.get()
        if preset:
            user_id = preset

    # Resolve identity via our UserManager (knows DR signed auth-context in
    # addition to NAT's standard extractors). Done explicitly rather than
    # via a module rebind so the default-user fallback below stays scoped
    # to this session() — DRAgentUserManager remains a pure identity
    # resolver that other callers can safely use.
    if user_id is None and isinstance(http_connection, (Request, WebSocket)):
        user_info = DRAgentUserManager.extract_user_from_connection(http_connection)
        if user_info is not None:
            user_id = user_info.get_user_id()

    # Per-user workflows need *a* key in the per-user-builder dict.
    # Fall back to a constant so they do not crash when no identity is
    # available (e.g. locust against a deployed agent). This is a
    # workflow-keying concern, not an identity claim — all such requests
    # share one builder/state.
    if user_id is None and self.is_workflow_per_user:
        logger.debug(
            "No identity resolved for per-user workflow — using %s as key",
            DEFAULT_DR_AGENT_USER_ID,
        )
        user_id = DEFAULT_DR_AGENT_USER_ID

    # Inject A2A request headers BEFORE super().session() so they are
    # available during per-user builder creation (agent card discovery
    # reads Context.get().metadata.headers). This is safe because NAT's
    # session() does NOT call set_metadata_from_http_request() when
    # http_connection is None (the A2A case).
    token_metadata = None
    preset_headers = _a2a_headers.get()
    if preset_headers:
        attrs = _build_metadata_from_headers(preset_headers)
        token_metadata = self._context_state._metadata.set(attrs)

    # Wrap the entire super().session() in try/finally so _metadata is
    # always reset — even if super().session().__aenter__() raises
    # (e.g. per-user builder creation failure).
    try:
        async with super().session(
            user_id=user_id,
            http_connection=http_connection,
            user_message_id=user_message_id,
            conversation_id=conversation_id,
            user_input_callback=user_input_callback,
            user_authentication_callback=user_authentication_callback,
        ) as sess:
            yield sess
    finally:
        if token_metadata is not None:
            self._context_state._metadata.reset(token_metadata)

get_workflow_input_schema 

get_workflow_input_schema() -> type[BaseModel]

Get workflow input schema for OpenAPI documentation.

Source code in datarobot_genai/dragent/frontends/session.py 
179
180
181
def get_workflow_input_schema(self) -> type[BaseModel]:
    """Get workflow input schema for OpenAPI documentation."""
    return DRAgentRunAgentInput

get_workflow_streaming_output_schema 

get_workflow_streaming_output_schema() -> type[BaseModel]

Get workflow streaming output schema for OpenAPI documentation.

Source code in datarobot_genai/dragent/frontends/session.py 
183
184
185
def get_workflow_streaming_output_schema(self) -> type[BaseModel]:
    """Get workflow streaming output schema for OpenAPI documentation."""
    return DRAgentEventResponse