Class: Legate::ToolContext

Inherits:
Object
  • Object
show all
Includes:
Auth::ToolContextExtension
Defined in:
lib/legate/tool_context.rb

Overview

Provides contextual information to Legate::Tool#perform_execution Includes session details and a reference to the agent’s tool registry. Read-only.

Instance Attribute Summary collapse

Instance Method Summary collapse

Methods included from Auth::ToolContextExtension

#auth_runner, #auth_session, #cancel_auth_flow, #handle_auth_response, #with_authentication

Constructor Details

#initialize(session_id:, user_id:, app_name:, tool_registry: nil, session_service: nil, logger: Legate.logger, invocation_id: nil, agent_auth_config: nil) ⇒ ToolContext

Returns a new instance of ToolContext.

Parameters:

  • session_id (String) —

    The ID of the current session.

  • user_id (String) —

    The user ID associated with the session.

  • app_name (String) —

    The application/agent name associated with the session.

  • tool_registry (Legate::ToolRegistry) (defaults to: nil) —

    The tool registry instance of the agent executing the tool.

  • session_service (Legate::SessionService::Base, nil) (defaults to: nil) —

    The session service instance.

  • logger (Logger, nil) (defaults to: Legate.logger) —

    The logger instance.

  • invocation_id (String, nil) (defaults to: nil) —

    The ID of the current agent invocation.

  • agent_auth_config (Hash, nil) (defaults to: nil) —

    Agent-specific authentication configuration.



26
27
28
29
30
31
32
33
34
35
36
 
# File 'lib/legate/tool_context.rb', line 26

def initialize(session_id:, user_id:, app_name:, tool_registry: nil, session_service: nil, logger: Legate.logger, invocation_id: nil, agent_auth_config: nil)
  @session_id = session_id
  @user_id = user_id
  @app_name = app_name
  @tool_registry = tool_registry
  @session_service = session_service
  @invocation_id = invocation_id
  @pending_state_delta = {}
  @token_manager = nil
  @agent_auth_config = agent_auth_config
end

Instance Attribute Details

#agent_auth_config ⇒ Hash? (readonly)

Agent-specific authentication configuration

Returns:

  • (Hash, nil) —

    The agent’s auth config or nil if not configured



16
17
18
 
# File 'lib/legate/tool_context.rb', line 16

def agent_auth_config
  @agent_auth_config
end

#app_name ⇒ Object (readonly)

Returns the value of attribute app_name.



9
10
11
 
# File 'lib/legate/tool_context.rb', line 9

def app_name
  @app_name
end

#invocation_id ⇒ Object (readonly)

Returns the value of attribute invocation_id.



9
10
11
 
# File 'lib/legate/tool_context.rb', line 9

def invocation_id
  @invocation_id
end

#logger ⇒ Object (readonly)

Returns the value of attribute logger.



9
10
11
 
# File 'lib/legate/tool_context.rb', line 9

def logger
  @logger
end

#pending_state_delta ⇒ Object (readonly)

Expose pending state delta for inspection but not direct modification



12
13
14
 
# File 'lib/legate/tool_context.rb', line 12

def pending_state_delta
  @pending_state_delta
end

#session_id ⇒ Object (readonly)

Returns the value of attribute session_id.



9
10
11
 
# File 'lib/legate/tool_context.rb', line 9

def session_id
  @session_id
end

#session_service ⇒ Object (readonly)

Returns the value of attribute session_service.



9
10
11
 
# File 'lib/legate/tool_context.rb', line 9

def session_service
  @session_service
end

#tool_registry ⇒ Object (readonly)

Returns the value of attribute tool_registry.



9
10
11
 
# File 'lib/legate/tool_context.rb', line 9

def tool_registry
  @tool_registry
end

#user_id ⇒ Object (readonly)

Returns the value of attribute user_id.



9
10
11
 
# File 'lib/legate/tool_context.rb', line 9

def user_id
  @user_id
