diff --git a/pkg/cmd/agent-task/capi/sessions.go b/pkg/cmd/agent-task/capi/sessions.go index a3989a862..4d109ea26 100644 --- a/pkg/cmd/agent-task/capi/sessions.go +++ b/pkg/cmd/agent-task/capi/sessions.go @@ -27,24 +27,25 @@ var ErrSessionNotFound = errors.New("not found") // session is an in-flight agent task type session struct { - ID string `json:"id"` - Name string `json:"name"` - UserID int64 `json:"user_id"` - AgentID int64 `json:"agent_id"` - Logs string `json:"logs"` - State string `json:"state"` - OwnerID uint64 `json:"owner_id"` - RepoID uint64 `json:"repo_id"` - ResourceType string `json:"resource_type"` - ResourceID int64 `json:"resource_id"` - LastUpdatedAt time.Time `json:"last_updated_at,omitempty"` - CreatedAt time.Time `json:"created_at,omitempty"` - CompletedAt time.Time `json:"completed_at,omitempty"` - EventURL string `json:"event_url"` - EventType string `json:"event_type"` - PremiumRequests float64 `json:"premium_requests"` - WorkflowRunID uint64 `json:"workflow_run_id,omitempty"` - Error *struct { + ID string `json:"id"` + Name string `json:"name"` + UserID int64 `json:"user_id"` + AgentID int64 `json:"agent_id"` + Logs string `json:"logs"` + State string `json:"state"` + OwnerID uint64 `json:"owner_id"` + RepoID uint64 `json:"repo_id"` + ResourceType string `json:"resource_type"` + ResourceID int64 `json:"resource_id"` + ResourceGlobalID string `json:"resource_global_id"` + LastUpdatedAt time.Time `json:"last_updated_at,omitempty"` + CreatedAt time.Time `json:"created_at,omitempty"` + CompletedAt time.Time `json:"completed_at,omitempty"` + EventURL string `json:"event_url"` + EventType string `json:"event_type"` + PremiumRequests float64 `json:"premium_requests"` + WorkflowRunID uint64 `json:"workflow_run_id,omitempty"` + Error *struct { Code string `json:"code"` Message string `json:"message"` } `json:"error,omitempty"` @@ -100,6 +101,26 @@ type SessionError struct { Message string } +type resource struct { + ID string `json:"id"` + UserID uint64 `json:"user_id"` + ResourceType string `json:"resource_type"` + ResourceID int64 `json:"resource_id"` + ResourceGlobalID string `json:"resource_global_id"` + SessionCount int `json:"session_count"` + SessionLastUpdatedAt int64 `json:"last_updated_at"` + SessionState string `json:"state,omitempty"` + ResourceState string `json:"resource_state"` + Sessions []resourceSession `json:"sessions"` +} + +type resourceSession struct { + SessionID string `json:"id"` + Name string `json:"name"` + SessionState string `json:"state,omitempty"` + SessionLastUpdatedAt int64 `json:"last_updated_at"` +} + // ListLatestSessionsForViewer lists all agent sessions for the // authenticated user up to limit. func (c *CAPIClient) ListLatestSessionsForViewer(ctx context.Context, limit int) ([]*Session, error) { @@ -258,46 +279,42 @@ func (c *CAPIClient) ListSessionsByResourceID(ctx context.Context, resourceType return nil, nil } - url := fmt.Sprintf("%s/agents/sessions/resource/%s/%d", baseCAPIURL, url.PathEscape(resourceType), resourceID) - pageSize := defaultSessionsPerPage + url := fmt.Sprintf("%s/agents/resource/%s/%d", baseCAPIURL, url.PathEscape(resourceType), resourceID) - sessions := make([]session, 0, limit+pageSize) - - for page := 1; ; page++ { - req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, http.NoBody) - if err != nil { - return nil, err - } - - q := req.URL.Query() - q.Set("page_size", strconv.Itoa(pageSize)) - q.Set("page_number", strconv.Itoa(page)) - req.URL.RawQuery = q.Encode() - - res, err := c.httpClient.Do(req) - if err != nil { - return nil, err - } - defer res.Body.Close() - if res.StatusCode != http.StatusOK { - return nil, fmt.Errorf("failed to list sessions: %s", res.Status) - } - var response struct { - Sessions []session `json:"sessions"` - } - if err := json.NewDecoder(res.Body).Decode(&response); err != nil { - return nil, fmt.Errorf("failed to decode sessions response: %w", err) - } - - sessions = append(sessions, response.Sessions...) - if len(response.Sessions) < pageSize || len(sessions) >= limit { - break - } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, http.NoBody) + if err != nil { + return nil, err } - // Drop any above the limit - if len(sessions) > limit { - sessions = sessions[:limit] + res, err := c.httpClient.Do(req) + if err != nil { + return nil, err + } + defer res.Body.Close() + if res.StatusCode != http.StatusOK { + return nil, fmt.Errorf("failed to list sessions: %s", res.Status) + } + + var response resource + if err := json.NewDecoder(res.Body).Decode(&response); err != nil { + return nil, fmt.Errorf("failed to decode sessions response: %w", err) + } + + sessions := make([]session, 0, len(response.Sessions)) + for _, s := range response.Sessions { + session := session{ + ID: s.SessionID, + Name: s.Name, + UserID: int64(response.UserID), + ResourceType: response.ResourceType, + ResourceID: response.ResourceID, + ResourceGlobalID: response.ResourceGlobalID, + State: s.SessionState, + } + if s.SessionLastUpdatedAt != 0 { + session.LastUpdatedAt = time.Unix(s.SessionLastUpdatedAt, 0).UTC() + } + sessions = append(sessions, session) } result, err := c.hydrateSessionPullRequestsAndUsers(sessions) @@ -317,7 +334,12 @@ func (c *CAPIClient) hydrateSessionPullRequestsAndUsers(sessions []session) ([]* userNodeIds := make([]string, 0, len(sessions)) for _, session := range sessions { if session.ResourceType == "pull" { - prNodeID := generatePullRequestNodeID(int64(session.RepoID), session.ResourceID) + prNodeID := session.ResourceGlobalID + // TODO: probably this can be dropped since the API should always + // keep returning the resource global ID. + if session.ResourceGlobalID == "" { + prNodeID = generatePullRequestNodeID(int64(session.RepoID), session.ResourceID) + } if !slices.Contains(prNodeIds, prNodeID) { prNodeIds = append(prNodeIds, prNodeID) } diff --git a/pkg/cmd/agent-task/capi/sessions_test.go b/pkg/cmd/agent-task/capi/sessions_test.go index daccbeaed..848fd716a 100644 --- a/pkg/cmd/agent-task/capi/sessions_test.go +++ b/pkg/cmd/agent-task/capi/sessions_test.go @@ -1200,9 +1200,10 @@ func TestListSessionsByResourceIDRequiresResource(t *testing.T) { } func TestListSessionsByResourceID(t *testing.T) { - sampleDateString := "2025-08-29T00:00:00Z" + sampleDateString := "2025-08-29T07:00:00Z" sampleDate, err := time.Parse(time.RFC3339, sampleDateString) require.NoError(t, err) + sampleDateTimestamp := sampleDate.Unix() resourceID := int64(999) resourceType := "pull" @@ -1221,54 +1222,47 @@ func TestListSessionsByResourceID(t *testing.T) { wantOut: nil, }, { - name: "no sessions", + // If the given pull request does not exist or the pull request has no sessions, + // the API endpoint returns 404 with different messages. We should treat them + // the same though. + name: "no sessions or no pull request", limit: 10, httpStubs: func(t *testing.T, reg *httpmock.Registry) { reg.Register( - httpmock.WithHost( - httpmock.QueryMatcher("GET", "agents/sessions/resource/pull/999", url.Values{ - "page_number": {"1"}, - "page_size": {"50"}, - }), - "api.githubcopilot.com", - ), - httpmock.StringResponse(`{"sessions":[]}`), + httpmock.WithHost(httpmock.REST("GET", "agents/resource/pull/999"), "api.githubcopilot.com"), + + httpmock.StatusStringResponse(404, "{}"), ) }, - wantOut: nil, + wantErr: "failed to list sessions", }, { name: "single session", limit: 10, httpStubs: func(t *testing.T, reg *httpmock.Registry) { reg.Register( - httpmock.WithHost( - httpmock.QueryMatcher("GET", "agents/sessions/resource/pull/999", url.Values{ - "page_number": {"1"}, - "page_size": {"50"}, - }), - "api.githubcopilot.com", - ), + httpmock.WithHost(httpmock.REST("GET", "agents/resource/pull/999"), "api.githubcopilot.com"), httpmock.StringResponse(heredoc.Docf(` { + "id": "resource:pull:2000", + "user_id": 1, + "resource_global_id": "PR_kwDNA-jNB9A", + "resource_type": "pull", + "resource_id": 2000, + "session_count": 1, + "last_updated_at": %[1]d, + "state": "completed", + "resource_state": "draft", "sessions": [ { "id": "sess1", "name": "Build artifacts", - "user_id": 1, - "agent_id": 2, - "logs": "", "state": "completed", - "owner_id": 10, - "repo_id": 1000, - "resource_type": "pull", - "resource_id": 2000, - "created_at": "%[1]s", - "premium_requests": 0.1 + "last_updated_at": %[1]d } ] }`, - sampleDateString, + sampleDateTimestamp, )), ) // GraphQL hydration @@ -1311,19 +1305,14 @@ func TestListSessionsByResourceID(t *testing.T) { }, wantOut: []*Session{ { - - ID: "sess1", - Name: "Build artifacts", - UserID: 1, - AgentID: 2, - Logs: "", - State: "completed", - OwnerID: 10, - RepoID: 1000, - ResourceType: "pull", - ResourceID: 2000, - CreatedAt: sampleDate, - PremiumRequests: 0.1, + ID: "sess1", + CreatedAt: time.Time{}, + LastUpdatedAt: sampleDate, + Name: "Build artifacts", + UserID: 1, + State: "completed", + ResourceType: "pull", + ResourceID: 2000, PullRequest: &api.PullRequest{ ID: "PR_node", FullDatabaseID: "2000", @@ -1348,20 +1337,23 @@ func TestListSessionsByResourceID(t *testing.T) { }, }, { - name: "multiple sessions, paginated", + name: "multiple sessions", perPage: 1, // to enforce pagination limit: 2, httpStubs: func(t *testing.T, reg *httpmock.Registry) { reg.Register( - httpmock.WithHost( - httpmock.QueryMatcher("GET", "agents/sessions/resource/pull/999", url.Values{ - "page_number": {"1"}, - "page_size": {"1"}, - }), - "api.githubcopilot.com", - ), + httpmock.WithHost(httpmock.REST("GET", "agents/resource/pull/999"), "api.githubcopilot.com"), httpmock.StringResponse(heredoc.Docf(` { + "id": "resource:pull:2000", + "user_id": 1, + "resource_global_id": "PR_kwDNA-jNB9A", + "resource_type": "pull", + "resource_id": 2000, + "session_count": 1, + "last_updated_at": %[1]d, + "state": "completed", + "resource_state": "draft", "sessions": [ { "id": "sess1", @@ -1374,27 +1366,9 @@ func TestListSessionsByResourceID(t *testing.T) { "repo_id": 1000, "resource_type": "pull", "resource_id": 2000, - "created_at": "%[1]s", + "created_at": %[1]d, "premium_requests": 0.1 - } - ] - }`, - sampleDateString, - )), - ) - - // Second page - reg.Register( - httpmock.WithHost( - httpmock.QueryMatcher("GET", "agents/sessions/resource/pull/999", url.Values{ - "page_number": {"2"}, - "page_size": {"1"}, - }), - "api.githubcopilot.com", - ), - httpmock.StringResponse(heredoc.Docf(` - { - "sessions": [ + }, { "id": "sess2", "name": "Build artifacts", @@ -1406,12 +1380,12 @@ func TestListSessionsByResourceID(t *testing.T) { "repo_id": 1000, "resource_type": "pull", "resource_id": 2001, - "created_at": "%[1]s", + "created_at": %[1]d, "premium_requests": 0.1 } ] }`, - sampleDateString, + sampleDateTimestamp, )), ) // GraphQL hydration @@ -1437,22 +1411,6 @@ func TestListSessionsByResourceID(t *testing.T) { "nameWithOwner": "OWNER/REPO" } }, - { - "__typename": "PullRequest", - "id": "PR_node", - "fullDatabaseId": "2001", - "number": 43, - "title": "Improve docs", - "state": "OPEN", - "isDraft": true, - "url": "https://github.com/OWNER/REPO/pull/43", - "body": "", - "createdAt": "%[1]s", - "updatedAt": "%[1]s", - "repository": { - "nameWithOwner": "OWNER/REPO" - } - }, { "__typename": "User", "login": "octocat", @@ -1464,24 +1422,18 @@ func TestListSessionsByResourceID(t *testing.T) { }`, sampleDateString, ), func(q string, vars map[string]interface{}) { - assert.Equal(t, []interface{}{"PR_kwDNA-jNB9A", "PR_kwDNA-jNB9E", "U_kgAB"}, vars["ids"]) + assert.Equal(t, []interface{}{"PR_kwDNA-jNB9A", "U_kgAB"}, vars["ids"]) }), ) }, wantOut: []*Session{ { - ID: "sess1", - Name: "Build artifacts", - UserID: 1, - AgentID: 2, - Logs: "", - State: "completed", - OwnerID: 10, - RepoID: 1000, - ResourceType: "pull", - ResourceID: 2000, - CreatedAt: sampleDate, - PremiumRequests: 0.1, + ID: "sess1", + Name: "Build artifacts", + UserID: 1, + State: "completed", + ResourceType: "pull", + ResourceID: 2000, PullRequest: &api.PullRequest{ ID: "PR_node", FullDatabaseID: "2000", @@ -1504,26 +1456,20 @@ func TestListSessionsByResourceID(t *testing.T) { }, }, { - ID: "sess2", - Name: "Build artifacts", - UserID: 1, - AgentID: 2, - Logs: "", - State: "completed", - OwnerID: 10, - RepoID: 1000, - ResourceType: "pull", - ResourceID: 2001, - CreatedAt: sampleDate, - PremiumRequests: 0.1, + ID: "sess2", + Name: "Build artifacts", + UserID: 1, + State: "completed", + ResourceType: "pull", + ResourceID: 2000, PullRequest: &api.PullRequest{ ID: "PR_node", - FullDatabaseID: "2001", - Number: 43, + FullDatabaseID: "2000", + Number: 42, Title: "Improve docs", State: "OPEN", IsDraft: true, - URL: "https://github.com/OWNER/REPO/pull/43", + URL: "https://github.com/OWNER/REPO/pull/42", Body: "", CreatedAt: sampleDate, UpdatedAt: sampleDate, @@ -1544,13 +1490,7 @@ func TestListSessionsByResourceID(t *testing.T) { limit: 10, httpStubs: func(t *testing.T, reg *httpmock.Registry) { reg.Register( - httpmock.WithHost( - httpmock.QueryMatcher("GET", "agents/sessions/resource/pull/999", url.Values{ - "page_number": {"1"}, - "page_size": {"50"}, - }), - "api.githubcopilot.com", - ), + httpmock.WithHost(httpmock.REST("GET", "agents/resource/pull/999"), "api.githubcopilot.com"), httpmock.StatusStringResponse(500, "{}"), ) }, @@ -1560,13 +1500,7 @@ func TestListSessionsByResourceID(t *testing.T) { limit: 10, httpStubs: func(t *testing.T, reg *httpmock.Registry) { reg.Register( - httpmock.WithHost( - httpmock.QueryMatcher("GET", "agents/sessions/resource/pull/999", url.Values{ - "page_number": {"1"}, - "page_size": {"50"}, - }), - "api.githubcopilot.com", - ), + httpmock.WithHost(httpmock.REST("GET", "agents/resource/pull/999"), "api.githubcopilot.com"), httpmock.StringResponse(heredoc.Docf(` { "sessions": [ diff --git a/pkg/cmd/agent-task/view/view.go b/pkg/cmd/agent-task/view/view.go index b002c2a3d..60fa892bf 100644 --- a/pkg/cmd/agent-task/view/view.go +++ b/pkg/cmd/agent-task/view/view.go @@ -224,10 +224,6 @@ func viewRun(opts *ViewOptions) error { prURL = pr.URL } - // TODO(babakks): currently we just fetch a pre-defined number of - // matching sessions to avoid hitting the API too many times, but it's - // technically possible for a PR to be associated with lots of sessions - // (i.e. above our selected limit). sessions, err := capiClient.ListSessionsByResourceID(ctx, "pull", prID, defaultLimit) if err != nil { return fmt.Errorf("failed to list sessions for pull request: %w", err)