@theia/ai-copilot 1.71.0-next.8 → 1.71.0

package/src/node/copilot-auth-service-impl.spec.ts

@@ -18,6 +18,7 @@ import { expect } from 'chai';
 import * as sinon from 'sinon';
 import { Container } from '@theia/core/shared/inversify';
 import { KeyStoreService } from '@theia/core/lib/common/key-store';
+import { CopilotOAuthConfig, DEFAULT_COPILOT_OAUTH_CONFIG } from '../common/copilot-oauth-config';
 import { CopilotAuthServiceImpl } from './copilot-auth-service-impl';
 
 describe('CopilotAuthServiceImpl', () => {
@@ -40,6 +41,7 @@ describe('CopilotAuthServiceImpl', () => {
         };
 
         container.bind(KeyStoreService).toConstantValue(keyStoreService);
+        container.bind(CopilotOAuthConfig).toConstantValue(DEFAULT_COPILOT_OAUTH_CONFIG);
         container.bind(CopilotAuthServiceImpl).toSelf().inSingletonScope();
 
         authService = container.get(CopilotAuthServiceImpl);
@@ -183,3 +185,60 @@ describe('CopilotAuthServiceImpl', () => {
         });
     });
 });
+
+describe('CopilotAuthServiceImpl with custom CopilotOAuthConfig', () => {
+
+    const customConfig: CopilotOAuthConfig = {
+        clientId: 'CustomClientId',
+        userAgent: 'MyApp/1.0.0',
+        keystoreService: 'myapp-copilot',
+        keystoreAccount: 'myapp-github'
+    };
+
+    let authService: CopilotAuthServiceImpl;
+    let getPasswordStub: sinon.SinonStub;
+    let deletePasswordStub: sinon.SinonStub;
+
+    beforeEach(() => {
+        const container = new Container();
+
+        getPasswordStub = sinon.stub().resolves(undefined);
+        deletePasswordStub = sinon.stub().resolves(true);
+        const keyStoreService: KeyStoreService = {
+            setPassword: sinon.stub().resolves() as KeyStoreService['setPassword'],
+            getPassword: getPasswordStub as KeyStoreService['getPassword'],
+            deletePassword: deletePasswordStub as KeyStoreService['deletePassword'],
+            findPassword: sinon.stub().resolves(undefined) as KeyStoreService['findPassword'],
+            findCredentials: sinon.stub().resolves([]) as KeyStoreService['findCredentials'],
+            keys: sinon.stub().resolves([]) as KeyStoreService['keys']
+        };
+
+        container.bind(KeyStoreService).toConstantValue(keyStoreService);
+        container.bind(CopilotOAuthConfig).toConstantValue(customConfig);
+        container.bind(CopilotAuthServiceImpl).toSelf().inSingletonScope();
+
+        authService = container.get(CopilotAuthServiceImpl);
+    });
+
+    afterEach(() => {
+        sinon.restore();
+    });
+
+    it('should read credentials from the custom keystore service and account', async () => {
+        getPasswordStub.resolves(JSON.stringify({
+            accessToken: 'gho_custom_token',
+            accountLabel: 'customuser'
+        }));
+
+        const token = await authService.getAccessToken();
+
+        expect(token).to.equal('gho_custom_token');
+        expect(getPasswordStub.calledOnceWith(customConfig.keystoreService, customConfig.keystoreAccount)).to.be.true;
+    });
+
+    it('should delete credentials from the custom keystore service and account on sign out', async () => {
+        await authService.signOut();
+
+        expect(deletePasswordStub.calledOnceWith(customConfig.keystoreService, customConfig.keystoreAccount)).to.be.true;
+    });
+});

package/src/node/copilot-auth-service-impl.ts

@@ -23,13 +23,10 @@ import {
     CopilotAuthState,
     DeviceCodeResponse
 } from '../common/copilot-auth-service';
-import { COPILOT_USER_AGENT } from '../common';
+import { CopilotOAuthConfig } from '../common/copilot-oauth-config';
 
