feat: server error handlers

build.gradle.kts

@@ -10,8 +10,9 @@ plugins {
   kotlin("jvm") version "2.1.0"
   kotlin("plugin.serialization") version "2.1.0"
   id("com.diffplug.spotless") version "7.0.3"
-  id("pl.allegro.tech.build.axion-release") version "1.18.7"
+  id("io.kotest.multiplatform") version "5.0.2"
   id("org.jreleaser") version "1.17.0"
+  id("pl.allegro.tech.build.axion-release") version "1.18.7"
 }
 
 scmVersion {
@@ -28,6 +29,7 @@ description = "MCP Server for GitHub Code Repositories Analysis"
 dependencies {
   val ktorVersion = "3.0.2"
   val coroutinesVersion = "1.10.2"
+  val kotestVersion = "5.9.1"
 
   // Kotlin standard library
   implementation(kotlin("stdlib"))
@@ -63,6 +65,10 @@ dependencies {
   testImplementation("io.mockk:mockk:1.14.2")
   testImplementation("io.ktor:ktor-client-mock-jvm:$ktorVersion")
   testImplementation("org.jetbrains.kotlinx:kotlinx-coroutines-test:$coroutinesVersion")
+
+  testImplementation("io.kotest:kotest-runner-junit5:$kotestVersion")
+  testImplementation("io.kotest:kotest-assertions-core:$kotestVersion")
+  testImplementation("io.kotest:kotest-property:$kotestVersion")
 }
 
 application { mainClass.set("MainKt") }

src/main/kotlin/Main.kt

@@ -37,9 +37,10 @@ data class AppCommand(val type: CommandType, val port: Int)
 
 /**
  * Enum representing the application command type.
- *
- * STDIO - Standard input/output mode for the MCP server SSE_KTOR - Server-Sent Events mode using the Ktor plugin
- * SSE_PLAIN - Server-Sent Events mode using plain configuration UNKNOWN - Unknown or unsupported command
+ * * STDIO - Standard input/output mode for the MCP server
+ * * SSE_KTOR - Server-Sent Events mode using the Ktor plugin
+ * * SSE_PLAIN - Server-Sent Events mode using plain configuration
+ * * UNKNOWN - Unknown or unsupported command
  */
 enum class CommandType {
   STDIO,
@@ -67,13 +68,13 @@ enum class CommandType {
 }
 
 /**
- * Runs the application by parsing arguments, ensuring necessary directories exist, and executing the appropriate server
- * command.
+ * Runs the application by parsing arguments, ensuring the necessary directories exist, and executing the appropriate
+ * server command.
  *
  * @param args Command-line arguments to determine server behavior.
  * @return Result wrapping Unit, with success if the application runs successfully, or failure with the exception if an
  *   error occurs.
- * @throws IOException If required directories cannot be created.
+ * @throws IOException If the required directories cannot be created.
  */
 fun runApplication(args: Array<String>): Result<Unit> = runCatching {
   val logger = LoggerFactory.getLogger("Main")

src/main/kotlin/mcp/code/analysis/server/Server.kt

@@ -7,15 +7,17 @@ import io.ktor.server.engine.*
 import io.ktor.server.response.*
 import io.ktor.server.routing.*
 import io.ktor.server.sse.*
-import io.ktor.sse.ServerSentEvent
+import io.ktor.sse.*
 import io.ktor.util.collections.*
 import io.modelcontextprotocol.kotlin.sdk.*
-import io.modelcontextprotocol.kotlin.sdk.server.*
 import io.modelcontextprotocol.kotlin.sdk.server.Server as SdkServer
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.delay
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.runBlocking
+import io.modelcontextprotocol.kotlin.sdk.server.ServerOptions
+import io.modelcontextprotocol.kotlin.sdk.server.SseServerTransport
+import io.modelcontextprotocol.kotlin.sdk.server.StdioServerTransport
+import io.modelcontextprotocol.kotlin.sdk.server.mcp
+import kotlinx.coroutines.*
+import kotlinx.coroutines.flow.*
+import kotlinx.io.IOException
 import kotlinx.io.asSink
 import kotlinx.io.asSource
 import kotlinx.io.buffered
@@ -79,24 +81,41 @@ class Server(
 
         routing {
           sse("/sse") {
-            launch {
-              while (true) {
-                send(ServerSentEvent(event = "heartbeat"))
-                delay(15_000)
-              }
-            }
-
             val transport = SseServerTransport("/message", this)
             val server = configureServer()
-
             servers[transport.sessionId] = server
 
+            val heartbeatJob = launch {
+              flow {
+                  while (true) {
+                    emit(Unit)
+                    delay(15_000)
+                  }
+                }
+                .onEach { send(ServerSentEvent(event = "heartbeat")) }
+                .catch { e ->
+                  when (e) {
+                    is IOException -> logger.debug("Client disconnected during heartbeat: ${e.message}")
+                    else -> logger.error("Heartbeat error: ${e.message}", e)
+                  }
+                }
+                .onCompletion { logger.debug("Heartbeat job terminated for session: ${transport.sessionId}") }
+                .collect()
+            }
+
             server.onClose {
               logger.info("Server closed")
               servers.remove(transport.sessionId)
             }
 
             server.connect(transport)
+
+            try {
+              awaitCancellation()
+            } finally {
+              heartbeatJob.cancel()
+              logger.info("SSE connection closed for session: ${transport.sessionId}")
+            }
           }
 
           post("/message") {

src/main/kotlin/mcp/code/analysis/service/CodeAnalyzer.kt

@@ -1,9 +1,19 @@
 package mcp.code.analysis.service
 
 import java.io.File
+import kotlin.text.lines
 import org.slf4j.Logger
 import org.slf4j.LoggerFactory
 
+data class LanguagePatterns(
+  val definitionPattern: Regex,
+  val commentPrefixes: List<String>,
+  val blockCommentStart: String,
+  val blockCommentEnd: String,
+)
+
+data class State(val lines: List<String> = emptyList(), val inCommentBlock: Boolean = false)
+
 /**
  * Responsible for analyzing the structure of a codebase. Identifies files, directories, and their respective metadata
  * such as size, language, imports, and declarations.
@@ -23,37 +33,182 @@ data class CodeAnalyzer(
   fun analyzeStructure(repoDir: File): Map<String, Any> = processDirectory(repoDir, repoDir.absolutePath)
 
   /**
-   * Collects all code snippets from the repository.
+   * Collects summarized code snippets from the repository.
    *
    * @param repoDir The root directory of the repository
-   * @return List of code snippets with metadata including file path and language
+   * @param maxLines The maximum number of lines to include per file summary
+   * @return List of code summaries with metadata
    */
-  fun collectAllCodeSnippets(repoDir: File): List<String> =
+  fun collectSummarizedCodeSnippets(repoDir: File, maxLines: Int = 100): List<String> =
     findCodeFiles(repoDir)
       .filter { file ->
-        file.extension.lowercase() in setOf("kt", "java", "scala", "py", "rb", "js", "ts", "go", "c", "cpp", "rust") &&
+        file.extension.lowercase() in setOf("kt", "java", "scala", "py", "rb", "js", "ts", "go", "c", "cpp", "rs") &&
           !file.absolutePath.contains("test", ignoreCase = true)
       }
       .map { file ->
-        val relativePath = file.absolutePath.substring(repoDir.absolutePath.length + 1)
+        val relativePath = file.absolutePath.removePrefix(repoDir.absolutePath).removePrefix("/")
         val lang = getLanguageFromExtension(file.extension)
-        val content = file.readLines().joinToString("\n")
-        """|--- File: $relativePath
-           |~~~$lang
-           |$content
-           |~~~"""
-          .trimMargin()
+        val content = file.readText()
+        summarizeCodeContent(relativePath, lang, content, maxLines)
       }
-      .toList()
       .also { snippets ->
-        logger.info("Collected ${snippets.size} code snippets from ${repoDir.absolutePath}")
+        logger.info("Collected ${snippets.size} summarized snippets from ${repoDir.absolutePath}")
         logger.debug(
           """|Snippets Found:
-             |${snippets.joinToString("\n")}"""
+             |${snippets.joinToString("\n")}
+             |"""
             .trimMargin()
         )
       }
 
+  /**
+   * Summarizes the content of a file.
+   *
+   * @param path The path of the file
+   * @param language The language of the file
+   * @param content The content of the file
+   * @param maxLines The maximum number of lines to include in the summary
+   */
+  fun summarizeCodeContent(path: String, language: String, content: String, maxLines: Int = 250): String {
+
+    val patterns =
+      when (language.lowercase()) {
+        "kotlin" ->
+          LanguagePatterns(
+            Regex(
+              """(class|interface|object|enum class|data class|sealed class|fun|val|var|const|typealias|annotation class).*"""
+            ),
+            listOf("//"),
+            "/*",
+            "*/",
+          )
+
+        "scala" ->
+          LanguagePatterns(
+            Regex(
+              """(class|object|trait|case class|case object|def|val|var|lazy val|type|implicit|sealed|abstract|override|package object).*"""
+            ),
+            listOf("//"),
+            "/*",
+            "*/",
+          )
+
+        "java" ->
+          LanguagePatterns(
+            Regex(
+              """(class|interface|enum|@interface|record|public|private|protected|static|abstract|final|synchronized|volatile|native|transient|strictfp).*"""
+            ),
+            listOf("//"),
+            "/*",
+            "*/",
+          )
+
+        "python" ->
+          LanguagePatterns(Regex("""(def|class|async def|@|import|from).*"""), listOf("#"), "\"\"\"", "\"\"\"")
+
+        "ruby" ->
+          LanguagePatterns(
+            Regex("""(def|class|module|attr_|require|include|extend).*"""),
+            listOf("#"),
+            "=begin",
+            "=end",
+          )
+
+        "javascript",
+        "typescript" ->
+          LanguagePatterns(
+            Regex("""(function|class|const|let|var|import|export|interface|type|enum|namespace).*"""),
+            listOf("//"),
+            "/*",
+            "*/",
+          )
+
+        "go" ->
+          LanguagePatterns(
+            Regex("""(func|type|struct|interface|package|import|var|const).*"""),
+            listOf("//"),
+            "/*",
+            "*/",
+          )
+
+        "rust" ->
+          LanguagePatterns(
+            Regex("""(fn|struct|enum|trait|impl|pub|use|mod|const|static|type|async|unsafe).*"""),
+            listOf("//"),
+            "/*",
+            "*/",
+          )
+
+        "c",
+        "cpp" ->
+          LanguagePatterns(
+            Regex("""(class|struct|enum|typedef|namespace|template|void|int|char|bool|auto|extern|static|virtual).*"""),
+            listOf("//"),
+            "/*",
+            "*/",
+          )
+
+        // Default fallback for other languages
+        else ->
+          LanguagePatterns(
+            Regex("""(class|interface|object|enum|fun|def|function|public|private|protected|static).*"""),
+            listOf("//", "#"),
+            "/*",
+            "*/",
+          )
+      }
+
+    val definitionPattern = patterns.definitionPattern
+    val commentPrefixes = patterns.commentPrefixes
+    val blockCommentStart = patterns.blockCommentStart
+    val blockCommentEnd = patterns.blockCommentEnd
+
+    val isDefinition: (String) -> Boolean = { line -> line.trim().matches(definitionPattern) }
+
+    val isCommentLine: (String) -> Boolean = { line ->
+      val trimmed = line.trim()
+      commentPrefixes.any { trimmed.startsWith(it) } || trimmed.startsWith(blockCommentStart) || trimmed.startsWith("*")
+    }
+
+    val processDefinitionLine: (String) -> String = { line ->
+      val trimmed = line.trim()
+      if (trimmed.contains("{") && !trimmed.contains("}")) "$trimmed }" else trimmed
+    }
+
+    val finalState =
+      content.lines().fold(State()) { state, line ->
+        if (state.lines.size >= maxLines) return@fold state
+        val trimmed = line.trim()
+        val nextInCommentBlock =
+          when {
+            trimmed.startsWith(blockCommentStart) -> true
+            trimmed.endsWith(blockCommentEnd) -> false
+            language.lowercase() == "python" && trimmed == "\"\"\"" -> !state.inCommentBlock
+            else -> state.inCommentBlock
+          }
+
+        val shouldIncludeLine = isDefinition(line) || isCommentLine(line) || state.inCommentBlock
+        val updatedLines =
+          if (shouldIncludeLine) {
+            if (isDefinition(line)) {
+              // Apply processing to definition lines to ensure braces are complete
+              state.lines + processDefinitionLine(line)
+            } else {
+              state.lines + line
+            }
+          } else state.lines
+
+        State(updatedLines, nextInCommentBlock)
+      }
+
+    // Ensure we're using the correct file path and language
+    return """|### File: $path
+              |~~~$language
+              |${finalState.lines.joinToString("\n")}
+              |~~~"""
+      .trimMargin()
+  }
+
   /**
    * Finds the README file in the repository.
    *