Skip to content

Commit c9c45fc

Browse files
aryairaniaryairani
authored
Merge pull request #5840 from unisonweb/cp/share-http-clients
2 parents 84a0aa9 + 13fa34a commit c9c45fc

File tree

5 files changed

+96
-61
lines changed

5 files changed

+96
-61
lines changed

‎unison-cli/src/Unison/Auth/Tokens.hs‎

Lines changed: 26 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import Data.Time.Clock (getCurrentTime)
88
import Network.HTTP.Client qualified as HTTP
99
import Network.HTTP.Client.TLS qualified as HTTP
1010
import Network.HTTP.Types qualified as Network
11+
import System.Environment (lookupEnv)
1112
import Unison.Auth.CredentialManager
1213
import Unison.Auth.Discovery (fetchDiscoveryDoc)
1314
import Unison.Auth.Types
@@ -20,21 +21,34 @@ import UnliftIO qualified
2021
-- The TokenProvider may automatically refresh access tokens if we have a refresh token.
2122
type TokenProvider = CodeserverId -> IO (Either CredentialFailure AccessToken)
2223

24+
-- | If provided, this access token will be used on all
25+
-- requests which use the Authenticated HTTP Client; i.e. all codeserver interactions.
26+
--
27+
-- It's useful in scripted contexts or when running transcripts against a codeserver.
28+
accessTokenEnvVarKey :: String
29+
accessTokenEnvVarKey = "UNISON_SHARE_ACCESS_TOKEN"
30+
2331
-- | Creates a 'TokenProvider' using the given 'CredentialManager'
2432
newTokenProvider :: CredentialManager -> TokenProvider
2533
newTokenProvider manager host = UnliftIO.try @_ @CredentialFailure $ do
26-
creds@CodeserverCredentials {tokens, discoveryURI} <- throwEitherM $ getCredentials manager host
27-
let Tokens {accessToken = currentAccessToken} = tokens
28-
expired <- isExpired creds
29-
if expired
30-
then do
31-
discoveryDoc <- throwEitherM $ fetchDiscoveryDoc discoveryURI
32-
fetchTime <- getCurrentTime
33-
newTokens@(Tokens {accessToken = newAccessToken}) <- throwEitherM $ performTokenRefresh discoveryDoc tokens
34-
userInfo <- throwEitherM $ getUserInfo discoveryDoc newAccessToken
35-
saveCredentials manager host (codeserverCredentials discoveryURI newTokens fetchTime userInfo)
36-
pure $ newAccessToken
37-
else pure currentAccessToken
34+
mayShareAccessToken <- fmap Text.pack <$> lookupEnv accessTokenEnvVarKey
35+
case mayShareAccessToken of
36+
Just accessToken -> do
37+
-- If the access token is provided via environment variable, we don't need to refresh it.
38+
pure accessToken
39+
Nothing -> do
40+
creds@CodeserverCredentials {tokens, discoveryURI} <- throwEitherM $ getCredentials manager host
41+
let Tokens {accessToken = currentAccessToken} = tokens
42+
expired <- isExpired creds
43+
if expired
44+
then do
45+
discoveryDoc <- throwEitherM $ fetchDiscoveryDoc discoveryURI
46+
fetchTime <- getCurrentTime
47+
newTokens@(Tokens {accessToken = newAccessToken}) <- throwEitherM $ performTokenRefresh discoveryDoc tokens
48+
userInfo <- throwEitherM $ getUserInfo discoveryDoc newAccessToken
49+
saveCredentials manager host (codeserverCredentials discoveryURI newTokens fetchTime userInfo)
50+
pure newAccessToken
51+
else pure currentAccessToken
3852

3953
-- | Don't yet support automatically refreshing tokens.
4054
--

‎unison-cli/src/Unison/Codebase/Transcript/Runner.hs‎

Lines changed: 29 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@ import Data.Text.Encoding qualified as Text
2424
import Data.These (These (..))
2525
import Data.UUID.V4 qualified as UUID
2626
import Network.HTTP.Client qualified as HTTP
27-
import System.Environment (lookupEnv)
2827
import System.IO qualified as IO
2928
import Text.Megaparsec qualified as P
3029
import U.Codebase.Sqlite.DbId qualified as Db
@@ -79,13 +78,6 @@ import Prelude hiding (readFile, writeFile)
7978
terminalWidth :: Pretty.Width
8079
terminalWidth = 65
8180

