lex-llm-ledger 0.4.2 → 0.6.0

data/lib/legion/extensions/llm/ledger/runners/retention_purge.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+require 'legion/logging'
+
+module Legion
+  module Extensions
+    module Llm
+      module Ledger
+        module Runners
+          module RetentionPurge
+            extend self
+            extend Legion::Logging::Helper
+
+            PURGEABLE_TABLES = %i[
+              llm_conversations
+            ].freeze
+
+            BATCH_SIZE = 500
+
+            def purge_expired
+              db = ::Legion::Data.connection
+              total_deleted = 0
+
+              PURGEABLE_TABLES.each do |table|
+                next unless db.table_exists?(table)
+
+                deleted = purge_table(db, table)
+                total_deleted += deleted
+                log.info("[ledger] retention_purge: deleted #{deleted} expired rows from #{table}") if deleted.positive?
+              end
+
+              { result: :ok, deleted: total_deleted }
+            rescue StandardError => e
+              handle_exception(e, level: :error, handled: true, operation: 'retention_purge')
+              { result: :error, error: e.message }
+            end
+
+            private
+
+            def purge_table(db, table)
+              deleted = 0
+              loop do
+                ids = db[table]
+                      .where { expires_at <= Time.now.utc }
+                      .where(Sequel.~(expires_at: nil))
+                      .select(:id)
+                      .limit(BATCH_SIZE)
+                      .select_map(:id)
+                break if ids.empty?
+
+                db[table].where(id: ids).delete
+                deleted += ids.size
+                break if ids.size < BATCH_SIZE
+              end
+              deleted
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/ledger/runners/skills.rb

@@ -0,0 +1,99 @@
+# frozen_string_literal: true
+
+require 'digest'
+require 'securerandom'
+require_relative '../helpers/caller_identity'
+require_relative '../helpers/json'
+require_relative '../helpers/persistence_logging'
+
+module Legion
+  module Extensions
+    module Llm
+      module Ledger
+        module Runners
+          module Skills
+            extend self
+
+            def write_skill_record(payload = nil, metadata = {}, **message)
+              payload, metadata = normalize_runner_args(payload, metadata, message)
+              headers = Helpers::SubscriptionMessage.extract_headers(payload, metadata)
+              props   = metadata[:properties] || {}
+
+              body = payload.is_a?(Hash) ? payload : Helpers::Decryption.decrypt_if_needed(payload, metadata)
+
+              db = ::Legion::Data.connection
+              record = build_skill_record(db, body, props, headers)
+
+              Helpers::PersistenceLogging.insert_row(
+                db, :llm_skill_events, record,
+                operation: 'write_skill_record'
+              )
+              { result: :ok }
+            rescue Sequel::UniqueConstraintViolation => e
+              log.warn("write_skill_record duplicate insert ignored: #{e.message}")
+              { result: :duplicate }
+            rescue Helpers::DecryptionUnavailable => e
+              handle_exception(e, level: :warn, handled: true, operation: 'write_skill_record.decrypt')
+              raise
+            rescue Helpers::DecryptionFailed => e
+              handle_exception(e, level: :error, handled: true, operation: 'write_skill_record.decrypt')
+              raise
+            rescue StandardError => e
+              handle_exception(e, level: :error, handled: true, operation: 'write_skill_record')
+              { result: :error, error: e.message }
+            end
+
+            private
+
+            def normalize_runner_args(payload, metadata, message)
+              Helpers::SubscriptionMessage.runner_args(payload, metadata, message)
+            end
+
+            def build_skill_record(db, body, props, headers)
+              identity = Helpers::CallerIdentity.normalize(
+                caller_raw: body[:caller], identity: body[:identity], headers: headers
+              )
+              skill = body[:skill] || {}
+
+              {
+                uuid:                    stable_uuid(props[:message_id] || body[:event_id] || SecureRandom.uuid),
+                conversation_id:         resolve_conversation_id(db, body, headers),
+                request_ref:             body[:request_id] || props[:correlation_id],
+                skill_name:              skill[:name] || body[:skill_name],
+                skill_version:           skill[:version],
+                trigger:                 skill[:trigger] || body[:trigger],
+                status:                  body[:status] || 'completed',
+                duration_ms:             body[:duration_ms].to_i,
+                identity_canonical_name: identity[:identity],
+                identity_principal_id:   identity[:principal_id],
+                identity_id:             identity[:identity_id],
+                recorded_at:             body[:recorded_at] || body[:timestamp] || Time.now.utc,
+                inserted_at:             Time.now.utc
+              }
+            end
+
+            def resolve_conversation_id(db, body, headers)
+              conv_ref = body[:conversation_id] || headers['x-legion-llm-conversation-id']
+              return nil unless conv_ref # rubocop:disable Legion/Extension/RunnerReturnHash
+
+              conv = db[:llm_conversations].where(uuid: stable_uuid(conv_ref)).first ||
+                     db[:llm_conversations].where(uuid: conv_ref).first
+              conv&.[](:id)
+            end
+
+            def stable_uuid(value)
+              raw = value.to_s
+              return raw if raw.length <= 36 # rubocop:disable Legion/Extension/RunnerReturnHash
+
+              hex = Digest::SHA256.hexdigest(raw)[0, 32]
+              "#{hex[0, 8]}-#{hex[8, 4]}-#{hex[12, 4]}-#{hex[16, 4]}-#{hex[20, 12]}"
+            end
+
+            include Legion::Extensions::Helpers::Lex if Legion::Extensions.const_defined?(:Helpers, false) &&
+                                                        Legion::Extensions::Helpers.const_defined?(:Lex, false)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/ledger/runners/tools.rb

