queen-mq 0.12.4 → 0.14.0

package/README.md

@@ -250,6 +250,47 @@ for (const message of messages) {
 }
 ```
 
+### Multi-Partition Pop (Drain Many Partitions Per Call)
+
+```javascript
+// One round-trip drains up to 200 messages spread across up to 50 partitions.
+// batch(200) is the GLOBAL cap on total messages; partitions(50) is the
+// hard cap on partitions claimed. All claimed partitions share one leaseId
+// — a single renew() call extends every partition's lease atomically.
+const messages = await queen.queue('events')
+  .batch(200)
+  .partitions(50)
+  .wait(true)
+  .pop()
+
+// Each message carries its own partition info (per-message partitionId,
+// partition name, leaseId, consumerGroup) — ACK and renew always work
+// message-by-message regardless of how many partitions the batch spans.
+for (const m of messages) {
+  console.log(`from ${m.partition}:`, m.data)
+}
+
+// Same builder works on .consume() for long-running workers
+await queen.queue('events')
+  .batch(100)
+  .partitions(8)
+  .consume(async (msgs) => {
+    for (const m of msgs) await process(m.data)
+  })
+```
+
+**When to use:** queues with many partitions where each partition only has
+a handful of new messages per polling interval (per-customer event streams,
+per-tenant work queues, per-device telemetry). Reduces network round-trips
+from O(P) to O(P / N) while preserving per-partition FIFO ordering.
+
+**When not to use:** few partitions, or each one busy enough to fill
+`batch(B)` on its own. Default is `partitions(1)` which preserves the
+legacy single-partition behaviour.
+
+`.partitions(N)` only applies to **wildcard** pops; specifying
+`.partition('name')` ignores the cap.
+
 ### Transactions (Atomic Operations)
 
 ```javascript
