fix(agent): support restoring roster versions (#37734)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

api/controllers/console/agent/roster.py

@@ -36,6 +36,7 @@ from extensions.ext_database import db
 from fields.agent_fields import (
     AgentConfigSnapshotDetailResponse,
     AgentConfigSnapshotListResponse,
+    AgentConfigSnapshotRestoreResponse,
     AgentInviteOptionsResponse,
     AgentLogListResponse,
     AgentLogMessageListResponse,
@@ -223,6 +224,7 @@ register_response_schema_models(
     AgentAppPartial,
     AgentConfigSnapshotDetailResponse,
     AgentConfigSnapshotListResponse,
+    AgentConfigSnapshotRestoreResponse,
     AgentInviteOptionsResponse,
     AgentLogListResponse,
     AgentLogMessageListResponse,
@@ -649,3 +651,24 @@ class AgentRosterVersionDetailApi(Resource):
                 version_id=str(version_id),
             ),
         )
+
+
+@console_ns.route("/agent/<uuid:agent_id>/versions/<uuid:version_id>/restore")
+class AgentRosterVersionRestoreApi(Resource):
+    @console_ns.response(200, "Agent version restored", console_ns.models[AgentConfigSnapshotRestoreResponse.__name__])
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @edit_permission_required
+    @with_current_user
+    @with_current_tenant_id
+    def post(self, tenant_id: str, current_user: Account, agent_id: UUID, version_id: UUID):
+        return dump_response(
+            AgentConfigSnapshotRestoreResponse,
+            _agent_roster_service().restore_agent_version(
+                tenant_id=tenant_id,
+                agent_id=str(agent_id),
+                version_id=str(version_id),
+                account_id=current_user.id,
+            ),
+        )

api/fields/agent_fields.py

@@ -291,6 +291,11 @@ class AgentConfigSnapshotListResponse(ResponseModel):
     data: list[AgentConfigSnapshotSummaryResponse]
 
 
+class AgentConfigSnapshotRestoreResponse(ResponseModel):
+    result: Literal["success"]
+    active_config_snapshot_id: str
+
+
 class AgentComposerAgentResponse(ResponseModel):
     id: str
     name: str

api/models/agent.py

@@ -83,6 +83,8 @@ class AgentConfigRevisionOperation(StrEnum):
     SAVE_NEW_AGENT = "save_new_agent"
     # Promotes a workflow-only Agent into the reusable Agent Roster.
     SAVE_TO_ROSTER = "save_to_roster"
+    # Switches the Agent's current published config back to an existing version.
+    RESTORE_VERSION = "restore_version"
 
 
 class WorkflowAgentBindingType(StrEnum):

api/openapi/markdown/console-openapi.md

@@ -936,6 +936,20 @@ Infer CLI tool + ENV suggestions from a standardized Agent App skill
 | ---- | ----------- | ------ |
 | 200 | Agent version detail | **application/json**: [AgentConfigSnapshotDetailResponse](#agentconfigsnapshotdetailresponse)<br> |
 
+### [POST] /agent/{agent_id}/versions/{version_id}/restore
+#### Parameters
+
+| Name | Located in | Description | Required | Schema |
+| ---- | ---------- | ----------- | -------- | ------ |
+| agent_id | path |  | Yes | string (uuid) |
+| version_id | path |  | Yes | string (uuid) |
+
+#### Responses
+
+| Code | Description | Schema |
+| ---- | ----------- | ------ |
+| 200 | Agent version restored | **application/json**: [AgentConfigSnapshotRestoreResponse](#agentconfigsnapshotrestoreresponse)<br> |
+
 ### [GET] /all-workspaces
 #### Parameters
 
@@ -12405,6 +12419,13 @@ Audit operation recorded for Agent Soul version/revision changes.
 | ---- | ---- | ----------- | -------- |
 | data | [ [AgentConfigSnapshotSummaryResponse](#agentconfigsnapshotsummaryresponse) ] |  | Yes |
 
+#### AgentConfigSnapshotRestoreResponse
+
+| Name | Type | Description | Required |
+| ---- | ---- | ----------- | -------- |
+| active_config_snapshot_id | string |  | Yes |
+| result | string |  | Yes |
+
 #### AgentConfigSnapshotSummaryResponse
 
 | Name | Type | Description | Required |

api/services/agent/roster_service.py

@@ -666,12 +666,16 @@ class AgentRosterService:
     @staticmethod
     def _visible_version_operations(agent: Agent) -> set[AgentConfigRevisionOperation]:
         if agent.source == AgentSource.AGENT_APP:
-            return {AgentConfigRevisionOperation.SAVE_NEW_VERSION}
+            return {
+                AgentConfigRevisionOperation.SAVE_NEW_VERSION,
+                AgentConfigRevisionOperation.RESTORE_VERSION,
+            }
         return {
             AgentConfigRevisionOperation.CREATE_VERSION,
             AgentConfigRevisionOperation.SAVE_NEW_VERSION,
             AgentConfigRevisionOperation.SAVE_NEW_AGENT,
             AgentConfigRevisionOperation.SAVE_TO_ROSTER,
+            AgentConfigRevisionOperation.RESTORE_VERSION,
         }
 
     def active_config_is_published(self, *, tenant_id: str, agent: Agent) -> bool:
@@ -764,6 +768,46 @@ class AgentRosterService:
         ]
         return result
 
+    def restore_agent_version(
+        self, *, tenant_id: str, agent_id: str, version_id: str, account_id: str
+    ) -> dict[str, Any]:
+        agent = self._get_agent(tenant_id=tenant_id, agent_id=agent_id, roster_only=True)
+        visible_version_ids = self._visible_version_ids_stmt(tenant_id=tenant_id, agent_id=agent_id, agent=agent)
+        visible_version_id = self._session.scalar(
+            select(AgentConfigSnapshot.id)
+            .where(
+                AgentConfigSnapshot.tenant_id == tenant_id,
+                AgentConfigSnapshot.agent_id == agent_id,
+                AgentConfigSnapshot.id == version_id,
+                AgentConfigSnapshot.id.in_(select(visible_version_ids.c.current_snapshot_id)),
+            )
+            .limit(1)
+        )
+        if not visible_version_id:
+            raise AgentVersionNotFoundError()
+
+        version = self._get_version(tenant_id=tenant_id, agent_id=agent_id, version_id=version_id)
+        if agent.active_config_snapshot_id == version.id:
+            return {"result": "success", "active_config_snapshot_id": version.id}
+
+        previous_snapshot_id = agent.active_config_snapshot_id
+        agent.active_config_snapshot_id = version.id
+        agent.active_config_has_model = agent_soul_has_model(version.config_snapshot)
+        agent.updated_by = account_id
+        self._session.add(
+            AgentConfigRevision(
+                tenant_id=tenant_id,
+                agent_id=agent_id,
+                previous_snapshot_id=previous_snapshot_id,
+                current_snapshot_id=version.id,
+                revision=self._next_revision(tenant_id=tenant_id, agent_id=agent_id),
+                operation=AgentConfigRevisionOperation.RESTORE_VERSION,
+                created_by=account_id,
+            )
+        )
+        self._session.commit()
+        return {"result": "success", "active_config_snapshot_id": version.id}
+
     def _get_agent(self, *, tenant_id: str, agent_id: str, roster_only: bool = False) -> Agent:
         stmt = select(Agent).where(Agent.tenant_id == tenant_id, Agent.id == agent_id)
         if roster_only:
@@ -789,6 +833,17 @@ class AgentRosterService:
             raise AgentVersionNotFoundError()
         return version
 
+    def _next_revision(self, *, tenant_id: str, agent_id: str) -> int:
+        return (
+            self._session.scalar(
+                select(func.max(AgentConfigRevision.revision)).where(
+                    AgentConfigRevision.tenant_id == tenant_id,
+                    AgentConfigRevision.agent_id == agent_id,
+                )
+            )
+            or 0
+        ) + 1
+
     def _load_published_active_snapshot_agent_ids(self, *, tenant_id: str, agents: list[Agent]) -> set[str]:
         predicates = [
             and_(

api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py

@@ -28,6 +28,7 @@ from controllers.console.agent.roster import (
     AgentLogsApi,
     AgentLogSourcesApi,
     AgentRosterVersionDetailApi,
+    AgentRosterVersionRestoreApi,
     AgentRosterVersionsApi,
     AgentStatisticsSummaryApi,
 )
@@ -158,6 +159,9 @@ def test_agent_v2_console_routes_are_agent_id_first() -> None:
         "/agent/<uuid:agent_id>/logs/<uuid:conversation_id>/messages",
         "/agent/<uuid:agent_id>/log-sources",
         "/agent/<uuid:agent_id>/statistics/summary",
+        "/agent/<uuid:agent_id>/versions",
+        "/agent/<uuid:agent_id>/versions/<uuid:version_id>",
+        "/agent/<uuid:agent_id>/versions/<uuid:version_id>/restore",
         "/agent/invite-options",
     ):
         assert route in paths
@@ -513,6 +517,13 @@ def test_agent_versions_call_services(app: Flask, monkeypatch: pytest.MonkeyPatc
             ],
         },
     )
+    captured_restore: dict[str, object] = {}
+
+    def restore_agent_version(_self, **kwargs):
+        captured_restore.update(kwargs)
+        return {"result": "success", "active_config_snapshot_id": kwargs["version_id"]}
+
+    monkeypatch.setattr(roster_controller.AgentRosterService, "restore_agent_version", restore_agent_version)
 
     assert (
         unwrap(AgentRosterVersionsApi.get)(AgentRosterVersionsApi(), "tenant-1", agent_id)["data"][0]["id"]
@@ -523,6 +534,16 @@ def test_agent_versions_call_services(app: Flask, monkeypatch: pytest.MonkeyPatc
     )
     assert version_detail["id"] == version_id
     assert version_detail["agent_id"] == agent_id
+    restored = unwrap(AgentRosterVersionRestoreApi.post)(
+        AgentRosterVersionRestoreApi(), "tenant-1", SimpleNamespace(id="account-1"), agent_id, version_id
+    )
+    assert restored == {"result": "success", "active_config_snapshot_id": version_id}
+    assert captured_restore == {
+        "tenant_id": "tenant-1",
+        "agent_id": agent_id,
+        "version_id": version_id,
+        "account_id": "account-1",
+    }
 
 
 def test_agent_observability_routes_resolve_app_from_agent_id(

api/tests/unit_tests/models/test_agent.py

@@ -33,6 +33,7 @@ def test_agent_enums_match_prd_boundaries():
     assert AgentStatus.ACTIVE.value == "active"
     assert AgentStatus.ARCHIVED.value == "archived"
     assert AgentConfigRevisionOperation.SAVE_CURRENT_VERSION.value == "save_current_version"
+    assert AgentConfigRevisionOperation.RESTORE_VERSION.value == "restore_version"
     assert WorkflowAgentBindingType.ROSTER_AGENT.value == "roster_agent"
     assert WorkflowAgentBindingType.INLINE_AGENT.value == "inline_agent"
 

api/tests/unit_tests/services/agent/test_agent_services.py

@@ -1046,12 +1046,93 @@ def test_agent_app_visible_versions_exclude_draft_saves():
     agent_app_operations = AgentRosterService._visible_version_operations(agent_app)
     roster_operations = AgentRosterService._visible_version_operations(roster_agent)
 
-    assert agent_app_operations == {AgentConfigRevisionOperation.SAVE_NEW_VERSION}
+    assert agent_app_operations == {
+        AgentConfigRevisionOperation.SAVE_NEW_VERSION,
+        AgentConfigRevisionOperation.RESTORE_VERSION,
+    }
     assert AgentConfigRevisionOperation.SAVE_CURRENT_VERSION not in agent_app_operations
     assert AgentConfigRevisionOperation.CREATE_VERSION in roster_operations
+    assert AgentConfigRevisionOperation.RESTORE_VERSION in roster_operations
     assert AgentConfigRevisionOperation.SAVE_CURRENT_VERSION not in roster_operations
 
 
+def test_restore_roster_agent_version_switches_active_snapshot(monkeypatch: pytest.MonkeyPatch):
+    fake_session = FakeSession(scalar=["version-2", 6])
+    service = AgentRosterService(fake_session)
+    agent = Agent(
+        id="agent-1",
+        tenant_id="tenant-1",
+        name="Analyst",
+        description="old",
+        agent_kind=AgentKind.DIFY_AGENT,
+        scope=AgentScope.ROSTER,
+        source=AgentSource.AGENT_APP,
+        status=AgentStatus.ACTIVE,
+        active_config_snapshot_id="version-4",
+    )
+    version = AgentConfigSnapshot(
+        id="version-2",
+        tenant_id="tenant-1",
+        agent_id="agent-1",
+        version=2,
+        config_snapshot=_agent_soul_with_model(),
+    )
+
+    monkeypatch.setattr(service, "_get_agent", lambda **kwargs: agent)
+    monkeypatch.setattr(service, "_get_version", lambda **kwargs: version)
+
+    restored = service.restore_agent_version(
+        tenant_id="tenant-1",
+        agent_id="agent-1",
+        version_id="version-2",
+        account_id="account-1",
+    )
+
+    assert restored == {"result": "success", "active_config_snapshot_id": "version-2"}
+    assert agent.active_config_snapshot_id == "version-2"
+    assert agent.active_config_has_model is True
+    assert agent.updated_by == "account-1"
+    assert fake_session.commits == 1
+    revision = fake_session.added[0]
+    assert revision.tenant_id == "tenant-1"
+    assert revision.agent_id == "agent-1"
+    assert revision.previous_snapshot_id == "version-4"
+    assert revision.current_snapshot_id == "version-2"
+    assert revision.revision == 7
+    assert revision.operation == AgentConfigRevisionOperation.RESTORE_VERSION
+    assert revision.created_by == "account-1"
+
+
+def test_restore_roster_agent_version_rejects_invisible_versions(monkeypatch: pytest.MonkeyPatch):
+    fake_session = FakeSession(scalar=[None])
+    service = AgentRosterService(fake_session)
+    agent = Agent(
+        id="agent-1",
+        tenant_id="tenant-1",
+        name="Analyst",
+        description="old",
+        agent_kind=AgentKind.DIFY_AGENT,
+        scope=AgentScope.ROSTER,
+        source=AgentSource.AGENT_APP,
+        status=AgentStatus.ACTIVE,
+        active_config_snapshot_id="version-4",
+    )
+
+    monkeypatch.setattr(service, "_get_agent", lambda **kwargs: agent)
+
+    with pytest.raises(roster_service.AgentVersionNotFoundError):
+        service.restore_agent_version(
+            tenant_id="tenant-1",
+            agent_id="agent-1",
+            version_id="version-2",
+            account_id="account-1",
+        )
+
+    assert agent.active_config_snapshot_id == "version-4"
+    assert fake_session.added == []
+    assert fake_session.commits == 0
+
+
 def test_app_list_all_excludes_agent_apps_by_default():
     filters = AppService._build_app_list_filters(
         "account-1", "tenant-1", AppListParams(mode="all"), FakeSession(scalar=None, scalars=None)

packages/contracts/generated/api/console/agent/orpc.gen.ts

@@ -90,6 +90,8 @@ import {
   zPostAgentByAgentIdSkillsUploadBody,
   zPostAgentByAgentIdSkillsUploadPath,
   zPostAgentByAgentIdSkillsUploadResponse,
+  zPostAgentByAgentIdVersionsByVersionIdRestorePath,
+  zPostAgentByAgentIdVersionsByVersionIdRestoreResponse,
   zPostAgentResponse,
   zPutAgentByAgentIdBody,
   zPutAgentByAgentIdComposerBody,
@@ -738,6 +740,21 @@ export const statistics = {
   summary,
 }
 
+export const post10 = oc
+  .route({
+    inputStructure: 'detailed',
+    method: 'POST',
+    operationId: 'postAgentByAgentIdVersionsByVersionIdRestore',
+    path: '/agent/{agent_id}/versions/{version_id}/restore',
+    tags: ['console'],
+  })
+  .input(z.object({ params: zPostAgentByAgentIdVersionsByVersionIdRestorePath }))
+  .output(zPostAgentByAgentIdVersionsByVersionIdRestoreResponse)
+
+export const restore = {
+  post: post10,
+}
+
 export const get19 = oc
   .route({
     inputStructure: 'detailed',
@@ -751,6 +768,7 @@ export const get19 = oc
 
 export const byVersionId = {
   get: get19,
+  restore,
 }
 
 export const get20 = oc
@@ -835,7 +853,7 @@ export const get22 = oc
   .input(z.object({ query: zGetAgentQuery.optional() }))
   .output(zGetAgentResponse)
 
-export const post10 = oc
+export const post11 = oc
   .route({
     inputStructure: 'detailed',
     method: 'POST',
@@ -849,7 +867,7 @@ export const post10 = oc
 
 export const agent = {
   get: get22,
-  post: post10,
+  post: post11,
   inviteOptions,
   byAgentId,
 }

packages/contracts/generated/api/console/agent/types.gen.ts

@@ -315,6 +315,11 @@ export type AgentConfigSnapshotDetailResponse = {
   version_note?: string | null
 }
 
+export type AgentConfigSnapshotRestoreResponse = {
+  active_config_snapshot_id: string
+  result: 'success'
+}
+
 export type AgentAppPartial = {
   access_mode?: string | null
   active_config_is_published?: boolean
@@ -1176,6 +1181,7 @@ export type AgentUserSatisfactionRateStatisticResponse = {
 
 export type AgentConfigRevisionOperation
   = | 'create_version'
+    | 'restore_version'
     | 'save_current_version'
     | 'save_new_agent'
     | 'save_new_version'
@@ -2274,3 +2280,20 @@ export type GetAgentByAgentIdVersionsByVersionIdResponses = {
 
 export type GetAgentByAgentIdVersionsByVersionIdResponse
   = GetAgentByAgentIdVersionsByVersionIdResponses[keyof GetAgentByAgentIdVersionsByVersionIdResponses]
+
+export type PostAgentByAgentIdVersionsByVersionIdRestoreData = {
+  body?: never
+  path: {
+    agent_id: string
+    version_id: string
+  }
+  query?: never
+  url: '/agent/{agent_id}/versions/{version_id}/restore'
+}
+
+export type PostAgentByAgentIdVersionsByVersionIdRestoreResponses = {
+  200: AgentConfigSnapshotRestoreResponse
+}
+
+export type PostAgentByAgentIdVersionsByVersionIdRestoreResponse
+  = PostAgentByAgentIdVersionsByVersionIdRestoreResponses[keyof PostAgentByAgentIdVersionsByVersionIdRestoreResponses]

packages/contracts/generated/api/console/agent/zod.gen.ts

@@ -78,6 +78,14 @@ export const zAgentSandboxUploadPayload = z.object({
   path: z.string().min(1),
 })
 
+/**
+ * AgentConfigSnapshotRestoreResponse
+ */
+export const zAgentConfigSnapshotRestoreResponse = z.object({
+  active_config_snapshot_id: z.string(),
+  result: z.literal('success'),
+})
+
 /**
  * IconType
  */
@@ -1183,6 +1191,7 @@ export const zAgentStatisticSummaryEnvelopeResponse = z.object({
  */
 export const zAgentConfigRevisionOperation = z.enum([
   'create_version',
+  'restore_version',
   'save_current_version',
   'save_new_agent',
   'save_new_version',
@@ -2616,3 +2625,14 @@ export const zGetAgentByAgentIdVersionsByVersionIdPath = z.object({
  * Agent version detail
  */
 export const zGetAgentByAgentIdVersionsByVersionIdResponse = zAgentConfigSnapshotDetailResponse
+
+export const zPostAgentByAgentIdVersionsByVersionIdRestorePath = z.object({
+  agent_id: z.uuid(),
+  version_id: z.uuid(),
+})
+
+/**
+ * Agent version restored
+ */
+export const zPostAgentByAgentIdVersionsByVersionIdRestoreResponse
+  = zAgentConfigSnapshotRestoreResponse