Skip to content

feat(BA-3947): Add sessions resolver to strawberry AgentV2 schema #9620

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jopemachine merged 5 commits into from
Mar 6, 2026
+70 −3
Merged

feat(BA-3947): Add sessions resolver to strawberry AgentV2 schema #9620

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a24c18a
feat: Implement session node resolver to Agent
jopemachine Mar 4, 2026
b79b507
chore: update api schema dump
jopemachine Mar 4, 2026
f406948
docs: Add news fragment
jopemachine Mar 4, 2026
469c24e
wip
jopemachine Mar 4, 2026
7f55470
chore: update api schema dump
jopemachine Mar 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changes/9620.feature.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add `sessions` resolver to strawberry `AgentV2` schema
5 changes: 5 additions & 0 deletions docs/manager/graphql-reference/supergraph.graphql
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -520,6 +520,11 @@ type AgentV2 implements Node
Added in 26.2.0. List of kernels running on this agent with pagination support.
"""
kernels(filter: KernelV2Filter = null, orderBy: [KernelV2OrderBy!] = null, before: String = null, after: String = null, first: Int = null, last: Int = null, limit: Int = null, offset: Int = null): KernelV2Connection!

"""
Added in 26.3.0. List of sessions running on this agent with pagination support.
"""
sessions(filter: SessionV2Filter = null, orderBy: [SessionV2OrderBy!] = null, before: String = null, after: String = null, first: Int = null, last: Int = null, limit: Int = null, offset: Int = null): SessionV2Connection!
}

"""
Expand Down
5 changes: 5 additions & 0 deletions docs/manager/graphql-reference/v2-schema.graphql
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -316,6 +316,11 @@ type AgentV2 implements Node {
Added in 26.2.0. List of kernels running on this agent with pagination support.
"""
kernels(filter: KernelV2Filter = null, orderBy: [KernelV2OrderBy!] = null, before: String = null, after: String = null, first: Int = null, last: Int = null, limit: Int = null, offset: Int = null): KernelV2Connection!

"""
Added in 26.3.0. List of sessions running on this agent with pagination support.
"""
sessions(filter: SessionV2Filter = null, orderBy: [SessionV2OrderBy!] = null, before: String = null, after: String = null, first: Int = null, last: Int = null, limit: Int = null, offset: Int = null): SessionV2Connection!
}

"""
Expand Down
49 changes: 48 additions & 1 deletion src/ai/backend/manager/api/gql/agent/types.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,11 @@
KernelV2FilterGQL,
KernelV2OrderByGQL,
)
from ai.backend.manager.api.gql.session.types import (
SessionV2ConnectionGQL,
SessionV2FilterGQL,
SessionV2OrderByGQL,
)
from ai.backend.manager.api.gql.utils import dedent_strip
from ai.backend.manager.data.agent.types import AgentDetailData, AgentStatus
from ai.backend.manager.models.rbac.permission_defs import AgentPermission
Expand All @@ -32,6 +37,7 @@
)
from ai.backend.manager.repositories.scheduler.options import (
KernelConditions,
SessionConditions,
)


Expand Down Expand Up @@ -439,7 +445,48 @@ async def kernels(
last=last,
limit=limit,
offset=offset,
base_conditions=[KernelConditions.by_agent_id(str(self._agent_id))],
base_conditions=[KernelConditions.by_agent_id(self._agent_id)],
)

@strawberry.field( # type: ignore[misc]
description="Added in 26.3.0. List of sessions running on this agent with pagination support."
)
async def sessions(
self,
info: Info[StrawberryGQLContext],
filter: Annotated[
SessionV2FilterGQL, strawberry.lazy("ai.backend.manager.api.gql.session.types")
]
| None = None,
order_by: list[
Annotated[
SessionV2OrderByGQL, strawberry.lazy("ai.backend.manager.api.gql.session.types")
]
]
| None = None,
before: str | None = None,
after: str | None = None,
first: int | None = None,
last: int | None = None,
limit: int | None = None,
offset: int | None = None,
) -> Annotated[
SessionV2ConnectionGQL, strawberry.lazy("ai.backend.manager.api.gql.session.types")
]:
"""Fetch sessions associated with this agent."""
from ai.backend.manager.api.gql.session.fetcher.session import fetch_sessions

return await fetch_sessions(
info=info,
filter=filter,
order_by=order_by,
before=before,
after=after,
first=first,
last=last,
limit=limit,
offset=offset,
base_conditions=[SessionConditions.by_agent_id(self._agent_id)],
)
Comment on lines +451 to 490
Copy link

Copilot AI Mar 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AgentV2GQL.sessions calls fetch_sessions() directly per agent instance. In queries that return a list/connection of agents and request sessions for each node, this will trigger one session-search per agent (N+1 pattern) and can become expensive. Consider introducing a batching approach (e.g., a DataLoader keyed by (agent_id, pagination/filter args) with internal grouping, or a dedicated processor/repository method that can fetch sessions for multiple agent IDs in one call when args match) to keep the number of DB/service calls bounded per GraphQL request.

Copilot uses AI. Check for mistakes.
Comment on lines +451 to 490
Copy link

Copilot AI Mar 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new AgentV2GQL.sessions field introduces a new query surface and filtering/pagination wiring via base_conditions=[SessionConditions.by_agent_id(...)], but there are no unit tests covering this resolver behavior. Adding a resolver-level unit test (similar to other GraphQL resolver tests that patch the fetcher) would help ensure the agent-id condition is always applied and pagination args are forwarded correctly.

Copilot uses AI. Check for mistakes.

@classmethod
Expand Down
13 changes: 11 additions & 2 deletions src/ai/backend/manager/repositories/scheduler/options.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
)
from ai.backend.manager.api.gql.kernel.types import KernelStatusInMatchSpec

from ai.backend.common.types import KernelId, SessionId
from ai.backend.common.types import AgentId, KernelId, SessionId
from ai.backend.manager.data.kernel.types import KernelStatus
from ai.backend.manager.data.session.types import KernelMatchType, SessionStatus
from ai.backend.manager.models.image import ImageRow
Expand Down Expand Up @@ -456,6 +456,15 @@ def inner() -> sa.sql.expression.ColumnElement[bool]:

return inner

@staticmethod
def by_agent_id(agent_id: AgentId) -> QueryCondition:
"""Filter sessions that have kernels running on the given agent."""

def inner() -> sa.sql.expression.ColumnElement[bool]:
return sa.literal(agent_id) == sa.any_(SessionRow.agent_ids)

return inner

@staticmethod
def by_cursor_forward(cursor_id: str) -> QueryCondition:
"""Cursor condition for forward pagination (after cursor).
Expand Down Expand Up @@ -619,7 +628,7 @@ def inner() -> sa.sql.expression.ColumnElement[bool]:
return inner

@staticmethod
def by_agent_id(agent_id: str) -> QueryCondition:
def by_agent_id(agent_id: AgentId) -> QueryCondition:
"""Filter kernels by agent ID."""

def inner() -> sa.sql.expression.ColumnElement[bool]:
Expand Down
Loading