agent-session_context 0.1.0

data/lib/agent/session_context/evidence_packet.rb

@@ -0,0 +1,271 @@
+# frozen_string_literal: true
+
+module Agent
+  module SessionContext
+    class EvidencePacket
+      MAX_BYTES = 65_536
+      Result = Data.define(:chunks, :warnings, :source_refs) do
+        def source_refs_for(chunk)
+          ref_index = source_refs.to_h do |source_ref|
+            [EvidencePacket.safe_ref_text(source_ref), source_ref]
+          end
+          represented_refs = []
+          seen_refs = {}
+
+          String(chunk).each_line(chomp: true) do |line|
+            prefix, = line.split(" ", 2)
+            source_ref = ref_index[prefix]
+            next unless source_ref
+            next if seen_refs.key?(prefix)
+
+            seen_refs[prefix] = true
+            represented_refs << source_ref
+          end
+
+          represented_refs.freeze
+        end
+      end
+      TRUNCATABLE_KEYS = {
+        message: ["text"].freeze,
+        tool_use: %w[input name].freeze,
+        observed: ["label"].freeze
+      }.freeze
+      UNSAFE_SOURCE_REF_TEXT_PATTERN = /(?:\p{Space}|\p{Cntrl})/
+
+      class << self
+        def safe_ref_text(source_ref)
+          text = source_ref.to_s
+          normalized_text = normalize_ref_text(text, source_ref)
+
+          if normalized_text.match?(UNSAFE_SOURCE_REF_TEXT_PATTERN)
+            raise ArgumentError,
+                  "Source ref text cannot contain whitespace or control characters (#{source_ref_type(source_ref)})"
+          end
+
+          normalized_text
+        end
+
+        private
+
+        def normalize_ref_text(text, source_ref)
+          return text if text.encoding == Encoding::UTF_8 && text.valid_encoding?
+
+          utf8_text = duplicate_as_utf8(text)
+          return utf8_text if utf8_text.valid_encoding?
+          return utf8_text if utf8_bytes?(text)
+
+          raise ArgumentError,
+                "Source ref text must be valid UTF-8 or ASCII-only in a safely " \
+                "convertible encoding (#{source_ref_type(source_ref)})"
+        rescue Encoding::UndefinedConversionError, Encoding::InvalidByteSequenceError
+          raise ArgumentError,
+                "Source ref text must be valid UTF-8 or ASCII-only in a safely " \
+                "convertible encoding (#{source_ref_type(source_ref)})"
+        end
+
+        def duplicate_as_utf8(text)
+          String.new(text, encoding: text.encoding).dup.force_encoding(Encoding::UTF_8)
+        end
+
+        def utf8_bytes?(text)
+          text.bytes.all? { |byte| byte < 128 }
+        end
+
+        def source_ref_type(source_ref)
+          source_ref.class.name || source_ref.class.to_s
+        end
+      end
+
+      def call(transcript:, observed:)
+        warnings = []
+        carried_refs = []
+        seen_refs = {}
+        lines = []
+
+        transcript.entries.each do |entry|
+          entry.parts.each do |part|
+            line = serialize_transcript_part(entry.role, part, warnings)
+            next unless line
+
+            remember_ref(seen_refs, carried_refs, part.source_ref)
+            lines << line
+          end
+        end
+
+        observed_lines(observed).each do |source_ref, payload|
+          line = bounded_serialized_line(source_ref, payload, TRUNCATABLE_KEYS.fetch(:observed), warnings)
+          remember_ref(seen_refs, carried_refs, source_ref)
+          lines << line
+        end
+
+        Result.new(
+          chunks: chunk_lines(lines).freeze,
+          warnings: warnings.map { |warning| String.new(warning).freeze }.freeze,
+          source_refs: carried_refs.freeze
+        )
+      end
+
+      private
+
+      def serialize_transcript_part(role, part, warnings)
+        return if part.injected
+
+        case part.type
+        when :text
+          return unless %i[user assistant system].include?(role)
+
+          bounded_serialized_line(
+            part.source_ref,
+            {
+              "kind" => "message",
+              "role" => role.to_s,
+              "text" => part.text.to_s
+            },
+            TRUNCATABLE_KEYS.fetch(:message),
+            warnings
+          )
+        when :tool_use
+          bounded_serialized_line(
+            part.source_ref,
+            {
+              "kind" => "tool_use",
+              "role" => role.to_s,
+              "name" => part.name || "(unknown)",
+              "input" => part.text.to_s
+            },
+            TRUNCATABLE_KEYS.fetch(:tool_use),
+            warnings
+          )
+        end
+      end
+
+      def observed_lines(observed)
+        lines = []
+
+        observed.files.each do |item|
+          classifications = ["file"]
+          classifications.unshift("document") if EvidenceCollector.document_item?(item)
+          item.source_refs.each do |source_ref|
+            lines << [source_ref, observed_payload(item, classifications)]
+          end
+        end
+
+        lines
+      end
+
+      def observed_payload(item, classifications)
+        {
+          "kind" => "observed",
+          "classifications" => classifications.freeze,
+          "label" => item.label,
+          "action" => item.attributes.fetch(:action).to_s
+        }
+      end
+
+      def serialize_line(source_ref, payload)
+        source_ref_text = validated_source_ref_text(source_ref)
+        "#{source_ref_text} #{JSON.generate(payload)}"
+      end
+
+      def bounded_serialized_line(source_ref, payload, truncatable_keys, warnings)
+        fail_if_fixed_overhead_too_large(source_ref, payload, truncatable_keys)
+
+        line = serialize_line(source_ref, payload)
+        return line if line.bytesize <= MAX_BYTES
+
+        bounded_payload = payload.dup
+
+        truncatable_keys.each do |key|
+          next unless bounded_payload[key].is_a?(String)
+
+          bounded_value = maximal_fitting_value(source_ref, bounded_payload, key)
+          bounded_payload[key] = bounded_value
+
+          line = serialize_line(source_ref, bounded_payload)
+          return warning_line(source_ref, line, warnings) if line.bytesize <= MAX_BYTES
+        end
+
+        raise ArgumentError, "Evidence line for #{source_ref} cannot fit within #{MAX_BYTES} bytes"
+      end
+
+      def fail_if_fixed_overhead_too_large(source_ref, payload, truncatable_keys)
+        minimal_payload = payload.each_with_object({}) do |(key, value), normalized|
+          normalized[key] = truncatable_keys.include?(key) && value.is_a?(String) ? "" : value
+        end
+
+        return if serialize_line(source_ref, minimal_payload).bytesize <= MAX_BYTES
+
+        raise ArgumentError,
+              "Source ref and fixed payload overhead cannot fit within #{MAX_BYTES} bytes for #{source_ref}"
+      end
+
+      def maximal_fitting_value(source_ref, payload, key)
+        original = payload.fetch(key)
+        low = 0
+        high = original.bytesize
+        best = ""
+
+        while low <= high
+          middle = (low + high) / 2
+          candidate = utf8_prefix(original, middle)
+          payload[key] = candidate
+
+          if serialize_line(source_ref, payload).bytesize <= MAX_BYTES
+            best = candidate
+            low = middle + 1
+          else
+            high = middle - 1
+          end
+        end
+
+        payload[key] = original
+        best
+      end
+
+      def utf8_prefix(text, max_bytes)
+        String.new(text.byteslice(0, max_bytes), encoding: Encoding::UTF_8).scrub("")
+      end
+
+      def warning_line(source_ref, line, warnings)
+        warnings << "Truncated oversized evidence line for #{source_ref}"
+        line
+      end
+
+      def chunk_lines(lines)
+        chunks = []
+        current_chunk = String.new
+
+        lines.each do |line|
+          if current_chunk.empty?
+            current_chunk << line
+            next
+          end
+
+          projected_size = current_chunk.bytesize + 1 + line.bytesize
+
+          if projected_size <= MAX_BYTES
+            current_chunk << "\n" << line
+          else
+            chunks << current_chunk.freeze
+            current_chunk = String.new(line)
+          end
+        end
+
+        chunks << current_chunk.freeze unless current_chunk.empty?
+        chunks
+      end
+
+      def remember_ref(seen_refs, carried_refs, source_ref)
+        key = validated_source_ref_text(source_ref)
+        return if seen_refs.key?(key)
+
+        seen_refs[key] = true
+        carried_refs << source_ref
+      end
+
+      def validated_source_ref_text(source_ref)
+        self.class.safe_ref_text(source_ref)
+      end
+    end
+  end
+end