end

Instance Method Details

#authenticate_request(request, scheme, credential) ⇒ Hash

Apply authentication to a request using the specified scheme and credential

Parameters:

Returns:

  • (Hash) —

    The authenticated request



86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
 
# File 'lib/legate/tool_context.rb', line 86

def authenticate_request(request, scheme, credential)
  require_relative 'auth/tool_integration'

  # Try to use token manager if available
  token_manager = get_token_manager
  if token_manager
    # First get the token from token manager
    token = token_manager.get_token(scheme, credential, force_refresh: false)

    # Then apply authentication
    return Legate::Auth::ToolIntegration.apply_authentication(
      request,
      scheme,
      credential,
      nil,
      token_manager
    )
  end

  # Fall back to token store if token manager not available
  token_store = get_token_store

  # Apply authentication
  Legate::Auth::ToolIntegration.apply_authentication(request, scheme, credential, token_store)
end

#authentication_error?(response) ⇒ Boolean

Check if a response indicates an authentication error

Parameters:

  • response (Hash) —

    The response to check

Returns:

  • (Boolean) —

    True if the response indicates an authentication error



115
116
117
118
 
# File 'lib/legate/tool_context.rb', line 115

def authentication_error?(response)
  require_relative 'auth/tool_integration'
  Legate::Auth::ToolIntegration.authentication_error?(response)
end

#clear_pending_state_delta! ⇒ Object

Clears any accumulated pending state changes within this context instance.



75
76
77
 
# File 'lib/legate/tool_context.rb', line 75

def clear_pending_state_delta!
  @pending_state_delta = {}
end

#clear_token(scheme, credential) ⇒ Boolean

Clear a cached authentication token

Parameters:

Returns:

  • (Boolean) —

    True if the token was cleared successfully



191
192
193
194
195
196
197
198
199
 
# File 'lib/legate/tool_context.rb', line 191

def clear_token(scheme, credential)
  token_store = get_token_store
  return false unless token_store

  require_relative 'auth/tool_integration'
  cache_key = Legate::Auth::ToolIntegration.generate_cache_key(scheme, credential)
  token_store.clear(cache_key)
  true
end

#find_agent_auth_for_url(url, auth_manager) ⇒ Array<Legate::Auth::Scheme, Legate::Auth::Credential>?

Find authentication for a URL using agent-specific mappings

Parameters:

  • url (String) —

    The URL to find authentication for

  • auth_manager (Legate::Auth::Manager) —

    The auth manager to resolve schemes/credentials

Returns:



273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
 
# File 'lib/legate/tool_context.rb', line 273

def find_agent_auth_for_url(url, auth_manager)
  return nil unless @agent_auth_config && @agent_auth_config[:url_mappings]

  @agent_auth_config[:url_mappings].each do |mapping|
    pattern = mapping[:pattern]
    next unless pattern

    matched = if pattern.is_a?(Regexp)
                !!(url =~ pattern)
              elsif pattern.is_a?(String)
                if pattern.include?('*')
                  # Convert glob pattern to regex
                  regex = Regexp.new('^' + Regexp.escape(pattern).gsub('\\*', '.*') + '$')
                  !!(url =~ regex)
                else
                  url == pattern || url.start_with?(pattern)
                end
              else
                false
              end

    next unless matched

    scheme_name = mapping[:scheme_name]
    credential_name = mapping[:credential_name]

    # Resolve from Auth::Manager
    scheme = auth_manager.get_scheme(scheme_name)
    credential = auth_manager.get_credential(credential_name)

    return [scheme, credential] if scheme && credential

    Legate.logger.warn { "[ToolContext] Agent auth mapping matched but scheme '#{scheme_name}' or credential '#{credential_name}' not found in Auth::Manager" }
  end

  nil
end

#get_token(scheme, credential, force_refresh: false) ⇒ Legate::Auth::ExchangedCredential?

Get an authentication token from the session cache

