Skip to content

Commit f78fee2

Browse files
kpavlovkpavlov
committed
Refactor imports and update schema handling for MCP SDK 0.8.0
1 parent 327db6a commit f78fee2

File tree

10 files changed

+115
-68
lines changed

10 files changed

+115
-68
lines changed

‎agents/agents-mcp-server/src/commonMain/kotlin/ai/koog/agents/mcp/server/McpServer.kt‎

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,15 @@ import io.ktor.server.engine.ApplicationEngineFactory
1010
import io.ktor.server.engine.EmbeddedServer
1111
import io.ktor.server.engine.EngineConnectorConfig
1212
import io.ktor.server.engine.embeddedServer
13-
import io.modelcontextprotocol.kotlin.sdk.CallToolResult
14-
import io.modelcontextprotocol.kotlin.sdk.Implementation
15-
import io.modelcontextprotocol.kotlin.sdk.ServerCapabilities
16-
import io.modelcontextprotocol.kotlin.sdk.TextContent
1713
import io.modelcontextprotocol.kotlin.sdk.server.Server
1814
import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions
1915
import io.modelcontextprotocol.kotlin.sdk.server.mcp
16+
import io.modelcontextprotocol.kotlin.sdk.types.CallToolResult
17+
import io.modelcontextprotocol.kotlin.sdk.types.EmptyJsonObject
18+
import io.modelcontextprotocol.kotlin.sdk.types.Implementation
19+
import io.modelcontextprotocol.kotlin.sdk.types.ServerCapabilities
20+
import io.modelcontextprotocol.kotlin.sdk.types.TextContent
21+
import io.modelcontextprotocol.kotlin.sdk.types.ToolSchema
2022
import kotlinx.coroutines.coroutineScope
2123
import kotlinx.coroutines.isActive
2224
import kotlinx.serialization.json.JsonObject
@@ -26,10 +28,11 @@ import kotlinx.serialization.json.buildJsonObject
2628
import kotlinx.serialization.json.put
2729
import kotlinx.serialization.json.putJsonArray
2830
import kotlinx.serialization.json.putJsonObject
29-
import io.modelcontextprotocol.kotlin.sdk.Tool as SdkTool
31+
import io.modelcontextprotocol.kotlin.sdk.types.Tool as SdkTool
3032