data/lib/agent/session_context/immutable_value.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module Agent
+  module SessionContext
+    module ImmutableValue
+      CYCLIC_ERROR = "cyclic arrays and hashes are not supported"
+      private_constant :CYCLIC_ERROR
+
+      class Copier
+        def initialize
+          @copies = {}.compare_by_identity
+          @active = {}.compare_by_identity
+        end
+
+        def copy(value)
+          case value
+          when String
+            String.new(value).freeze
+          when Array
+            copy_array(value)
+          when Hash
+            copy_hash(value)
+          else
+            value
+          end
+        end
+
+        private
+
+        def copy_array(array)
+          detect_cycle!(array)
+          return @copies.fetch(array) if @copies.key?(array)
+
+          duplicate = []
+          @copies[array] = duplicate
+          @active[array] = true
+          array.each { |entry| duplicate << copy(entry) }
+          duplicate.freeze
+        ensure
+          @active.delete(array)
+        end
+
+        def copy_hash(hash)
+          detect_cycle!(hash)
+          return @copies.fetch(hash) if @copies.key?(hash)
+
+          duplicate = {}
+          @copies[hash] = duplicate
+          @active[hash] = true
+          hash.each do |key, value|
+            duplicate[copy(key)] = copy(value)
+          end
+          duplicate.freeze
+        ensure
+          @active.delete(hash)
+        end
+
+        def detect_cycle!(value)
+          raise ArgumentError, CYCLIC_ERROR if @active[value]
+        end
+      end
+      private_constant :Copier
+
+      module_function
+
+      def copy(value)
+        Copier.new.copy(value)
+      end
+    end
+  end
+end