82-
-- | If provided, this access token will be used on all
83-
-- requests which use the Authenticated HTTP Client; i.e. all codeserver interactions.
84-
--
85-
-- It's useful in scripted contexts or when running transcripts against a codeserver.
86-
accessTokenEnvVarKey :: String
87-
accessTokenEnvVarKey = "UNISON_SHARE_ACCESS_TOKEN"
88-
8981
type Runner =
9082
-- | The name of the transcript to run.
9183
String ->
@@ -104,14 +96,18 @@ withRunner ::
10496
(Runner -> m r) ->
10597
m r
10698
withRunner isTest verbosity ucmVersion action = do
99+
credMan <- AuthN.newCredentialManager
100+
authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient credMan
101+
107102
-- If we're in a transcript test, configure the environment to use a non-existent fzf binary
108103
-- so that errors are consistent.
109104
-- This also prevents automated transcript tests from mistakenly opening fzf and waiting for user input.
110105
when isTest $ do
111106
liftIO $ setEnv Fuzzy.fzfPathEnvVar "NONE"
112107
withRuntimes \runtime sbRuntime ->
113108
action \transcriptName transcriptSrc codebase -> do
114-
mcpServerConfig <- MCP.initServer codebase runtime sbRuntime Nothing ucmVersion
109+
let workDir = Nothing
110+
mcpServerConfig <- MCP.initServer codebase runtime sbRuntime workDir ucmVersion authenticatedHTTPClient
115111
Server.startServer
116112
isTest
117113
Backend.BackendEnv {Backend.useNamesIndex = False}
@@ -121,19 +117,33 @@ withRunner isTest verbosity ucmVersion action = do
121117
(MCP.mcpServer mcpServerConfig)
122118
\case
123119
Nothing -> pure $ Left PortBindingFailure
124-
Just baseUrl ->
125-
either
126-
(pure . Left . ParseError)
127-
( run isTest verbosity codebase runtime sbRuntime ucmVersion $
128-
tShow @Server.BaseUrl baseUrl
129-
)
130-
$ Transcript.parse transcriptName transcriptSrc
120+
Just baseUrl -> do
121+
let baseUrlText = tShow @Server.BaseUrl baseUrl
122+
case (Transcript.parse transcriptName transcriptSrc) of
123+
Left parseError -> pure $ Left (ParseError parseError)
124+
Right stanzas ->
125+
run
126+
isTest
127+
verbosity
128+
codebase
129+
runtime
130+
sbRuntime
131+
ucmVersion
132+
baseUrlText
133+
authenticatedHTTPClient
134+
credMan
135+
stanzas
131136
where
132137
withRuntimes :: (Runtime.Runtime Symbol -> Runtime.Runtime Symbol -> m a) -> m a
133138
withRuntimes action =
134139
RTI.withRuntime False RTI.Persistent ucmVersion \runtime ->
135140
RTI.withRuntime True RTI.Persistent ucmVersion \sbRuntime ->
136141
action runtime sbRuntime
142+
initTranscriptAuthenticatedHTTPClient :: AuthN.CredentialManager -> m AuthN.AuthenticatedHttpClient
143+
initTranscriptAuthenticatedHTTPClient credMan = liftIO $ do
144+
let tokenProvider :: AuthN.TokenProvider
145+
tokenProvider = AuthN.newTokenProvider credMan
146+
AuthN.newAuthenticatedHTTPClient tokenProvider ucmVersion
137147

138148
isGeneratedBlock :: ProcessedBlock -> Bool
139149
isGeneratedBlock = generated . getCommonInfoTags
@@ -147,9 +157,11 @@ run ::
147157
Runtime.Runtime Symbol ->
148158
UCMVersion ->
149159
Text ->
160+
AuthN.AuthenticatedHttpClient ->
161+
AuthN.CredentialManager ->
150162
Transcript ->
151163
IO (Either Error Transcript)
152-
run isTest verbosity codebase runtime sbRuntime ucmVersion baseURL transcript = UnliftIO.try do
164+
run isTest verbosity codebase runtime sbRuntime ucmVersion baseURL authenticatedHTTPClient credMan transcript = UnliftIO.try do
153165
let behaviors = extractBehaviors $ settings transcript
154166
let stanzas' = stanzas transcript
155167
httpManager <- HTTP.newManager HTTP.defaultManagerSettings
@@ -163,14 +175,6 @@ run isTest verbosity codebase runtime sbRuntime ucmVersion baseURL transcript =
163175
"Running the provided transcript file...",
164176
""
165177
]
166-
mayShareAccessToken <- fmap Text.pack <$> lookupEnv accessTokenEnvVarKey
167-
credMan <- AuthN.newCredentialManager
168-
let tokenProvider :: AuthN.TokenProvider
169-
tokenProvider =
170-
maybe
171-
(AuthN.newTokenProvider credMan)
172-
(\accessToken _codeserverID -> pure $ Right accessToken)
173-
mayShareAccessToken
174178
-- Queue of Stanzas and Just index, or Nothing if the stanza was programmatically generated
175179
-- e.g. a unison-file update by a command like 'edit'
176180
inputQueue <-
@@ -510,8 +514,6 @@ run isTest verbosity codebase runtime sbRuntime ucmVersion baseURL transcript =
510514
\issues."
511515
(_, _, _) -> pure ()
512516

