1import { Endpoint, ErrorCode, Headers, Model } from './constants.js';
2import { GemMixin } from './components/gem-mixin.js';
3import {
4  APIError,
5  AuthError,
6  GeminiError,
7  ImageGenerationError,
8  ModelInvalid,
9  TemporarilyBlocked,
10  TimeoutError,
11  UsageLimitExceeded,
12} from './exceptions.js';
13import { Candidate, Gem, GeneratedImage, ModelOutput, RPCData, WebImage } from './types/index.js';
14import {
15  extract_json_from_response,
16  get_access_token,
17  get_nested_value,
18  logger,
19  parse_file_name,
20  rotate_1psidts,
21  rotate_tasks,
22  fetch_with_timeout,
23  sleep,
24  upload_file,
25  write_cookie_file,
26  resolveGeminiWebCookiePath,
27} from './utils/index.js';
28 
29type InitOptions = {
30  timeout?: number;
31  auto_close?: boolean;
32  close_delay?: number;
33  auto_refresh?: boolean;
34  refresh_interval?: number;
35  verbose?: boolean;
36};
37 
38type RequestKwargs = RequestInit & { timeout_ms?: number };
39 
40const GENERATED_IMAGE_URL_PREFIX = 'https://lh3.googleusercontent.com/gg-dl/';
41 
42function normalize_headers(h?: HeadersInit): Record<string, string> {
43  if (!h) return {};
44  if (Array.isArray(h)) return Object.fromEntries(h.map(([k, v]) => [k, v]));
45  if (h instanceof Headers) {
46    const out: Record<string, string> = {};
47    h.forEach((v, k) => {
48      out[k] = v;
49    });
50    return out;
51  }
52  return { ...(h as Record<string, string>) };
53}
54 
55function collect_strings(root: unknown, accept: (s: string) => boolean, limit: number = 20): string[] {
56  const out: string[] = [];
57  const seen = new Set<string>();
58  const stack: unknown[] = [root];
59 
60  while (stack.length > 0 && out.length < limit) {
61    const v = stack.pop();
62    if (typeof v === 'string') {
63      if (accept(v) && !seen.has(v)) {
64        seen.add(v);
65        out.push(v);
66      }
67      continue;
68    }
69 
70    if (Array.isArray(v)) {
71      for (let i = 0; i < v.length; i++) stack.push(v[i]);
72      continue;
73    }
74 
75    if (v && typeof v === 'object') {
76      for (const val of Object.values(v as Record<string, unknown>)) stack.push(val);
77    }
78  }
79 
80  return out;
81}
82 
83function collect_generated_image_urls(root: unknown, limit: number = 4): string[] {
84  return collect_strings(root, (s) => s.startsWith(GENERATED_IMAGE_URL_PREFIX), limit);
85}
86 
87function parse_response_part_body(part: unknown): unknown[] | null {
88  if (!Array.isArray(part)) return null;
89  const part_body = get_nested_value<string | null>(part, [2], null);
90  if (!part_body) return null;
91  try {
92    const part_json = JSON.parse(part_body) as unknown;
93    return Array.isArray(part_json) ? part_json : null;
94  } catch {
95    return null;
96  }
97}
98 
99function find_generated_image_part(
100  response_json: unknown[],
101  body_index: number,
102  candidate_index: number,
103  limit: number = 4,
104): { body: unknown[]; urls: string[] } | null {
105  for (let part_index = body_index; part_index < response_json.length; part_index++) {
106    const part_json = parse_response_part_body(response_json[part_index]);
107    if (!part_json) continue;
108    const cand = get_nested_value<unknown>(part_json, [4, candidate_index], null);
109    if (!cand) continue;
110    const urls = collect_generated_image_urls(cand, limit);
111    if (urls.length > 0) return { body: part_json, urls };
112  }
113  return null;
114}
115 
116export function collect_generated_image_urls_from_response_parts(
117  response_json: unknown[],
118  body_index: number,
119  candidate_index: number,
120  limit: number = 4,
121): string[] {
122  return find_generated_image_part(response_json, body_index, candidate_index, limit)?.urls ?? [];
123}
124 
125function push_generated_images(
126  generated_images: GeneratedImage[],
127  urls: string[],
128  proxy: string | null,
129  cookies: Record<string, string>,
130): void {
131  for (const url of urls) {
132    generated_images.push(new GeneratedImage(url, '[Generated Image]', '', proxy, cookies));
133  }
134}
135 
136export class GeminiClient extends GemMixin {
137  public cookies: Record<string, string> = {};
138  public proxy: string | null = null;
139  public _running: boolean = false;
140  public access_token: string | null = null;
141  public timeout: number = 300;
142  public auto_close: boolean = false;
143  public close_delay: number = 300;
144  public auto_refresh: boolean = true;
145  public refresh_interval: number = 540;
146  public kwargs: RequestInit;
147 
148  private close_timer: ReturnType<typeof setTimeout> | null = null;
149  private refresh_abort: AbortController | null = null;
150 
151  constructor(
152    secure_1psid: string | null = null,
153    secure_1psidts: string | null = null,
154    proxy: string | null = null,
155    kwargs: RequestInit = {},
156  ) {
157    super();
158    this.proxy = proxy;
159    this.kwargs = kwargs;
160 
161    if (secure_1psid) {
162      this.cookies['__Secure-1PSID'] = secure_1psid;
163      if (secure_1psidts) this.cookies['__Secure-1PSIDTS'] = secure_1psidts;
164    }
165  }
166 
167  async init(
168    timeoutOrOpts: number | InitOptions = 300,
169    auto_close: boolean = false,
170    close_delay: number = 300,
171    auto_refresh: boolean = true,
172    refresh_interval: number = 540,
173    verbose: boolean = true,
174  ): Promise<void> {
175    const opts: InitOptions =
176      typeof timeoutOrOpts === 'object'
177        ? timeoutOrOpts
178        : { timeout: timeoutOrOpts, auto_close, close_delay, auto_refresh, refresh_interval, verbose };
179 
180    const timeout = opts.timeout ?? 300;
181    const ac = opts.auto_close ?? false;
182    const cd = opts.close_delay ?? 300;
183    const ar = opts.auto_refresh ?? true;
184    const ri = opts.refresh_interval ?? 540;
185    const vb = opts.verbose ?? true;
186 
187    try {
188      const [token, valid] = await get_access_token(this.cookies, this.proxy, vb);
189      this.access_token = token;
190      this.cookies = valid;
191      this._running = true;
192 
193      this.timeout = timeout;
194      this.auto_close = ac;
195      this.close_delay = cd;
196      if (this.auto_close) await this.reset_close_task();
197 
198      this.auto_refresh = ar;
199      this.refresh_interval = ri;
200 
201      const sid = this.cookies['__Secure-1PSID'];
202      if (sid) {
203        const existing = rotate_tasks.get(sid);
204        if (existing && existing instanceof AbortController) existing.abort();
205        rotate_tasks.delete(sid);
206      }
207 
208      if (this.auto_refresh && sid) {
209        const ctl = new AbortController();
210        this.refresh_abort?.abort();
211        this.refresh_abort = ctl;
212        rotate_tasks.set(sid, ctl);
213        void this.start_auto_refresh(ctl.signal);
214      }
215 
216      await write_cookie_file(this.cookies, resolveGeminiWebCookiePath(), 'client').catch(() => {});
217 
218      if (vb) logger.success('Gemini client initialized successfully.');
219    } catch (e) {
220      await this.close();
221      throw e;
222    }
223  }
224 
225  async close(delay: number = 0): Promise<void> {
226    if (delay > 0) await sleep(delay * 1000);
227    this._running = false;
228 
229    if (this.close_timer) {
230      clearTimeout(this.close_timer);
231      this.close_timer = null;
232    }
233 
234    this.refresh_abort?.abort();
235    this.refresh_abort = null;
236 
237    const sid = this.cookies['__Secure-1PSID'];
238    const t = sid ? rotate_tasks.get(sid) : null;
239    if (t && t instanceof AbortController) t.abort();
240    if (sid) rotate_tasks.delete(sid);
241  }
242 
243  async reset_close_task(): Promise<void> {
244    if (this.close_timer) {
245      clearTimeout(this.close_timer);
246      this.close_timer = null;
247    }
248 
249    this.close_timer = setTimeout(() => {
250      void this.close(0);
251    }, this.close_delay * 1000);
252    this.close_timer.unref?.();
253  }
254 
255  async start_auto_refresh(signal: AbortSignal): Promise<void> {
256    while (!signal.aborted) {
257      let newTs: string | null = null;
258      try {
259        newTs = await rotate_1psidts(this.cookies, this.proxy);
260      } catch (e) {
261        if (e instanceof AuthError) {
262          logger.warning('AuthError: Failed to refresh cookies. Auto refresh task canceled.');
263          return;
264        }
265        logger.warning(`Unexpected error while refreshing cookies: ${e instanceof Error ? e.message : String(e)}`);
266      }
267 
268      if (newTs) {
269        this.cookies['__Secure-1PSIDTS'] = newTs;
270        await write_cookie_file(this.cookies, resolveGeminiWebCookiePath(), 'refresh').catch(() => {});
271        logger.debug('Cookies refreshed. New __Secure-1PSIDTS applied.');
272      }
273 
274      await sleep(this.refresh_interval * 1000, signal);
275    }
276  }
277 
278  protected async _run<T>(fn: () => Promise<T>, retry: number): Promise<T> {
279    try {
280      if (!this._running) {
281        await this.init({
282          timeout: this.timeout,
283          auto_close: this.auto_close,
284          close_delay: this.close_delay,
285          auto_refresh: this.auto_refresh,
286          refresh_interval: this.refresh_interval,
287          verbose: false,
288        });
289 
290        if (!this._running) {
291          throw new APIError('Client initialization failed.');
292        }
293      }
294 
295      return await fn();
296    } catch (e) {
297      let r = retry;
298      if (e instanceof ImageGenerationError) r = Math.min(1, r);
299      if (e instanceof APIError && r > 0) {
300        await sleep(1000);
301        return await this._run(fn, r - 1);
302      }
303      throw e;
304    }
305  }
306 
307  async generate_content(
308    prompt: string,
309    files: string[] | null = null,
310    model: Model | string | Record<string, unknown> = Model.UNSPECIFIED,
311    gem: Gem | string | null = null,
312    chat: ChatSession | null = null,
313    kwargs: RequestKwargs = {},
314  ): Promise<ModelOutput> {
315    return await this._run(async () => {
316      if (!prompt) throw new Error('Prompt cannot be empty.');
317 
318      let mdl: Model;
319      if (typeof model === 'string') mdl = Model.from_name(model);
320      else if (model instanceof Model) mdl = model;
321      else if (model && typeof model === 'object') mdl = Model.from_dict(model);
322      else throw new TypeError(`'model' must be a Model instance, string, or dictionary; got ${typeof model}`);
323 
324      const gem_id = gem instanceof Gem ? gem.id : gem;
325 
326      if (this.auto_close) await this.reset_close_task();
327 
328      if (!this.access_token) throw new APIError('Missing access token.');
329 
330      const f = files?.length ? files : null;
331      const uploaded =
332        f &&
333        (await Promise.all(
334          f.map(async (p) => [[await upload_file(p, this.proxy)], parse_file_name(p)] as [string[], string]),
335        ));
336 
337      const first = uploaded ? [prompt, 0, null, uploaded] : [prompt];
338      const inner: unknown[] = [first, null, chat ? chat.metadata : null];
339 
340      if (gem_id) {
341        for (let i = 0; i < 16; i++) inner.push(null);
342        inner.push(gem_id);
343      }
344 
345      const f_req = JSON.stringify([null, JSON.stringify(inner)]);
346      const body = new URLSearchParams({ at: this.access_token, 'f.req': f_req }).toString();
347 
348      const h0 = { ...Headers.GEMINI, ...mdl.model_header, Cookie: Object.entries(this.cookies).map(([k, v]) => `${k}=${v}`).join('; ') };
349      const h1 = { ...h0, ...normalize_headers(kwargs.headers) };
350 
351      let res: Response;
352      try {
353        const timeout_ms = typeof kwargs.timeout_ms === 'number' ? kwargs.timeout_ms : this.timeout * 1000;
354        const { timeout_ms: _t, ...rest } = kwargs;
355        res = await fetch_with_timeout(Endpoint.GENERATE, {
356          method: 'POST',
357          headers: h1,
358          body,
359          redirect: 'follow',
360          ...this.kwargs,
361          ...rest,
362          timeout_ms,
363        });
364      } catch (e) {
365        throw new TimeoutError(
366          `Generate content request timed out, please try again. If the problem persists, consider setting a higher 'timeout' value when initializing GeminiClient. (${e instanceof Error ? e.message : String(e)})`,
367        );
368      }
369 
370      if (res.status !== 200) {
371        await this.close();
372        throw new APIError(`Failed to generate contents. Request failed with status code ${res.status}`);
373      }
374 
375      const txt = await res.text();
376      const response_json = extract_json_from_response(txt);
377 
378      let body_json: unknown[] | null = null;
379      let body_index = 0;
380 
381      try {
382        if (!Array.isArray(response_json)) throw new Error('Invalid JSON');
383        for (let part_index = 0; part_index < response_json.length; part_index++) {
384          const part = response_json[part_index];
385          if (!Array.isArray(part)) continue;
386          const part_body = get_nested_value<string | null>(part, [2], null);
387          if (!part_body) continue;
388          try {
389            const part_json = JSON.parse(part_body) as unknown[];
390            if (get_nested_value(part_json, [4], null)) {
391              body_index = part_index;
392              body_json = part_json;
393              break;
394            }
395          } catch {}
396        }
397        if (!body_json) throw new Error('No body');
398      } catch {
399        await this.close();
400        try {
401          const code = get_nested_value<number>(response_json, [0, 5, 2, 0, 1, 0], -1);
402          if (code === ErrorCode.USAGE_LIMIT_EXCEEDED) {
403            throw new UsageLimitExceeded(
404              `Failed to generate contents. Usage limit of ${mdl.model_name} model has exceeded. Please try switching to another model.`,
405            );
406          }
407          if (code === ErrorCode.MODEL_INCONSISTENT) {
408            throw new ModelInvalid(
409              'Failed to generate contents. The specified model is inconsistent with the chat history. Please make sure to pass the same `model` parameter when starting a chat session with previous metadata.',
410            );
411          }
412          if (code === ErrorCode.MODEL_HEADER_INVALID) {
413            throw new ModelInvalid(
414              'Failed to generate contents. The specified model is not available. Please update gemini_webapi to the latest version. If the error persists and is caused by the package, please report it on GitHub.',
415            );
416          }
417          if (code === ErrorCode.IP_TEMPORARILY_BLOCKED) {
418            throw new TemporarilyBlocked(
419              'Failed to generate contents. Your IP address is temporarily blocked by Google. Please try using a proxy or waiting for a while.',
420            );
421          }
422        } catch (e) {
423          if (e instanceof GeminiError) throw e;
424        }
425 
426        logger.debug(`Invalid response: ${txt.slice(0, 500)}`);
427        throw new APIError('Failed to generate contents. Invalid response data received. Client will try to re-initialize on next request.');
428      }
429 
430      try {
431        const candidate_list = get_nested_value<unknown[]>(body_json, [4], []);
432        const out: Candidate[] = [];
433 
434        for (let candidate_index = 0; candidate_index < candidate_list.length; candidate_index++) {
435          const candidate = candidate_list[candidate_index];
436          if (!Array.isArray(candidate)) continue;
437 
438          const rcid = get_nested_value<string | null>(candidate, [0], null);
439          if (!rcid) continue;
440 
441          let text = String(get_nested_value(candidate, [1, 0], ''));
442          if (/^http:\/\/googleusercontent\.com\/card_content\/\d+/.test(text)) {
443            text = String(get_nested_value(candidate, [22, 0], text));
444          }
445 
446          const thoughts = get_nested_value<string | null>(candidate, [37, 0, 0], null);
447 
448          const web_images: WebImage[] = [];
449          for (const w of get_nested_value<unknown[]>(candidate, [12, 1], [])) {
450            if (!Array.isArray(w)) continue;
451            const url = get_nested_value<string | null>(w, [0, 0, 0], null);
452            if (!url) continue;
453            web_images.push(new WebImage(url, String(get_nested_value(w, [7, 0], '')), String(get_nested_value(w, [0, 4], '')), this.proxy));
454          }
455 
456          const generated_images: GeneratedImage[] = [];
457          const wants_generated =
458            get_nested_value(candidate, [12, 7, 0], null) != null ||
459            /http:\/\/googleusercontent\.com\/image_generation_content\/\d+/.test(text);
460 
461          if (wants_generated) {
462            const image_part = find_generated_image_part(response_json as unknown[], body_index, candidate_index, 1);
463            const img_body = image_part?.body ?? null;
464 
465            if (!img_body) {
466              throw new ImageGenerationError(
467                'Failed to parse generated images. Please update gemini_webapi to the latest version. If the error persists and is caused by the package, please report it on GitHub.',
468              );
469            }
470 
471            const img_candidate = get_nested_value<unknown[]>(img_body, [4, candidate_index], []);
472            const finished = get_nested_value<string | null>(img_candidate, [1, 0], null);
473            if (finished) {
474              text = finished.replace(/http:\/\/googleusercontent\.com\/image_generation_content\/\d+/g, '').trimEnd();
475            }
476 
477            const gen = get_nested_value<unknown[]>(img_candidate, [12, 7, 0], []);
478            for (let img_index = 0; img_index < gen.length; img_index++) {
479              const g = gen[img_index];
480              if (!Array.isArray(g)) continue;
481              const url = get_nested_value<string | null>(g, [0, 3, 3], null);
482              if (!url) continue;
483              const img_num = get_nested_value<number | null>(g, [3, 6], null);
484              const title = img_num ? `[Generated Image ${img_num}]` : '[Generated Image]';
485              const alt_list = get_nested_value<unknown[]>(g, [3, 5], []);
486              const alt =
487                (typeof alt_list[img_index] === 'string' ? (alt_list[img_index] as string) : null) ??
488                (typeof alt_list[0] === 'string' ? (alt_list[0] as string) : '') ??
489                '';
490              generated_images.push(new GeneratedImage(url, title, alt, this.proxy, this.cookies));
491            }
492 
493            if (generated_images.length === 0) {
494              push_generated_images(generated_images, collect_generated_image_urls(img_candidate), this.proxy, this.cookies);
495            }
496          }
497 
498          // Fallback: unconditionally scan all response parts for generated image URLs.
499          // The `wants_generated` detection above relies on `candidate[12][7][0]` and an old
500          // `googleusercontent.com/image_generation_content/` URL pattern, both of which no
501          // longer appear in the current Gemini Web API response format.
502          // When Gemini does generate images, their URLs now start with
503          // `https://lh3.googleusercontent.com/gg-dl/` and are present somewhere in the
504          // response parts — this fallback finds them so images are not silently dropped.
505          if (generated_images.length === 0) {
506            const urls = collect_generated_image_urls_from_response_parts(response_json as unknown[], body_index, candidate_index);
507            push_generated_images(generated_images, urls, this.proxy, this.cookies);
508          }
509 
510          out.push(new Candidate({ rcid, text, thoughts, web_images, generated_images }));
511        }
512 
513        if (out.length === 0) {
514          throw new GeminiError('Failed to generate contents. No output data found in response.');
515        }
516 
517        const metadata = get_nested_value<string[]>(body_json, [1], []);
518        const output = new ModelOutput({ metadata, candidates: out });
519 
520        if (chat instanceof ChatSession) chat.last_output = output;
521        return output;
522      } catch (e) {
523        if (e instanceof GeminiError || e instanceof APIError) throw e;
524        throw new APIError('Failed to parse response body. Data structure is invalid.');
525      }
526    }, 2);
527  }
528 
529  async generateContent(
530    prompt: string,
531    files?: string[] | null,
532    model?: Model | string | Record<string, unknown>,
533    gem?: Gem | string | null,
534    chat?: ChatSession | null,
535    kwargs?: RequestKwargs,
536  ): Promise<ModelOutput> {
537    return await this.generate_content(prompt, files ?? null, model ?? Model.UNSPECIFIED, gem ?? null, chat ?? null, kwargs ?? {});
538  }
539 
540  start_chat(opts?: ConstructorParameters<typeof ChatSession>[1]): ChatSession {
541    return new ChatSession(this, opts);
542  }
543 
544  startChat(opts?: ConstructorParameters<typeof ChatSession>[1]): ChatSession {
545    return this.start_chat(opts);
546  }
547 
548  protected async _batch_execute(payloads: RPCData[], opts: RequestInit = {}): Promise<Response> {
549    if (!this.access_token) throw new APIError('Missing access token.');
550 
551    const f_req = JSON.stringify([payloads.map((p) => p.serialize())]);
552    const body = new URLSearchParams({ at: this.access_token, 'f.req': f_req }).toString();
553 
554    const h0 = { ...Headers.GEMINI, Cookie: Object.entries(this.cookies).map(([k, v]) => `${k}=${v}`).join('; ') };
555    const h1 = { ...h0, ...normalize_headers(opts.headers) };
556 
557    const res = await fetch_with_timeout(Endpoint.BATCH_EXEC, {
558      method: 'POST',
559      headers: h1,
560      body,
561      redirect: 'follow',
562      ...this.kwargs,
563      ...opts,
564      timeout_ms: this.timeout * 1000,
565    });
566 
567    if (res.status !== 200) {
568      await this.close();
569      throw new APIError(`Batch execution failed with status code ${res.status}`);
570    }
571 
572    return res;
573  }
574}
575 
576export class ChatSession {
577  private __metadata: Array<string | null> = [null, null, null];
578  public geminiclient: GeminiClient;
579  private _last_output: ModelOutput | null = null;
580  public model: Model | string | Record<string, unknown>;
581  public gem: Gem | string | null;
582 
583  constructor(
584    geminiclient: GeminiClient,
585    opts: {
586      metadata?: Array<string | null>;
587      cid?: string | null;
588      rid?: string | null;
589      rcid?: string | null;
590      model?: Model | string | Record<string, unknown>;
591      gem?: Gem | string | null;
592    } = {},
593  ) {
594    this.geminiclient = geminiclient;
595    this.model = opts.model ?? Model.UNSPECIFIED;
596    this.gem = opts.gem ?? null;
597 
598    if (opts.metadata) this.metadata = opts.metadata;
599    if (opts.cid) this.cid = opts.cid;
600    if (opts.rid) this.rid = opts.rid;
601    if (opts.rcid) this.rcid = opts.rcid;
602  }
603 
604  toString(): string {
605    return `ChatSession(cid='${this.cid}', rid='${this.rid}', rcid='${this.rcid}')`;
606  }
607 
608  get last_output(): ModelOutput | null {
609    return this._last_output;
610  }
611 
612  set last_output(v: ModelOutput | null) {
613    this._last_output = v;
614    if (v) {
615      this.metadata = (v.metadata ?? []) as Array<string | null>;
616      this.rcid = v.rcid;
617    }
618  }
619 
620  async send_message(prompt: string, files: string[] | null = null, kwargs: RequestKwargs = {}): Promise<ModelOutput> {
621    return await this.geminiclient.generate_content(prompt, files, this.model, this.gem, this, kwargs);
622  }
623 
624  async sendMessage(prompt: string, files?: string[] | null, kwargs?: RequestKwargs): Promise<ModelOutput> {
625    return await this.send_message(prompt, files ?? null, kwargs ?? {});
626  }
627 
628  choose_candidate(index: number): ModelOutput {
629    if (!this.last_output) throw new Error('No previous output data found in this chat session.');
630    if (index >= this.last_output.candidates.length) {
631      throw new Error(`Index ${index} exceeds the number of candidates in last model output.`);
632    }
633    this.last_output.chosen = index;
634    this.rcid = this.last_output.rcid;
635    return this.last_output;
636  }
637 
638  chooseCandidate(index: number): ModelOutput {
639    return this.choose_candidate(index);
640  }
641 
642  get metadata(): Array<string | null> {
643    return this.__metadata;
644  }
645 
646  set metadata(v: Array<string | null>) {
647    if (v.length > 3) throw new Error('metadata cannot exceed 3 elements');
648    this.__metadata = [null, null, null];
649    for (let i = 0; i < v.length; i++) this.__metadata[i] = v[i] ?? null;
650  }
651 
652  get cid(): string | null {
653    return this.__metadata[0];
654  }
655 
656  set cid(v: string | null) {
657    this.__metadata[0] = v;
658  }
659 
660  get rid(): string | null {
661    return this.__metadata[1];
662  }
663 
664  set rid(v: string | null) {
665    this.__metadata[1] = v;
666  }
667 
668  get rcid(): string | null {
669    return this.__metadata[2];
670  }
671 
672  set rcid(v: string | null) {
673    this.__metadata[2] = v;
674  }
675}
676