@stainless-api/docs 0.1.0-beta.137 → 0.1.0-beta.139

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,192 @@
1
+ import path from 'path';
2
+ import type { SpecLoaderFn, SpecLoaderParams } from './utils';
3
+ import Stainless, { APIError } from '@stainless-api/sdk';
4
+ import { DocsLanguage } from '@stainless-api/docs-ui/routing';
5
+ import { bold } from '../../shared/terminalUtils';
6
+ import { mkdir, readdir, readFile, rm, writeFile } from 'fs/promises';
7
+ import { generateSpecFromStrings, previewWorkerCode } from '@stainless/sdk-json/spec';
8
+ import { Spec, SpecLanguage } from '@stainless/sdk-json';
9
+ import crypto from 'crypto';
10
+
11
+ function resolvePath(inputPath: string) {
12
+ return path.resolve(process.cwd(), inputPath);
13
+ }
14
+
15
+ function getLocalFilePaths() {
16
+ // eslint-disable-next-line turbo/no-undeclared-env-vars
17
+ const oasPath = process.env.OPENAPI_PATH;
18
+ // eslint-disable-next-line turbo/no-undeclared-env-vars
19
+ const configPath = process.env.STAINLESS_CONFIG_PATH;
20
+
21
+ if (!oasPath || !configPath) {
22
+ return null;
23
+ }
24
+
25
+ return {
26
+ oasPath: resolvePath(oasPath),
27
+ configPath: resolvePath(configPath),
28
+ };
29
+ }
30
+
31
+ async function fetchVersionInfo(project: string, apiKey: string): Promise<Record<DocsLanguage, string>> {
32
+ const data = await fetch(`https://api.stainless.com/api/projects/${project}/package-versions`, {
33
+ headers: { Authorization: `Bearer ${apiKey}` },
34
+ });
35
+
36
+ const content = await data.text();
37
+ return JSON.parse(content) as Record<DocsLanguage, string>;
38
+ }
39
+
40
+ function redactApiKey(apiKey: string) {
41
+ return apiKey
42
+ .split('')
43
+ .map((char, index) => (index < 10 ? char : '*'))
44
+ .join('');
45
+ }
46
+
47
+ async function loadInputs({ apiKey, logger, stainlessProject, branch }: SpecLoaderParams) {
48
+ const localFilePaths = getLocalFilePaths();
49
+
50
+ if (localFilePaths) {
51
+ try {
52
+ const oasStr = await readFile(localFilePaths.oasPath, 'utf8');
53
+ const configStr = await readFile(localFilePaths.configPath, 'utf8');
54
+ return {
55
+ oasStr,
56
+ configStr,
57
+ versionInfo: null,
58
+ };
59
+ } catch (e) {
60
+ logger.error(bold('Failed to load spec inputs from files:'));
61
+ logger.error(e instanceof Error ? e.message : String(e));
62
+ process.exit(1);
63
+ }
64
+ }
65
+
66
+ if (!apiKey) {
67
+ logger.error(
68
+ [
69
+ bold(
70
+ 'No Stainless credentials found. Please choose one of the following options to authenticate with Stainless:',
71
+ ),
72
+ '- Run `stl auth login` to authenticate via the Stainless CLI',
73
+ '- Provide a Stainless API key via the `STAINLESS_API_KEY` environment variable (eg. in a .env file)',
74
+ '- Set the `apiKey` option in the Stainless Docs config',
75
+ ].join('\n'),
76
+ );
77
+ process.exit(1);
78
+ }
79
+
80
+ try {
81
+ const client = new Stainless({ apiKey });
82
+ const configs = await client.projects.configs.retrieve({
83
+ project: stainlessProject,
84
+ branch: branch,
85
+ include: 'openapi',
86
+ });
87
+ const versionInfo = await fetchVersionInfo(stainlessProject, apiKey);
88
+
89
+ const configYML = Object.values(configs)[0] as { content: unknown };
90
+ const oasJson = Object.values(configs)[1] as { content: unknown };
91
+ const oasStr = oasJson['content'];
92
+ const configStr = configYML['content'];
93
+
94
+ if (typeof oasStr !== 'string' || typeof configStr !== 'string') {
95
+ logger.error('Received invalid OAS or config from Stainless');
96
+ process.exit(1);
97
+ }
98
+
99
+ return {
100
+ oasStr,
101
+ configStr,
102
+ versionInfo,
103
+ };
104
+ } catch (e) {
105
+ if (e instanceof APIError && e.status >= 400 && e.status < 500) {
106
+ logger.error(`Failed to load requested project slug: "${stainlessProject}"`);
107
+ if (apiKey) {
108
+ logger.error(`API key: "${redactApiKey(apiKey)}"`);
109
+ }
110
+ logger.error(
111
+ `This error can usually be corrected by re-authenticating with the Stainless. Use the CLI (stl auth login) or verify that the Stainless API key you're using can access the project mentioned above.`,
112
+ );
113
+ }
114
+ process.exit(1);
115
+ }
116
+ }
117
+
118
+ async function maybeLoadJSONFile<T>(filePath: string): Promise<T | null> {
119
+ try {
120
+ const fileContents = await readFile(filePath, 'utf8');
121
+ return JSON.parse(fileContents) as T;
122
+ } catch {
123
+ return null;
124
+ }
125
+ }
126
+
127
+ async function cleanupDirectory(directory: string, filesToKeep: string[]) {
128
+ const allFiles = await readdir(directory);
129
+ const unusedFiles = allFiles.filter((file) => !filesToKeep.includes(file));
130
+ await Promise.all(unusedFiles.map((file) => rm(path.join(directory, file))));
131
+ return {
132
+ deletedCount: unusedFiles.length,
133
+ };
134
+ }
135
+
136
+ export const defaultSpecLoader: SpecLoaderFn = async (params) => {
137
+ const { createCodegenDir } = params;
138
+
139
+ const inputs = await loadInputs(params);
140
+
141
+ const specsDirectory = path.join(createCodegenDir().pathname, 'specs2');
142
+ await mkdir(specsDirectory, { recursive: true });
143
+
144
+ const fileName =
145
+ crypto
146
+ .createHash('sha256')
147
+ .update(JSON.stringify(inputs) + previewWorkerCode)
148
+ .digest('hex')
149
+ .slice(0, 10) + '.json';
150
+
151
+ const filePath = path.join(specsDirectory, fileName);
152
+
153
+ const cachedSpec = await maybeLoadJSONFile<{ languages: SpecLanguage[]; sdkJson: Spec }>(filePath);
154
+ // skip generation since we already have a cached spec
155
+ if (cachedSpec) {
156
+ params.logger.info(`Loaded cached spec: ${fileName}`);
157
+ return [
158
+ {
159
+ filePath,
160
+ languages: cachedSpec.languages,
161
+ sdkJson: cachedSpec.sdkJson,
162
+ },
163
+ ];
164
+ }
165
+
166
+ const result = await generateSpecFromStrings({
167
+ oasStr: inputs.oasStr,
168
+ configStr: inputs.configStr,
169
+ languageOverrides: {
170
+ mode: 'exclude',
171
+ list: params.excludeLanguages ?? [],
172
+ },
173
+ versionInfo: inputs.versionInfo,
174
+ stainlessProject: params.stainlessProject,
175
+ });
176
+
177
+ await writeFile(filePath, JSON.stringify(result), 'utf8');
178
+ params.logger.info(`Generated: ${fileName}`);
179
+
180
+ const { deletedCount } = await cleanupDirectory(specsDirectory, [fileName]);
181
+ if (deletedCount > 0) {
182
+ params.logger.info(`Cleaned up ${deletedCount} unused spec file(s)`);
183
+ }
184
+
185
+ return [
186
+ {
187
+ filePath,
188
+ languages: result.languages.filter((language) => language !== 'sql' && language !== 'openapi'),
189
+ sdkJson: result.sdkJson,
190
+ },
191
+ ];
192
+ };
@@ -1,7 +1,7 @@
1
1
  import { readFile } from 'fs/promises';