513-
authenticatedHTTPClient <- AuthN.newAuthenticatedHTTPClient tokenProvider ucmVersion
514-
515517
seedRef <- newIORef (0 :: Int)
516518

517519
let env =

‎unison-cli/src/Unison/CommandLine/Main.hs‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,9 @@ import System.Console.Haskeline.History qualified as Line
2121
import System.FSNotify qualified as FSNotify
2222
import System.IO (hGetEcho, hPutStrLn, hSetEcho, stderr, stdin)
2323
import System.IO.Error (isDoesNotExistError)
24-
import Unison.Auth.CredentialManager (newCredentialManager)
24+
import Unison.Auth.CredentialManager qualified as AuthN
2525
import Unison.Auth.HTTPClient (AuthenticatedHttpClient)
2626
import Unison.Auth.HTTPClient qualified as AuthN
27-
import Unison.Auth.Tokens qualified as AuthN
2827
import Unison.Cli.Monad qualified as Cli
2928
import Unison.Cli.Pretty qualified as P
3029
import Unison.Cli.ProjectUtils qualified as ProjectUtils
@@ -143,10 +142,12 @@ main ::
143142
Codebase IO Symbol Ann ->
144143
Maybe Server.BaseUrl ->
145144
UCMVersion ->
145+
AuthN.AuthenticatedHttpClient ->
146+
AuthN.CredentialManager ->
146147
(PP.ProjectPathIds -> IO ()) ->
147148
ShouldWatchFiles ->
148149
IO ()
149-
main dir welcome ppIds initialInputs runtime sbRuntime codebase serverBaseUrl ucmVersion lspCheckForChanges shouldWatchFiles = do
150+
main dir welcome ppIds initialInputs runtime sbRuntime codebase serverBaseUrl ucmVersion authHTTPClient credentialManager lspCheckForChanges shouldWatchFiles = do
150151
-- we don't like FSNotify's debouncing (it seems to drop later events)
151152
-- so we will be doing our own instead
152153
let config = FSNotify.defaultConfig
@@ -175,9 +176,6 @@ main dir welcome ppIds initialInputs runtime sbRuntime codebase serverBaseUrl uc
175176
initialInputsRef <- newIORef $ Welcome.run welcome ++ initialInputs
176177
pageOutput <- newIORef True
177178

178-
credentialManager <- newCredentialManager
179-
let tokenProvider = AuthN.newTokenProvider credentialManager
180-
authHTTPClient <- AuthN.newAuthenticatedHTTPClient tokenProvider ucmVersion
181179
initialEcho <- hGetEcho stdin
182180
let restoreEcho = (\currentEcho -> when (currentEcho /= initialEcho) $ hSetEcho stdin initialEcho)
183181
let getInput :: Cli.LoopState -> IO Input

‎unison-cli/src/Unison/MCP.hs‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,7 @@ import Network.MCP.Server qualified as MCP
44
import Network.MCP.Server.StdIO qualified as MCP
55
import Network.MCP.Types
66
import Text.RawString.QQ (r)
7-
import Unison.Auth.CredentialManager qualified as AuthN
87
import Unison.Auth.HTTPClient qualified as AuthN
9-
import Unison.Auth.Tokens qualified as AuthN
108
import Unison.Codebase (Codebase)
119
import Unison.Codebase.Runtime (Runtime)
1210
import Unison.MCP.Prompts (prompts)
@@ -28,12 +26,8 @@ serverDescription =
2826
Before doing any work in unison please read the file://unison-guide resource for information on how to write unison.
2927
|]
3028

