Fix logger propagation (#85)

Bumps the actions group with 2 updates: [actions/checkout](https://github.com/actions/checkout) and [golangci/golangci-lint-action](https://github.com/golangci/golangci-lint-action).


Updates `actions/checkout` from 5 to 6
- [Release notes](https://github.com/actions/checkout/releases)
- [Changelog](https://github.com/actions/checkout/blob/main/CHANGELOG.md)
- [Commits](https://github.com/actions/checkout/compare/v5...v6)

Updates `golangci/golangci-lint-action` from 8.0.0 to 9.2.0
- [Release notes](https://github.com/golangci/golangci-lint-action/releases)
- [Commits](https://github.com/golangci/golangci-lint-action/compare/4afd733a84b1f43292c63897423277bb7f4313a9...1e7e51e771db61008b38414a730f564565cf7c20)

---
updated-dependencies:
- dependency-name: actions/checkout
  dependency-version: '6'
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
- dependency-name: golangci/golangci-lint-action
  dependency-version: 9.2.0
  dependency-type: direct:production
  update-type: version-update:semver-major
  dependency-group: actions
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

.devcontainer/Dockerfile

@@ -0,0 +1,3 @@
+FROM mcr.microsoft.com/devcontainers/go:1.25-bookworm
+
+USER vscode

.devcontainer/devcontainer.json

@@ -0,0 +1,33 @@
+{
+  "build": {
+    "dockerfile": "Dockerfile"
+  },
+  "features": {
+    "ghcr.io/devcontainers/features/sshd:1": "latest",
+    "ghcr.io/devcontainers/features/github-cli:1": {},
+    // Node is here only to support Copilot extension
+    "ghcr.io/devcontainers/features/node:1": {}
+  },
+  "hostRequirements": {
+    "cpus": 8,
+    "memory": "16gb"
+  },
+  "customizations": {
+    "vscode": {
+      // Add the IDs of extensions you want installed when the container is created.
+      "extensions": [
+        "golang.Go",
+        "github.copilot"
+      ],
+      // Set *default* container specific settings.json values on container create.
+      "settings": {
+        "go.toolsManagement.checkForUpdates": "local",
+        "go.useLanguageServer": true,
+        "go.gopath": "/go",
+        "gopls": {
+          "formatting.gofumpt": true
+        }
+      }
+    }
+  }
+}

.github/CODEOWNERS

@@ -0,0 +1 @@
+* @actions/actions-runtime @nikola-jokic

.github/dependabot.yml

@@ -0,0 +1,18 @@
+version: 2
+updates:
+  - package-ecosystem: "gomod" # See documentation for possible values
+    directory: "/" # Location of package manifests
+    schedule:
+      interval: "weekly"
+    groups:
+      gomod:
+        patterns:
+          - "*"
+  - package-ecosystem: github-actions
+    directory: "/"
+    schedule:
+      interval: "weekly"
+    groups:
+      actions:
+        patterns:
+          - "*"

.github/workflows/e2e.yaml

@@ -0,0 +1,55 @@
+name: E2E
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+env:
+  E2E_WORKFLOW_TARGET_ORG: "scaleset-canary"
+  E2E_WORKFLOW_TARGET_REPO: "e2e"
+
+jobs:
+  basic-e2e:
+    name: Basic E2E
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    env:
+      E2E_WORKFLOW_TARGET_FILE: "basic.yaml"
+    steps:
+      - uses: actions/checkout@v6
+        with:
+          persist-credentials: true
+      - uses: actions/setup-go@v6
+        with:
+          go-version-file: "go.mod"
+      - name: Get configure token
+        id: config-token
+        uses: peter-murray/workflow-application-token-action@d17e3a9a36850ea89f35db16c1067dd2b68ee343
+        with:
+          application_id: ${{ secrets.E2E_TESTS_ACCESS_CLIENT_ID }}
+          application_private_key: ${{ secrets.E2E_TESTS_ACCESS_PK }}
+          organization: ${{ env.E2E_WORKFLOW_TARGET_ORG }}
+      - name: Run simple test
+        run: |
+          git config --global url."https://x-access-token:${{ secrets.GITHUB_TOKEN }}@github.com/actions/scaleset".insteadOf "https://github.com/actions/scaleset"
+          E2E_SCALESET_NAME="basic-$(date +'%M%S')$(((RANDOM + 100) % 100 + 1))" go test
+        working-directory: examples/dockerscaleset
+        env:
+          GOPRIVATE: github.com/actions/scaleset
+          GONOSUMDB: github.com/actions/scaleset
+          E2E_SCALESET_GITHUB_TOKEN: "${{steps.config-token.outputs.token}}"
+          E2E_WORKFLOW_GITHUB_TOKEN: "${{steps.config-token.outputs.token}}"
+          E2E_SCALESET_URL: "https://github.com/scaleset-canary/e2e"
+          GITHUB_TOKEN: "${{ secrets.GITHUB_TOKEN }}"
+          E2E: "true"

.github/workflows/go.yaml

@@ -19,7 +19,7 @@ jobs:
   fmt:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - uses: actions/setup-go@v6
         with:
           go-version-file: "go.mod"
@@ -29,16 +29,29 @@ jobs:
       - name: Check diff
         run: git diff --exit-code
 
+  mocks:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-go@v6
+        with:
+          go-version-file: "go.mod"
+          cache: false
+      - name: "Run mockery"
+        run: go tool github.com/vektra/mockery/v3
+      - name: Check diff
+        run: git diff --exit-code
+
   lint:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - uses: actions/setup-go@v6
         with:
           go-version-file: "go.mod"
           cache: false
       - name: golangci-lint
-        uses: golangci/golangci-lint-action@4afd733a84b1f43292c63897423277bb7f4313a9
+        uses: golangci/golangci-lint-action@1e7e51e771db61008b38414a730f564565cf7c20
         with:
           only-new-issues: true
           version: v2.5.0
@@ -46,10 +59,10 @@ jobs:
   test:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v5
+      - uses: actions/checkout@v6
       - uses: actions/setup-go@v6
         with:
           go-version-file: "go.mod"
-          cache: false
+          cache: true
       - name: Run tests
-        run: go test ./...
+        run: go test ./... -race

.mockery.yml

@@ -0,0 +1,14 @@
+all: false
+dir: "{{.InterfaceDir}}"
+filename: mocks_test.go
+force-file-write: true
+formatter: goimports
+log-level: info
+structname: "{{.Mock}}{{.InterfaceName}}"
+pkgname: "{{.SrcPackageName}}"
+recursive: true
+template: testify
+packages:
+  github.com/actions/scaleset/listener:
+    config:
+      all: true

CODE_OF_CONDUCT.md

@@ -0,0 +1,74 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+In the interest of fostering an open and welcoming environment, we as
+contributors and maintainers pledge to making participation in our project and
+our community a harassment-free experience for everyone, regardless of age, body
+size, disability, ethnicity, gender identity and expression, level of experience,
+nationality, personal appearance, race, religion, or sexual identity and
+orientation.
+
+## Our Standards
+
+Examples of behavior that contributes to creating a positive environment
+include:
+
+* Using welcoming and inclusive language
+* Being respectful of differing viewpoints and experiences
+* Gracefully accepting constructive criticism
+* Focusing on what is best for the community
+* Showing empathy towards other community members
+
+Examples of unacceptable behavior by participants include:
+
+* The use of sexualized language or imagery and unwelcome sexual attention or
+advances
+* Trolling, insulting/derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or electronic
+  address, without explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Our Responsibilities
+
+Project maintainers are responsible for clarifying the standards of acceptable
+behavior and are expected to take appropriate and fair corrective action in
+response to any instances of unacceptable behavior.
+
+Project maintainers have the right and responsibility to remove, edit, or
+reject comments, commits, code, wiki edits, issues, and other contributions
+that are not aligned to this Code of Conduct, or to ban temporarily or
+permanently any contributor for other behaviors that they deem inappropriate,
+threatening, offensive, or harmful.
+
+## Scope
+
+This Code of Conduct applies both within project spaces and in public spaces
+when an individual is representing the project or its community. Examples of
+representing a project or community include using an official project e-mail
+address, posting via an official social media account, or acting as an appointed
+representative at an online or offline event. Representation of a project may be
+further defined and clarified by project maintainers.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported by contacting the project team at <opensource@github.com>. All
+complaints will be reviewed and investigated and will result in a response that
+is deemed necessary and appropriate to the circumstances. The project team is
+obligated to maintain confidentiality with regard to the reporter of an incident.
+Further details of specific enforcement policies may be posted separately.
+
+Project maintainers who do not follow or enforce the Code of Conduct in good
+faith may face temporary or permanent repercussions as determined by other
+members of the project's leadership.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
+available at [http://contributor-covenant.org/version/1/4][version]
+
+[homepage]: http://contributor-covenant.org
+[version]: http://contributor-covenant.org/version/1/4/

CONTRIBUTING.md

@@ -0,0 +1,41 @@
+## Contributing
+
+[fork]: https://github.com/actions/scaleset/fork
+[pr]: https://github.com/actions/scaleset/compare
+[style]: https://github.com/actions/scaleset/blob/main/.golangci.yaml
+
+Hi there! We're thrilled that you'd like to contribute to this project. Your help is essential for keeping it great.
+
+Contributions to this project are [released](https://help.github.com/articles/github-terms-of-service/#6-contributions-under-repository-license) to the public under the [project's open source license](LICENSE).
+
+Please note that this project is released with a [Contributor Code of Conduct](CODE_OF_CONDUCT.md). By participating in this project you agree to abide by its terms.
+
+## Prerequisites for running and testing code
+
+These are one time installations required to be able to test your changes locally as part of the pull request (PR) submission process.
+
+1. install Go [through download](https://go.dev/doc/install) | [through Homebrew](https://formulae.brew.sh/formula/go)
+1. [install golangci-lint](https://golangci-lint.run/welcome/install/#local-installation)
+
+## Submitting a pull request
+
+1. [Fork][fork] and clone the repository
+1. Make sure the tests pass on your machine: `go test -v ./...`
+1. Make sure linter passes on your machine: `script/lint`
+1. Create a new branch: `git checkout -b my-branch-name`
+1. Make your change, add tests, and make sure the tests and linter still pass
+1. Push to your fork and [submit a pull request][pr]
+1. Pat yourself on the back and wait for your pull request to be reviewed and merged.
+
+Here are a few things you can do that will increase the likelihood of your pull request being accepted:
+
+- Follow the [style guide][style].
+- Write tests.
+- Keep your change as focused as possible. If there are multiple changes you would like to make that are not dependent upon each other, consider submitting them as separate pull requests.
+- Write a [good commit message](http://tbaggery.com/2008/04/19/a-note-about-git-commit-messages.html).
+
+## Resources
+
+- [How to Contribute to Open Source](https://opensource.guide/how-to-contribute/)
+- [Using Pull Requests](https://help.github.com/articles/about-pull-requests/)
+- [GitHub Help](https://help.github.com)

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright GitHub, Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,185 @@
+# GitHub Actions Runner Scale Set Client (Public Preview)
+
+> Status: **Public Preview** – While the API is stable, interfaces and examples in this repository may change.
+
+This repository provides a standalone Go client for the GitHub Actions **Runner Scale Set** APIs. It is extracted from the `actions-runner-controller` project so that platform teams, integrators, and infrastructure providers can build **their own custom autoscaling solutions** for GitHub Actions runners.
+
+You do *not* need to adopt the full controller (and Kubernetes) to take advantage of scale sets. This package contains all the primitives you need: create/update/delete scale sets, generate just‑in‑time (JIT) runner configs, and manage message sessions.
+
+---
+
+## What is a Scale Set?
+
+A runner scale set is a group of self-hosted runners that autoscales based on workflow demand. Here's how it works:
+
+1. **Registration**: You create a scale set with a name, which also serves as the label workflows use to target it (e.g., `runs-on: my-scale-set`). Multiple labels can be assigned per scale set. Like regular self-hosted runners, scale sets can be registered at the repository, organization, or enterprise level.
+2. **Polling**: Your scale set client continuously polls the API, reporting its maximum capacity (how many runners it can produce).
+3. **Job matching**: GitHub matches jobs to your scale set based on the label and runner group policies, just like regular self-hosted runners.
+4. **Scaling signal**: The API responds with how many runners your scale set needs online (`statistics.TotalAssignedJobs`).
+5. **Runner provisioning**: Your client creates or maintains enough runners to meet demand. Runners can be created just-in-time as jobs arrive, or pre-provisioned ahead of demand to reduce latency.
+6. **Job assignment**: GitHub assigns a pending job to any idle runner in the scale set.
+
+Runners in a scale set are ephemeral by default: each runner executes one job and is then removed. This ensures a clean environment for every job.
+
+---
+
+## High-Level Flow
+
+1. Create a `Client` with either a GitHub App credential (recommended) or a PAT.
+2. Create a Runner Scale Set with a name.
+3. Start a message session and poll for scaling events. The `listener` package handles this for you.
+4. When the API indicates runners are needed:
+   - Call `GenerateJitRunnerConfig` to get a JIT config for a new runner.
+   - Start your runner (process, container, VM, etc.) with the JIT config.
+5. Idle runners are assigned jobs automatically by GitHub.
+
+You can also pre-provision runners before jobs arrive to reduce startup latency. See [`examples/dockerscaleset`](./examples/dockerscaleset) for a complete example that supports both `minRunners` (pre-provisioned) and just-in-time scaling.
+
+---
+
+## Autoscaling
+
+Use `statistics.TotalAssignedJobs` from each message response to determine how many runners your scale set needs online. This value represents the total number of jobs assigned to your scale set, including both jobs waiting for a runner and jobs already running (`TotalAssignedJobs >= TotalRunningJobs`).
+
+Do not count individual job messages (`JobAssigned`, `JobStarted`, `JobCompleted`) in the response body to determine scaling:
+
+- Responses contain at most 50 messages. Large backlogs will be truncated.
+- The `statistics` field is always current and reflects the true state of your scale set.
+
+When polling for messages, include your scale set's maximum capacity via the `maxCapacity` parameter (sent as the `X-ScaleSetMaxCapacity` header). This allows the backend to assign jobs accurately and avoid creating backlogs your scale set cannot fulfill.
+
+Here's a simplified polling loop:
+
+```go
+var lastMessageID int
+for {
+    msg, err := client.GetMessage(ctx, lastMessageID, maxCapacity)
+    if err != nil {
+        return err
+    }
+
+    if msg == nil {
+        // No messages available (202 response), poll again
+        continue
+    }
+
+    lastMessageID = msg.MessageID
+
+    // Scale based on statistics, not message counts
+    desiredRunners := msg.Statistics.TotalAssignedJobs
+    scaleToDesired(desiredRunners)
+
+    // Acknowledge the message
+    if err := client.DeleteMessage(ctx, msg.MessageID); err != nil {
+        return err
+    }
+}
+```
+
+The `listener` package provides a ready-to-use implementation of this pattern, handling session management, polling, and acknowledgment. See [`listener/listener.go`](./listener/listener.go).
+
+### Job lifecycle messages
+
+Individual job messages (`JobStarted`, `JobCompleted`, etc.) are useful for purposes beyond scaling. For example, [actions-runner-controller](https://github.com/actions/actions-runner-controller) uses `JobStarted` to mark runner pods as busy, preventing premature cleanup during scale-down. These messages can also be used for metrics or logging.
+
+See [`types.go`](./types.go) for payload definitions.
+
+---
+
+## How the Message API Works
+
+### Long Polling
+
+`GetMessage` uses long polling:
+
+1. If messages are available, they are returned immediately.
+2. Otherwise, the request blocks for up to ~50 seconds.
+3. If no messages arrive, a 202 response is returned (`nil, nil` in the Go client).
+
+Poll again immediately after handling each response.
+
+### Message Acknowledgment
+
+Call `DeleteMessage` after processing a message. This acts as an acknowledgment:
+
+- Unacknowledged messages are redelivered on the next poll.
+- This prevents message loss if your client crashes mid-processing.
+
+### Message ID Tracking
+
+Pass the ID of the last processed message to `GetMessage`. Omitting this (or passing 0) returns the first available message, potentially causing reprocessing.
+
+### Job Reassignment
+
+Jobs may appear multiple times as `JobAssigned` followed by `JobCompleted` (with `result: "canceled"`). This occurs when a job is assigned to your scale set but not acquired by a runner in time—GitHub cancels the assignment and requeues the job. This can happen up to 3 times with incremental delays.
+
+Each attempt generates new messages, but they represent the same workflow job. This is why `statistics.TotalAssignedJobs` is the correct scaling metric: it reflects the current state, not the message history.
+
+---
+
+## Getting Started
+
+```bash
+go get github.com/actions/scaleset@latest
+```
+
+Import:
+
+```go
+import "github.com/actions/scaleset"
+```
+
+### Using Without Go Experience
+
+If you are not a Go developer, you can still:
+
+- Treat this repo as reference documentation to design an API integration in another language.
+- Vendor the code and compile a minimal binary that exposes a simpler CLI.
+- Use the example CLI (`examples/dockerscaleset`) as inspiration—its flags show required inputs.
+- Copilot can also help you translate this Go code into your language of choice.
+
+---
+
+## Authentication
+
+Two options:
+
+1. **GitHub App (preferred):** Stronger scoping & rotation. Provide: `ClientID`, `InstallationID`, `PrivateKey`.
+2. **PAT (personal access token):** Simpler but broader scoped.
+
+The client automatically exchanges credentials for a registration token + admin token behind the scenes and refreshes them before expiry.
+
+You can find more details on required permissions in the [GitHub Docs](https://docs.github.com/en/actions/tutorials/use-actions-runner-controller/authenticate-to-the-api).
+
+GitHub Enterprise Server (GHES) is supported out of the box—just use your GHES URL when creating the client.
+
+---
+
+## Security Notes
+
+- Always prefer GitHub App credentials; rotate PATs if you must use them.
+- Treat JIT configs as secrets until consumed.
+
+---
+
+## Requirements
+
+- Go 1.25 or later
+
+---
+
+## License
+
+This project is licensed under the terms of the MIT open source license. Please refer to [LICENSE](./LICENSE) for the full terms.
+
+---
+
+## Maintainers
+
+See [CODEOWNERS](./.github/CODEOWNERS) for the list of maintainers.
+
+---
+
+## Support
+
+Please refer to [SUPPORT.md](./SUPPORT.md) for information on how to get help with this project.

SECURITY.md

@@ -0,0 +1,31 @@
+Thanks for helping make GitHub safe for everyone.
+
+# Security
+
+GitHub takes the security of our software products and services seriously, including all of the open source code repositories managed through our GitHub organizations, such as [GitHub](https://github.com/GitHub).
+
+Even though [open source repositories are outside of the scope of our bug bounty program](https://bounty.github.com/index.html#scope) and therefore not eligible for bounty rewards, we will ensure that your finding gets passed along to the appropriate maintainers for remediation. 
+
+## Reporting Security Issues
+
+If you believe you have found a security vulnerability in any GitHub-owned repository, please report it to us through coordinated disclosure.
+
+**Please do not report security vulnerabilities through public GitHub issues, discussions, or pull requests.**
+
+Instead, please send an email to opensource-security[@]github.com.
+
+Please include as much of the information listed below as you can to help us better understand and resolve the issue:
+
+  * The type of issue (e.g., buffer overflow, SQL injection, or cross-site scripting)
+  * Full paths of source file(s) related to the manifestation of the issue
+  * The location of the affected source code (tag/branch/commit or direct URL)
+  * Any special configuration required to reproduce the issue
+  * Step-by-step instructions to reproduce the issue
+  * Proof-of-concept or exploit code (if possible)
+  * Impact of the issue, including how an attacker might exploit the issue
+
+This information will help us triage your report more quickly.
+
+## Policy
+
+See [GitHub's Safe Harbor Policy](https://docs.github.com/en/site-policy/security-policies/github-bug-bounty-program-legal-safe-harbor#1-safe-harbor-terms)

SUPPORT.md

@@ -0,0 +1,13 @@
+# Support 
+
+## How to file issues and get help
+
+This project uses GitHub issues to track bugs and feature requests. Please search the existing issues before filing new issues to avoid duplicates. For new issues, file your bug or feature request as a new issue.
+
+For help or questions about using this project, please open a new issue.
+
+**actions/scaleset** is under active development and maintained by GitHub staff **AND THE COMMUNITY**. We will do our best to respond to support, feature requests, and community questions in a timely manner.
+
+## GitHub Support Policy
+
+Support for this project is limited to the resources listed above.

common_client.go

@@ -0,0 +1,229 @@
+package scaleset
+
+import (
+	"bytes"
+	"crypto/tls"
+	"crypto/x509"
+	"encoding/json"
+	"fmt"
+	"io"
+	"log/slog"
+	"net/http"
+	"time"
+
+	"github.com/hashicorp/go-retryablehttp"
+)
+
+const (
+	headerActionsActivityID = "ActivityId"
+	headerGitHubRequestID   = "X-GitHub-Request-Id"
+)
+
+type commonClient struct {
+	httpClient *http.Client
+
+	systemInfo SystemInfo // never set directly, use setSystemInfoUnlocked
+
+	userAgent string
+
+	httpClientOption
+}
+
+func newCommonClient(systemInfo SystemInfo, httpClientOption httpClientOption) *commonClient {
+	c := &commonClient{
+		httpClientOption: httpClientOption,
+	}
+	c.setSystemInfo(systemInfo)
+
+	retryableHTTPClient, err := httpClientOption.newRetryableHTTPClient()
+	if err != nil {
+		panic(fmt.Sprintf("failed to create retryable HTTP client: %v", err))
+	}
+	c.httpClient = retryableHTTPClient.StandardClient()
+
+	return c
+}
+
+func (c *commonClient) newRetryableHTTPClient() (*retryablehttp.Client, error) {
+	return c.httpClientOption.newRetryableHTTPClient()
+}
+
+func (c *commonClient) do(req *http.Request) (*http.Response, error) {
+	return sendRequest(c.httpClient, req)
+}
+
+// sendRequest ensures that the request is sent and the response body is fully read and closed.
+// It trims the BOM when present in the response body.
+//
+// Make sure to use this function instead of http.Client.Do directly to avoid issues.
+func sendRequest(c *http.Client, req *http.Request) (*http.Response, error) {
+	resp, err := c.Do(req)
+	if err != nil {
+		return nil, newRequestResponseError(req, resp, fmt.Errorf("failed to send request: %w", err))
+	}
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, newRequestResponseError(req, resp, fmt.Errorf("failed to read the response body: %w", err))
+	}
+	if err := resp.Body.Close(); err != nil {
+		return nil, newRequestResponseError(req, resp, fmt.Errorf("failed to close the response body: %w", err))
+	}
+
+	body = trimByteOrderMark(body)
+	resp.Body = io.NopCloser(bytes.NewReader(body))
+	return resp, nil
+}
+
+type httpClientOption struct {
+	logger *slog.Logger
+
+	// Options for built-in retryable HTTP client.
+	// Ignored if a custom retryable HTTP client is provided via WithRetryableHTTPClint.
+	retryMax     int
+	retryWaitMax time.Duration
+
+	// fields added to the transport if specified
+	rootCAs               *x509.CertPool
+	tlsInsecureSkipVerify bool
+	proxyFunc             ProxyFunc
+	timeout               time.Duration
+
+	retryableHTTPClient *retryablehttp.Client
+}
+
+func (o *httpClientOption) defaults() {
+	if o.logger == nil {
+		o.logger = slog.New(slog.DiscardHandler)
+	}
+	if o.retryMax == 0 {
+		o.retryMax = 4
+	}
+	if o.retryWaitMax == 0 {
+		o.retryWaitMax = 30 * time.Second
+	}
+	if o.timeout == 0 {
+		o.timeout = 5 * time.Minute
+	}
+}
+
+func (o *httpClientOption) newRetryableHTTPClient() (*retryablehttp.Client, error) {
+	var retryClient *retryablehttp.Client
+	if o.retryableHTTPClient != nil {
+		retryClient = o.retryableHTTPClient
+	} else {
+		retryClient = retryablehttp.NewClient()
+		retryClient.RetryMax = o.retryMax
+		retryClient.RetryWaitMax = o.retryWaitMax
+	}
+
+	if retryClient.HTTPClient.Timeout == 0 {
+		retryClient.HTTPClient.Timeout = o.timeout
+	}
+
+	retryClient.Logger = o.logger
+
+	transport, ok := retryClient.HTTPClient.Transport.(*http.Transport)
+	if !ok {
+		// this should always be true, because retryablehttp.NewClient() uses
+		// cleanhttp.DefaultPooledTransport()
+		return nil, fmt.Errorf("failed to get http transport from retryablehttp client")
+	}
+	if transport.TLSClientConfig == nil {
+		transport.TLSClientConfig = &tls.Config{}
+	}
+
+	if o.rootCAs != nil {
+		transport.TLSClientConfig.RootCAs = o.rootCAs
+	}
+
+	if o.tlsInsecureSkipVerify {
+		transport.TLSClientConfig.InsecureSkipVerify = true
+	}
+
+	if o.proxyFunc != nil {
+		transport.Proxy = o.proxyFunc
+	}
+
+	retryClient.HTTPClient.Transport = transport
+
+	return retryClient, nil
+}
+
+func (c *commonClient) setSystemInfo(info SystemInfo) {
+	c.systemInfo = info
+	c.setUserAgent()
+}
+
+func (c *commonClient) setUserAgent() {
+	b, _ := json.Marshal(userAgent{
+		SystemInfo:     c.systemInfo,
+		BuildVersion:   buildInfo.version,
+		BuildCommitSHA: buildInfo.commitSHA,
+		Kind:           "scaleset",
+	})
+	c.userAgent = string(b)
+}
+
+// HTTPOption defines a functional option for configuring the Client.
+type HTTPOption func(*httpClientOption)
+
+// WithRetryableHTTPClint allows users to provide a custom retryable HTTP client.
+// If not set, a default client will be used with the specified retry and timeout settings.
+func WithRetryableHTTPClint(client *retryablehttp.Client) HTTPOption {
+	return func(c *httpClientOption) {
+		c.retryableHTTPClient = client
+	}
+}
+
+// WithLogger sets a custom logger for the Client.
+// If nil is passed, a discard logger will be used.
+func WithLogger(logger *slog.Logger) HTTPOption {
+	return func(c *httpClientOption) {
+		if logger == nil {
+			logger = slog.New(slog.DiscardHandler)
+		}
+		c.logger = logger
+	}
+}
+
+// WithRetryMax sets the maximum number of retries for the Client.
+func WithRetryMax(retryMax int) HTTPOption {
+	return func(c *httpClientOption) {
+		c.retryMax = retryMax
+	}
+}
+
+// WithRetryWaitMax sets the maximum wait time between retries for the Client.
+func WithRetryWaitMax(retryWaitMax time.Duration) HTTPOption {
+	return func(c *httpClientOption) {
+		c.retryWaitMax = retryWaitMax
+	}
+}
+
+// WithRootCAs sets custom root certificate authorities for the Client.
+func WithRootCAs(rootCAs *x509.CertPool) HTTPOption {
+	return func(c *httpClientOption) {
+		c.rootCAs = rootCAs
+	}
+}
+
+// WithoutTLSVerify disables TLS certificate verification for the Client.
+func WithoutTLSVerify() HTTPOption {
+	return func(c *httpClientOption) {
+		c.tlsInsecureSkipVerify = true
+	}
+}
+
+// WithProxy sets a custom proxy function for the Client.
+func WithProxy(proxyFunc ProxyFunc) HTTPOption {
+	return func(c *httpClientOption) {
+		c.proxyFunc = proxyFunc
+	}
+}
+
+// WithTimeout sets a timeout for the Client.
+func WithTimeout(duration time.Duration) HTTPOption {
+	return func(c *httpClientOption) {
+		c.timeout = duration
+	}
+}

common_client_test.go

@@ -0,0 +1,296 @@
+package scaleset
+
+import (
+	"encoding/json"
+	"io"
+	"log/slog"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+	"time"
+
+	"github.com/actions/scaleset/internal/testserver"
+	"github.com/hashicorp/go-retryablehttp"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/net/http/httpproxy"
+)
+
+func defaultHTTPClientOption() httpClientOption {
+	var opt httpClientOption
+	opt.defaults()
+	return opt
+}
+
+func TestClient_Do(t *testing.T) {
+	t.Run("trims byte order mark from response if present", func(t *testing.T) {
+		t.Run("when there is no body", func(t *testing.T) {
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			}))
+			defer server.Close()
+
+			client := newCommonClient(
+				testSystemInfo,
+				defaultHTTPClientOption(),
+			)
+
+			req, err := http.NewRequest("GET", server.URL, nil)
+			require.NoError(t, err)
+
+			resp, err := client.do(req)
+			require.NoError(t, err)
+
+			body, err := io.ReadAll(resp.Body)
+			require.NoError(t, err)
+			assert.Empty(t, string(body))
+		})
+
+		responses := []string{
+			"\xef\xbb\xbf{\"foo\":\"bar\"}",
+			"{\"foo\":\"bar\"}",
+		}
+
+		for _, response := range responses {
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("Content-Type", "application/json")
+				w.Write([]byte(response))
+			}))
+			defer server.Close()
+
+			client := newCommonClient(
+				testSystemInfo,
+				defaultHTTPClientOption(),
+			)
+
+			req, err := http.NewRequest("GET", server.URL, nil)
+			require.NoError(t, err)
+
+			resp, err := client.do(req)
+			require.NoError(t, err)
+
+			body, err := io.ReadAll(resp.Body)
+			require.NoError(t, err)
+			assert.Equal(t, "{\"foo\":\"bar\"}", string(body))
+		}
+	})
+}
+
+func TestClientProxy(t *testing.T) {
+	serverCalled := false
+
+	proxy := testserver.New(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		serverCalled = true
+	}))
+
+	proxyConfig := &httpproxy.Config{
+		HTTPProxy: proxy.URL,
+	}
+	proxyFunc := func(req *http.Request) (*url.URL, error) {
+		return proxyConfig.ProxyFunc()(req.URL)
+	}
+
+	opts := defaultHTTPClientOption()
+	WithProxy(proxyFunc)(&opts)
+
+	client := newCommonClient(
+		testSystemInfo,
+		opts,
+	)
+
+	req, err := http.NewRequest(http.MethodGet, "http://example.com", nil)
+	require.NoError(t, err)
+
+	_, err = client.do(req)
+	require.NoError(t, err)
+
+	assert.True(t, serverCalled)
+}
+
+func TestUserAgent(t *testing.T) {
+	version, sha := detectModuleVersionAndCommit()
+	userAgentInfo := SystemInfo{
+		System:     "actions-runner-controller",
+		Version:    "0.1.0",
+		CommitSHA:  "1234567890abcdef",
+		ScaleSetID: 10,
+		Subsystem:  "test",
+	}
+
+	client := newCommonClient(
+		testSystemInfo,
+		defaultHTTPClientOption(),
+	)
+
+	got := client.userAgent
+	wantInfo := userAgent{
+		SystemInfo:     testSystemInfo,
+		BuildCommitSHA: sha,
+		BuildVersion:   version,
+		Kind:           "scaleset",
+	}
+	b, err := json.Marshal(wantInfo)
+	require.NoError(t, err, "failed to marshal expected user agent")
+	want := string(b)
+
+	assert.Equal(t, want, got)
+
+	client.setSystemInfo(SystemInfo{
+		System:     "actions-runner-controller",
+		Version:    "0.1.0",
+		CommitSHA:  "1234567890abcdef",
+		ScaleSetID: 10,
+		Subsystem:  "test",
+	})
+
+	got = client.userAgent
+	wantInfo = userAgent{
+		SystemInfo:     userAgentInfo,
+		BuildCommitSHA: sha,
+		BuildVersion:   version,
+		Kind:           "scaleset",
+	}
+	b, err = json.Marshal(wantInfo)
+	require.NoError(t, err, "failed to marshal expected user agent after SetSystemInfo")
+	want = string(b)
+
+	assert.Equal(t, want, got)
+}
+
+func TestWithLogger(t *testing.T) {
+	newJSONHandler := func() slog.Handler {
+		return slog.NewJSONHandler(
+			io.Discard,
+			&slog.HandlerOptions{
+				AddSource: true,
+				Level:     slog.LevelError,
+			},
+		)
+	}
+	t.Run("WithLogger(nil) sets a discard logger on raw httpClientOption", func(t *testing.T) {
+		opts := httpClientOption{}
+		WithLogger(nil)(&opts)
+		require.NotNil(t, opts.logger, "WithLogger(nil) should set a discard logger, not leave it nil")
+		assert.Equal(t, slog.DiscardHandler, opts.logger.Handler(), "WithLogger(nil) should set a discard logger handler")
+	})
+
+	t.Run("WithLogger(customLogger) assigns the provided logger", func(t *testing.T) {
+		handler := newJSONHandler()
+		customLogger := slog.New(handler)
+		opts := httpClientOption{}
+		WithLogger(customLogger)(&opts)
+		require.Equal(t, customLogger, opts.logger, "WithLogger should assign the provided logger")
+		assert.Equal(t, handler, opts.logger.Handler(), "WithLogger should set the provided logger handler")
+	})
+
+	t.Run("WithLogger(nil) propagates discard logger to retryable HTTP client", func(t *testing.T) {
+		opts := httpClientOption{}
+		WithLogger(nil)(&opts)
+		client, err := opts.newRetryableHTTPClient()
+		require.NoError(t, err)
+		assert.NotNil(t, client.Logger, "retryable client should have logger set from WithLogger(nil)")
+
+		logger, ok := client.Logger.(*slog.Logger)
+		require.True(t, ok, "retryable client logger should be a *slog.Logger")
+		assert.Same(t, opts.logger, logger, "retryable client logger should be the same logger set by WithLogger(nil)")
+		assert.Equal(t, slog.DiscardHandler, logger.Handler(), "retryable client logger should be a discard logger from WithLogger(nil)")
+	})
+
+	t.Run("WithLogger(customLogger) propagates custom logger to retryable HTTP client", func(t *testing.T) {
+		handler := newJSONHandler()
+		customLogger := slog.New(handler)
+		opts := httpClientOption{}
+		WithLogger(customLogger)(&opts)
+		client, err := opts.newRetryableHTTPClient()
+		require.NoError(t, err)
+		assert.NotNil(t, client.Logger, "retryable client should have logger set")
+		logger, ok := client.Logger.(*slog.Logger)
+		require.True(t, ok, "retryable client logger should be a *slog.Logger")
+		assert.Equal(t, handler, logger.Handler(), "retryable client logger should be the custom logger from WithLogger")
+	})
+}
+
+// TestWithRetryableHTTPClient verifies that a custom retryable HTTP client
+// provided via WithRetryableHTTPClient is actually used instead of the built-in one
+func TestWithRetryableHTTPClient(t *testing.T) {
+	t.Run("uses custom retryable client instead of built-in", func(t *testing.T) {
+		attemptCount := 0
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			attemptCount++
+			if attemptCount == 1 {
+				w.WriteHeader(http.StatusServiceUnavailable)
+				return
+			}
+			w.WriteHeader(http.StatusOK)
+			w.Write([]byte(`{"result": "success"}`))
+		}))
+		defer server.Close()
+
+		// Create a custom retryable HTTP client with specific retry configuration
+		customRetryClient := retryablehttp.NewClient()
+		customRetryClient.RetryMax = 3
+		customRetryClient.RetryWaitMax = 10 * time.Millisecond
+
+		// Create options with the custom retryable client
+		opts := defaultHTTPClientOption()
+		WithRetryableHTTPClint(customRetryClient)(&opts)
+
+		// Verify that the custom client is set in options
+		assert.NotNil(t, opts.retryableHTTPClient)
+		assert.Equal(t, customRetryClient, opts.retryableHTTPClient)
+
+		// Create the common client with custom retryable client
+		client := newCommonClient(testSystemInfo, opts)
+
+		// Make a request that will trigger a retry
+		req, err := http.NewRequest("GET", server.URL, nil)
+		require.NoError(t, err)
+
+		resp, err := client.do(req)
+		require.NoError(t, err)
+
+		// Should succeed after retry
+		assert.Equal(t, http.StatusOK, resp.StatusCode)
+		assert.Equal(t, 2, attemptCount)
+
+		// Verify that the client used is the custom one by checking newRetryableHTTPClient
+		retrievedRetryClient, err := client.newRetryableHTTPClient()
+		require.NoError(t, err)
+		assert.Equal(t, customRetryClient, retrievedRetryClient, "should return the custom retryable client")
+	})
+
+	t.Run("respects custom client's retry configuration over built-in defaults", func(t *testing.T) {
+		attemptCount := 0
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			attemptCount++
+			w.WriteHeader(http.StatusServiceUnavailable)
+		}))
+		defer server.Close()
+
+		// Create custom client with limited retries
+		customRetryClient := retryablehttp.NewClient()
+		customRetryClient.RetryMax = 1 // Only 1 retry (2 total attempts)
+		customRetryClient.RetryWaitMax = 5 * time.Millisecond
+
+		opts := defaultHTTPClientOption()
+		WithRetryableHTTPClint(customRetryClient)(&opts)
+
+		client := newCommonClient(testSystemInfo, opts)
+
+		req, err := http.NewRequest("GET", server.URL, nil)
+		require.NoError(t, err)
+
+		resp, err := client.do(req)
+		// When all retries are exhausted with a retryable error, the client gives up
+		// and an error is returned
+		if err != nil {
+			// Expected: request failed after exhausting retries
+			assert.Contains(t, err.Error(), "giving up after 2 attempt(s)")
+		} else {
+			// Or the final response is returned
+			assert.Equal(t, http.StatusServiceUnavailable, resp.StatusCode)
+		}
+		// Should have tried 1 initial + 1 retry = 2 times total
+		assert.Equal(t, 2, attemptCount)
+	})
+}

config.go

@@ -9,37 +9,37 @@ import (
 
 var ErrInvalidGitHubConfigURL = fmt.Errorf("invalid config URL, should point to an enterprise, org, or repository")
 
-type GitHubScope int
+type gitHubScope int
 
 const (
-	GitHubScopeUnknown GitHubScope = iota
-	GitHubScopeEnterprise
-	GitHubScopeOrganization
-	GitHubScopeRepository
+	gitHubScopeUnknown gitHubScope = iota
+	gitHubScopeEnterprise
+	gitHubScopeOrganization
+	gitHubScopeRepository
 )
 
-type GitHubConfig struct {
-	ConfigURL *url.URL
-	Scope     GitHubScope
+type gitHubConfig struct {
+	configURL *url.URL
+	scope     gitHubScope
 
-	Enterprise   string
-	Organization string
-	Repository   string
+	enterprise   string
+	organization string
+	repository   string
 
-	IsHosted bool
+	isHosted bool
 }
 
-func ParseGitHubConfigFromURL(in string) (*GitHubConfig, error) {
+func parseGitHubConfigFromURL(in string) (*gitHubConfig, error) {
 	u, err := url.Parse(strings.Trim(in, "/"))
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to parse URL: %w", err)
 	}
 
 	isHosted := isHostedGitHubURL(u)
 
-	configURL := &GitHubConfig{
-		ConfigURL: u,
-		IsHosted:  isHosted,
+	configURL := &gitHubConfig{
+		configURL: u,
+		isHosted:  isHosted,
 	}
 
 	invalidURLError := fmt.Errorf("%q: %w", u.String(), ErrInvalidGitHubConfigURL)
@@ -52,19 +52,19 @@ func ParseGitHubConfigFromURL(in string) (*GitHubConfig, error) {
 			return nil, invalidURLError
 		}
 
-		configURL.Scope = GitHubScopeOrganization
-		configURL.Organization = pathParts[0]
+		configURL.scope = gitHubScopeOrganization
+		configURL.organization = pathParts[0]
 
 	case 2: // Repository or enterprise
 		if strings.ToLower(pathParts[0]) == "enterprises" {
-			configURL.Scope = GitHubScopeEnterprise
-			configURL.Enterprise = pathParts[1]
+			configURL.scope = gitHubScopeEnterprise
+			configURL.enterprise = pathParts[1]
 			break
 		}
 
-		configURL.Scope = GitHubScopeRepository
-		configURL.Organization = pathParts[0]
-		configURL.Repository = pathParts[1]
+		configURL.scope = gitHubScopeRepository
+		configURL.organization = pathParts[0]
+		configURL.repository = pathParts[1]
 	default:
 		return nil, invalidURLError
 	}
@@ -72,20 +72,20 @@ func ParseGitHubConfigFromURL(in string) (*GitHubConfig, error) {
 	return configURL, nil
 }
 
-func (c *GitHubConfig) GitHubAPIURL(path string) *url.URL {
+func (c *gitHubConfig) gitHubAPIURL(path string) *url.URL {
 	result := &url.URL{
-		Scheme: c.ConfigURL.Scheme,
-		Host:   c.ConfigURL.Host, // default for Enterprise mode
+		Scheme: c.configURL.Scheme,
+		Host:   c.configURL.Host, // default for Enterprise mode
 		Path:   "/api/v3",        // default for Enterprise mode
 	}
 
-	isHosted := isHostedGitHubURL(c.ConfigURL)
+	isHosted := isHostedGitHubURL(c.configURL)
 
 	if isHosted {
-		result.Host = fmt.Sprintf("api.%s", c.ConfigURL.Host)
+		result.Host = fmt.Sprintf("api.%s", c.configURL.Host)
 		result.Path = ""
 
-		if strings.EqualFold("www.github.com", c.ConfigURL.Host) {
+		if strings.EqualFold("www.github.com", c.configURL.Host) {
 			// re-routing www.github.com to api.github.com
 			result.Host = "api.github.com"
 		}

config_test.go

@@ -15,116 +15,116 @@ func TestGitHubConfig(t *testing.T) {
 	t.Run("when given a valid URL", func(t *testing.T) {
 		tests := []struct {
 			configURL string
-			expected  *GitHubConfig
+			expected  *gitHubConfig
 		}{
 			{
 				configURL: "https://github.com/org/repo",
-				expected: &GitHubConfig{
-					Scope:        GitHubScopeRepository,
-					Enterprise:   "",
-					Organization: "org",
-					Repository:   "repo",
-					IsHosted:     true,
+				expected: &gitHubConfig{
+					scope:        gitHubScopeRepository,
+					enterprise:   "",
+					organization: "org",
+					repository:   "repo",
+					isHosted:     true,
 				},
 			},
 			{
 				configURL: "https://github.com/org/repo/",
-				expected: &GitHubConfig{
-					Scope:        GitHubScopeRepository,
-					Enterprise:   "",
-					Organization: "org",
-					Repository:   "repo",
-					IsHosted:     true,
+				expected: &gitHubConfig{
+					scope:        gitHubScopeRepository,
+					enterprise:   "",
+					organization: "org",
+					repository:   "repo",
+					isHosted:     true,
 				},
 			},
 			{
 				configURL: "https://github.com/org",
-				expected: &GitHubConfig{
-					Scope:        GitHubScopeOrganization,
-					Enterprise:   "",
-					Organization: "org",
-					Repository:   "",
-					IsHosted:     true,
+				expected: &gitHubConfig{
+					scope:        gitHubScopeOrganization,
+					enterprise:   "",
+					organization: "org",
+					repository:   "",
+					isHosted:     true,
 				},
 			},
 			{
 				configURL: "https://github.com/enterprises/my-enterprise",
-				expected: &GitHubConfig{
-					Scope:        GitHubScopeEnterprise,
-					Enterprise:   "my-enterprise",
-					Organization: "",
-					Repository:   "",
-					IsHosted:     true,
+				expected: &gitHubConfig{
+					scope:        gitHubScopeEnterprise,
+					enterprise:   "my-enterprise",
+					organization: "",
+					repository:   "",
+					isHosted:     true,
 				},
 			},
 			{
 				configURL: "https://github.com/enterprises/my-enterprise/",
-				expected: &GitHubConfig{
-					Scope:        GitHubScopeEnterprise,
-					Enterprise:   "my-enterprise",
-					Organization: "",
-					Repository:   "",
-					IsHosted:     true,
+				expected: &gitHubConfig{
+					scope:        gitHubScopeEnterprise,
+					enterprise:   "my-enterprise",
+					organization: "",
+					repository:   "",
+					isHosted:     true,
 				},
 			},
 			{
 				configURL: "https://www.github.com/org",
-				expected: &GitHubConfig{
-					Scope:        GitHubScopeOrganization,
-					Enterprise:   "",
-					Organization: "org",
-					Repository:   "",
-					IsHosted:     true,
+				expected: &gitHubConfig{
+					scope:        gitHubScopeOrganization,
+					enterprise:   "",
+					organization: "org",
+					repository:   "",
+					isHosted:     true,
 				},
 			},
 			{
 				configURL: "https://www.github.com/org/",
-				expected: &GitHubConfig{
-					Scope:        GitHubScopeOrganization,
-					Enterprise:   "",
-					Organization: "org",
-					Repository:   "",
-					IsHosted:     true,
+				expected: &gitHubConfig{
+					scope:        gitHubScopeOrganization,
+					enterprise:   "",
+					organization: "org",
+					repository:   "",
+					isHosted:     true,
 				},
 			},
 			{
 				configURL: "https://github.localhost/org",
-				expected: &GitHubConfig{
-					Scope:        GitHubScopeOrganization,
-					Enterprise:   "",
-					Organization: "org",
-					Repository:   "",
-					IsHosted:     true,
+				expected: &gitHubConfig{
+					scope:        gitHubScopeOrganization,
+					enterprise:   "",
+					organization: "org",
+					repository:   "",
+					isHosted:     true,
 				},
 			},
 			{
 				configURL: "https://my-ghes.com/org",
-				expected: &GitHubConfig{
-					Scope:        GitHubScopeOrganization,
-					Enterprise:   "",
-					Organization: "org",
-					Repository:   "",
-					IsHosted:     false,
+				expected: &gitHubConfig{
+					scope:        gitHubScopeOrganization,
+					enterprise:   "",
+					organization: "org",
+					repository:   "",
+					isHosted:     false,
 				},
 			},
 			{
 				configURL: "https://my-ghes.com/org/",
-				expected: &GitHubConfig{
-					Scope:        GitHubScopeOrganization,
-					Enterprise:   "",
-					Organization: "org",
-					Repository:   "",
-					IsHosted:     false,
+				expected: &gitHubConfig{
+					scope:        gitHubScopeOrganization,
+					enterprise:   "",
+					organization: "org",
+					repository:   "",
+					isHosted:     false,
 				},
 			},
 			{
 				configURL: "https://my-ghes.ghe.com/org/",
-				expected: &GitHubConfig{
-					Scope:        GitHubScopeOrganization,
-					Enterprise:   "",
-					Organization: "org",
-					Repository:   "",
-					IsHosted:     true,
+				expected: &gitHubConfig{
+					scope:        gitHubScopeOrganization,
+					enterprise:   "",
+					organization: "org",
+					repository:   "",
+					isHosted:     true,
 				},
 			},
 		}
@@ -133,9 +133,9 @@ func TestGitHubConfig(t *testing.T) {
 			t.Run(test.configURL, func(t *testing.T) {
 				parsedURL, err := url.Parse(strings.Trim(test.configURL, "/"))
 				require.NoError(t, err)
-				test.expected.ConfigURL = parsedURL
+				test.expected.configURL = parsedURL
 
-				cfg, err := ParseGitHubConfigFromURL(test.configURL)
+				cfg, err := parseGitHubConfigFromURL(test.configURL)
 				require.NoError(t, err)
 				assert.Equal(t, test.expected, cfg)
 			})
@@ -150,7 +150,7 @@ func TestGitHubConfig(t *testing.T) {
 		}
 
 		for _, u := range invalidURLs {
-			_, err := ParseGitHubConfigFromURL(u)
+			_, err := parseGitHubConfigFromURL(u)
 			require.Error(t, err)
 			assert.True(t, errors.Is(err, ErrInvalidGitHubConfigURL))
 		}
@@ -159,37 +159,37 @@ func TestGitHubConfig(t *testing.T) {
 
 func TestGitHubConfig_GitHubAPIURL(t *testing.T) {
 	t.Run("when hosted", func(t *testing.T) {
-		config, err := ParseGitHubConfigFromURL("https://github.com/org/repo")
+		config, err := parseGitHubConfigFromURL("https://github.com/org/repo")
 		require.NoError(t, err)
-		assert.True(t, config.IsHosted)
+		assert.True(t, config.isHosted)
 
-		result := config.GitHubAPIURL("/some/path")
+		result := config.gitHubAPIURL("/some/path")
 		assert.Equal(t, "https://api.github.com/some/path", result.String())
 	})
 	t.Run("when hosted with ghe.com", func(t *testing.T) {
-		config, err := ParseGitHubConfigFromURL("https://github.ghe.com/org/repo")
+		config, err := parseGitHubConfigFromURL("https://github.ghe.com/org/repo")
 		require.NoError(t, err)
-		assert.True(t, config.IsHosted)
+		assert.True(t, config.isHosted)
 
-		result := config.GitHubAPIURL("/some/path")
+		result := config.gitHubAPIURL("/some/path")
 		assert.Equal(t, "https://api.github.ghe.com/some/path", result.String())
 	})
 	t.Run("when not hosted", func(t *testing.T) {
-		config, err := ParseGitHubConfigFromURL("https://ghes.com/org/repo")
+		config, err := parseGitHubConfigFromURL("https://ghes.com/org/repo")
 		require.NoError(t, err)
-		assert.False(t, config.IsHosted)
+		assert.False(t, config.isHosted)
 
-		result := config.GitHubAPIURL("/some/path")
+		result := config.gitHubAPIURL("/some/path")
 		assert.Equal(t, "https://ghes.com/api/v3/some/path", result.String())
 	})
 	t.Run("when not hosted with ghe.com", func(t *testing.T) {
 		os.Setenv("GITHUB_ACTIONS_FORCE_GHES", "1")
 		defer os.Unsetenv("GITHUB_ACTIONS_FORCE_GHES")
-		config, err := ParseGitHubConfigFromURL("https://test.ghe.com/org/repo")
+		config, err := parseGitHubConfigFromURL("https://test.ghe.com/org/repo")
 		require.NoError(t, err)
-		assert.False(t, config.IsHosted)
+		assert.False(t, config.isHosted)
 
-		result := config.GitHubAPIURL("/some/path")
+		result := config.gitHubAPIURL("/some/path")
 		assert.Equal(t, "https://test.ghe.com/api/v3/some/path", result.String())
 	})
 }

error.go

@@ -1,124 +0,0 @@
-package scaleset
-
-import (
-	"encoding/json"
-	"errors"
-	"fmt"
-	"io"
-	"net/http"
-	"strings"
-)
-
-// Header names for request IDs
-const (
-	headerActionsActivityID = "ActivityId"
-	headerGitHubRequestID   = "X-GitHub-Request-Id"
-)
-
-type GitHubAPIError struct {
-	StatusCode int
-	RequestID  string
-	Err        error
-}
-
-func (e *GitHubAPIError) Error() string {
-	return fmt.Sprintf("github api error: StatusCode %d, RequestID %q: %v", e.StatusCode, e.RequestID, e.Err)
-}
-
-func (e *GitHubAPIError) Unwrap() error {
-	return e.Err
-}
-
-type ActionsError struct {
-	ActivityID string
-	StatusCode int
-	Err        error
-}
-
-func (e *ActionsError) Error() string {
-	return fmt.Sprintf("actions error: StatusCode %d, AcivityId %q: %v", e.StatusCode, e.ActivityID, e.Err)
-}
-
-func (e *ActionsError) Unwrap() error {
-	return e.Err
-}
-
-func (e *ActionsError) IsException(target string) bool {
-	if ex, ok := e.Err.(*ActionsExceptionError); ok {
-		return strings.Contains(ex.ExceptionName, target)
-	}
-	return false
-}
-
-type ActionsExceptionError struct {
-	ExceptionName string `json:"typeName,omitempty"`
-	Message       string `json:"message,omitempty"`
-}
-
-func (e *ActionsExceptionError) Error() string {
-	return fmt.Sprintf("%s: %s", e.ExceptionName, e.Message)
-}
-
-func parseActionsErrorFromResponse(response *http.Response) error {
-	if response.ContentLength == 0 {
-		return &ActionsError{
-			ActivityID: response.Header.Get(headerActionsActivityID),
-			StatusCode: response.StatusCode,
-			Err:        errors.New("unknown exception"),
-		}
-	}
-
-	body, err := io.ReadAll(response.Body)
-	if err != nil {
-		return &ActionsError{
-			ActivityID: response.Header.Get(headerActionsActivityID),
-			StatusCode: response.StatusCode,
-			Err:        err,
-		}
-	}
-
-	body = trimByteOrderMark(body)
-	contentType, ok := response.Header["Content-Type"]
-	if ok && len(contentType) > 0 && strings.Contains(contentType[0], "text/plain") {
-		message := string(body)
-		return &ActionsError{
-			ActivityID: response.Header.Get(headerActionsActivityID),
-			StatusCode: response.StatusCode,
-			Err:        errors.New(message),
-		}
-	}
-
-	var exception ActionsExceptionError
-	if err := json.Unmarshal(body, &exception); err != nil {
-		return &ActionsError{
-			ActivityID: response.Header.Get(headerActionsActivityID),
-			StatusCode: response.StatusCode,
-			Err:        err,
-		}
-	}
-
-	return &ActionsError{
-		ActivityID: response.Header.Get(headerActionsActivityID),
-		StatusCode: response.StatusCode,
-		Err:        &exception,
-	}
-}
-
-type MessageQueueTokenExpiredError struct {
-	activityID string
-	statusCode int
-	msg        string
-}
-
-func (e *MessageQueueTokenExpiredError) Error() string {
-	return fmt.Sprintf("MessageQueueTokenExpiredError: AcivityId %q, StatusCode %d: %s", e.activityID, e.statusCode, e.msg)
-}
-
-type HttpClientSideError struct {
-	msg  string
-	Code int
-}
-
-func (e *HttpClientSideError) Error() string {
-	return e.msg
-}

errors.go

@@ -0,0 +1,98 @@
+package scaleset
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+)
+
+type scalesetError string
+
+func (e scalesetError) Error() string {
+	return string(e)
+}
+
+var (
+	RunnerNotFoundError           = scalesetError("runner not found")
+	RunnerExistsError             = scalesetError("runner exists")
+	JobStillRunningError          = scalesetError("job still running")
+	MessageQueueTokenExpiredError = scalesetError("message queue token expired")
+)
+
+type actionsExceptionError struct {
+	ExceptionName string `json:"typeName,omitempty"`
+	Message       string `json:"message,omitempty"`
+}
+
+func (e actionsExceptionError) Error() string {
+	return fmt.Sprintf("%s: %s", e.ExceptionName, e.Message)
+}
+
+// newRequestResponseError creates a detailed error message based on the HTTP request and response,
+// including parsing the response body for known error formats.
+//
+// The sendRequest already parses errors using this method, so use this error if the client doesn't
+// return an error, but the error is happening on the application logic level.
+//
+// Prefer creating errors using this function instead of manually constructing error messages since it automatically
+// includes useful metadata like activity IDs and request IDs, and handles well-known error cases.
+func newRequestResponseError(req *http.Request, resp *http.Response, err error) error {
+	var sb strings.Builder
+	fmt.Fprintf(&sb, "request %s %s failed", req.Method, req.URL.String())
+
+	if resp == nil {
+		return fmt.Errorf("%s: %w", sb.String(), err)
+	}
+
+	sb.WriteRune('(')
+	fmt.Fprintf(&sb, "status=%q", resp.Status)
+	if resp.Header.Get(headerActionsActivityID) != "" {
+		fmt.Fprintf(&sb, ", activity_id=%q", resp.Header.Get(headerActionsActivityID))
+	}
+
+	if resp.Header.Get(headerGitHubRequestID) != "" {
+		fmt.Fprintf(&sb, ", github_request_id=%q", resp.Header.Get(headerGitHubRequestID))
+	}
+	sb.WriteRune(')')
+
+	if resp.Body == nil || resp.ContentLength == 0 {
+		return fmt.Errorf("%s: %w: unknown error", sb.String(), err)
+	}
+
+	body, bodyErr := io.ReadAll(resp.Body)
+	if bodyErr != nil {
+		return fmt.Errorf("%s: %w: failed to read error response body: %w", sb.String(), err, bodyErr)
+	}
+	if len(body) == 0 {
+		return fmt.Errorf("%s: %w: unknown error", sb.String(), err)
+	}
+
+	var scalesetErr scalesetError
+	if errors.As(err, &scalesetErr) {
+		return fmt.Errorf("%s: %w: %s", sb.String(), err, string(body))
+	}
+
+	contentType := resp.Header.Get("Content-Type")
+	if len(contentType) > 0 && strings.Contains(contentType, "text/plain") {
+		return fmt.Errorf("%s: %w: %s", sb.String(), err, string(body))
+	}
+
+	var exception actionsExceptionError
+	if err := json.Unmarshal(body, &exception); err != nil {
+		return fmt.Errorf("%s: %w: failed to unmarshal error response body: %q", sb.String(), err, string(body))
+	}
+
+	switch {
+	case strings.Contains(exception.ExceptionName, "AgentExistsException"):
+		return fmt.Errorf("%s: %w: %s", sb.String(), RunnerExistsError, exception.Message)
+	case strings.Contains(exception.ExceptionName, "AgentNotFoundException"):
+		return fmt.Errorf("%s: %w: %s", sb.String(), RunnerNotFoundError, exception.Message)
+	case strings.Contains(exception.ExceptionName, "JobStillRunningException"):
+		return fmt.Errorf("%s: %w: %s", sb.String(), JobStillRunningError, exception.Message)
+	default:
+		return fmt.Errorf("%s: %w: %w", sb.String(), err, exception)
+	}
+}

errors_test.go

@@ -2,8 +2,10 @@ package scaleset
 
 import (
 	"errors"
+	"fmt"
 	"io"
 	"net/http"
+	"net/url"
 	"strings"
 	"testing"
 
@@ -11,78 +13,14 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestActionsError(t *testing.T) {
-	t.Run("contains the status code, activity ID, and error", func(t *testing.T) {
-		err := &ActionsError{
-			ActivityID: "activity-id",
-			StatusCode: 404,
-			Err:        errors.New("example error description"),
-		}
+type readErrCloser struct{}
 
-		s := err.Error()
-		assert.Contains(t, s, "StatusCode 404")
-		assert.Contains(t, s, "AcivityId \"activity-id\"")
-		assert.Contains(t, s, "example error description")
-	})
-
-	t.Run("unwraps the error", func(t *testing.T) {
-		err := &ActionsError{
-			ActivityID: "activity-id",
-			StatusCode: 404,
-			Err: &ActionsExceptionError{
-				ExceptionName: "exception-name",
-				Message:       "example error message",
-			},
-		}
-
-		assert.Equal(t, err.Unwrap(), err.Err)
-	})
-
-	t.Run("is exception is ok", func(t *testing.T) {
-		err := &ActionsError{
-			ActivityID: "activity-id",
-			StatusCode: 404,
-			Err: &ActionsExceptionError{
-				ExceptionName: "exception-name",
-				Message:       "example error message",
-			},
-		}
-
-		var exception *ActionsExceptionError
-		assert.True(t, errors.As(err, &exception))
-
-		assert.True(t, err.IsException("exception-name"))
-	})
-
-	t.Run("is exception is not ok", func(t *testing.T) {
-		tt := map[string]*ActionsError{
-			"not an exception": {
-				ActivityID: "activity-id",
-				StatusCode: 404,
-				Err:        errors.New("example error description"),
-			},
-			"not target exception": {
-				ActivityID: "activity-id",
-				StatusCode: 404,
-				Err: &ActionsExceptionError{
-					ExceptionName: "exception-name",
-					Message:       "example error message",
-				},
-			},
-		}
-
-		targetException := "target-exception"
-		for name, err := range tt {
-			t.Run(name, func(t *testing.T) {
-				assert.False(t, err.IsException(targetException))
-			})
-		}
-	})
-}
+func (readErrCloser) Read([]byte) (int, error) { return 0, fmt.Errorf("read failed") }
+func (readErrCloser) Close() error             { return nil }
 
 func TestActionsExceptionError(t *testing.T) {
 	t.Run("contains the exception name and message", func(t *testing.T) {
-		err := &ActionsExceptionError{
+		err := actionsExceptionError{
 			ExceptionName: "exception-name",
 			Message:       "example error message",
 		}
@@ -93,113 +31,223 @@ func TestActionsExceptionError(t *testing.T) {
 	})
 }
 
-func TestGitHubAPIError(t *testing.T) {
-	t.Run("contains the status code, request ID, and error", func(t *testing.T) {
-		err := &GitHubAPIError{
-			StatusCode: 404,
-			RequestID:  "request-id",
-			Err:        errors.New("example error description"),
-		}
+func TestNewRequestResponseError(t *testing.T) {
+	req := func(t *testing.T) *http.Request {
+		t.Helper()
+		u, err := url.Parse("https://example.com/org/repo")
+		require.NoError(t, err)
+		return &http.Request{Method: http.MethodGet, URL: u}
+	}
 
-		s := err.Error()
-		assert.Contains(t, s, "StatusCode 404")
-		assert.Contains(t, s, "RequestID \"request-id\"")
-		assert.Contains(t, s, "example error description")
+	t.Run("resp is nil", func(t *testing.T) {
+		base := errors.New("base")
+		err := newRequestResponseError(req(t), nil, base)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "request GET https://example.com/org/repo failed")
+		assert.True(t, errors.Is(err, base))
 	})
 
-	t.Run("unwraps the error", func(t *testing.T) {
-		err := &GitHubAPIError{
-			StatusCode: 404,
-			RequestID:  "request-id",
-			Err:        errors.New("example error description"),
+	t.Run("resp body is nil", func(t *testing.T) {
+		base := errors.New("base")
+		resp := &http.Response{
+			Status:        "500 Internal Server Error",
+			StatusCode:    http.StatusInternalServerError,
+			ContentLength: 123,
+			Header:        make(http.Header),
+			Body:          nil,
 		}
 
-		assert.Equal(t, err.Unwrap(), err.Err)
+		err := newRequestResponseError(req(t), resp, base)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "unknown error")
+		assert.True(t, errors.Is(err, base))
 	})
-}
 
-func ParseActionsErrorFromResponse(t *testing.T) {
-	t.Run("empty content length", func(t *testing.T) {
-		response := &http.Response{
+	t.Run("empty body returns unknown error", func(t *testing.T) {
+		base := errors.New("base")
+		resp := &http.Response{
+			Status:        "404 Not Found",
+			StatusCode:    http.StatusNotFound,
 			ContentLength: 0,
-			Header: http.Header{
-				headerActionsActivityID: []string{"activity-id"},
-			},
-			StatusCode: 404,
+			Header:        make(http.Header),
 		}
+		resp.Header.Set(headerActionsActivityID, "activity-id")
+		resp.Header.Set(headerGitHubRequestID, "request-id")
 
-		err := parseActionsErrorFromResponse(response)
+		err := newRequestResponseError(req(t), resp, base)
 		require.Error(t, err)
-		assert.Equal(t, err.(*ActionsError).ActivityID, "activity-id")
-		assert.Equal(t, err.(*ActionsError).StatusCode, 404)
-		assert.Equal(t, err.(*ActionsError).Err.Error(), "unknown exception")
+		assert.Contains(t, err.Error(), "status=\"404 Not Found\"")
+		assert.Contains(t, err.Error(), "activity_id=\"activity-id\"")
+		assert.Contains(t, err.Error(), "github_request_id=\"request-id\"")
+		assert.Contains(t, err.Error(), "unknown error")
+		assert.True(t, errors.Is(err, base))
 	})
 
-	t.Run("contains text plain error", func(t *testing.T) {
-		errorMessage := "example error message"
-		response := &http.Response{
-			ContentLength: int64(len(errorMessage)),
-			Header: http.Header{
-				headerActionsActivityID: []string{"activity-id"},
-				"Content-Type":          []string{"text/plain"},
-			},
-			StatusCode: 404,
-			Body:       io.NopCloser(strings.NewReader(errorMessage)),
+	t.Run("read body failure includes read error", func(t *testing.T) {
+		base := errors.New("base")
+		resp := &http.Response{
+			Status:        "400 Bad Request",
+			StatusCode:    http.StatusBadRequest,
+			ContentLength: 1,
+			Header:        make(http.Header),
+			Body:          io.NopCloser(readErrCloser{}),
 		}
 
-		err := parseActionsErrorFromResponse(response)
+		err := newRequestResponseError(req(t), resp, base)
 		require.Error(t, err)
-		var actionsError *ActionsError
-		assert.ErrorAs(t, err, &actionsError)
-		assert.Equal(t, actionsError.ActivityID, "activity-id")
-		assert.Equal(t, actionsError.StatusCode, 404)
-		assert.Equal(t, actionsError.Err.Error(), errorMessage)
+		assert.Contains(t, err.Error(), "failed to read error response body")
+		assert.True(t, errors.Is(err, base))
+		assert.Contains(t, err.Error(), "read failed")
 	})
 
-	t.Run("contains json error", func(t *testing.T) {
-		errorMessage := `{"typeName":"exception-name","message":"example error message"}`
-		response := &http.Response{
-			ContentLength: int64(len(errorMessage)),
-			Header: http.Header{
-				headerActionsActivityID: []string{"activity-id"},
-				"Content-Type":          []string{"application/json"},
-			},
-			StatusCode: 404,
-			Body:       io.NopCloser(strings.NewReader(errorMessage)),
+	t.Run("unknown content length and empty body returns unknown error", func(t *testing.T) {
+		base := errors.New("base")
+		resp := &http.Response{
+			Status:        "400 Bad Request",
+			StatusCode:    http.StatusBadRequest,
+			ContentLength: -1,
+			Header:        make(http.Header),
+			Body:          io.NopCloser(strings.NewReader("")),
 		}
 
-		err := parseActionsErrorFromResponse(response)
+		err := newRequestResponseError(req(t), resp, base)
 		require.Error(t, err)
-		var actionsError *ActionsError
-		assert.ErrorAs(t, err, &actionsError)
-		assert.Equal(t, actionsError.ActivityID, "activity-id")
-		assert.Equal(t, actionsError.StatusCode, 404)
-
-		inner, ok := actionsError.Err.(*ActionsExceptionError)
-		require.True(t, ok)
-		assert.Equal(t, inner.ExceptionName, "exception-name")
-		assert.Equal(t, inner.Message, "example error message")
+		assert.Contains(t, err.Error(), "unknown error")
+		assert.True(t, errors.Is(err, base))
 	})
 
-	t.Run("wrapped exception error", func(t *testing.T) {
-		errorMessage := `{"typeName":"exception-name","message":"example error message"}`
-		response := &http.Response{
-			ContentLength: int64(len(errorMessage)),
-			Header: http.Header{
-				headerActionsActivityID: []string{"activity-id"},
-				"Content-Type":          []string{"application/json"},
-			},
-			StatusCode: 404,
-			Body:       io.NopCloser(strings.NewReader(errorMessage)),
+	t.Run("text/plain body is included", func(t *testing.T) {
+		base := errors.New("base")
+		body := "example plain text error"
+		resp := &http.Response{
+			Status:        "400 Bad Request",
+			StatusCode:    http.StatusBadRequest,
+			ContentLength: int64(len(body)),
+			Header:        make(http.Header),
+			Body:          io.NopCloser(strings.NewReader(body)),
 		}
+		resp.Header.Set("Content-Type", "text/plain")
+		resp.Header.Set(headerActionsActivityID, "activity-id")
 
-		err := parseActionsErrorFromResponse(response)
+		err := newRequestResponseError(req(t), resp, base)
 		require.Error(t, err)
+		assert.Contains(t, err.Error(), body)
+		assert.True(t, errors.Is(err, base))
+	})
 
-		var actionsExceptionError *ActionsExceptionError
-		assert.ErrorAs(t, err, &actionsExceptionError)
+	t.Run("scalesetError in error chain uses raw body (no JSON parsing)", func(t *testing.T) {
+		wrapped := fmt.Errorf("wrapped: %w", RunnerNotFoundError)
+		body := `{"typeName":"AgentExistsException","message":"should not be parsed"}`
+		resp := &http.Response{
+			Status:        "404 Not Found",
+			StatusCode:    http.StatusNotFound,
+			ContentLength: int64(len(body)),
+			Header:        make(http.Header),
+			Body:          io.NopCloser(strings.NewReader(body)),
+		}
+		resp.Header.Set("Content-Type", "application/json")
 
-		assert.Equal(t, actionsExceptionError.ExceptionName, "exception-name")
-		assert.Equal(t, actionsExceptionError.Message, "example error message")
+		err := newRequestResponseError(req(t), resp, wrapped)
+		require.Error(t, err)
+		assert.True(t, errors.Is(err, RunnerNotFoundError))
+		assert.Contains(t, err.Error(), body)
+	})
+
+	t.Run("known actions exception maps to sentinel error", func(t *testing.T) {
+		base := errors.New("base")
+		jsonBody := `{"typeName":"AgentExistsException","message":"runner already exists"}`
+		resp := &http.Response{
+			Status:        "409 Conflict",
+			StatusCode:    http.StatusConflict,
+			ContentLength: int64(len(jsonBody)),
+			Header:        make(http.Header),
+			Body:          io.NopCloser(strings.NewReader(jsonBody)),
+		}
+		resp.Header.Set("Content-Type", "application/json")
+
+		err := newRequestResponseError(req(t), resp, base)
+		require.Error(t, err)
+		assert.True(t, errors.Is(err, RunnerExistsError))
+		assert.False(t, errors.Is(err, base), "base error should not be wrapped for mapped exceptions")
+		assert.Contains(t, err.Error(), "runner already exists")
+	})
+
+	t.Run("agent not found exception maps to sentinel error", func(t *testing.T) {
+		base := errors.New("base")
+		jsonBody := `{"typeName":"AgentNotFoundException","message":"missing"}`
+		resp := &http.Response{
+			Status:        "404 Not Found",
+			StatusCode:    http.StatusNotFound,
+			ContentLength: int64(len(jsonBody)),
+			Header:        make(http.Header),
+			Body:          io.NopCloser(strings.NewReader(jsonBody)),
+		}
+		resp.Header.Set("Content-Type", "application/json")
+
+		err := newRequestResponseError(req(t), resp, base)
+		require.Error(t, err)
+		assert.True(t, errors.Is(err, RunnerNotFoundError))
+		assert.False(t, errors.Is(err, base))
+		assert.Contains(t, err.Error(), "missing")
+	})
+
+	t.Run("job still running exception maps to sentinel error", func(t *testing.T) {
+		base := errors.New("base")
+		jsonBody := `{"typeName":"JobStillRunningException","message":"still running"}`
+		resp := &http.Response{
+			Status:        "409 Conflict",
+			StatusCode:    http.StatusConflict,
+			ContentLength: int64(len(jsonBody)),
+			Header:        make(http.Header),
+			Body:          io.NopCloser(strings.NewReader(jsonBody)),
+		}
+		resp.Header.Set("Content-Type", "application/json")
+
+		err := newRequestResponseError(req(t), resp, base)
+		require.Error(t, err)
+		assert.True(t, errors.Is(err, JobStillRunningError))
+		assert.False(t, errors.Is(err, base))
+		assert.Contains(t, err.Error(), "still running")
+	})
+
+	t.Run("invalid json returns unmarshal error and includes body", func(t *testing.T) {
+		base := errors.New("base")
+		bad := "not-json"
+		resp := &http.Response{
+			Status:        "400 Bad Request",
+			StatusCode:    http.StatusBadRequest,
+			ContentLength: int64(len(bad)),
+			Header:        make(http.Header),
+			Body:          io.NopCloser(strings.NewReader(bad)),
+		}
+		resp.Header.Set("Content-Type", "application/json")
+
+		err := newRequestResponseError(req(t), resp, base)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "failed to unmarshal error response body")
+		assert.Contains(t, err.Error(), "not-json")
+		assert.False(t, errors.Is(err, base), "base error is not wrapped on JSON unmarshal failures")
+	})
+
+	t.Run("unknown json error wraps exception", func(t *testing.T) {
+		base := errors.New("base")
+		jsonBody := `{"typeName":"SomeException","message":"example error message"}`
+		resp := &http.Response{
+			Status:        "500 Internal Server Error",
+			StatusCode:    http.StatusInternalServerError,
+			ContentLength: int64(len(jsonBody)),
+			Header:        make(http.Header),
+			Body:          io.NopCloser(strings.NewReader(jsonBody)),
+		}
+		resp.Header.Set("Content-Type", "application/json")
+
+		err := newRequestResponseError(req(t), resp, base)
+		require.Error(t, err)
+		assert.True(t, errors.Is(err, base))
+
+		var ex actionsExceptionError
+		assert.True(t, errors.As(err, &ex))
+		assert.Equal(t, "SomeException", ex.ExceptionName)
+		assert.Equal(t, "example error message", ex.Message)
 	})
 }

examples/dockerscaleset/README.md

@@ -0,0 +1,55 @@
+# Docker Runner Scale Set Example
+
+This example showcases a Docker implementation of GitHub Actions runner scale sets, using the `github.com/actions/scaleset` client to provision ephemeral GitHub Actions runners as Docker containers.
+
+The goal of this example is to show how simple and powerful it is when you only need to focus on the core logic of scaling runners up and down, while the client handles all the API interactions.
+
+> [!WARNING]
+> This is a simplified example meant for demonstration and learning purposes. It is not intended for production use.
+
+> [!NOTE]
+> When exiting normally all runners and the scale set itself are cleaned up automatically.
+
+## Getting started
+
+You can install the example with:
+
+```bash
+go install github.com/actions/scaleset/examples/dockerscaleset@latest
+```
+
+If this fails you should also try running the command with
+
+```bash
+GONOSUMDB=github.com/actions/scaleset GOPRIVATE=github.com/actions/scaleset go
+install github.com/actions/scaleset/examples/dockerscaleset@latest
+```
+
+You'll then need:
+
+- Docker installed and running on your machine.
+- A URL for the target repository, organization, or enterprise where you want to register your scale set.
+- [Credentials that have access to the above target](https://docs.github.com/en/actions/tutorials/use-actions-runner-controller/authenticate-to-the-api): you can use either a GitHub App (recommended) or a Personal Access Token (PAT).
+- A name for your scale set (this must be unique within the runner group the scale set is created in).
+
+---
+
+## Flags
+
+| Flag | Required | Description |
+|------|----------|-------------|
+| `--url` | Yes | Registration target (org, repo, or enterprise URL, e.g. `https://github.com/org/repo`). |
+| `--name` | Yes | Runner scale set name (must be unique within the runner group). |
+| `--labels` | No | Labels for workflow targeting (comma-separated or repeated). Defaults to `--name` if not provided. |
+| `--max-runners` | No | Upper bound of concurrently provisioned runners (default 10). |
+| `--min-runners` | No | Lower bound to maintain (default 0). |
+| `--runner-group` | No | Runner group name (default `default`). |
+| `--app-client-id` | Cond.* | GitHub App Client (App) ID. |
+| `--app-installation-id` | Cond.* | GitHub App Installation ID. |
+| `--app-private-key` | Cond.* | GitHub App private key PEM contents. |
+| `--token` | Cond.* | Personal Access Token (alternative to App). |
+| `--log-level` | No | `debug`, `info`, `warn`, `error` (default `info`). |
+| `--log-format` | No | `text`, `json`, or `none` (any invalid → no logs). |
+| `--runner-image` | No | Override container image (defaults to latest official). |
+
+*Provide either App credentials (all three) OR a PAT.*

examples/dockerscaleset/config.go

@@ -15,6 +15,7 @@ type Config struct {
 	MaxRunners      int
 	MinRunners      int
 	ScaleSetName    string
+	Labels          []string
 	RunnerGroup     string
 	GitHubApp       scaleset.GitHubAppAuth
 	Token           string
@@ -47,6 +48,11 @@ func (c *Config) Validate() error {
 	if c.ScaleSetName == "" {
 		return fmt.Errorf("scale set name is required")
 	}
+	for i, label := range c.Labels {
+		if strings.TrimSpace(label) == "" {
+			return fmt.Errorf("label at index %d is empty", i)
+		}
+	}
 	if c.MaxRunners < c.MinRunners {
 		return fmt.Errorf("max runners cannot be less than min-runners")
 	}
@@ -59,16 +65,35 @@ func (c *Config) Validate() error {
 	return nil
 }
 
-func (c *Config) ActionsAuth() *scaleset.ActionsAuth {
+// systemInfo serves as a base system info
+func systemInfo(scaleSetID int) scaleset.SystemInfo {
+	return scaleset.SystemInfo{
+		System:     "dockerscaleset",
+		Subsystem:  "dockerscaleset",
+		CommitSHA:  "NA",    // You can leverage build flags to set commit SHA
+		Version:    "0.1.0", // You can leverage build flags to set version
+		ScaleSetID: scaleSetID,
+	}
+}
+
+func (c *Config) ScalesetClient() (*scaleset.Client, error) {
 	if err := c.GitHubApp.Validate(); err == nil {
-		return &scaleset.ActionsAuth{
-			App: &c.GitHubApp,
-		}
+		return scaleset.NewClientWithGitHubApp(
+			scaleset.ClientWithGitHubAppConfig{
+				GitHubConfigURL: c.RegistrationURL,
+				GitHubAppAuth:   c.GitHubApp,
+				SystemInfo:      systemInfo(0),
+			},
+		)
 	}
 
-	return &scaleset.ActionsAuth{
-		Token: c.Token,
-	}
+	return scaleset.NewClientWithPersonalAccessToken(
+		scaleset.NewClientWithPersonalAccessTokenConfig{
+			GitHubConfigURL:     c.RegistrationURL,
+			PersonalAccessToken: c.Token,
+			SystemInfo:          systemInfo(0),
+		},
+	)
 }
 
 func (c *Config) Logger() *slog.Logger {
@@ -101,3 +126,16 @@ func (c *Config) Logger() *slog.Logger {
 		return slog.New(slog.DiscardHandler)
 	}
 }
+
+// BuildLabels returns the labels to use for the runner scale set.
+// If custom labels are provided, those are used; otherwise, the scale set name is used as the label.
+func (c *Config) BuildLabels() []scaleset.Label {
+	if len(c.Labels) > 0 {
+		labels := make([]scaleset.Label, len(c.Labels))
+		for i, name := range c.Labels {
+			labels[i] = scaleset.Label{Name: strings.TrimSpace(name)}
+		}
+		return labels
+	}
+	return []scaleset.Label{{Name: c.ScaleSetName}}
+}

examples/dockerscaleset/e2e_test.go

@@ -0,0 +1,322 @@
+package main
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strconv"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/google/go-github/v79/github"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestE2E(t *testing.T) {
+	if os.Getenv("E2E") != "true" {
+		t.Skip("Skipping E2E test; set E2E=true to run")
+	}
+
+	configURL := mustGetEnv(t, "E2E_SCALESET_URL")
+	name := mustGetEnv(t, "E2E_SCALESET_NAME")
+
+	workflowEnv := mustE2EWorkflowEnv(t, name)
+	runArgs := mustE2ECommandArgs(t, configURL, name)
+
+	tempDir, err := os.MkdirTemp("", "e2e-dockerscaleset-")
+	require.NoError(t, err, "Failed to create temp dir")
+	defer os.RemoveAll(tempDir)
+
+	binaryPath := filepath.Join(tempDir, "dockerscaleset")
+
+	// Build the dockerscaleset binary in temp dir
+	{
+		cmd := exec.Command("go", "build", "-o", binaryPath, ".")
+		output, err := cmd.CombinedOutput()
+		require.NoError(t, err, "Failed to build dockerscaleset: %s", output)
+	}
+
+	// Fatal channel
+	testErrCh := make(chan error, 2)
+
+	runCmd := exec.Command(binaryPath, runArgs...)
+	stdout, err := runCmd.StdoutPipe()
+	runCmd.Stderr = os.Stderr
+	require.NoError(t, err, "Failed to get stdout pipe")
+	err = runCmd.Start()
+	require.NoError(t, err, "Failed to start dockerscaleset")
+
+	// Command exit error
+	cmdCh := make(chan error, 1)
+	t.Cleanup(func() {
+		_ = runCmd.Process.Signal(os.Interrupt)
+		<-cmdCh
+	})
+
+	// Wait for log line
+	waitCh := make(chan struct{}, 1)
+
+	var (
+		bufMu sync.Mutex
+		buf   bytes.Buffer
+	)
+	go func() {
+		scanner := bufio.NewScanner(stdout)
+		for scanner.Scan() {
+			line := scanner.Text()
+			bufMu.Lock()
+			buf.WriteString(line + "\n")
+			bufMu.Unlock()
+			if strings.Contains(line, "Getting next message") {
+				close(waitCh)
+				break
+			}
+		}
+		if err := scanner.Err(); err != nil {
+			testErrCh <- fmt.Errorf("error reading dockerscaleset stdout: %w", err)
+			return
+		}
+
+		cmdCh <- runCmd.Wait()
+		close(cmdCh)
+	}()
+
+	runID, err := workflowEnv.triggerWorkflowDispatch(t, t.Context())
+	require.NoError(t, err, "Failed to trigger workflow")
+
+	statusCh := make(chan *WorkflowRun, 1)
+	go func() {
+		select {
+		case <-waitCh:
+		case <-time.After(30 * time.Second):
+			bufMu.Lock()
+			logs := buf.String()
+			bufMu.Unlock()
+			testErrCh <- fmt.Errorf("timeout waiting for dockerscaleset to be ready; logs:\n%s", logs)
+			return
+		}
+		status, err := workflowEnv.waitForWorkflowCompletion(t, t.Context(), runID, 10*time.Minute)
+		if err != nil {
+			testErrCh <- fmt.Errorf("failed to wait for workflow completion: %w", err)
+			return
+		}
+		statusCh <- status
+	}()
+
+	select {
+	case err := <-cmdCh:
+		select {
+		case status := <-statusCh:
+			assert.Equal(t, "completed", status.Status)
+			assert.Equal(t, "success", status.Conclusion)
+		case <-time.After(30 * time.Second):
+			bufMu.Lock()
+			logs := buf.String()
+			bufMu.Unlock()
+			t.Fatalf("Timeout waiting for workflow status after dockerscaleset exited\nexit: %v\nlogs:%s\n", err, logs)
+		}
+	case status := <-statusCh:
+		assert.NotNil(t, status, "WorkflowRun status is nil")
+		assert.Equal(t, "completed", status.Status)
+		assert.Equal(t, "success", status.Conclusion)
+		return
+	case err := <-testErrCh:
+		t.Fatal(err)
+	}
+}
+
+type e2eWorkflowEnv struct {
+	targetOrg  string
+	targetRepo string
+	targetFile string
+
+	scalesetName string
+	client       *github.Client
+}
+
+func mustE2EWorkflowEnv(t *testing.T, scalesetName string) *e2eWorkflowEnv {
+	return &e2eWorkflowEnv{
+		targetOrg:    mustGetEnv(t, "E2E_WORKFLOW_TARGET_ORG"),
+		targetRepo:   mustGetEnv(t, "E2E_WORKFLOW_TARGET_REPO"),
+		targetFile:   mustGetEnv(t, "E2E_WORKFLOW_TARGET_FILE"),
+		scalesetName: scalesetName,
+		client:       github.NewClient(nil).WithAuthToken(mustGetEnv(t, "E2E_WORKFLOW_GITHUB_TOKEN")),
+	}
+}
+
+func mustE2ECommandArgs(t *testing.T, configURL, name string) []string {
+	args := []string{
+		"--url", configURL,
+		"--name", name,
+		"--log-level", "debug",
+	}
+
+	// GitHub App credentials
+	var (
+		clientID       string
+		installationID int
+		privateKeyPath string
+	)
+
+	// GitHub token
+	var token string
+
+	clientID = os.Getenv("E2E_SCALESET_GITHUB_APP_CLIENT_ID")
+	installationIDStr := os.Getenv("E2E_SCALESET_GITHUB_APP_INSTALLATION_ID")
+	privateKeyPath = os.Getenv("E2E_SCALESET_GITHUB_APP_PRIVATE_KEY_PATH")
+
+	if clientID != "" && installationIDStr != "" && privateKeyPath != "" {
+		id, err := strconv.Atoi(installationIDStr)
+		require.NoError(t, err, "Invalid E2E_SCALESET_GITHUB_APP_INSTALLATION_ID")
+		installationID = id
+		args = append(args,
+			"--app-client-id", clientID,
+			"--app-installation-id", fmt.Sprintf("%d", installationID),
+			"--app-private-key", privateKeyPath,
+		)
+	} else {
+		token = os.Getenv("E2E_SCALESET_GITHUB_TOKEN")
+		require.NotEmpty(t, token, "E2E_SCALESET_GITHUB_TOKEN must be set if GitHub App credentials are not provided")
+		args = append(args,
+			"--token", token,
+		)
+	}
+
+	runnerGroup := os.Getenv("E2E_SCALESET_RUNNER_GROUP")
+	if runnerGroup != "" {
+		args = append(args,
+			"--runner-group", runnerGroup,
+		)
+	}
+
+	minRunners := 0
+	if minRunnersStr := os.Getenv("E2E_SCALESET_MIN_RUNNERS"); minRunnersStr != "" {
+		m, err := strconv.Atoi(minRunnersStr)
+		require.NoError(t, err, "Invalid E2E_SCALESET_MIN_RUNNERS")
+		minRunners = m
+		require.GreaterOrEqual(t, minRunners, 0, "E2E_SCALESET_MIN_RUNNERS must be >= 0")
+	}
+
+	maxRunners := 10
+	if maxRunnersStr := os.Getenv("E2E_SCALESET_MAX_RUNNERS"); maxRunnersStr != "" {
+		m, err := strconv.Atoi(maxRunnersStr)
+		require.NoError(t, err, "Invalid E2E_SCALESET_MAX_RUNNERS")
+		maxRunners = m
+		require.GreaterOrEqual(t, maxRunners, 0, "E2E_SCALESET_MAX_RUNNERS must be >= 0")
+	}
+
+	require.GreaterOrEqual(t, maxRunners, minRunners, "E2E_SCALESET_MAX_RUNNERS must be >= E2E_SCALESET_MIN_RUNNERS")
+
+	args = append(args,
+		"--min-runners", strconv.Itoa(minRunners),
+		"--max-runners", strconv.Itoa(maxRunners),
+	)
+
+	return args
+}
+
+type WorkflowRun struct {
+	ID         int    `json:"id"`
+	Status     string `json:"status"`
+	Conclusion string `json:"conclusion"`
+	CreatedAt  string `json:"created_at"`
+}
+
+func (env *e2eWorkflowEnv) triggerWorkflowDispatch(t *testing.T, ctx context.Context) (int, error) {
+	dispatchTime := time.Now().UTC()
+
+	resp, err := env.client.Actions.CreateWorkflowDispatchEventByFileName(
+		ctx,
+		env.targetOrg,
+		env.targetRepo,
+		env.targetFile,
+		github.CreateWorkflowDispatchEventRequest{
+			Ref: "main",
+			Inputs: map[string]any{
+				"scaleset_name": env.scalesetName,
+			},
+		},
+	)
+	require.NoError(t, err, "Failed to create workflow dispatch")
+	require.Equal(t, 204, resp.StatusCode, "Unexpected status code from workflow dispatch")
+
+	// Wait a bit for the run to be created
+	time.Sleep(10 * time.Second)
+
+	// List runs with event=workflow_dispatch and since=dispatchTime
+	opts := &github.ListWorkflowRunsOptions{
+		Event:   "workflow_dispatch",
+		Created: ">=" + dispatchTime.Format(time.RFC3339),
+		ListOptions: github.ListOptions{
+			PerPage: 10,
+		},
+	}
+	runs, _, err := env.client.Actions.ListWorkflowRunsByFileName(
+		t.Context(),
+		env.targetOrg,
+		env.targetRepo,
+		env.targetFile,
+		opts,
+	)
+	require.NoError(t, err, "Failed to list workflow runs")
+	require.Greater(t, len(runs.WorkflowRuns), 0, "No workflow runs found after dispatch")
+
+	// Sort by created_at desc, take the first (most recent)
+	var latestRun *github.WorkflowRun
+	var latestTime time.Time
+	for _, run := range runs.WorkflowRuns {
+		createdAt := run.CreatedAt.Time
+		if createdAt.After(latestTime) {
+			latestTime = createdAt
+			latestRun = run
+		}
+	}
+
+	if latestRun == nil {
+		return 0, fmt.Errorf("no workflow runs found after dispatch")
+	}
+
+	return int(latestRun.GetID()), nil
+}
+
+func (env *e2eWorkflowEnv) waitForWorkflowCompletion(t *testing.T, ctx context.Context, runID int, timeout time.Duration) (*WorkflowRun, error) {
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
+	ticker := time.NewTicker(5 * time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-ticker.C:
+			run, _, err := env.client.Actions.GetWorkflowRunByID(ctx, env.targetOrg, env.targetRepo, int64(runID))
+			require.NoError(t, err, "Failed to get workflow run by ID")
+
+			if run.GetStatus() == "completed" {
+				return &WorkflowRun{
+					ID:         int(run.GetID()),
+					Status:     run.GetStatus(),
+					Conclusion: run.GetConclusion(),
+					CreatedAt:  run.GetCreatedAt().Format(time.RFC3339),
+				}, nil
+			}
+		}
+	}
+}
+
+func mustGetEnv(t *testing.T, key string) string {
+	value := os.Getenv(key)
+	if value == "" {
+		t.Fatalf("Environment variable %s not set", key)
+	}
+	return value
+}

examples/dockerscaleset/main.go

@@ -12,7 +12,8 @@ import (
 	"github.com/actions/scaleset"
 	"github.com/actions/scaleset/listener"
 	"github.com/docker/docker/api/types/image"
-	dockerclient "github.com/moby/moby/client"
+	dockerclient "github.com/docker/docker/client"
+	"github.com/google/uuid"
 	"github.com/spf13/cobra"
 )
 
@@ -22,6 +23,7 @@ func init() {
 	flags.IntVar(&cfg.MaxRunners, "max-runners", 10, "Maximum number of runners")
 	flags.IntVar(&cfg.MinRunners, "min-runners", 0, "Minimum number of runners")
 	flags.StringVar(&cfg.ScaleSetName, "name", "", "REQUIRED: Name of your scale set")
+	flags.StringSliceVar(&cfg.Labels, "labels", nil, "Labels for workflow targeting (comma-separated or repeated). Defaults to --name if not provided.")
 	flags.StringVar(&cfg.RunnerGroup, "runner-group", scaleset.DefaultRunnerGroup, "Name of the runner group your scale set should belong to")
 	flags.StringVar(&cfg.GitHubApp.ClientID, "app-client-id", "", "GitHub App client id")
 	flags.Int64Var(&cfg.GitHubApp.InstallationID, "app-installation-id", 0, "GitHub App installation ID")
@@ -54,7 +56,7 @@ func run(ctx context.Context, c Config) error {
 	logger := c.Logger()
 
 	// Create a new scaleset scalesetClient
-	scalesetClient, err := scaleset.NewClient(c.RegistrationURL, c.ActionsAuth())
+	scalesetClient, err := c.ScalesetClient()
 	if err != nil {
 		return fmt.Errorf("failed to create scaleset client: %w", err)
 	}
@@ -76,14 +78,8 @@ func run(ctx context.Context, c Config) error {
 	scaleSet, err := scalesetClient.CreateRunnerScaleSet(ctx, &scaleset.RunnerScaleSet{
 		Name:          c.ScaleSetName,
 		RunnerGroupID: runnerGroupID,
-		Labels: []scaleset.Label{
-			{
-				Name: c.ScaleSetName,
-				Type: "System",
-			},
-		},
+		Labels:        c.BuildLabels(),
 		RunnerSetting: scaleset.RunnerSetting{
-			Ephemeral:     true,
 			DisableUpdate: true,
 		},
 	})
@@ -92,13 +88,7 @@ func run(ctx context.Context, c Config) error {
 	}
 
 	// Set the user agent for the scaleset client now that we have the scale set ID
-	scalesetClient.SetUserAgent(scaleset.UserAgentInfo{
-		System:     "dockerscaleset",
-		Version:    "0.1.0",
-		CommitSHA:  "unknown",
-		ScaleSetID: scaleSet.ID,
-		Subsystem:  scaleSet.Name,
-	})
+	scalesetClient.SetSystemInfo(systemInfo(scaleSet.ID))
 
 	defer func() {
 		logger.Info(
@@ -137,10 +127,22 @@ func run(ctx context.Context, c Config) error {
 		return fmt.Errorf("failed to close image pull: %w", err)
 	}
 
+	// Get the name of the client which will be used as the owner
+	hostname, err := os.Hostname()
+	if err != nil {
+		hostname = uuid.NewString()
+		logger.Info("Failed to get hostname, fallback to uuid", "uuid", hostname, "error", err)
+	}
+
+	sessionClient, err := scalesetClient.MessageSessionClient(ctx, scaleSet.ID, hostname)
+	if err != nil {
+		return fmt.Errorf("failed to create message session client: %w", err)
+	}
+	defer sessionClient.Close(context.Background())
+
 	logger.Info("Initializing listener")
-	listener, err := listener.New(scalesetClient, listener.Config{
+	listener, err := listener.New(sessionClient, listener.Config{
 		ScaleSetID: scaleSet.ID,
-		MinRunners: c.MinRunners,
 		MaxRunners: c.MaxRunners,
 		Logger:     logger.WithGroup("listener"),
 	})
@@ -156,6 +158,7 @@ func run(ctx context.Context, c Config) error {
 		},
 		runnerImage:    c.RunnerImage,
 		minRunners:     c.MinRunners,
+		maxRunners:     c.MaxRunners,
 		dockerClient:   dockerClient,
 		scalesetClient: scalesetClient,
 		scaleSetID:     scaleSet.ID,

examples/dockerscaleset/scaler.go

@@ -9,8 +9,8 @@ import (
 	"github.com/actions/scaleset"
 	"github.com/actions/scaleset/listener"
 	"github.com/docker/docker/api/types/container"
+	dockerclient "github.com/docker/docker/client"
 	"github.com/google/uuid"
-	dockerclient "github.com/moby/moby/client"
 )
 
 type Scaler struct {
@@ -20,12 +20,13 @@ type Scaler struct {
 	dockerClient   *dockerclient.Client
 	scalesetClient *scaleset.Client
 	minRunners     int
+	maxRunners     int
 	logger         *slog.Logger
 }
 
 func (a *Scaler) HandleDesiredRunnerCount(ctx context.Context, count int) (int, error) {
 	currentCount := a.runners.count()
-	targetRunnerCount := min(a.minRunners + count)
+	targetRunnerCount := min(a.maxRunners, a.minRunners+count)
 
 	switch {
 	case targetRunnerCount == currentCount:
@@ -41,17 +42,12 @@ func (a *Scaler) HandleDesiredRunnerCount(ctx context.Context, count int) (int,
 			slog.Int("scaleUp", scaleUp),
 		)
 
-		var errs []error
 		for range scaleUp {
 			if _, err := a.startRunner(ctx); err != nil {
-				errs = append(errs, err)
+				return 0, fmt.Errorf("failed to start runner: %w", err)
 			}
 		}
 
-		for _, err := range errs {
-			// TODO: should we return error?
-			a.logger.Error("Failed to start runner", slog.String("error", err.Error()))
-		}
 		return a.runners.count(), nil
 	default:
 		// No need to handle scale down events, since:
@@ -146,7 +142,7 @@ func (a *Scaler) shutdown(ctx context.Context) {
 	clear(a.runners.busy)
 }
 
-var _ listener.Handler = (*Scaler)(nil)
+var _ listener.Scaler = (*Scaler)(nil)
 
 type runnerState struct {
 	mu   sync.Mutex

go.mod

@@ -3,20 +3,90 @@ module github.com/actions/scaleset
 go 1.25.3
 
 require (
+	github.com/docker/docker v28.5.2+incompatible
 	github.com/golang-jwt/jwt/v4 v4.5.2
+	github.com/google/go-github/v79 v79.0.0
 	github.com/google/uuid v1.6.0
 	github.com/hashicorp/go-retryablehttp v0.7.8
+	github.com/spf13/cobra v1.10.2
 	github.com/stretchr/testify v1.11.1
-	golang.org/x/net v0.46.0
+	golang.org/x/net v0.52.0
 )
 
 require (
-	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/Microsoft/go-winio v0.6.2 // indirect
+	github.com/brunoga/deep v1.2.4 // indirect
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/containerd/errdefs v1.0.0 // indirect
+	github.com/containerd/errdefs/pkg v0.3.0 // indirect
+	github.com/containerd/log v0.1.0 // indirect
+	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
+	github.com/distribution/reference v0.6.0 // indirect
+	github.com/docker/go-connections v0.6.0 // indirect
+	github.com/docker/go-units v0.5.0 // indirect
+	github.com/fatih/structs v1.1.0 // indirect
+	github.com/felixge/httpsnoop v1.0.4 // indirect
+	github.com/fsnotify/fsnotify v1.8.0 // indirect
+	github.com/go-logr/logr v1.4.3 // indirect
+	github.com/go-logr/stdr v1.2.2 // indirect
+	github.com/go-viper/mapstructure/v2 v2.4.0 // indirect
+	github.com/google/go-querystring v1.1.0 // indirect
 	github.com/hashicorp/go-cleanhttp v0.5.2 // indirect
-	github.com/kr/pretty v0.3.1 // indirect
-	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/rogpeppe/go-internal v1.13.1 // indirect
-	golang.org/x/text v0.30.0 // indirect
-	gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect
+	github.com/huandu/xstrings v1.5.0 // indirect
+	github.com/inconshreveable/mousetrap v1.1.0 // indirect
+	github.com/jedib0t/go-pretty/v6 v6.6.7 // indirect
+	github.com/knadh/koanf/maps v0.1.2 // indirect
+	github.com/knadh/koanf/parsers/yaml v0.1.0 // indirect
+	github.com/knadh/koanf/providers/env v1.0.0 // indirect
+	github.com/knadh/koanf/providers/file v1.1.2 // indirect
+	github.com/knadh/koanf/providers/posflag v0.1.0 // indirect
+	github.com/knadh/koanf/providers/structs v0.1.0 // indirect
+	github.com/knadh/koanf/v2 v2.3.0 // indirect
+	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/mattn/go-runewidth v0.0.16 // indirect
+	github.com/mitchellh/copystructure v1.2.0 // indirect
+	github.com/mitchellh/reflectwalk v1.0.2 // indirect
+	github.com/moby/docker-image-spec v1.3.1 // indirect
+	github.com/moby/sys/atomicwriter v0.1.0 // indirect
+	github.com/moby/term v0.5.2 // indirect
+	github.com/morikuni/aec v1.0.0 // indirect
+	github.com/opencontainers/go-digest v1.0.0 // indirect
+	github.com/opencontainers/image-spec v1.1.1 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
+	github.com/rivo/uniseg v0.4.7 // indirect
+	github.com/rs/zerolog v1.33.0 // indirect
+	github.com/spf13/pflag v1.0.10 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
+	github.com/vektra/mockery/v3 v3.6.1 // indirect
+	github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f // indirect
+	github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 // indirect
+	github.com/xeipuuv/gojsonschema v1.2.0 // indirect
+	go.opentelemetry.io/auto/sdk v1.2.1 // indirect
+	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 // indirect
+	go.opentelemetry.io/otel v1.39.0 // indirect
+	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0 // indirect
+	go.opentelemetry.io/otel/metric v1.39.0 // indirect
+	go.opentelemetry.io/otel/sdk v1.39.0 // indirect
+	go.opentelemetry.io/otel/sdk/metric v1.39.0 // indirect
+	go.opentelemetry.io/otel/trace v1.39.0 // indirect
+	golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56 // indirect
+	golang.org/x/mod v0.33.0 // indirect
+	golang.org/x/sync v0.20.0 // indirect
+	golang.org/x/sys v0.42.0 // indirect
+	golang.org/x/telemetry v0.0.0-20260209163413-e7419c687ee4 // indirect
+	golang.org/x/term v0.41.0 // indirect
+	golang.org/x/text v0.35.0 // indirect
+	golang.org/x/time v0.14.0 // indirect
+	golang.org/x/tools v0.42.0 // indirect
+	google.golang.org/grpc v1.79.3 // indirect
+	google.golang.org/protobuf v1.36.10 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
+	gotest.tools/v3 v3.5.2 // indirect
+)
+
+tool (
+	github.com/vektra/mockery/v3
+	golang.org/x/tools/cmd/deadcode
 )

go.sum

@@ -1,42 +1,209 @@
-github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
-github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
-github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c h1:udKWzYgxTojEKWjV8V+WSxDXJ4NFATAsZjh8iIbsQIg=
+github.com/Azure/go-ansiterm v0.0.0-20250102033503-faa5f7b0171c/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E=
+github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
+github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
+github.com/brunoga/deep v1.2.4 h1:Aj9E9oUbE+ccbyh35VC/NHlzzjfIVU69BXu2mt2LmL8=
+github.com/brunoga/deep v1.2.4/go.mod h1:GDV6dnXqn80ezsLSZ5Wlv1PdKAWAO4L5PnKYtv2dgaI=
+github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
+github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
+github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
+github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI=
+github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M=
+github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE=
+github.com/containerd/errdefs/pkg v0.3.0/go.mod h1:NJw6s9HwNuRhnjJhM7pylWwMyAkmCQvQ4GpJHEqRLVk=
+github.com/containerd/log v0.1.0 h1:TCJt7ioM2cr/tfR8GPbGf9/VRAX8D2B4PjzCpfX540I=
+github.com/containerd/log v0.1.0/go.mod h1:VRRf09a7mHDIRezVKTRCrOq78v577GXq3bSa3EhrzVo=
+github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
+github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
+github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk=
+github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E=
+github.com/docker/docker v28.5.2+incompatible h1:DBX0Y0zAjZbSrm1uzOkdr1onVghKaftjlSWt4AFexzM=
+github.com/docker/docker v28.5.2+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
+github.com/docker/go-connections v0.6.0 h1:LlMG9azAe1TqfR7sO+NJttz1gy6KO7VJBh+pMmjSD94=
+github.com/docker/go-connections v0.6.0/go.mod h1:AahvXYshr6JgfUJGdDCs2b5EZG/vmaMAntpSFH5BFKE=
+github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4=
+github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
 github.com/fatih/color v1.16.0 h1:zmkK9Ngbjj+K0yRhTVONQh1p/HknKYSlNT+vZCzyokM=
 github.com/fatih/color v1.16.0/go.mod h1:fL2Sau1YI5c0pdGEVCbKQbLXB6edEj1ZgiY4NijnWvE=
+github.com/fatih/structs v1.1.0 h1:Q7juDM0QtcnhCpeyLGQKyg4TOIghuNXrkL32pHAUMxo=
+github.com/fatih/structs v1.1.0/go.mod h1:9NiDSp5zOcgEDl+j00MP/WkGVPOlPRLejGD8Ga6PJ7M=
+github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
+github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
+github.com/fsnotify/fsnotify v1.8.0 h1:dAwr6QBTBZIkG8roQaJjGof0pp0EeF+tNV7YBP3F/8M=
+github.com/fsnotify/fsnotify v1.8.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0=
+github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
+github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
+github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
+github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
+github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
+github.com/go-viper/mapstructure/v2 v2.4.0 h1:EBsztssimR/CONLSZZ04E8qAkxNYq4Qp9LvH92wZUgs=
+github.com/go-viper/mapstructure/v2 v2.4.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
+github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/golang-jwt/jwt/v4 v4.5.2 h1:YtQM7lnr8iZ+j5q71MGKkNw9Mn7AjHM68uc9g5fXeUI=
 github.com/golang-jwt/jwt/v4 v4.5.2/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0=
+github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
+github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
+github.com/google/go-github/v79 v79.0.0 h1:MdodQojuFPBhmtwHiBcIGLw/e/wei2PvFX9ndxK0X4Y=
+github.com/google/go-github/v79 v79.0.0/go.mod h1:OAFbNhq7fQwohojb06iIIQAB9CBGYLq999myfUFnrS4=
+github.com/google/go-querystring v1.1.0 h1:AnCroh3fv4ZBgVIf1Iwtovgjaw/GiKJo8M8yD/fhyJ8=
+github.com/google/go-querystring v1.1.0/go.mod h1:Kcdr2DB4koayq7X8pmAG4sNG59So17icRSOU623lUBU=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1 h1:e9Rjr40Z98/clHv5Yg79Is0NtosR5LXRvdr7o/6NwbA=
+github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1/go.mod h1:tIxuGz/9mpox++sgp9fJjHO0+q1X9/UOWd798aAm22M=
 github.com/hashicorp/go-cleanhttp v0.5.2 h1:035FKYIWjmULyFRBKPs8TBQoi0x6d9G4xc9neXJWAZQ=
 github.com/hashicorp/go-cleanhttp v0.5.2/go.mod h1:kO/YDlP8L1346E6Sodw+PrpBSV4/SoxCXGY6BqNFT48=
 github.com/hashicorp/go-hclog v1.6.3 h1:Qr2kF+eVWjTiYmU7Y31tYlP1h0q/X3Nl3tPGdaB11/k=
 github.com/hashicorp/go-hclog v1.6.3/go.mod h1:W4Qnvbt70Wk/zYJryRzDRU/4r0kIg0PVHBcfoyhpF5M=
 github.com/hashicorp/go-retryablehttp v0.7.8 h1:ylXZWnqa7Lhqpk0L1P1LzDtGcCR0rPVUrx/c8Unxc48=
 github.com/hashicorp/go-retryablehttp v0.7.8/go.mod h1:rjiScheydd+CxvumBsIrFKlx3iS0jrZ7LvzFGFmuKbw=
+github.com/huandu/xstrings v1.5.0 h1:2ag3IFq9ZDANvthTwTiqSSZLjDc+BedvHPAp5tJy2TI=
+github.com/huandu/xstrings v1.5.0/go.mod h1:y5/lhBue+AyNmUVz9RLU9xbLR0o4KIIExikq4ovT0aE=
+github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
+github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
+github.com/jedib0t/go-pretty/v6 v6.6.7 h1:m+LbHpm0aIAPLzLbMfn8dc3Ht8MW7lsSO4MPItz/Uuo=
+github.com/jedib0t/go-pretty/v6 v6.6.7/go.mod h1:YwC5CE4fJ1HFUDeivSV1r//AmANFHyqczZk+U6BDALU=
+github.com/knadh/koanf/maps v0.1.2 h1:RBfmAW5CnZT+PJ1CVc1QSJKf4Xu9kxfQgYVQSu8hpbo=
+github.com/knadh/koanf/maps v0.1.2/go.mod h1:npD/QZY3V6ghQDdcQzl1W4ICNVTkohC8E73eI2xW4yI=
+github.com/knadh/koanf/parsers/yaml v0.1.0 h1:ZZ8/iGfRLvKSaMEECEBPM1HQslrZADk8fP1XFUxVI5w=
+github.com/knadh/koanf/parsers/yaml v0.1.0/go.mod h1:cvbUDC7AL23pImuQP0oRw/hPuccrNBS2bps8asS0CwY=
+github.com/knadh/koanf/providers/env v1.0.0 h1:ufePaI9BnWH+ajuxGGiJ8pdTG0uLEUWC7/HDDPGLah0=
+github.com/knadh/koanf/providers/env v1.0.0/go.mod h1:mzFyRZueYhb37oPmC1HAv/oGEEuyvJDA98r3XAa8Gak=
+github.com/knadh/koanf/providers/file v1.1.2 h1:aCC36YGOgV5lTtAFz2qkgtWdeQsgfxUkxDOe+2nQY3w=
+github.com/knadh/koanf/providers/file v1.1.2/go.mod h1:/faSBcv2mxPVjFrXck95qeoyoZ5myJ6uxN8OOVNJJCI=
+github.com/knadh/koanf/providers/posflag v0.1.0 h1:mKJlLrKPcAP7Ootf4pBZWJ6J+4wHYujwipe7Ie3qW6U=
+github.com/knadh/koanf/providers/posflag v0.1.0/go.mod h1:SYg03v/t8ISBNrMBRMlojH8OsKowbkXV7giIbBVgbz0=
+github.com/knadh/koanf/providers/structs v0.1.0 h1:wJRteCNn1qvLtE5h8KQBvLJovidSdntfdyIbbCzEyE0=
+github.com/knadh/koanf/providers/structs v0.1.0/go.mod h1:sw2YZ3txUcqA3Z27gPlmmBzWn1h8Nt9O6EP/91MkcWE=
+github.com/knadh/koanf/v2 v2.3.0 h1:Qg076dDRFHvqnKG97ZEsi9TAg2/nFTa9hCdcSa1lvlM=
+github.com/knadh/koanf/v2 v2.3.0/go.mod h1:gRb40VRAbd4iJMYYD5IxZ6hfuopFcXBpc9bbQpZwo28=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
 github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
+github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
+github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
 github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
-github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
-github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc=
+github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
+github.com/mitchellh/copystructure v1.2.0 h1:vpKXTN4ewci03Vljg/q9QvCGUDttBOGBIa15WveJJGw=
+github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HKCj9FbZEVFJRxO9s=
+github.com/mitchellh/reflectwalk v1.0.2 h1:G2LzWKi524PWgd3mLHV8Y5k7s6XUvT0Gef6zxSIeXaQ=
+github.com/mitchellh/reflectwalk v1.0.2/go.mod h1:mSTlrgnPZtwu0c4WaC2kGObEpuNDbx0jmZXqmk4esnw=
+github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0=
+github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo=
+github.com/moby/sys/atomicwriter v0.1.0 h1:kw5D/EqkBwsBFi0ss9v1VG3wIkVhzGvLklJ+w3A14Sw=
+github.com/moby/sys/atomicwriter v0.1.0/go.mod h1:Ul8oqv2ZMNHOceF643P6FKPXeCmYtlQMvpizfsSoaWs=
+github.com/moby/sys/sequential v0.6.0 h1:qrx7XFUd/5DxtqcoH1h438hF5TmOvzC/lspjy7zgvCU=
+github.com/moby/sys/sequential v0.6.0/go.mod h1:uyv8EUTrca5PnDsdMGXhZe6CCe8U/UiTWd+lL+7b/Ko=
+github.com/moby/term v0.5.2 h1:6qk3FJAFDs6i/q3W/pQ97SX192qKfZgGjCQqfCJkgzQ=
+github.com/moby/term v0.5.2/go.mod h1:d3djjFCrjnB+fl8NJux+EJzu0msscUP+f8it8hPkFLc=
+github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A=
+github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc=
+github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
+github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
+github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040=
+github.com/opencontainers/image-spec v1.1.1/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M=
+github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
+github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs=
-github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
-github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
+github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
+github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
+github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
+github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
+github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
+github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
+github.com/rs/xid v1.5.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg=
+github.com/rs/zerolog v1.33.0 h1:1cU2KZkvPxNyfgEmhHAz/1A9Bz+llsdYzklWFzgp0r8=
+github.com/rs/zerolog v1.33.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss=
+github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
+github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
+github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU=
+github.com/spf13/cobra v1.10.2/go.mod h1:7C1pvHqHw5A4vrJfjNwvOdzYu0Gml16OCs2GRiTUUS4=
+github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
+github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk=
+github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
+github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
-golang.org/x/net v0.46.0 h1:giFlY12I07fugqwPuWJi68oOnpfqFnJIJzaIIm2JVV4=
-golang.org/x/net v0.46.0/go.mod h1:Q9BGdFy1y4nkUwiLvT5qtyhAnEHgnQ/zd8PfU6nc210=
-golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ=
-golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
-golang.org/x/text v0.30.0 h1:yznKA/E9zq54KzlzBEAWn1NXSQ8DIp/NYMy88xJjl4k=
-golang.org/x/text v0.30.0/go.mod h1:yDdHFIX9t+tORqspjENWgzaCVXgk0yYnYuSZ8UzzBVM=
+github.com/vektra/mockery/v3 v3.6.1 h1:YyqAXihdNML8y6SJnvPKYr+2HAHvBjdvqFu/fMYlX8g=
+github.com/vektra/mockery/v3 v3.6.1/go.mod h1:Oti3Df0WP8wwT31yuVri3QNsDeMUQU5Q4QEg8EabaBw=
+github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f h1:J9EGpcZtP0E/raorCMxlFGSTBrsSlaDGf3jU/qvAE2c=
+github.com/xeipuuv/gojsonpointer v0.0.0-20180127040702-4e3ac2762d5f/go.mod h1:N2zxlSyiKSe5eX1tZViRH5QA0qijqEDrYZiPEAiq3wU=
+github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415 h1:EzJWgHovont7NscjpAxXsDA8S8BMYve8Y5+7cuRE7R0=
+github.com/xeipuuv/gojsonreference v0.0.0-20180127040603-bd5ef7bd5415/go.mod h1:GwrjFmJcFw6At/Gs6z4yjiIwzuJ1/+UwLxMQDVQXShQ=
+github.com/xeipuuv/gojsonschema v1.2.0 h1:LhYJRs+L4fBtjZUfuSZIKGeVu0QRy8e5Xi7D17UxZ74=
+github.com/xeipuuv/gojsonschema v1.2.0/go.mod h1:anYRn/JVcOK2ZgGU+IjEV4nwlhoK5sQluxsYJ78Id3Y=
+go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
+go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
+go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 h1:sbiXRNDSWJOTobXh5HyQKjq6wUC5tNybqjIqDpAY4CU=
+go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0/go.mod h1:69uWxva0WgAA/4bu2Yy70SLDBwZXuQ6PbBpbsa5iZrQ=
+go.opentelemetry.io/otel v1.39.0 h1:8yPrr/S0ND9QEfTfdP9V+SiwT4E0G7Y5MO7p85nis48=
+go.opentelemetry.io/otel v1.39.0/go.mod h1:kLlFTywNWrFyEdH0oj2xK0bFYZtHRYUdv1NklR/tgc8=
+go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.35.0 h1:1fTNlAIJZGWLP5FVu0fikVry1IsiUnXjf7QFvoNN3Xw=
+go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.35.0/go.mod h1:zjPK58DtkqQFn+YUMbx0M2XV3QgKU0gS9LeGohREyK4=
+go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0 h1:xJ2qHD0C1BeYVTLLR9sX12+Qb95kfeD/byKj6Ky1pXg=
+go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0/go.mod h1:u5BF1xyjstDowA1R5QAO9JHzqK+ublenEW/dyqTjBVk=
+go.opentelemetry.io/otel/metric v1.39.0 h1:d1UzonvEZriVfpNKEVmHXbdf909uGTOQjA0HF0Ls5Q0=
+go.opentelemetry.io/otel/metric v1.39.0/go.mod h1:jrZSWL33sD7bBxg1xjrqyDjnuzTUB0x1nBERXd7Ftcs=
+go.opentelemetry.io/otel/sdk v1.39.0 h1:nMLYcjVsvdui1B/4FRkwjzoRVsMK8uL/cj0OyhKzt18=
+go.opentelemetry.io/otel/sdk v1.39.0/go.mod h1:vDojkC4/jsTJsE+kh+LXYQlbL8CgrEcwmt1ENZszdJE=
+go.opentelemetry.io/otel/sdk/metric v1.39.0 h1:cXMVVFVgsIf2YL6QkRF4Urbr/aMInf+2WKg+sEJTtB8=
+go.opentelemetry.io/otel/sdk/metric v1.39.0/go.mod h1:xq9HEVH7qeX69/JnwEfp6fVq5wosJsY1mt4lLfYdVew=
+go.opentelemetry.io/otel/trace v1.39.0 h1:2d2vfpEDmCJ5zVYz7ijaJdOF59xLomrvj7bjt6/qCJI=
+go.opentelemetry.io/otel/trace v1.39.0/go.mod h1:88w4/PnZSazkGzz/w84VHpQafiU4EtqqlVdxWy+rNOA=
+go.opentelemetry.io/proto/otlp v1.5.0 h1:xJvq7gMzB31/d406fB8U5CBdyQGw4P399D1aQWU/3i4=
+go.opentelemetry.io/proto/otlp v1.5.0/go.mod h1:keN8WnHxOy8PG0rQZjJJ5A2ebUoafqWp0eVQ4yIXvJ4=
+go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
+golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56 h1:2dVuKD2vS7b0QIHQbpyTISPd0LeHDbnYEryqj5Q1ug8=
+golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56/go.mod h1:M4RDyNAINzryxdtnbRXRL/OHtkFuWGRjvuhBJpk2IlY=
+golang.org/x/mod v0.33.0 h1:tHFzIWbBifEmbwtGz65eaWyGiGZatSrT9prnU8DbVL8=
+golang.org/x/mod v0.33.0/go.mod h1:swjeQEj+6r7fODbD2cqrnje9PnziFuw4bmLbBZFrQ5w=
+golang.org/x/net v0.52.0 h1:He/TN1l0e4mmR3QqHMT2Xab3Aj3L9qjbhRm78/6jrW0=
+golang.org/x/net v0.52.0/go.mod h1:R1MAz7uMZxVMualyPXb+VaqGSa3LIaUqk0eEt3w36Sw=
+golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
+golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
+golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo=
+golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
+golang.org/x/telemetry v0.0.0-20260209163413-e7419c687ee4 h1:bTLqdHv7xrGlFbvf5/TXNxy/iUwwdkjhqQTJDjW7aj0=
+golang.org/x/telemetry v0.0.0-20260209163413-e7419c687ee4/go.mod h1:g5NllXBEermZrmR51cJDQxmJUHUOfRAaNyWBM+R+548=
+golang.org/x/term v0.41.0 h1:QCgPso/Q3RTJx2Th4bDLqML4W6iJiaXFq2/ftQF13YU=
+golang.org/x/term v0.41.0/go.mod h1:3pfBgksrReYfZ5lvYM0kSO0LIkAl4Yl2bXOkKP7Ec2A=
+golang.org/x/text v0.35.0 h1:JOVx6vVDFokkpaq1AEptVzLTpDe9KGpj5tR4/X+ybL8=
+golang.org/x/text v0.35.0/go.mod h1:khi/HExzZJ2pGnjenulevKNX1W67CUy0AsXcNubPGCA=
+golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI=
+golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4=
+golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k=
+golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0=
+golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 h1:fCvbg86sFXwdrl5LgVcTEvNC+2txB5mgROGmRL5mrls=
+google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217/go.mod h1:+rXWjjaukWZun3mLfjmVnQi18E1AsFbDN9QdJ5YXLto=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 h1:gRkg/vSppuSQoDjxyiGfN4Upv/h/DQmIR10ZU8dh4Ww=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk=
+google.golang.org/grpc v1.79.3 h1:sybAEdRIEtvcD68Gx7dmnwjZKlyfuc61Dyo9pGXXkKE=
+google.golang.org/grpc v1.79.3/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ=
+google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE=
+google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
-gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gotest.tools/v3 v3.5.2 h1:7koQfIKdy+I8UTetycgUqXWSDwpgv193Ka+qRsmBY8Q=
+gotest.tools/v3 v3.5.2/go.mod h1:LtdLGcnqToBH83WByAAi/wiwSFCArdFIUV/xxN4pcjA=

listener/listener.go

@@ -3,27 +3,24 @@ package listener
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"log/slog"
-	"net/http"
-	"os"
-	"time"
+	"math"
+	"sync/atomic"
 
 	"github.com/actions/scaleset"
 	"github.com/google/uuid"
 )
 
-const (
-	sessionCreationMaxRetries = 10
-)
-
+// Config holds the configuration for the Listener.
 type Config struct {
+	// ScaleSetID is the ID of the runner scale set to listen to.
 	ScaleSetID int
-	MinRunners int
+	// MaxRunners is the capacity of runners that can be handled at once.
 	MaxRunners int
-	Logger     *slog.Logger
+	// Logger is the logger to use for logging. Default is a no-op logger.
+	Logger *slog.Logger
 }
 
 func (c *Config) defaults() {
@@ -32,99 +29,146 @@ func (c *Config) defaults() {
 	}
 }
 
+// Validate returns an error if the configuration is invalid.
 func (c *Config) Validate() error {
 	c.defaults()
 
 	if c.ScaleSetID == 0 {
 		return errors.New("scaleSetID is required")
 	}
-	if c.MinRunners < 0 {
-		return errors.New("minRunners must be greater than or equal to 0")
-	}
-	if c.MaxRunners < 0 {
-		return errors.New("maxRunners must be greater than or equal to 0")
-	}
-	if c.MaxRunners > 0 && c.MinRunners > c.MaxRunners {
-		return errors.New("minRunners must be less than or equal to maxRunners")
+	if c.MaxRunners < 0 || c.MaxRunners > math.MaxInt32 {
+		return errors.New("maxRunners must be between 0 and MaxInt32")
 	}
 	return nil
 }
 
+// Client defines the interface for communicating with the scaleset API.
+// In most cases, it should be scaleset.Client from the scaleset package.
+// This interface is defined to allow for easier testing and mocking, as well
+// as allowing wrappers around the scaleset client if needed.
+type Client interface {
+	GetMessage(ctx context.Context, lastMessageID, maxCapacity int) (*scaleset.RunnerScaleSetMessage, error)
+	DeleteMessage(ctx context.Context, messageID int) error
+	AcquireJobs(ctx context.Context, requestIDs []int64) ([]int64, error)
+	Session() scaleset.RunnerScaleSetSession
+}
+
+// MetricsRecorder defines the hook methods that will be called by the listener at
+// various points in the message handling process. This allows for custom
+// metrics to be collected without coupling the listener to a specific metrics
+// implementation. The methods in this interface will be called with relevant
+// information about the message handling process, such as the number of jobs
+// started/completed, the desired runner count, and any errors that occur.
+// Implementers can use this information to track the performance and behavior
+// of the listener and the scaleset service.
+type MetricsRecorder interface {
+	RecordStatistics(statistics *scaleset.RunnerScaleSetStatistic)
+	RecordJobStarted(msg *scaleset.JobStarted)
+	RecordJobCompleted(msg *scaleset.JobCompleted)
+	RecordDesiredRunners(count int)
+}
+
+type discardMetricsRecorder struct{}
+
+func (d *discardMetricsRecorder) RecordStatistics(statistics *scaleset.RunnerScaleSetStatistic) {}
+func (d *discardMetricsRecorder) RecordJobStarted(msg *scaleset.JobStarted)                     {}
+func (d *discardMetricsRecorder) RecordJobCompleted(msg *scaleset.JobCompleted)                 {}
+func (d *discardMetricsRecorder) RecordDesiredRunners(count int)                                {}
+
+// Listener listens for messages from the scaleset service and handles them. It automatically handles session
+// creation/deletion/refreshing and message polling and acking.
 type Listener struct {
 	// The main client responsible for communicating with the scaleset service
-	client *scaleset.Client
+	client          Client
+	metricsRecorder MetricsRecorder
 
 	// Configuration for the listener
-	scaleSetID int
-	minRunners int
-	maxRunners int
-
-	// lastMessageID keeps track of the last processed message ID
-	lastMessageID int64
-	// hostname of the current machine
-	hostname string
-	// session represents the current message session
-	session *scaleset.RunnerScaleSetSession
+	scaleSetID       int
+	maxRunners       atomic.Uint32
+	latestStatistics *scaleset.RunnerScaleSetStatistic
 
 	// configuration for the listener
 	logger *slog.Logger
 }
 
-func New(client *scaleset.Client, config Config) (*Listener, error) {
+type Option func(*Listener)
+
+// WithMetricsRecorder sets the MetricsRecorder for the listener. If not set, a no-op recorder will be used.
+// If the nil is passed, the MetricsRecorder will not be updated and the existing one will be used (which is a no-op by default).
+func WithMetricsRecorder(recorder MetricsRecorder) Option {
+	return func(l *Listener) {
+		if recorder == nil {
+			return
+		}
+		l.metricsRecorder = recorder
+	}
+}
+
+// SetMaxRunners sets the capacity of the scaleset. It is concurrently
+// safe to update the max runners during listener.Run.
+func (l *Listener) SetMaxRunners(count int) {
+	l.maxRunners.Store(uint32(count))
+}
+
+// New creates a new Listener with the given configuration.
+func New(client Client, config Config, options ...Option) (*Listener, error) {
 	if client == nil {
 		return nil, errors.New("client is required")
 	}
 
 	if err := config.Validate(); err != nil {
-		return nil, err
+		return nil, fmt.Errorf("invalid config: %w", err)
 	}
 
-	hostname, err := os.Hostname()
-	if err != nil {
-		hostname = uuid.NewString()
-		config.Logger.Info("Failed to get hostname, fallback to uuid", "uuid", hostname, "error", err)
+	listener := &Listener{
+		client:          client,
+		metricsRecorder: &discardMetricsRecorder{},
+		scaleSetID:      config.ScaleSetID,
+		logger:          config.Logger,
+	}
+	listener.SetMaxRunners(config.MaxRunners)
+
+	for _, option := range options {
+		option(listener)
 	}
 
-	return &Listener{
-		client:     client,
-		scaleSetID: config.ScaleSetID,
-		minRunners: config.MinRunners,
-		maxRunners: config.MaxRunners,
-		hostname:   hostname,
-		logger:     config.Logger,
-	}, nil
+	return listener, nil
 }
 
+// Scaler defines the interface for handling scale set messages.
 type Scaler interface {
 	HandleJobStarted(ctx context.Context, jobInfo *scaleset.JobStarted) error
 	HandleJobCompleted(ctx context.Context, jobInfo *scaleset.JobCompleted) error
 	HandleDesiredRunnerCount(ctx context.Context, count int) (int, error)
 }
 
-func (l *Listener) Run(ctx context.Context, handler Scaler) error {
-	l.logger.Info("Creating message session")
-	if err := l.createSession(ctx); err != nil {
-		return fmt.Errorf("failed to create session: %w", err)
-	}
+// Run starts the listener and processes messages using the provided scaler.
+func (l *Listener) Run(ctx context.Context, scaler Scaler) error {
+	{
+		initialSession := l.client.Session()
 
-	defer func() {
-		l.logger.Debug("Deleting message session")
-		if err := l.deleteMessageSession(); err != nil {
-			l.logger.Error("failed to delete message session", "error", err.Error())
+		if initialSession.SessionID == uuid.Nil {
+			return fmt.Errorf("initial session is nil")
 		}
-	}()
 
-	if l.session.Statistics == nil {
-		return fmt.Errorf("session statistics is nil")
-	}
-
-	l.logger.Info("Message session created; listening for messages", "sessionID", l.session.SessionID)
-
-	// Handle initial statistics
-	if _, err := handler.HandleDesiredRunnerCount(ctx, l.session.Statistics.TotalAssignedJobs); err != nil {
-		return fmt.Errorf("handling initial message failed: %w", err)
+		if initialSession.Statistics == nil {
+			return fmt.Errorf("session statistics is nil")
+		}
+
+		l.handleStatistics(ctx, initialSession.Statistics)
+
+		l.logger.Info(
+			"Handling initial session statistics",
+			slog.Int("totalAssignedJobs", initialSession.Statistics.TotalAssignedJobs),
+		)
+		desiredCount, err := scaler.HandleDesiredRunnerCount(ctx, initialSession.Statistics.TotalAssignedJobs)
+		if err != nil {
+			return fmt.Errorf("handling initial message failed: %w", err)
+		}
+		l.metricsRecorder.RecordDesiredRunners(desiredCount)
 	}
 
+	var lastMessageID int
 	for {
 		select {
 		case <-ctx.Done():
@@ -132,13 +176,18 @@ func (l *Listener) Run(ctx context.Context, handler Scaler) error {
 		default:
 		}
 
-		msg, err := l.getMessage(ctx)
+		l.logger.Info("Getting next message", slog.Int("lastMessageID", lastMessageID))
+		msg, err := l.client.GetMessage(
+			ctx,
+			lastMessageID,
+			int(l.maxRunners.Load()),
+		)
 		if err != nil {
 			return fmt.Errorf("failed to get message: %w", err)
 		}
 
 		if msg == nil {
-			_, err := handler.HandleDesiredRunnerCount(ctx, 0)
+			_, err := scaler.HandleDesiredRunnerCount(ctx, l.latestStatistics.TotalAssignedJobs)
 			if err != nil {
 				return fmt.Errorf("handling nil message failed: %w", err)
 			}
@@ -146,260 +195,68 @@ func (l *Listener) Run(ctx context.Context, handler Scaler) error {
 			continue
 		}
 
+		lastMessageID = msg.MessageID
+
 		// Remove cancellation from the context to avoid cancelling the message handling.
-		if err := l.handleMessage(context.WithoutCancel(ctx), handler, msg); err != nil {
+		if err := l.handleMessage(context.WithoutCancel(ctx), scaler, msg); err != nil {
 			return fmt.Errorf("failed to handle message: %w", err)
 		}
 	}
 }
 
 func (l *Listener) handleMessage(ctx context.Context, handler Scaler, msg *scaleset.RunnerScaleSetMessage) error {
-	parsedMsg, err := l.parseMessage(ctx, msg)
-	if err != nil {
-		return fmt.Errorf("failed to parse message: %w", err)
-	}
+	l.handleStatistics(ctx, msg.Statistics)
 
-	l.lastMessageID = msg.MessageID
-
-	if err := l.deleteLastMessage(ctx); err != nil {
+	if err := l.client.DeleteMessage(ctx, msg.MessageID); err != nil {
 		return fmt.Errorf("failed to delete message: %w", err)
 	}
 
-	for _, jobStarted := range parsedMsg.jobsStarted {
+	if len(msg.JobAvailableMessages) > 0 {
+		if err := l.acquireAvailableJobs(ctx, msg.JobAvailableMessages); err != nil {
+			return fmt.Errorf("failed to acquire available jobs: %w", err)
+		}
+	}
+
+	for _, jobStarted := range msg.JobStartedMessages {
+		l.metricsRecorder.RecordJobStarted(jobStarted)
 		if err := handler.HandleJobStarted(ctx, jobStarted); err != nil {
 			return fmt.Errorf("failed to handle job started: %w", err)
 		}
 	}
-	for _, jobCompleted := range parsedMsg.jobsCompleted {
+	for _, jobCompleted := range msg.JobCompletedMessages {
+		l.metricsRecorder.RecordJobCompleted(jobCompleted)
 		if err := handler.HandleJobCompleted(ctx, jobCompleted); err != nil {
 			return fmt.Errorf("failed to handle job completed: %w", err)
 		}
 	}
 
-	if _, err := handler.HandleDesiredRunnerCount(ctx, parsedMsg.statistics.TotalAssignedJobs); err != nil {
+	desiredCount, err := handler.HandleDesiredRunnerCount(ctx, msg.Statistics.TotalAssignedJobs)
+	if err != nil {
 		return fmt.Errorf("failed to handle desired runner count: %w", err)
 	}
+	l.metricsRecorder.RecordDesiredRunners(desiredCount)
 
 	return nil
 }
 
-func (l *Listener) createSession(ctx context.Context) error {
-	var session *scaleset.RunnerScaleSetSession
-	var retries int
-
-	for {
-		var err error
-		session, err = l.client.CreateMessageSession(ctx, l.scaleSetID, l.hostname)
-		if err == nil {
-			break
-		}
-
-		clientErr := &scaleset.HttpClientSideError{}
-		if !errors.As(err, &clientErr) {
-			return fmt.Errorf("failed to create session: %w", err)
-		}
-
-		if clientErr.Code != http.StatusConflict {
-			return fmt.Errorf("failed to create session: %w", err)
-		}
-
-		retries++
-		if retries >= sessionCreationMaxRetries {
-			return fmt.Errorf("failed to create session after %d retries: %w", retries, err)
-		}
-
-		l.logger.Info("Unable to create message session. Will try again in 30 seconds", "error", err.Error())
-
-		select {
-		case <-ctx.Done():
-			return fmt.Errorf("context cancelled: %w", ctx.Err())
-		case <-time.After(30 * time.Second):
-		}
+func (l *Listener) acquireAvailableJobs(ctx context.Context, jobsAvailable []*scaleset.JobAvailable) error {
+	ids := make([]int64, 0, len(jobsAvailable))
+	for _, job := range jobsAvailable {
+		ids = append(ids, job.RunnerRequestID)
 	}
 
-	statistics, err := json.Marshal(session.Statistics)
+	l.logger.Info("Acquiring jobs", slog.Int("count", len(ids)))
+
+	acquired, err := l.client.AcquireJobs(ctx, ids)
 	if err != nil {
-		return fmt.Errorf("failed to marshal statistics: %w", err)
+		return fmt.Errorf("acquiring jobs: %w", err)
 	}
-	l.logger.Info("Current runner scale set statistics.", "statistics", string(statistics))
-
-	l.session = session
 
+	l.logger.Info("Jobs acquired", slog.Int("count", len(acquired)))
 	return nil
 }
 
-func (l *Listener) getMessage(ctx context.Context) (*scaleset.RunnerScaleSetMessage, error) {
-	l.logger.Info("Getting next message", "lastMessageID", l.lastMessageID)
-	msg, err := l.client.GetMessage(
-		ctx,
-		l.session.MessageQueueURL,
-		l.session.MessageQueueAccessToken,
-		l.lastMessageID,
-		l.maxRunners,
-	)
-	if err == nil { // if NO error
-		return msg, nil
-	}
-
-	expiredError := &scaleset.MessageQueueTokenExpiredError{}
-	if !errors.As(err, &expiredError) {
-		return nil, fmt.Errorf("failed to get next message: %w", err)
-	}
-
-	if err := l.refreshSession(ctx); err != nil {
-		return nil, err
-	}
-
-	l.logger.Info("Getting next message", "lastMessageID", l.lastMessageID)
-
-	msg, err = l.client.GetMessage(
-		ctx,
-		l.session.MessageQueueURL,
-		l.session.MessageQueueAccessToken,
-		l.lastMessageID,
-		l.maxRunners,
-	)
-	if err != nil { // if error
-		return nil, fmt.Errorf("failed to get next message after message session refresh: %w", err)
-	}
-
-	return msg, nil
-}
-
-func (l *Listener) deleteLastMessage(ctx context.Context) error {
-	l.logger.Info("Deleting last message", "lastMessageID", l.lastMessageID)
-	err := l.client.DeleteMessage(
-		ctx,
-		l.session.MessageQueueURL,
-		l.session.MessageQueueAccessToken,
-		l.lastMessageID,
-	)
-	if err == nil { // if NO error
-		return nil
-	}
-
-	expiredError := &scaleset.MessageQueueTokenExpiredError{}
-	if !errors.As(err, &expiredError) {
-		return fmt.Errorf("failed to delete last message: %w", err)
-	}
-
-	if err := l.refreshSession(ctx); err != nil {
-		return err
-	}
-
-	err = l.client.DeleteMessage(
-		ctx,
-		l.session.MessageQueueURL,
-		l.session.MessageQueueAccessToken,
-		l.lastMessageID,
-	)
-	if err != nil {
-		return fmt.Errorf("failed to delete last message after message session refresh: %w", err)
-	}
-
-	return nil
-}
-
-type parsedMessage struct {
-	statistics    *scaleset.RunnerScaleSetStatistic
-	jobsStarted   []*scaleset.JobStarted
-	jobsCompleted []*scaleset.JobCompleted
-}
-
-func (l *Listener) parseMessage(ctx context.Context, msg *scaleset.RunnerScaleSetMessage) (*parsedMessage, error) {
-	if msg.MessageType != "RunnerScaleSetJobMessages" {
-		l.logger.Info("Skipping message", "messageType", msg.MessageType)
-		return nil, fmt.Errorf("invalid message type: %s", msg.MessageType)
-	}
-
-	l.logger.Info("Processing message", "messageId", msg.MessageID, "messageType", msg.MessageType)
-	if msg.Statistics == nil {
-		return nil, fmt.Errorf("invalid message: statistics is nil")
-	}
-
-	l.logger.Info("New runner scale set statistics.", "statistics", msg.Statistics)
-
-	var batchedMessages []json.RawMessage
-	if len(msg.Body) > 0 {
-		if err := json.Unmarshal([]byte(msg.Body), &batchedMessages); err != nil {
-			return nil, fmt.Errorf("failed to unmarshal batched messages: %w", err)
-		}
-	}
-
-	parsedMsg := &parsedMessage{
-		statistics: msg.Statistics,
-	}
-
-	for _, msg := range batchedMessages {
-		var messageType scaleset.JobMessageType
-		if err := json.Unmarshal(msg, &messageType); err != nil {
-			return nil, fmt.Errorf("failed to decode job message type: %w", err)
-		}
-
-		switch messageType.MessageType {
-		case scaleset.MessageTypeJobAssigned:
-			var jobAssigned scaleset.JobAssigned
-			if err := json.Unmarshal(msg, &jobAssigned); err != nil {
-				return nil, fmt.Errorf("failed to decode job assigned: %w", err)
-			}
-
-			l.logger.Info("Job assigned message received", "jobId", jobAssigned.JobID)
-
-		case scaleset.MessageTypeJobStarted:
-			var jobStarted scaleset.JobStarted
-			if err := json.Unmarshal(msg, &jobStarted); err != nil {
-				return nil, fmt.Errorf("could not decode job started message. %w", err)
-			}
-			l.logger.Info("Job started message received.", "JobID", jobStarted.JobID, "RunnerId", jobStarted.RunnerID)
-			parsedMsg.jobsStarted = append(parsedMsg.jobsStarted, &jobStarted)
-
-		case scaleset.MessageTypeJobCompleted:
-			var jobCompleted scaleset.JobCompleted
-			if err := json.Unmarshal(msg, &jobCompleted); err != nil {
-				return nil, fmt.Errorf("failed to decode job completed: %w", err)
-			}
-
-			l.logger.Info(
-				"Job completed message received.",
-				"JobID", jobCompleted.JobID,
-				"Result", jobCompleted.Result,
-				"RunnerId", jobCompleted.RunnerID,
-				"RunnerName", jobCompleted.RunnerName,
-			)
-			parsedMsg.jobsCompleted = append(parsedMsg.jobsCompleted, &jobCompleted)
-
-		default:
-			l.logger.Info("unknown job message type.", "messageType", messageType.MessageType)
-		}
-	}
-
-	return parsedMsg, ctx.Err()
-}
-
-func (l *Listener) refreshSession(ctx context.Context) error {
-	l.logger.Info("Message queue token is expired during GetNextMessage, refreshing...")
-	session, err := l.client.RefreshMessageSession(
-		ctx,
-		l.session.RunnerScaleSet.ID,
-		l.session.SessionID,
-	)
-	if err != nil {
-		return fmt.Errorf("refresh message session failed. %w", err)
-	}
-
-	l.session = session
-	return nil
-}
-
-func (l *Listener) deleteMessageSession() error {
-	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
-	defer cancel()
-
-	l.logger.Info("Deleting message session")
-
-	if err := l.client.DeleteMessageSession(ctx, l.session.RunnerScaleSet.ID, l.session.SessionID); err != nil {
-		return fmt.Errorf("failed to delete message session: %w", err)
-	}
-
-	return nil
+func (l *Listener) handleStatistics(ctx context.Context, msg *scaleset.RunnerScaleSetStatistic) {
+	l.latestStatistics = msg
+	l.metricsRecorder.RecordStatistics(msg)
 }

listener/listener_test.go

@@ -0,0 +1,180 @@
+package listener
+
+import (
+	"context"
+	"math"
+	"testing"
+
+	"github.com/actions/scaleset"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNew(t *testing.T) {
+	t.Parallel()
+	t.Run("invalid config", func(t *testing.T) {
+		t.Parallel()
+		var config Config
+		assert.Error(t, config.Validate())
+	})
+
+	t.Run("valid config", func(t *testing.T) {
+		t.Parallel()
+		config := Config{
+			ScaleSetID: 1,
+		}
+		assert.NoError(t, config.Validate())
+	})
+
+	t.Run("invalid max runners", func(t *testing.T) {
+		t.Parallel()
+		config := Config{
+			ScaleSetID: 1,
+			MaxRunners: -1,
+		}
+		assert.Error(t, config.Validate())
+	})
+
+	t.Run("zero max runners", func(t *testing.T) {
+		t.Parallel()
+		config := Config{
+			ScaleSetID: 1,
+			MaxRunners: math.MaxInt32 + 1,
+		}
+		assert.Error(t, config.Validate())
+	})
+
+	t.Run("creates listener", func(t *testing.T) {
+		t.Parallel()
+		config := Config{
+			ScaleSetID: 1,
+			MaxRunners: 5,
+		}
+
+		client := NewMockClient(t)
+		l, err := New(client, config)
+		require.Nil(t, err)
+		assert.Equal(t, config.ScaleSetID, l.scaleSetID)
+		assert.Equal(t, uint32(config.MaxRunners), l.maxRunners.Load())
+	})
+}
+
+func TestListener_Run(t *testing.T) {
+	t.Parallel()
+
+	t.Run("call handle regardless of initial message", func(t *testing.T) {
+		t.Parallel()
+		ctx, cancel := context.WithCancel(context.Background())
+
+		config := Config{
+			ScaleSetID: 1,
+		}
+
+		client := NewMockClient(t)
+
+		uuid := uuid.New()
+		initialStatistics := &scaleset.RunnerScaleSetStatistic{
+			TotalAssignedJobs: 2,
+		}
+		session := scaleset.RunnerScaleSetSession{
+			SessionID:               uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &scaleset.RunnerScaleSet{},
+			MessageQueueURL:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              initialStatistics,
+		}
+
+		client.On("Session").Return(session).Once()
+
+		metricsRecorder := NewMockMetricsRecorder(t)
+		metricsRecorder.On("RecordStatistics", initialStatistics).Once()
+		metricsRecorder.On("RecordDesiredRunners", initialStatistics.TotalAssignedJobs).
+			Return(initialStatistics.TotalAssignedJobs, nil).
+			Run(func(mock.Arguments) { cancel() }).
+			Once()
+
+		l, err := New(client, config, WithMetricsRecorder(metricsRecorder))
+		require.Nil(t, err)
+
+		handler := NewMockScaler(t)
+		handler.On("HandleDesiredRunnerCount", mock.Anything, mock.Anything).
+			Return(initialStatistics.TotalAssignedJobs, nil).
+			Once()
+
+		err = l.Run(ctx, handler)
+		assert.ErrorIs(t, err, context.Canceled)
+	})
+
+	t.Run("cancel context after get message", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := context.WithCancel(context.Background())
+
+		config := Config{
+			ScaleSetID: 1,
+			MaxRunners: 10,
+		}
+
+		uuid := uuid.New()
+		initialStatistics := &scaleset.RunnerScaleSetStatistic{
+			TotalAssignedJobs: 2,
+		}
+
+		session := scaleset.RunnerScaleSetSession{
+			SessionID:               uuid,
+			OwnerName:               "example",
+			RunnerScaleSet:          &scaleset.RunnerScaleSet{},
+			MessageQueueURL:         "https://example.com",
+			MessageQueueAccessToken: "1234567890",
+			Statistics:              initialStatistics,
+		}
+
+		msg := &scaleset.RunnerScaleSetMessage{
+			MessageID: 1,
+			Statistics: &scaleset.RunnerScaleSetStatistic{
+				TotalAssignedJobs: 3,
+			},
+		}
+
+		metricsRecorder := NewMockMetricsRecorder(t)
+		client := NewMockClient(t)
+		handler := NewMockScaler(t)
+
+		client.On("Session").Return(session).Once()
+		metricsRecorder.On("RecordStatistics", initialStatistics).Once()
+		metricsRecorder.On("RecordDesiredRunners", initialStatistics.TotalAssignedJobs).
+			Return(initialStatistics.TotalAssignedJobs, nil).
+			Once()
+		handler.On("HandleDesiredRunnerCount", mock.Anything, initialStatistics.TotalAssignedJobs).
+			Return(initialStatistics.TotalAssignedJobs, nil).
+			Once()
+
+		client.On("GetMessage", ctx, mock.Anything, 10).
+			Return(msg, nil).
+			Run(func(mock.Arguments) { cancel() }).
+			Once()
+
+		metricsRecorder.On("RecordStatistics", msg.Statistics).Once()
+		// Ensure delete message is called without cancel
+		client.On("DeleteMessage", context.WithoutCancel(ctx), mock.Anything).
+			Return(nil).
+			Once()
+
+		metricsRecorder.On("RecordDesiredRunners", msg.Statistics.TotalAssignedJobs).
+			Return(msg.Statistics.TotalAssignedJobs, nil).
+			Once()
+
+		handler.On("HandleDesiredRunnerCount", mock.Anything, msg.Statistics.TotalAssignedJobs).
+			Return(msg.Statistics.TotalAssignedJobs, nil).
+			Once()
+
+		l, err := New(client, config, WithMetricsRecorder(metricsRecorder))
+		require.Nil(t, err)
+
+		err = l.Run(ctx, handler)
+		assert.ErrorIs(t, context.Canceled, err)
+	})
+}

listener/mocks_test.go

@@ -0,0 +1,676 @@
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
+
+package listener
+
+import (
+	"context"
+
+	"github.com/actions/scaleset"
+	mock "github.com/stretchr/testify/mock"
+)
+
+// NewMockClient creates a new instance of MockClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockClient {
+	mock := &MockClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
+// MockClient is an autogenerated mock type for the Client type
+type MockClient struct {
+	mock.Mock
+}
+
+type MockClient_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockClient) EXPECT() *MockClient_Expecter {
+	return &MockClient_Expecter{mock: &_m.Mock}
+}
+
+// AcquireJobs provides a mock function for the type MockClient
+func (_mock *MockClient) AcquireJobs(ctx context.Context, requestIDs []int64) ([]int64, error) {
+	ret := _mock.Called(ctx, requestIDs)
+
+	if len(ret) == 0 {
+		panic("no return value specified for AcquireJobs")
+	}
+
+	var r0 []int64
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, []int64) ([]int64, error)); ok {
+		return returnFunc(ctx, requestIDs)
+	}
+	if returnFunc, ok := ret.Get(0).(func(context.Context, []int64) []int64); ok {
+		r0 = returnFunc(ctx, requestIDs)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]int64)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func(context.Context, []int64) error); ok {
+		r1 = returnFunc(ctx, requestIDs)
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// MockClient_AcquireJobs_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AcquireJobs'
+type MockClient_AcquireJobs_Call struct {
+	*mock.Call
+}
+
+// AcquireJobs is a helper method to define mock.On call
+//   - ctx context.Context
+//   - requestIDs []int64
+func (_e *MockClient_Expecter) AcquireJobs(ctx interface{}, requestIDs interface{}) *MockClient_AcquireJobs_Call {
+	return &MockClient_AcquireJobs_Call{Call: _e.mock.On("AcquireJobs", ctx, requestIDs)}
+}
+
+func (_c *MockClient_AcquireJobs_Call) Run(run func(ctx context.Context, requestIDs []int64)) *MockClient_AcquireJobs_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 []int64
+		if args[1] != nil {
+			arg1 = args[1].([]int64)
+		}
+		run(
+			arg0,
+			arg1,
+		)
+	})
+	return _c
+}
+
+func (_c *MockClient_AcquireJobs_Call) Return(int64s []int64, err error) *MockClient_AcquireJobs_Call {
+	_c.Call.Return(int64s, err)
+	return _c
+}
+
+func (_c *MockClient_AcquireJobs_Call) RunAndReturn(run func(ctx context.Context, requestIDs []int64) ([]int64, error)) *MockClient_AcquireJobs_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// DeleteMessage provides a mock function for the type MockClient
+func (_mock *MockClient) DeleteMessage(ctx context.Context, messageID int) error {
+	ret := _mock.Called(ctx, messageID)
+
+	if len(ret) == 0 {
+		panic("no return value specified for DeleteMessage")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, int) error); ok {
+		r0 = returnFunc(ctx, messageID)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// MockClient_DeleteMessage_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'DeleteMessage'
+type MockClient_DeleteMessage_Call struct {
+	*mock.Call
+}
+
+// DeleteMessage is a helper method to define mock.On call
+//   - ctx context.Context
+//   - messageID int
+func (_e *MockClient_Expecter) DeleteMessage(ctx interface{}, messageID interface{}) *MockClient_DeleteMessage_Call {
+	return &MockClient_DeleteMessage_Call{Call: _e.mock.On("DeleteMessage", ctx, messageID)}
+}
+
+func (_c *MockClient_DeleteMessage_Call) Run(run func(ctx context.Context, messageID int)) *MockClient_DeleteMessage_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 int
+		if args[1] != nil {
+			arg1 = args[1].(int)
+		}
+		run(
+			arg0,
+			arg1,
+		)
+	})
+	return _c
+}
+
+func (_c *MockClient_DeleteMessage_Call) Return(err error) *MockClient_DeleteMessage_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *MockClient_DeleteMessage_Call) RunAndReturn(run func(ctx context.Context, messageID int) error) *MockClient_DeleteMessage_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// GetMessage provides a mock function for the type MockClient
+func (_mock *MockClient) GetMessage(ctx context.Context, lastMessageID int, maxCapacity int) (*scaleset.RunnerScaleSetMessage, error) {
+	ret := _mock.Called(ctx, lastMessageID, maxCapacity)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetMessage")
+	}
+
+	var r0 *scaleset.RunnerScaleSetMessage
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, int, int) (*scaleset.RunnerScaleSetMessage, error)); ok {
+		return returnFunc(ctx, lastMessageID, maxCapacity)
+	}
+	if returnFunc, ok := ret.Get(0).(func(context.Context, int, int) *scaleset.RunnerScaleSetMessage); ok {
+		r0 = returnFunc(ctx, lastMessageID, maxCapacity)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*scaleset.RunnerScaleSetMessage)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func(context.Context, int, int) error); ok {
+		r1 = returnFunc(ctx, lastMessageID, maxCapacity)
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// MockClient_GetMessage_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetMessage'
+type MockClient_GetMessage_Call struct {
+	*mock.Call
+}
+
+// GetMessage is a helper method to define mock.On call
+//   - ctx context.Context
+//   - lastMessageID int
+//   - maxCapacity int
+func (_e *MockClient_Expecter) GetMessage(ctx interface{}, lastMessageID interface{}, maxCapacity interface{}) *MockClient_GetMessage_Call {
+	return &MockClient_GetMessage_Call{Call: _e.mock.On("GetMessage", ctx, lastMessageID, maxCapacity)}
+}
+
+func (_c *MockClient_GetMessage_Call) Run(run func(ctx context.Context, lastMessageID int, maxCapacity int)) *MockClient_GetMessage_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 int
+		if args[1] != nil {
+			arg1 = args[1].(int)
+		}
+		var arg2 int
+		if args[2] != nil {
+			arg2 = args[2].(int)
+		}
+		run(
+			arg0,
+			arg1,
+			arg2,
+		)
+	})
+	return _c
+}
+
+func (_c *MockClient_GetMessage_Call) Return(runnerScaleSetMessage *scaleset.RunnerScaleSetMessage, err error) *MockClient_GetMessage_Call {
+	_c.Call.Return(runnerScaleSetMessage, err)
+	return _c
+}
+
+func (_c *MockClient_GetMessage_Call) RunAndReturn(run func(ctx context.Context, lastMessageID int, maxCapacity int) (*scaleset.RunnerScaleSetMessage, error)) *MockClient_GetMessage_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Session provides a mock function for the type MockClient
+func (_mock *MockClient) Session() scaleset.RunnerScaleSetSession {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Session")
+	}
+
+	var r0 scaleset.RunnerScaleSetSession
+	if returnFunc, ok := ret.Get(0).(func() scaleset.RunnerScaleSetSession); ok {
+		r0 = returnFunc()
+	} else {
+		r0 = ret.Get(0).(scaleset.RunnerScaleSetSession)
+	}
+	return r0
+}
+
+// MockClient_Session_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Session'
+type MockClient_Session_Call struct {
+	*mock.Call
+}
+
+// Session is a helper method to define mock.On call
+func (_e *MockClient_Expecter) Session() *MockClient_Session_Call {
+	return &MockClient_Session_Call{Call: _e.mock.On("Session")}
+}
+
+func (_c *MockClient_Session_Call) Run(run func()) *MockClient_Session_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *MockClient_Session_Call) Return(runnerScaleSetSession scaleset.RunnerScaleSetSession) *MockClient_Session_Call {
+	_c.Call.Return(runnerScaleSetSession)
+	return _c
+}
+
+func (_c *MockClient_Session_Call) RunAndReturn(run func() scaleset.RunnerScaleSetSession) *MockClient_Session_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewMockMetricsRecorder creates a new instance of MockMetricsRecorder. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockMetricsRecorder(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockMetricsRecorder {
+	mock := &MockMetricsRecorder{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
+// MockMetricsRecorder is an autogenerated mock type for the MetricsRecorder type
+type MockMetricsRecorder struct {
+	mock.Mock
+}
+
+type MockMetricsRecorder_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockMetricsRecorder) EXPECT() *MockMetricsRecorder_Expecter {
+	return &MockMetricsRecorder_Expecter{mock: &_m.Mock}
+}
+
+// RecordDesiredRunners provides a mock function for the type MockMetricsRecorder
+func (_mock *MockMetricsRecorder) RecordDesiredRunners(count int) {
+	_mock.Called(count)
+	return
+}
+
+// MockMetricsRecorder_RecordDesiredRunners_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecordDesiredRunners'
+type MockMetricsRecorder_RecordDesiredRunners_Call struct {
+	*mock.Call
+}
+
+// RecordDesiredRunners is a helper method to define mock.On call
+//   - count int
+func (_e *MockMetricsRecorder_Expecter) RecordDesiredRunners(count interface{}) *MockMetricsRecorder_RecordDesiredRunners_Call {
+	return &MockMetricsRecorder_RecordDesiredRunners_Call{Call: _e.mock.On("RecordDesiredRunners", count)}
+}
+
+func (_c *MockMetricsRecorder_RecordDesiredRunners_Call) Run(run func(count int)) *MockMetricsRecorder_RecordDesiredRunners_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 int
+		if args[0] != nil {
+			arg0 = args[0].(int)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *MockMetricsRecorder_RecordDesiredRunners_Call) Return() *MockMetricsRecorder_RecordDesiredRunners_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockMetricsRecorder_RecordDesiredRunners_Call) RunAndReturn(run func(count int)) *MockMetricsRecorder_RecordDesiredRunners_Call {
+	_c.Run(run)
+	return _c
+}
+
+// RecordJobCompleted provides a mock function for the type MockMetricsRecorder
+func (_mock *MockMetricsRecorder) RecordJobCompleted(msg *scaleset.JobCompleted) {
+	_mock.Called(msg)
+	return
+}
+
+// MockMetricsRecorder_RecordJobCompleted_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecordJobCompleted'
+type MockMetricsRecorder_RecordJobCompleted_Call struct {
+	*mock.Call
+}
+
+// RecordJobCompleted is a helper method to define mock.On call
+//   - msg *scaleset.JobCompleted
+func (_e *MockMetricsRecorder_Expecter) RecordJobCompleted(msg interface{}) *MockMetricsRecorder_RecordJobCompleted_Call {
+	return &MockMetricsRecorder_RecordJobCompleted_Call{Call: _e.mock.On("RecordJobCompleted", msg)}
+}
+
+func (_c *MockMetricsRecorder_RecordJobCompleted_Call) Run(run func(msg *scaleset.JobCompleted)) *MockMetricsRecorder_RecordJobCompleted_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 *scaleset.JobCompleted
+		if args[0] != nil {
+			arg0 = args[0].(*scaleset.JobCompleted)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *MockMetricsRecorder_RecordJobCompleted_Call) Return() *MockMetricsRecorder_RecordJobCompleted_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockMetricsRecorder_RecordJobCompleted_Call) RunAndReturn(run func(msg *scaleset.JobCompleted)) *MockMetricsRecorder_RecordJobCompleted_Call {
+	_c.Run(run)
+	return _c
+}
+
+// RecordJobStarted provides a mock function for the type MockMetricsRecorder
+func (_mock *MockMetricsRecorder) RecordJobStarted(msg *scaleset.JobStarted) {
+	_mock.Called(msg)
+	return
+}
+
+// MockMetricsRecorder_RecordJobStarted_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecordJobStarted'
+type MockMetricsRecorder_RecordJobStarted_Call struct {
+	*mock.Call
+}
+
+// RecordJobStarted is a helper method to define mock.On call
+//   - msg *scaleset.JobStarted
+func (_e *MockMetricsRecorder_Expecter) RecordJobStarted(msg interface{}) *MockMetricsRecorder_RecordJobStarted_Call {
+	return &MockMetricsRecorder_RecordJobStarted_Call{Call: _e.mock.On("RecordJobStarted", msg)}
+}
+
+func (_c *MockMetricsRecorder_RecordJobStarted_Call) Run(run func(msg *scaleset.JobStarted)) *MockMetricsRecorder_RecordJobStarted_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 *scaleset.JobStarted
+		if args[0] != nil {
+			arg0 = args[0].(*scaleset.JobStarted)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *MockMetricsRecorder_RecordJobStarted_Call) Return() *MockMetricsRecorder_RecordJobStarted_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockMetricsRecorder_RecordJobStarted_Call) RunAndReturn(run func(msg *scaleset.JobStarted)) *MockMetricsRecorder_RecordJobStarted_Call {
+	_c.Run(run)
+	return _c
+}
+
+// RecordStatistics provides a mock function for the type MockMetricsRecorder
+func (_mock *MockMetricsRecorder) RecordStatistics(statistics *scaleset.RunnerScaleSetStatistic) {
+	_mock.Called(statistics)
+	return
+}
+
+// MockMetricsRecorder_RecordStatistics_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecordStatistics'
+type MockMetricsRecorder_RecordStatistics_Call struct {
+	*mock.Call
+}
+
+// RecordStatistics is a helper method to define mock.On call
+//   - statistics *scaleset.RunnerScaleSetStatistic
+func (_e *MockMetricsRecorder_Expecter) RecordStatistics(statistics interface{}) *MockMetricsRecorder_RecordStatistics_Call {
+	return &MockMetricsRecorder_RecordStatistics_Call{Call: _e.mock.On("RecordStatistics", statistics)}
+}
+
+func (_c *MockMetricsRecorder_RecordStatistics_Call) Run(run func(statistics *scaleset.RunnerScaleSetStatistic)) *MockMetricsRecorder_RecordStatistics_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 *scaleset.RunnerScaleSetStatistic
+		if args[0] != nil {
+			arg0 = args[0].(*scaleset.RunnerScaleSetStatistic)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *MockMetricsRecorder_RecordStatistics_Call) Return() *MockMetricsRecorder_RecordStatistics_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *MockMetricsRecorder_RecordStatistics_Call) RunAndReturn(run func(statistics *scaleset.RunnerScaleSetStatistic)) *MockMetricsRecorder_RecordStatistics_Call {
+	_c.Run(run)
+	return _c
+}
+
+// NewMockScaler creates a new instance of MockScaler. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewMockScaler(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MockScaler {
+	mock := &MockScaler{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
+// MockScaler is an autogenerated mock type for the Scaler type
+type MockScaler struct {
+	mock.Mock
+}
+
+type MockScaler_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *MockScaler) EXPECT() *MockScaler_Expecter {
+	return &MockScaler_Expecter{mock: &_m.Mock}
+}
+
+// HandleDesiredRunnerCount provides a mock function for the type MockScaler
+func (_mock *MockScaler) HandleDesiredRunnerCount(ctx context.Context, count int) (int, error) {
+	ret := _mock.Called(ctx, count)
+
+	if len(ret) == 0 {
+		panic("no return value specified for HandleDesiredRunnerCount")
+	}
+
+	var r0 int
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, int) (int, error)); ok {
+		return returnFunc(ctx, count)
+	}
+	if returnFunc, ok := ret.Get(0).(func(context.Context, int) int); ok {
+		r0 = returnFunc(ctx, count)
+	} else {
+		r0 = ret.Get(0).(int)
+	}
+	if returnFunc, ok := ret.Get(1).(func(context.Context, int) error); ok {
+		r1 = returnFunc(ctx, count)
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// MockScaler_HandleDesiredRunnerCount_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'HandleDesiredRunnerCount'
+type MockScaler_HandleDesiredRunnerCount_Call struct {
+	*mock.Call
+}
+
+// HandleDesiredRunnerCount is a helper method to define mock.On call
+//   - ctx context.Context
+//   - count int
+func (_e *MockScaler_Expecter) HandleDesiredRunnerCount(ctx interface{}, count interface{}) *MockScaler_HandleDesiredRunnerCount_Call {
+	return &MockScaler_HandleDesiredRunnerCount_Call{Call: _e.mock.On("HandleDesiredRunnerCount", ctx, count)}
+}
+
+func (_c *MockScaler_HandleDesiredRunnerCount_Call) Run(run func(ctx context.Context, count int)) *MockScaler_HandleDesiredRunnerCount_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 int
+		if args[1] != nil {
+			arg1 = args[1].(int)
+		}
+		run(
+			arg0,
+			arg1,
+		)
+	})
+	return _c
+}
+
+func (_c *MockScaler_HandleDesiredRunnerCount_Call) Return(n int, err error) *MockScaler_HandleDesiredRunnerCount_Call {
+	_c.Call.Return(n, err)
+	return _c
+}
+
+func (_c *MockScaler_HandleDesiredRunnerCount_Call) RunAndReturn(run func(ctx context.Context, count int) (int, error)) *MockScaler_HandleDesiredRunnerCount_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// HandleJobCompleted provides a mock function for the type MockScaler
+func (_mock *MockScaler) HandleJobCompleted(ctx context.Context, jobInfo *scaleset.JobCompleted) error {
+	ret := _mock.Called(ctx, jobInfo)
+
+	if len(ret) == 0 {
+		panic("no return value specified for HandleJobCompleted")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, *scaleset.JobCompleted) error); ok {
+		r0 = returnFunc(ctx, jobInfo)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// MockScaler_HandleJobCompleted_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'HandleJobCompleted'
+type MockScaler_HandleJobCompleted_Call struct {
+	*mock.Call
+}
+
+// HandleJobCompleted is a helper method to define mock.On call
+//   - ctx context.Context
+//   - jobInfo *scaleset.JobCompleted
+func (_e *MockScaler_Expecter) HandleJobCompleted(ctx interface{}, jobInfo interface{}) *MockScaler_HandleJobCompleted_Call {
+	return &MockScaler_HandleJobCompleted_Call{Call: _e.mock.On("HandleJobCompleted", ctx, jobInfo)}
+}
+
+func (_c *MockScaler_HandleJobCompleted_Call) Run(run func(ctx context.Context, jobInfo *scaleset.JobCompleted)) *MockScaler_HandleJobCompleted_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 *scaleset.JobCompleted
+		if args[1] != nil {
+			arg1 = args[1].(*scaleset.JobCompleted)
+		}
+		run(
+			arg0,
+			arg1,
+		)
+	})
+	return _c
+}
+
+func (_c *MockScaler_HandleJobCompleted_Call) Return(err error) *MockScaler_HandleJobCompleted_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *MockScaler_HandleJobCompleted_Call) RunAndReturn(run func(ctx context.Context, jobInfo *scaleset.JobCompleted) error) *MockScaler_HandleJobCompleted_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// HandleJobStarted provides a mock function for the type MockScaler
+func (_mock *MockScaler) HandleJobStarted(ctx context.Context, jobInfo *scaleset.JobStarted) error {
+	ret := _mock.Called(ctx, jobInfo)
+
+	if len(ret) == 0 {
+		panic("no return value specified for HandleJobStarted")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, *scaleset.JobStarted) error); ok {
+		r0 = returnFunc(ctx, jobInfo)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// MockScaler_HandleJobStarted_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'HandleJobStarted'
+type MockScaler_HandleJobStarted_Call struct {
+	*mock.Call
+}
+
+// HandleJobStarted is a helper method to define mock.On call
+//   - ctx context.Context
+//   - jobInfo *scaleset.JobStarted
+func (_e *MockScaler_Expecter) HandleJobStarted(ctx interface{}, jobInfo interface{}) *MockScaler_HandleJobStarted_Call {
+	return &MockScaler_HandleJobStarted_Call{Call: _e.mock.On("HandleJobStarted", ctx, jobInfo)}
+}
+
+func (_c *MockScaler_HandleJobStarted_Call) Run(run func(ctx context.Context, jobInfo *scaleset.JobStarted)) *MockScaler_HandleJobStarted_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 *scaleset.JobStarted
+		if args[1] != nil {
+			arg1 = args[1].(*scaleset.JobStarted)
+		}
+		run(
+			arg0,
+			arg1,
+		)
+	})
+	return _c
+}
+
+func (_c *MockScaler_HandleJobStarted_Call) Return(err error) *MockScaler_HandleJobStarted_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *MockScaler_HandleJobStarted_Call) RunAndReturn(run func(ctx context.Context, jobInfo *scaleset.JobStarted) error) *MockScaler_HandleJobStarted_Call {
+	_c.Call.Return(run)
+	return _c
+}

script/lint

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+
+set -eEuo pipefail
+
+golangci-lint run ./...
+
+go tool deadcode -test ./...
+

script/test

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+
+set -eEuo pipefail
+
+go test ./... "$@"
+

session_client.go

@@ -0,0 +1,327 @@
+package scaleset
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"strconv"
+	"sync"
+
+	"github.com/google/uuid"
+)
+
+// MessageSessionClient is a client used to interact with a message session for a runner scale set.
+// It provides methods to Get and Delete messages from the message queue associated with the session,
+// handling session token expiration and refreshing as needed.
+//
+// It is safe for concurrent use by multiple goroutines.
+// Please do not forget to call Close when done to clean up the session.
+type MessageSessionClient struct {
+	mu sync.Mutex
+	// inner client is the parent of the message session, allowing session refreshing
+	// use this client to create (and potentially refresh the session) requests.
+	innerClient *Client
+	// commonClient uses different options than the original client
+	// use this client for message session requests
+	commonClient *commonClient
+	scaleSetID   int
+	owner        string
+	session      *RunnerScaleSetSession
+}
+
+// Close deletes the message session associated with this client.
+func (c *MessageSessionClient) Close(ctx context.Context) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.deleteMessageSession(ctx, c.scaleSetID, c.session.SessionID)
+}
+
+func (c *MessageSessionClient) createMessageSession(ctx context.Context) error {
+	path := fmt.Sprintf("/%s/%d/sessions", scaleSetEndpoint, c.scaleSetID)
+
+	newSession := &RunnerScaleSetSession{
+		OwnerName: c.owner,
+	}
+
+	requestData, err := json.Marshal(newSession)
+	if err != nil {
+		return fmt.Errorf("failed to marshal new session: %w", err)
+	}
+
+	var createdSession RunnerScaleSetSession
+	if err = c.doSessionRequest(
+		ctx,
+		http.MethodPost,
+		path,
+		bytes.NewBuffer(requestData),
+		http.StatusOK,
+		&createdSession,
+	); err != nil {
+		return fmt.Errorf("failed to do the session request: %w", err)
+	}
+
+	c.session = &createdSession
+
+	return nil
+}
+
+// DeleteMessageSession deletes a message session for the specified runner scale set.
+func (c *MessageSessionClient) deleteMessageSession(ctx context.Context, runnerScaleSetID int, sessionID uuid.UUID) error {
+	path := fmt.Sprintf("/%s/%d/sessions/%s", scaleSetEndpoint, runnerScaleSetID, sessionID.String())
+	return c.doSessionRequest(ctx, http.MethodDelete, path, nil, http.StatusNoContent, nil)
+}
+
+// RefreshMessageSession refreshes a message session for the specified runner scale set.
+// This should be used when a MessageQueueTokenExpiredError is encountered.
+func (c *MessageSessionClient) refreshMessageSession(ctx context.Context) error {
+	path := fmt.Sprintf("/%s/%d/sessions/%s", scaleSetEndpoint, c.scaleSetID, c.session.SessionID.String())
+	refreshedSession := &RunnerScaleSetSession{}
+	if err := c.doSessionRequest(ctx, http.MethodPatch, path, nil, http.StatusOK, refreshedSession); err != nil {
+		return fmt.Errorf("failed to do the session request: %w", err)
+	}
+	c.session = refreshedSession
+	return nil
+}
+
+// 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.
+func (c *MessageSessionClient) GetMessage(ctx context.Context, lastMessageID int, maxCapacity int) (*RunnerScaleSetMessage, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	message, err := c.getMessage(
+		ctx,
+		lastMessageID,
+		maxCapacity,
+	)
+	if err == nil {
+		return message, nil
+	}
+
+	if !errors.Is(err, MessageQueueTokenExpiredError) {
+		return nil, fmt.Errorf("failed to get next message: %w", err)
+	}
+
+	if err := c.refreshMessageSession(ctx); err != nil {
+		return nil, fmt.Errorf("failed to refresh message session: %w", err)
+	}
+
+	return c.getMessage(
+		ctx,
+		lastMessageID,
+		maxCapacity,
+	)
+}
+
+func (c *MessageSessionClient) getMessage(ctx context.Context, lastMessageID int, maxCapacity int) (*RunnerScaleSetMessage, error) {
+	u, err := url.Parse(c.session.MessageQueueURL)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse message queue url: %w", err)
+	}
+
+	if lastMessageID > 0 {
+		q := u.Query()
+		q.Set("lastMessageId", strconv.Itoa(lastMessageID))
+		u.RawQuery = q.Encode()
+	}
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create new request with context: %w", err)
+	}
+
+	req.Header.Set("Accept", "application/json; api-version=6.0-preview")
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.session.MessageQueueAccessToken))
+	req.Header.Set("User-Agent", c.commonClient.userAgent)
+	req.Header.Set(HeaderScaleSetMaxCapacity, strconv.Itoa(maxCapacity))
+
+	resp, err := c.commonClient.do(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to issue the request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	switch resp.StatusCode {
+	case http.StatusAccepted:
+		return nil, nil
+
+	case http.StatusOK:
+		message, err := parseRunnerScaleSetMessageResponse(resp.Body)
+		if err != nil {
+			return nil, newRequestResponseError(req, resp, fmt.Errorf("failed to parse message response: %w", err))
+		}
+		return message, nil
+
+	case http.StatusUnauthorized:
+		return nil, newRequestResponseError(req, resp, MessageQueueTokenExpiredError)
+
+	default:
+		return nil, newRequestResponseError(req, resp, fmt.Errorf("unexpected status code %s", resp.Status))
+	}
+}
+
+// DeleteMessage deletes a message from the runner scale set message queue.
+// This should typically be done after processing the message and acts as an acknowledgment.
+// If the current session token is expired, it refreshes the session and tries one more time.
+func (c *MessageSessionClient) DeleteMessage(ctx context.Context, messageID int) error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	err := c.deleteMessage(ctx, messageID)
+	if err == nil {
+		return nil
+	}
+
+	if !errors.Is(err, MessageQueueTokenExpiredError) {
+		return fmt.Errorf("failed to delete message: %w", err)
+	}
+
+	if err := c.refreshMessageSession(ctx); err != nil {
+		return fmt.Errorf("failed to refresh message session: %w", err)
+	}
+
+	return c.deleteMessage(ctx, messageID)
+}
+
+func (c *MessageSessionClient) deleteMessage(ctx context.Context, messageID int) error {
+	u, err := url.Parse(c.session.MessageQueueURL)
+	if err != nil {
+		return fmt.Errorf("failed to parse message queue url: %w", err)
+	}
+
+	u.Path = fmt.Sprintf("%s/%d", u.Path, messageID)
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodDelete, u.String(), nil)
+	if err != nil {
+		return fmt.Errorf("failed to create new request with context: %w", err)
+	}
+
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.session.MessageQueueAccessToken))
+	req.Header.Set("User-Agent", c.commonClient.userAgent)
+
+	resp, err := c.commonClient.do(req)
+	if err != nil {
+		return fmt.Errorf("failed to issue the request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode == http.StatusNoContent {
+		return nil
+	}
+
+	if resp.StatusCode != http.StatusUnauthorized {
+		return newRequestResponseError(req, resp, fmt.Errorf("unexpected status code %s", resp.Status))
+	}
+
+	return newRequestResponseError(req, resp, MessageQueueTokenExpiredError)
+}
+
+func (c *MessageSessionClient) Session() RunnerScaleSetSession {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.session == nil {
+		return RunnerScaleSetSession{}
+	}
+
+	return *c.session
+}
+
+// AcquireJobs acquires the given job request IDs from the runner scale set.
+// If the current session token is expired, it refreshes the session and tries one more time.
+func (c *MessageSessionClient) AcquireJobs(ctx context.Context, requestIDs []int64) ([]int64, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	ids, err := c.acquireJobs(ctx, requestIDs)
+	if err == nil {
+		return ids, nil
+	}
+
+	if !errors.Is(err, MessageQueueTokenExpiredError) {
+		return nil, fmt.Errorf("failed to acquire jobs: %w", err)
+	}
+
+	if err := c.refreshMessageSession(ctx); err != nil {
+		return nil, fmt.Errorf("failed to refresh message session: %w", err)
+	}
+
+	return c.acquireJobs(ctx, requestIDs)
+}
+
+func (c *MessageSessionClient) acquireJobs(ctx context.Context, requestIDs []int64) ([]int64, error) {
+	body, err := json.Marshal(requestIDs)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal request ids: %w", err)
+	}
+
+	path := fmt.Sprintf("/%s/%d/acquirejobs", scaleSetEndpoint, c.scaleSetID)
+
+	c.innerClient.mu.Lock()
+	req, err := c.innerClient.newActionsServiceRequest(ctx, http.MethodPost, path, bytes.NewBuffer(body))
+	c.innerClient.mu.Unlock()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create acquire jobs request: %w", err)
+	}
+
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.session.MessageQueueAccessToken))
+
+	resp, err := c.commonClient.do(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to issue acquire jobs request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode == http.StatusUnauthorized {
+		return nil, newRequestResponseError(req, resp, MessageQueueTokenExpiredError)
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, newRequestResponseError(req, resp, fmt.Errorf("unexpected status code %s", resp.Status))
+	}
+
+	var result acquireJobsResponse
+	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return nil, newRequestResponseError(req, resp, fmt.Errorf("failed to decode acquire jobs response: %w", err))
+	}
+
+	return result.Value, nil
+}
+
+func (c *MessageSessionClient) doSessionRequest(ctx context.Context, method, path string, requestData io.Reader, expectedResponseStatusCode int, responseUnmarshalTarget any) error {
+	c.innerClient.mu.Lock()
+	defer c.innerClient.mu.Unlock()
+
+	req, err := c.innerClient.newActionsServiceRequest(ctx, method, path, requestData)
+	if err != nil {
+		return fmt.Errorf("failed to create new actions service request: %w", err)
+	}
+
+	// use potentially modified client to issue a request
+	resp, err := c.commonClient.do(req)
+	if err != nil {
+		return fmt.Errorf("failed to issue the request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != expectedResponseStatusCode {
+		return newRequestResponseError(req, resp, fmt.Errorf("unexpected status code %s", resp.Status))
+	}
+
+	if responseUnmarshalTarget == nil {
+		return nil
+	}
+
+	if err := json.NewDecoder(resp.Body).Decode(responseUnmarshalTarget); err != nil {
+		return newRequestResponseError(req, resp, fmt.Errorf("failed to unmarshal response body: %w", err))
+	}
+
+	return nil
+}

session_client_test.go

@@ -0,0 +1,879 @@
+package scaleset
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"net/http"
+	"strconv"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func newTestSessionRequestHandler(t *testing.T, session RunnerScaleSetSession) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		srv := r.Context().Value(ctxKeyServer).(*actionsServer)
+		session.MessageQueueURL = srv.URL
+		resp, err := json.Marshal(session)
+		require.NoError(t, err)
+		w.Header().Set("Content-Type", "application/json")
+		w.Write(resp)
+	}
+}
+
+func TestCreateMessageSession(t *testing.T) {
+	ctx := context.Background()
+	auth := actionsAuth{
+		token: "token",
+	}
+
+	t.Run("CreateMessageSession unmarshals correctly", func(t *testing.T) {
+		runnerScaleSet := RunnerScaleSet{
+			ID:            1,
+			Name:          "ScaleSet",
+			CreatedOn:     time.Date(1, time.January, 1, 0, 0, 0, 0, time.UTC),
+			RunnerSetting: RunnerSetting{},
+		}
+
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			handleSessionRequest(w, r)
+		}))
+		want := server.testRunnerScaleSetSession()
+		handleSessionRequest = newTestSessionRequestHandler(t, want)
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, runnerScaleSet.ID, "my-org")
+		require.NoError(t, err)
+
+		session := sessionClient.Session()
+		require.NotEqual(t, session.SessionID, uuid.Nil)
+		assert.Equal(t, want, session)
+	})
+
+	t.Run("CreateMessageSession includes actions exception details", func(t *testing.T) {
+		owner := "foo"
+		runnerScaleSet := RunnerScaleSet{
+			ID:            1,
+			Name:          "ScaleSet",
+			CreatedOn:     time.Date(1, time.January, 1, 0, 0, 0, 0, time.UTC),
+			RunnerSetting: RunnerSetting{},
+		}
+
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			w.Header().Set("Content-Type", "application/json")
+			w.Header().Set(headerActionsActivityID, exampleRequestID)
+			w.WriteHeader(http.StatusBadRequest)
+			resp := []byte(`{"typeName": "CSharpExceptionNameHere","message": "could not do something"}`)
+			w.Write(resp)
+		}))
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(context.Background(), runnerScaleSet.ID, owner)
+		assert.Nil(t, sessionClient)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "status=\"400 Bad Request\"")
+		assert.Contains(t, err.Error(), "activity_id=\""+exampleRequestID+"\"")
+
+		var ex actionsExceptionError
+		assert.True(t, errors.As(err, &ex))
+		assert.Equal(t, "CSharpExceptionNameHere", ex.ExceptionName)
+		assert.Equal(t, "could not do something", ex.Message)
+	})
+
+	t.Run("CreateMessageSession call is retried the correct amount of times", func(t *testing.T) {
+		owner := "foo"
+		runnerScaleSet := RunnerScaleSet{
+			ID:            1,
+			Name:          "ScaleSet",
+			CreatedOn:     time.Date(1, time.January, 1, 0, 0, 0, 0, time.UTC),
+			RunnerSetting: RunnerSetting{},
+		}
+
+		gotRetries := 0
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusInternalServerError)
+			gotRetries++
+		}))
+
+		retryMax := 3
+		retryWaitMax := 1 * time.Microsecond
+
+		wantRetries := retryMax + 1
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+			WithRetryMax(retryMax),
+			WithRetryWaitMax(retryWaitMax),
+		)
+		require.NoError(t, err)
+
+		_, err = client.MessageSessionClient(
+			ctx,
+			runnerScaleSet.ID,
+			owner,
+			WithRetryMax(retryMax),
+			WithRetryWaitMax(retryWaitMax),
+		)
+		assert.NotNil(t, err)
+		assert.Equalf(t, gotRetries, wantRetries, "CreateMessageSession got unexpected retry count: got=%v, want=%v", gotRetries, wantRetries)
+	})
+}
+
+func TestGetMessage(t *testing.T) {
+	ctx := context.Background()
+	auth := actionsAuth{
+		token: "token",
+	}
+
+	runnerScaleSetMessage := &RunnerScaleSetMessage{
+		MessageID: 1,
+	}
+
+	t.Run("Get Runner Scale Set Message", func(t *testing.T) {
+		want := runnerScaleSetMessage
+		response := []byte(`{"messageId":1,"messageType":"RunnerScaleSetJobMessages"}`)
+
+		var handleSessionRequest http.HandlerFunc
+		s := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			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)
+
+		got, err := sessionClient.GetMessage(ctx, 0, 10)
+		require.NoError(t, err)
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("GetMessage sets the last message id if not 0", func(t *testing.T) {
+		want := runnerScaleSetMessage
+		response := []byte(`{"messageId":1,"messageType":"RunnerScaleSetJobMessages"}`)
+		var handleSessionRequest http.HandlerFunc
+		s := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			q := r.URL.Query()
+			assert.Equal(t, "1", q.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)
+
+		got, err := sessionClient.GetMessage(ctx, 1, 10)
+		require.NoError(t, err)
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("Default retries on server error", func(t *testing.T) {
+		retryMax := 1
+
+		actualRetry := 0
+		expectedRetry := retryMax + 1
+
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			w.WriteHeader(http.StatusServiceUnavailable)
+			actualRetry++
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+			WithRetryMax(retryMax),
+			WithRetryWaitMax(1*time.Millisecond),
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(
+			ctx,
+			1,
+			"my-org",
+			WithRetryMax(retryMax),
+			WithRetryWaitMax(1*time.Millisecond),
+		)
+		require.NoError(t, err)
+
+		msg, err := sessionClient.GetMessage(ctx, 0, 10)
+		assert.Nil(t, msg)
+		assert.NotNil(t, err)
+		assert.Equalf(t, actualRetry, expectedRetry, "A retry was expected after the first request but got: %v", actualRetry)
+	})
+
+	t.Run("Message token expired", func(t *testing.T) {
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// create session
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			// refresh
+			if strings.Contains(r.URL.Path, "/sessions/") {
+				// just set the same session
+				handleSessionRequest(w, r)
+				return
+			}
+			w.WriteHeader(http.StatusUnauthorized)
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		msg, err := sessionClient.GetMessage(ctx, 0, 10)
+		assert.Nil(t, msg)
+		assert.ErrorIs(t, err, MessageQueueTokenExpiredError, "expected error to be MessageQueueTokenExpiredError but got: %v", err)
+	})
+
+	t.Run("Message token refreshed", func(t *testing.T) {
+		want := runnerScaleSetMessage
+		afterRefreshResponse := []byte(`{"messageId":1,"messageType":"RunnerScaleSetJobMessages"}`)
+		var handleSessionRequest http.HandlerFunc
+		type state int
+		const (
+			createSession state = iota
+			firstGetMessage
+			refreshToken
+			secondGetMessage
+		)
+		currentState := createSession
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// create session
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				require.Equal(t, createSession, currentState)
+				handleSessionRequest(w, r)
+				currentState = firstGetMessage
+				return
+			}
+			// refresh
+			if strings.Contains(r.URL.Path, "/sessions/") {
+				// just set the same session
+				require.Equal(t, refreshToken, currentState)
+				handleSessionRequest(w, r)
+				currentState = secondGetMessage
+				return
+			}
+			if currentState == firstGetMessage {
+				w.WriteHeader(http.StatusUnauthorized)
+				currentState = refreshToken
+				return
+			}
+			require.Equal(t, secondGetMessage, currentState)
+			w.Write(afterRefreshResponse)
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		got, err := sessionClient.GetMessage(ctx, 0, 10)
+		require.NoError(t, err)
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("Status code not found", func(t *testing.T) {
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			w.WriteHeader(http.StatusNotFound)
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		msg, err := sessionClient.GetMessage(ctx, 0, 10)
+		assert.Nil(t, msg)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "status=\"404 Not Found\"")
+		assert.Contains(t, err.Error(), "unknown error")
+	})
+
+	t.Run("Error when Content-Type is text/plain", func(t *testing.T) {
+		plainBody := "example plain text error"
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			w.Header().Set("Content-Type", "text/plain")
+			w.WriteHeader(http.StatusBadRequest)
+			w.Write([]byte(plainBody))
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		msg, err := sessionClient.GetMessage(ctx, 0, 10)
+		assert.Nil(t, msg)
+		assert.NotNil(t, err)
+		assert.Contains(t, err.Error(), "status=\"400 Bad Request\"")
+		assert.Contains(t, err.Error(), plainBody)
+	})
+
+	t.Run("Capacity error handling", func(t *testing.T) {
+		plainBody := "capacity error"
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			hc := r.Header.Get(HeaderScaleSetMaxCapacity)
+			c, err := strconv.Atoi(hc)
+			require.NoError(t, err)
+			assert.GreaterOrEqual(t, c, 0)
+			w.Header().Set("Content-Type", "text/plain")
+			w.WriteHeader(http.StatusBadRequest)
+			w.Write([]byte(plainBody))
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		msg, err := sessionClient.GetMessage(ctx, 0, 0)
+		assert.Nil(t, msg)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "status=\"400 Bad Request\"")
+		assert.Contains(t, err.Error(), plainBody)
+	})
+}
+
+func TestDeleteMessage(t *testing.T) {
+	ctx := context.Background()
+	auth := actionsAuth{
+		token: "token",
+	}
+
+	runnerScaleSetMessage := &RunnerScaleSetMessage{
+		MessageID: 1,
+	}
+
+	t.Run("Delete existing message", func(t *testing.T) {
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			w.WriteHeader(http.StatusNoContent)
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		err = sessionClient.DeleteMessage(ctx, runnerScaleSetMessage.MessageID)
+		assert.Nil(t, err)
+	})
+
+	t.Run("Message token expired", func(t *testing.T) {
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// create session
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+
+			// refresh
+			if strings.Contains(r.URL.Path, "/sessions/") {
+				// just set the same session
+				handleSessionRequest(w, r)
+				return
+			}
+			w.WriteHeader(http.StatusUnauthorized)
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		err = sessionClient.DeleteMessage(ctx, 0)
+		require.NotNil(t, err)
+		assert.ErrorIs(t, err, MessageQueueTokenExpiredError, "expected error to be MessageQueueTokenExpiredError but got: %v", err)
+	})
+
+	t.Run("message token refreshed", func(t *testing.T) {
+		type state int
+		const (
+			createSession state = iota
+			firstDeleteMessage
+			refreshToken
+			secondDeleteMessage
+		)
+		currentState := createSession
+
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// create session
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				require.Equal(t, createSession, currentState)
+				handleSessionRequest(w, r)
+				currentState = firstDeleteMessage
+				return
+			}
+			// refresh
+			if strings.Contains(r.URL.Path, "/sessions/") {
+				// just set the same session
+				require.Equal(t, refreshToken, currentState)
+				handleSessionRequest(w, r)
+				currentState = secondDeleteMessage
+				return
+			}
+			if currentState == firstDeleteMessage {
+				w.WriteHeader(http.StatusUnauthorized)
+				currentState = refreshToken
+				return
+			}
+			require.Equal(t, secondDeleteMessage, currentState)
+			w.WriteHeader(http.StatusNoContent)
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		err = sessionClient.DeleteMessage(ctx, 0)
+		require.NoError(t, err)
+	})
+
+	t.Run("Error when Content-Type is text/plain", func(t *testing.T) {
+		plainBody := "example plain text error"
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			w.Header().Set("Content-Type", "text/plain")
+			w.WriteHeader(http.StatusBadRequest)
+			w.Write([]byte(plainBody))
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		err = sessionClient.DeleteMessage(ctx, runnerScaleSetMessage.MessageID)
+		require.NotNil(t, err)
+		assert.Contains(t, err.Error(), "status=\"400 Bad Request\"")
+		assert.Contains(t, err.Error(), plainBody)
+	})
+
+	t.Run("Default retries on server error", func(t *testing.T) {
+		actualRetry := 0
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			w.WriteHeader(http.StatusServiceUnavailable)
+			actualRetry++
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		retryMax := 1
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+			WithRetryMax(retryMax),
+			WithRetryWaitMax(1*time.Nanosecond),
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(
+			ctx,
+			1,
+			"my-org",
+			WithRetryMax(retryMax),
+			WithRetryWaitMax(1*time.Nanosecond),
+		)
+		require.NoError(t, err)
+
+		err = sessionClient.DeleteMessage(ctx, runnerScaleSetMessage.MessageID)
+		assert.NotNil(t, err)
+		expectedRetry := retryMax + 1
+		assert.Equalf(t, actualRetry, expectedRetry, "A retry was expected after the first request but got: %v", actualRetry)
+	})
+
+	t.Run("No message found", func(t *testing.T) {
+		want := (*RunnerScaleSetMessage)(nil)
+		rsl, err := json.Marshal(want)
+		require.NoError(t, err)
+
+		var handleSessionRequest http.HandlerFunc
+
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			w.Write(rsl)
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		err = sessionClient.DeleteMessage(ctx, runnerScaleSetMessage.MessageID+1)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "unexpected status code")
+	})
+}
+
+func TestAcquireJobs(t *testing.T) {
+	ctx := context.Background()
+	auth := actionsAuth{
+		token: "token",
+	}
+
+	t.Run("Acquire jobs successfully", func(t *testing.T) {
+		requestIDs := []int64{1, 2}
+		want := []int64{1, 2}
+		response := []byte(`{"count":2,"value":[1,2]}`)
+
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.Contains(r.URL.Path, "acquirejobs") {
+				assert.Equal(t, http.MethodPost, r.Method)
+				assert.Contains(t, r.URL.Path, "acquirejobs")
+				assert.True(t, strings.HasPrefix(r.Header.Get("Authorization"), "Bearer"), "expected Bearer authorization header")
+
+				var gotIDs []int64
+				err := json.NewDecoder(r.Body).Decode(&gotIDs)
+				require.NoError(t, err)
+				assert.Equal(t, requestIDs, gotIDs)
+
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(http.StatusOK)
+				w.Write(response)
+				return
+			}
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			if strings.Contains(r.URL.Path, "/sessions/") {
+				handleSessionRequest(w, r)
+				return
+			}
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		got, err := sessionClient.AcquireJobs(ctx, requestIDs)
+		require.NoError(t, err)
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("Message token expired", func(t *testing.T) {
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.Contains(r.URL.Path, "acquirejobs") {
+				w.WriteHeader(http.StatusUnauthorized)
+				return
+			}
+			// create session
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			// refresh
+			if strings.Contains(r.URL.Path, "/sessions/") {
+				handleSessionRequest(w, r)
+				return
+			}
+			w.WriteHeader(http.StatusUnauthorized)
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		got, err := sessionClient.AcquireJobs(ctx, []int64{1})
+		assert.Nil(t, got)
+		assert.ErrorIs(t, err, MessageQueueTokenExpiredError, "expected error to be MessageQueueTokenExpiredError but got: %v", err)
+	})
+
+	t.Run("Message token refreshed", func(t *testing.T) {
+		want := []int64{1, 2}
+		afterRefreshResponse := []byte(`{"count":2,"value":[1,2]}`)
+
+		var handleSessionRequest http.HandlerFunc
+		type state int
+		const (
+			createSession state = iota
+			firstAcquire
+			refreshToken
+			secondAcquire
+		)
+		currentState := createSession
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.Contains(r.URL.Path, "acquirejobs") {
+				if currentState == firstAcquire {
+					w.WriteHeader(http.StatusUnauthorized)
+					currentState = refreshToken
+					return
+				}
+				require.Equal(t, secondAcquire, currentState)
+				w.Header().Set("Content-Type", "application/json")
+				w.Write(afterRefreshResponse)
+				return
+			}
+			// create session
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				require.Equal(t, createSession, currentState)
+				handleSessionRequest(w, r)
+				currentState = firstAcquire
+				return
+			}
+			// refresh
+			if strings.Contains(r.URL.Path, "/sessions/") {
+				require.Equal(t, refreshToken, currentState)
+				handleSessionRequest(w, r)
+				currentState = secondAcquire
+				return
+			}
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		got, err := sessionClient.AcquireJobs(ctx, []int64{1, 2})
+		require.NoError(t, err)
+		assert.Equal(t, want, got)
+	})
+
+	t.Run("Server error", func(t *testing.T) {
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.Contains(r.URL.Path, "acquirejobs") {
+				w.WriteHeader(http.StatusInternalServerError)
+				return
+			}
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			if strings.Contains(r.URL.Path, "/sessions/") {
+				handleSessionRequest(w, r)
+				return
+			}
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		retryMax := 1
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+			WithRetryMax(retryMax),
+			WithRetryWaitMax(1*time.Nanosecond),
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(
+			ctx,
+			1,
+			"my-org",
+			WithRetryMax(retryMax),
+			WithRetryWaitMax(1*time.Nanosecond),
+		)
+		require.NoError(t, err)
+
+		got, err := sessionClient.AcquireJobs(ctx, []int64{1})
+		assert.Nil(t, got)
+		assert.NotNil(t, err)
+	})
+
+	t.Run("Empty request IDs", func(t *testing.T) {
+		response := []byte(`{"count":0,"value":[]}`)
+
+		var handleSessionRequest http.HandlerFunc
+		server := newActionsServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if strings.Contains(r.URL.Path, "acquirejobs") {
+				assert.Equal(t, http.MethodPost, r.Method)
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(http.StatusOK)
+				w.Write(response)
+				return
+			}
+			if strings.HasSuffix(r.URL.Path, "sessions") {
+				handleSessionRequest(w, r)
+				return
+			}
+			if strings.Contains(r.URL.Path, "/sessions/") {
+				handleSessionRequest(w, r)
+				return
+			}
+		}))
+		handleSessionRequest = newTestSessionRequestHandler(t, server.testRunnerScaleSetSession())
+
+		client, err := newClient(
+			testSystemInfo,
+			server.configURLForOrg("my-org"),
+			auth,
+		)
+		require.NoError(t, err)
+
+		sessionClient, err := client.MessageSessionClient(ctx, 1, "my-org")
+		require.NoError(t, err)
+
+		got, err := sessionClient.AcquireJobs(ctx, []int64{})
+		require.NoError(t, err)
+		assert.Empty(t, got)
+	})
+}

types.go

@@ -12,16 +12,12 @@ type MessageType string
 
 // message types
 const (
+	MessageTypeJobAvailable MessageType = "JobAvailable"
 	MessageTypeJobAssigned  MessageType = "JobAssigned"
 	MessageTypeJobStarted   MessageType = "JobStarted"
 	MessageTypeJobCompleted MessageType = "JobCompleted"
 )
 
-type Int64List struct {
-	Count int     `json:"count"`
-	Value []int64 `json:"value"`
-}
-
 type JobAvailable struct {
 	AcquireJobURL string `json:"acquireJobUrl"`
 	JobMessageBase
@@ -73,7 +69,7 @@ type Label struct {
 type RunnerGroup struct {
 	ID        int    `json:"id"`
 	Name      string `json:"name"`
-	Size      int64  `json:"size"`
+	Size      int    `json:"size"`
 	IsDefault bool   `json:"isDefaultGroup"`
 }
 
@@ -99,18 +95,32 @@ type RunnerScaleSetJitRunnerSetting struct {
 	WorkFolder string `json:"workFolder"`
 }
 
-type RunnerScaleSetMessage struct {
-	MessageID   int64                    `json:"messageId"`
+type runnerScaleSetMessageResponse struct {
+	MessageID   int                      `json:"messageId"`
 	MessageType string                   `json:"messageType"`
 	Body        string                   `json:"body"`
 	Statistics  *RunnerScaleSetStatistic `json:"statistics"`
 }
 
+type RunnerScaleSetMessage struct {
+	MessageID            int
+	Statistics           *RunnerScaleSetStatistic
+	JobAvailableMessages []*JobAvailable
+	JobAssignedMessages  []*JobAssigned
+	JobStartedMessages   []*JobStarted
+	JobCompletedMessages []*JobCompleted
+}
+
 type runnerScaleSetsResponse struct {
 	Count           int              `json:"count"`
 	RunnerScaleSets []RunnerScaleSet `json:"value"`
 }
 
+type acquireJobsResponse struct {
+	Count int     `json:"count"`
+	Value []int64 `json:"value"`
+}
+
 type RunnerScaleSetSession struct {
 	SessionID               uuid.UUID                `json:"sessionId,omitempty"`
 	OwnerName               string                   `json:"ownerName,omitempty"`
@@ -131,8 +141,6 @@ type RunnerScaleSetStatistic struct {
 }
 
 type RunnerSetting struct {
-	Ephemeral     bool `json:"ephemeral,omitempty"`
-	IsElastic     bool `json:"isElastic,omitempty"`
 	DisableUpdate bool `json:"disableUpdate,omitempty"`
 }