feat: delegated auth support — token forwarding and principal-based resolver

lib/ai.rb

@@ -47,7 +47,7 @@ class Error < StandardError
   LogProbs = T.type_alias { T.anything }
   ProviderMetadata = T.type_alias { T.anything }
 
-  config_accessor :origin, :client, :api_key
+  config_accessor :origin, :client, :api_key, :delegated_token_resolver
 
   sig { params(content: String).returns(Ai::Message) }
   def self.user_message(content)

lib/ai/agent.rb

@@ -22,15 +22,17 @@ def initialize(agent_name:, client: Ai.client)
         runtime_context: T::Hash[String, T.anything],
         max_retries: Integer,
         max_steps: Integer,
-        telemetry: Ai::TelemetrySettings
+        telemetry: Ai::TelemetrySettings,
+        delegated_auth: T.nilable(Object)
       ).returns(Ai::GenerateTextResult)
     end
     def generate_text(
       messages:,
       runtime_context: {},
       max_retries: 2,
       max_steps: 5,
-      telemetry: Ai::TelemetrySettings.new
+      telemetry: Ai::TelemetrySettings.new,
+      delegated_auth: nil
     )
       options = {
         runtime_context: runtime_context,
@@ -39,7 +41,8 @@ def generate_text(
         telemetry: telemetry
       }
 
-      data = client.generate(agent_name, messages: messages, options: options)
+      delegated_token = resolve_delegated_token(delegated_auth)
+      data = client.generate(agent_name, messages: messages, options: options, delegated_token: delegated_token)
       TypeCoerce[Ai::GenerateTextResult].new.from(data, raise_coercion_error: false)
     end
 
@@ -51,7 +54,8 @@ def generate_text(
           runtime_context: T::Hash[String, T.anything],
           max_retries: Integer,
           max_steps: Integer,
-          telemetry: Ai::TelemetrySettings
+          telemetry: Ai::TelemetrySettings,
+          delegated_auth: T.nilable(Object)
         )
         .returns(GenerateObjectResult[T.type_parameter(:O)])
     end
@@ -61,7 +65,8 @@ def generate_object(
       runtime_context: {},
       max_retries: 2,
       max_steps: 5,
-      telemetry: Ai::TelemetrySettings.new
+      telemetry: Ai::TelemetrySettings.new,
+      delegated_auth: nil
     )
       schema = Ai::StructToJsonSchema.convert(T.cast(output_class, T.class_of(T::Struct)))
 
@@ -75,13 +80,26 @@ def generate_object(
         telemetry: telemetry
       }
 
-      data = client.generate(agent_name, messages: messages, options: options)
+      delegated_token = resolve_delegated_token(delegated_auth)
+      data = client.generate(agent_name, messages: messages, options: options, delegated_token: delegated_token)
 
       object = TypeCoerce[output_class].from(data['object'])
       TypeCoerce[GenerateObjectResult]
         .new
         .from(data, raise_coercion_error: false)
         .with(object: object)
     end
+
+    private
+
+    sig { params(delegated_auth: T.nilable(Object)).returns(T.nilable(String)) }
+    def resolve_delegated_token(delegated_auth)
+      return nil if delegated_auth.nil?
+
+      resolver = Ai.config.delegated_token_resolver
+      raise Ai::Error, 'delegated_token_resolver is not configured. Set Ai.delegated_token_resolver in your initializer.' if resolver.nil?
+
+      resolver.call(delegated_auth)
+    end
   end
 end

lib/ai/client.rb

@@ -32,15 +32,16 @@ def agent_names
         .params(
           agent_name: String,
           messages: T::Array[Ai::Message],
-          options: T::Hash[Symbol, T.anything]
+          options: T::Hash[Symbol, T.anything],
+          delegated_token: T.nilable(String)
         )
         .returns(T::Hash[String, T.anything])
     end
-    def generate(agent_name, messages:, options: {})
+    def generate(agent_name, messages:, options: {}, delegated_token: nil)
     end
 
-    sig { abstract.params(workflow_name: String, input: T::Struct).returns(ApiResponse) }
-    def run_workflow(workflow_name, input:)
+    sig { abstract.params(workflow_name: String, input: T::Struct, delegated_token: T.nilable(String)).returns(ApiResponse) }
+    def run_workflow(workflow_name, input:, delegated_token: nil)
     end
 
     sig { abstract.params(workflow_name: String).returns(SchemaHash) }

lib/ai/clients/mastra.rb

@@ -60,13 +60,14 @@ def agent_names
           .params(
             agent_name: String,
             messages: T::Array[Ai::Message],
-            options: T::Hash[Symbol, T.anything]
+            options: T::Hash[Symbol, T.anything],
+            delegated_token: T.nilable(String)
           )
           .returns(T::Hash[String, T.anything])
       end
-      def generate(agent_name, messages:, options: {})
+      def generate(agent_name, messages:, options: {}, delegated_token: nil)
         url = URI.join(@base_uri, "api/agents/#{agent_name}/generate")
-        generated_response = response(url: url, messages: messages, options: options)
+        generated_response = response(url: url, messages: messages, options: options, delegated_token: delegated_token)
 
         parsed_response =
           JSON.parse(generated_response.body || '').deep_transform_keys(&:underscore)
@@ -83,15 +84,15 @@ def generate(agent_name, messages:, options: {})
       end
 
       sig do
-        override.params(workflow_name: String, input: T::Struct).returns(Ai::Client::ApiResponse)
+        override.params(workflow_name: String, input: T::Struct, delegated_token: T.nilable(String)).returns(Ai::Client::ApiResponse)
       end
-      def run_workflow(workflow_name, input:)
+      def run_workflow(workflow_name, input:, delegated_token: nil)
         run_id = SecureRandom.uuid
 
         # Step 1: Create a new run for the workflow
         create_url =
           URI.join(@base_uri, "api/workflows/#{workflow_name}/create-run?runId=#{run_id}")
-        create_response = http_post(create_url, body: '{}')
+        create_response = http_post(create_url, body: '{}', delegated_token: delegated_token)
 
         unless create_response.is_a?(Net::HTTPSuccess)
           raise Ai::Error, "Mastra error – could not create workflow run: #{create_response.body}"
@@ -104,7 +105,7 @@ def run_workflow(workflow_name, input:)
         stream_error_body = T.let(nil, T.nilable(String))
         stream_body_chunks = T.let([], T::Array[String])
         stream_response =
-          http_post(stream_url, body: stream_request_body, stream: true) do |response|
+          http_post(stream_url, body: stream_request_body, stream: true, delegated_token: delegated_token) do |response|
             if response.is_a?(Net::HTTPSuccess)
               response.read_body do |chunk|
                 # Capture the stream body to check for workflow failures
@@ -144,6 +145,7 @@ def run_workflow(workflow_name, input:)
             .config
             .api_key
             .present?
+          result_request['X-Factorial-Delegated-Bearer'] = delegated_token if delegated_token.present?
           http = build_http
           result_response = http.request(result_request)
 
@@ -244,13 +246,15 @@ def deep_camelize_keys(options)
           url: URI::Generic,
           body: T.nilable(String),
           stream: T::Boolean,
+          delegated_token: T.nilable(String),
           blk: T.nilable(T.proc.params(response: Net::HTTPResponse).void)
         ).returns(Net::HTTPResponse)
       end
