allow subscribing to OpenCtx controller before it's initialized

Previously, if some statically initialized value tried to read the `openCtx.controller` before it was set in `exposeOpenCtxClient`, it would throw. This required some contortions to ensure that it was initialized only after `exposeOpenCtxClient` was called. Now, anything can subscribe to `openctxController` at any time, and it will receive an emission of the OpenCtx controller when it's ready.

lib/shared/src/context/openctx/api.ts

@@ -1,29 +1,42 @@
 import type { Client } from '@openctx/client'
+import type { Observable } from 'observable-fns'
 import type * as vscode from 'vscode'
+import { fromLateSetSource, shareReplay, storeLastValue } from '../../misc/observable'
 
-type OpenCtxController = Pick<
+export type OpenCtxController = Pick<
     Client<vscode.Range>,
     'meta' | 'metaChanges' | 'mentions' | 'mentionsChanges' | 'items'
-> & {}
+>
 
-interface OpenCtx {
-    controller?: OpenCtxController
-    disposable?: vscode.Disposable
-}
+const _openctxController = fromLateSetSource<OpenCtxController>()
 
-export const openCtx: OpenCtx = {}
+export const openctxController: Observable<OpenCtxController> = _openctxController.observable.pipe(
+    shareReplay({ shouldCountRefs: false })
+)
 
 /**
- * Set the handle to the OpenCtx. If there is an existing handle it will be
- * disposed and replaced.
+ * Set the observable that will be used to provide the global {@link openctxController}.
  */
