fix: prevent MCP server hanging by adding concurrency, timeouts, and context propagation (#45)

Root cause: synchronous MCP request processing combined with missing
context propagation to the embedding layer caused indefinite hangs when
ONNX inference was slow or the database was contended.

Changes:
- MCP server: dispatch each request in its own goroutine with semaphore
  (cap 10) and WaitGroup for clean shutdown drain
- Embedding: add context-aware mutex acquisition (acquireMutex) so
  callers can bail out instead of blocking forever on a stuck ONNX model
- Vector client: propagate context through getOrComputeEmbedding and
  replace bare RLock() calls with context-aware acquireRLockWithContext
- Worker handlers: add 15s request-scoped timeouts to all search/context
  handlers (handleSearchByPrompt, handleContextInject, handleFileContext,
  handleContextCount, handleGetObservations/Summaries/Prompts)
- Worker HTTP server: set WriteTimeout=60s (was 0); SSE endpoint extends
  deadline per-request via http.ResponseController

Fixes #45
This commit is contained in:
2026-05-26 14:29:34 +01:00
parent 56616d0616
commit 29d57857ff
7 changed files with 244 additions and 40 deletions
+24 -6
View File
@@ -129,13 +129,22 @@ func (s *Server) Run(ctx context.Context) error {
}
}()
// Semaphore limits concurrent request goroutines.
const maxConcurrent = 10
sem := make(chan struct{}, maxConcurrent)
var wg sync.WaitGroup
for {
select {
case <-ctx.Done():
// Drain in-flight requests before returning.
wg.Wait()
return ctx.Err()
case line, ok := <-lines:
if !ok {
// Scanner finished, check for errors
// Scanner finished — drain in-flight requests, then check for errors.
wg.Wait()
err := <-scanErr
if err != nil {
if errors.Is(err, bufio.ErrTooLong) {
@@ -154,18 +163,27 @@ func (s *Server) Run(ctx context.Context) error {
var req Request
if err := json.Unmarshal([]byte(line), &req); err != nil {
// Parse errors are cheap — send inline, no goroutine needed.
if werr := s.sendError(nil, -32700, "Parse error", err.Error()); werr != nil {
return fmt.Errorf("write error: %w", werr)
}
continue
}
resp := s.handleRequest(ctx, &req)
if resp != nil {
if werr := s.sendResponse(resp); werr != nil {
return fmt.Errorf("write error: %w", werr)
// Dispatch request to its own goroutine.
wg.Add(1)
sem <- struct{}{} // acquire semaphore slot
go func(r Request) {
defer wg.Done()
defer func() { <-sem }() // release semaphore slot
resp := s.handleRequest(ctx, &r)
if resp != nil {
if werr := s.sendResponse(resp); werr != nil {
log.Error().Err(werr).Msg("Failed to send response")
}
}
}
}(req)
}
}
}