Feature/server enhancements

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

@@ -1,7 +1,6 @@
 package mcp.code.analysis.processor
 
 import java.io.File
-import kotlin.text.lines
 import mcp.code.analysis.service.ModelContextService
 import org.slf4j.Logger
 import org.slf4j.LoggerFactory
@@ -13,8 +12,6 @@ data class LanguagePatterns(
   val blockCommentEnd: String,
 )
 
-data class ProcessingState(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.
@@ -126,15 +123,14 @@ data class CodeAnalyzer(
     findCodeFiles(repoDir)
       .filter { it.isRelevantCodeFile() && !it.isTestFile() }
       .mapNotNull { file ->
-        try {
-          val relativePath = file.getRelativePathFrom(repoDir)
-          val language = getLanguageFromExtension(file.extension)
-          val content = file.readText()
-          summarizeCodeContent(relativePath, language, content, maxLines)
-        } catch (e: Exception) {
-          logger.warn("Error processing file ${file.absolutePath}: ${e.message}")
-          null
-        }
+        runCatching {
+            val relativePath = file.getRelativePathFrom(repoDir)
+            val language = getLanguageFromExtension(file.extension)
+            val content = file.readText()
+            summarizeCodeContent(relativePath, language, content, maxLines)
+          }
+          .onFailure { e -> logger.warn("Error processing file ${file.absolutePath}: ${e.message}") }
+          .getOrNull()
       }
       .also { logCollectionResults(it, repoDir) }
 

src/main/kotlin/mcp/code/analysis/processor/CodeContentProcessor.kt

@@ -13,32 +13,67 @@ internal class CodeContentProcessor(private val patterns: LanguagePatterns, priv
    * Processes a list of code lines and returns those that should be included in the summary. Lines are included if they
    * are definitions, comments, or part of a block comment.
    *
+   * The processing follows a two-pass approach:
+   * 1. First pass: Decides which original lines should be included, tracking comment block state
+   * 2. Second pass: Builds output with explicit separators between non-contiguous regions
+   *
+   * During the first pass, lines are included if they are definitions, comments, or part of a block comment. The
+   * comment block state is tracked to ensure multi-line comments are properly captured.
+   *
+   * During the second pass, separators ("...") are added between non-contiguous regions to indicate omitted code
+   * sections. The maxLines limit is strictly respected, prioritizing code lines over separators.
+   *
    * @param lines The lines of code to process.
-   * @return A list of lines selected for summarization.
+   * @return A list of lines selected for summarization, with separators indicating omitted sections.
    */
   fun processContent(lines: List<String>): List<String> {
-    val finalState =
-      lines.fold(ProcessingState()) { state, line ->
-        if (state.lines.size >= maxLines) return@fold state
+    if (lines.isEmpty()) return emptyList()
+
+    // First pass: compute inclusion flags functionally while tracking the comment block state
+    data class Pass1(val flags: MutableList<Boolean>, val inBlock: Boolean)
 
+    val pass1 =
+      lines.foldIndexed(Pass1(mutableListOf<Boolean>(), false)) { idx, acc, line ->
         val trimmed = line.trim()
-        val nextInCommentBlock = determineCommentBlockState(trimmed, state.inCommentBlock)
-        val shouldIncludeLine = isDefinition(line) || isCommentLine(line) || state.inCommentBlock
-
-        val updatedLines =
-          if (shouldIncludeLine) {
-            when {
-              isDefinition(line) -> state.lines + processDefinitionLine(line)
-              else -> state.lines + line
-            }
-          } else {
-            state.lines
-          }
-
-        ProcessingState(updatedLines, nextInCommentBlock)
+        val shouldInclude = isDefinition(line) || isCommentLine(line) || acc.inBlock
+        acc.flags.add(shouldInclude)
+        val nextInCommentBlock = determineCommentBlockState(trimmed, acc.inBlock)
+        acc.copy(flags = acc.flags, inBlock = nextInCommentBlock)
+      }
+
+    val includeFlags: List<Boolean> = pass1.flags
+
+    // Second pass: build output with separators between non-contiguous regions
+    // Accumulates second-pass output and the index of the last included source line
+    data class OutputAcc(val result: MutableList<String>, val lastIdx: Int)
+
+    fun maybeAddSeparatorFn(state: OutputAcc, nextIndex: Int): OutputAcc {
+      if (state.result.isEmpty()) return state
+      val isGap = nextIndex != state.lastIdx + 1
+      if (!isGap) return state
+      if (state.result.size + 2 > maxLines) return state
+      state.result.add("...")
+      return state
+    }
+
+    val finalAcc: OutputAcc =
+      lines.indices.fold(OutputAcc(mutableListOf(), -2)) { acc, i ->
+        if (!includeFlags[i]) return@fold acc
+
+        val afterSep = maybeAddSeparatorFn(acc, i)
+
+        val line = lines[i]
+        val toAdd = if (isDefinition(line)) processDefinitionLine(line) else line
+
+        if (afterSep.result.size >= maxLines) return@fold afterSep
+
+        afterSep.result.add(toAdd)
+        val updated = afterSep.copy(result = afterSep.result, lastIdx = i)
+
+        if (updated.result.size >= maxLines) updated else updated
       }
 
-    return finalState.lines
+    return finalAcc.result
   }
 
   private fun isDefinition(line: String): Boolean = patterns.definitionPattern.containsMatchIn(line.trim())

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

@@ -7,17 +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.*
 import io.ktor.util.collections.*
 import io.modelcontextprotocol.kotlin.sdk.*
 import io.modelcontextprotocol.kotlin.sdk.server.Server as SdkServer
 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.coroutines.Job
+import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.withTimeout
 import kotlinx.io.asSink
 import kotlinx.io.asSource
 import kotlinx.io.buffered
@@ -85,59 +85,50 @@ class Mcp(
             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")
+              logger.info("Server closed for session: ${transport.sessionId}")
               servers.remove(transport.sessionId)
             }
 
-            server.connect(transport)
-
             try {
+              server.connect(transport)
+              logger.info("Server connected for session: ${transport.sessionId}")
               awaitCancellation()
+            } catch (e: Exception) {
+              logger.error("Connection error: ${e.message}", e)
             } finally {
-              heartbeatJob.cancel()
+              servers.remove(transport.sessionId)
               logger.info("SSE connection closed for session: ${transport.sessionId}")
             }
           }
 
           post("/message") {
             try {
-
               val sessionId =
                 call.request.queryParameters["sessionId"]
                   ?: return@post call.respond(HttpStatusCode.BadRequest, "Missing sessionId parameter")
 
-              val transport = servers[sessionId]?.transport as? SseServerTransport
-              if (transport == null) {
+              val server = servers[sessionId]
+              if (server == null) {
+                logger.warn("Session not found: $sessionId")
                 call.respond(HttpStatusCode.NotFound, "Session not found")
                 return@post
               }
 
+              val transport = server.transport as? SseServerTransport
+              if (transport == null) {
+                logger.warn("Invalid transport for session: $sessionId")
+                call.respond(HttpStatusCode.InternalServerError, "Invalid transport")
+                return@post
+              }
+
               logger.debug("Handling message for session: $sessionId")
-              transport.handlePostMessage(call)
+              withTimeout(3_600_000) { transport.handlePostMessage(call) }
 
               call.respond(HttpStatusCode.OK)
             } catch (e: Exception) {
               logger.error("Error handling message: ${e.message}", e)
-              call.respond(HttpStatusCode.InternalServerError, "Error handling message: ${e.message}")
+              call.respond(HttpStatusCode.InternalServerError, "Error: ${e.message}")
             }
           }
         }

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

@@ -16,22 +16,16 @@ data class RepositoryAnalysisService(
    * @param branch The branch of the repository to analyze.
    * @return A summary of the analysis results.
    */
-  suspend fun analyzeRepository(repoUrl: String, branch: String): String {
-    return try {
-      val repoDir = gitService.cloneRepository(repoUrl, branch)
-
-      val readme = codeAnalyzer.findReadmeFile(repoDir)
-      val codeSnippets = codeAnalyzer.collectSummarizedCodeSnippets(repoDir)
-
-      val insightsPrompt = modelContextService.buildInsightsPrompt(codeSnippets, readme)
-      val insightsResponse = modelContextService.generateResponse(insightsPrompt)
-
-      val summaryPrompt = modelContextService.buildSummaryPrompt(insightsResponse)
-      val summaryResponse = modelContextService.generateResponse(summaryPrompt)
-
-      summaryResponse
-    } catch (e: Exception) {
-      throw Exception("Error analyzing repository: ${e.message}", e)
-    }
-  }
+  suspend fun analyzeRepository(repoUrl: String, branch: String): String =
+    runCatching { gitService.cloneRepository(repoUrl, branch) }
+      .mapCatching { repoDir ->
+        val readme = codeAnalyzer.findReadmeFile(repoDir)
+        val codeSnippets = codeAnalyzer.collectSummarizedCodeSnippets(repoDir)
+        val insightsPrompt = modelContextService.buildInsightsPrompt(codeSnippets, readme)
+        insightsPrompt
+      }
+      .mapCatching { insightsPrompt -> modelContextService.generateResponse(insightsPrompt) }
+      .mapCatching { insightsResponse -> modelContextService.buildSummaryPrompt(insightsResponse) }
+      .mapCatching { summaryPrompt -> modelContextService.generateResponse(summaryPrompt) }
+      .getOrElse { e -> throw Exception("Error analyzing repository: ${e.message}", e) }
 }

src/test/kotlin/mcp/code/analysis/config/AppConfigTest.kt

@@ -1,6 +1,6 @@
 package mcp.code.analysis.config
 
-import org.junit.jupiter.api.Assertions.*
+import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Test
 
 class AppConfigTest {