Rate limit improvements

Sources/Strand/Activity.swift

@@ -74,80 +74,6 @@ extension ActivityCancellationType: PostgresCodable {
     }
 }
 
-// MARK: - RateLimit
-
-/// Leaky-bucket rate limit for an activity.  When set, each enqueue atomically
-/// claims a time slot; the run waits in PENDING until its slot arrives.
-///
-/// Only enforced when explicitly set — `nil` (the default) means no rate limit
-/// and the activity is enqueued at full speed, exactly as without this option.
-///
-/// - `limit` / `period`: together define the slot interval.
-///   `.init(limit: 10, period: .seconds(1))` → one slot every 100 ms.
-///   `.init(limit: 2, period: .minutes(1))` → one slot every 30 s.
-///
-/// - `key`: optional partition key that gives each entity its own independent
-///   bucket.  `nil` = one global bucket shared by all executions of this
-///   activity type on this queue.  Use the workflow's own computed value —
-///   there is no server-side expression language; Swift code is the expression.
-///
-/// ```swift
-/// // Global: 10 payment activities per second across all customers
-/// options: .init(rateLimit: .init(limit: 10, period: .seconds(1)))
-///
-/// // Per-customer: 2 enrichment calls per minute per customer ID
-/// options: .init(
-///     rateLimit: .init(limit: 2, period: .minutes(1), key: "customer:\(customerID)")
-/// )
-/// ```
-public struct RateLimit: Sendable {
-    /// Maximum executions allowed within `period`.
-    public var limit: Double
-    /// Time window over which `limit` applies. Default: `.seconds(1)`.
-    public var period: Duration
-    /// Optional partition key.  `nil` = one global bucket for this activity
-    /// type on this queue.  A non-nil value gives each entity its own bucket.
-    public var key: String?
-
-    /// When set, this string is used verbatim as the slot key — no activity-name
-    /// prefix is applied. Use this to share a single rate-limit bucket across
-    /// multiple activity types (e.g. all HHAX activities sharing "hhax").
-    /// When both `sharedKey` and `key` are set, `sharedKey` wins.
-    public var sharedKey: String?
-
-    /// Number of tasks that may be dispatched immediately from a cold bucket (or
-    /// from any bucket that still has burst tokens remaining) before the regular
-    /// per-interval drip begins.  Default `0` = no burst.
-    public var burst: Int
-
-    /// Returns `nil` when `limit` is not positive; no crash.
-    public init?(limit: Double, period: Duration = .seconds(1), key: String? = nil, sharedKey: String? = nil, burst: Int = 0) {
-        guard limit > 0 else { return nil }
-        self.limit = limit
-        self.period = period
-        self.key = key
-        self.sharedKey = sharedKey
-        self.burst = burst
-    }
-
-    // MARK: - Internal
-
-    /// Slot-allocation parameters for the `strand.rate_limit_slots` SQL.
-    ///
-    /// - Parameter activityName: the activity's registered name.
-    ///   Used as the bucket prefix when `key` is `nil` (global bucket),
-    ///   or as the prefix in `"ActivityName:entityKey"` (per-entity bucket).
-    /// - Returns: `slotKey` — the row key in `strand.rate_limit_slots`;
-    ///   `intervalMs` — milliseconds between consecutive slots,
-    ///   derived from `period / limit` and clamped to at least 1 ms;
-    ///   `burstSlots` — number of tasks eligible for immediate dispatch.
-    internal func slotParams(for activityName: String) -> (slotKey: String, intervalMs: Int, burstSlots: Int) {
-        let slotKey = sharedKey ?? key.map { "\(activityName):\($0)" } ?? activityName
-        let intervalMs = max(1, Int(Double(period.milliseconds) / limit))
-        return (slotKey: slotKey, intervalMs: intervalMs, burstSlots: max(0, burst))
-    }
-}
-
 // MARK: - ActivityOptions
 
 /// Per-activity execution configuration: timeout, retry, routing, priority, and fairness.

Sources/Strand/Client.swift