-export function setOpenCtx({ controller, disposable }: OpenCtx): void {
-    const { disposable: oldDisposable } = openCtx
+export function setOpenCtxControllerObservable(input: Observable<OpenCtxController>): void {
+    _openctxController.setSource(input)
+}
 
-    openCtx.controller = controller
-    openCtx.disposable = disposable
+const { value: syncValue } = storeLastValue(openctxController)
 
-    oldDisposable?.dispose()
+/**
+ * The current OpenCtx controller. Callers should use {@link openctxController} instead so that
+ * they react to changes. This function is provided for old call sites that haven't been updated
+ * to use an Observable.
+ *
+ * Callers should take care to avoid race conditions and prefer observing {@link openctxController}.
+ *
+ * Throws if the OpenCtx controller is not yet set.
+ */
+export function currentOpenCtxController(): OpenCtxController {
+    if (!syncValue.isSet) {
+        throw new Error('OpenCtx controller is not initialized')
+    }
+    return syncValue.last
 }
 
 export const REMOTE_REPOSITORY_PROVIDER_URI = 'internal-remote-repository-search'
@@ -32,3 +45,4 @@ export const REMOTE_DIRECTORY_PROVIDER_URI = 'internal-remote-directory-search'
 export const WEB_PROVIDER_URI = 'internal-web-provider'
 export const GIT_OPENCTX_PROVIDER_URI = 'internal-git-openctx-provider'
 export const CODE_SEARCH_PROVIDER_URI = 'internal-code-search-provider'
+export const RULES_PROVIDER_URI = 'internal-rules-provider'

lib/shared/src/context/openctx/context.ts

@@ -1,14 +1,14 @@
 import { URI } from 'vscode-uri'
 import { type ContextItemOpenCtx, ContextItemSource } from '../../codebase-context/messages'
-import { openCtx } from './api'
+import { currentOpenCtxController } from './api'
 
 // getContextForChatMessage returns context items for a given chat message from the OpenCtx providers.
 export const getContextForChatMessage = async (
     message: string,
     signal?: AbortSignal
 ): Promise<ContextItemOpenCtx[]> => {
     try {
-        const openCtxClient = openCtx.controller
+        const openCtxClient = currentOpenCtxController()
         if (!openCtxClient) {
             return []
         }
@@ -52,7 +52,7 @@ export const getContextForChatMessage = async (
                         content: item.ai?.content || '',
                         provider: 'openctx',
                         source: ContextItemSource.User, // To indicate that this is a user-added item.
-                    }) as ContextItemOpenCtx
+                    }) satisfies ContextItemOpenCtx
             )
     } catch {
         return []

lib/shared/src/index.ts

@@ -352,14 +352,16 @@ export * from './token'
 export * from './token/constants'
 export * from './configuration'
 export {
-    setOpenCtx,
-    openCtx,
+    setOpenCtxControllerObservable,
+    openctxController,
+    type OpenCtxController,
     REMOTE_REPOSITORY_PROVIDER_URI,
     REMOTE_FILE_PROVIDER_URI,
     REMOTE_DIRECTORY_PROVIDER_URI,
     WEB_PROVIDER_URI,
     GIT_OPENCTX_PROVIDER_URI,
     CODE_SEARCH_PROVIDER_URI,
+    currentOpenCtxController,
 } from './context/openctx/api'
 export * from './context/openctx/context'
 export * from './lexicalEditor/editorState'

lib/shared/src/mentions/api.test.ts

@@ -1,4 +1,6 @@
-import { describe, expect, it } from 'vitest'
+import { Observable } from 'observable-fns'
+import { describe, expect, it, vi } from 'vitest'
+import * as openctxAPI from '../context/openctx/api'
 import { firstValueFrom } from '../misc/observable'
 import {
     FILE_CONTEXT_MENTION_PROVIDER,
@@ -7,6 +9,15 @@ import {
 } from './api'
 
 describe('mentionProvidersMetadata', () => {
+    vi.spyOn(openctxAPI, 'openctxController', 'get').mockReturnValue(
+        Observable.of({
+            metaChanges: () => Observable.of([]),
+        } satisfies Pick<
+            openctxAPI.OpenCtxController,
+            'metaChanges'
+        > as unknown as openctxAPI.OpenCtxController)
+    )
+
     it('should return all providers when no options are provided', async () => {
         const providers = await firstValueFrom(mentionProvidersMetadata())
         expect(providers.length).toBeGreaterThanOrEqual(2)

lib/shared/src/mentions/api.ts

@@ -1,7 +1,7 @@
 import type { MetaResult } from '@openctx/client'
-import { Observable, map } from 'observable-fns'
-import { openCtx } from '../context/openctx/api'
-import { distinctUntilChanged } from '../misc/observable'
+import { type Observable, map } from 'observable-fns'
+import { openctxController } from '../context/openctx/api'
+import { distinctUntilChanged, switchMap } from '../misc/observable'
 
 /**
  * Props required by context item providers to return possible context items.
@@ -73,18 +73,17 @@ export function openCtxProviderMetadata(
 }
 
 function openCtxMentionProviders(): Observable<ContextMentionProviderMetadata[]> {
-    const controller = openCtx.controller
-    if (!controller) {
-        return Observable.of([])
-    }
-
-    return controller.metaChanges({}, {}).pipe(
-        map(providers =>
-            providers
-                .filter(provider => !!provider.mentions)
-                .map(openCtxProviderMetadata)
-                .sort((a, b) => (a.title > b.title ? 1 : -1))
-        ),
-        distinctUntilChanged()
+    return openctxController.pipe(
+        switchMap(c =>
+            c.metaChanges({}, {}).pipe(
+                map(providers =>
+                    providers
+                        .filter(provider => !!provider.mentions)
+                        .map(openCtxProviderMetadata)
+                        .sort((a, b) => (a.title > b.title ? 1 : -1))
+                ),
+                distinctUntilChanged()
+            )
+        )
     )
 }

vscode/src/chat/agentic/CodyTool.ts

@@ -6,9 +6,9 @@ import {
     type ContextMentionProviderMetadata,
     ProcessType,
     PromptString,
+    currentOpenCtxController,
     firstValueFrom,
     logDebug,
-    openCtx,
     parseMentionQuery,
     pendingOperation,
     ps,
@@ -305,7 +305,7 @@ export class OpenCtxTool extends CodyTool {
 
     async execute(span: Span, queries: string[]): Promise<ContextItem[]> {
         span.addEvent('executeOpenCtxTool')
-        const openCtxClient = openCtx.controller
+        const openCtxClient = currentOpenCtxController()
         if (!queries?.length || !openCtxClient) {
             return []
         }

vscode/src/chat/agentic/CodyToolProvider.test.ts

@@ -1,7 +1,8 @@
-import { type ContextItem, ContextItemSource, openCtx, ps } from '@sourcegraph/cody-shared'
+import { type ContextItem, ContextItemSource, ps } from '@sourcegraph/cody-shared'
 import { Observable } from 'observable-fns'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 import { URI } from 'vscode-uri'
+import * as openctxAPI from '../../../../lib/shared/src/context/openctx/api'
 import { mockLocalStorage } from '../../services/LocalStorageProvider'
 import type { ContextRetriever } from '../chat-view/ContextRetriever'
 import { CodyTool, type CodyToolConfig } from './CodyTool'
@@ -62,7 +63,7 @@ describe('CodyToolProvider', () => {
     beforeEach(() => {
         vi.clearAllMocks()
         CodyToolProvider.initialize(mockContextRetriever)
-        openCtx.controller = mockController
+        vi.spyOn(openctxAPI, 'openctxController', 'get').mockReturnValue(Observable.of(mockController))
     })
 
     it('should register default tools on initialization', () => {
@@ -73,9 +74,9 @@ describe('CodyToolProvider', () => {
     })
 
     it('should set up OpenCtx provider listener and build OpenCtx tools from provider metadata', async () => {
-        openCtx.controller = mockController
         CodyToolProvider.setupOpenCtxProviderListener()
-        expect(openCtx.controller?.metaChanges).toHaveBeenCalled()
+        await new Promise(resolve => setTimeout(resolve, 0))
+        expect(mockController.metaChanges).toHaveBeenCalled()
         // Wait for the observable to emit
         await new Promise(resolve => setTimeout(resolve, 0))
 

vscode/src/chat/agentic/CodyToolProvider.ts

@@ -4,9 +4,10 @@ import {
     PromptString,
     type Unsubscribable,
     isDefined,
-    openCtx,
     openCtxProviderMetadata,
+    openctxController,
     ps,
+    switchMap,
 } from '@sourcegraph/cody-shared'
 import { map } from 'observable-fns'
 import type { ContextRetriever } from '../chat-view/ContextRetriever'
@@ -190,10 +191,19 @@ export class CodyToolProvider {
         if (provider && !CodyToolProvider.configSubscription) {
             CodyToolProvider.configSubscription = toolboxManager.observable.subscribe({})
         }
-        if (provider && !CodyToolProvider.openCtxSubscription && openCtx.controller) {
-            CodyToolProvider.openCtxSubscription = openCtx.controller
-                .metaChanges({}, {})
-                .pipe(map(providers => providers.filter(p => !!p.mentions).map(openCtxProviderMetadata)))
+        if (provider && !CodyToolProvider.openCtxSubscription) {
+            CodyToolProvider.openCtxSubscription = openctxController
+                .pipe(
+                    switchMap(c =>
+                        c
+                            .metaChanges({}, {})
+                            .pipe(
+                                map(providers =>
+                                    providers.filter(p => !!p.mentions).map(openCtxProviderMetadata)
+                                )
+                            )
+                    )
+                )
                 .subscribe(providerMeta => provider.factory.createOpenCtxTools(providerMeta))
         }
     }

vscode/src/chat/context/chatContext.ts

@@ -11,12 +11,12 @@ import {
     SYMBOL_CONTEXT_MENTION_PROVIDER,
     clientCapabilities,
     combineLatest,
+    currentOpenCtxController,
     firstResultFromOperation,
     fromVSCodeEvent,
     isAbortError,
     isError,
     mentionProvidersMetadata,
-    openCtx,
     pendingOperation,
     promiseFactoryToObservable,
     skipPendingOperation,
@@ -153,11 +153,7 @@ export async function getChatContextItemsForMention(
         }
 
         default: {
-            if (!openCtx.controller) {
-                return []
-            }
-
-            const items = await openCtx.controller.mentions(
+            const items = await currentOpenCtxController().mentions(
                 {
                     query: mentionQuery.text,
                     ...(await firstResultFromOperation(activeEditorContextForOpenCtxMentions)),

vscode/src/chat/initialContext.ts

@@ -20,7 +20,7 @@ import {
     fromVSCodeEvent,
     isDotCom,
     isError,
-    openCtx,
+    openctxController,
     pendingOperation,
     shareReplay,
     startWith,
@@ -265,53 +265,52 @@ export function getCorpusContextItemsForEditorState(): Observable<
 }
 
 function getOpenCtxContextItems(): Observable<ContextItem[] | typeof pendingOperation> {
-    const openctxController = openCtx.controller
-    if (!openctxController) {
-        return Observable.of([])
-    }
-
-    return openctxController.metaChanges({}).pipe(
-        switchMap((providers): Observable<ContextItem[] | typeof pendingOperation> => {
-            const providersWithAutoInclude = providers.filter(meta => meta.mentions?.autoInclude)
-            if (providersWithAutoInclude.length === 0) {
-                return Observable.of([])
-            }
-
-            return activeTextEditor.pipe(
-                debounceTime(50),
-                switchMap(() => activeEditorContextForOpenCtxMentions),
-                switchMap(activeEditorContext => {
-                    if (activeEditorContext === pendingOperation) {
-                        return Observable.of(pendingOperation)
-                    }
-                    if (isError(activeEditorContext)) {
+    return openctxController.pipe(
+        switchMap(c =>
+            c.metaChanges({}, {}).pipe(
+                switchMap((providers): Observable<ContextItem[] | typeof pendingOperation> => {
+                    const providersWithAutoInclude = providers.filter(meta => meta.mentions?.autoInclude)
+                    if (providersWithAutoInclude.length === 0) {
                         return Observable.of([])
                     }
-                    return combineLatest(
-                        ...providersWithAutoInclude.map(provider =>
-                            openctxController.mentionsChanges(
-                                { ...activeEditorContext, autoInclude: true },
-                                provider
-                            )
-                        )
-                    ).pipe(
-                        map(mentionsResults =>
-                            mentionsResults.flat().map(
-                                mention =>
-                                    ({
-                                        ...mention,
-                                        provider: 'openctx',
-                                        type: 'openctx',
-                                        uri: URI.parse(mention.uri),
-                                        source: ContextItemSource.Initial,
-                                        mention, // include the original mention to pass to `items` later
-                                    }) satisfies ContextItem
+
+                    return activeTextEditor.pipe(
+                        debounceTime(50),
+                        switchMap(() => activeEditorContextForOpenCtxMentions),
+                        switchMap(activeEditorContext => {
+                            if (activeEditorContext === pendingOperation) {
+                                return Observable.of(pendingOperation)
+                            }
+                            if (isError(activeEditorContext)) {
+                                return Observable.of([])
+                            }
+                            return combineLatest(
+                                ...providersWithAutoInclude.map(provider =>
+                                    c.mentionsChanges(
+                                        { ...activeEditorContext, autoInclude: true },
+                                        provider
+                                    )
+                                )
+                            ).pipe(
+                                map(mentionsResults =>
+                                    mentionsResults.flat().map(
+                                        mention =>
+                                            ({
+                                                ...mention,
+                                                provider: 'openctx',
+                                                type: 'openctx',
+                                                uri: URI.parse(mention.uri),
+                                                source: ContextItemSource.Initial,
+                                                mention, // include the original mention to pass to `items` later
+                                            }) satisfies ContextItem
+                                    )
+                                ),
+                                startWith(pendingOperation)
                             )
-                        ),
-                        startWith(pendingOperation)
+                        })
                     )
                 })
             )
-        })
+        )
     )
 }