Parameters:

Returns:



133
134
135
136
137
138
139
140
141
142
143
144
145
 
# File 'lib/legate/tool_context.rb', line 133

def get_token(scheme, credential, force_refresh: false)
  # Try to use token manager if available
  token_manager = get_token_manager
  return token_manager.get_token(scheme, credential, force_refresh: force_refresh) if token_manager

  # Fall back to old mechanism if token manager not available
  token_store = get_token_store
  return nil unless token_store

  require_relative 'auth/tool_integration'
  cache_key = Legate::Auth::ToolIntegration.generate_cache_key(scheme, credential)
  token_store.get(cache_key)
end

#get_token_manager ⇒ Legate::Auth::TokenManager?

Create or get the token manager for this context

Returns:



313
314
315
316
317
318
319
320
321
 
# File 'lib/legate/tool_context.rb', line 313

def get_token_manager
  return @token_manager if @token_manager

  token_store = get_token_store
  return nil unless token_store

  require_relative 'auth/token_manager'
  @token_manager = Legate::Auth::TokenManager.new(token_store)
end

#handle_request_auth(request, options = {}) ⇒ Hash

Handle authentication for a request, automatically selecting the appropriate scheme and credential Checks agent-specific auth config first, then falls back to global Auth::Manager

Parameters:

  • request (Hash) —

    The request to authenticate

  • options (Hash) (defaults to: {}) —

    Options for authentication (e.g., credential_name, scheme_type)

Returns:

  • (Hash) —

    The authenticated request or the original request if authentication not possible



232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
 
# File 'lib/legate/tool_context.rb', line 232

def handle_request_auth(request, options = {})
  # Skip if request doesn't need authentication
  return request unless requires_authentication?(request)

  require_relative 'auth/manager'

  begin
    auth_manager = Legate::Auth::Manager.instance
    scheme = nil
    credential = nil

    # First, check agent-specific URL mappings
    if @agent_auth_config && @agent_auth_config[:url_mappings]&.any?
      scheme, credential = find_agent_auth_for_url(request[:url], auth_manager)
      if scheme && credential
        Legate.logger.debug { "[ToolContext] Using agent-specific auth for URL: #{request[:url]}" }
        return authenticate_request(request, scheme, credential)
      end
    end

    # Fall back to global Auth::Manager lookup
    scheme, credential = auth_manager.find_scheme_and_credential(
      url: request[:url],
      scheme_type: options[:scheme_type],
      credential_name: options[:credential_name]
    )

    # Apply authentication if found
    return authenticate_request(request, scheme, credential) if scheme && credential
  rescue StandardError => e
    Legate.logger.error("Error in automatic authentication: #{e.message}")
  end

  # Return the original request if no authentication applied
  request
end

#refresh_token(scheme, credential, token = nil) ⇒ Legate::Auth::ExchangedCredential?

Refresh an authentication token

Parameters:

Returns:



152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
 
# File 'lib/legate/tool_context.rb', line 152

def refresh_token(scheme, credential, token = nil)
  # Try to use token manager if available
  token_manager = get_token_manager
  return token_manager.refresh_token(scheme, credential, token) if token_manager

  # Fall back to direct refresh if token manager not available
  if token && scheme.supports_refresh? && token.refreshable?
    begin
      refreshed = scheme.refresh_token(token, credential)
      store_token(scheme, credential, refreshed)
      return refreshed
    rescue Legate::Auth::TokenRefreshError => e
      Legate.logger.error("Failed to refresh token: #{e.message}")
      return nil
    end
  end

  nil
end

#requires_authentication?(request) ⇒ Boolean

Check if a request likely requires authentication

Parameters:

  • request (Hash) —

    The request to check

Returns:

  • (Boolean) —

    True if the request likely requires authentication



123
124
125
126
 
# File 'lib/legate/tool_context.rb', line 123

def requires_authentication?(request)
  require_relative 'auth/tool_integration'
  Legate::Auth::ToolIntegration.requires_authentication?(request)
