lex-llm-ledger 0.5.0 → 0.7.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

@@ -23,11 +23,6 @@ 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]
@@ -125,6 +120,8 @@ module Legion
               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,
@@ -135,6 +132,12 @@ module Legion
                                            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,
@@ -149,7 +152,7 @@ module Legion
               [row, false]
             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)
               return nil unless tool_call_row # rubocop:disable Legion/Extension/RunnerReturnHash
 
               tool_call_id = tool_call_row[:id]
@@ -167,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

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.5.0'
+        VERSION = '0.7.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),
@@ -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,18 @@ 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:            if request_payload(body)
+                                                                      phi_protect(json_dump(request_payload(body)),
+                                                                                  contains_phi?(body))
+                                                                    end,
                                            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')
@@ -187,13 +198,28 @@ module Legion
             end
 
             def find_or_create_response(db, request, response_message, body)
-              response_uuid = stable_uuid(reference(body, :provider_response_ref) || "response:#{request_ref(body)}")
+              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
+
+              # Fallback: if we couldn't find a response by UUID, check if a response
+              # already exists for this request (e.g., metering arrived first and created
+              # a response with a different UUID). Enrich it instead of creating a duplicate.
+              unless existing
+                existing = db[:llm_message_inference_responses]
+                           .where(message_inference_request_id: request[:id])
+                           .first
+                log.debug("[ledger] response fallback: found existing response id=#{existing[:id]} for request_id=#{request[:id]}") if existing
+              end
+
               if existing
                 enrich_response!(db, existing, response_message, body)
                 return existing
               end
 