data/lib/agent/session_context/injected_context.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+module Agent
+  module SessionContext
+    InjectedContext = Data.define(:kind, :bytes, :occurrences, :source_refs, :text) do
+      def initialize(kind:, bytes:, occurrences:, source_refs:, text: nil)
+        normalized_kind = normalize_symbol(kind, :kind)
+        normalized_bytes = normalize_nonnegative_integer(bytes, :bytes)
+        normalized_occurrences = normalize_positive_integer(occurrences, :occurrences)
+        normalized_source_refs = normalize_source_refs(source_refs)
+        normalized_text = normalize_optional_string(text, :text)
+        validate_consistency!(
+          bytes: normalized_bytes,
+          occurrences: normalized_occurrences,
+          source_refs: normalized_source_refs,
+          text: normalized_text
+        )
+
+        super(
+          kind: normalized_kind,
+          bytes: normalized_bytes,
+          occurrences: normalized_occurrences,
+          source_refs: normalized_source_refs,
+          text: normalized_text
+        )
+      end
+
+      private
+
+      def normalize_symbol(value, name)
+        return value if value.is_a?(Symbol)
+        return value.to_sym if value.respond_to?(:to_sym)
+
+        raise TypeError, "#{name} must be symbolizable"
+      end
+
+      def normalize_nonnegative_integer(value, name)
+        integer = normalize_integer(value, name)
+        raise ArgumentError, "#{name} must be greater than or equal to 0" if integer.negative?
+
+        integer
+      end
+
+      def normalize_positive_integer(value, name)
+        integer = normalize_integer(value, name)
+        raise ArgumentError, "#{name} must be greater than or equal to 1" if integer < 1
+
+        integer
+      end
+
+      def normalize_integer(value, name)
+        integer =
+          if value.is_a?(Integer)
+            value
+          elsif value.respond_to?(:to_int)
+            value.to_int
+          elsif value.is_a?(String)
+            Integer(value, exception: false)
+          end
+
+        raise TypeError, "#{name} must be an Integer or integer-like value" if integer.nil?
+
+        integer
+      end
+
+      def normalize_source_refs(value)
+        Array(value).map do |source_ref|
+          unless source_ref.is_a?(SourceRef)
+            raise TypeError,
+                  "source_refs must contain only Agent::SessionContext::SourceRef values"
+          end
+
+          source_ref
+        end.freeze
+      end
+
+      def normalize_optional_string(value, name)
+        return if value.nil?
+        raise TypeError, "#{name} must be a String" unless value.respond_to?(:to_str)
+
+        String.new(value.to_str).freeze
+      end
+
+      def validate_consistency!(bytes:, occurrences:, source_refs:, text:)
+        raise ArgumentError, "occurrences must equal source_refs length" unless occurrences == source_refs.length
+        return if text.nil? || bytes == text.bytesize
+
+        raise ArgumentError, "bytes must equal text bytesize"
+      end
+    end
+  end
+end

