module Unison.MCP.Cli
( handleInputMCP,
ppForProjectContext,
cliToMCP,
)
where
import Control.Monad.Except (ExceptT (..), throwError)
import Control.Monad.Reader
import Crypto.Random qualified as Random
import Data.Aeson
import Data.IORef
import Data.Sequence qualified as Seq
import Data.Text qualified as Text
import U.Codebase.Sqlite.Queries qualified as Queries
import Unison.Auth.CredentialManager qualified as AuthN
import Unison.Auth.HTTPClient qualified as AuthN
import Unison.Auth.Tokens qualified as AuthN
import Unison.Cli.Monad qualified as Cli
import Unison.Codebase qualified as Codebase
import Unison.Codebase.Editor.HandleInput qualified as HandleInput
import Unison.Codebase.Editor.Input (Event, Input)
import Unison.Codebase.Path qualified as Path
import Unison.Codebase.ProjectPath qualified as PP
import Unison.CommandLine.OutputMessages qualified as Output
import Unison.MCP.Types
import Unison.MCP.Types qualified as MCP
import Unison.Prelude
import Unison.Sqlite (Transaction)
import Unison.Syntax.Parser qualified as Parser
import Unison.Util.Pretty qualified as Pretty
import UnliftIO.STM
import Prelude hiding (readFile, writeFile)
data CliOutput = CliOutput
{ CliOutput -> [Text]
sourceCodeUpdates :: [Text],
CliOutput -> [Text]
outputMessages :: [Text]
}
deriving (CliOutput -> CliOutput -> Bool
(CliOutput -> CliOutput -> Bool)
-> (CliOutput -> CliOutput -> Bool) -> Eq CliOutput
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: CliOutput -> CliOutput -> Bool
== :: CliOutput -> CliOutput -> Bool
$c/= :: CliOutput -> CliOutput -> Bool
/= :: CliOutput -> CliOutput -> Bool
Eq, Int -> CliOutput -> ShowS
[CliOutput] -> ShowS
CliOutput -> String
(Int -> CliOutput -> ShowS)
-> (CliOutput -> String)
-> ([CliOutput] -> ShowS)
-> Show CliOutput
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CliOutput -> ShowS
showsPrec :: Int -> CliOutput -> ShowS
$cshow :: CliOutput -> String
show :: CliOutput -> String
$cshowList :: [CliOutput] -> ShowS
showList :: [CliOutput] -> ShowS
Show)
instance Semigroup CliOutput where
CliOutput [Text]
src1 [Text]
out1 <> :: CliOutput -> CliOutput -> CliOutput
<> CliOutput [Text]
src2 [Text]
out2 =
[Text] -> [Text] -> CliOutput
CliOutput ([Text]
src1 [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> [Text]
src2) ([Text]
out1 [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> [Text]
out2)
instance Monoid CliOutput where
mempty :: CliOutput
mempty = [Text] -> [Text] -> CliOutput
CliOutput [] []
instance ToJSON CliOutput where
toJSON :: CliOutput -> Value
toJSON (CliOutput [Text]
sourceCodeUpdates [Text]
outputMessages) =
[Pair] -> Value
object
[ Key
"sourceCodeUpdates" Key -> [Text] -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= [Text]
sourceCodeUpdates,
Key
"outputMessages" Key -> [Text] -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= [Text]
outputMessages
]
ppForProjectContext :: ProjectContext -> ExceptT Text Transaction PP.ProjectPath
ppForProjectContext :: ProjectContext -> ExceptT Text Transaction ProjectPath
ppForProjectContext ProjectContext {ProjectName
projectName :: ProjectName
$sel:projectName:ProjectContext :: ProjectContext -> ProjectName
projectName, ProjectBranchName
branchName :: ProjectBranchName
$sel:branchName:ProjectContext :: ProjectContext -> ProjectBranchName
branchName} = do
Project
project <-
Transaction (Maybe Project)
-> ExceptT Text Transaction (Maybe Project)
forall (m :: * -> *) a. Monad m => m a -> ExceptT Text m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ProjectName -> Transaction (Maybe Project)
Queries.loadProjectByName ProjectName
projectName) ExceptT Text Transaction (Maybe Project)
-> (ExceptT Text Transaction (Maybe Project)
-> ExceptT Text Transaction Project)
-> ExceptT Text Transaction Project
forall a b. a -> (a -> b) -> b
& ExceptT Text Transaction Project
-> ExceptT Text Transaction (Maybe Project)
-> ExceptT Text Transaction Project
forall (m :: * -> *) a. Monad m => m a -> m (Maybe a) -> m a
onNothingM do
Text -> ExceptT Text Transaction Project
forall a. Text -> ExceptT Text Transaction a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Text -> ExceptT Text Transaction Project)
-> Text -> ExceptT Text Transaction Project
forall a b. (a -> b) -> a -> b
$ Text
"Project not found: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> forall target source. From source target => source -> target
into @Text ProjectName
projectName
ProjectBranch
branch <-
Transaction (Maybe ProjectBranch)
-> ExceptT Text Transaction (Maybe ProjectBranch)
forall (m :: * -> *) a. Monad m => m a -> ExceptT Text m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ProjectId -> ProjectBranchName -> Transaction (Maybe ProjectBranch)
Queries.loadProjectBranchByName Project
project.projectId ProjectBranchName
branchName) ExceptT Text Transaction (Maybe ProjectBranch)
-> (Maybe ProjectBranch -> ExceptT Text Transaction ProjectBranch)
-> ExceptT Text Transaction ProjectBranch
forall a b.
ExceptT Text Transaction a
-> (a -> ExceptT Text Transaction b) -> ExceptT Text Transaction b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
Maybe ProjectBranch
Nothing -> Text -> ExceptT Text Transaction ProjectBranch
forall a. Text -> ExceptT Text Transaction a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError (Text -> ExceptT Text Transaction ProjectBranch)
-> Text -> ExceptT Text Transaction ProjectBranch
forall a b. (a -> b) -> a -> b
$ Text
"Branch not found: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> forall target source. From source target => source -> target
into @Text ProjectBranchName
branchName
Just ProjectBranch
projectBranch -> ProjectBranch -> ExceptT Text Transaction ProjectBranch
forall a. a -> ExceptT Text Transaction a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ProjectBranch
projectBranch
ProjectPath -> ExceptT Text Transaction ProjectPath
forall a. a -> ExceptT Text Transaction a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ProjectPath -> ExceptT Text Transaction ProjectPath)
-> ProjectPath -> ExceptT Text Transaction ProjectPath
forall a b. (a -> b) -> a -> b
$ ProjectAndBranch Project ProjectBranch -> Absolute -> ProjectPath
PP.fromProjectAndBranch (Project -> ProjectBranch -> ProjectAndBranch Project ProjectBranch
forall a b. a -> b -> ProjectAndBranch a b
PP.ProjectAndBranch Project
project ProjectBranch
branch) Absolute
Path.Root
handleInputMCP :: ProjectContext -> [Either Event Input] -> ExceptT Text MCP CliOutput
handleInputMCP :: ProjectContext
-> [Either Event Input] -> ExceptT Text MCP CliOutput
handleInputMCP ProjectContext
projectContext [Either Event Input]
input = do
case [Either Event Input]
input of
(Either Event Input
inp : [Either Event Input]
rest) -> do
(Maybe ()
_, CliOutput
cliOutput) <- ProjectContext -> Cli () -> ExceptT Text MCP (Maybe (), CliOutput)
forall a.
ProjectContext -> Cli a -> ExceptT Text MCP (Maybe a, CliOutput)
cliToMCP ProjectContext
projectContext (Either Event Input -> Cli ()
HandleInput.loop Either Event Input
inp)
(CliOutput
cliOutput CliOutput -> CliOutput -> CliOutput
forall a. Semigroup a => a -> a -> a
<>) (CliOutput -> CliOutput)
-> ExceptT Text MCP CliOutput -> ExceptT Text MCP CliOutput
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ProjectContext
-> [Either Event Input] -> ExceptT Text MCP CliOutput
handleInputMCP ProjectContext
projectContext [Either Event Input]
rest
[] -> CliOutput -> ExceptT Text MCP CliOutput
forall a. a -> ExceptT Text MCP a
forall (f :: * -> *) a. Applicative f => a -> f a
pure CliOutput
forall a. Monoid a => a
mempty
cliToMCP :: ProjectContext -> Cli.Cli a -> ExceptT Text MCP (Maybe a, CliOutput)
cliToMCP :: forall a.
ProjectContext -> Cli a -> ExceptT Text MCP (Maybe a, CliOutput)
cliToMCP ProjectContext
projCtx Cli a
cli = do
MCP.Env {Text
ucmVersion :: Text
$sel:ucmVersion:Env :: Env -> Text
ucmVersion, Codebase IO Symbol Ann
codebase :: Codebase IO Symbol Ann
$sel:codebase:Env :: Env -> Codebase IO Symbol Ann
codebase, Runtime Symbol
runtime :: Runtime Symbol
$sel:runtime:Env :: Env -> Runtime Symbol
runtime, Maybe String
workDir :: Maybe String
$sel:workDir:Env :: Env -> Maybe String
workDir} <- ExceptT Text MCP Env
forall r (m :: * -> *). MonadReader r m => m r
ask
ProjectPath
initialPP <- MCP (Either Text ProjectPath) -> ExceptT Text MCP ProjectPath
forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT (MCP (Either Text ProjectPath) -> ExceptT Text MCP ProjectPath)
-> (IO (Either Text ProjectPath) -> MCP (Either Text ProjectPath))
-> IO (Either Text ProjectPath)
-> ExceptT Text MCP ProjectPath
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Either Text ProjectPath) -> MCP (Either Text ProjectPath)
forall a. IO a -> MCP a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either Text ProjectPath) -> ExceptT Text MCP ProjectPath)
-> IO (Either Text ProjectPath) -> ExceptT Text MCP ProjectPath
forall a b. (a -> b) -> a -> b
$ Codebase IO Symbol Ann
-> ExceptT Text Transaction ProjectPath
-> IO (Either Text ProjectPath)
forall (m :: * -> *) v a e b.
MonadIO m =>
Codebase m v a -> ExceptT e Transaction b -> m (Either e b)
Codebase.runTransactionExceptT Codebase IO Symbol Ann
codebase (ExceptT Text Transaction ProjectPath
-> IO (Either Text ProjectPath))
-> ExceptT Text Transaction ProjectPath
-> IO (Either Text ProjectPath)
forall a b. (a -> b) -> a -> b
$ do
ProjectContext -> ExceptT Text Transaction ProjectPath
ppForProjectContext ProjectContext
projCtx
CredentialManager
credMan <- ExceptT Text MCP CredentialManager
forall (m :: * -> *). MonadIO m => m CredentialManager
AuthN.newCredentialManager
let tokenProvider :: AuthN.TokenProvider
tokenProvider :: TokenProvider
tokenProvider = CredentialManager -> TokenProvider
AuthN.newTokenProvider CredentialManager
credMan
AuthenticatedHttpClient
authenticatedHTTPClient <- TokenProvider -> Text -> ExceptT Text MCP AuthenticatedHttpClient
forall (m :: * -> *).
MonadIO m =>
TokenProvider -> Text -> m AuthenticatedHttpClient
AuthN.newAuthenticatedHTTPClient TokenProvider
tokenProvider Text
ucmVersion
TVar (Seq Pretty)
outputVar <- Seq Pretty -> ExceptT Text MCP (TVar (Seq Pretty))
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Seq Pretty
forall a. Seq a
Seq.empty
TVar (Seq Text)
sourceCodeUpdatesVar <- Seq Text -> ExceptT Text MCP (TVar (Seq Text))
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Seq Text
forall a. Seq a
Seq.empty
let notify :: Output -> IO ()
notify Output
output = do
Pretty
pretty <- Maybe String -> Output -> IO Pretty
Output.notifyUser Maybe String
workDir Output
output
STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Seq Pretty) -> (Seq Pretty -> Seq Pretty) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar TVar (Seq Pretty)
outputVar (Seq Pretty -> Seq Pretty -> Seq Pretty
forall a. Semigroup a => a -> a -> a
<> Pretty -> Seq Pretty
forall a. a -> Seq a
Seq.singleton Pretty
pretty)
let notifyNumbered :: NumberedOutput -> IO NumberedArgs
notifyNumbered NumberedOutput
output = do
let (Pretty
pretty, NumberedArgs
nargs) = NumberedOutput -> (Pretty, NumberedArgs)
Output.notifyNumbered NumberedOutput
output
STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Seq Pretty) -> (Seq Pretty -> Seq Pretty) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar TVar (Seq Pretty)
outputVar (Seq Pretty -> Seq Pretty -> Seq Pretty
forall a. Semigroup a => a -> a -> a
<> Pretty -> Seq Pretty
forall a. a -> Seq a
Seq.singleton Pretty
pretty)
NumberedArgs -> IO NumberedArgs
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NumberedArgs
nargs
let loadSource :: a
loadSource = String -> a
forall a. HasCallStack => String -> a
error String
"loadSource is not implemented for the MCP server."
let writeSource :: Text -> Text -> Bool -> IO ()
writeSource Text
_sourceName Text
content Bool
replace = do
if Bool
replace
then do
STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Seq Text) -> Seq Text -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Seq Text)
sourceCodeUpdatesVar (Text -> Seq Text
forall a. a -> Seq a
Seq.singleton Text
content)
else do
STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Seq Text) -> (Seq Text -> Seq Text) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar TVar (Seq Text)
sourceCodeUpdatesVar (Seq Text -> Seq Text -> Seq Text
forall a. Semigroup a => a -> a -> a
<> Text -> Seq Text
forall a. a -> Seq a
Seq.singleton Text
content)
IORef Int
seedRef <- IO (IORef Int) -> ExceptT Text MCP (IORef Int)
forall a. IO a -> ExceptT Text MCP a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef Int) -> ExceptT Text MCP (IORef Int))
-> IO (IORef Int) -> ExceptT Text MCP (IORef Int)
forall a b. (a -> b) -> a -> b
$ Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef (Int
0 :: Int)
let cliEnv :: Env
cliEnv =
Cli.Env
{ $sel:authHTTPClient:Env :: AuthenticatedHttpClient
authHTTPClient = AuthenticatedHttpClient
authenticatedHTTPClient,
Codebase IO Symbol Ann
codebase :: Codebase IO Symbol Ann
$sel:codebase:Env :: Codebase IO Symbol Ann
codebase,
$sel:credentialManager:Env :: CredentialManager
credentialManager = CredentialManager
credMan,
$sel:generateUniqueName:Env :: IO UniqueName
generateUniqueName = do
Int
i <- IORef Int -> (Int -> (Int, Int)) -> IO Int
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef Int
seedRef \Int
i -> let !i' :: Int
i' = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 in (Int
i', Int
i)
UniqueName -> IO UniqueName
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ChaChaDRG -> UniqueName
forall gen. DRG gen => gen -> UniqueName
Parser.uniqueBase32Namegen (Seed -> ChaChaDRG
Random.drgNewSeed (Integer -> Seed
Random.seedFromInteger (Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i)))),
Text -> IO LoadSourceResult
forall {a}. a
loadSource :: forall {a}. a
$sel:loadSource:Env :: Text -> IO LoadSourceResult
loadSource,
$sel:lspCheckForChanges:Env :: ProjectPathIds -> IO ()
lspCheckForChanges = \ProjectPathIds
_ -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
Text -> Text -> Bool -> IO ()
writeSource :: Text -> Text -> Bool -> IO ()
$sel:writeSource:Env :: Text -> Text -> Bool -> IO ()
writeSource,
Output -> IO ()
notify :: Output -> IO ()
$sel:notify:Env :: Output -> IO ()
notify,
NumberedOutput -> IO NumberedArgs
notifyNumbered :: NumberedOutput -> IO NumberedArgs
$sel:notifyNumbered:Env :: NumberedOutput -> IO NumberedArgs
notifyNumbered,
Runtime Symbol
runtime :: Runtime Symbol
$sel:runtime:Env :: Runtime Symbol
runtime,
$sel:sandboxedRuntime:Env :: Runtime Symbol
sandboxedRuntime = String -> Runtime Symbol
forall a. HasCallStack => String -> a
error String
"Sandboxed runtime not implemented in MCP Server",
$sel:serverBaseUrl:Env :: Maybe BaseUrl
serverBaseUrl = Maybe BaseUrl
forall a. Maybe a
Nothing,
Text
ucmVersion :: Text
$sel:ucmVersion:Env :: Text
ucmVersion,
$sel:isTranscriptTest:Env :: Bool
isTranscriptTest = Bool
False
}
let startState :: LoopState
startState = (ProjectPathIds -> LoopState
Cli.loopState0 (ProjectPath -> ProjectPathIds
PP.toIds ProjectPath
initialPP))
(ReturnType a
cliResult, LoopState
_loopState) <- IO (ReturnType a, LoopState)
-> ExceptT Text MCP (ReturnType a, LoopState)
forall a. IO a -> ExceptT Text MCP a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Env -> LoopState -> Cli a -> IO (ReturnType a, LoopState)
forall a. Env -> LoopState -> Cli a -> IO (ReturnType a, LoopState)
Cli.runCli Env
cliEnv LoopState
startState Cli a
cli)
CliOutput
cliOut <- STM CliOutput -> ExceptT Text MCP CliOutput
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM CliOutput -> ExceptT Text MCP CliOutput)
-> STM CliOutput -> ExceptT Text MCP CliOutput
forall a b. (a -> b) -> a -> b
$ do
Seq Pretty
msgs <- TVar (Seq Pretty) -> STM (Seq Pretty)
forall a. TVar a -> STM a
readTVar TVar (Seq Pretty)
outputVar
[Text]
sourceCodeUpdates <- Seq Text -> [Text]
forall a. Seq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Seq Text -> [Text]) -> STM (Seq Text) -> STM [Text]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Seq Text) -> STM (Seq Text)
forall a. TVar a -> STM a
readTVar TVar (Seq Text)
sourceCodeUpdatesVar
let outputMessages :: [Text]
outputMessages =
Seq Pretty
msgs
Seq Pretty -> (Seq Pretty -> Seq Text) -> Seq Text
forall a b. a -> (a -> b) -> b
& (Pretty -> Text) -> Seq Pretty -> Seq Text
forall a b. (a -> b) -> Seq a -> Seq b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (String -> Text
Text.pack (String -> Text) -> (Pretty -> String) -> Pretty -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pretty -> String
Pretty.toPlainUnbroken)
Seq Text -> (Seq Text -> [Text]) -> [Text]
forall a b. a -> (a -> b) -> b
& Seq Text -> [Text]
forall a. Seq a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList
CliOutput -> STM CliOutput
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CliOutput -> STM CliOutput) -> CliOutput -> STM CliOutput
forall a b. (a -> b) -> a -> b
$
( CliOutput
{ [Text]
$sel:sourceCodeUpdates:CliOutput :: [Text]
sourceCodeUpdates :: [Text]
sourceCodeUpdates,
[Text]
$sel:outputMessages:CliOutput :: [Text]
outputMessages :: [Text]
outputMessages
}
)
case ReturnType a
cliResult of
ReturnType a
Cli.Continue -> (Maybe a, CliOutput) -> ExceptT Text MCP (Maybe a, CliOutput)
forall a. a -> ExceptT Text MCP a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a
forall a. Maybe a
Nothing, CliOutput
cliOut)
ReturnType a
Cli.HaltRepl -> (Maybe a, CliOutput) -> ExceptT Text MCP (Maybe a, CliOutput)
forall a. a -> ExceptT Text MCP a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a
forall a. Maybe a
Nothing, CliOutput
cliOut)
Cli.Success a
a -> (Maybe a, CliOutput) -> ExceptT Text MCP (Maybe a, CliOutput)
forall a. a -> ExceptT Text MCP a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Maybe a
forall a. a -> Maybe a
Just a
a, CliOutput
cliOut)