hatchet-sdk 0.1.0.pre.alpha → 0.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 39266ae76d015db9cfae9ab50c45b45772a71a6e83196a8f50d696427df1188f
-  data.tar.gz: '0989914579fb062642a4ffb31242632a0bfc5fd4bc510fbfac1c8b91a256de74'
+  metadata.gz: '00528f759273d232b37cc2427e1afb1c70fdb95eeaf9ef8251353bfe3a0af224'
+  data.tar.gz: fa24d7738b08f554d9e8b7f0579a81fb866d968590f18549415bea5b9472ab1d
 SHA512:
-  metadata.gz: 610bc9c45ce11f4615d9858e66348c4fcae1f0af0f85be8763d6e5b43d6eba26b2ad3af7b521b6b058357b39cd4e386c9c13f3046043f0d056dcf86d52232387
-  data.tar.gz: 53f3136c6ff7f902df3e71b8626391ef02294c21d0333f24c6de2ae924e2d12a04754340b2b9b8f9b3d3046bc5fe438e917b8799b79e849259023d3b0cf2d14e
+  metadata.gz: 3501969e6906c6a9cd84c4120d6cd2eb3f63c3204df8d673ac04754ae2152bd25a806cf23f2d84e9f219aa8f9a9bf9c09b4050b9b6959163c1a00618370128be
+  data.tar.gz: f81aed6b15226d78bdff0b1176263514153da55e13c81778c18a598f6a6e4181ef7b4cca01fb2f63c08b462520d95465d9e3b83d53a3e58b2b00d98db8d44664

data/.rubocop.yml

@@ -1,8 +1,153 @@
+# frozen_string_literal: true
+
+require:
+  - ./lib/rubocop/cop/hatchet/rbs_signature_exists
+
 AllCops:
   TargetRubyVersion: 3.1
+  NewCops: enable
+  SuggestExtensions: false
+  Exclude:
+    - "lib/hatchet/clients/rest/**/*"
+    - "lib/hatchet/contracts/**/*"
+    - "vendor/**/*"
+    - "bin/**/*"
+    - "scripts/**/*"
+
+# ─── Style ───────────────────────────────────────────────────────────
 
 Style/StringLiterals:
   EnforcedStyle: double_quotes
 
 Style/StringLiteralsInInterpolation:
   EnforcedStyle: double_quotes
+
+Style/FrozenStringLiteralComment:
+  EnforcedStyle: always
+
+Style/Documentation:
+  Enabled: false
+
+Style/TrailingCommaInArguments:
+  EnforcedStyleForMultiline: consistent_comma
+
+Style/TrailingCommaInArrayLiteral:
+  EnforcedStyleForMultiline: consistent_comma
+
+Style/TrailingCommaInHashLiteral:
+  EnforcedStyleForMultiline: consistent_comma
+
+# ─── Metrics ─────────────────────────────────────────────────────────
+
+Metrics/MethodLength:
+  Max: 55
+  CountAsOne:
+    - array
+    - hash
+    - heredoc
+    - method_call
+
+Metrics/ClassLength:
+  Max: 300
+  CountAsOne:
+    - array
+    - hash
+    - heredoc
+    - method_call
+
+Metrics/ModuleLength:
+  Max: 300
+  CountAsOne:
+    - array
+    - hash
+    - heredoc
+    - method_call
+
+Metrics/BlockLength:
+  Max: 35
+  Exclude:
+    - "spec/**/*"
+    - "Rakefile"
+    - "*.gemspec"
+
+Metrics/AbcSize:
+  Max: 65
+
+Metrics/CyclomaticComplexity:
+  Max: 35
+
+Metrics/PerceivedComplexity:
+  Max: 35
+
+Metrics/ParameterLists:
+  Max: 10
+  CountKeywordArgs: false
+
+# ─── Layout ──────────────────────────────────────────────────────────
+
+Layout/LineLength:
+  Max: 140
+  AllowedPatterns:
+    - "^\\s*#"
+  Exclude:
+    - "*.gemspec"
+
+Layout/MultilineMethodCallIndentation:
+  EnforcedStyle: indented
+
+Layout/FirstHashElementIndentation:
+  EnforcedStyle: consistent
+
+Layout/FirstArrayElementIndentation:
+  EnforcedStyle: consistent
+
+# ─── Naming ──────────────────────────────────────────────────────────
+
+Naming/FileName:
+  Exclude:
+    - "lib/hatchet-sdk.rb"
+
+Naming/MethodParameterName:
+  AllowedNames:
+    - id
+    - ip
+    - to
+    - by
+    - on
+    - in
+    - at
+    - fn
+
+Naming/PredicatePrefix:
+  Enabled: false
+
+Naming/AccessorMethodName:
+  Enabled: false
+
+Naming/VariableNumber:
+  Enabled: false
+
+# ─── Gemspec ─────────────────────────────────────────────────────────
+
+Gemspec/DevelopmentDependencies:
+  Enabled: false
+
+# ─── Lint ────────────────────────────────────────────────────────────
+
+Lint/UnusedMethodArgument:
+  AllowUnusedKeywordArguments: true
+
+Lint/MissingSuper:
+  Enabled: false
+
+Lint/DuplicateBranch:
+  Enabled: false
+
+# ─── Custom ─────────────────────────────────────────────────────────
+
+Hatchet/RbsSignatureExists:
+  Enabled: true
+  Exclude:
+    - "lib/hatchet/clients/rest/**/*"
+    - "lib/hatchet/contracts/**/*"
+    - "lib/rubocop/**/*"

