Dispatch MCP requests concurrently per connection

handleConn processed requests serially, so a slow tool (e.g.
wait_for_pattern with a 300s timeout) monopolized the single per-agent
MCP connection and every queued call timed out behind it. Handle each
request in its own goroutine, serialize responses through a per-conn
write mutex (full response written atomically, partial writes handled),
copy the request line before handing it off (bufio reuses its buffer),
and wait on a WaitGroup before closing the conn so in-flight handlers
finish cleanly. Greeting stays sequential; notifications still get no
response.

Resolves the [MCP TIMEOUT] TODO item.
This commit is contained in:
2026-05-25 12:39:31 +01:00
parent 53f06b604f
commit 7b5a22618f
4 changed files with 226 additions and 24 deletions

View File

@@ -96,10 +96,34 @@ func (s *Server) acceptLoop() {
// identity token (SPEC §10); we resolve it to a child id and stash that
// as the caller for every subsequent tool call.
func (s *Server) handleConn(conn net.Conn) {
defer conn.Close()
var writeMu sync.Mutex
var wg sync.WaitGroup
defer func() {
wg.Wait()
_ = conn.Close()
}()
r := bufio.NewReader(conn)
var callerID string
writeResp := func(resp []byte) bool {
if resp == nil {
return true
}
resp = append(resp, '\n')
writeMu.Lock()
defer writeMu.Unlock()
for len(resp) > 0 {
n, err := conn.Write(resp)
if err != nil {
return false
}
if n == 0 {
return false
}
resp = resp[n:]
}
return true
}
greeting, err := r.ReadBytes('\n')
if err != nil {
@@ -115,24 +139,21 @@ func (s *Server) handleConn(conn net.Conn) {
} else {
// Treat as a real request from an unknown caller.
resp := s.dispatch("", greeting)
if resp != nil {
resp = append(resp, '\n')
if _, werr := conn.Write(resp); werr != nil {
return
}
if !writeResp(resp) {
return
}
}
for {
line, err := r.ReadBytes('\n')
if len(line) > 0 {
resp := s.dispatch(callerID, line)
if resp != nil {
resp = append(resp, '\n')
if _, werr := conn.Write(resp); werr != nil {
return
}
}
req := append([]byte(nil), line...)
wg.Add(1)
go func() {
defer wg.Done()
resp := s.dispatch(callerID, req)
_ = writeResp(resp)
}()
}
if err != nil {
return