end

#revoke_token(scheme, credential, token) ⇒ Boolean

Revoke a token with the authentication provider

Parameters:

Returns:

  • (Boolean) —

    True if the token was revoked successfully



206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
 
# File 'lib/legate/tool_context.rb', line 206

def revoke_token(scheme, credential, token)
  # Try to use token manager if available
  token_manager = get_token_manager
  return token_manager.revoke_token(scheme, credential, token) if token_manager

  # Fall back to direct revocation if token manager not available
  unless scheme.respond_to?(:revoke_token)
    Legate.logger.warn("Scheme #{scheme.scheme_type} does not support token revocation")
    return false
  end

  begin
    result = scheme.revoke_token(token, credential)
    clear_token(scheme, credential) if result
    result
  rescue Legate::Auth::Error => e
    Legate.logger.error("Failed to revoke token: #{e.message}")
    false
  end
end

#state_get(key) ⇒ Object?

Retrieves a value from the session state via the session_service.

Parameters:

  • key (Symbol, String) —

    The key to retrieve

Returns:

  • (Object, nil) —

    The value or nil if not found



41
42
43
44
45
46
47
48
49
50
51
52
 
# File 'lib/legate/tool_context.rb', line 41

def state_get(key)
  unless @session_service
    Legate.logger.warn { '[ToolContext] state_get called but no session_service available.' }
    return nil
  end

  Legate.logger.debug { "[ToolContext] state_get for key: #{key} in session: #{@session_id}" }
  @session_service.get_state(session_id: @session_id, key: key)
rescue StandardError => e
  Legate.logger.error { "[ToolContext] Error in state_get for key '#{key}': #{e.message}" }
  nil
end

#state_set(key, value) ⇒ Object

Sets a value in the pending state delta for this context.

Parameters:

  • key (Symbol, String) —

    The key to set

  • value (Object) —

    The value to store (should be serializable)



57
58
59
60
 
# File 'lib/legate/tool_context.rb', line 57

def state_set(key, value)
  Legate.logger.debug { "[ToolContext] state_set for key: #{key} to value: #{value.inspect} (pending)" }
  @pending_state_delta[key.to_sym] = value
end

#state_update(hash_to_merge) ⇒ Object

Merges a hash into the pending state delta for this context.

Parameters:

  • hash_to_merge (Hash) —

    The hash to merge into the pending state delta



64
65
66
67
68
69
70
71
72
 
# File 'lib/legate/tool_context.rb', line 64

def state_update(hash_to_merge)
  unless hash_to_merge.is_a?(Hash)
    Legate.logger.warn { "[ToolContext] state_update called with non-hash: #{hash_to_merge.class}" }
    return
  end

  Legate.logger.debug { "[ToolContext] state_update with hash: #{hash_to_merge.inspect} (pending)" }
  @pending_state_delta.merge!(hash_to_merge.transform_keys(&:to_sym))
end

#store_token(scheme, credential, token) ⇒ Boolean

Store an authentication token in the session cache

Parameters:

Returns:

  • (Boolean) —

    True if the token was stored successfully



177
178
179
180
181
182
183
184
185
 
# File 'lib/legate/tool_context.rb', line 177

def store_token(scheme, credential, token)
  token_store = get_token_store
  return false unless token_store

  require_relative 'auth/tool_integration'
  cache_key = Legate::Auth::ToolIntegration.generate_cache_key(scheme, credential)
  token_store.store(cache_key, token)
  true
end

#to_h ⇒ Object 



323
324
325
326
327
328
329
330
331
332
 
# File 'lib/legate/tool_context.rb', line 323

def to_h
  {
    session_id: @session_id,
    user_id: @user_id,
    app_name: @app_name,
    invocation_id: @invocation_id,
    tool_registry_object_id: @tool_registry&.object_id,
    session_service_present: !@session_service.nil?
  }
end