|
| 1 | +/*--------------------------------------------------------------------------------------------- |
| 2 | + * Copyright (c) Microsoft Corporation. All rights reserved. |
| 3 | + * Licensed under the MIT License. See License.txt in the project root for license information. |
| 4 | + *--------------------------------------------------------------------------------------------*/ |
| 5 | + |
| 6 | +import { sumBy } from '../../../../../base/common/arrays.js'; |
| 7 | +import { TaskQueue, timeout } from '../../../../../base/common/async.js'; |
| 8 | +import { Lazy } from '../../../../../base/common/lazy.js'; |
| 9 | +import { Disposable, DisposableStore, toDisposable } from '../../../../../base/common/lifecycle.js'; |
| 10 | +import { autorun, mapObservableArrayCached, observableValue, runOnChange } from '../../../../../base/common/observable.js'; |
| 11 | +import { AnnotatedStringEdit } from '../../../../../editor/common/core/edits/stringEdit.js'; |
| 12 | +import { isAiEdit, isUserEdit } from '../../../../../editor/common/textModelEditSource.js'; |
| 13 | +import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js'; |
| 14 | +import { IStorageService, StorageScope, StorageTarget } from '../../../../../platform/storage/common/storage.js'; |
| 15 | +import { AnnotatedDocuments } from '../helpers/annotatedDocuments.js'; |
| 16 | +import { AiStatsStatusBar } from './aiStatsStatusBar.js'; |
| 17 | + |
| 18 | +export class AiStatsFeature extends Disposable { |
| 19 | + private readonly _data: IValue<IData>; |
| 20 | + private readonly _dataVersion = observableValue(this, 0); |
| 21 | + |
| 22 | + constructor( |
| 23 | + annotatedDocuments: AnnotatedDocuments, |
| 24 | + @IStorageService private readonly _storageService: IStorageService, |
| 25 | + @IInstantiationService private readonly _instantiationService: IInstantiationService, |
| 26 | + ) { |
| 27 | + super(); |
| 28 | + |
| 29 | + const storedValue = getStoredValue<IData>(this._storageService, 'aiStats', StorageScope.WORKSPACE, StorageTarget.USER); |
| 30 | + this._data = rateLimitWrite<IData>(storedValue, 1 / 60, this._store); |
| 31 | + |
| 32 | + this.aiRate.recomputeInitiallyAndOnChange(this._store); |
| 33 | + |
| 34 | + this._register(autorun(reader => { |
| 35 | + reader.store.add(this._instantiationService.createInstance(AiStatsStatusBar.hot.read(reader), this.aiRate)); |
| 36 | + })); |
| 37 | + |
| 38 | + const obs = mapObservableArrayCached(this, annotatedDocuments.documents, (doc, store) => { |
| 39 | + store.add(runOnChange(doc.documentWithAnnotations.value, (_val, _prev, edit) => { |
| 40 | + const e = AnnotatedStringEdit.compose(edit.map(e => e.edit)); |
| 41 | + |
| 42 | + const curSession = new Lazy(() => this._getDataAndSession()); |
| 43 | + |
| 44 | + for (const r of e.replacements) { |
| 45 | + if (isAiEdit(r.data.editSource)) { |
| 46 | + curSession.value.currentSession.aiCharacters += r.newText.length; |
| 47 | + } else if (isUserEdit(r.data.editSource)) { |
| 48 | + curSession.value.currentSession.typedCharacters += r.newText.length; |
| 49 | + } |
| 50 | + } |
| 51 | + |
| 52 | + if (curSession.hasValue) { |
| 53 | + this._data.writeValue(curSession.value.data); |
| 54 | + this._dataVersion.set(this._dataVersion.get() + 1, undefined); |
| 55 | + } |
| 56 | + })); |
| 57 | + }); |
| 58 | + |
| 59 | + obs.recomputeInitiallyAndOnChange(this._store); |
| 60 | + |
| 61 | + } |
| 62 | + |
| 63 | + public readonly aiRate = this._dataVersion.map(() => { |
| 64 | + const val = this._data.getValue(); |
| 65 | + if (!val) { |
| 66 | + return 0; |
| 67 | + } |
| 68 | + |
| 69 | + const r = average(val.sessions, session => { |
| 70 | + const sum = session.typedCharacters + session.aiCharacters; |
| 71 | + if (sum === 0) { |
| 72 | + return 0; |
| 73 | + } |
| 74 | + return session.aiCharacters / sum; |
| 75 | + }); |
| 76 | + |
| 77 | + return r; |
| 78 | + }); |
| 79 | + |
| 80 | + private _getDataAndSession(): { data: IData; currentSession: ISession } { |
| 81 | + const state = this._data.getValue() ?? { sessions: [] }; |
| 82 | + |
| 83 | + const sessionLengthMs = 5 * 60 * 1000; // 5 minutes |
| 84 | + |
| 85 | + let lastSession = state.sessions.at(-1); |
| 86 | + const nowTime = Date.now(); |
| 87 | + if (!lastSession || nowTime - lastSession.startTime > sessionLengthMs) { |
| 88 | + state.sessions.push({ |
| 89 | + startTime: nowTime, |
| 90 | + typedCharacters: 0, |
| 91 | + aiCharacters: 0 |
| 92 | + }); |
| 93 | + lastSession = state.sessions.at(-1)!; |
| 94 | + |
| 95 | + const dayMs = 24 * 60 * 60 * 1000; // 24h |
| 96 | + // Clean up old sessions, keep only the last 24h worth of sessions |
| 97 | + while (state.sessions.length > dayMs / sessionLengthMs) { |
| 98 | + state.sessions.shift(); |
| 99 | + } |
| 100 | + } |
| 101 | + return { data: state, currentSession: lastSession }; |
| 102 | + } |
| 103 | +} |
| 104 | + |
| 105 | +interface IData { |
| 106 | + sessions: ISession[]; |
| 107 | +} |
| 108 | + |
| 109 | +// 5 min window |
| 110 | +interface ISession { |
| 111 | + startTime: number; |
| 112 | + typedCharacters: number; |
| 113 | + aiCharacters: number; |
| 114 | +} |
| 115 | + |
| 116 | + |
| 117 | +function average<T>(arr: T[], selector: (item: T) => number): number { |
| 118 | + if (arr.length === 0) { |
| 119 | + return 0; |
| 120 | + } |
| 121 | + const s = sumBy(arr, selector); |
| 122 | + return s / arr.length; |
| 123 | +} |
| 124 | + |
| 125 | + |
| 126 | +interface IValue<T> { |
| 127 | + writeValue(value: T | undefined): void; |
| 128 | + getValue(): T | undefined; |
| 129 | +} |
| 130 | + |
| 131 | +function rateLimitWrite<T>(targetValue: IValue<T>, maxWritesPerSecond: number, store: DisposableStore): IValue<T> { |
| 132 | + const queue = new TaskQueue(); |
| 133 | + let _value: T | undefined = undefined; |
| 134 | + let valueVersion = 0; |
| 135 | + let savedVersion = 0; |
| 136 | + store.add(toDisposable(() => { |
| 137 | + if (valueVersion !== savedVersion) { |
| 138 | + targetValue.writeValue(_value); |
| 139 | + savedVersion = valueVersion; |
| 140 | + } |
| 141 | + })); |
| 142 | + |
| 143 | + return { |
| 144 | + writeValue(value: T | undefined): void { |
| 145 | + valueVersion++; |
| 146 | + const v = valueVersion; |
| 147 | + _value = value; |
| 148 | + |
| 149 | + queue.clearPending(); |
| 150 | + queue.schedule(async () => { |
| 151 | + targetValue.writeValue(value); |
| 152 | + savedVersion = v; |
| 153 | + await timeout(5000); |
| 154 | + }); |
| 155 | + }, |
| 156 | + getValue(): T | undefined { |
| 157 | + if (valueVersion > 0) { |
| 158 | + return _value; |
| 159 | + } |
| 160 | + return targetValue.getValue(); |
| 161 | + } |
| 162 | + }; |
| 163 | +} |
| 164 | + |
| 165 | +function getStoredValue<T>(service: IStorageService, key: string, scope: StorageScope, target: StorageTarget): IValue<T> { |
| 166 | + let lastValue: T | undefined = undefined; |
| 167 | + let hasLastValue = false; |
| 168 | + return { |
| 169 | + writeValue(value: T | undefined): void { |
| 170 | + if (value === undefined) { |
| 171 | + service.remove(key, scope); |
| 172 | + } else { |
| 173 | + service.store(key, JSON.stringify(value), scope, target); |
| 174 | + } |
| 175 | + lastValue = value; |
| 176 | + }, |
| 177 | + getValue(): T | undefined { |
| 178 | + if (hasLastValue) { |
| 179 | + return lastValue; |
| 180 | + } |
| 181 | + const strVal = service.get(key, scope); |
| 182 | + lastValue = strVal === undefined ? undefined : JSON.parse(strVal) as T | undefined; |
| 183 | + hasLastValue = true; |
| 184 | + return lastValue; |
| 185 | + } |
| 186 | + }; |
| 187 | +} |
0 commit comments