feat(agent): support add request headers in config. (#351)

* feat(agent): support add request headers in config.

* fix(agent): fix authRequired event condition.
release-0.0
Zhiming Ma 2023-08-11 19:39:17 +08:00 committed by GitHub
parent ad85befc7c
commit d497217427
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 143 additions and 54 deletions

View File

@ -36,7 +36,8 @@
"axios": "^1.4.0",
"chokidar": "^3.5.3",
"deep-equal": "^2.2.1",
"deepmerge": "^4.3.1",
"deepmerge-ts": "^5.1.0",
"dot-prop": "^8.0.2",
"fast-levenshtein": "^3.0.0",
"form-data": "^4.0.0",
"fs-extra": "^11.1.1",

View File

@ -4,12 +4,12 @@ import {
CompletionResponse as ApiCompletionResponse,
} from "./generated";
import { AgentConfig } from "./AgentConfig";
import { AgentConfig, PartialAgentConfig } from "./AgentConfig";
export type AgentInitOptions = {
config: Partial<AgentConfig>;
export type AgentInitOptions = Partial<{
config: PartialAgentConfig;
client: string;
};
}>;
export type CompletionRequest = {
filepath: string;
@ -24,19 +24,44 @@ export type CompletionResponse = ApiCompletionResponse;
export type LogEventRequest = ApiLogEventRequest;
/**
* `notInitialized`: When the agent is not initialized.
* `ready`: When the agent get a valid response from the server, and is ready to use.
* `disconnected`: When the agent failed to connect to the server.
* `unauthorized`: When the server is set to a Tabby Cloud endpoint that requires auth,
* and no `Authorization` request header is provided in the agent config,
* and user has not completed the auth flow or the auth token is expired.
* See also `requestAuthUrl` and `waitForAuthToken`.
*/
export type AgentStatus = "notInitialized" | "ready" | "disconnected" | "unauthorized";
export interface AgentFunction {
initialize(options: Partial<AgentInitOptions>): Promise<boolean>;
updateConfig(config: Partial<AgentConfig>): Promise<boolean>;
/**
* Initialize agent. Client should call this method before calling any other methods.
* @param options
*/
initialize(options: AgentInitOptions): Promise<boolean>;
/**
* @returns the current config
*
* Configuration precedence:
* The agent configuration has the following levels, will be deep merged in the order:
* 1. Default config
* 2. User config file `~/.tabby/agent/config.toml` (not available in browser)
* 3. Agent `initialize` and `updateConfig` methods
*
* This method will update the 3rd level config.
* @param key the key of the config to update, can be nested with dot, e.g. `server.endpoint`
* @param value the value to set
*/
updateConfig(key: string, value: any): Promise<boolean>;
/**
* Clear the 3rd level config.
* @param key the key of the config to clear, can be nested with dot, e.g. `server.endpoint`
*/
clearConfig(key: string): Promise<boolean>;
/**
* @returns the current config
*/
getConfig(): AgentConfig;
@ -46,15 +71,16 @@ export interface AgentFunction {
getStatus(): AgentStatus;
/**
* @returns the auth url for redirecting, and the code for next step `waitingForAuth`, only return value when
* `AgentStatus` is `unauthorized`, return null otherwise
* Request auth url for Tabby Cloud endpoint. Only return value when the `AgentStatus` is `unauthorized`.
* Otherwise, return null. See also `AgentStatus`.
* @returns the auth url for redirecting, and the code for next step `waitingForAuth`
* @throws Error if agent is not initialized
*/
requestAuthUrl(): CancelablePromise<{ authUrl: string; code: string } | null>;
/**
* Wait for auth token to be ready after redirecting user to auth url,
* returns nothing, but `AgentStatus` will change to `ready` if resolved successfully
* returns nothing, but `AgentStatus` will change to `ready` if resolved successfully.
* @param code from `requestAuthUrl`
* @throws Error if agent is not initialized
*/
@ -83,6 +109,10 @@ export type ConfigUpdatedEvent = {
event: "configUpdated";
config: AgentConfig;
};
/**
* This event is emitted when the server is set to a Tabby Cloud endpoint that requires auth,
* and no `Authorization` request header is provided in the agent config.
*/
export type AuthRequiredEvent = {
event: "authRequired";
server: AgentConfig["server"];

View File

@ -3,6 +3,7 @@ import { isBrowser } from "./env";
export type AgentConfig = {
server: {
endpoint: string;
requestHeaders: Record<string, string>;
};
completion: {
maxPrefixLines: number;
@ -16,9 +17,20 @@ export type AgentConfig = {
};
};
type RecursivePartial<T> = {
[P in keyof T]?: T[P] extends (infer U)[]
? RecursivePartial<U>[]
: T[P] extends object | undefined
? RecursivePartial<T[P]>
: T[P];
};
export type PartialAgentConfig = RecursivePartial<AgentConfig>;
export const defaultAgentConfig: AgentConfig = {
server: {
endpoint: "http://localhost:8080",
requestHeaders: {},
},
completion: {
maxPrefixLines: 20,
@ -42,7 +54,7 @@ export const userAgentConfig = isBrowser
class ConfigFile extends EventEmitter {
filepath: string;
data: Partial<AgentConfig> = {};
data: PartialAgentConfig = {};
watcher: ReturnType<typeof chokidar.watch> | null = null;
logger = require("./logger").rootLogger.child({ component: "ConfigFile" });
@ -51,7 +63,7 @@ export const userAgentConfig = isBrowser
this.filepath = filepath;
}
get config(): Partial<AgentConfig> {
get config(): PartialAgentConfig {
return this.data;
}

View File

@ -1,7 +1,8 @@
import { EventEmitter } from "events";
import { v4 as uuid } from "uuid";
import deepEqual from "deep-equal";
import deepMerge from "deepmerge";
import { deepmerge } from "deepmerge-ts";
import { getProperty, setProperty, deleteProperty } from "dot-prop";
import { TabbyApi, CancelablePromise } from "./generated";
import { cancelable, splitLines, isBlank } from "./utils";
import {
@ -14,7 +15,7 @@ import {
LogEventRequest,
} from "./Agent";
import { Auth } from "./Auth";
import { AgentConfig, defaultAgentConfig, userAgentConfig } from "./AgentConfig";
import { AgentConfig, PartialAgentConfig, defaultAgentConfig, userAgentConfig } from "./AgentConfig";
import { CompletionCache } from "./CompletionCache";
import { DataStore } from "./dataStore";
import { postprocess, preCacheProcess } from "./postprocess";
@ -26,15 +27,15 @@ import { AnonymousUsageLogger } from "./AnonymousUsageLogger";
* so it is not suitable for cli, but only used when imported as module by other js project.
*/
export type TabbyAgentOptions = {
dataStore: DataStore;
dataStore?: DataStore;
};
export class TabbyAgent extends EventEmitter implements Agent {
private readonly logger = rootLogger.child({ component: "TabbyAgent" });
private anonymousUsageLogger: AnonymousUsageLogger;
private config: AgentConfig = defaultAgentConfig;
private userConfig: Partial<AgentConfig> = {}; // config from `~/.tabby/agent/config.toml`
private clientConfig: Partial<AgentConfig> = {}; // config from `initialize` and `updateConfig` method
private userConfig: PartialAgentConfig = {}; // config from `~/.tabby/agent/config.toml`
private clientConfig: PartialAgentConfig = {}; // config from `initialize` and `updateConfig` method
private status: AgentStatus = "notInitialized";
private api: TabbyApi;
private auth: Auth;
@ -54,7 +55,7 @@ export class TabbyAgent extends EventEmitter implements Agent {
}, TabbyAgent.tryConnectInterval);
}
static async create(options?: Partial<TabbyAgentOptions>): Promise<TabbyAgent> {
static async create(options?: TabbyAgentOptions): Promise<TabbyAgent> {
const agent = new TabbyAgent();
agent.dataStore = options?.dataStore;
agent.anonymousUsageLogger = await AnonymousUsageLogger.create({ dataStore: options?.dataStore });
@ -62,13 +63,17 @@ export class TabbyAgent extends EventEmitter implements Agent {
}
private async applyConfig() {
this.config = deepMerge.all<AgentConfig>([defaultAgentConfig, this.userConfig, this.clientConfig]);
this.config = deepmerge(defaultAgentConfig, this.userConfig, this.clientConfig);
allLoggers.forEach((logger) => (logger.level = this.config.logs.level));
this.anonymousUsageLogger.disabled = this.config.anonymousUsageTracking.disable;
if (this.config.server.requestHeaders["Authorization"] === undefined) {
if (this.config.server.endpoint !== this.auth?.endpoint) {
this.auth = await Auth.create({ endpoint: this.config.server.endpoint, dataStore: this.dataStore });
this.auth.on("updated", this.setupApi.bind(this));
}
} else {
this.auth = null;
}
await this.setupApi();
}
@ -76,6 +81,7 @@ export class TabbyAgent extends EventEmitter implements Agent {
this.api = new TabbyApi({
BASE: this.config.server.endpoint.replace(/\/+$/, ""), // remove trailing slash
TOKEN: this.auth?.token,
HEADERS: this.config.server.requestHeaders,
});
await this.healthCheck();
}
@ -86,12 +92,20 @@ export class TabbyAgent extends EventEmitter implements Agent {
const event: AgentEvent = { event: "statusChanged", status };
this.logger.debug({ event }, "Status changed");
super.emit("statusChanged", event);
if (this.status === "unauthorized") {
this.emitAuthRequired();
}
if (this.status == "ready") {
this.anonymousUsageLogger.uniqueEvent("AgentConnected");
}
}
}
private emitAuthRequired() {
const event: AgentEvent = { event: "authRequired", server: this.config.server };
super.emit("authRequired", event);
}
private callApi<Request, Response>(
api: (request: Request) => CancelablePromise<Response>,
request: Request,
@ -108,7 +122,12 @@ export class TabbyAgent extends EventEmitter implements Agent {
.catch((error) => {
if (!!error.isCancelled) {
this.logger.debug({ api: api.name, error }, "API request canceled");
} else if (error.name === "ApiError" && [401, 403, 405].indexOf(error.status) !== -1) {
} else if (
error.name === "ApiError" &&
[401, 403, 405].indexOf(error.status) !== -1 &&
new URL(this.config.server.endpoint).hostname.endsWith("app.tabbyml.com") &&
this.config.server.requestHeaders["Authorization"] === undefined
) {
this.logger.debug({ api: api.name, error }, "API unauthorized");
this.changeStatus("unauthorized");
} else if (error.name === "ApiError") {
@ -144,7 +163,7 @@ export class TabbyAgent extends EventEmitter implements Agent {
};
}
public async initialize(options: Partial<AgentInitOptions>): Promise<boolean> {
public async initialize(options: AgentInitOptions): Promise<boolean> {
if (options.client) {
// Client info is only used in logging for now
// `pino.Logger.setBindings` is not present in the browser
@ -161,35 +180,40 @@ export class TabbyAgent extends EventEmitter implements Agent {
userAgentConfig.watch();
}
if (options.config) {
this.clientConfig = deepMerge(this.clientConfig, options.config);
this.clientConfig = options.config;
}
await this.applyConfig();
if (this.status === "unauthorized") {
const event: AgentEvent = { event: "authRequired", server: this.config.server };
super.emit("authRequired", event);
}
await this.anonymousUsageLogger.uniqueEvent("AgentInitialized");
this.logger.debug({ options }, "Initialized");
return this.status !== "notInitialized";
}
public async updateConfig(config: Partial<AgentConfig>): Promise<boolean> {
const mergedConfig = deepMerge(this.clientConfig, config);
if (!deepEqual(this.clientConfig, mergedConfig)) {
const serverUpdated = !deepEqual(this.config.server, mergedConfig.server);
this.clientConfig = mergedConfig;
public async updateConfig(key: string, value: any): Promise<boolean> {
const current = getProperty(this.clientConfig, key);
if (!deepEqual(current, value)) {
if (value === undefined) {
deleteProperty(this.clientConfig, key);
} else {
setProperty(this.clientConfig, key, value);
}
const prevStatus = this.status;
await this.applyConfig();
// If status unchanged, `authRequired` will not be emitted when `applyConfig`,
// so we need to emit it manually.
if (key.startsWith("server") && prevStatus === "unauthorized" && this.status === "unauthorized") {
this.emitAuthRequired();
}
const event: AgentEvent = { event: "configUpdated", config: this.config };
this.logger.debug({ event }, "Config updated");
super.emit("configUpdated", event);
if (serverUpdated && this.status === "unauthorized") {
const event: AgentEvent = { event: "authRequired", server: this.config.server };
super.emit("authRequired", event);
}
}
return true;
}
public async clearConfig(key: string): Promise<boolean> {
return await this.updateConfig(key, undefined);
}
public getConfig(): AgentConfig {
return this.config;
}

View File

@ -11,6 +11,6 @@ export {
LogEventRequest,
agentEventNames,
} from "./Agent";
export { AgentConfig } from "./AgentConfig";
export { AgentConfig, PartialAgentConfig } from "./AgentConfig";
export { DataStore } from "./dataStore";
export { CancelablePromise } from "./generated";

View File

@ -1,9 +1,9 @@
import { ExtensionContext, workspace, env, version } from "vscode";
import { TabbyAgent, AgentConfig, DataStore } from "tabby-agent";
import { TabbyAgent, PartialAgentConfig, DataStore } from "tabby-agent";
function getWorkspaceConfiguration(): Partial<AgentConfig> {
function getWorkspaceConfiguration(): PartialAgentConfig {
const configuration = workspace.getConfiguration("tabby");
const config: Partial<AgentConfig> = {};
const config: PartialAgentConfig = {};
const endpoint = configuration.get<string>("api.endpoint");
if (endpoint && endpoint.trim().length > 0) {
config.server = {
@ -38,14 +38,24 @@ export async function createAgentInstance(context: ExtensionContext): Promise<Ta
},
};
const agent = await TabbyAgent.create({ dataStore: env.appHost === "desktop" ? undefined : extensionDataStore });
agent.initialize({
const initPromise = agent.initialize({
config: getWorkspaceConfiguration(),
client: `${env.appName} ${env.appHost} ${version}, ${context.extension.id} ${context.extension.packageJSON.version}`,
});
workspace.onDidChangeConfiguration((event) => {
if (event.affectsConfiguration("tabby")) {
const config = getWorkspaceConfiguration();
agent.updateConfig(config);
workspace.onDidChangeConfiguration(async (event) => {
await initPromise;
const configuration = workspace.getConfiguration("tabby");
if (event.affectsConfiguration("tabby.api.endpoint")) {
const endpoint = configuration.get<string>("api.endpoint");
if (endpoint && endpoint.trim().length > 0) {
agent.updateConfig("server.endpoint", endpoint);
} else {
agent.clearConfig("server.endpoint");
}
}
if (event.affectsConfiguration("tabby.usage.anonymousUsageTracking")) {
const anonymousUsageTrackingDisabled = configuration.get<boolean>("usage.anonymousUsageTracking", false);
agent.updateConfig("anonymousUsageTracking.disable", anonymousUsageTrackingDisabled);
}
});
instance = agent;

View File

@ -1244,10 +1244,10 @@ deep-is@^0.1.3:
resolved "https://registry.yarnpkg.com/deep-is/-/deep-is-0.1.4.tgz#a6f2dce612fadd2ef1f519b73551f17e85199831"
integrity sha512-oIPzksmTg4/MriiaYGO+okXDT7ztn/w3Eptv/+gSIdMdKsJo0u4CfYNFJPy+4SKMuCqGw2wxnA+URMg3t8a/bQ==
deepmerge@^4.3.1:
version "4.3.1"
resolved "https://registry.yarnpkg.com/deepmerge/-/deepmerge-4.3.1.tgz#44b5f2147cd3b00d4b56137685966f26fd25dd4a"
integrity sha512-3sUqbMEc77XqpdNO7FRyRog+eW3ph+GYCbj+rK+uYyRMuwsVy0rMiVtPn+QJlKFvWP/1PYpapqYn0Me2knFn+A==
deepmerge-ts@^5.1.0:
version "5.1.0"
resolved "https://registry.yarnpkg.com/deepmerge-ts/-/deepmerge-ts-5.1.0.tgz#c55206cc4c7be2ded89b9c816cf3608884525d7a"
integrity sha512-eS8dRJOckyo9maw9Tu5O5RUi/4inFLrnoLkBe3cPfDMx3WZioXtmOew4TXQaxq7Rhl4xjDtR7c6x8nNTxOvbFw==
define-properties@^1.1.3, define-properties@^1.1.4, define-properties@^1.2.0:
version "1.2.0"
@ -1341,6 +1341,13 @@ domutils@^3.0.1:
domelementtype "^2.3.0"
domhandler "^5.0.3"
dot-prop@^8.0.2:
version "8.0.2"
resolved "https://registry.yarnpkg.com/dot-prop/-/dot-prop-8.0.2.tgz#afda6866610684dd155a96538f8efcdf78a27f18"
integrity sha512-xaBe6ZT4DHPkg0k4Ytbvn5xoxgpG0jOS1dYxSOwAHPuNLjP3/OzN0gH55SrLqpx8cBfSaVt91lXYkApjb+nYdQ==
dependencies:
type-fest "^3.8.0"
duplexify@^3.5.0, duplexify@^3.6.0:
version "3.7.1"
resolved "https://registry.yarnpkg.com/duplexify/-/duplexify-3.7.1.tgz#2a4df5317f6ccfd91f86d6fd25d8d8a103b88309"
@ -4041,6 +4048,11 @@ type-fest@^1.0.1, type-fest@^1.2.1, type-fest@^1.2.2:
resolved "https://registry.yarnpkg.com/type-fest/-/type-fest-1.4.0.tgz#e9fb813fe3bf1744ec359d55d1affefa76f14be1"
integrity sha512-yGSza74xk0UG8k+pLh5oeoYirvIiWo5t0/o3zHHAO2tRDiZcxWP7fywNlXhqb6/r6sWvwi+RsyQMWhVLe4BVuA==
type-fest@^3.8.0:
version "3.13.1"
resolved "https://registry.yarnpkg.com/type-fest/-/type-fest-3.13.1.tgz#bb744c1f0678bea7543a2d1ec24e83e68e8c8706"
integrity sha512-tLq3bSNx+xSpwvAJnzrK0Ep5CLNWjvFTOp71URMaAEWBfRb9nnJiBoUe0tF8bI4ZFO3omgBR6NvnbzVUT3Ly4g==
type-is@^1.6.16:
version "1.6.18"
resolved "https://registry.yarnpkg.com/type-is/-/type-is-1.6.18.tgz#4e552cd05df09467dcbc4ef739de89f2cf37c131"