-const COPILOT_CLIENT_ID = 'Ov23liS2vINy9VOAweyv'; // 'Theia Copilot OAuth Access' Client ID
 const COPILOT_SCOPE = 'read:user';
 const COPILOT_GRANT_TYPE = 'urn:ietf:params:oauth:grant-type:device_code';
-const KEYSTORE_SERVICE = 'theia-copilot-auth';
-const KEYSTORE_ACCOUNT = 'github-copilot';
 
 /**
  * Maximum number of polling attempts for token retrieval.
@@ -53,6 +50,9 @@ export class CopilotAuthServiceImpl implements CopilotAuthService {
     @inject(KeyStoreService)
     protected readonly keyStoreService: KeyStoreService;
 
+    @inject(CopilotOAuthConfig)
+    protected readonly oauthConfig: CopilotOAuthConfig;
+
     protected client: CopilotAuthServiceClient | undefined;
     protected cachedState: CopilotAuthState | undefined;
 
@@ -87,10 +87,10 @@ export class CopilotAuthServiceImpl implements CopilotAuthService {
             headers: {
                 'Accept': 'application/json',
                 'Content-Type': 'application/json',
-                'User-Agent': COPILOT_USER_AGENT
+                'User-Agent': this.oauthConfig.userAgent
             },
             body: JSON.stringify({
-                client_id: COPILOT_CLIENT_ID,
+                client_id: this.oauthConfig.clientId,
                 scope: COPILOT_SCOPE
             })
         });
@@ -117,10 +117,10 @@ export class CopilotAuthServiceImpl implements CopilotAuthService {
                 headers: {
                     'Accept': 'application/json',
                     'Content-Type': 'application/json',
-                    'User-Agent': COPILOT_USER_AGENT
+                    'User-Agent': this.oauthConfig.userAgent
                 },
                 body: JSON.stringify({
-                    client_id: COPILOT_CLIENT_ID,
+                    client_id: this.oauthConfig.clientId,
                     device_code: deviceCode,
                     grant_type: COPILOT_GRANT_TYPE
                 })
@@ -149,8 +149,8 @@ export class CopilotAuthServiceImpl implements CopilotAuthService {
                 };
 
                 await this.keyStoreService.setPassword(
-                    KEYSTORE_SERVICE,
-                    KEYSTORE_ACCOUNT,
+                    this.oauthConfig.keystoreService,
+                    this.oauthConfig.keystoreAccount,
                     JSON.stringify(credentials)
                 );
 
@@ -199,7 +199,7 @@ export class CopilotAuthServiceImpl implements CopilotAuthService {
             const response = await fetch(`${apiBaseUrl}/user`, {
                 headers: {
                     'Authorization': `Bearer ${accessToken}`,
-                    'User-Agent': COPILOT_USER_AGENT,
+                    'User-Agent': this.oauthConfig.userAgent,
                     'Accept': 'application/vnd.github.v3+json'
                 }
             });
@@ -220,13 +220,13 @@ export class CopilotAuthServiceImpl implements CopilotAuthService {
         }
 
         try {
-            const stored = await this.keyStoreService.getPassword(KEYSTORE_SERVICE, KEYSTORE_ACCOUNT);
+            const stored = await this.keyStoreService.getPassword(this.oauthConfig.keystoreService, this.oauthConfig.keystoreAccount);
             if (stored) {
                 const credentials: StoredCredentials = JSON.parse(stored);
                 // Tokens from the current OAuth App start with 'gho_'; other prefixes (e.g. 'ghu_') indicate a token from the previous GitHub App (Iv-prefixed client ID).
                 if (!credentials.accessToken.startsWith('gho_')) {
                     console.info('Copilot: clearing outdated GitHub App token. Please sign in again.');
-                    await this.keyStoreService.deletePassword(KEYSTORE_SERVICE, KEYSTORE_ACCOUNT);
+                    await this.keyStoreService.deletePassword(this.oauthConfig.keystoreService, this.oauthConfig.keystoreAccount);
                     this.cachedState = { isAuthenticated: false, migrationRequired: true };
                     return this.cachedState;
                 }
@@ -247,7 +247,7 @@ export class CopilotAuthServiceImpl implements CopilotAuthService {
 
     async getAccessToken(): Promise<string | undefined> {
         try {
-            const stored = await this.keyStoreService.getPassword(KEYSTORE_SERVICE, KEYSTORE_ACCOUNT);
+            const stored = await this.keyStoreService.getPassword(this.oauthConfig.keystoreService, this.oauthConfig.keystoreAccount);
             if (stored) {
                 const credentials: StoredCredentials = JSON.parse(stored);
                 return credentials.accessToken;
@@ -260,7 +260,7 @@ export class CopilotAuthServiceImpl implements CopilotAuthService {
 
     async signOut(): Promise<void> {
         try {
-            await this.keyStoreService.deletePassword(KEYSTORE_SERVICE, KEYSTORE_ACCOUNT);
+            await this.keyStoreService.deletePassword(this.oauthConfig.keystoreService, this.oauthConfig.keystoreAccount);
         } catch (error) {
             console.warn('Failed to delete Copilot credentials:', error);
         }

package/src/node/copilot-backend-module.ts

@@ -24,6 +24,7 @@ import {
     COPILOT_AUTH_SERVICE_PATH,
     CopilotAuthServiceClient
 } from '../common';
+import { CopilotOAuthConfig, DEFAULT_COPILOT_OAUTH_CONFIG } from '../common/copilot-oauth-config';
 import { CopilotLanguageModelsManagerImpl } from './copilot-language-models-manager-impl';
 import { CopilotAuthServiceImpl } from './copilot-auth-service-impl';
 
@@ -54,5 +55,6 @@ const copilotConnectionModule = ConnectionContainerModule.create(({ bind }) => {
 });
 
 export default new ContainerModule(bind => {
+    bind(CopilotOAuthConfig).toConstantValue(DEFAULT_COPILOT_OAUTH_CONFIG);
     bind(ConnectionContainerModule).toConstantValue(copilotConnectionModule);
 });

package/src/node/copilot-language-model.ts

@@ -23,7 +23,6 @@ import {
     LanguageModelResponse,
     LanguageModelStatus,
     LanguageModelTextResponse,
-    TokenUsageService,
     UserRequest
 } from '@theia/ai-core';
 import { CancellationToken } from '@theia/core';
@@ -31,7 +30,7 @@ import OpenAI from 'openai';
 import { RunnableToolFunctionWithoutParse } from 'openai/lib/RunnableFunction';
 import { ChatCompletionMessageParam } from 'openai/resources';
 import { StreamingAsyncIterator } from '@theia/ai-openai/lib/node/openai-streaming-iterator';
-import { COPILOT_PROVIDER_ID, COPILOT_USER_AGENT, getCopilotApiBaseUrl } from '../common';
+import { COPILOT_PROVIDER_ID, getCopilotApiBaseUrl } from '../common';
 import type { RunnerOptions } from 'openai/lib/AbstractChatCompletionRunner';
 import type { ChatCompletionStream } from 'openai/lib/ChatCompletionStream';
 
@@ -54,7 +53,7 @@ export class CopilotLanguageModel implements LanguageModel {
         public maxRetries: number,
         protected readonly accessTokenProvider: () => Promise<string | undefined>,
         protected readonly enterpriseUrlProvider: () => string | undefined,
-        protected readonly tokenUsageService?: TokenUsageService
+        protected readonly userAgentProvider: () => string,
     ) { }
 
     protected getSettings(request: LanguageModelRequest): Record<string, unknown> {
@@ -107,7 +106,7 @@ export class CopilotLanguageModel implements LanguageModel {
         }
 
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        return { stream: new StreamingAsyncIterator(runner as any, request.requestId, cancellationToken, this.tokenUsageService, this.id) };
+        return { stream: new StreamingAsyncIterator(runner as any, cancellationToken) };
     }
 
     protected async handleNonStreamingRequest(openai: OpenAI, request: UserRequest): Promise<LanguageModelTextResponse> {
@@ -120,19 +119,12 @@ export class CopilotLanguageModel implements LanguageModel {
 
         const message = response.choices[0].message;
 
-        if (this.tokenUsageService && response.usage) {
-            await this.tokenUsageService.recordTokenUsage(
-                this.id,
-                {
-                    inputTokens: response.usage.prompt_tokens,
-                    outputTokens: response.usage.completion_tokens,
-                    requestId: request.requestId
-                }
-            );
-        }
-
         return {
-            text: message.content ?? ''
+            text: message.content ?? '',
+            usage: response.usage ? {
+                input_tokens: response.usage.prompt_tokens,
+                output_tokens: response.usage.completion_tokens,
+            } : undefined
         };
     }
 
@@ -150,20 +142,13 @@ export class CopilotLanguageModel implements LanguageModel {
             console.error('Error in Copilot chat completion:', JSON.stringify(message));
         }
 
-        if (this.tokenUsageService && result.usage) {
-            await this.tokenUsageService.recordTokenUsage(
-                this.id,
-                {
-                    inputTokens: result.usage.prompt_tokens,
-                    outputTokens: result.usage.completion_tokens,
-                    requestId: request.requestId
-                }
-            );
-        }
-
         return {
             content: message.content ?? '',
-            parsed: message.parsed
+            parsed: message.parsed,
+            usage: result.usage ? {
+                input_tokens: result.usage.prompt_tokens,
+                output_tokens: result.usage.completion_tokens,
+            } : undefined
         };
     }
 
@@ -191,7 +176,7 @@ export class CopilotLanguageModel implements LanguageModel {
             apiKey: accessToken,
             baseURL,
             defaultHeaders: {
-                'User-Agent': COPILOT_USER_AGENT,
+                'User-Agent': this.userAgentProvider(),
                 'Openai-Intent': 'conversation-edits',
                 'X-Initiator': 'user'
             }

package/src/node/copilot-language-models-manager-impl.ts

@@ -14,10 +14,11 @@
 // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
 // *****************************************************************************
 
-import { LanguageModelRegistry, LanguageModelStatus, TokenUsageService } from '@theia/ai-core';
+import { LanguageModelRegistry, LanguageModelStatus } from '@theia/ai-core';
 import { Disposable, DisposableCollection } from '@theia/core';
 import { inject, injectable, postConstruct } from '@theia/core/shared/inversify';
-import { CopilotLanguageModelsManager, CopilotModelDescription, COPILOT_PROVIDER_ID, COPILOT_USER_AGENT, getCopilotApiBaseUrl } from '../common';
+import { CopilotLanguageModelsManager, CopilotModelDescription, COPILOT_PROVIDER_ID, getCopilotApiBaseUrl } from '../common';
+import { CopilotOAuthConfig } from '../common/copilot-oauth-config';
 import { CopilotLanguageModel } from './copilot-language-model';
 import { CopilotAuthServiceImpl } from './copilot-auth-service-impl';
 
@@ -31,12 +32,12 @@ export class CopilotLanguageModelsManagerImpl implements CopilotLanguageModelsMa
     @inject(LanguageModelRegistry)
     protected readonly languageModelRegistry: LanguageModelRegistry;
 
-    @inject(TokenUsageService)
-    protected readonly tokenUsageService: TokenUsageService;
-
     @inject(CopilotAuthServiceImpl)
     protected readonly authService: CopilotAuthServiceImpl;
 
+    @inject(CopilotOAuthConfig)
+    protected readonly oauthConfig: CopilotOAuthConfig;
+
     protected enterpriseUrl: string | undefined;
     protected readonly toDispose = new DisposableCollection();
 
@@ -92,7 +93,7 @@ export class CopilotLanguageModelsManagerImpl implements CopilotLanguageModelsMa
                         modelDescription.maxRetries,
                         () => this.authService.getAccessToken(),
                         () => this.enterpriseUrl,
-                        this.tokenUsageService
+                        () => this.oauthConfig.userAgent
                     )
                 ]);
             }
@@ -128,7 +129,7 @@ export class CopilotLanguageModelsManagerImpl implements CopilotLanguageModelsMa
             const response = await fetch(`${baseURL}/models`, {
                 headers: {
                     'Authorization': `Bearer ${accessToken}`,
-                    'User-Agent': COPILOT_USER_AGENT,
+                    'User-Agent': this.oauthConfig.userAgent,
                     'Accept': 'application/json'
                 }
             });