{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Hermes gateway api_server client: one persistent server-side session as
-- a 'Host' seat.
--
-- The api_server (part of @hermes gateway run@) offers @POST \/api\/sessions@
-- and @POST \/api\/sessions\/{id}\/chat\/stream@ — SSE event streams with
-- protocol-posted halt marks (@run.completed@, @done@).  Framing is the
-- 'frameAgent' mark machine at 'sseMarks': the blank line ends an event
-- frame, the event names are the halt alphabet.  No polling, no quiet
-- inference — the server posts its own marks.
--
-- Auth is a Bearer token read from the environment variable named by
-- 'gwKeyEnv' (default @API_SERVER_KEY@, the platform's own convention).
--
-- Probe card: @coffee\/loom\/next-yin.md@ ("hermes server (api_server :8642)").
module Free.Agent.Gateway
  ( -- * Configuration
    GatewayConfig (..),
    defaultGatewayConfig,

    -- * Client
    GatewayClient (..),
    openGateway,

    -- * Turns
    gatewayChat,

    -- * Host seat
    gatewayHost,

    -- * SSE framing (exposed for pure tests)
    parseSseFrame,
  )
where

import Circuit.Agent (run1)
import Circuit.Agent.StdPorts (frameAgent, sseMarks)
import Circuit.Parser.Json (Json (..), decodeJson, encodeJson)
import Control.Exception (throwIO)
import Data.ByteString qualified as BS
import Data.ByteString.Char8 qualified as BSC
import Data.ByteString.Lazy qualified as BL
import Data.Foldable (foldl')
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding (decodeUtf8, encodeUtf8)
import Free.Agent.Host (BodyMode (..), Host (..))
import Free.Agent.Json (jobject, jtext)
import Network.HTTP.Client
import Network.HTTP.Types (RequestHeaders, statusCode)
import System.Environment (lookupEnv)
import Prelude

-- ---------------------------------------------------------------------------
-- Configuration
-- ---------------------------------------------------------------------------

-- | Where the api_server is and how to authenticate to it.
data GatewayConfig = GatewayConfig
  { -- | Base URL, e.g. @http:\/\/127.0.0.1:8642@.
    GatewayConfig -> Text
gwBaseUrl :: Text,
    -- | Environment variable holding the Bearer key.
    GatewayConfig -> String
gwKeyEnv :: String,
    -- | Working directory pinned on the server-side session.
    GatewayConfig -> String
gwCwd :: FilePath
  }
  deriving (Int -> GatewayConfig -> ShowS
[GatewayConfig] -> ShowS
GatewayConfig -> String
(Int -> GatewayConfig -> ShowS)
-> (GatewayConfig -> String)
-> ([GatewayConfig] -> ShowS)
-> Show GatewayConfig
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> GatewayConfig -> ShowS
showsPrec :: Int -> GatewayConfig -> ShowS
$cshow :: GatewayConfig -> String
show :: GatewayConfig -> String
$cshowList :: [GatewayConfig] -> ShowS
showList :: [GatewayConfig] -> ShowS
Show, GatewayConfig -> GatewayConfig -> Bool
(GatewayConfig -> GatewayConfig -> Bool)
-> (GatewayConfig -> GatewayConfig -> Bool) -> Eq GatewayConfig
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: GatewayConfig -> GatewayConfig -> Bool
== :: GatewayConfig -> GatewayConfig -> Bool
$c/= :: GatewayConfig -> GatewayConfig -> Bool
/= :: GatewayConfig -> GatewayConfig -> Bool
Eq)

defaultGatewayConfig :: GatewayConfig
defaultGatewayConfig :: GatewayConfig
defaultGatewayConfig =
  GatewayConfig
    { gwBaseUrl :: Text
gwBaseUrl = Text
"http://127.0.0.1:8642/p/pit",
      gwKeyEnv :: String
gwKeyEnv = String
"API_SERVER_KEY",
      gwCwd :: String
gwCwd = String
"/Users/tonyday567/pit"
    }

-- ---------------------------------------------------------------------------
-- Client
-- ---------------------------------------------------------------------------

-- | A live api_server session.
data GatewayClient = GatewayClient
  { GatewayClient -> GatewayConfig
gcConfig :: GatewayConfig,
    GatewayClient -> Manager
gcManager :: Manager,
    GatewayClient -> ByteString
gcKey :: BS.ByteString,
    GatewayClient -> Text
gcSessionId :: Text
  }

-- | Read the Bearer key from 'gwKeyEnv' and create a server-side session.
openGateway :: GatewayConfig -> IO GatewayClient
openGateway :: GatewayConfig -> IO GatewayClient
openGateway GatewayConfig
cfg = do
  mKey <- String -> IO (Maybe String)
lookupEnv (GatewayConfig -> String
gwKeyEnv GatewayConfig
cfg)
  key <- case mKey of
    Maybe String
Nothing -> IOError -> IO ByteString
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (String -> IOError
userError (String
"🔴 " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> GatewayConfig -> String
gwKeyEnv GatewayConfig
cfg String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" environment variable not set"))
    Just String
k -> ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> ByteString
BSC.pack String
k)
  mgr <- newManager defaultManagerSettings {managerResponseTimeout = responseTimeoutNone}
  req0 <- parseRequest (T.unpack (gwBaseUrl cfg) <> "/api/sessions")
  let req =
        Request
req0
          { method = "POST",
            requestHeaders = hdrs key,
            requestBody = RequestBodyLBS (BL.fromStrict (encodeJson (jobject [("cwd", jtext (T.pack (gwCwd cfg)))])))
          }
  resp <- httpLbs req mgr
  case decodeJson (BL.toStrict (responseBody resp)) of
    Left String
e -> IOError -> IO GatewayClient
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (String -> IOError
userError (String
"🔴 api_server session create: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
e))
    Right Json
v -> do
      let sid :: Text
sid = [Text] -> Json -> Text
textAt [Text
"session", Text
"id"] Json
v
      if Text -> Bool
T.null Text
sid
        then IOError -> IO GatewayClient
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (String -> IOError
userError String
"🔴 api_server session create: no session.id in response")
        else GatewayClient -> IO GatewayClient
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (GatewayConfig -> Manager -> ByteString -> Text -> GatewayClient
GatewayClient GatewayConfig
cfg Manager
mgr ByteString
key Text
sid)

-- ---------------------------------------------------------------------------
-- Turns
-- ---------------------------------------------------------------------------

-- | One turn: POST @chat\/stream@, consume SSE events until the @done@
-- mark, return the completed reply text.
--
-- The reply is the @assistant.completed@ content when it arrives, else the
-- accumulated @assistant.delta@ stream.  HTTP and stream failures come back
-- as 🔴-prefixed text (the 'Free.Agent.Host.chatCompletion' convention), so
-- a failing gateway surfaces as a bus post, not a crash.
gatewayChat :: GatewayClient -> Text -> IO Text
gatewayChat :: GatewayClient -> Text -> IO Text
gatewayChat GatewayClient
c Text
msg = do
  req0 <- String -> IO Request
forall (m :: * -> *). MonadThrow m => String -> m Request
parseRequest (Text -> String
T.unpack (GatewayConfig -> Text
gwBaseUrl (GatewayClient -> GatewayConfig
gcConfig GatewayClient
c)) String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"/api/sessions/" String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack (GatewayClient -> Text
gcSessionId GatewayClient
c) String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"/chat/stream")
  let req =
        Request
req0
          { method = "POST",
            requestHeaders = hdrs (gcKey c),
            requestBody = RequestBodyLBS (BL.fromStrict (encodeJson (jobject [("message", jtext msg)])))
          }
  withResponse req (gcManager c) $ \Response (IO ByteString)
resp -> do
    let status :: Int
status = Status -> Int
statusCode (Response (IO ByteString) -> Status
forall body. Response body -> Status
responseStatus Response (IO ByteString)
resp)
    if Int
status Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
200 Bool -> Bool -> Bool
|| Int
status Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
300
      then do
        body <- IO ByteString -> Int -> IO ByteString
brReadSome (Response (IO ByteString) -> IO ByteString
forall body. Response body -> body
responseBody Response (IO ByteString)
resp) Int
4096
        pure ("🔴 HTTP " <> T.pack (show status) <> ": " <> decodeUtf8 (BL.toStrict body))
      else IO ByteString
-> (ByteString, [ByteString]) -> (Text, Maybe Text) -> IO Text
go (Response (IO ByteString) -> IO ByteString
forall body. Response body -> body
responseBody Response (IO ByteString)
resp) (ByteString
BS.empty, []) (Text, Maybe Text)
forall a. Monoid a => a
mempty
  where
    sys :: Agent
  (->) (ByteString, [ByteString]) (Maybe ByteString) [ByteString]
sys = ProcMarks
-> (ByteString -> ByteString)
-> Agent
     (->) (ByteString, [ByteString]) (Maybe ByteString) [ByteString]
forall a.
ProcMarks
-> (ByteString -> a)
-> Agent (->) (ByteString, [a]) (Maybe ByteString) [a]
frameAgent ProcMarks
sseMarks ByteString -> ByteString
forall a. a -> a
id

    go :: IO ByteString
-> (ByteString, [ByteString]) -> (Text, Maybe Text) -> IO Text
go IO ByteString
body (ByteString, [ByteString])
st (Text
deltas, Maybe Text
completed) = do
      chunk <- IO ByteString -> IO ByteString
brRead IO ByteString
body
      let input = if ByteString -> Bool
BS.null ByteString
chunk then Maybe ByteString
forall a. Maybe a
Nothing else ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
chunk
          (frames, st') = run1 sys st input
          evs = (ByteString -> Maybe (Text, Json))
-> [ByteString] -> [(Text, Json)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ByteString -> Maybe (Text, Json)
parseSseFrame [ByteString]
frames
          (deltas', completed') = foldl' step (deltas, completed) evs
      if any ((== "done") . fst) evs || BS.null chunk
        then pure (finish (deltas', completed'))
        else go body st' (deltas', completed')

    step :: (Text, Maybe Text) -> (a, Json) -> (Text, Maybe Text)
step (Text
deltas, Maybe Text
completed) (a
ev, Json
v) = case a
ev of
      a
"assistant.delta" -> (Text
deltas Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Text] -> Json -> Text
textAt [Text
"delta"] Json
v, Maybe Text
completed)
      a
"assistant.completed" -> (Text
deltas, Text -> Maybe Text
forall a. a -> Maybe a
Just ([Text] -> Json -> Text
textAt [Text
"content"] Json
v))
      a
"run.failed" -> (Text
deltas, Text -> Maybe Text
forall a. a -> Maybe a
Just (Text
"🔴 run.failed: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Text] -> Json -> Text
textAt [Text
"error"] Json
v))
      a
_ -> (Text
deltas, Maybe Text
completed)

    finish :: (Text, Maybe Text) -> Text
finish (Text
deltas, Maybe Text
completed) =
      let r :: Text
r = Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe Text
deltas Maybe Text
completed
       in if Text -> Bool
T.null (Text -> Text
T.strip Text
r) then Text
"🔴 empty reply stream" else Text
r

-- ---------------------------------------------------------------------------
-- Host seat
-- ---------------------------------------------------------------------------

-- | A 'Host' backed by the gateway session.  All post bodies of a wake
-- cycle join into one user message (the 'Free.Agent.Host.hermesHostBatch'
-- convention); the reply is one text.
gatewayHost ::
  -- | Host name (used as the 'from' field of reply posts).
  Text ->
  -- | System prompt text prepended to every message.
  Text ->
  GatewayClient ->
  Host
gatewayHost :: Text -> Text -> GatewayClient -> Host
gatewayHost Text
name Text
systemPrompt GatewayClient
c =
  Host
    { hostName :: Text
hostName = Text
name,
      hostBodyMode :: BodyMode
hostBodyMode = BodyMode
BodyWhole,
      hostRun :: [Text] -> IO [Text]
hostRun = \[Text]
bodies ->
        (Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: []) (Text -> [Text]) -> IO Text -> IO [Text]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GatewayClient -> Text -> IO Text
gatewayChat GatewayClient
c (Text
systemPrompt Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n\nUser messages:\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Text] -> Text
T.unlines [Text]
bodies)
    }

-- ---------------------------------------------------------------------------
-- SSE framing
-- ---------------------------------------------------------------------------

-- | Parse one SSE event frame (@event: x\\ndata: {…}@) into the event name
-- and decoded JSON payload.  'Nothing' for comment\/keep-alive frames.
--
-- >>> parseSseFrame (Data.ByteString.Char8.pack "event: done\ndata: {\"seq\": 7}")
-- Just ("done",JObject [("seq",JNumber 7.0)])
parseSseFrame :: BS.ByteString -> Maybe (Text, Json)
parseSseFrame :: ByteString -> Maybe (Text, Json)
parseSseFrame ByteString
f = do
  let ls :: [ByteString]
ls = ByteString -> [ByteString]
BSC.lines ByteString
f
  ev <- [ByteString] -> Maybe ByteString
forall {a}. [a] -> Maybe a
listToMaybe [Int -> ByteString -> ByteString
BSC.drop Int
6 ByteString
l | ByteString
l <- [ByteString]
ls, ByteString
"event:" ByteString -> ByteString -> Bool
`BS.isPrefixOf` ByteString
l]
  dat <- listToMaybe [BSC.drop 5 l | l <- ls, "data:" `BS.isPrefixOf` l]
  v <- either (const Nothing) Just (decodeJson (BSC.dropWhile (== ' ') dat))
  pure (decodeUtf8 (BSC.dropWhile (== ' ') ev), v)
  where
    listToMaybe :: [a] -> Maybe a
listToMaybe [] = Maybe a
forall a. Maybe a
Nothing
    listToMaybe (a
x : [a]
_) = a -> Maybe a
forall a. a -> Maybe a
Just a
x

-- ---------------------------------------------------------------------------
-- Internal helpers
-- ---------------------------------------------------------------------------

hdrs :: BS.ByteString -> RequestHeaders
hdrs :: ByteString -> RequestHeaders
hdrs ByteString
key = [(HeaderName
"Authorization", ByteString
"Bearer " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
key), (HeaderName
"Content-Type", ByteString
"application/json")]

textAt :: [Text] -> Json -> Text
textAt :: [Text] -> Json -> Text
textAt [Text]
path Json
v = [Text] -> Json -> Text
go [Text]
path Json
v
  where
    go :: [Text] -> Json -> Text
go [] (JString Text
t) = Text
t
    go [] Json
_ = Text
""
    go (Text
p : [Text]
ps) (JObject [(Text, Json)]
o) = Text -> (Json -> Text) -> Maybe Json -> Text
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Text
"" ([Text] -> Json -> Text
go [Text]
ps) (Text -> [(Text, Json)] -> Maybe Json
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Text
p [(Text, Json)]
o)
    go [Text]
_ Json
_ = Text
""