data/CHANGELOG.md

@@ -1,6 +1,27 @@
 ## Hatchet SDK Changelog
 
-## [0.1.0] - 2025-08-05
+## [0.1.0] - 2025-02-15
 
-- Initial release
-- Events and Runs Rest Clients
+- Initial release of the Ruby SDK for Hatchet
+- Task orchestration with simple tasks, DAGs, and child/fanout workflows
+- Durable execution with durable tasks, durable events, and durable sleep
+- Concurrency control (limit, round-robin, cancel in progress, cancel newest, multiple keys, workflow-level)
+- Rate limiting
+- Event-driven workflows
+- Cron and scheduled workflows
+- Retries with configurable backoff strategies
+- Timeout management with refresh support
+- On-failure and on-success callbacks
+- Streaming support
+- Webhook integration
+- Bulk operations (fanout, replay)
+- Priority scheduling
+- Sticky and affinity worker assignment
+- Deduplication
+- Manual slot release
+- Dependency injection
+- Unit testing helpers
+- Logging integration
+- Run detail inspection
+- RBS type signatures for IDE support
+- REST and gRPC client support

data/CLAUDE.md

@@ -33,7 +33,7 @@ The SDK includes comprehensive documentation and type signatures for excellent I
 
 **REST API Generation:**
 - `rake api:generate` - Generate REST API client from OpenAPI specification
-- `rake api:clean` - Remove generated REST API client files  
+- `rake api:clean` - Remove generated REST API client files
 - `rake api:regenerate` - Clean and regenerate REST API client
 - `rake api:validate` - Validate OpenAPI specification
 - `rake api:info` - Show OpenAPI specification information
@@ -57,8 +57,7 @@ The SDK includes comprehensive documentation and type signatures for excellent I
 - `Hatchet::Clients` - Factory for creating REST and other protocol clients
 - `Hatchet::Clients::Rest::*` - Generated REST API clients (WorkflowApi, EventApi, etc.)
 - `Hatchet::TLSConfig` - TLS configuration for secure connections
-- `Hatchet::HealthcheckConfig` - Worker health monitoring configuration  
-- `Hatchet::OpenTelemetryConfig` - Observability configuration
+- `Hatchet::HealthcheckConfig` - Worker health monitoring configuration
 - `Hatchet::Error` - Base error class for gem-specific exceptions
 
 **Configuration Sources (priority order):**
@@ -79,7 +78,7 @@ The codebase uses frozen string literals and follows Ruby 3.1+ requirements.
 
 **When adding new configuration options:**
 1. Add the parameter to `Config#initialize` method
-2. Update the `@option` YARD documentation in both `Client` and `Config` classes  
+2. Update the `@option` YARD documentation in both `Client` and `Config` classes
 3. Add the parameter to RBS type signatures in `sig/hatchet-sdk.rbs`
 4. Add comprehensive test coverage in `spec/hatchet/config_spec.rb`
 5. Update this CLAUDE.md file with any architectural changes
@@ -96,4 +95,4 @@ The codebase uses frozen string literals and follows Ruby 3.1+ requirements.
 - The config extracts tenant_id from the 'sub' field in JWT tokens
 - Test both explicit tenant_id override and JWT extraction scenarios
 
-Keep the CLAUDE.md instructions up to date as the project continues to develop.
+Keep the CLAUDE.md instructions up to date as the project continues to develop.

