diff --git a/internal/daemon/daemon.go b/internal/daemon/daemon.go index 7dc34d5..310ff60 100644 --- a/internal/daemon/daemon.go +++ b/internal/daemon/daemon.go @@ -1,16 +1,11 @@ package daemon import ( - "bufio" - "bytes" "context" - "encoding/json" "errors" "fmt" - "io" "log" "net" - "net/http" "net/url" "strings" "sync" @@ -61,8 +56,6 @@ const ( SDKClaude = "claude" ) -const mcpProbeID = "agynd-mcp-ready" - type Daemon struct { cfg config.Config sdk string @@ -848,7 +841,6 @@ type tcpServiceTarget struct { type mcpServiceTarget struct { label string addr string - url string } func waitForTCPService(ctx context.Context, addr string, timeout time.Duration, label string) error { @@ -891,7 +883,6 @@ func waitForMCPServers(ctx context.Context, servers []config.MCPServer, timeout targets = append(targets, mcpServiceTarget{ label: fmt.Sprintf("MCP server %s", server.Name), addr: addr, - url: mcpEndpoint(server.Port), }) } return waitForMCPServices(ctx, targets, timeout) @@ -907,12 +898,11 @@ func waitForMCPServices(ctx context.Context, targets []mcpServiceTarget, timeout } deadline := time.NewTimer(timeout) defer deadline.Stop() - client := &http.Client{Timeout: 10 * time.Second} for _, target := range targets { attempt := 0 for { attempt++ - if err := probeMCPService(ctx, client, target.url); err == nil { + if err := probeMCPService(ctx, target.addr); err == nil { log.Printf("%s at %s ready after %d attempt(s)", target.label, target.addr, attempt) break } else { @@ -930,106 +920,13 @@ func waitForMCPServices(ctx context.Context, targets []mcpServiceTarget, timeout return nil } -func probeMCPService(ctx context.Context, client *http.Client, endpoint string) error { - body := []byte(fmt.Sprintf(`{"jsonrpc":"2.0","id":%q,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"agynd","version":"mcp-ready"}}}`, mcpProbeID)) - req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body)) +func probeMCPService(ctx context.Context, addr string) error { + dialer := net.Dialer{Timeout: 1 * time.Second} + conn, err := dialer.DialContext(ctx, "tcp", addr) if err != nil { return err } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Accept", "application/json, text/event-stream") - resp, err := client.Do(req) - if err != nil { - return err - } - defer resp.Body.Close() - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - return fmt.Errorf("http status %s", resp.Status) - } - return readMCPProbeResult(resp.Body) -} - -type mcpProbeResponse struct { - JSONRPC string `json:"jsonrpc"` - ID string `json:"id"` - Result json.RawMessage `json:"result"` - Error json.RawMessage `json:"error"` -} - -func readMCPProbeResult(body io.Reader) error { - reader := bufio.NewReader(io.LimitReader(body, 64*1024)) - var raw bytes.Buffer - var eventData []string - var lastEventErr error - sawSSE := false - - for { - line, err := reader.ReadString('\n') - if line != "" { - raw.WriteString(line) - trimmed := strings.TrimRight(line, "\r\n") - if data, ok := strings.CutPrefix(trimmed, "data:"); ok { - sawSSE = true - eventData = append(eventData, strings.TrimSpace(data)) - } else if sawSSE && strings.TrimSpace(trimmed) == "" { - if len(eventData) > 0 { - if err := validateMCPProbePayload([]byte(strings.Join(eventData, "\n"))); err == nil { - return nil - } else { - lastEventErr = err - } - } - eventData = nil - } - } - if err == nil { - continue - } - if errors.Is(err, io.EOF) { - break - } - return err - } - - if sawSSE { - if len(eventData) > 0 { - if err := validateMCPProbePayload([]byte(strings.Join(eventData, "\n"))); err == nil { - return nil - } else { - lastEventErr = err - } - } - if lastEventErr != nil { - return lastEventErr - } - return fmt.Errorf("initialize SSE response missing data") - } - return validateMCPProbePayload(raw.Bytes()) -} - -func validateMCPProbePayload(payload []byte) error { - payload = bytes.TrimSpace(payload) - if len(payload) == 0 { - return fmt.Errorf("initialize response is empty") - } - var response mcpProbeResponse - if err := json.Unmarshal(payload, &response); err != nil { - return fmt.Errorf("parse initialize response: %w", err) - } - if response.JSONRPC != "2.0" { - return fmt.Errorf("initialize response jsonrpc %q does not match 2.0", response.JSONRPC) - } - if response.ID != mcpProbeID { - return fmt.Errorf("initialize response id %q does not match %q", response.ID, mcpProbeID) - } - if len(bytes.TrimSpace(response.Error)) > 0 && !bytes.Equal(bytes.TrimSpace(response.Error), []byte("null")) { - return fmt.Errorf("initialize response contains error") - } - result := bytes.TrimSpace(response.Result) - if len(result) == 0 || bytes.Equal(result, []byte("null")) { - return fmt.Errorf("initialize response missing result") - } - return nil + return conn.Close() } type codexThreadDefaults struct { diff --git a/internal/daemon/daemon_mcp_test.go b/internal/daemon/daemon_mcp_test.go index 9d2f9cd..b6971b7 100644 --- a/internal/daemon/daemon_mcp_test.go +++ b/internal/daemon/daemon_mcp_test.go @@ -3,11 +3,8 @@ package daemon import ( "context" "errors" - "fmt" "net" - "net/http" "strings" - "sync/atomic" "testing" "time" @@ -15,8 +12,10 @@ import ( ) func TestWaitForMCPServersReady(t *testing.T) { - _, portA := startMCPReadyServer(t) - _, portB := startMCPReadyServer(t) + listenerA, portA := startAcceptingTCPListener(t) + defer listenerA.Close() + listenerB, portB := startAcceptingTCPListener(t) + defer listenerB.Close() ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) defer cancel() @@ -30,75 +29,42 @@ func TestWaitForMCPServersReady(t *testing.T) { } } -func TestWaitForMCPServersWaitsForInitializeResponse(t *testing.T) { - var requests atomic.Int32 - server := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - attempt := requests.Add(1) - if attempt < 3 { - http.Error(w, "starting", http.StatusServiceUnavailable) +func TestWaitForMCPServersDoesNotSendInitialize(t *testing.T) { + readBytes := make(chan int, 1) + listener, port := startTCPListener(t) + defer listener.Close() + + go func() { + conn, err := listener.Accept() + if err != nil { return } - _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":"agynd-mcp-ready","result":{"protocolVersion":"2025-06-18"}}`)) - })} - listener, port := startHTTPListener(t, server) - defer listener.Close() + defer conn.Close() - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + if err := conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond)); err != nil { + readBytes <- -1 + return + } + buf := make([]byte, 1) + n, _ := conn.Read(buf) + readBytes <- n + }() + + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) defer cancel() servers := []config.MCPServer{{Name: "memory", Port: port}} - if err := waitForMCPServers(ctx, servers, 5*time.Second); err != nil { - t.Fatalf("expected MCP server to become ready, got %v", err) - } - if got := requests.Load(); got < 3 { - t.Fatalf("expected readiness probe retries, got %d request(s)", got) - } -} - -func TestReadMCPProbeResultRejectsInvalidPayloads(t *testing.T) { - tests := []struct { - name string - payload string - }{ - { - name: "json rpc error containing result text", - payload: `{"jsonrpc":"2.0","id":"agynd-mcp-ready","error":{"message":"no result yet"}}`, - }, - { - name: "unrelated payload containing result text", - payload: `{"message":"result is almost ready"}`, - }, - { - name: "wrong response id", - payload: `{"jsonrpc":"2.0","id":"other","result":{"protocolVersion":"2025-06-18"}}`, - }, - { - name: "null result", - payload: `{"jsonrpc":"2.0","id":"agynd-mcp-ready","result":null}`, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if err := readMCPProbeResult(strings.NewReader(tt.payload)); err == nil { - t.Fatal("expected invalid MCP initialize response to be rejected") - } - }) + if err := waitForMCPServers(ctx, servers, 500*time.Millisecond); err != nil { + t.Fatalf("expected MCP listener to be ready, got %v", err) } -} -func TestReadMCPProbeResultAcceptsSSEInitializeResult(t *testing.T) { - payload := strings.Join([]string{ - `event: message`, - `data: {"jsonrpc":"2.0","id":"unrelated","result":{"protocolVersion":"2025-06-18"}}`, - ``, - `event: message`, - `data: {"jsonrpc":"2.0","id":"agynd-mcp-ready","result":{"protocolVersion":"2025-06-18"}}`, - ``, - }, "\n") - - if err := readMCPProbeResult(strings.NewReader(payload)); err != nil { - t.Fatalf("expected SSE initialize result to be accepted, got %v", err) + select { + case got := <-readBytes: + if got != 0 { + t.Fatalf("expected readiness probe to send no payload, got %d byte(s)", got) + } + case <-time.After(500 * time.Millisecond): + t.Fatal("listener did not observe readiness connection") } } @@ -150,37 +116,19 @@ func startTCPListener(t *testing.T) (net.Listener, int) { return listener, addr.Port } -func startMCPReadyServer(t *testing.T) (*http.Server, int) { +func startAcceptingTCPListener(t *testing.T) (net.Listener, int) { t.Helper() - server := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != "/mcp" { - http.NotFound(w, r) - return - } - _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":"agynd-mcp-ready","result":{"protocolVersion":"2025-06-18"}}`)) - })} - listener, port := startHTTPListener(t, server) - t.Cleanup(func() { _ = listener.Close() }) - return server, port -} - -func startHTTPListener(t *testing.T, server *http.Server) (net.Listener, int) { - t.Helper() - listener, err := net.Listen("tcp", "127.0.0.1:0") - if err != nil { - t.Fatalf("listen: %v", err) - } - addr, ok := listener.Addr().(*net.TCPAddr) - if !ok { - _ = listener.Close() - t.Fatalf("unexpected listener address type %T", listener.Addr()) - } + listener, port := startTCPListener(t) go func() { - if err := server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) && !strings.Contains(err.Error(), "use of closed network connection") { - panic(fmt.Sprintf("serve MCP ready test server: %v", err)) + for { + conn, err := listener.Accept() + if err != nil { + return + } + _ = conn.Close() } }() - return listener, addr.Port + return listener, port } func unusedTCPPort(t *testing.T) int { diff --git a/internal/daemon/daemon_test.go b/internal/daemon/daemon_test.go index 2848dbd..9896920 100644 --- a/internal/daemon/daemon_test.go +++ b/internal/daemon/daemon_test.go @@ -4,10 +4,9 @@ import ( "context" "errors" "fmt" - "net/http" + "net" "strings" "sync" - "sync/atomic" "testing" "time" @@ -570,17 +569,21 @@ func TestRunRetriesReadbackTransportFailure(t *testing.T) { } func TestSyncMessagesWaitsForMCPReadyBeforeHandling(t *testing.T) { - var requests atomic.Int32 - server := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - attempt := requests.Add(1) - if attempt < 3 { - http.Error(w, "starting", http.StatusServiceUnavailable) + port := unusedTCPPort(t) + listenerReady := make(chan struct{}) + go func() { + time.Sleep(50 * time.Millisecond) + listener, err := net.Listen("tcp", net.JoinHostPort(mcpLoopbackHost, fmt.Sprintf("%d", port))) + if err != nil { return } - _, _ = w.Write([]byte(`{"jsonrpc":"2.0","id":"agynd-mcp-ready","result":{"protocolVersion":"2025-06-18"}}`)) - })} - listener, port := startHTTPListener(t, server) - defer listener.Close() + defer listener.Close() + close(listenerReady) + conn, err := listener.Accept() + if err == nil { + _ = conn.Close() + } + }() consumer := &handlingMessageConsumer{message: platform.Message{ID: "msg-1", ThreadID: "thread-1", Body: "hello"}} daemon := &Daemon{ @@ -589,12 +592,17 @@ func TestSyncMessagesWaitsForMCPReadyBeforeHandling(t *testing.T) { consumer: consumer, } - err := daemon.syncMessages(context.Background()) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + err := daemon.syncMessages(ctx) if err == nil || !strings.Contains(err.Error(), "unknown sdk") { t.Fatalf("expected handler to run after MCP readiness, got %v", err) } - if got := requests.Load(); got < 3 { - t.Fatalf("expected MCP readiness retries before handling, got %d request(s)", got) + select { + case <-listenerReady: + default: + t.Fatal("expected MCP listener to start before message handling") } if calls := consumer.Calls(); calls != 1 { t.Fatalf("expected one sync call, got %d", calls)