2
2
 
3
3
  import { api } from 'virtual:stainless-apis-manifest';
4
- import { SpecWithAuth } from './generateSpec';
4
+ import type { SpecWithAuth } from '@stainless/sdk-json/spec';
5
5
  import { DocsLanguage } from '@stainless-api/docs-ui/routing';
6
6
 
7
7
  const cachedSpecWithAuth: Record<string, SpecWithAuth> = {};
@@ -0,0 +1,86 @@
1
+ import { Spec } from '@stainless/sdk-json';
2
+ import { readFile } from 'fs/promises';
3
+ import { DocsLanguage } from '@stainless-api/docs-ui/routing';
4
+ import { AstroIntegrationLogger } from 'astro';
5
+
6
+ type PossibleLanguage = NonNullable<NonNullable<NonNullable<Spec['docs']>['languages']>[number]>;
7
+
8
+ export type SpecLoaderParams = {
9
+ /**
10
+ * The slug of your Stainless project.
11
+ */
12
+ stainlessProject: string;
13
+ /**
14
+ * The branch of your Stainless project.
15
+ */
16
+ branch: string;
17
+ /**
18
+ * The Stainless API key. This can be used to make requests against the Stainless API.
19
+ */
20
+ apiKey: string | null;
21
+ /**
22
+ * The languages that the user has explicitly asked to be excluded from the API reference.
23
+ */
24
+ excludeLanguages: DocsLanguage[] | null;
25
+ /**
26
+ * An Astro logger. This should be used for logging messages.
27
+ */
28
+ logger: AstroIntegrationLogger;
29
+ /**
30
+ * A function that creates a directory in .astro. See: https://docs.astro.build/en/reference/integrations-reference/#createcodegendir
31
+ */
32
+ createCodegenDir: () => URL;
33
+ };
34
+
35
+ type SpecLoaderResult = {
36
+ /**
37
+ * The file path to the loaded spec. The spec MUST be written to a path on disk.
38
+ * If you are not sure where to place it, create a directory in .astro using the `createCodegenDir` function passed in the parameters of the spec loader function.
39
+ */
40
+ filePath: string;
41
+ /**
42
+ * The languages that are represented in the spec.
43
+ */
44
+ languages: PossibleLanguage[];
45
+ /**
46
+ * Optionally, if you already have already the spec in memory, you can provide it. If not provided, the contents of the file at `filePath` will be read.
47
+ * IMPORTANT: This should be equivalent to the contents of the file at `filePath`. If not, bugs and inconsistencies will arise.
48
+ */
49
+ sdkJson?: Spec;
50
+ };
51
+
52
+ export type SpecLoaderFn = (opts: SpecLoaderParams) => Promise<SpecLoaderResult[]>;
53
+
54
+ async function readSpecFromFile(filePath: string) {
55
+ const txt = await readFile(filePath, 'utf8');
56
+ const json = JSON.parse(txt) as Spec;
57
+ return json;
58
+ }
59
+
60
+ export async function loadAllSpecs(specLoaderResultsPromise: Promise<SpecLoaderResult[]>) {
61
+ const specLoaderResults = await specLoaderResultsPromise;
62
+ const specs = await Promise.all(
63
+ specLoaderResults.map(async (result) => {
64
+ return {
65
+ filePath: result.filePath,
66
+ languages: result.languages,
67
+ sdkJson: result.sdkJson ?? (await readSpecFromFile(result.filePath)),
68
+ };
69
+ }),
70
+ );
71
+ return specs;
72
+ }
73
+
74
+ export type LoadedSpecs = Awaited<ReturnType<typeof loadAllSpecs>>;
75
+
76
+ export function flatSpecsList(specs: LoadedSpecs) {
77
+ return specs
78
+ .map((s) =>
79
+ s.languages.map((language) => ({
80
+ language,
81
+ sdkJson: s.sdkJson,
82
+ filePath: s.filePath,
83
+ })),
84
+ )
85
+ .flat();
86
+ }
@@ -10,23 +10,6 @@ export function buildVirtualModuleString<T extends Record<string, unknown>>(vars
10
10
  .join('\n');
11
11
  }