3133
/**
32-
* Starts a new MCP server with the passed [tools] that listens to and writes to the specified [port] on the passed [host].
34+
* Starts a new MCP server with the passed [tools] that listens to and writes
35+
* to the specified [port] on the passed [host].
3336
* A port can be obtained from the returned list of [EngineConnectorConfig].
3437
*/
3538
public suspend fun startSseMcpServer(
@@ -40,7 +43,8 @@ public suspend fun startSseMcpServer(
4043
): Server = doStartSseMcpServer(factory, port, host, tools, true).first
4144

4245
/**
43-
* Starts a new MCP server with the passed [tools] that listens to and writes to the allocated port on the passed [host].
46+
* Starts a new MCP server with the passed [tools] that listens to and writes
47+
* to the allocated port on the passed [host].
4448
* A port can be obtained from the returned list of [EngineConnectorConfig].
4549
*/
4650
public suspend fun startSseMcpServer(
@@ -119,7 +123,7 @@ public fun Server.addTool(
119123
tool: Tool<*, *>,
120124
) {
121125
addTool(tool.descriptor.asSdkTool()) { request ->
122-
val args = tool.decodeArgs(request.arguments)
126+
val args = tool.decodeArgs(request.arguments ?: EmptyJsonObject)
123127
val result = tool.executeUnsafe(args)
124128

125129
CallToolResult(
@@ -132,7 +136,7 @@ private fun ToolDescriptor.asSdkTool(): SdkTool {
132136
return SdkTool(
133137
name = name,
134138
description = description,
135-
inputSchema = SdkTool.Input(
139+
inputSchema = ToolSchema(
136140
properties = buildJsonObject {
137141
(requiredParameters + optionalParameters).forEach { param ->
138142
put(param.name, param.toJsonSchema())
@@ -155,10 +159,15 @@ private fun ToolParameterDescriptor.toJsonSchema(): JsonObject = buildJsonObject
155159
private fun JsonObjectBuilder.fillJsonSchema(type: ToolParameterType) {
156160
when (type) {
157161
ToolParameterType.Boolean -> put("type", "boolean")
162+
158163
ToolParameterType.Float -> put("type", "number")
164+
159165
ToolParameterType.Integer -> put("type", "integer")
166+
160167
ToolParameterType.String -> put("type", "string")
168+
161169
ToolParameterType.Null -> put("type", "null")
170+
162171
is ToolParameterType.Enum -> {
163172
put("type", "string")
164173
putJsonArray("enum") {

‎agents/agents-mcp-server/src/jvmTest/kotlin/ai/koog/agents/mcp/server/KoogToolAsMcpToolTest.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@ import ai.koog.agents.testing.network.NetUtil.isPortAvailable
99
import ai.koog.agents.testing.tools.RandomNumberTool
1010
import io.github.oshai.kotlinlogging.KotlinLogging
1111
import io.ktor.server.cio.CIO
12-
import io.modelcontextprotocol.kotlin.sdk.EmptyJsonObject
13-
import io.modelcontextprotocol.kotlin.sdk.TextContent
12+
import io.modelcontextprotocol.kotlin.sdk.types.EmptyJsonObject
13+
import io.modelcontextprotocol.kotlin.sdk.types.TextContent
1414
import kotlinx.coroutines.Dispatchers
1515
import kotlinx.coroutines.delay
1616
import kotlinx.coroutines.test.runTest

‎agents/agents-mcp/src/commonMain/kotlin/ai/koog/agents/mcp/McpTool.kt‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,8 @@ package ai.koog.agents.mcp
22

33
import ai.koog.agents.core.tools.Tool
44
import ai.koog.agents.core.tools.ToolDescriptor
5-
import io.modelcontextprotocol.kotlin.sdk.CallToolResultBase
65
import io.modelcontextprotocol.kotlin.sdk.client.Client
6+
import io.modelcontextprotocol.kotlin.sdk.types.CallToolResult
77
import kotlinx.serialization.KSerializer
88
import kotlinx.serialization.builtins.nullable
99
import kotlinx.serialization.json.JsonElement
@@ -23,12 +23,12 @@ import kotlinx.serialization.json.jsonObject
2323
public class McpTool(
2424
private val mcpClient: Client,
2525
override val descriptor: ToolDescriptor,
26-
) : Tool<JsonObject, CallToolResultBase?>() {
26+
) : Tool<JsonObject, CallToolResult?>() {
2727
override val name: String = descriptor.name
2828
override val description: String = descriptor.description
2929

3030
override val argsSerializer: KSerializer<JsonObject> = JsonObject.serializer()
31-
override val resultSerializer: KSerializer<CallToolResultBase?> = CallToolResultBase.serializer().nullable
31+
override val resultSerializer: KSerializer<CallToolResult?> = CallToolResult.serializer().nullable
3232

3333
/**
3434
* Executes the MCP tool with the given arguments.
@@ -39,7 +39,7 @@ public class McpTool(
3939
* @param args The arguments for the MCP tool call.
4040
* @return The result of the MCP tool call.
4141
*/
42-
override suspend fun execute(args: JsonObject): CallToolResultBase? {
42+
override suspend fun execute(args: JsonObject): CallToolResult {
4343
return mcpClient.callTool(
4444
name = descriptor.name,
4545
arguments = args
@@ -49,7 +49,7 @@ public class McpTool(
4949
/**
5050
* Postprocess result string representation for LLMs a bit, removing unnecessary meta fields.
5151
*/
52-
override fun encodeResultToString(result: CallToolResultBase?): String {
52+
override fun encodeResultToString(result: CallToolResult?): String {
5353
val preparedResultJson: JsonElement = result
5454
?.let {
5555
JsonObject(

‎agents/agents-mcp/src/commonMain/kotlin/ai/koog/agents/mcp/McpToolDefinitionParser.kt‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,14 @@ package ai.koog.agents.mcp
33
import ai.koog.agents.core.tools.ToolDescriptor
44
import ai.koog.agents.core.tools.ToolParameterDescriptor
55
import ai.koog.agents.core.tools.ToolParameterType
6+
import io.modelcontextprotocol.kotlin.sdk.types.EmptyJsonObject
67
import kotlinx.serialization.json.JsonObject
78
import kotlinx.serialization.json.JsonPrimitive
89
import kotlinx.serialization.json.boolean
910
import kotlinx.serialization.json.jsonArray
1011
import kotlinx.serialization.json.jsonObject
1112
import kotlinx.serialization.json.jsonPrimitive
12-
import io.modelcontextprotocol.kotlin.sdk.Tool as SDKTool
13+
import io.modelcontextprotocol.kotlin.sdk.types.Tool as SDKTool
1314

1415
/**
1516
* Parsers tool definition from MCP SDK to our tool descriptor format.
@@ -42,7 +43,7 @@ public object DefaultMcpToolDescriptorParser : McpToolDescriptorParser {
4243
*/
4344
override fun parse(sdkTool: SDKTool): ToolDescriptor {
4445
// Parse all parameters from the input schema
45-
val parameters = parseParameters(sdkTool.inputSchema.properties)
46+
val parameters = parseParameters(sdkTool.inputSchema.properties ?: EmptyJsonObject)
4647

4748
// Get the list of required parameters
4849
val requiredParameters = sdkTool.inputSchema.required ?: emptyList()
@@ -142,9 +143,13 @@ public object DefaultMcpToolDescriptorParser : McpToolDescriptorParser {
142143
return when (typeStr.lowercase()) {
143144
// Primitive types
144145
"string" -> ToolParameterType.String
146+
145147
"integer" -> ToolParameterType.Integer
148+
146149
"number" -> ToolParameterType.Float
150+
147151
"boolean" -> ToolParameterType.Boolean
152+
148153
"enum" -> ToolParameterType.Enum(
149154
element.getValue("enum").jsonArray.map { it.jsonPrimitive.content }.toTypedArray()
150155
)
@@ -188,6 +193,7 @@ public object DefaultMcpToolDescriptorParser : McpToolDescriptorParser {
188193
element.getValue("additionalProperties").jsonObject,
189194
depth + 1
190195
)
196+
191197
else -> null
192198
}
193199
} else {

‎agents/agents-mcp/src/commonMain/kotlin/ai/koog/agents/mcp/McpToolRegistryProvider.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,10 @@ import ai.koog.agents.core.tools.ToolRegistry
44
import io.github.oshai.kotlinlogging.KotlinLogging
55
import io.ktor.client.HttpClient
66
import io.ktor.client.plugins.sse.SSE
7-
import io.modelcontextprotocol.kotlin.sdk.Implementation
87
import io.modelcontextprotocol.kotlin.sdk.client.Client
98
import io.modelcontextprotocol.kotlin.sdk.client.SseClientTransport
109
import io.modelcontextprotocol.kotlin.sdk.shared.Transport
10+
import io.modelcontextprotocol.kotlin.sdk.types.Implementation
1111

1212
/**
1313
* A provider for creating tool registries that connect to Model Context Protocol (MCP) servers.
@@ -60,7 +60,7 @@ public object McpToolRegistryProvider {
6060
mcpClient: Client,
6161
mcpToolParser: McpToolDescriptorParser = DefaultMcpToolDescriptorParser,
6262
): ToolRegistry {
63-
val sdkTools = mcpClient.listTools()?.tools.orEmpty()
63+
val sdkTools = mcpClient.listTools().tools
6464
return ToolRegistry {
6565
sdkTools.forEach { sdkTool ->
6666
try {

‎agents/agents-mcp/src/commonTest/kotlin/ai/koog/agents/mcp/DefaultMcpToolDescriptorParserTest.kt‎

Lines changed: 52 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,8 @@ package ai.koog.agents.mcp
33
import ai.koog.agents.core.tools.ToolDescriptor
44
import ai.koog.agents.core.tools.ToolParameterDescriptor
55
import ai.koog.agents.core.tools.ToolParameterType
6-
import io.modelcontextprotocol.kotlin.sdk.Tool
6+
import io.modelcontextprotocol.kotlin.sdk.types.Tool
7+
import io.modelcontextprotocol.kotlin.sdk.types.ToolSchema
78
import kotlinx.serialization.json.JsonObject
89
import kotlinx.serialization.json.add
910
import kotlinx.serialization.json.addJsonArray
@@ -625,8 +626,16 @@ class DefaultMcpToolDescriptorParserTest {
625626
description = "String or number parameter",
626627
type = ToolParameterType.AnyOf(
627628
types = arrayOf(
628-
ToolParameterDescriptor(name = "", description = "String option", type = ToolParameterType.String),
629-
ToolParameterDescriptor(name = "", description = "Number option", type = ToolParameterType.Float)
629+
ToolParameterDescriptor(
630+
name = "",
631+
description = "String option",
632+
type = ToolParameterType.String
633+
),
634+
ToolParameterDescriptor(
635+
name = "",
636+
description = "Number option",
637+
type = ToolParameterType.Float
638+
)
630639
)
631640
)
632641
)
@@ -654,7 +663,10 @@ class DefaultMcpToolDescriptorParserTest {
654663
}
655664
putJsonObject("operator") {
656665
put("type", "string")
657-
put("description", "The operator to use for the filter. Supported operators: '==', 'IN'")
666+
put(
667+
"description",
668+
"The operator to use for the filter. Supported operators: '==', 'IN'"
669+
)
658670
}
659671
putJsonObject("value") {
660672
put("description", "The value(s) to filter by")
@@ -731,23 +743,51 @@ class DefaultMcpToolDescriptorParserTest {
731743
description = "The value(s) to filter by",
732744
type = ToolParameterType.AnyOf(
733745
types = listOf(
734-
ToolParameterDescriptor(name = "", description = "", type = ToolParameterType.Boolean),
735-
ToolParameterDescriptor(name = "", description = "", type = ToolParameterType.Float),
746+
ToolParameterDescriptor(
747+
name = "",
748+
description = "",
749+
type = ToolParameterType.Boolean
750+
),
751+
ToolParameterDescriptor(
752+
name = "",
753+
description = "",
754+
type = ToolParameterType.Float
755+
),
736756
ToolParameterDescriptor(
737757
name = "",
738758
description = "",
739759
type = ToolParameterType.List(
740760
itemsType = ToolParameterType.AnyOf(
741761
types = listOf(
742-
ToolParameterDescriptor(name = "", description = "", type = ToolParameterType.String),
743-
ToolParameterDescriptor(name = "", description = "", type = ToolParameterType.Float),
744-
ToolParameterDescriptor(name = "", description = "", type = ToolParameterType.Boolean),
745-
ToolParameterDescriptor(name = "", description = "", type = ToolParameterType.Null)
762+
ToolParameterDescriptor(
763+
name = "",
764+
description = "",
765+
type = ToolParameterType.String
766+
),
767+
ToolParameterDescriptor(
768+
name = "",
769+
description = "",
770+
type = ToolParameterType.Float
771+
),
772+
ToolParameterDescriptor(
773+
name = "",
774+
description = "",
775+
type = ToolParameterType.Boolean
776+
),
777+
ToolParameterDescriptor(
778+
name = "",
779+
description = "",
780+
type = ToolParameterType.Null
781+
)
746782
).toTypedArray()
747783
)
748784
)
749785
),
750-
ToolParameterDescriptor(name = "", description = "", type = ToolParameterType.String),
786+
ToolParameterDescriptor(
787+
name = "",
788+
description = "",
789+
type = ToolParameterType.String
790+
),
751791
).toTypedArray()
752792
)
753793
)
@@ -771,7 +811,7 @@ class DefaultMcpToolDescriptorParserTest {
771811
return Tool(
772812
name = name,
773813
description = description,
774-
inputSchema = Tool.Input(
814+
inputSchema = ToolSchema(
775815
properties = properties,
776816
required = required
777817
),

0 commit comments

Comments
 (0)