+              vis_resp = visible_response(body)
+              think_resp = thinking_response(body)
+              phi = contains_phi?(body)
+
               id = insert_with_savepoint(db, :llm_message_inference_responses, {
                                            uuid:                         response_uuid,
                                            message_inference_request_id: request[:id],
@@ -209,9 +235,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:                vis_resp ? phi_protect(json_dump(vis_resp), phi) : nil,
+                                           response_thinking_json:       think_resp ? phi_protect(json_dump(think_resp), phi) : nil,
                                            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),
@@ -239,11 +271,21 @@ module Legion
               update_if_missing(updates, existing, :dispatch_path, body[:dispatch_path] || body[:tier])
               update_if_missing(updates, existing, :identity_canonical_name, identity_canonical_name(body))
 
-              response_json = json_dump(visible_response(body))
-              update_if_placeholder(updates, existing, :response_json, response_json)
+              vis = visible_response(body)
+              if vis
+                response_json = json_dump(vis)
+                update_if_placeholder(updates, existing, :response_json, response_json)
+              elsif existing[:response_json].nil?
+                # Nothing to add, but also don't overwrite existing data with nil
+              end
 
-              thinking_json = json_dump(thinking_response(body))
-              update_if_placeholder(updates, existing, :response_thinking_json, thinking_json)
+              think = thinking_response(body)
+              if think
+                thinking_json = json_dump(think)
+                update_if_placeholder(updates, existing, :response_thinking_json, thinking_json)
+              elsif existing[:response_thinking_json].nil?
+                # Nothing to add
+              end
 
               return if updates.empty?
 
@@ -251,12 +293,27 @@ module Legion
               log.info("[ledger] enriched response id=#{existing[:id]} fields=#{updates.keys.join(',')}")
             end
 
+            # Core guard: never upsert a value that is nil, empty string, or empty JSON object.
+            # This prevents a leaner message (e.g., metering) from overwriting valid data
+            # written by a richer message (e.g., prompt audit).
+            def upsert_guard?(value)
+              return false if value.nil?
+              return false if value.is_a?(String) && value.strip.empty?
+              return false if value.to_s == '{}'
+
+              true
+            end
+
             def update_if_missing(updates, existing, key, value)
-              updates[key] = value if existing[key].nil? && present?(value)
+              updates[key] = value if existing[key].nil? && upsert_guard?(value)
             end
 
             def update_if_placeholder(updates, existing, key, value)
-              updates[key] = value if existing[key].to_s == '{}' && value != '{}'
+              return unless upsert_guard?(value)
+
+              existing_val = existing[key]
+              is_placeholder = ['{}', 'null'].include?(existing_val.to_s)
+              updates[key] = value if is_placeholder
             end
 
             def find_or_create_metric(db, request, response, body)
@@ -324,15 +381,19 @@ module Legion
               updates = {}
               update_if_missing(updates, existing, :latest_message_id, latest_message&.dig(:id))
               caller_refs = caller_identity_refs(db, body)
-              updates[:caller_identity_id] = caller_refs[:identity_id] if existing[:caller_identity_id].nil? && caller_refs[:identity_id]
-              updates[:caller_principal_id] = caller_refs[:principal_id] if existing[:caller_principal_id].nil? && caller_refs[:principal_id]
-              updates[:runtime_caller_type] = caller_type(body) if existing[:runtime_caller_type].nil? && caller_type(body)
+              update_if_missing(updates, existing, :caller_identity_id, caller_refs[:identity_id])
+              update_if_missing(updates, existing, :caller_principal_id, caller_refs[:principal_id])
+              update_if_missing(updates, existing, :runtime_caller_type, caller_type(body))
               update_if_missing(updates, existing, :runtime_caller_class, runtime_caller_class(body))
               update_if_missing(updates, existing, :runtime_caller_client, runtime_caller_client(body))
               update_if_missing(updates, existing, :identity_canonical_name, identity_canonical_name(body))
 
-              request_json = json_dump(request_payload(body))
-              updates[:request_json] = request_json if existing[:request_json].to_s == '{}' && request_json != '{}'
+              request_json = request_payload(body) ? json_dump(request_payload(body)) : nil
+              if request_json
+                update_if_placeholder(updates, existing, :request_json, request_json)
+              elsif existing[:request_json].nil?
+                # Nothing to add
+              end
 
               msg_count = Array(body.dig(:request, :messages) || body[:messages]).size
               updates[:context_message_count] = msg_count if existing[:context_message_count].to_i.zero? && msg_count.positive?
@@ -377,6 +438,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 +619,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) || body.dig(:caller, :parent_request_ref)
+              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
@@ -599,7 +676,7 @@ module Legion
             end
 
             def request_payload(body)
-              body[:request] || body[:messages] || {}
+              body[:request] || body[:messages]
             end
 
             def request_content(body)
@@ -610,7 +687,9 @@ module Legion
             end
 
             def visible_response(body)
-              response = body[:response] || body[:response_content] || body[:content] || {}
+              response = body[:response] || body[:response_content] || body[:content]
+              return nil if response.nil? || (response.is_a?(Hash) && response.empty?)
+
               if response.is_a?(String)
                 clean, _thinking = extract_inline_thinking(response)
                 return { content: clean }
@@ -630,10 +709,10 @@ module Legion
               end
 
               content_str = body[:response_content] || body[:response] || body[:content]
-              return {} unless content_str.is_a?(String)
+              return nil unless content_str.is_a?(String)
 
               _clean, extracted = extract_inline_thinking(content_str)
-              extracted ? { content: extracted } : {}
+              extracted ? { content: extracted } : nil
             end
 
             def extract_inline_thinking(text)
@@ -646,7 +725,10 @@ module Legion
             end
 
             def response_content(body)
-              stringify_content(visible_response(body)[:content] || visible_response(body).dig(:message, :content))
+              vis = visible_response(body)
+              return nil unless vis
+
+              stringify_content(vis[:content] || vis.dig(:message, :content))
             end
 
             def finish_reason(body)
@@ -656,8 +738,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 +794,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 +829,18 @@ 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]
+              present?(val) ? val.to_i : 0
+            end
           end
         end
       end