Press n or j to go to the next uncovered block, b, p or k for the previous block.
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 | 11x 11x 11x 11x 535x 11x 347x 38x 11x 217x 217x 217x 217x 217x 217x 217x 136x 136x 133x 3x 3x 38x 38x 38x 38x 38x 38x 38x 38x 38x 38x 38x 38x 38x 75x 75x 75x 75x 1x 75x 149x 149x 149x 149x 149x 6x 143x 143x 5x 2x 5x 4x 139x 139x 139x 6x 6x 6x 139x 2x 137x 2x 135x 135x 1x 134x 134x 24x 24x 6x 6x 6x 4x 2x 2x 2x 1x 19x 5x 14x 2x 12x 12x 21x 16x 16x 16x 16x 16x 16x 16x 5x 5x 5x 5x 5x 5x 5x 6x 6x 4x 5x 302x 280x 22x 22x 22x 149x 149x 181x 181x 76x 76x 10x 66x 66x 66x | import { createHash } from "node:crypto";
import { SECRET_STORAGE_KEY } from "../shared/constants";
export interface ApiKeySecretStorage {
get(key: string): PromiseLike<string | undefined>;
}
export type ApiKeyResolutionSource = "configured" | "runtime" | "legacy";
export interface ResolvedApiKey {
value: string;
source: ApiKeyResolutionSource;
}
export interface ToolApiKeyResolutionOptions {
/** Fingerprint of the cache whose model will be used by the tool. */
cacheKeyFingerprint?: string;
/** Manual refresh may intentionally replace a cache whose former owner disappeared. */
allowUnmatchedRuntimeKey?: boolean;
}
const DEFAULT_RUNTIME_GROUP = "<runtime>";
const MODEL_KEY_BINDING_PROPERTY = "__nvidiaNimRuntimeKeyBinding";
interface RuntimeModelBinding {
bindingId: string;
}
function normalizeApiKey(value: unknown): string | undefined {
return typeof value === "string" && value.trim().length > 0 ? value.trim() : undefined;
}
/** Return a non-reversible cache owner marker for an API key. */
export function getApiKeyFingerprint(apiKey: string): string {
return createHash("sha256").update(apiKey).digest("hex");
}
function getRuntimeBindingId(groupId: string, apiKey: string): string {
return createHash("sha256").update(groupId).update("\0").update(apiKey).digest("hex");
}
/**
* Resolves provider-group keys and legacy SecretStorage keys without putting
* the raw key in LanguageModelChatInformation. Each provider group receives
* an opaque binding token so cloned model objects can still resolve the
* correct key when several groups expose the same model ID.
*/
export class NvidiaApiKeyResolver {
private modelBindings = new WeakMap<object, RuntimeModelBinding>();
private readonly runtimeKeysByGroup = new Map<string, string>();
private readonly runtimeKeysByBinding = new Map<string, string>();
private readonly bindingGroups = new Map<string, string>();
private readonly bindingsByGroup = new Map<string, Set<string>>();
private readonly bindingsByModelId = new Map<string, Set<string>>();
constructor(private readonly secrets: ApiKeySecretStorage) {}
normalize(value: unknown): string | undefined {
return normalizeApiKey(value);
}
async resolveConfiguredOrLegacy(configuredApiKey?: unknown): Promise<ResolvedApiKey | undefined> {
const configured = normalizeApiKey(configuredApiKey);
if (configured) {
return { value: configured, source: "configured" };
}
const legacy = normalizeApiKey(await this.secrets.get(SECRET_STORAGE_KEY));
return legacy ? { value: legacy, source: "legacy" } : undefined;
}
registerModelKey(model: object, apiKey: string, groupId = DEFAULT_RUNTIME_GROUP): void {
const normalized = normalizeApiKey(apiKey);
Iif (!normalized) {
return;
}
this.rememberRuntimeKey(normalized, groupId);
const bindingId = getRuntimeBindingId(groupId, normalized);
this.runtimeKeysByBinding.set(bindingId, normalized);
this.bindingGroups.set(bindingId, groupId);
this.getOrCreateSet(this.bindingsByGroup, groupId).add(bindingId);
this.modelBindings.set(model, { bindingId });
// The opaque token is intentionally enumerable because VS Code can clone
// provider model information before invoking the response callback.
try {
Object.defineProperty(model, MODEL_KEY_BINDING_PROPERTY, {
value: bindingId,
configurable: true,
enumerable: true,
});
} catch {
// Frozen model objects still work through the WeakMap and unique-ID
// fallback. Ambiguous cloned IDs fail closed instead of using a wrong key.
}
const id = this.getModelId(model);
if (id) {
this.getOrCreateSet(this.bindingsByModelId, id).add(bindingId);
}
}
rememberRuntimeKey(apiKey: unknown, groupId = DEFAULT_RUNTIME_GROUP): void {
const normalized = normalizeApiKey(apiKey);
Iif (!normalized) {
return;
}
const previousKey = this.runtimeKeysByGroup.get(groupId);
if (previousKey && previousKey !== normalized) {
this.clearRuntimeBindings(groupId);
}
this.runtimeKeysByGroup.set(groupId, normalized);
}
async resolveForModel(model: object): Promise<ResolvedApiKey | undefined> {
const weakBinding = this.modelBindings.get(model)?.bindingId;
const serializedBinding = this.getSerializedBindingId(model);
const hasBindingMarker = Boolean(weakBinding || serializedBinding);
const directRuntime =
this.resolveBindingKey(weakBinding) ?? this.resolveBindingKey(serializedBinding);
if (directRuntime) {
return { value: directRuntime, source: "runtime" };
}
const modelId = this.getModelId(model);
// A cloned model can retain a binding token after its provider group was
// removed or its key changed. Never let that stale object fall through to
// an active binding from a different group that happens to expose the
// same model ID. When no runtime group remains at all, legacy SecretStorage
// fallback is still allowed for backwards compatibility; during a key
// change, however, do not use a legacy key while the replacement binding
// is still being resolved.
if (hasBindingMarker) {
const hasActiveBindingForModel =
modelId !== undefined &&
[...(this.bindingsByModelId.get(modelId) ?? [])].some((bindingId) =>
Boolean(this.resolveBindingKey(bindingId)),
);
if (modelId === undefined || hasActiveBindingForModel || this.runtimeKeysByGroup.size > 0) {
return undefined;
}
}
if (modelId) {
const matchingKeys = new Set<string>();
for (const bindingId of this.bindingsByModelId.get(modelId) ?? []) {
const key = this.resolveBindingKey(bindingId);
if (key) {
matchingKeys.add(key);
}
}
if (matchingKeys.size === 1) {
return { value: matchingKeys.values().next().value as string, source: "runtime" };
}
if (matchingKeys.size > 1) {
// A cloned model without its binding token is ambiguous when several
// provider groups expose the same model ID. Never fall back to a
// legacy key that could belong to a different provider group.
return undefined;
}
}
// Compatibility for model objects supplied by older callers/tests. New
// discovery results never include this property.
const legacyModelKey = normalizeApiKey((model as { apiKey?: unknown }).apiKey);
if (legacyModelKey) {
return { value: legacyModelKey, source: "runtime" };
}
const legacy = normalizeApiKey(await this.secrets.get(SECRET_STORAGE_KEY));
return legacy ? { value: legacy, source: "legacy" } : undefined;
}
async resolveForTool(
options: ToolApiKeyResolutionOptions = {},
): Promise<ResolvedApiKey | undefined> {
const runtimeKeys = [...new Set(this.runtimeKeysByGroup.values())];
if (options.cacheKeyFingerprint) {
const matchingRuntime = runtimeKeys.find(
(apiKey) => getApiKeyFingerprint(apiKey) === options.cacheKeyFingerprint,
);
if (matchingRuntime) {
return { value: matchingRuntime, source: "runtime" };
}
const legacy = normalizeApiKey(await this.secrets.get(SECRET_STORAGE_KEY));
Iif (legacy && getApiKeyFingerprint(legacy) === options.cacheKeyFingerprint) {
return { value: legacy, source: "legacy" };
}
if (!options.allowUnmatchedRuntimeKey) {
return undefined;
}
}
if (runtimeKeys.length === 1) {
return { value: runtimeKeys[0], source: "runtime" };
}
if (runtimeKeys.length > 1) {
// Without a cache-owner fingerprint there is no safe way to decide
// which provider group supplied the cached vision model list. Falling
// back to an unrelated legacy key could send the request with the wrong
// account, so ambiguous runtime ownership must fail closed.
return undefined;
}
const legacy = normalizeApiKey(await this.secrets.get(SECRET_STORAGE_KEY));
return legacy ? { value: legacy, source: "legacy" } : undefined;
}
async resolveLegacy(): Promise<string | undefined> {
return normalizeApiKey(await this.secrets.get(SECRET_STORAGE_KEY));
}
clearRuntimeBindings(groupId?: string): void {
if (groupId === undefined) {
this.modelBindings = new WeakMap<object, RuntimeModelBinding>();
this.runtimeKeysByGroup.clear();
this.runtimeKeysByBinding.clear();
this.bindingGroups.clear();
this.bindingsByGroup.clear();
this.bindingsByModelId.clear();
return;
}
this.runtimeKeysByGroup.delete(groupId);
const bindingIds = this.bindingsByGroup.get(groupId);
Iif (!bindingIds) {
return;
}
for (const bindingId of bindingIds) {
this.runtimeKeysByBinding.delete(bindingId);
this.bindingGroups.delete(bindingId);
for (const [modelId, modelBindings] of this.bindingsByModelId) {
modelBindings.delete(bindingId);
if (modelBindings.size === 0) {
this.bindingsByModelId.delete(modelId);
}
}
}
this.bindingsByGroup.delete(groupId);
}
private resolveBindingKey(bindingId: string | undefined): string | undefined {
if (!bindingId) {
return undefined;
}
const groupId = this.bindingGroups.get(bindingId);
const key = this.runtimeKeysByBinding.get(bindingId);
return groupId && key && this.runtimeKeysByGroup.get(groupId) === key ? key : undefined;
}
private getSerializedBindingId(model: object): string | undefined {
const value = (model as Record<string, unknown>)[MODEL_KEY_BINDING_PROPERTY];
return typeof value === "string" && value.length > 0 ? value : undefined;
}
private getModelId(model: object): string | undefined {
const id = (model as { id?: unknown }).id;
return typeof id === "string" && id.length > 0 ? id : undefined;
}
private getOrCreateSet(map: Map<string, Set<string>>, key: string): Set<string> {
const existing = map.get(key);
if (existing) {
return existing;
}
const created = new Set<string>();
map.set(key, created);
return created;
}
}
|