From ea9dfae3bf16b6730e4d5e23e6ca5caefd8b7c68 Mon Sep 17 00:00:00 2001
From: Kynan Ware <47394200+BagToad@users.noreply.github.com>
Date: Fri, 29 Aug 2025 14:23:43 -0600
Subject: [PATCH] Optimize session pull request hydration logic

Replaces linear search with a map for associating sessions with pull requests, improving performance and simplifying code in hydrateSessionPullRequests.
---
 pkg/cmd/agent-task/capi/sessions.go | 20 +++++++-------------
 1 file changed, 7 insertions(+), 13 deletions(-)

diff --git a/pkg/cmd/agent-task/capi/sessions.go b/pkg/cmd/agent-task/capi/sessions.go
index 6d1c96e4b..13b665c90 100644
--- a/pkg/cmd/agent-task/capi/sessions.go
+++ b/pkg/cmd/agent-task/capi/sessions.go
@@ -147,9 +147,9 @@ func (c *CAPIClient) hydrateSessionPullRequests(sessions []session) ([]*Session,
 		return nil, err
 	}
 
-	prs := make([]*api.PullRequest, 0, len(prNodeIds))
+	prMap := make(map[string]*api.PullRequest, len(prNodeIds))
 	for _, node := range resp.Nodes {
-		prs = append(prs, &api.PullRequest{
+		prMap[node.PullRequest.FullDatabaseID] = &api.PullRequest{
 			ID:             node.PullRequest.ID,
 			FullDatabaseID: node.PullRequest.FullDatabaseID,
 			Number:         node.PullRequest.Number,
@@ -162,21 +162,15 @@ func (c *CAPIClient) hydrateSessionPullRequests(sessions []session) ([]*Session,
 			ClosedAt:       node.PullRequest.ClosedAt,
 			MergedAt:       node.PullRequest.MergedAt,
 			Repository:     node.PullRequest.Repository,
-		})
+		}
 	}
 
 	newSessions := make([]*Session, 0, len(sessions))
-	// For each session, we need to attach the Pull Request
 	for _, s := range sessions {
-		// For each Pull Request, check if it matches the session
-		for _, pr := range prs {
-			if strconv.FormatInt(s.ResourceID, 10) == pr.FullDatabaseID {
-				newSessions = append(newSessions, &Session{
-					session:     s,
-					PullRequest: pr,
-				})
-			}
-		}
+		newSessions = append(newSessions, &Session{
+			session:     s,
+			PullRequest: prMap[strconv.FormatInt(s.ResourceID, 10)],
+		})
 	}
 
 	return newSessions, nil