1 files changed, 24 insertions, 1 deletions
diff --git a/Omni/Agent/Engine.hs b/Omni/Agent/Engine.hs
index 69edd36..2704606 100644
--- a/Omni/Agent/Engine.hs
+++ b/Omni/Agent/Engine.hs
@@ -130,7 +130,30 @@ test =
       Test.unit "estimateCost calculates correctly" <| do
         let gpt4oCost = estimateCost "gpt-4o" 1000
             gpt4oMiniCost = estimateCost "gpt-4o-mini" 1000
-        (gpt4oCost >= gpt4oMiniCost) Test.@=? True
+        (gpt4oCost >= gpt4oMiniCost) Test.@=? True,
+      Test.unit "ToolCall JSON roundtrip" <| do
+        let tc =
+              ToolCall
+                { tcId = "call_123",
+                  tcType = "function",
+                  tcFunction = FunctionCall "read_file" "{\"path\":\"/tmp/test\"}"
+                }
+        case Aeson.decode (Aeson.encode tc) of
+          Nothing -> Test.assertFailure "Failed to decode ToolCall"
+          Just decoded -> tcId decoded Test.@=? "call_123",
+      Test.unit "FunctionCall JSON roundtrip" <| do
+        let fc = FunctionCall "test_func" "{\"arg\":\"value\"}"
+        case Aeson.decode (Aeson.encode fc) of
+          Nothing -> Test.assertFailure "Failed to decode FunctionCall"
+          Just decoded -> do
+            fcName decoded Test.@=? "test_func"
+            fcArguments decoded Test.@=? "{\"arg\":\"value\"}",
+      Test.unit "Role JSON roundtrip for all roles" <| do
+        let roles = [System, User, Assistant, ToolRole]
+        forM_ roles <| \role ->
+          case Aeson.decode (Aeson.encode role) of
+            Nothing -> Test.assertFailure ("Failed to decode Role: " <> show role)
+            Just decoded -> decoded Test.@=? role
     ]
 
 data Tool = Tool