@@ -104,7 +104,8 @@ public struct StrandClient: Sendable {
         input: A.Input,
         options: ActivityOptions = .init()
     ) async throws -> EnqueueResult {
-        let rlParams = options.rateLimit.map { $0.slotParams(for: A.name) }
+        // Resolve any dynamic key extractor while the typed input is still available.
+        let rlParams = options.rateLimit.map { $0.resolving(input: input).slotParams(for: A.name) }
         let resolvedID = options.id ?? A.generateActivityID()
         return try await _enqueue(
             queue: options.queue ?? queueName,

Sources/Strand/RateLimit.swift

@@ -0,0 +1,215 @@
+#if canImport(FoundationEssentials)
+import FoundationEssentials
+#else
+import Foundation
+#endif
+
+// MARK: - RateLimit
+
+/// Leaky-bucket rate limit applied when enqueuing an activity, workflow, or
+/// generic task.  When set, each enqueue atomically claims a time slot; the
+/// run waits in PENDING until its slot arrives.
+///
+/// Only enforced when explicitly set — `nil` (the default) means no rate limit
+/// and the task is enqueued at full speed, exactly as without this option.
+///
+/// - `limit` / `period`: together define the slot interval.
+///   `.init(limit: 10, period: .seconds(1))` → one slot every 100 ms.
+///   `.init(limit: 2, period: .minutes(1))` → one slot every 30 s.
+///
+/// - `key`: static partition key — one bucket per unique value.
+///
+/// ```swift
+/// // Global: 10 payment activities per second across all customers
+/// options: .init(rateLimit: .init(limit: 10, period: .seconds(1)))
+///
+/// // Per-customer static key — computed at the call site
+/// options: .init(
+///     rateLimit: .init(limit: 2, period: .minutes(1), key: "customer:\(customerID)")
+/// )
+/// ```
+public struct RateLimit: Sendable {
+    /// Maximum executions allowed within `period`.
+    public var limit: Double
+    /// Time window over which `limit` applies. Default: `.seconds(1)`.
+    public var period: Duration
+    /// Optional partition key.  `nil` = one global bucket for this task
+    /// type on this queue.  A non-nil value gives each entity its own bucket.
+    public var key: String?
+
+    /// When set, this string is used as the slot-key prefix — no task-name
+    /// prefix is applied.  Use this to share a rate-limit bucket across multiple
+    /// task types (e.g. all payment tasks sharing `"payments"`).
+    ///
+    /// - When only `sharedKey` is set the slot key is `sharedKey` verbatim
+    ///   (one global bucket shared by all matching tasks).
+    /// - When BOTH `sharedKey` and `key` are set the slot key is
+    ///   `"sharedKey:key"` — per-entity buckets within the shared namespace.
+    ///   Example: `sharedKey: "payments", key: "tenant-acme"` → `"payments:tenant-acme"`.
+    public var sharedKey: String?
+
+    /// Number of tasks that may be dispatched immediately from a cold bucket (or
+    /// from any bucket that still has burst tokens remaining) before the regular
+    /// per-interval drip begins.  Default `0` = no burst.
+    public var burst: Int
+
+    /// Number of rate-limit tokens this dispatch consumes from the shared bucket.
+    /// Default `1` = one dispatch charges one slot.  Set to the number of
+    /// external API calls made inside the activity so the bucket stays
+    /// calibrated to actual outbound call volume rather than dispatch count.
+    ///
+    /// Example: a batch activity that calls a third-party API 50 times per run
+    /// should set `units: 50` so the cursor advances 50× further on each
+    /// dispatch, correctly throttling peer tasks on the same bucket.
+    public var units: Int
+
+    /// Derives the per-entity partition key from the task's typed input at
+    /// enqueue time.  Set via the ``keyed(_:period:sharedKey:burst:units:by:)``
+    /// factory, which wraps the `any Sendable` cast in a typed closure and
+    /// adds a debug assertion for type mismatches.
+    ///
+    /// `internal` — callers should use `keyed(...)`, not set this directly.
+    ///
+    /// `nonisolated(unsafe)` because the closure is called exactly once,
+    /// synchronously, inside `resolving(input:)` before anything crosses an
+    /// isolation boundary.  The closure must be a pure function; `keyed`
+    /// enforces this by construction (KeyPath reads and field accesses).
+    nonisolated(unsafe) var keyExtractor: ((any Sendable) -> String?)?
+
+    /// Returns `nil` when `limit` is not positive; no crash.
+    public init?(limit: Double, period: Duration = .seconds(1), key: String? = nil, sharedKey: String? = nil, burst: Int = 0, units: Int = 1) {
+        guard limit > 0 else { return nil }
+        self.limit = limit
+        self.period = period
+        self.key = key
+        self.sharedKey = sharedKey
+        self.burst = burst
+        self.units = max(1, units)
+        self.keyExtractor = nil
+    }
+
+    // MARK: - Dynamic key factories
+
+    /// Creates a rate limit whose per-entity key is derived from the task's
+    /// input at enqueue time via a typed closure.
+    ///
+    /// The closure is evaluated while `Input` is still a concrete Swift
+    /// value — before it is JSON-encoded.  This lets you declare the key schema
+    /// once rather than repeating `key: input.tenantID` at every call site.
+    ///
+    /// ```swift
+    /// // Declare once — Input type is inferred from the closure:
+    /// static let rateLimitOpts = ActivityOptions(
+    ///     rateLimit: .keyed(limit: 100, period: .minutes(1)) {
+    ///         (input: ProcessPaymentInput) in input.customerID
+    ///     }
+    /// )
+    ///
+    /// // Every call site just passes the shared options:
+    /// try await context.runActivity(ProcessPayment.self, input: input, options: .rateLimitOpts)
+    /// ```
+    static func keyed<Input: Sendable>(
+        limit: Double,
+        period: Duration = .seconds(1),
+        sharedKey: String? = nil,
+        burst: Int = 0,
+        units: Int = 1,
+        by keyExtractor: @escaping (Input) -> String?
+    ) -> RateLimit? {
+        // Optional.map on the failable init handles the limit <= 0 guard;
+        // the transform sets keyExtractor on the valid copy and returns it.
+        RateLimit(limit: limit, period: period, sharedKey: sharedKey, burst: burst, units: units).map { base in
+            var rl = base
+            rl.keyExtractor = { anyInput in
+                guard let typed = anyInput as? Input else {
+                    assertionFailure(
+                        "RateLimit keyExtractor: expected \(Input.self), got \(type(of: anyInput)). "
+                            + "Ensure the RateLimit is used with the matching task's Input type."
+                    )
+                    return nil
+                }
+                return keyExtractor(typed)
+            }
+            return rl
+        }
+    }
+
+    /// KeyPath overload — derives the per-entity key from a non-optional property.
+    ///
+    /// ```swift
+    /// rateLimit: .keyed(limit: 100, period: .minutes(1), by: \ProcessPaymentInput.customerID)
+    /// ```
+    static func keyed<Input: Sendable>(
+        limit: Double,
+        period: Duration = .seconds(1),
+        sharedKey: String? = nil,
+        burst: Int = 0,
+        units: Int = 1,
+        by keyPath: KeyPath<Input, String>
+    ) -> RateLimit? {
+        keyed(limit: limit, period: period, sharedKey: sharedKey, burst: burst, units: units) {
+            (input: Input) in input[keyPath: keyPath]
+        }
+    }
+
+    /// KeyPath overload — derives the per-entity key from an optional property.
+    /// Returns `nil` (global bucket) when the property itself is `nil`.
+    ///
+    /// ```swift
+    /// rateLimit: .keyed(limit: 100, period: .minutes(1), by: \SyncVisitsPageInput.officeID)
+    /// ```
+    static func keyed<Input: Sendable>(
+        limit: Double,
+        period: Duration = .seconds(1),
+        sharedKey: String? = nil,
+        burst: Int = 0,
+        units: Int = 1,
+        by keyPath: KeyPath<Input, String?>
+    ) -> RateLimit? {
+        keyed(limit: limit, period: period, sharedKey: sharedKey, burst: burst, units: units) {
+            (input: Input) in input[keyPath: keyPath]
+        }
+    }
+
+    // MARK: - Internal
+
+    /// Slot-allocation parameters for the `strand.rate_limit_slots` SQL.
+    ///
+    /// - Parameter taskName: the task's registered name (activity or workflow name).
+    ///   Used as the bucket prefix when `key` is `nil` (global bucket),
+    ///   or as the prefix in `"TaskName:entityKey"` (per-entity bucket).
+    /// - Returns: `slotKey` — the row key in `strand.rate_limit_slots`;
+    ///   `intervalMs` — milliseconds between consecutive slots,
+    ///   derived from `(period / limit) × units` and clamped to at least 1 ms;
+    ///   `burstSlots` — number of tasks eligible for immediate dispatch.
+    ///
+    /// - Note: Uses `self.key` directly.  When a `keyExtractor` is present,
+    ///   call sites must resolve it first via `resolving(input:)` so the
+    ///   per-entity key is written into `self.key` before this is invoked.
+    internal func slotParams(for taskName: String) -> (slotKey: String, intervalMs: Int, burstSlots: Int) {
+        let slotKey: String
+        if let sk = sharedKey {
+            // sharedKey as namespace: append per-entity key when present.
+            slotKey = key.map { "\(sk):\($0)" } ?? sk
+        } else {
+            slotKey = key.map { "\(taskName):\($0)" } ?? taskName
+        }
+        let intervalMs = max(1, Int(Double(period.milliseconds) / limit)) * units
+        return (slotKey: slotKey, intervalMs: intervalMs, burstSlots: max(0, burst))
+    }
+
+    /// Resolves `key` in-place from `input` via `keyExtractor`.
+    /// No-op when no extractor is set.
+    internal mutating func resolve(input: any Sendable) {
+        guard let extractor = keyExtractor else { return }
+        key = extractor(input)
+    }
+
+    /// Returns a copy with `key` resolved from `input`.
+    /// Used where a mutating call isn’t possible (e.g. inside `Optional.map`).
+    internal func resolving(input: any Sendable) -> RateLimit {
+        var copy = self
+        copy.resolve(input: input)
+        return copy
+    }
+}

Sources/Strand/WorkflowContext.swift

@@ -902,6 +902,10 @@ public struct WorkflowContext<W: Workflow>: Sendable {
         // throw CancellationError here rather than scheduling a new activity that would
         // never be needed. Fast paths above are unaffected — already-completed work replays.
         try Task.checkCancellation()
+        // Resolve any dynamic rate-limit key while the typed input is still available.
+        // WorkflowScheduling only has a ByteBuffer at that point, so this is the
+        // last moment a KeyPath / closure extractor can be evaluated.
+        let resolvedOptions = options.resolvingRateLimit(input)
         let inputBuffer = try _impl.codec.encode(input)
         // Idempotency key is always derived from the run UUID and sequence number.
         // Run UUID (not task UUID) scopes the key to this run so that continueAsNew
@@ -921,7 +925,7 @@ public struct WorkflowContext<W: Workflow>: Sendable {
             .scheduleActivity(
                 name: A.name,
                 input: inputBuffer,
-                options: options,
+                options: resolvedOptions,
                 seqNum: seqNum,
                 idempotencyKey: idempotencyKey
             )
@@ -1869,3 +1873,21 @@ public struct ChildWorkflowHandle<CW: Workflow>: Sendable {
         try await _result()
     }
 }
+
+// MARK: - fileprivate helpers
+
+extension ActivityOptions {
+    /// Returns a copy of these options with the rate-limit key resolved from
+    /// `input` via `RateLimit.keyExtractor`, or `self` unchanged when no
+    /// extractor is set.
+    ///
+    /// Only needed in `WorkflowContext.runActivity`, which must pass a fully
+    /// resolved `ActivityOptions` into the `WorkflowCommand` — the typed input
+    /// is unavailable once it has been encoded to a `ByteBuffer`.
+    fileprivate func resolvingRateLimit(_ input: any Sendable) -> ActivityOptions {
+        guard rateLimit?.keyExtractor != nil else { return self }
+        var copy = self
+        copy.rateLimit?.resolve(input: input)
+        return copy
+    }
+}