12
12
 
13
- export function makeVirtualModPlugin(bareId: string, content: string): VitePlugin {
14
- return {
15
- name: `stl-virtual-module-loader-${bareId}`,
16
- resolveId(id) {
17
- // The '\0' prefix tells Vite “this is a virtual module” and prevents it from being resolved again.
18
- if (id === bareId) {
19
- return `\0${bareId}`;
20
- }
21
- },
22
- load(id) {
23
- if (id === `\0${bareId}`) {
24
- return content;
25
- }
26
- },
27
- };
28
- }
29
-
30
13
  export function makeAsyncVirtualModPlugin<T extends Record<string, unknown>>(
31
14
  bareId: string,
32
15
  contentLoader: () => Promise<T | string>,
@@ -1,95 +1,31 @@
1
- import { AstroIntegrationLogger } from 'astro';
2
- import z from 'zod';
3
- import {
4
- buildVirtualModuleString,
5
- makeAsyncVirtualModPlugin,
6
- makeVirtualModPlugin,
7
- } from '../shared/virtualModule';
1
+ import { buildVirtualModuleString } from '../shared/virtualModule';
8
2
  import type * as virtualExampleModule from 'virtual:stl-docs-ai-chat-examples';
9
3
  type VirtualExampleModule = typeof virtualExampleModule;
10
4
 
11
- const exampleSchema = z.array(
12
- z.object({
13
- shortPrompt: z.string(),
14
- longPrompt: z.string(),
15
- icon: z.string(),
16
- }),
17
- );
18
- export type ExamplePromptResponse = z.infer<typeof exampleSchema>;
5
+ export type ExamplePromptResponse = {
6
+ shortPrompt: string;
7
+ longPrompt: string;
8
+ icon: string;
9
+ }[];
19
10
 
20
- // handles actually retrieving the information via the Stainless API
21
- async function loadExamples(
22
- projectName: string,
23
- logger: AstroIntegrationLogger,
24
- ): Promise<ExamplePromptResponse | undefined> {
25
- try {
26
- const response = await fetch(`https://api.stainless.com/api/ai/steelie-examples/${projectName}`, {
27
- method: 'GET',
28
- });
29
-
30
- const text = await response.text();
31
- if (!response.ok) {
32
- logger.error(`failed to fetch AI chat examples: ${text}`);
33
- return undefined;
34
- }
35
- const examples = exampleSchema.parse(JSON.parse(text));
36
- return examples;
37
- } catch (error) {
38
- if (error instanceof Error) logger.error(`failed to fetch AI chat examples: ${error.message}`);
39
- else throw error;
40
- return undefined;
41
- }
42
- }
43
-
44
- export default async function generateExamplesPlugin({
45
- projectName,
46
- logger,
47
- exampleOverrides,
48
- }: {
49
- projectName: string | undefined;
50
- logger: AstroIntegrationLogger;
51
- exampleOverrides?: ExamplePromptResponse;
52
- }) {
53
- // if the user has specified any examples, return those immediately
54
- // instead of loading them via the web.
55
- if (exampleOverrides) {
56
- return makeVirtualModPlugin(
57
- 'virtual:stl-docs-ai-chat-examples',
58
- generateVirtualModuleString(exampleOverrides),
59
- );
11
+ export function generateExamplesVirtualModule(exampleOverrides: ExamplePromptResponse | undefined): string {
12
+ if (!exampleOverrides) {
13
+ return buildVirtualModuleString({ examples: undefined } satisfies VirtualExampleModule);
60
14
  }
61
- // if we don't have a defined project name, don't try to fetch examples
62
- if (!projectName) {
63
- return makeVirtualModPlugin(
64
- 'virtual:stl-docs-ai-chat-examples',
65
- buildVirtualModuleString({ examples: undefined } satisfies VirtualExampleModule),
66
- );
67
- }
68
-
69
- // otherwise, promise to get the right examples at some point later on
70
- const examplesPromise = loadExamples(projectName, logger);
71
- return makeAsyncVirtualModPlugin<VirtualExampleModule>('virtual:stl-docs-ai-chat-examples', async () => {
72
- const examples = await examplesPromise;
73
- return generateVirtualModuleString(examples);
74
- });
75
- }
76
-
77
- function generateVirtualModuleString(examples: ExamplePromptResponse | undefined) {
78
- if (!examples) return 'export const examples = undefined;';
79
15
 
80
16
  // Generate icon imports
81
17
  // prettier-ignore
82
18
  const pascalToKebab = (str: string) => str.split(/(?=[A-Z])/).join('-').toLowerCase();
83
19
  const iconImportPath = (iconName: string) =>
84
20
  import.meta.resolve(`lucide-react/dist/esm/icons/${pascalToKebab(iconName)}.js`);
85
- const iconImports = examples.map(
21
+ const iconImports = exampleOverrides.map(
86
22
  ({ icon }) => `import ${icon} from ${JSON.stringify(iconImportPath(icon))}`,
87
23
  );
88
24
 
89
25
  // Reference icon imports in `examples` exported object
90
26
  // "icon":"Sparkles" -> "icon":Sparkles
91
27
  const iconStringsToIdents = (jsonBlob: string) => jsonBlob.replace(/"icon":\s*"(\w+)"/g, '"icon":$1');
92
- const exportBody = `export const examples = ${iconStringsToIdents(JSON.stringify(examples))};`;
28
+ const exportBody = `export const examples = ${iconStringsToIdents(JSON.stringify(exampleOverrides))};`;
93
29
 
94
30
  return [...iconImports, exportBody].join('\n');
95
31
  }
package/stl-docs/index.ts CHANGED
@@ -20,13 +20,12 @@ import {
20
20
  import { buildVirtualModuleString } from '../shared/virtualModule';
21
21
  import { resolveSrcFile } from '../resolveSrcFile';
22
22
  import { stainlessDocsMarkdownRenderer } from './proseMarkdown/proseMarkdownIntegration';
23
- import { getSharedLogger, setSharedLogger } from '../shared/getSharedLogger';
23
+ import { setSharedLogger } from '../shared/getSharedLogger';
24
24
  import { stainlessDocsVectorProseIndexing } from './proseDocSync';
25
- import { stainlessDocsAlgoliaProseIndexing } from './proseSearchIndexing';
26
25
  import { stainlessStarlight } from '../plugin';
27
26
  import { getFontRoles, flattenFonts } from './fonts';
28
27
  import conditionalIntegration from '../shared/conditionalIntegration';
29
- import generateExamplesPlugin from './aiChatExamples';
28
+ import { generateExamplesVirtualModule } from './aiChatExamples';
30
29
  import { ogImageStarlightPlugin } from './og-image';
31
30
 
32
31
  export * from '../plugin';
@@ -172,8 +171,7 @@ function stainlessDocsIntegration(
172
171
  return {
173
172
  name: 'stl-docs-astro',
174
173
  hooks: {
175
- 'astro:config:setup': async ({ updateConfig, command, config: astroConfig, logger: localLogger }) => {
176
- const logger = getSharedLogger({ fallback: localLogger });
174
+ 'astro:config:setup': ({ updateConfig, command, config: astroConfig }) => {
177
175
  // we only handle redirects for builds
178
176
  // in dev, Astro handles them for us
179
177
  if (command === 'build' && astroConfig.redirects) {
@@ -195,30 +193,35 @@ function stainlessDocsIntegration(
195
193
  vmAiChatHandlerExport = `export { default as AI_CHAT_HANDLER } from '${handlerEntrypoint}';`;
196
194
  }
197
195
 
198
- const virtualModules = new Map(
199
- Object.entries({
200
- 'virtual:stl-docs-virtual-module': [
201
- buildVirtualModuleString({
202
- TABS: config.tabs.map((tab) => ({ ...tab, link: withBase(tab.link) })),
203
- SPLIT_TABS_ENABLED: config.splitTabsEnabled,
204
- HEADER_LINKS: config.header.links.map((link) => ({ ...link, link: withBase(link.link) })),
205
- HEADER_LAYOUT: config.header.layout,
206
- ENABLE_CLIENT_ROUTER: config.enableClientRouter,
207
- API_REFERENCE_BASE_PATH: apiReferenceBasePath ?? '/api',
208
- ENABLE_PROSE_MARKDOWN_RENDERING: config.enableProseMarkdownRendering,
209
- ENABLE_CONTEXT_MENU: !!config.contextMenu, // TODO: do not duplicate this between both virtual modules
210
- CONTEXT_MENU_ENABLE_THIRD_PARTY:
211
- (typeof config.contextMenu === 'object' ? config.contextMenu.thirdParty : null) ?? true,
212
- RENDER_PAGE_DESCRIPTIONS: config.renderPageDescriptions,
213
- FONTS: getFontRoles(config.fonts),
214
- LINK_GROUP_TITLES_TO_OVERVIEW_PAGES: config.linkGroupTitlesToOverviewPages,
215
- RENDER_CREDITS: config.credits,
216
- SITE_TITLE: config.siteTitle,
217
- }),
218
- ].join('\n'),
219
- 'virtual:stl-docs-ai-chat': vmAiChatHandlerExport,
220
- }),
221
- );
196
+ const virtualModules = new Map<string, string>([
197
+ [
198
+ 'virtual:stl-docs-virtual-module',
199
+ buildVirtualModuleString({
200
+ TABS: config.tabs.map((tab) => ({ ...tab, link: withBase(tab.link) })),
201
+ SPLIT_TABS_ENABLED: config.splitTabsEnabled,
202
+ HEADER_LINKS: config.header.links.map((link) => ({ ...link, link: withBase(link.link) })),
203
+ HEADER_LAYOUT: config.header.layout,
204
+ ENABLE_CLIENT_ROUTER: config.enableClientRouter,
205
+ API_REFERENCE_BASE_PATH: apiReferenceBasePath ?? '/api',
206
+ ENABLE_PROSE_MARKDOWN_RENDERING: config.enableProseMarkdownRendering,
207
+ ENABLE_CONTEXT_MENU: !!config.contextMenu, // TODO: do not duplicate this between both virtual modules
208
+ CONTEXT_MENU_ENABLE_THIRD_PARTY:
209
+ (typeof config.contextMenu === 'object' ? config.contextMenu.thirdParty : null) ?? true,
210
+ RENDER_PAGE_DESCRIPTIONS: config.renderPageDescriptions,
211
+ FONTS: getFontRoles(config.fonts),
212
+ LINK_GROUP_TITLES_TO_OVERVIEW_PAGES: config.linkGroupTitlesToOverviewPages,
213
+ RENDER_CREDITS: config.credits,
214
+ SITE_TITLE: config.siteTitle,
215
+ }),
216
+ ],
217
+ ['virtual:stl-docs-ai-chat', vmAiChatHandlerExport],
218
+ ]);
219
+ if (config.aiChat) {
220
+ virtualModules.set(
221
+ 'virtual:stl-docs-ai-chat-examples',
222
+ generateExamplesVirtualModule(config.aiChat.examples),
223
+ );
224
+ }
222
225
 
223
226
  updateConfig({
224
227
  fonts: [...flattenFonts(config.fonts), ...(astroConfig?.fonts ?? [])],
@@ -239,17 +242,6 @@ function stainlessDocsIntegration(
239
242
  if (virtualModules.has(bare)) return virtualModules.get(bare);
240
243
  },
241
244
  },
242
- // Separate plugin for the examples because it has async resolution; not a simple string
243
- // like the above plugins
244
- ...(config.aiChat
245
- ? [
246
- await generateExamplesPlugin({
247
- projectName: config.apiReference?.stainlessProject ?? undefined,
248
- logger,
249
- exampleOverrides: config.aiChat.examples,
250
- }),
251
- ]
252
- : []),
253
245
  ],
254
246
  },
255
247
  build: {
@@ -311,11 +303,6 @@ export function stainlessDocs(config: StainlessDocsUserConfig): AstroIntegration
311
303
  integration: stainlessDocsMarkdownRenderer({ apiReferenceBasePath }),
312
304
  reason: 'disabled by experimental config "disableProseMarkdownRendering"',
313
305
  }),
314
- conditionalIntegration({
315
- condition: !config.experimental?.disableStainlessProseIndexing,
316
- integration: stainlessDocsAlgoliaProseIndexing({ apiReferenceBasePath }),
317
- reason: 'disabled by experimental config "disableStainlessProseIndexing"',
318
- }),
319
306
  conditionalIntegration({
320
307
  condition: !config.experimental?.disableStainlessProseIndexing,
321
308
  integration: stainlessDocsVectorProseIndexing(normalizedConfig, apiReferenceBasePath),
package/tsconfig.json CHANGED
@@ -7,7 +7,7 @@
7
7
  ".stl-docs/**/*",
8
8
  "./plugin/**/*"
9
9
  ],
10
- "exclude": ["dist", "**/preview.worker.docs.js"],
10
+ "exclude": ["dist"],
11
11
  "compilerOptions": {
12
12
  "jsx": "react-jsx",
13
13
  "jsxImportSource": "react"
@@ -1,99 +0,0 @@
1
- import crypto from 'crypto';
2
- import path from 'path';
3
- import { readdir, readFile, rm, writeFile } from 'fs/promises';
4
-
5
- type CacheResultSource = 'memory' | 'disk' | 'generation';
6
-
7
- type CacheGetResult<T> = {
8
- resultSource: CacheResultSource;
9
- data: T;
10
- filePath: string;
11
- };
12
-
13
- export class FileCache<Inputs extends Record<string, unknown>, Output> {
14
- private memoryCache: Map<string, Output> = new Map();
15
-
16
- private cacheDirectory: string | null = null;
17
-
18
- public setCacheDirectory(cacheDirectory: string) {
19
- this.cacheDirectory = cacheDirectory;
20
- }
21
- public getCacheDirectory() {
22
- if (!this.cacheDirectory) {
23
- console.error(`Tried to retrieve entry from cache, but no cache directory was set.`);
24
- process.exit(1);
25
- }
26
- return this.cacheDirectory;
27
- }
28
-
29
- private hashInputs(inputs: Inputs): string {
30
- return crypto
31
- .createHash('sha256')
32
- .update(JSON.stringify(inputs) + this.config.globalHashBits)
33
- .digest('hex')
34
- .slice(0, 10);
35
- }
36
-
37
- private getFileName(hash: string) {
38
- return `${hash}.json`;
39
- }
40
-
41
- public async cleanupUnusedFiles() {
42
- const allFiles = await readdir(this.getCacheDirectory());
43
- const usedFiles = Array.from(this.memoryCache.keys()).map((key) => this.getFileName(key));
44
- const unusedFiles = allFiles.filter((file) => !usedFiles.includes(file));
45
- await Promise.all(unusedFiles.map((file) => rm(path.join(this.getCacheDirectory(), file))));
46
- return {
47
- deletedCount: unusedFiles.length,
48
- };
49
- }
50
-
51
- public async get(inputs: Inputs): Promise<CacheGetResult<Output>> {
52
- const hash = this.hashInputs(inputs);
53
- const filePath = path.join(this.getCacheDirectory(), this.getFileName(hash));
54
-
55
- const memoryCacheResult = this.memoryCache.get(hash);
56
- if (memoryCacheResult) {
57
- return {
58
- resultSource: 'memory',
59
- data: memoryCacheResult,
60
- filePath,
61
- };
62
- }
63
-
64
- const getFromFileOrGenerate = async () => {
65
- try {
66
- const fileContents = await readFile(filePath, 'utf8');
67
- return {
68
- resultSource: 'disk' as const,
69
- data: JSON.parse(fileContents) as Output,
70
- filePath,
71
- };
72
- } catch {
73
- const data = await this.config.generate(inputs);
74
- await writeFile(filePath, JSON.stringify(data), 'utf8');
75
- return {
76
- resultSource: 'generation' as const,
77
- data,
78
- filePath,
79
- };
80
- }
81
- };
82
-
83
- const result = await getFromFileOrGenerate();
84
- this.memoryCache.set(hash, result.data);
85
- return result;
86
- }
87
-
88
- constructor(
89
- private config: {
90
- /**
91
- * Additional information to include in the hash of the inputs.
92
- * This is useful for cases where the inputs are the same, but the global state is different.
93
- * Eg: The preview worker source can be used here.
94
- */
95
- globalHashBits: string;
96
- generate: (inputs: Inputs) => Promise<Output>;
97
- },
98
- ) {}
99
- }