diff --git a/cpg-mcp/src/integrationTest/kotlin/mcp/tools/CpgDfgBackwardToolTest.kt b/cpg-mcp/src/integrationTest/kotlin/mcp/tools/CpgDfgBackwardToolTest.kt new file mode 100644 index 00000000000..c6934c65010 --- /dev/null +++ b/cpg-mcp/src/integrationTest/kotlin/mcp/tools/CpgDfgBackwardToolTest.kt @@ -0,0 +1,83 @@ +/* + * Copyright (c) 2026, Fraunhofer AISEC. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * $$$$$$\ $$$$$$$\ $$$$$$\ + * $$ __$$\ $$ __$$\ $$ __$$\ + * $$ / \__|$$ | $$ |$$ / \__| + * $$ | $$$$$$$ |$$ |$$$$\ + * $$ | $$ ____/ $$ |\_$$ | + * $$ | $$\ $$ | $$ | $$ | + * \$$$$$ |$$ | \$$$$$ | + * \______/ \__| \______/ + * + */ +package de.fraunhofer.aisec.cpg.mcp.tools + +import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.addDfgBackwardTool +import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.listCalls +import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.runCpgAnalyze +import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.utils.CallInfo +import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.utils.CpgAnalyzePayload +import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.utils.NodeInfo +import de.fraunhofer.aisec.cpg.mcp.utils.withClient +import io.modelcontextprotocol.kotlin.sdk.types.TextContent +import kotlin.test.Test +import kotlin.test.assertIs +import kotlin.test.assertNotNull +import kotlin.test.assertTrue +import kotlinx.serialization.json.Json +import org.junit.jupiter.api.BeforeEach + +class CpgDfgBackwardToolTest { + @BeforeEach + fun setAnalysisResult() { + val payload = + CpgAnalyzePayload( + content = "def hello():\n foo = bar\n print(foo)", + extension = "py", + ) + runCpgAnalyze(payload, runPasses = true, cleanup = true) + } + + @Test + fun dfgBackwardToolTest() = + withClient( + registerTools = { + listCalls() + addDfgBackwardTool() + } + ) { client -> + val callsResult = client.callTool(name = "cpg_list_calls", arguments = emptyMap()) + assertNotNull(callsResult) + assertTrue(callsResult.content.isNotEmpty()) + + val callInfo = + Json.decodeFromString((callsResult.content.first() as TextContent).text) + + val result = + client.callTool( + name = "cpg_dfg_backward", + arguments = mapOf("id" to callInfo.nodeId), + ) + assertNotNull(result) + assertTrue(result.content.isNotEmpty()) + + val content = result.content.single() + assertIs(content) + + val nodes = Json.decodeFromString>(content.text) + assertTrue(nodes.isNotEmpty()) + } +} diff --git a/cpg-mcp/src/main/kotlin/de/fraunhofer/aisec/cpg/mcp/mcpserver/McpServer.kt b/cpg-mcp/src/main/kotlin/de/fraunhofer/aisec/cpg/mcp/mcpserver/McpServer.kt index bcc38471db6..4959a27f88d 100644 --- a/cpg-mcp/src/main/kotlin/de/fraunhofer/aisec/cpg/mcp/mcpserver/McpServer.kt +++ b/cpg-mcp/src/main/kotlin/de/fraunhofer/aisec/cpg/mcp/mcpserver/McpServer.kt @@ -29,6 +29,7 @@ import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.addCpgAnalyzeTool import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.addCpgApplyConceptsTool import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.addCpgDataflowTool import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.addCpgTranslate +import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.addDfgBackwardTool import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.addListPasses import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.addRunPass import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.addSuggestConceptsPrompt @@ -64,6 +65,7 @@ fun configureServer( this.getAllArgs() this.getArgByIndexOrName() this.listConceptsAndOperations() + this.addDfgBackwardTool() // PROMPTS this.addSuggestConceptsPrompt() this diff --git a/cpg-mcp/src/main/kotlin/de/fraunhofer/aisec/cpg/mcp/mcpserver/tools/CpgDfgBackwardTool.kt b/cpg-mcp/src/main/kotlin/de/fraunhofer/aisec/cpg/mcp/mcpserver/tools/CpgDfgBackwardTool.kt new file mode 100644 index 00000000000..402d1bb8dc1 --- /dev/null +++ b/cpg-mcp/src/main/kotlin/de/fraunhofer/aisec/cpg/mcp/mcpserver/tools/CpgDfgBackwardTool.kt @@ -0,0 +1,68 @@ +/* + * Copyright (c) 2025, Fraunhofer AISEC. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * $$$$$$\ $$$$$$$\ $$$$$$\ + * $$ __$$\ $$ __$$\ $$ __$$\ + * $$ / \__|$$ | $$ |$$ / \__| + * $$ | $$$$$$$ |$$ |$$$$\ + * $$ | $$ ____/ $$ |\_$$ | + * $$ | $$\ $$ | $$ | $$ | + * \$$$$$ |$$ | \$$$$$ | + * \______/ \__| \______/ + * + */ +package de.fraunhofer.aisec.cpg.mcp.mcpserver.tools + +import de.fraunhofer.aisec.cpg.TranslationResult +import de.fraunhofer.aisec.cpg.graph.collectAllPrevDFGPaths +import de.fraunhofer.aisec.cpg.graph.nodes +import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.utils.CpgIdPayload +import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.utils.NodeInfo +import de.fraunhofer.aisec.cpg.mcp.mcpserver.tools.utils.addTool +import io.modelcontextprotocol.kotlin.sdk.server.Server +import io.modelcontextprotocol.kotlin.sdk.types.CallToolResult +import io.modelcontextprotocol.kotlin.sdk.types.TextContent +import kotlin.uuid.Uuid +import kotlinx.serialization.json.Json + +fun Server.addDfgBackwardTool() { + val toolDescription = + """ + Traverse the Data Flow Graph (DFG) backwards from a given node to find where data originates. + + Uses the CPG's built-in dataflow analysis with backward direction to trace data sources. + The analysis follows prevDFG edges and stops at nodes that have no further incoming data flows. + + Example usage: + - "Where does this value come from?" + """ + .trimIndent() + + this.addTool(name = "cpg_dfg_backward", description = toolDescription) { + result: TranslationResult, + payload: CpgIdPayload -> + val startId = Uuid.parse(payload.id) + val startNode = + result.nodes.find { it.id == startId } + ?: return@addTool CallToolResult( + content = listOf(TextContent("No node found with ID ${payload.id}")) + ) + + val paths = startNode.collectAllPrevDFGPaths() + val nodes = paths.flatMap { it.nodes }.map { NodeInfo(it) } + + CallToolResult(content = listOf(TextContent(Json.encodeToString(nodes)))) + } +} diff --git a/cpg-mcp/src/test/kotlin/de/fraunhofer/aisec/cpg/mcp/CpgAnalyzeToolTest.kt b/cpg-mcp/src/test/kotlin/de/fraunhofer/aisec/cpg/mcp/CpgAnalyzeToolTest.kt index 51bc4f828a3..50c518a8378 100644 --- a/cpg-mcp/src/test/kotlin/de/fraunhofer/aisec/cpg/mcp/CpgAnalyzeToolTest.kt +++ b/cpg-mcp/src/test/kotlin/de/fraunhofer/aisec/cpg/mcp/CpgAnalyzeToolTest.kt @@ -61,6 +61,7 @@ class CpgAnalyzeToolTest { "cpg_list_call_arg_by_name_or_index", "cpg_list_available_concepts", "cpg_list_available_operations", + "cpg_dfg_backward", ), testServer.tools.keys, )