@@ -10,8 +10,6 @@ module Legion
   module Extensions
     module Llm
       module Ledger
-        class ResponseNotReady < StandardError; end
-
         module Runners
           module Tools
             extend self
@@ -25,17 +23,14 @@ module Legion
               ctx  = body[:message_context] || {}
               tool = body[:tool_call]       || {}
 
-              Helpers::Retention.resolve(
-                retention:    headers['x-legion-retention'],
-                contains_phi: headers['x-legion-contains-phi'] == 'true'
-              )
-
               db = ::Legion::Data.connection
+              response = find_or_resolve_response_with_retry(db, body, ctx, props, headers)
               write_result = [:ok]
               db.transaction do
-                response                     = find_or_resolve_response(db, body, ctx, props, headers)
                 identity_attrs               = extract_identity_attrs(body, headers, db)
-                tool_call_row, new_tool_call = find_or_create_tool_call(db, response, body, ctx, tool, headers, identity_attrs)
+                conversation_id              = resolve_conversation_id(db, body, ctx, headers)
+                tool_call_row, new_tool_call = find_or_create_tool_call(db, response, body, ctx, tool, headers,
+                                                                        identity_attrs, conversation_id)
                 if tool_call_row && !new_tool_call
                   write_result[0] = :duplicate
                 elsif new_tool_call
@@ -53,9 +48,6 @@ module Legion
             rescue Helpers::DecryptionFailed => e
               handle_exception(e, level: :error, handled: true, operation: 'write_tool_record.decrypt')
               raise
-            rescue ResponseNotReady => e
-              log.warn("[ledger] write_tool_record: parent response not yet written, dead-lettering message (#{e.message})")
-              raise unrecoverable_message_error(e)
             rescue StandardError => e
               handle_exception(e, level: :error, handled: true, operation: 'write_tool_record')
               { result: :error, error: e.message }
@@ -67,16 +59,26 @@ module Legion
               Helpers::SubscriptionMessage.runner_args(payload, metadata, message)
             end
 
