fix: re-create the message session on refresh 404 instead of failing

session_client.go

@@ -97,6 +97,41 @@ func (c *MessageSessionClient) refreshMessageSession(ctx context.Context, expire
 	return nil
 }
 
+// refreshOrRecreateSession refreshes the message session; when the service no
+// longer knows the session (404 on the refresh), it creates a fresh session in
+// place instead of failing. Without this, a broker-side session eviction is
+// fatal to the caller on every token refresh even though a new session would
+// work fine — and the caller's message cursor (lastMessageId, passed per
+// request) survives re-creation. Only the 404 path re-creates; every other
+// refresh failure surfaces exactly as before.
+func (c *MessageSessionClient) refreshOrRecreateSession(ctx context.Context, expiredSession RunnerScaleSetSession) error {
+	refreshErr := c.refreshMessageSession(ctx, expiredSession)
+	if refreshErr == nil {
+		return nil
+	}
+	if !errors.Is(refreshErr, NotFoundError) {
+		return refreshErr
+	}
+	if createErr := c.recreateMessageSession(ctx, expiredSession); createErr != nil {
+		return fmt.Errorf("%w (session was gone; re-create also failed: %v)", refreshErr, createErr)
+	}
+	return nil
+}
+
+// recreateMessageSession mirrors refreshMessageSession's locking and
+// double-check: if another goroutine already replaced the session, do nothing.
+func (c *MessageSessionClient) recreateMessageSession(ctx context.Context, expiredSession RunnerScaleSetSession) error {
+	c.refreshMu.Lock()
+	defer c.refreshMu.Unlock()
+
+	session := c.Session()
+	if session.SessionID != expiredSession.SessionID {
+		return nil
+	}
+
+	return c.createMessageSession(ctx)
+}
+
 // GetMessage fetches a message from the runner scale set message queue. If there are no messages available, it returns (nil, nil).
 // Unless a message is deleted after being processed (using DeleteMessage), it will be returned again in subsequent calls.
 // If the current session token is expired, it refreshes the session and tries one more time.
@@ -116,7 +151,7 @@ func (c *MessageSessionClient) GetMessage(ctx context.Context, lastMessageID int
 		return nil, fmt.Errorf("failed to get next message: %w", err)
 	}
 
-	if err := c.refreshMessageSession(ctx, session); err != nil {
+	if err := c.refreshOrRecreateSession(ctx, session); err != nil {
 		return nil, fmt.Errorf("failed to refresh message session: %w", err)
 	}
 
@@ -189,7 +224,7 @@ func (c *MessageSessionClient) DeleteMessage(ctx context.Context, messageID int)
 		return fmt.Errorf("failed to delete message: %w", err)
 	}
 