data/lib/agent/session_context/injected_context_collector.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module Agent
+  module SessionContext
+    class InjectedContextCollector
+      def call(transcript, include_text: false)
+        validate_include_text!(include_text)
+        groups = {}
+
+        transcript.entries.each do |entry|
+          entry.parts.each do |part|
+            next unless part.injected
+
+            text = injected_text(part)
+            group = groups[text] ||= {
+              kind: Transcript.injection_kind(transcript.session.agent, text) || :provider_meta,
+              text:,
+              source_refs: []
+            }
+            group.fetch(:source_refs) << part.source_ref
+          end
+        end
+
+        groups.values.map do |group|
+          source_refs = group.fetch(:source_refs)
+          text = group.fetch(:text)
+          InjectedContext.new(
+            kind: group.fetch(:kind),
+            bytes: text.bytesize,
+            occurrences: source_refs.length,
+            source_refs:,
+            text: include_text ? text : nil
+          )
+        end.freeze
+      end
+
+      private
+
+      def validate_include_text!(value)
+        return if [true, false].include?(value)
+
+        raise ArgumentError, "include_text must be true or false"
+      end
+
+      def injected_text(part)
+        text = part.text
+        raise TypeError, "injected parts must contain text" unless text.respond_to?(:to_str)
+
+        text.to_str
+      end
+    end
+  end
+end

data/lib/agent/session_context/item.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+module Agent
+  module SessionContext
+    Item = Data.define(:kind, :label, :detail, :evidence, :source_refs, :attributes) do
+      EVIDENCE_VALUES = %i[observed explicit inferred].freeze
+
+      def initialize(kind:, label:, evidence:, source_refs:, detail: nil, attributes: {})
+        unless EVIDENCE_VALUES.include?(evidence)
+          raise ArgumentError,
+                "evidence must be one of: #{EVIDENCE_VALUES.join(", ")}"
+        end
+
+        super(
+          kind: kind.to_sym,
+          label: normalize_string(label, :label),
+          detail: normalize_optional_string(detail, :detail),
+          evidence: evidence,
+          source_refs: ImmutableValue.copy(Array(source_refs)),
+          attributes: normalize_attributes(attributes)
+        )
+      end
+
+      private
+
+      def normalize_attributes(attributes)
+        raise TypeError, "attributes must be a Hash" unless attributes.is_a?(Hash)
+
+        attributes.each_with_object({}) do |(key, value), normalized|
+          normalized[normalize_attribute_key(key)] = ImmutableValue.copy(value)
+        end.freeze
+      end
+
+      def normalize_attribute_key(key)
+        return key if key.is_a?(Symbol)
+        return key.to_sym if key.respond_to?(:to_sym)
+
+        raise TypeError, "attribute keys must be symbolizable"
+      end
+
+      def normalize_string(value, name)
+        raise TypeError, "#{name} must be a String" unless value.respond_to?(:to_str)
+
+        String.new(value.to_str).freeze
+      end
+
+      def normalize_optional_string(value, name)
+        return if value.nil?
+
+        normalize_string(value, name)
+      end
+    end
+  end
+end

data/lib/agent/session_context/prompt.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module Agent
+  module SessionContext
+    Prompt = Data.define(:index, :at, :text, :source_refs) do
+      def initialize(index:, at:, text:, source_refs:)
+        super(
+          index: normalize_positive_index(index),
+          at: at,
+          text: normalize_string(text, :text),
+          source_refs: ImmutableValue.copy(Array(source_refs))
+        )
+      end
+
+      private
+
+      def normalize_positive_index(value)
+        index =
+          if value.is_a?(Integer)
+            value
+          elsif value.respond_to?(:to_int)
+            value.to_int
+          elsif value.is_a?(String)
+            Integer(value, exception: false)
+          end
+
+        raise TypeError, "index must be an Integer or integer-like value" if index.nil?
+        raise ArgumentError, "index must be greater than or equal to 1" if index < 1
+
+        index
+      end
+
+      def normalize_string(value, name)
+        raise TypeError, "#{name} must be a String" unless value.respond_to?(:to_str)
+
+        String.new(value.to_str).freeze
+      end
+    end
+  end
+end

data/lib/agent/session_context/prompt_extractor.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Agent
+  module SessionContext
+    class PromptExtractor
+      def call(transcript)
+        prompt_index = 0
+
+        prompts = transcript.entries.each_with_object([]) do |entry, collected|
+          next unless entry.role == :user
+
+          contributing_parts = entry.parts.select { |part| part.type == :text && !part.injected }
+          next if contributing_parts.empty?
+
+          text = contributing_parts.map(&:text).join
+          next if text.empty?
+
+          prompt_index += 1
+          collected << Prompt.new(
+            index: prompt_index,
+            at: entry.at,
+            text: text,
+            source_refs: contributing_parts.map(&:source_ref)
+          )
+        end
+
+        prompts.freeze
+      end
+    end
+  end
+end