-            def unrecoverable_message_error(error)
-              if defined?(Legion::Extensions::Actors::UnrecoverableMessageError)
-                Legion::Extensions::Actors::UnrecoverableMessageError.new(error.message)
-              else
-                error
+            def find_or_resolve_response_with_retry(db, body, ctx, props, headers)
+              response = find_or_resolve_response(db, body, ctx, props, headers)
+              return response if response # rubocop:disable Legion/Extension/RunnerReturnHash
+
+              retry_attempts = tool_write_setting(:response_retry_attempts, 3)
+              retry_delay    = tool_write_setting(:response_retry_delay, 1)
+
+              retry_attempts.times do |attempt|
+                sleep retry_delay
+                response = find_or_resolve_response(db, body, ctx, props, headers)
+                if response
+                  log.debug("[ledger] write_tool_record: response found on retry #{attempt + 1}")
+                  return response # rubocop:disable Legion/Extension/RunnerReturnHash
+                end
               end
+
+              log.info('[ledger] write_tool_record: response not available after retries, proceeding with null response_id')
+              nil
             end
 
-            # Resolve the llm_message_inference_responses row this tool call belongs to.
-            # Returns nil if we cannot link at all.
             def find_or_resolve_response(db, body, ctx, props, headers)
               request_ref = ctx[:request_id] || body[:request_id] ||
                             props[:correlation_id] || headers['x-legion-llm-request-id']
@@ -89,31 +91,53 @@ module Legion
                 .where(message_inference_request_id: request[:id]).first
             end
 
-            def find_or_create_tool_call(db, response, body, ctx, tool, headers, identity_attrs)
+            def resolve_conversation_id(db, body, ctx, headers)
+              conv_ref = ctx[:conversation_id] || body[:conversation_id] ||
+                         headers['x-legion-llm-conversation-id']
+              return nil unless conv_ref # rubocop:disable Legion/Extension/RunnerReturnHash
+
+              conv = db[:llm_conversations].where(uuid: stable_uuid(conv_ref)).first ||
+                     db[:llm_conversations].where(uuid: conv_ref).first
+              conv&.[](:id)
+            end
+
+            def find_or_create_tool_call(db, response, body, ctx, tool, headers, identity_attrs, conversation_id) # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
               tool_uuid = derive_tool_call_uuid(body, ctx, tool, headers)
               existing  = db[:llm_tool_calls].where(uuid: tool_uuid).first
               return [existing, false] if existing # rubocop:disable Legion/Extension/RunnerReturnHash
 
-              response_id = resolve_response_id(db, response, body, ctx, headers, tool_uuid)
-              return [nil, false] unless response_id # rubocop:disable Legion/Extension/RunnerReturnHash
+              response_id = response&.[](:id)
 
-              next_index = db[:llm_tool_calls]
-                           .where(message_inference_response_id: response_id)
-                           .max(:tool_call_index).to_i + 1
+              next_index = if response_id
+                             db[:llm_tool_calls]
+                               .where(message_inference_response_id: response_id)
+                               .max(:tool_call_index).to_i + 1
+                           else
+                             0
+                           end
 
               src    = tool[:source] || {}
               status = tool[:status] || headers['x-legion-tool-status'] || 'success'
               ts     = body[:timestamps] || {}
 
