Internal.hs

{-# LANGUAGE OverloadedStrings #-}

-- | Runtime support for the code generated by "MCP.Server.Derive".
--
-- Two families of parsers exist because the MCP specification types the two
-- argument channels differently: prompt arguments are always strings, while
-- tool arguments are full JSON values. The Value-based parsers are lenient:
-- they accept the native JSON type and fall back to parsing the string
-- representation, since many clients send numbers and booleans as strings.
module MCP.Server.Derive.Internal
  ( -- * Text-based parsers (prompt arguments)
    parseText
  , parseInt
  , parseInteger
  , parseDouble
  , parseFloat
  , parseBool
    -- * Value-based parsers (tool arguments)
  , valueText
  , valueInt
  , valueInteger
  , valueDouble
  , valueFloat
  , valueBool
  , valueArray
  , renderValue
    -- * Argument/field extraction
  , requiredArg
  , optionalArg
  , requiredTextArg
  , optionalTextArg
  , requiredField
  , optionalField
    -- * Resource template matching
  , matchTemplateSegments
  , templateField
    -- * Structured output
  , resolveToolOutput
  , omitNothing
  , canonicalJsonText
  ) where

import           Data.Aeson           (ToJSON (..), Value (..), encode)
import qualified Data.Aeson.Key       as Key
import qualified Data.Aeson.KeyMap    as KM
import           Data.Aeson.Types     (Pair)
import qualified Data.ByteString.Lazy as BSL
import           Data.Foldable        (toList)
import           Data.Map             (Map)
import qualified Data.Map             as Map
import           Data.Scientific      (base10Exponent, floatingOrInteger,
                                       toBoundedInteger, toRealFloat)
import           Data.Text            (Text)
import qualified Data.Text            as T
import qualified Data.Text.Encoding   as TE
import           Network.URI          (unEscapeString)
import           Text.Read            (readMaybe)

import           MCP.Server.Types     (Content (..), Error (..),
                                       ToolOutput (..), ToolResult (..),
                                       toolError, toolResult)

-- ---------------------------------------------------------------------------
-- Text-based parsers (prompt arguments are strings per the MCP spec)
-- ---------------------------------------------------------------------------

parseText :: Text -> Either Text Text
parseText = Right

parseInt :: Text -> Either Text Int
parseInt t = case readMaybe (T.unpack t) of
  Just v  -> Right v
  Nothing -> Left $ "Failed to parse Int from: " <> t

parseInteger :: Text -> Either Text Integer
parseInteger t = case readMaybe (T.unpack t) of
  Just v  -> Right v
  Nothing -> Left $ "Failed to parse Integer from: " <> t

parseDouble :: Text -> Either Text Double
parseDouble t = case readMaybe (T.unpack t) of
  Just v  -> Right v
  Nothing -> Left $ "Failed to parse Double from: " <> t

parseFloat :: Text -> Either Text Float
parseFloat t = case readMaybe (T.unpack t) of
  Just v  -> Right v
  Nothing -> Left $ "Failed to parse Float from: " <> t

parseBool :: Text -> Either Text Bool
parseBool t = case T.toLower t of
  "true"  -> Right True
  "false" -> Right False
  _       -> Left $ "Failed to parse Bool from: " <> t

-- ---------------------------------------------------------------------------
-- Value-based parsers (tool arguments are full JSON values)
-- ---------------------------------------------------------------------------

-- | A short rendering of a JSON value for error messages.
renderValue :: Value -> Text
renderValue (String t) = t
renderValue v          = TE.decodeUtf8 $ BSL.toStrict $ encode v

valueText :: Value -> Either Text Text
valueText (String t) = Right t
valueText v          = Left $ "Failed to parse Text from: " <> renderValue v

valueInt :: Value -> Either Text Int
valueInt (Number n) = case toBoundedInteger n of
  Just i  -> Right i
  Nothing -> Left $ "Failed to parse Int from: " <> renderValue (Number n)
valueInt (String t) = parseInt t
valueInt v          = Left $ "Failed to parse Int from: " <> renderValue v

valueInteger :: Value -> Either Text Integer
valueInteger (Number n)
  -- Bound the exponent before materializing the Integer: aeson keeps huge
  -- exponents compact, so a tiny payload like 1e1000000000 would otherwise
  -- allocate a billion-digit Integer (same bound aeson itself uses for its
  -- FromJSON Integer instance).
  | base10Exponent n > 1024 =
      Left "Failed to parse Integer: exponent too large"
  | otherwise = case floatingOrInteger n :: Either Double Integer of
      Right i -> Right i
      Left _  -> Left $ "Failed to parse Integer from: " <> renderValue (Number n)
valueInteger (String t) = parseInteger t
valueInteger v          = Left $ "Failed to parse Integer from: " <> renderValue v

valueDouble :: Value -> Either Text Double
valueDouble (Number n) = Right (toRealFloat n)
valueDouble (String t) = parseDouble t
valueDouble v          = Left $ "Failed to parse Double from: " <> renderValue v

valueFloat :: Value -> Either Text Float
valueFloat (Number n) = Right (toRealFloat n)
valueFloat (String t) = parseFloat t
valueFloat v          = Left $ "Failed to parse Float from: " <> renderValue v

valueBool :: Value -> Either Text Bool
valueBool (Bool b)   = Right b
valueBool (String t) = parseBool t
valueBool v          = Left $ "Failed to parse Bool from: " <> renderValue v

valueArray :: (Value -> Either Text a) -> Value -> Either Text [a]
valueArray p (Array xs) = traverse p (toList xs)
valueArray _ v          = Left $ "Failed to parse array from: " <> renderValue v

-- ---------------------------------------------------------------------------
-- Argument extraction from the top-level arguments map
-- ---------------------------------------------------------------------------

prefixed :: Text -> Either Text a -> Either Error a
prefixed name = either (Left . InvalidParams . (("field '" <> name <> "': ") <>)) Right

-- | A required tool argument: missing yields 'MissingRequiredParams', a parse
-- failure yields 'InvalidParams' with the field name prefixed.
requiredArg :: Text -> (Value -> Either Text a) -> Map Text Value -> Either Error a
requiredArg name p args = case Map.lookup name args of
  Nothing -> Left $ MissingRequiredParams $ "field '" <> name <> "' is missing"
  Just v  -> prefixed name (p v)

-- | An optional tool argument: absent (or JSON null) becomes 'Nothing'.
optionalArg :: Text -> (Value -> Either Text a) -> Map Text Value -> Either Error (Maybe a)
optionalArg name p args = case Map.lookup name args of
  Nothing   -> Right Nothing
  Just Null -> Right Nothing
  Just v    -> Just <$> prefixed name (p v)

-- | A required prompt argument (string-valued).
requiredTextArg :: Text -> (Text -> Either Text a) -> Map Text Text -> Either Error a
requiredTextArg name p args = case Map.lookup name args of
  Nothing -> Left $ MissingRequiredParams $ "field '" <> name <> "' is missing"
  Just t  -> prefixed name (p t)

-- | An optional prompt argument (string-valued).
optionalTextArg :: Text -> (Text -> Either Text a) -> Map Text Text -> Either Error (Maybe a)
optionalTextArg name p args = case Map.lookup name args of
  Nothing -> Right Nothing
  Just t  -> Just <$> prefixed name (p t)

-- ---------------------------------------------------------------------------
-- Field extraction from nested JSON objects
-- ---------------------------------------------------------------------------

-- | A required field of a nested object argument.
requiredField :: Text -> (Value -> Either Text a) -> KM.KeyMap Value -> Either Text a
requiredField name p o = case KM.lookup (Key.fromText name) o of
  Nothing -> Left $ "field '" <> name <> "' is missing"
  Just v  -> either (Left . (("field '" <> name <> "': ") <>)) Right (p v)

-- | An optional field of a nested object argument: absent or null is 'Nothing'.
optionalField :: Text -> (Value -> Either Text a) -> KM.KeyMap Value -> Either Text (Maybe a)
optionalField name p o = case KM.lookup (Key.fromText name) o of
  Nothing   -> Right Nothing
  Just Null -> Right Nothing
  Just v    -> either (Left . (("field '" <> name <> "': ") <>)) (Right . Just) (p v)

-- ---------------------------------------------------------------------------
-- Resource template matching
-- ---------------------------------------------------------------------------

-- | Match a URI against a template of the form @\<prefix\>{a}\/{b}\/…@: the
-- rendered URI must start with the prefix and continue with exactly @n@
-- non-empty, slash-separated segments, which are returned percent-decoded.
-- Any query or fragment on the URI is ignored — templates declare only path
-- variables (a literal @?@ or @#@ inside a segment value arrives
-- percent-encoded and is unaffected).
matchTemplateSegments :: Text -> Int -> String -> Maybe [Text]
matchTemplateSegments prefix n uriStr = do
  let base = T.takeWhile (\c -> c /= '?' && c /= '#') (T.pack uriStr)
  rest <- T.stripPrefix prefix base
  let segs = T.splitOn "/" rest
  if length segs == n && all (not . T.null) segs
    then Just (map (T.pack . unEscapeString . T.unpack) segs)
    else Nothing

-- | Decode one matched template segment into a field value.
templateField :: Text -> (Text -> Either Text a) -> [Text] -> Int -> Either Error a
templateField name p segs i = case drop i segs of
  (s:_) -> either (Left . InvalidParams . (("template field '" <> name <> "': ") <>)) Right (p s)
  []    -> Left $ InvalidParams $ "missing segment for template field '" <> name <> "'"

-- ---------------------------------------------------------------------------
-- Structured output
-- ---------------------------------------------------------------------------

-- | Turn a 'ToolOutput' into the 'ToolResult' that goes on the wire, given
-- the generated serializer for the output type. A plain 'ToolOutput' also
-- carries the serialized JSON as a text content block, per the spec's
-- recommendation for clients that predate structured output.
resolveToolOutput :: (o -> Value) -> ToolOutput o -> ToolResult
resolveToolOutput serialize out = case out of
  ToolOutput o ->
    let v = serialize o
    in (toolResult [ContentText (canonicalJsonText v)]) { toolResultStructured = Just v }
  ToolOutputWith content o ->
    (toolResult content) { toolResultStructured = Just (serialize o) }
  ToolOutputError msg -> toolError msg
  ToolOutputRaw result -> result

-- | Compact JSON with objects in sorted-key order, independent of the
-- aeson\/hashable pair in the build plan. Used wherever JSON ends up
-- /inside a string/ (the structured-output text block): structural
-- comparison cannot see through a string, so those bytes must be
-- deterministic — which also keeps repeated tool results stable for client
-- prompt caching.
canonicalJsonText :: Value -> Text
canonicalJsonText = TE.decodeUtf8 . BSL.toStrict . encode . Canonical

-- Encoding-only wrapper: objects are emitted via a sorted 'Map', arrays
-- and atoms recurse. 'encode' uses 'toEncoding', so no KeyMap is rebuilt.
newtype Canonical = Canonical Value

instance ToJSON Canonical where
  toJSON (Canonical v) = v
  toEncoding (Canonical (Object o)) =
    toEncoding (Map.fromList [(Key.toText k, Canonical v) | (k, v) <- KM.toList o])
  toEncoding (Canonical (Array xs)) = toEncoding (fmap Canonical xs)
  toEncoding (Canonical v) = toEncoding v

-- | A record field for the generated serializers: absent when 'Nothing',
-- mirroring how the generated schema marks 'Maybe' fields optional.
omitNothing :: Text -> (a -> Value) -> Maybe a -> [Pair]
omitNothing _ _ Nothing        = []
omitNothing name f (Just x)    = [(Key.fromText name, f x)]