31-
initServer :: Codebase IO Symbol Ann -> Runtime Symbol -> Runtime Symbol -> Maybe FilePath -> Text -> IO MCP.Server
32-
initServer codebase runtime sbRuntime workDir ucmVersion = do
33-
credMan <- AuthN.newCredentialManager
34-
let tokenProvider :: AuthN.TokenProvider
35-
tokenProvider = AuthN.newTokenProvider credMan
36-
authenticatedHTTPClient <- AuthN.newAuthenticatedHTTPClient tokenProvider ucmVersion
29+
initServer :: Codebase IO Symbol Ann -> Runtime Symbol -> Runtime Symbol -> Maybe FilePath -> Text -> AuthN.AuthenticatedHttpClient -> IO MCP.Server
30+
initServer codebase runtime sbRuntime workDir ucmVersion authenticatedHTTPClient = do
3731
let env =
3832
Env
3933
{ codebase,
@@ -49,8 +43,8 @@ initServer codebase runtime sbRuntime workDir ucmVersion = do
4943
runMCP env $ MCPWrapper.mkServer serverInfo serverDescription staticResources tools prompts
5044

5145
-- | Run the MCP server until we hit EOF.
52-
runOnStdIO :: Codebase IO Symbol Ann -> Runtime Symbol -> Runtime Symbol -> FilePath -> Text -> IO ()
53-
runOnStdIO codebase runtime sbRuntime workDir ucmVersion = do
54-
server <- initServer codebase runtime sbRuntime (pure workDir) ucmVersion
46+
runOnStdIO :: Codebase IO Symbol Ann -> Runtime Symbol -> Runtime Symbol -> FilePath -> Text -> AuthN.AuthenticatedHttpClient -> IO ()
47+
runOnStdIO codebase runtime sbRuntime workDir ucmVersion authenticatedHTTPClient = do
48+
server <- initServer codebase runtime sbRuntime (Just workDir) ucmVersion authenticatedHTTPClient
5549
-- Start the server with StdIO transport
5650
MCP.runServerWithSTDIO server

‎unison-cli/src/Unison/Main.hs‎

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -60,10 +60,14 @@ import System.IO.Temp qualified as Temp
6060
import System.Path qualified as Path
6161
import Text.Megaparsec qualified as MP
6262
import U.Codebase.Sqlite.Queries qualified as Queries
63+
import Unison.Auth.CredentialManager qualified as AuthN
64+
import Unison.Auth.HTTPClient qualified as AuthN
65+
import Unison.Auth.Tokens qualified as AuthN
6366
import Unison.Cli.ProjectUtils qualified as ProjectUtils
6467
import Unison.Codebase (Codebase, CodebasePath)
6568
import Unison.Codebase qualified as Codebase
6669
import Unison.Codebase.Editor.Input qualified as Input
70+
import Unison.Codebase.Editor.UCMVersion (UCMVersion)
6771
import Unison.Codebase.Execute (execute)
6872
import Unison.Codebase.Init (CodebaseInitOptions (..), InitError (..), InitResult (..), SpecifiedCodebase (..))
6973
import Unison.Codebase.Init qualified as CodebaseInit
@@ -151,9 +155,12 @@ main version = do
151155
PrintVersion ->
152156
Text.putStrLn $ Text.pack progName <> " version: " <> Version.gitDescribeWithDate version
153157
MCPServer -> do
158+
let ucmVersion = Version.gitDescribeWithDate version
159+
credMan <- AuthN.newCredentialManager
160+
authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient ucmVersion credMan
154161
getCodebaseOrExit mCodePathOption SC.DontLock (SC.MigrateAfterPrompt SC.Backup SC.Vacuum) \(_initRes, _, theCodebase) -> do
155162
withRuntimes RTI.Persistent \(runtime, sbRuntime) -> do
156-
MCP.runOnStdIO theCodebase runtime sbRuntime currentDir (Version.gitDescribeWithDate version)
163+
MCP.runOnStdIO theCodebase runtime sbRuntime currentDir ucmVersion authenticatedHTTPClient
157164
Init -> do
158165
exitError
159166
( P.lines
@@ -184,6 +191,9 @@ main version = do
184191
let fileEvent = Input.UnisonFileChanged (Text.pack file) contents
185192
let noOpCheckForChanges _ = pure ()
186193
let serverUrl = Nothing
194+
let ucmVersion = Version.gitDescribeWithDate version
195+
credMan <- liftIO $ AuthN.newCredentialManager
196+
authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient ucmVersion credMan
187197
startProjectPath <- Codebase.runTransaction theCodebase Codebase.expectCurrentProjectPath
188198
launch
189199
version
@@ -192,6 +202,8 @@ main version = do
192202
sbrt
193203
theCodebase
194204
[Left fileEvent, Right $ Input.ExecuteI mainName args, Right Input.QuitI]
205+
authenticatedHTTPClient
206+
credMan
195207
serverUrl
196208
(PP.toIds startProjectPath)
197209
initRes
@@ -207,6 +219,9 @@ main version = do
207219
let fileEvent = Input.UnisonFileChanged (Text.pack "<standard input>") contents
208220
let noOpCheckForChanges _ = pure ()
209221
let serverUrl = Nothing
222+
let ucmVersion = Version.gitDescribeWithDate version
223+
credMan <- liftIO $ AuthN.newCredentialManager
224+
authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient ucmVersion credMan
210225
startProjectPath <- Codebase.runTransaction theCodebase Codebase.expectCurrentProjectPath
211226
launch
212227
version
@@ -215,6 +230,8 @@ main version = do
215230
sbrt
216231
theCodebase
217232
[Left fileEvent, Right $ Input.ExecuteI mainName args, Right Input.QuitI]
233+
authenticatedHTTPClient
234+
credMan
218235
serverUrl
219236
(PP.toIds startProjectPath)
220237
initRes
@@ -322,8 +339,10 @@ main version = do
322339
-- https://gitlab.haskell.org/ghc/ghc/-/merge_requests/1224
323340
void . Ki.fork scope $ LSP.spawnLsp lspFormattingConfig theCodebase runtime changeSignal
324341
let isTest = False
325-
mcpServerConfig <-
326-
MCP.initServer theCodebase runtime sbRuntime (pure currentDir) $ Version.gitDescribeWithDate version
342+
let ucmVersion = Version.gitDescribeWithDate version
343+
credMan <- liftIO $ AuthN.newCredentialManager
344+
authenticatedHTTPClient <- initTranscriptAuthenticatedHTTPClient ucmVersion credMan
345+
mcpServerConfig <- MCP.initServer theCodebase runtime sbRuntime (Just currentDir) ucmVersion authenticatedHTTPClient
327346
Server.startServer
328347
isTest
329348
Backend.BackendEnv {Backend.useNamesIndex = False}
@@ -358,14 +377,15 @@ main version = do
358377
takeMVar mvar
359378
WithCLI -> do
360379
PT.putPrettyLn $ P.string "Now starting the Unison Codebase Manager (UCM)..."
361-
362380
launch
363381
version
364382
currentDir
365383
runtime
366384
sbRuntime
367385
theCodebase
368386
[]
387+
authenticatedHTTPClient
388+
credMan
369389
mayBaseUrl
370390
(PP.toIds startingProjectPath)
371391
initRes
@@ -379,6 +399,9 @@ main version = do
379399
RTI.withRuntime False mode (Version.gitDescribeWithDate version) \runtime -> do
380400
RTI.withRuntime True mode (Version.gitDescribeWithDate version) \sbRuntime ->
381401
action (runtime, sbRuntime)
402+
initTranscriptAuthenticatedHTTPClient :: UCMVersion -> AuthN.CredentialManager -> IO AuthN.AuthenticatedHttpClient
403+
initTranscriptAuthenticatedHTTPClient ucmVersion credMan = do
404+
AuthN.newAuthenticatedHTTPClient (AuthN.newTokenProvider credMan) ucmVersion
382405

383406
isExitSuccess :: SomeException -> Bool
384407
isExitSuccess =
@@ -583,13 +606,15 @@ launch ::
583606
Rt.Runtime Symbol ->
584607
Codebase.Codebase IO Symbol Ann ->
585608
[Either Input.Event Input.Input] ->
609+
AuthN.AuthenticatedHttpClient ->
610+
AuthN.CredentialManager ->
586611
Maybe Server.BaseUrl ->
587612
PP.ProjectPathIds ->
588613
InitResult ->
589614
(PP.ProjectPathIds -> IO ()) ->
590615
CommandLine.ShouldWatchFiles ->
591616
IO ()
592-
launch version dir runtime sbRuntime codebase inputs serverBaseUrl startingPath initResult lspCheckForChanges shouldWatchFiles = do
617+
launch version dir runtime sbRuntime codebase inputs authenticatedHTTPClient credMan serverBaseUrl startingPath initResult lspCheckForChanges shouldWatchFiles = do
593618
showWelcomeHint <- Codebase.runTransaction codebase Queries.doProjectsExist
594619
let isNewCodebase = case initResult of
595620
CreatedCodebase -> NewlyCreatedCodebase
@@ -606,6 +631,8 @@ launch version dir runtime sbRuntime codebase inputs serverBaseUrl startingPath
606631
codebase
607632
serverBaseUrl
608633
ucmVersion
634+
authenticatedHTTPClient
635+
credMan
609636
lspCheckForChanges
610637
shouldWatchFiles
611638

0 commit comments

Comments
 (0)