+              result_value = tool[:result] || body[:result]
+              has_result = tool[:result] || body.key?(:result)
               id = insert_with_savepoint(db, :llm_tool_calls, {
                                            uuid:                          tool_uuid,
                                            message_inference_response_id: response_id,
+                                           conversation_id:               conversation_id,
                                            tool_call_index:               next_index,
                                            provider_tool_call_ref:        tool[:id],
                                            tool_name:                     tool[:name] || headers['x-legion-tool-name'],
                                            tool_source_type:              src[:type] || headers['x-legion-tool-source-type'],
                                            tool_source_server:            src[:server] || headers['x-legion-tool-source-server'],
                                            status:                        status,
+                                           tool_arguments_json:           tool[:arguments] ? Helpers::Json.dump(tool[:arguments]) : nil,
+                                           tool_result_json:              has_result ? Helpers::Json.dump(result_value) : nil,
+                                           tool_category:                 tool[:category] || tool[:tool_category],
+                                           data_handling_classification:  tool[:data_handling_classification],
+                                           policy_decision:               tool[:policy_decision],
+                                           requires_human_approval:       tool[:requires_human_approval],
                                            requested_at:                  ts[:tool_start] || tool[:started_at],
                                            completed_at:                  ts[:tool_end] || tool[:finished_at],
                                            **identity_attrs,
@@ -128,35 +152,7 @@ module Legion
               [row, false]
             end
 
-            # Extract or fall back to find a response_id for linking the tool call.
-            def resolve_response_id(db, response, body, ctx, headers, tool_uuid)
-              return response[:id] if response # rubocop:disable Legion/Extension/RunnerReturnHash
-
-              fallback = fallback_response_for_conversation(db, body, ctx, headers)
-              return fallback[:id] if fallback # rubocop:disable Legion/Extension/RunnerReturnHash
-
-              raise ResponseNotReady, "no response row found for tool call uuid=#{tool_uuid}"
-            end
-
-            def fallback_response_for_conversation(db, body, ctx, headers)
-              conv_id = ctx[:conversation_id] || body[:conversation_id] ||
-                        headers['x-legion-llm-conversation-id']
-              return nil unless conv_id # rubocop:disable Legion/Extension/RunnerReturnHash
-
-              conv = db[:llm_conversations].where(uuid: stable_uuid(conv_id)).first ||
-                     db[:llm_conversations].where(uuid: conv_id).first
-              return nil unless conv # rubocop:disable Legion/Extension/RunnerReturnHash
-
-              db[:llm_message_inference_responses]
-                .join(:llm_message_inference_requests,
-                      id: :message_inference_request_id)
-                .where(Sequel[:llm_message_inference_requests][:conversation_id] => conv[:id])
-                .order(Sequel.desc(Sequel[:llm_message_inference_responses][:id]))
-                .select_all(:llm_message_inference_responses)
-                .first
-            end
-
-            def find_or_create_tool_call_attempt(db, tool_call_row, tool, body, props, headers, identity_attrs) # rubocop:disable Metrics/CyclomaticComplexity
+            def find_or_create_tool_call_attempt(db, tool_call_row, tool, body, props, headers, identity_attrs) # rubocop:disable Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity, Metrics/AbcSize
               return nil unless tool_call_row # rubocop:disable Legion/Extension/RunnerReturnHash
 
               tool_call_id = tool_call_row[:id]
@@ -174,21 +170,24 @@ module Legion
               runner_ref = body[:worker_id] || body[:runner_ref] || props[:app_id]
 
               id = insert_with_savepoint(db, :llm_tool_call_attempts, {
-                                           uuid:           attempt_uuid,
-                                           tool_call_id:   tool_call_id,
-                                           attempt_no:     attempt_no,
-                                           runner_ref:     runner_ref,
-                                           status:         status,
-                                           error_category: error_hash[:category] || error_hash[:type],
-                                           error_code:     error_hash[:code],
-                                           error_message:  error_info.is_a?(String) ? error_info : error_hash[:message],
-                                           duration_ms:    tool[:duration_ms].to_i,
-                                           arguments_ref:  sha256_ref(tool[:arguments]),
-                                           result_ref:     sha256_ref(tool[:result] || body[:result]),
-                                           started_at:     ts[:tool_start] || tool[:started_at],
-                                           ended_at:       ts[:tool_end] || tool[:finished_at],
+                                           uuid:                attempt_uuid,
+                                           tool_call_id:        tool_call_id,
+                                           attempt_no:          attempt_no,
+                                           runner_ref:          runner_ref,
+                                           status:              status,
+                                           error_category:      error_hash[:category] || error_hash[:type],
+                                           error_code:          error_hash[:code],
+                                           error_message:       error_info.is_a?(String) ? error_info : error_hash[:message],
+                                           duration_ms:         tool[:duration_ms].to_i,
+                                           arguments_ref:       sha256_ref(tool[:arguments]),
+                                           result_ref:          sha256_ref(tool[:result] || body[:result]),
+                                           attempt_input_json:  tool[:arguments] ? Helpers::Json.dump(tool[:arguments]) : nil,
+                                           attempt_output_json: Helpers::Json.dump(tool[:result] || body[:result]),
+                                           error_details_json:  tool[:error] ? Helpers::Json.dump(tool[:error]) : nil,
+                                           started_at:          ts[:tool_start] || tool[:started_at],
+                                           ended_at:            ts[:tool_end] || tool[:finished_at],
                                            **identity_attrs,
-                                           inserted_at:    Time.now.utc
+                                           inserted_at:         Time.now.utc
                                          }, operation: 'write_tool_record.attempt')
               db[:llm_tool_call_attempts][id: id]
             rescue Sequel::UniqueConstraintViolation => e
@@ -265,6 +264,12 @@ module Legion
               "#{hex[0, 8]}-#{hex[8, 4]}-#{hex[12, 4]}-#{hex[16, 4]}-#{hex[20, 12]}"
             end
 
+            def tool_write_setting(key, default)
+              ledger = Legion::Settings.dig(:extensions, :llm, :ledger) || {}
+              tool_write = ledger[:tool_write] || {}
+              (tool_write[key] || default).to_i
+            end
+
             def insert_with_savepoint(db, table, attributes, operation:)
               db.transaction(savepoint: true) do
                 Helpers::PersistenceLogging.insert_row(db, table, attributes,

data/lib/legion/extensions/llm/ledger/transport/exchanges/escalation.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Llm
+      module Ledger
+        module Transport
+          module Exchanges
+            class Escalation < ::Legion::Transport::Exchange
+              def exchange_name
+                'llm.escalation'
+              end
+
+              def default_type
+                'topic'
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/ledger/transport/queues/audit_escalations.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Llm
+      module Ledger
+        module Transport
+          module Queues
+            class AuditEscalations < Legion::Transport::Queue
+              def queue_name
+                'llm.audit.escalations'
+              end
+
+              def queue_options
+                { durable: true }
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/ledger/transport/queues/audit_skills.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Llm
+      module Ledger
+        module Transport
+          module Queues
+            class AuditSkills < Legion::Transport::Queue
+              def queue_name
+                'llm.audit.skills'
+              end
+
+              def queue_options
+                { durable: true }
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/ledger/transport/transport.rb

@@ -3,6 +3,7 @@
 require 'legion/extensions/transport'
 require_relative 'exchanges/metering'
 require_relative 'exchanges/audit'
+require_relative 'exchanges/escalation'
 require_relative 'exchanges/registry'
 
 module Legion
@@ -29,6 +30,16 @@ module Legion
                 to:          Legion::Extensions::Llm::Ledger::Transport::Queues::AuditTools,
                 routing_key: 'audit.tool.#'
               },
+              {
+                from:        Legion::Extensions::Llm::Ledger::Transport::Exchanges::Audit,
+                to:          Legion::Extensions::Llm::Ledger::Transport::Queues::AuditSkills,
+                routing_key: 'audit.skill.#'
+              },
+              {
+                from:        Legion::Extensions::Llm::Ledger::Transport::Exchanges::Escalation,
+                to:          Legion::Extensions::Llm::Ledger::Transport::Queues::AuditEscalations,
+                routing_key: '#'
+              },
               {
                 from:        Legion::Extensions::Llm::Ledger::Transport::Exchanges::Registry,
                 to:          Legion::Extensions::Llm::Ledger::Transport::Queues::RegistryAvailability,

data/lib/legion/extensions/llm/ledger/version.rb

@@ -4,7 +4,7 @@ module Legion
   module Extensions
     module Llm
       module Ledger
-        VERSION = '0.4.2'
+        VERSION = '0.6.0'
       end
     end
   end

data/lib/legion/extensions/llm/ledger/writers/official_record_writer.rb

@@ -63,6 +63,8 @@ module Legion
                                            classification_level:    classification_level(body),
                                            contains_phi:            contains_phi?(body),
                                            contains_pii:            contains_pii?(body),
+                                           pii_types_json:          json_dump(Array(body.dig(:classification, :pii_types))),
+                                           jurisdictions_json:      json_dump(Array(body.dig(:classification, :jurisdictions) || body[:jurisdictions])),
                                            retention_policy:        body[:retention_policy] || 'default',
                                            expires_at:              body[:expires_at],
                                            identity_canonical_name: identity_canonical_name(body),
@@ -110,7 +112,7 @@ module Legion
               end
             end
 
-            def find_or_create_request(db, conversation, latest_message, body)
+            def find_or_create_request(db, conversation, latest_message, body) # rubocop:disable Metrics/AbcSize
               request_id = request_ref(body)
               existing = db[:llm_message_inference_requests].where(request_ref: request_id).first
               return enrich_request!(db, existing, body, latest_message) if existing
@@ -121,6 +123,7 @@ module Legion
                                            uuid:                    stable_uuid(request_id),
                                            conversation_id:         conversation[:id],
                                            latest_message_id:       latest_message&.dig(:id),
+                                           parent_request_id:       resolve_parent_request_id(db, body),
                                            caller_principal_id:     caller_refs[:principal_id],
                                            caller_identity_id:      caller_refs[:identity_id],
                                            identity_canonical_name: identity_canonical_name(body),
@@ -137,10 +140,15 @@ module Legion
                                            status:                  'responded',
                                            context_message_count:   Array(body.dig(:request, :messages) || body[:messages]).size,
                                            request_capture_mode:    'full',
-                                           request_json:            json_dump(request_payload(body)),
+                                           request_json:            phi_protect(json_dump(request_payload(body)), contains_phi?(body)),
                                            classification_level:    classification_level(body),
                                            cost_center:             billing(body)[:cost_center],
                                            budget_key:              billing(body)[:budget_id] || billing(body)[:budget_key],
+                                           injected_tool_count:     Array(body.dig(:audit, :injected_tools) || body[:injected_tools]).size,
+                                           context_tokens:          resolve_context_tokens(body),
+                                           request_content_hash:    compute_content_hash(body.dig(:request, :content) || body.dig(:audit, :request_content)),
+                                           curation_strategy:       body[:curation_strategy] || body.dig(:audit, :curation_strategy),
+                                           tool_policy:             body[:tool_policy] || body.dig(:audit, :tool_policy),
                                            requested_at:            recorded_at(body),
                                            inserted_at:             Time.now.utc
                                          }, operation: 'official_record_writer.inference_request')
@@ -186,14 +194,15 @@ module Legion
               end
             end
 
-            def find_or_create_response(db, request, response_message, body)
-              response_uuid = stable_uuid(reference(body, :provider_response_ref) || "response:#{request_ref(body)}")
+            def find_or_create_response(db, request, response_message, body) # rubocop:disable Metrics/AbcSize
+              response_uuid = stable_uuid(reference(body, :provider_response_ref) || "response:#{request_ref(body)}:#{body[:provider] || 'unknown'}")
               existing = db[:llm_message_inference_responses].where(uuid: response_uuid).first
               if existing
                 enrich_response!(db, existing, response_message, body)
                 return existing
               end
 
+              phi = contains_phi?(body)
               id = insert_with_savepoint(db, :llm_message_inference_responses, {
                                            uuid:                         response_uuid,
                                            message_inference_request_id: request[:id],
@@ -209,9 +218,15 @@ module Legion
                                            latency_ms:                   integer(body[:latency_ms]),
                                            wall_clock_ms:                integer(body[:wall_clock_ms]),
                                            response_capture_mode:        'full',
-                                           response_json:                json_dump(visible_response(body)),
-                                           response_thinking_json:       json_dump(thinking_response(body)),
+                                           response_json:                phi_protect(json_dump(visible_response(body)), phi),
+                                           response_thinking_json:       phi_protect(json_dump(thinking_response(body)), phi),
                                            dispatch_path:                body[:dispatch_path] || body[:tier],
+                                           error_category:               body[:error_category] || body.dig(:error, :category),
+                                           error_code:                   body[:error_code] || body.dig(:error, :code),
+                                           error_message:                body[:error_message] || body.dig(:error, :message),
+                                           response_content_hash:        compute_content_hash(body[:response_content] || body.dig(:audit, :response_content)),
+                                           route_attempts:               (body[:route_attempts] || body.dig(:audit, :route_attempts)).to_i,
+                                           escalation_chain_ref:         body[:escalation_chain_ref],
                                            identity_principal_id:        caller_identity_refs(db, body)[:principal_id],
                                            identity_id:                  caller_identity_refs(db, body)[:identity_id],
                                            identity_canonical_name:      identity_canonical_name(body),
@@ -377,6 +392,10 @@ module Legion
                 explicit_identity_id = integer_or_nil(body[:caller_identity_id] || body.dig(:caller, :requested_by, :id))
                 explicit_principal_id = integer_or_nil(body[:caller_principal_id] ||
                                                        body.dig(:caller, :requested_by, :principal_id))
+
+                explicit_identity_id ||= integer_or_nil(body[:__header_identity_id])
+                explicit_principal_id ||= integer_or_nil(body[:__header_principal_id])
+
                 refs = { principal_id: explicit_principal_id, identity_id: explicit_identity_id }.compact
                 unless refs[:principal_id] && refs[:identity_id]
                   if explicit_identity_id && !explicit_principal_id && identity_tables_available?(db)
@@ -554,6 +573,18 @@ module Legion
               int.positive? ? int : nil
             end
 
+            def resolve_parent_request_id(db, body)
+              parent_ref = body[:parent_request_id] || body.dig(:context, :parent_request_id)
+              return nil unless present?(parent_ref)
+
+              if parent_ref.is_a?(Integer)
+                parent_ref
+              else
+                parent = db[:llm_message_inference_requests].where(request_ref: parent_ref.to_s).first
+                parent&.dig(:id)
+              end
+            end
+
             def correlation_id(body)
               reference(body, :correlation_id, :correlation_ref) || body.dig(:tracing, :correlation_id)
             end
@@ -656,8 +687,14 @@ module Legion
               body.dig(:response, :finish_reason) || body.dig(:response, :stop, :reason)
             end
 
+            ALLOWED_CLASSIFICATION_LEVELS = %w[public internal confidential restricted].freeze
+
             def classification_level(body)
-              body[:classification_level] || body.dig(:classification, :level)
+              raw = body[:classification_level] || body.dig(:classification, :level)
+              return 'internal' if raw.nil? || raw.to_s.empty?
+
+              normalized = raw.to_s.downcase
+              ALLOWED_CLASSIFICATION_LEVELS.include?(normalized) ? normalized : 'internal'
             end
 
             def contains_phi?(body)
@@ -706,6 +743,19 @@ module Legion
               json_dump(content)
             end
 
+            def phi_protect(json_string, is_phi)
+              return json_string unless is_phi && crypt_available?
+
+              Legion::Crypt.encrypt(json_string)
+            rescue StandardError => e
+              handle_exception(e, level: :warn, handled: true, operation: 'official_record_writer.phi_encrypt')
+              json_string
+            end
+
+            def crypt_available?
+              defined?(Legion::Crypt) && Legion::Crypt.respond_to?(:encrypt)
+            end
+
             def json_dump(value)
               Helpers::Json.dump(value)
             end
@@ -728,6 +778,20 @@ module Legion
             def present?(value)
               !value.nil? && !(value.respond_to?(:empty?) && value.empty?)
             end
+
+            def compute_content_hash(content)
+              return nil if content.nil? || content.to_s.empty?
+
+              Digest::SHA256.hexdigest(json_dump(content))[0..31]
+            end
+
+            def resolve_context_tokens(body)
+              raw = body[:tokens] || body[:audit] || body
+              val = raw[:input_tokens] || raw[:input] || raw[:context_tokens] || raw[:prompt_tokens]
+              return nil unless present?(val)
+
+              val.to_i
+            end
           end
         end
       end