Merge pull request #159 from keboola/KAB-1127-mcp-add-tracking-mechanism-to-flows

Kab 1127 add flow tracking mechanism
We add metadata tag to the flow creation and update tools as it is done in components. We also test the metadata tag in integration tests.

Author: mariankrotil

.github/workflows/release.yml

@@ -66,4 +66,4 @@ jobs:
 #        with:
 #          helm-chart: "mcp-server"
 #          image-tag: production-${{ github.sha }}
-#          github-app-private-key: ${{ secrets.GITOPS_KBC_STACKS_TRIGGER_APP_PVK }}
+#          github-app-private-key: ${{ secrets.GITOPS_KBC_STACKS_TRIGGER_APP_PVK }}

integtests/test_flow_tools.py

@@ -2,7 +2,8 @@
 from fastmcp import Context
 
 from integtests.conftest import ConfigDef
-from keboola_mcp_server.client import KeboolaClient
+from keboola_mcp_server.client import ORCHESTRATOR_COMPONENT_ID, KeboolaClient
+from keboola_mcp_server.config import MetadataField
 from keboola_mcp_server.tools.flow import (
     FlowToolResponse,
     create_flow,
@@ -55,6 +56,7 @@ async def test_create_and_retrieve_flow(mcp_context: Context, configs: list[Conf
         tasks=tasks,
     )
     flow_id = created.flow_id
+    client = KeboolaClient.from_state(mcp_context.session.state)
     try:
         assert isinstance(created, FlowToolResponse)
         assert created.description == flow_description
@@ -68,8 +70,20 @@ async def test_create_and_retrieve_flow(mcp_context: Context, configs: list[Conf
         assert detail.phases[0].name == 'Extract'
         assert detail.phases[1].name == 'Transform'
         assert detail.tasks[0].task['componentId'] == configs[0].component_id
+
+        # Verify the metadata - check that KBC.MCP.createdBy is set to 'true'
+        metadata = await client.storage_client.configuration_metadata_get(
+            component_id=ORCHESTRATOR_COMPONENT_ID,
+            configuration_id=flow_id
+        )
+
+        # Convert metadata list to dictionary for easier checking
+        # metadata is a list of dicts with 'key' and 'value' keys
+        assert isinstance(metadata, list)
+        metadata_dict = {item['key']: item['value'] for item in metadata if isinstance(item, dict)}
+        assert MetadataField.CREATED_BY_MCP in metadata_dict
+        assert metadata_dict[MetadataField.CREATED_BY_MCP] == 'true'
     finally:
-        client = KeboolaClient.from_state(mcp_context.session.state)
         await client.storage_client.flow_delete(flow_id, skip_trash=True)
 
 
@@ -106,6 +120,7 @@ async def test_update_flow(mcp_context: Context, configs: list[ConfigDef]) -> No
         tasks=tasks,
     )
     flow_id = created.flow_id
+    client = KeboolaClient.from_state(mcp_context.session.state)
     try:
         new_name = 'Updated Flow Name'
         new_description = 'Updated description.'
@@ -123,8 +138,21 @@ async def test_update_flow(mcp_context: Context, configs: list[ConfigDef]) -> No
         assert updated.description == new_description
         assert updated.success is True
         assert len(updated.links) == 3
+
+        # Verify the metadata - check that KBC.MCP.updatedBy.version.{version} is set to 'true'
+        metadata = await client.storage_client.configuration_metadata_get(
+            component_id=ORCHESTRATOR_COMPONENT_ID,
+            configuration_id=flow_id
+        )
+
+        assert isinstance(metadata, list)
+        metadata_dict = {item['key']: item['value'] for item in metadata if isinstance(item, dict)}
+        sync_flow = await client.storage_client.flow_detail(flow_id)
+        updated_by_md_key = f'{MetadataField.UPDATED_BY_MCP_PREFIX}{sync_flow["version"]}'
+        assert updated_by_md_key in metadata_dict
+        assert metadata_dict[updated_by_md_key] == 'true'
+
     finally:
-        client = KeboolaClient.from_state(mcp_context.session.state)
         await client.storage_client.flow_delete(flow_id, skip_trash=True)
 
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "keboola-mcp-server"
-version = "1.0.0"
+version = "1.0.1"
 description = "MCP server for interacting with Keboola Connection"
 readme = "README.md"
 requires-python = ">=3.10"

src/keboola_mcp_server/tools/flow.py

@@ -9,7 +9,7 @@
 from fastmcp import Context, FastMCP
 from pydantic import AliasChoices, BaseModel, Field
 
-from keboola_mcp_server.client import JsonDict, KeboolaClient
+from keboola_mcp_server.client import ORCHESTRATOR_COMPONENT_ID, JsonDict, KeboolaClient
 from keboola_mcp_server.errors import tool_errors
 from keboola_mcp_server.links import Link, ProjectLinksManager
 from keboola_mcp_server.mcp import with_session_state
@@ -20,6 +20,7 @@
     FlowTask,
     ReducedFlow,
 )
+from keboola_mcp_server.tools.components.tools import _set_cfg_creation_metadata, _set_cfg_update_metadata
 from keboola_mcp_server.tools.validation import validate_flow_configuration_against_schema
 
 LOG = logging.getLogger(__name__)
@@ -127,8 +128,14 @@ async def create_flow(
         name=name, description=description, flow_configuration=flow_configuration  # Direct configuration
     )
 
+    await _set_cfg_creation_metadata(
+        client,
+        component_id=ORCHESTRATOR_COMPONENT_ID,
+        configuration_id=str(new_raw_configuration['id']),
+    )
+
     flow_id = str(new_raw_configuration['id'])
-    flow_name = new_raw_configuration['name']
+    flow_name = str(new_raw_configuration['name'])
     flow_links = links_manager.get_flow_links(flow_id=flow_id, flow_name=flow_name)
     tool_response = FlowToolResponse.model_validate(new_raw_configuration | {'links': flow_links})
 
@@ -191,8 +198,15 @@ async def update_flow(
         flow_configuration=flow_configuration,  # Direct configuration
     )
 
+    await _set_cfg_update_metadata(
+        client,
+        component_id=ORCHESTRATOR_COMPONENT_ID,
+        configuration_id=str(updated_raw_configuration['id']),
+        configuration_version=cast(int, updated_raw_configuration['version']),
+    )
+
     flow_id = str(updated_raw_configuration['id'])
-    flow_name = updated_raw_configuration['name']
+    flow_name = str(updated_raw_configuration['name'])
     flow_links = links_manager.get_flow_links(flow_id=flow_id, flow_name=flow_name)
     tool_response = FlowToolResponse.model_validate(updated_raw_configuration | {'links': flow_links})