-      def http_post(url, body: nil, stream: false, &blk)
+      def http_post(url, body: nil, stream: false, delegated_token: nil, &blk)
         request = Net::HTTP::Post.new(url)
         request['Origin'] = Ai.config.origin
         request['Authorization'] = "Bearer #{Ai.config.api_key}" if Ai.config.api_key.present?
+        request['X-Factorial-Delegated-Bearer'] = delegated_token if delegated_token.present?
         if body
           request['Content-Type'] = 'application/json'
           request.body = body
@@ -282,14 +286,16 @@ def http_post(url, body: nil, stream: false, &blk)
         params(
           url: URI::Generic,
           messages: T::Array[Ai::Message],
-          options: T::Hash[Symbol, T.anything]
+          options: T::Hash[Symbol, T.anything],
+          delegated_token: T.nilable(String)
         ).returns(Net::HTTPResponse)
       end
-      def response(url:, messages:, options:)
+      def response(url:, messages:, options:, delegated_token: nil)
         request = Net::HTTP::Post.new(url)
         request['Content-Type'] = 'application/json'
         request['Origin'] = Ai.config.origin
         request['Authorization'] = "Bearer #{Ai.config.api_key}" if Ai.config.api_key.present?
+        request['X-Factorial-Delegated-Bearer'] = delegated_token if delegated_token.present?
 
         # convert to camelCase and unpacking for API compatibility
         camelized_options = deep_camelize_keys(options)

lib/ai/clients/test.rb

@@ -30,11 +30,12 @@ def set_returned_object(output) # rubocop:disable Naming/AccessorMethodName
           .params(
             agent_name: String,
             messages: T::Array[Ai::Message],
-            options: T::Hash[Symbol, T.anything]
+            options: T::Hash[Symbol, T.anything],
+            delegated_token: T.nilable(String)
           )
           .returns(T::Hash[String, T.anything])
       end
-      def generate(agent_name, messages:, options: {})
+      def generate(agent_name, messages:, options: {}, delegated_token: nil)
         output = options[:structured_output]
 
         # Use the first message content for testing purposes
