diff --git a/lib/roast/cogs/agent/providers/claude/message.rb b/lib/roast/cogs/agent/providers/claude/message.rb index e51b084f..739f372c 100644 --- a/lib/roast/cogs/agent/providers/claude/message.rb +++ b/lib/roast/cogs/agent/providers/claude/message.rb @@ -22,26 +22,24 @@ def from_json(json, raw_dump_file: nil) #: (Hash[Symbol, untyped]) -> Message def from_hash(hash) type = hash.delete(:type)&.to_sym - case type - when :assistant - Messages::AssistantMessage.new(type:, hash:) - when :result - Messages::ResultMessage.new(type:, hash:) - when :system - Messages::SystemMessage.new(type:, hash:) - when :text - Messages::TextMessage.new(type:, hash:) - when :thinking - Messages::ThinkingMessage.new(type:, hash:) - when :tool_result - Messages::ToolResultMessage.new(type:, hash:) - when :tool_use - Messages::ToolUseMessage.new(type:, hash:) - when :user - Messages::UserMessage.new(type:, hash:) + message_class = resolve_message_class(type) + message_class.new(type:, hash:) + end + + private + + #: (Symbol?) -> singleton(Message) + def resolve_message_class(type) + return Messages::UnknownMessage if type.nil? + + class_name = "#{type}_message".camelize + if Messages.const_defined?(class_name, false) + Messages.const_get(class_name, false) # rubocop:disable Sorbet/ConstantsFromStrings else - Messages::UnknownMessage.new(type:, hash:) + Messages::UnknownMessage end + rescue NameError + Messages::UnknownMessage end end diff --git a/test/roast/cogs/agent/providers/claude/message_test.rb b/test/roast/cogs/agent/providers/claude/message_test.rb index 3e2abdd5..070141ba 100644 --- a/test/roast/cogs/agent/providers/claude/message_test.rb +++ b/test/roast/cogs/agent/providers/claude/message_test.rb @@ -149,6 +149,13 @@ class MessageTest < ActiveSupport::TestCase assert_kind_of Messages::UnknownMessage, message end + test "from_hash with invalid constant name type creates UnknownMessage" do + hash = { type: :"some-hyphenated-type", foo: "bar" } + message = Message.from_hash(hash) + + assert_kind_of Messages::UnknownMessage, message + end + test "from_hash removes type from hash" do hash = { type: :text, text: "Hello" } Message.from_hash(hash)