diff --git a/api/controllers/console/agent/composer.py b/api/controllers/console/agent/composer.py index 4c9cb2d0db3..5ae35e2138a 100644 --- a/api/controllers/console/agent/composer.py +++ b/api/controllers/console/agent/composer.py @@ -1,8 +1,9 @@ from uuid import UUID +from flask import request from flask_restx import Resource -from controllers.common.schema import register_response_schema_models, register_schema_models +from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models from controllers.console import console_ns from controllers.console.agent.app_helpers import resolve_agent_app_model from controllers.console.app.wraps import get_app_model @@ -28,9 +29,15 @@ from libs.login import login_required from models.model import App, AppMode from services.agent.composer_service import AgentComposerService from services.agent.composer_validator import ComposerConfigValidator -from services.entities.agent_entities import ComposerSavePayload, WorkflowComposerCopyFromRosterPayload +from services.entities.agent_entities import ( + ComposerSavePayload, + WorkflowAgentComposerQuery, + WorkflowComposerCopyFromRosterPayload, +) -register_schema_models(console_ns, ComposerSavePayload, WorkflowComposerCopyFromRosterPayload) +register_schema_models( + console_ns, ComposerSavePayload, WorkflowAgentComposerQuery, WorkflowComposerCopyFromRosterPayload +) register_response_schema_models( console_ns, AgentAppComposerResponse, @@ -50,18 +57,21 @@ class WorkflowAgentComposerApi(Resource): @console_ns.response( 200, "Workflow agent composer state", console_ns.models[WorkflowAgentComposerResponse.__name__] ) + @console_ns.doc(params=query_params_from_model(WorkflowAgentComposerQuery)) @setup_required @login_required @account_initialization_required @get_app_model(mode=[AppMode.WORKFLOW, AppMode.ADVANCED_CHAT]) @with_current_tenant_id def get(self, tenant_id: str, app_model: App, node_id: str): + query = WorkflowAgentComposerQuery.model_validate(request.args.to_dict(flat=True)) return dump_response( WorkflowAgentComposerResponse, AgentComposerService.load_workflow_composer( tenant_id=tenant_id, app_id=app_model.id, node_id=node_id, + snapshot_id=query.snapshot_id, ), ) diff --git a/api/services/agent/composer_service.py b/api/services/agent/composer_service.py index a6fbfc73ef5..d2869c511ce 100644 --- a/api/services/agent/composer_service.py +++ b/api/services/agent/composer_service.py @@ -98,24 +98,62 @@ def _validate_composer_payload_for_strategy(payload: ComposerSavePayload) -> Non class AgentComposerService: @classmethod - def load_workflow_composer(cls, *, tenant_id: str, app_id: str, node_id: str) -> dict[str, Any]: + def load_workflow_composer( + cls, *, tenant_id: str, app_id: str, node_id: str, snapshot_id: str | None = None + ) -> dict[str, Any]: workflow = cls._get_draft_workflow(tenant_id=tenant_id, app_id=app_id) binding = cls._get_workflow_binding(tenant_id=tenant_id, workflow_id=workflow.id, node_id=node_id) if not binding: + if snapshot_id: + raise AgentVersionNotFoundError() return cls._empty_workflow_state(app_id=app_id, workflow_id=workflow.id, node_id=node_id) agent = cls._get_agent_if_present(tenant_id=tenant_id, agent_id=binding.agent_id) + version = cls._workflow_composer_version( + tenant_id=tenant_id, + binding=binding, + agent=agent, + snapshot_id=snapshot_id, + ) + return cls._serialize_workflow_state(binding=binding, agent=agent, version=version) + + @classmethod + def _workflow_composer_version( + cls, + *, + tenant_id: str, + binding: WorkflowAgentNodeBinding, + agent: Agent | None, + snapshot_id: str | None, + ) -> AgentConfigSnapshot | None: + if snapshot_id: + if agent is None: + raise AgentVersionNotFoundError() + if binding.binding_type == WorkflowAgentBindingType.ROSTER_AGENT: + if agent.scope != AgentScope.ROSTER: + raise AgentVersionNotFoundError() + elif binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT: + if ( + agent.scope != AgentScope.WORKFLOW_ONLY + or agent.app_id != binding.app_id + or agent.workflow_id != binding.workflow_id + or agent.workflow_node_id != binding.node_id + ): + raise AgentVersionNotFoundError() + else: + raise AgentVersionNotFoundError() + return cls._require_version(tenant_id=tenant_id, agent_id=agent.id, version_id=snapshot_id) + version_id = ( agent.active_config_snapshot_id if agent and binding.binding_type == WorkflowAgentBindingType.ROSTER_AGENT else binding.current_snapshot_id ) - version = cls._get_version_if_present( + return cls._get_version_if_present( tenant_id=tenant_id, agent_id=agent.id if agent else None, version_id=version_id, ) - return cls._serialize_workflow_state(binding=binding, agent=agent, version=version) @classmethod def save_workflow_composer( diff --git a/api/services/entities/agent_entities.py b/api/services/entities/agent_entities.py index a8634bceb09..3184455d0a4 100644 --- a/api/services/entities/agent_entities.py +++ b/api/services/entities/agent_entities.py @@ -31,6 +31,10 @@ class ComposerSoulLockPayload(BaseModel): unlocked_from_version_id: str | None = None +class WorkflowAgentComposerQuery(BaseModel): + snapshot_id: str | None = Field(default=None, max_length=255) + + class ComposerSavePayload(BaseModel): variant: ComposerVariant binding: ComposerBindingPayload | None = None diff --git a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py index 93372bd42c2..5b854eb7410 100644 --- a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py +++ b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py @@ -1094,10 +1094,11 @@ def test_workflow_composer_get_put_validate_candidates_impact_and_save( "save_strategy": ComposerSaveStrategy.NODE_JOB_ONLY.value, "binding": {"binding_type": "roster_agent", "current_snapshot_id": "version-1"}, } + captured_load: dict[str, object] = {} monkeypatch.setattr( composer_controller.AgentComposerService, "load_workflow_composer", - lambda **kwargs: _workflow_composer_response(node_id=kwargs["node_id"]), + lambda **kwargs: captured_load.update(kwargs) or _workflow_composer_response(node_id=kwargs["node_id"]), ) monkeypatch.setattr( composer_controller.AgentComposerService, @@ -1124,8 +1125,12 @@ def test_workflow_composer_get_put_validate_candidates_impact_and_save( }, ) - workflow_state = unwrap(WorkflowAgentComposerApi.get)(WorkflowAgentComposerApi(), "tenant-1", app_model, "node-1") + with app.test_request_context("?snapshot_id=preview-version"): + workflow_state = unwrap(WorkflowAgentComposerApi.get)( + WorkflowAgentComposerApi(), "tenant-1", app_model, "node-1" + ) assert workflow_state["node_id"] == "node-1" + assert captured_load["snapshot_id"] == "preview-version" with app.test_request_context(json=payload): saved_state = unwrap(WorkflowAgentComposerApi.put)( WorkflowAgentComposerApi(), "tenant-1", account_id, app_model, "node-1" diff --git a/api/tests/unit_tests/services/agent/test_agent_services.py b/api/tests/unit_tests/services/agent/test_agent_services.py index ebf9eb35025..f26b74c5e91 100644 --- a/api/tests/unit_tests/services/agent/test_agent_services.py +++ b/api/tests/unit_tests/services/agent/test_agent_services.py @@ -39,6 +39,7 @@ from services.agent.errors import ( AgentNameConflictError, AgentNotFoundError, AgentVersionConflictError, + AgentVersionNotFoundError, InvalidComposerConfigError, ) from services.agent.roster_service import AgentRosterService @@ -158,6 +159,96 @@ def test_load_workflow_composer_serializes_existing_binding(monkeypatch: pytest. assert result == {"agent": "agent-1", "version": "version-1"} +def test_load_workflow_composer_uses_roster_preview_snapshot(monkeypatch: pytest.MonkeyPatch): + binding = SimpleNamespace( + agent_id="agent-1", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + current_snapshot_id="binding-version", + ) + agent = SimpleNamespace(id="agent-1", scope=AgentScope.ROSTER, active_config_snapshot_id="active-version") + + monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", lambda **kwargs: SimpleNamespace(id="workflow-1")) + monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", lambda **kwargs: binding) + monkeypatch.setattr(AgentComposerService, "_get_agent_if_present", lambda **kwargs: agent) + monkeypatch.setattr( + AgentComposerService, + "_require_version", + lambda **kwargs: SimpleNamespace(id=kwargs["version_id"]), + ) + monkeypatch.setattr( + AgentComposerService, + "_serialize_workflow_state", + lambda **kwargs: { + "binding_snapshot_id": kwargs["binding"].current_snapshot_id, + "version": kwargs["version"].id, + }, + ) + + result = AgentComposerService.load_workflow_composer( + tenant_id="tenant-1", + app_id="app-1", + node_id="node-1", + snapshot_id="preview-version", + ) + + assert result == {"binding_snapshot_id": "binding-version", "version": "preview-version"} + + +def test_load_workflow_composer_uses_inline_preview_snapshot(monkeypatch: pytest.MonkeyPatch): + binding = SimpleNamespace( + agent_id="inline-agent-1", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + current_snapshot_id="inline-version-1", + app_id="app-1", + workflow_id="workflow-1", + node_id="node-1", + ) + agent = SimpleNamespace( + id="inline-agent-1", + scope=AgentScope.WORKFLOW_ONLY, + app_id="app-1", + workflow_id="workflow-1", + workflow_node_id="node-1", + active_config_snapshot_id="inline-version-1", + ) + + monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", lambda **kwargs: SimpleNamespace(id="workflow-1")) + monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", lambda **kwargs: binding) + monkeypatch.setattr(AgentComposerService, "_get_agent_if_present", lambda **kwargs: agent) + monkeypatch.setattr( + AgentComposerService, + "_require_version", + lambda **kwargs: SimpleNamespace(id=kwargs["version_id"]), + ) + monkeypatch.setattr( + AgentComposerService, + "_serialize_workflow_state", + lambda **kwargs: {"agent": kwargs["agent"].id, "version": kwargs["version"].id}, + ) + + result = AgentComposerService.load_workflow_composer( + tenant_id="tenant-1", + app_id="app-1", + node_id="node-1", + snapshot_id="inline-preview-version", + ) + + assert result == {"agent": "inline-agent-1", "version": "inline-preview-version"} + + +def test_load_workflow_composer_rejects_preview_without_binding(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(AgentComposerService, "_get_draft_workflow", lambda **kwargs: SimpleNamespace(id="workflow-1")) + monkeypatch.setattr(AgentComposerService, "_get_workflow_binding", lambda **kwargs: None) + + with pytest.raises(AgentVersionNotFoundError): + AgentComposerService.load_workflow_composer( + tenant_id="tenant-1", + app_id="app-1", + node_id="node-1", + snapshot_id="preview-version", + ) + + @pytest.mark.parametrize( ("strategy", "helper_name"), [ diff --git a/packages/contracts/generated/api/console/apps/orpc.gen.ts b/packages/contracts/generated/api/console/apps/orpc.gen.ts index ea72df28458..59783646462 100644 --- a/packages/contracts/generated/api/console/apps/orpc.gen.ts +++ b/packages/contracts/generated/api/console/apps/orpc.gen.ts @@ -185,6 +185,7 @@ import { zGetAppsByAppIdWorkflowsDraftNodesByNodeIdAgentComposerCandidatesPath, zGetAppsByAppIdWorkflowsDraftNodesByNodeIdAgentComposerCandidatesResponse, zGetAppsByAppIdWorkflowsDraftNodesByNodeIdAgentComposerPath, + zGetAppsByAppIdWorkflowsDraftNodesByNodeIdAgentComposerQuery, zGetAppsByAppIdWorkflowsDraftNodesByNodeIdAgentComposerResponse, zGetAppsByAppIdWorkflowsDraftNodesByNodeIdLastRunPath, zGetAppsByAppIdWorkflowsDraftNodesByNodeIdLastRunResponse, @@ -3569,7 +3570,12 @@ export const get62 = oc path: '/apps/{app_id}/workflows/draft/nodes/{node_id}/agent-composer', tags: ['console'], }) - .input(z.object({ params: zGetAppsByAppIdWorkflowsDraftNodesByNodeIdAgentComposerPath })) + .input( + z.object({ + params: zGetAppsByAppIdWorkflowsDraftNodesByNodeIdAgentComposerPath, + query: zGetAppsByAppIdWorkflowsDraftNodesByNodeIdAgentComposerQuery.optional(), + }), + ) .output(zGetAppsByAppIdWorkflowsDraftNodesByNodeIdAgentComposerResponse) export const put4 = oc diff --git a/packages/contracts/generated/api/console/apps/types.gen.ts b/packages/contracts/generated/api/console/apps/types.gen.ts index 146e5ee29c9..0a7db919b16 100644 --- a/packages/contracts/generated/api/console/apps/types.gen.ts +++ b/packages/contracts/generated/api/console/apps/types.gen.ts @@ -5497,7 +5497,9 @@ export type GetAppsByAppIdWorkflowsDraftNodesByNodeIdAgentComposerData = { app_id: string node_id: string } - query?: never + query?: { + snapshot_id?: string + } url: '/apps/{app_id}/workflows/draft/nodes/{node_id}/agent-composer' } diff --git a/packages/contracts/generated/api/console/apps/zod.gen.ts b/packages/contracts/generated/api/console/apps/zod.gen.ts index e67a718481b..c57f5e799b3 100644 --- a/packages/contracts/generated/api/console/apps/zod.gen.ts +++ b/packages/contracts/generated/api/console/apps/zod.gen.ts @@ -5507,6 +5507,10 @@ export const zGetAppsByAppIdWorkflowsDraftNodesByNodeIdAgentComposerPath = z.obj node_id: z.string(), }) +export const zGetAppsByAppIdWorkflowsDraftNodesByNodeIdAgentComposerQuery = z.object({ + snapshot_id: z.string().max(255).optional(), +}) + /** * Workflow agent composer state */