data/INTEGRATION_TESTING.md

@@ -63,7 +63,7 @@ The integration tests for the Runs feature test:
 
 #### ✅ Workflow Run Operations
 - Get workflow run details
-- Get workflow run status  
+- Get workflow run status
 - Get workflow run results
 - Data structure validation
 
@@ -173,4 +173,4 @@ bundle exec rspec --tag ~integration
 
 # Run integration tests with credentials (slower, requires setup)
 HATCHET_CLIENT_TOKEN=$SECRET_TOKEN bundle exec rspec --tag integration
-```
+```

data/README.md

@@ -71,4 +71,4 @@ We welcome contributions! Please check out our [contributing guidelines](https:/
 
 ## License
 
-This SDK is released under the MIT License. See [LICENSE](https://github.com/hatchet-dev/hatchet/blob/main/LICENSE) for details.
+This SDK is released under the MIT License. See [LICENSE](https://github.com/hatchet-dev/hatchet/blob/main/LICENSE) for details.

data/REST_API_GENERATION.md

@@ -254,4 +254,4 @@ The REST API generation setup requires minimal maintenance:
 2. **Dependency Updates** - Keep OpenAPI Generator CLI updated
 3. **Custom Patches** - Update patches when generation patterns change
 
-The generation process is designed to be idempotent and safe to run multiple times.
+The generation process is designed to be idempotent and safe to run multiple times.

data/Rakefile

@@ -13,9 +13,9 @@ RuboCop::RakeTask.new
 load File.expand_path("Rakefile.rest", __dir__)
 
 # Generate REST API client before building the gem
-task :build => "api:generate"
+task build: "api:generate"
 
 # Generate REST API client before running tests (for CI)
-task :spec => "api:generate_if_missing"
+task spec: "api:generate_if_missing"
 
 task default: %i[spec rubocop]

data/Rakefile.rest

@@ -1,44 +1,33 @@
 # frozen_string_literal: true
 
 # Additional Rake tasks for REST API client generation
+#
+# All generation goes through sdks/ruby/generate.sh, which is the single
+# entry point for both protobuf and REST generation.
 
 require "rake"
 
 namespace :api do
+  # Path to the unified generate script
+  GENERATE_SH = File.expand_path("../../generate.sh", __dir__)
+
   desc "Generate REST API client from OpenAPI specification"
   task :generate do
-    puts "Generating REST API client from OpenAPI specification..."
-    
-    ruby_script = File.join(__dir__, "scripts", "generate_api.rb")
-    
-    unless File.exist?(ruby_script)
-      puts "❌ Generator script not found at #{ruby_script}"
-      exit 1
-    end
-    
-    # Make the script executable
-    File.chmod(0755, ruby_script)
-    
-    # Execute the generation script
-    system("ruby #{ruby_script}") || begin
-      puts "❌ API generation failed"
-      exit 1
-    end
-    
-    puts "✅ REST API client generation completed"
+    puts "Generating REST API client..."
+    system("bash", GENERATE_SH, "rest") || abort("REST API generation failed")
   end
 
   desc "Clean generated REST API client files"
   task :clean do
     rest_client_dir = File.join(__dir__, "lib", "hatchet", "clients", "rest")
-    
+
     if Dir.exist?(rest_client_dir)
-      puts "🧹 Cleaning generated REST API client files..."
+      puts "Cleaning generated REST API client files..."
       require "fileutils"
       FileUtils.rm_rf(rest_client_dir)
-      puts "✅ REST API client files cleaned"
+      puts "REST API client files cleaned"
     else
-      puts "ℹ️  No REST API client files to clean"
+      puts "No REST API client files to clean"
     end
   end
 
@@ -49,85 +38,12 @@ namespace :api do
   task :generate_if_missing do
     rest_client_dir = File.join(__dir__, "lib", "hatchet", "clients", "rest")
     main_client_file = File.join(rest_client_dir, "lib", "hatchet-sdk-rest.rb")
-    
+
     unless File.exist?(main_client_file)
-      puts "🔄 REST API client not found, generating..."
+      puts "REST API client not found, generating..."
       Rake::Task["api:generate"].invoke
     else
-      puts "✅ REST API client already exists, skipping generation"
-    end
-  end
-
-  desc "Validate OpenAPI specification"
-  task :validate do
-    openapi_spec = File.join(__dir__, "..", "..", "..", "bin", "oas", "openapi.yaml")
-    
-    unless File.exist?(openapi_spec)
-      puts "❌ OpenAPI specification not found at #{openapi_spec}"
-      exit 1
-    end
-    
-    puts "🔍 Validating OpenAPI specification..."
-    
-    # Use openapi-generator-cli to validate the spec
-    cmd = ["openapi-generator-cli", "validate", "-i", openapi_spec]
-    
-    if system(*cmd)
-      puts "✅ OpenAPI specification is valid"
-    else
-      puts "❌ OpenAPI specification validation failed"
-      exit 1
-    end
-  end
-
-  desc "Install OpenAPI Generator CLI if not present"
-  task :install_generator do
-    puts "🔧 Checking OpenAPI Generator CLI installation..."
-    
-    if system("which openapi-generator-cli > /dev/null 2>&1")
-      puts "✅ OpenAPI Generator CLI is already installed"
-    else
-      puts "📦 Installing OpenAPI Generator CLI..."
-      
-      if system("npm install -g @openapitools/openapi-generator-cli@7.13.0")
-        puts "✅ OpenAPI Generator CLI installed successfully"
-      else
-        puts "❌ Failed to install OpenAPI Generator CLI"
-        puts "Please install it manually: npm install -g @openapitools/openapi-generator-cli"
-        exit 1
-      end
-    end
-  end
-
-  desc "Show information about the OpenAPI specification"
-  task :info do
-    openapi_spec = File.join(__dir__, "..", "..", "..", "bin", "oas", "openapi.yaml")
-    
-    unless File.exist?(openapi_spec)
-      puts "❌ OpenAPI specification not found at #{openapi_spec}"
-      exit 1
-    end
-    
-    puts "📊 OpenAPI Specification Information:"
-    puts "  File: #{openapi_spec}"
-    puts "  Size: #{File.size(openapi_spec)} bytes"
-    puts "  Modified: #{File.mtime(openapi_spec)}"
-    
-    # Try to parse and show basic info
-    begin
-      require "yaml"
-      spec = YAML.load_file(openapi_spec)
-      
-      puts "  OpenAPI Version: #{spec.dig('openapi')}"
-      puts "  API Version: #{spec.dig('info', 'version')}"
-      puts "  API Title: #{spec.dig('info', 'title')}"
-      puts "  Paths: #{spec.dig('paths')&.keys&.count || 0}"
-      puts "  Components: #{spec.dig('components', 'schemas')&.keys&.count || 0} schemas"
-    rescue => e
-      puts "  ⚠️  Could not parse specification: #{e.message}"
+      puts "REST API client already exists, skipping generation"
     end
   end
 end
-
-# Add REST API tasks to the default rake task
-task default: ["api:validate"]

data/config/openapi_generator_config.json

@@ -28,4 +28,4 @@
   "useActiveModel": false,
   "hideGenerationTimestamp": true,
   "withDefaultValues": true
-}
+}

data/lib/hatchet/clients/grpc/admin.rb

@@ -0,0 +1,264 @@
+# frozen_string_literal: true
+
+require "json"
+require "google/protobuf/timestamp_pb"
+
+module Hatchet
+  module Clients
+    module Grpc
+      # gRPC client for the Hatchet Admin service (workflow registration & triggering).
+      #
+      # Uses two stubs on the shared channel:
+      # - V1::AdminService::Stub (v1) for: PutWorkflow, GetRunDetails, CancelTasks, ReplayTasks
+      # - WorkflowService::Stub (v0) for: TriggerWorkflow, BulkTriggerWorkflow, ScheduleWorkflow, PutRateLimit
+      #
+      # The v0 WorkflowService is retained for triggering because it supports
+      # parent-child linking fields that the v1 TriggerWorkflowRun does not expose.
+      #
+      # @example
+      #   admin = Admin.new(config: hatchet_config, channel: channel)
+      #   admin.put_workflow(workflow.to_proto(config))
+      #   ref = admin.trigger_workflow("MyWorkflow", input: { "key" => "value" })
+      class Admin
+        BULK_TRIGGER_BATCH_SIZE = 1000
+
+        # @param config [Hatchet::Config] The Hatchet configuration
+        # @param channel [GRPC::Core::Channel] Shared gRPC channel
+        def initialize(config:, channel:)
+          @config = config
+          @logger = config.logger
+          @channel = channel
+          @v0_stub = nil
+          @v1_stub = nil
+        end
+
+        # Register a workflow definition with the server via v1 AdminService.
+        #
+        # @param workflow_proto [V1::CreateWorkflowVersionRequest] The workflow proto
+        # @return [V1::CreateWorkflowVersionResponse] Registration response
+        def put_workflow(workflow_proto)
+          ensure_connected!
+
+          response = @v1_stub.put_workflow(workflow_proto, metadata: @config.auth_metadata)
+          @logger.debug("Registered workflow: #{workflow_proto.name}")
+          response
+        end
+
+        # Trigger a workflow run via v0 WorkflowService.
+        #
+        # @param workflow_name [String] The workflow name (will be namespaced)
+        # @param input [Hash] Workflow input
+        # @param options [Hash] Trigger options
+        # @option options [String] :parent_id Parent workflow run ID
+        # @option options [String] :parent_task_run_external_id Parent step run ID
+        # @option options [Integer] :child_index Child workflow index
+        # @option options [String] :child_key Child workflow key
+        # @option options [Hash] :additional_metadata Additional metadata
+        # @option options [String] :desired_worker_id Desired worker for sticky dispatch
+        # @option options [Integer] :priority Priority value
+        # @return [String] The workflow run ID
+        # @raise [DedupeViolationError] If a deduplication violation occurs
+        def trigger_workflow(workflow_name, input: {}, options: {})
+          ensure_connected!
+
+          name = @config.apply_namespace(workflow_name)
+
+          request_args = {
+            name: name,
+            input: JSON.generate(input),
+          }
+
+          request_args[:parent_id] = options[:parent_id] if options[:parent_id]
+          request_args[:parent_task_run_external_id] = options[:parent_task_run_external_id] if options[:parent_task_run_external_id]
+          request_args[:child_index] = options[:child_index] if options[:child_index]
+          request_args[:child_key] = options[:child_key] if options[:child_key]
+          request_args[:desired_worker_id] = options[:desired_worker_id] if options[:desired_worker_id]
+          request_args[:priority] = options[:priority] if options[:priority]
+
+          if options[:additional_metadata]
+            request_args[:additional_metadata] = if options[:additional_metadata].is_a?(String)
+                                                   options[:additional_metadata]
+                                                 else
+                                                   JSON.generate(options[:additional_metadata])
+                                                 end
+          end
+
+          request = ::TriggerWorkflowRequest.new(**request_args)
+
+          begin
+            response = @v0_stub.trigger_workflow(request, metadata: @config.auth_metadata)
+            response.workflow_run_id
+          rescue ::GRPC::AlreadyExists => e
+            raise DedupeViolationError, "Deduplication violation: #{e.message}"
+          rescue ::GRPC::ResourceExhausted => e
+            raise ResourceExhaustedError, e.message
+          rescue ::GRPC::BadStatus => e
+            raise Error, "gRPC error triggering workflow: #{e.class}: #{e.message}"
+          end
+        end
+
+        # Trigger multiple workflow runs in bulk via v0 WorkflowService.
+        #
+        # Automatically batches requests in groups of 1000.
+        #
+        # @param workflow_name [String] The workflow name (will be namespaced)
+        # @param items [Array<Hash>] Array of { input:, options: } items
+        # @return [Array<String>] Array of workflow run IDs
+        def bulk_trigger_workflow(workflow_name, items)
+          ensure_connected!
+
+          name = @config.apply_namespace(workflow_name)
+
+          requests = items.map do |item|
+            input = item[:input] || {}
+            opts = item[:options] || {}
+
+            request_args = {
+              name: name,
+              input: JSON.generate(input),
+            }
+
+            request_args[:parent_id] = opts[:parent_id] if opts[:parent_id]
+            request_args[:parent_task_run_external_id] = opts[:parent_task_run_external_id] if opts[:parent_task_run_external_id]
+            request_args[:child_index] = opts[:child_index] if opts[:child_index]
+            request_args[:child_key] = opts[:child_key] if opts[:child_key]
+            request_args[:desired_worker_id] = opts[:desired_worker_id] if opts[:desired_worker_id]
+            request_args[:priority] = opts[:priority] if opts[:priority]
+
+            if opts[:additional_metadata]
+              request_args[:additional_metadata] = if opts[:additional_metadata].is_a?(String)
+                                                     opts[:additional_metadata]
+                                                   else
+                                                     JSON.generate(opts[:additional_metadata])
+                                                   end
+            end
+
+            ::TriggerWorkflowRequest.new(**request_args)
+          end
+
+          # Batch in groups of BULK_TRIGGER_BATCH_SIZE
+          all_run_ids = []
+          begin
+            requests.each_slice(BULK_TRIGGER_BATCH_SIZE) do |batch|
+              bulk_request = ::BulkTriggerWorkflowRequest.new(workflows: batch)
+              response = @v0_stub.bulk_trigger_workflow(bulk_request, metadata: @config.auth_metadata)
+              all_run_ids.concat(response.workflow_run_ids.to_a)
+            end
+          rescue ::GRPC::ResourceExhausted => e
+            raise ResourceExhaustedError, e.message
+          rescue ::GRPC::BadStatus => e
+            raise Error, "gRPC error triggering bulk workflow: #{e.class}: #{e.message}"
+          end
+
+          all_run_ids
+        end
+
+        # Schedule a workflow for future execution via v0 WorkflowService.
+        #
+        # @param workflow_name [String] The workflow name (will be namespaced)
+        # @param run_at [Time] When to run
+        # @param input [Hash] Workflow input
+        # @param options [Hash] Trigger options
+        # @return [WorkflowVersion] Schedule response
+        # @raise [DedupeViolationError] If a deduplication violation occurs
+        def schedule_workflow(workflow_name, run_at:, input: {}, options: {})
+          ensure_connected!
+
+          name = @config.apply_namespace(workflow_name)
+
+          schedule_timestamp = Google::Protobuf::Timestamp.new(
+            seconds: run_at.to_i,
+            nanos: run_at.respond_to?(:nsec) ? run_at.nsec : 0,
+          )
+
+          request_args = {
+            name: name,
+            schedules: [schedule_timestamp],
+            input: JSON.generate(input),
+          }
+
+          request_args[:parent_id] = options[:parent_id] if options[:parent_id]
+          request_args[:parent_task_run_external_id] = options[:parent_task_run_external_id] if options[:parent_task_run_external_id]
+          request_args[:child_index] = options[:child_index] if options[:child_index]
+          request_args[:child_key] = options[:child_key] if options[:child_key]
+          request_args[:priority] = options[:priority] if options[:priority]
+
+          if options[:additional_metadata]
+            request_args[:additional_metadata] = if options[:additional_metadata].is_a?(String)
+                                                   options[:additional_metadata]
+                                                 else
+                                                   JSON.generate(options[:additional_metadata])
+                                                 end
+          end
+
+          request = ::ScheduleWorkflowRequest.new(**request_args)
+
+          begin
+            @v0_stub.schedule_workflow(request, metadata: @config.auth_metadata)
+          rescue ::GRPC::AlreadyExists => e
+            raise DedupeViolationError, "Deduplication violation: #{e.message}"
+          rescue ::GRPC::ResourceExhausted => e
+            raise ResourceExhaustedError, e.message
+          rescue ::GRPC::BadStatus => e
+            raise Error, "gRPC error scheduling workflow: #{e.class}: #{e.message}"
+          end
+        end
+
+        # Get run details via v1 AdminService.
+        #
+        # @param external_id [String] The workflow run external ID
+        # @return [V1::GetRunDetailsResponse]
+        def get_run_details(external_id:)
+          ensure_connected!
+
+          request = ::V1::GetRunDetailsRequest.new(external_id: external_id)
+          @v1_stub.get_run_details(request, metadata: @config.auth_metadata)
+        end
+
+        # Put a rate limit via v0 WorkflowService.
+        #
+        # @param key [String] Rate limit key
+        # @param limit [Integer] Rate limit value
+        # @param duration [Symbol] Rate limit duration enum
+        # @return [PutRateLimitResponse]
+        def put_rate_limit(key:, limit:, duration:)
+          ensure_connected!
+
+          request = ::PutRateLimitRequest.new(
+            key: key,
+            limit: limit,
+            duration: duration,
+          )
+
+          @v0_stub.put_rate_limit(request, metadata: @config.auth_metadata)
+        end
+
+        # Close the connection.
+        def close
+          @v0_stub = nil
+          @v1_stub = nil
+        end
+
+        private
+
+        def ensure_connected!
+          return if @v0_stub && @v1_stub
+
+          @v0_stub = ::WorkflowService::Stub.new(
+            @config.host_port,
+            nil,
+            channel_override: @channel,
+          )
+
+          @v1_stub = ::V1::AdminService::Stub.new(
+            @config.host_port,
+            nil,
+            channel_override: @channel,
+          )
+
+          @logger.debug("Admin gRPC stubs (v0 + v1) connected via shared channel")
+        end
+      end
+    end
+  end
+end