-	if err := c.refreshMessageSession(ctx, session); err != nil {
+	if err := c.refreshOrRecreateSession(ctx, session); err != nil {
 		return fmt.Errorf("failed to refresh message session: %w", err)
 	}
 
@@ -252,7 +287,7 @@ func (c *MessageSessionClient) AcquireJobs(ctx context.Context, requestIDs []int
 		return nil, fmt.Errorf("failed to acquire jobs: %w", err)
 	}
 
-	if err := c.refreshMessageSession(ctx, session); err != nil {
+	if err := c.refreshOrRecreateSession(ctx, session); err != nil {
 		return nil, fmt.Errorf("failed to refresh message session: %w", err)
 	}
 
@@ -310,6 +345,13 @@ func (c *MessageSessionClient) doSessionRequest(ctx context.Context, method, pat
 	defer resp.Body.Close()
 
 	if resp.StatusCode != expectedResponseStatusCode {
+		// A 404 on a session request means the service no longer knows the
+		// session (or the scale set behind it). Surface the typed sentinel so
+		// callers can distinguish "session evicted, safe to re-create" from
+		// genuine failures.
+		if resp.StatusCode == http.StatusNotFound {
+			return newRequestResponseError(req, resp, NotFoundError)
+		}
 		return newRequestResponseError(req, resp, fmt.Errorf("unexpected status code %s", resp.Status))
 	}
 

session_client_test.go

@@ -952,3 +952,112 @@ func TestAcquireJobs(t *testing.T) {
 		assert.Empty(t, got)
 	})
 }
+
+// A broker-side session eviction makes every refresh 404 even though the scale
+// set still exists and a new session would work. The client must re-create the
+// session in place — preserving the caller's message cursor — rather than
+// surface a fatal error on every token refresh.
+func TestGetMessageSessionRecreate(t *testing.T) {
+	ctx := context.Background()
+	auth := actionsAuth{token: "token"}
+	response := []byte(`{"messageId":7,"messageType":"RunnerScaleSetJobMessages"}`)
+
+	t.Run("refresh 404 recreates the session and the retried get succeeds", func(t *testing.T) {
+		var creates, refreshes, gets int
+		var lastMessageIDSeen string
+		var handleSessionRequest http.HandlerFunc
+		s := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			switch {
+			case strings.Contains(r.URL.Path, "/sessions/") && r.Method == http.MethodPatch:
+				refreshes++
+				w.WriteHeader(http.StatusNotFound)
+				w.Write([]byte(`{"typeName":"RunnerAdminException","message":"runner referenced a session ID that doesn't exist in redis"}`))
+			case strings.HasSuffix(r.URL.Path, "sessions") && r.Method == http.MethodPost:
+				creates++
+				handleSessionRequest(w, r)
+			default: // message queue GET
+				gets++
+				if gets == 1 {
+					w.WriteHeader(http.StatusUnauthorized)
+					return
+				}
+				lastMessageIDSeen = r.URL.Query().Get("lastMessageId")
+				w.Write(response)
+			}
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, s.testRunnerScaleSetSession())
+
+		client, err := newClient(testSystemInfo, s.configURLForOrg("my-org"), auth)
+		require.NoError(t, err)
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+		creates = 0 // ignore the initial session creation
+
+		got, err := sessionClient.GetMessage(ctx, 42, 10)
+		require.NoError(t, err)
+		require.NotNil(t, got)
+		assert.Equal(t, 7, got.MessageID)
+		assert.Equal(t, 1, refreshes, "exactly one refresh attempt")
+		assert.Equal(t, 1, creates, "exactly one re-create after the 404")
+		assert.Equal(t, "42", lastMessageIDSeen, "caller-held cursor must survive re-creation")
+	})
+
+	t.Run("non-404 refresh failure does not recreate", func(t *testing.T) {
+		var creates int
+		var handleSessionRequest http.HandlerFunc
+		s := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			switch {
+			case strings.Contains(r.URL.Path, "/sessions/") && r.Method == http.MethodPatch:
+				w.WriteHeader(http.StatusInternalServerError)
+			case strings.HasSuffix(r.URL.Path, "sessions") && r.Method == http.MethodPost:
+				creates++
+				handleSessionRequest(w, r)
+			default:
+				w.WriteHeader(http.StatusUnauthorized)
+			}
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, s.testRunnerScaleSetSession())
+
+		client, err := newClient(testSystemInfo, s.configURLForOrg("my-org"), auth)
+		require.NoError(t, err)
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+		creates = 0
+
+		_, err = sessionClient.GetMessage(ctx, 0, 10)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "failed to refresh message session")
+		assert.Equal(t, 0, creates, "a 500 on refresh must not trigger re-create")
+	})
+
+	t.Run("refresh 404 with failing re-create surfaces both errors", func(t *testing.T) {
+		var creates int
+		var handleSessionRequest http.HandlerFunc
+		s := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			switch {
+			case strings.Contains(r.URL.Path, "/sessions/") && r.Method == http.MethodPatch:
+				w.WriteHeader(http.StatusNotFound)
+			case strings.HasSuffix(r.URL.Path, "sessions") && r.Method == http.MethodPost:
+				creates++
+				if creates == 1 {
+					handleSessionRequest(w, r)
+					return
+				}
+				w.WriteHeader(http.StatusInternalServerError)
+			default:
+				w.WriteHeader(http.StatusUnauthorized)
+			}
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, s.testRunnerScaleSetSession())
+
+		client, err := newClient(testSystemInfo, s.configURLForOrg("my-org"), auth)
+		require.NoError(t, err)
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		_, err = sessionClient.GetMessage(ctx, 0, 10)
+		require.Error(t, err)
+		assert.ErrorIs(t, err, NotFoundError)
+		assert.Contains(t, err.Error(), "re-create also failed")
+	})
+}