@@ -451,6 +492,7 @@ await queen.queue('q').buffer({ messageCount: 100, timeMillis: 1000 }).push([...
 const msgs = await queen.queue('q').pop()
 const msgs = await queen.queue('q').batch(10).pop()
 const msgs = await queen.queue('q').batch(10).wait(true).pop()
+const msgs = await queen.queue('q').batch(200).partitions(50).pop()  // multi-partition pop
 ```
 
 ### Consume

package/client-v2/Queen.js

@@ -350,6 +350,12 @@ export class Queen {
       }
     }
 
+    // Dedupe: with v4 multi-partition pop, all messages in one batch share
+    // the same leaseId (one renew_lease_v2 call extends every claimed
+    // partition_consumers row). Without this, callers passing the full
+    // messages array would issue N redundant identical HTTP calls.
+    leaseIds = [...new Set(leaseIds)]
+
     if (leaseIds.length === 0) {
       logger.warn('Queen.renew', 'No valid lease IDs found for renewal')
       return { success: false, error: 'No valid lease IDs found for renewal' }

package/client-v2/README.md

@@ -1605,6 +1605,11 @@ const msgs = await queen.queue('q').batch(10).wait(true).pop()
 
 // Pop from partition
 const msgs = await queen.queue('q').partition('p1').pop()
+
+// Multi-partition pop: drain up to 50 partitions in one round-trip,
+// capped at 200 total messages. All partitions share one leaseId.
+// Each returned message carries its own partitionId / partition / leaseId.
+const msgs = await queen.queue('q').batch(200).partitions(50).pop()
 ```
 
 ### Consume
@@ -1625,6 +1630,12 @@ await queen.queue('q').concurrency(5).consume(async (msg) => { /* 5 parallel wor
 // Consume from partition
 await queen.queue('q').partition('p1').consume(async (msg) => { /* process */ })
 
+// Multi-partition consume: drain up to 8 partitions per poll.
+// batch(100) is a global cap on total messages across all claimed partitions.
+await queen.queue('q').batch(100).partitions(8).consume(async (msgs) => {
+  for (const m of msgs) await process(m.data)  // m.partitionId / m.partition baked in
+})
+
 // Consume with consumer group
 await queen.queue('q').group('my-group').consume(async (msg) => { /* process */ })
 
@@ -1863,7 +1874,7 @@ Example log output:
 Here's a complete example showing many features together:
 
 ```javascript
-import { Queen } from './client-js/client-v2/index.js'
+import { Queen } from './index.js'
 
 const queen = new Queen('http://localhost:6632')
 

package/client-v2/builders/QueueBuilder.js

@@ -37,6 +37,7 @@ export class QueueBuilder {
   #subscriptionMode = CONSUME_DEFAULTS.subscriptionMode
   #subscriptionFrom = CONSUME_DEFAULTS.subscriptionFrom
   #each = false
+  #maxPartitions = 1
 
   // Buffer options
   #bufferOptions = null
@@ -190,6 +191,22 @@ export class QueueBuilder {
     return this
   }
 
+  /**
+   * Pop messages from up to N partitions in a single call (v4 multi-partition pop).
+   *
+   * Use this to drain many sparsely-loaded partitions efficiently. With
+   * `partitions(N)`, the global `batch(B)` budget is shared across all
+   * claimed partitions: at most B total messages, drawn from up to N
+   * partitions, in a single network round-trip. All N share one leaseId
+   * (renewing once extends them all).
+   *
+   * Default 1 = legacy single-partition behavior.
+   */
+  partitions(n) {
+    this.#maxPartitions = Math.max(1, n)
+    return this
+  }
+
   limit(count) {
     this.#limit = count
     return this
@@ -251,6 +268,7 @@ export class QueueBuilder {
       subscriptionMode: this.#subscriptionMode,
       subscriptionFrom: this.#subscriptionFrom,
       each: this.#each,
+      maxPartitions: this.#maxPartitions,
       signal: options.signal
     }
 
@@ -266,6 +284,20 @@ export class QueueBuilder {
     return this
   }
 
+  /**
+   * Set the long-poll timeout in milliseconds. Only meaningful when wait(true)
+   * is also set. Default is CONSUME_DEFAULTS.timeoutMillis (30000ms).
+   *
+   * Use a short timeout in tight consumer loops where you want the long-poll
+   * to release the request frequently (e.g. so a shutdown flag is observed
+   * promptly). The HTTP layer adds 5s of slack so the server-side timeout
+   * always fires first.
+   */
+  timeoutMillis(millis) {
+    this.#timeoutMillis = Math.max(1, millis)
+    return this
+  }
+
   async pop() {
     logger.log('QueueBuilder.pop', { queue: this.#queueName, partition: this.#partition, namespace: this.#namespace, task: this.#task, batch: this.#batch, wait: this.#wait, group: this.#group })
     
@@ -289,6 +321,7 @@ export class QueueBuilder {
       if (effectiveAutoAck) params.append('autoAck', 'true')
       if (this.#subscriptionMode) params.append('subscriptionMode', this.#subscriptionMode)
       if (this.#subscriptionFrom) params.append('subscriptionFrom', this.#subscriptionFrom)
+      if (this.#maxPartitions > 1) params.append('partitions', this.#maxPartitions.toString())
 
       // Generate affinity key for consistent routing to same backend
       const affinityKey = this.#getAffinityKey()

package/client-v2/consumer/ConsumerManager.js

@@ -48,6 +48,7 @@ export class ConsumerManager {
       subscriptionMode,
       subscriptionFrom,
       each,
+      maxPartitions,
       signal
     } = options
 
@@ -67,7 +68,7 @@ export class ConsumerManager {
 
     // Build the path and params for pop requests
     const path = this.#buildPath(queue, partition, namespace, task)
-    const baseParams = this.#buildParams(batch, wait, timeoutMillis, group, subscriptionMode, subscriptionFrom, namespace, task, autoAck)
+    const baseParams = this.#buildParams(batch, wait, timeoutMillis, group, subscriptionMode, subscriptionFrom, namespace, task, autoAck, maxPartitions)
     
     // Generate affinity key for consistent routing to same backend
     const affinityKey = this.#getAffinityKey(queue, partition, namespace, task, group)
@@ -389,7 +390,7 @@ export class ConsumerManager {
     throw new Error('Must specify queue, namespace, or task')
   }
 
-  #buildParams(batch, wait, timeoutMillis, group, subscriptionMode, subscriptionFrom, namespace, task, autoAck) {
+  #buildParams(batch, wait, timeoutMillis, group, subscriptionMode, subscriptionFrom, namespace, task, autoAck, maxPartitions) {
     const params = new URLSearchParams({
       batch: batch.toString(),
       wait: wait.toString(),
@@ -401,6 +402,8 @@ export class ConsumerManager {
     if (subscriptionFrom) params.append('subscriptionFrom', subscriptionFrom)
     if (namespace) params.append('namespace', namespace)
     if (task) params.append('task', task)
+    // v4 multi-partition pop: drain up to N sparse partitions per call.
+    if (maxPartitions && maxPartitions > 1) params.append('partitions', maxPartitions.toString())
     // NEVER send autoAck for consume - client always manages acking
     // autoAck is only for pop() where server auto-acks immediately
 

package/client-v2/utils/defaults.js

@@ -43,7 +43,8 @@ export const CONSUME_DEFAULTS = {
   renewLease: false,                   // No auto-renewal
   renewLeaseIntervalMillis: null,      // Auto-renewal interval when enabled
   subscriptionMode: null,              // No subscription mode (standard queue mode)
-  subscriptionFrom: null               // No subscription start point
+  subscriptionFrom: null,              // No subscription start point
+  maxPartitions: 1                     // v4 multi-partition pop cap (1 = legacy single-partition)
 }
 
 export const POP_DEFAULTS = {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "queen-mq",
-  "version": "0.12.4",
+  "version": "0.14.0",
   "type": "module",
   "description": "High-performance C++ message queue backed by PostgreSQL",
   "main": "client-v2/index.js",

package/test-v2/auth.js

@@ -0,0 +1,340 @@
+// Tests for server-stamped producer identity (issue #23, feature A).
+//
+// The goal of feature A is: when JWT auth is enabled, the server attaches the
+// authenticated client's `sub` claim to every pushed message, and a client
+// cannot spoof another producer's identity by supplying `producerSub` in the
+// request body.
+//
+// Three environments are exercised here:
+//
+//   1. No auth (default dev setup): `producer_sub` must be NULL; clients that
+//      send `producerSub` in the body must still get NULL (the field is not
+//      client-settable).
+//   2. JWT auth enabled (set JWT_SECRET env to match the server's
+//      JWT_SECRET): `producer_sub` must equal the `sub` claim.
+//   3. JWT auth enabled + spoofing attempt: body-supplied `producerSub` must
+//      be ignored; stored value is the JWT `sub`.
+//
+// Tests gate themselves on the `JWT_SECRET` env var so the suite can run
+// unattended in either server configuration.
+
+import crypto from 'crypto'
+import { dbPool } from './run.js'
+
+const SERVER_URL = process.env.QUEEN_URL || 'http://localhost:6632'
+const JWT_SECRET = process.env.JWT_SECRET || ''
+const JWT_ENABLED = JWT_SECRET.length > 0
+
+// ---------------------------------------------------------------------------
+// Minimal HS256 JWT signer (avoids pulling a dependency just for tests)
+// ---------------------------------------------------------------------------
+function b64url(buf) {
+    return Buffer.from(buf).toString('base64')
+        .replace(/=/g, '').replace(/\+/g, '-').replace(/\//g, '_')
+}
+
+function signHs256Jwt(payload, secret) {
+    const header = { alg: 'HS256', typ: 'JWT' }
+    const encHeader = b64url(JSON.stringify(header))
+    const encPayload = b64url(JSON.stringify(payload))
+    const signingInput = `${encHeader}.${encPayload}`
+    const sig = crypto.createHmac('sha256', secret).update(signingInput).digest()
+    return `${signingInput}.${b64url(sig)}`
+}
+
+function makeToken(sub, role = 'read-write') {
+    const now = Math.floor(Date.now() / 1000)
+    return signHs256Jwt(
+        { sub, username: sub, role, iat: now, exp: now + 3600 },
+        JWT_SECRET
+    )
+}
+
+// ---------------------------------------------------------------------------
+// Ground-truth helper: read producer_sub directly from Postgres, bypassing
+// the API so we can tell the difference between a server that actually stores
+// the value and one that just echoes the client's input.
+// ---------------------------------------------------------------------------
+async function getStoredProducerSub(queueName, transactionId) {
+    const res = await dbPool.query(
+        `SELECT m.producer_sub
+         FROM queen.messages m
+         JOIN queen.partitions p ON p.id = m.partition_id
+         JOIN queen.queues q ON q.id = p.queue_id
+         WHERE q.name = $1 AND m.transaction_id = $2`,
+        [queueName, transactionId]
+    )
+    return res.rows.length > 0 ? res.rows[0].producer_sub : null
+}
+
+// Raw HTTP push so we can control the Authorization header and inject
+// arbitrary extra fields (like a forged `producerSub`) per-test.
+async function httpPush({ queue, partition = 'Default', transactionId, data, bearerToken, extraItemFields = {} }) {
+    const headers = { 'Content-Type': 'application/json' }
+    if (bearerToken) headers['Authorization'] = `Bearer ${bearerToken}`
+
+    const body = {
+        items: [{
+            queue,
+            partition,
+            transactionId,
+            payload: data,
+            ...extraItemFields
+        }]
+    }
+
+    const res = await fetch(`${SERVER_URL}/api/v1/push`, {
+        method: 'POST',
+        headers,
+        body: JSON.stringify(body)
+    })
+    const text = await res.text()
+    if (!res.ok) {
+        throw new Error(`push failed ${res.status}: ${text}`)
+    }
+    return text ? JSON.parse(text) : null
+}
+
+// Raw HTTP configure (queue creation) with optional bearer.
+async function httpConfigure(queue, bearerToken) {
+    const headers = { 'Content-Type': 'application/json' }
+    if (bearerToken) headers['Authorization'] = `Bearer ${bearerToken}`
+    const res = await fetch(`${SERVER_URL}/api/v1/configure`, {
+        method: 'POST',
+        headers,
+        body: JSON.stringify({ queue, options: {} })
+    })
+    if (!res.ok) {
+        const text = await res.text()
+        throw new Error(`configure failed ${res.status}: ${text}`)
+    }
+}
+
+// Direct SP call via dbPool — bypasses HTTP and JWT entirely. Used by the
+// SP-level tests below so they work regardless of whether the server has
+// auth enabled or not.
+async function callPushMessagesV3(items) {
+    await dbPool.query(
+        'SELECT queen.push_messages_v3($1::jsonb)',
+        [JSON.stringify(items)]
+    )
+}
+
+// Direct SP pop. Returns the `messages` array for the first (only) partition.
+async function callPopSpecificBatch(queueName) {
+    const reqs = [{
+        idx: 0,
+        queue_name: queueName,
+        partition_name: 'Default',
+        consumer_group: '__QUEUE_MODE__',
+        lease_seconds: 60,
+        worker_id: `test-worker-${Date.now()}`,
+        batch_size: 10,
+        sub_mode: 'all',
+        sub_from: ''
+    }]
+    const res = await dbPool.query(
+        'SELECT queen.pop_specific_batch($1::jsonb) as out',
+        [JSON.stringify(reqs)]
+    )
+    const out = res.rows[0].out
+    return out[0]?.result?.messages || []
+}
+
+// ===========================================================================
+// TEST 1: Stored procedure round-trip (SQL-only, no HTTP, no JWT)
+// ---------------------------------------------------------------------------
+// Exercises the schema + SP contract end-to-end purely via SQL. If this fails,
+// the migration or 001_push.sql / 002_pop_unified.sql is wrong. Runs in any
+// server configuration because it talks to Postgres directly.
+// ===========================================================================
+export async function producerSubRoundTripViaStoredProcedure(client) {
+    const q = 'test-auth-sp-roundtrip'
+    const tx = `tx-sp-${Date.now()}`
+
+    await callPushMessagesV3([{
+        queue: q,
+        partition: 'Default',
+        transactionId: tx,
+        payload: { x: 1 },
+        producerSub: 'sp-test-sub'
+    }])
+
+    const stored = await getStoredProducerSub(q, tx)
+    if (stored !== 'sp-test-sub') {
+        return { success: false, message: `Expected producer_sub='sp-test-sub' in DB, got ${JSON.stringify(stored)}` }
+    }
+
+    const popped = await callPopSpecificBatch(q)
+    const target = popped.find(m => m.transactionId === tx)
+    if (!target) {
+        return { success: false, message: `Expected to find tx ${tx} in pop result` }
+    }
+    if (target.producerSub !== 'sp-test-sub') {
+        return { success: false, message: `Pop returned producerSub=${JSON.stringify(target.producerSub)}, expected 'sp-test-sub'` }
+    }
+
+    return { success: true, message: 'producer_sub round-trips through push_messages_v3 and pop exposes it' }
+}
+
+// ===========================================================================
+// TEST 2: SP with no producerSub => producer_sub is NULL
+// ---------------------------------------------------------------------------
+// Guards against an accidental "default to empty string" regression in the SP.
+// ===========================================================================
+export async function producerSubNullWhenNotProvided(client) {
+    const q = 'test-auth-sp-null'
+    const tx = `tx-sp-null-${Date.now()}`
+
+    await callPushMessagesV3([{
+        queue: q,
+        partition: 'Default',
+        transactionId: tx,
+        payload: { x: 1 }
+        // no producerSub
+    }])
+
+    const stored = await getStoredProducerSub(q, tx)
+    if (stored !== null) {
+        return { success: false, message: `Expected NULL producer_sub, got ${JSON.stringify(stored)}` }
+    }
+
+    const popped = await callPopSpecificBatch(q)
+    const target = popped.find(m => m.transactionId === tx)
+    if (!target) {
+        return { success: false, message: `Expected to find tx ${tx} in pop result` }
+    }
+    if (target.producerSub != null) {
+        return { success: false, message: `Pop returned producerSub=${JSON.stringify(target.producerSub)}, expected null` }
+    }
+
+    return { success: true, message: 'Omitting producerSub yields NULL in DB and null in pop response' }
+}
+
+// ===========================================================================
+// TEST 3: SP treats empty string producerSub as NULL
+// ---------------------------------------------------------------------------
+// The SP uses NULLIF(..., '') to avoid storing a literal empty string when a
+// client accidentally sends "". Verifies that invariant.
+// ===========================================================================
+export async function producerSubEmptyStringStoredAsNull(client) {
+    const q = 'test-auth-sp-empty'
+    const tx = `tx-sp-empty-${Date.now()}`
+
+    await callPushMessagesV3([{
+        queue: q,
+        partition: 'Default',
+        transactionId: tx,
+        payload: { x: 1 },
+        producerSub: ''
+    }])
+
+    const stored = await getStoredProducerSub(q, tx)
+    if (stored !== null) {
+        return { success: false, message: `Expected NULL for empty-string producerSub, got ${JSON.stringify(stored)}` }
+    }
+
+    return { success: true, message: 'Empty-string producerSub is stored as NULL' }
+}
+
+// ===========================================================================
+// TEST 4: HTTP push without JWT => producer_sub is NULL
+// ---------------------------------------------------------------------------
+// Runs only when the server is NOT configured with JWT (typical dev setup).
+// Also checks that a client sending `producerSub` in the push body does NOT
+// end up with it persisted (the field is not client-settable under any mode).
+// ===========================================================================
+export async function producerSubIgnoredFromBodyWithoutAuth(client) {
+    if (JWT_ENABLED) {
+        return { success: true, message: 'Skipped (JWT_SECRET is set - run producerSubStampedFromJwt instead)' }
+    }
+
+    const q = 'test-auth-http-no-jwt'
+    const tx = `tx-noauth-${Date.now()}`
+
+    const queue = await client.queue(q).create()
+    if (!queue.configured) {
+        return { success: false, message: 'Queue not created' }
+    }
+
+    // Spoofing attempt against an unauthenticated server.
+    await httpPush({
+        queue: q,
+        transactionId: tx,
+        data: { hello: 'world' },
+        extraItemFields: { producerSub: 'attacker-no-jwt' }
+    })
+
+    const stored = await getStoredProducerSub(q, tx)
+    if (stored !== null) {
+        return { success: false, message: `Expected NULL producer_sub (auth disabled), got ${JSON.stringify(stored)} - client was able to set it!` }
+    }
+
+    return { success: true, message: 'Body-supplied producerSub ignored when auth disabled; stored as NULL' }
+}
+
+// ===========================================================================
+// TEST 5: HTTP push WITH JWT => producer_sub = sub claim
+// ---------------------------------------------------------------------------
+// Requires:
+//   - Server started with JWT_ENABLED=true JWT_ALGORITHM=HS256 JWT_SECRET=<x>
+//   - Test run with JWT_SECRET=<x>
+// ===========================================================================
+export async function producerSubStampedFromJwt(client) {
+    if (!JWT_ENABLED) {
+        return { success: true, message: 'Skipped (set JWT_SECRET env var matching server to run)' }
+    }
+
+    const q = 'test-auth-http-jwt-stamp'
+    const tx = `tx-jwt-${Date.now()}`
+    const token = makeToken('alice-producer')
+
+    await httpConfigure(q, token)
+    await httpPush({
+        queue: q,
+        transactionId: tx,
+        data: { hello: 'world' },
+        bearerToken: token
+    })
+
+    const stored = await getStoredProducerSub(q, tx)
+    if (stored !== 'alice-producer') {
+        return { success: false, message: `Expected producer_sub='alice-producer', got ${JSON.stringify(stored)}` }
+    }
+
+    return { success: true, message: 'producer_sub stamped from authenticated JWT sub claim' }
+}
+
+// ===========================================================================
+// TEST 6: Spoofing producerSub in body is ignored even with valid JWT
+// ---------------------------------------------------------------------------
+// This is the core anti-impersonation invariant of issue #23.
+// ===========================================================================
+export async function producerSubSpoofingIgnoredWithJwt(client) {
+    if (!JWT_ENABLED) {
+        return { success: true, message: 'Skipped (set JWT_SECRET env var matching server to run)' }
+    }
+
+    const q = 'test-auth-http-jwt-spoof'
+    const tx = `tx-spoof-${Date.now()}`
+    const token = makeToken('legit-producer')
+
+    await httpConfigure(q, token)
+    await httpPush({
+        queue: q,
+        transactionId: tx,
+        data: { hello: 'world' },
+        bearerToken: token,
+        extraItemFields: { producerSub: 'attacker' }
+    })
+
+    const stored = await getStoredProducerSub(q, tx)
+    if (stored !== 'legit-producer') {
+        return {
+            success: false,
+            message: `Impersonation not prevented: stored producer_sub=${JSON.stringify(stored)}, expected 'legit-producer'`
+        }
+    }
+
+    return { success: true, message: 'Body-supplied producerSub ignored; stored sub is from validated JWT' }
+}

package/test-v2/consume.js

@@ -612,4 +612,100 @@ export async function autoRenewLease(client) {
     })    
 
     return { success: shouldFailHasFailed === false && shouldPassHasPass === true, message: 'Auto renew lease test completed successfully' }
+}
+
+
+// ============================================================================
+// v4 Multi-Partition Consume Tests
+// ============================================================================
+//
+// Exercises the .partitions(N) builder method against the high-level
+// consume() API (which calls pop_unified_batch_v4 under the hood).
+
+export async function testConsumerMultiPartition(client) {
+    const queueName = 'test-queue-v2-v4-consume-multi'
+    const queue = await client.queue(queueName).create()
+    if (!queue.configured) {
+        return { success: false, message: 'Queue not created' }
+    }
+
+    // Seed 3 partitions × 4 messages each = 12 total.
+    for (let p = 0; p < 3; p++) {
+        const payloads = []
+        for (let m = 0; m < 4; m++) payloads.push({ data: { p, m } })
+        await client.queue(queueName).partition(`p${p}`).push(payloads)
+    }
+
+    // Allow the post-push partition_lookup follow-up to settle.
+    await new Promise(r => setTimeout(r, 500))
+
+    let received = []
+
+    // .partitions(3).batch(100) → up to 100 msgs from up to 3 partitions
+    // in a single pop. .limit(1) ends the consume after one batch.
+    await client
+    .queue(queueName)
+    .batch(100)
+    .partitions(3)
+    .wait(false)
+    .limit(1)
+    .consume(async msgs => {
+        received = msgs
+    })
+
+    if (received.length !== 12) {
+        return { success: false, message: `Expected 12 messages across 3 partitions, got ${received.length}` }
+    }
+
+    const distinctPartitionIds = new Set(received.map(m => m.partitionId))
+    const distinctLeaseIds = new Set(received.map(m => m.leaseId).filter(Boolean))
+
+    if (distinctPartitionIds.size !== 3) {
+        return { success: false, message: `Expected 3 distinct partitionIds, got ${distinctPartitionIds.size}` }
+    }
+    // consume() uses client-side auto-ack which means leaseId may be
+    // empty if server-side autoAck is used. Accept either: shared lease
+    // id (1 distinct value) or all-empty (0).
+    if (distinctLeaseIds.size > 1) {
+        return { success: false, message: `All messages must share at most one leaseId, got ${distinctLeaseIds.size}` }
+    }
+
+    return { success: true, message: 'Consumer with .partitions(3) drained 3 partitions in one batch' }
+}
+
+export async function testConsumerMultiPartitionGlobalCap(client) {
+    const queueName = 'test-queue-v2-v4-consume-cap'
+    const queue = await client.queue(queueName).create()
+    if (!queue.configured) {
+        return { success: false, message: 'Queue not created' }
+    }
+
+    // Seed 5 partitions × 50 messages each = 250 total available.
+    for (let p = 0; p < 5; p++) {
+        const payloads = []
+        for (let m = 0; m < 50; m++) payloads.push({ data: { p, m } })
+        await client.queue(queueName).partition(`p${p}`).push(payloads)
+    }
+
+    await new Promise(r => setTimeout(r, 500))
+
+    let received = []
+
+    // batch(20) is the GLOBAL cap. With partitions(5) we can claim up to
+    // 5 partitions, but the total returned is bounded at 20 messages.
+    await client
+    .queue(queueName)
+    .batch(20)
+    .partitions(5)
+    .wait(false)
+    .limit(1)
+    .consume(async msgs => {
+        received = msgs
+    })
+
+    if (received.length !== 20) {
+        return { success: false, message: `batch(20) must be a hard global cap, got ${received.length}` }
+    }
+
+    return { success: true, message: 'Global batch cap respected across multi-partition consume' }
 }