@@ -101,9 +102,9 @@ def generate(agent_name, messages:, options: {})
       end
 
       sig do
-        override.params(workflow_name: String, input: T::Struct).returns(Ai::Client::ApiResponse)
+        override.params(workflow_name: String, input: T::Struct, delegated_token: T.nilable(String)).returns(Ai::Client::ApiResponse)
       end
-      def run_workflow(workflow_name, input:)
+      def run_workflow(workflow_name, input:, delegated_token: nil)
         @returned_object
       end
 

spec/lib/ai/agent_spec.rb

@@ -28,7 +28,8 @@
           max_retries: 2,
           max_steps: 5,
           telemetry: anything
-        }
+        },
+        delegated_token: nil
       ).and_call_original
 
       agent.generate_text(messages: [Ai.user_message('Hello')], runtime_context: runtime_context)
@@ -44,7 +45,8 @@
           max_retries: 5,
           max_steps: 10,
           telemetry: anything
-        }
+        },
+        delegated_token: nil
       ).and_call_original
 
       agent.generate_text(messages: [Ai.user_message('Hello')], max_retries: 5, max_steps: 10)
@@ -71,7 +73,8 @@
           max_retries: 2,
           max_steps: 5,
           telemetry: telemetry_settings
-        }
+        },
+        delegated_token: nil
       ).and_call_original
 
       agent.generate_text(messages: [Ai.user_message('Hello')], telemetry: telemetry_settings)
@@ -87,7 +90,8 @@
           max_retries: 2,
           max_steps: 5,
           telemetry: kind_of(Ai::TelemetrySettings)
-        }
+        },
+        delegated_token: nil
       ).and_call_original
 
       result = agent.generate_text(messages: [Ai.user_message('Hello')])
@@ -111,6 +115,45 @@
       expect(result.tool_results).to eq([])
       expect(result.steps).to eq([])
     end
+
+    describe 'delegated_auth' do
+      after { Ai.delegated_token_resolver = nil }
+
+      it 'resolves delegated_auth to a token via the configured resolver and passes it to the client' do
+        principal = instance_double(Object)
+        Ai.delegated_token_resolver = ->(p) { "resolved-token-for-#{p.class.name}" }
+
+        expect(client).to receive(:generate).with(
+          'test',
+          messages: anything,
+          options: anything,
+          delegated_token: "resolved-token-for-#{principal.class.name}"
+        ).and_call_original
+
+        agent.generate_text(messages: [Ai.user_message('Hello')], delegated_auth: principal)
+      end
+
+      it 'passes nil delegated_token when delegated_auth is nil' do
+        Ai.delegated_token_resolver = ->(_p) { raise 'should not be called' }
+
+        expect(client).to receive(:generate).with(
+          'test',
+          messages: anything,
+          options: anything,
+          delegated_token: nil
+        ).and_call_original
+
+        agent.generate_text(messages: [Ai.user_message('Hello')])
+      end
+
+      it 'raises Ai::Error when resolver is not configured but delegated_auth is provided' do
+        Ai.delegated_token_resolver = nil
+
+        expect do
+          agent.generate_text(messages: [Ai.user_message('Hello')], delegated_auth: instance_double(Object))
+        end.to raise_error(Ai::Error, /delegated_token_resolver is not configured/)
+      end
+    end
   end
 
   describe '#generate_object' do
@@ -153,7 +196,8 @@
             max_steps: 8,
             structured_output: hash_including(schema: anything),
             telemetry: anything
-          )
+          ),
+        delegated_token: nil
       ).and_call_original
 
       agent.generate_object(
@@ -184,7 +228,8 @@
           hash_including(
             telemetry: telemetry_settings,
             structured_output: hash_including(schema: anything)
-          )
+          ),
+        delegated_token: nil
       ).and_call_original
 
       agent.generate_object(
@@ -217,5 +262,36 @@
         expect { subject }.to raise_error(ArgumentError)
       end
     end
+
+    describe 'delegated_auth' do
+      after { Ai.delegated_token_resolver = nil }
+
+      before { client.set_returned_object({ 'name' => 'John Doe', 'age' => 30 }) }
+
+      let(:schema) do
+        Class.new(T::Struct) do
+          const :name, String
+          const :age, Integer
+        end
+      end
+
+      it 'resolves delegated_auth and passes delegated_token to the client' do
+        principal = instance_double(Object)
+        Ai.delegated_token_resolver = ->(_p) { 'company-token-abc' }
+
+        expect(client).to receive(:generate).with(
+          'test',
+          messages: anything,
+          options: anything,
+          delegated_token: 'company-token-abc'
+        ).and_call_original
+
+        agent.generate_object(
+          messages: [Ai.user_message('Create person')],
+          output_class: schema,
+          delegated_auth: principal
+        )
+      end
+    end
   end
 end