index.ts

import { Attributes } from "@opentelemetry/api";
import { ServerBackgroundWorker, TaskRunContext } from "../schemas/index.js";
import { SemanticInternalAttributes } from "../semanticInternalAttributes.js";
import { sdkScope } from "../sdkScope/index.js";
import { getGlobal, registerGlobal } from "../utils/globals.js";
import { TaskContext } from "./types.js";

const API_NAME = "task-context";

export class TaskContextAPI {
  private static _instance?: TaskContextAPI;
  private _runDisabled = false;
  private _conversationId?: string;

  private constructor() {}

  public static getInstance(): TaskContextAPI {
    if (!this._instance) {
      this._instance = new TaskContextAPI();
    }

    return this._instance;
  }

  get isInsideTask(): boolean {
    if (this.#isolatedFromContext()) return false;
    return this.#getTaskContext() !== undefined;
  }

  get isRunDisabled(): boolean {
    return this._runDisabled;
  }

  get ctx(): TaskRunContext | undefined {
    if (this.#isolatedFromContext()) return undefined;
    return this.#getTaskContext()?.ctx;
  }

  get worker(): ServerBackgroundWorker | undefined {
    if (this.#isolatedFromContext()) return undefined;
    return this.#getTaskContext()?.worker;
  }

  get isWarmStart(): boolean | undefined {
    if (this.#isolatedFromContext()) return undefined;
    return this.#getTaskContext()?.isWarmStart;
  }

  #isolatedFromContext(): boolean {
    const scope = sdkScope.getStore();
    return !!scope && !scope.inheritContext;
  }

  get attributes(): Attributes {
    if (this.ctx) {
      return {
        ...this.contextAttributes,
        ...this.workerAttributes,
        ...this.conversationAttributes,
        [SemanticInternalAttributes.WARM_START]: !!this.isWarmStart,
      };
    }

    return {};
  }

  get conversationAttributes(): Attributes {
    if (!this._conversationId) return {};
    return { [SemanticInternalAttributes.GEN_AI_CONVERSATION_ID]: this._conversationId };
  }

  get conversationId(): string | undefined {
    return this._conversationId;
  }

  public setConversationId(conversationId: string | undefined): void {
    this._conversationId = conversationId || undefined;
  }

  get resourceAttributes(): Attributes {
    if (this.ctx) {
      return {
        [SemanticInternalAttributes.ENVIRONMENT_ID]: this.ctx.environment.id,
        [SemanticInternalAttributes.ENVIRONMENT_TYPE]: this.ctx.environment.type,
        [SemanticInternalAttributes.ORGANIZATION_ID]: this.ctx.organization.id,
        [SemanticInternalAttributes.PROJECT_ID]: this.ctx.project.id,
        [SemanticInternalAttributes.PROJECT_REF]: this.ctx.project.ref,
        [SemanticInternalAttributes.PROJECT_NAME]: this.ctx.project.name,
        [SemanticInternalAttributes.ORGANIZATION_SLUG]: this.ctx.organization.slug,
        [SemanticInternalAttributes.ORGANIZATION_NAME]: this.ctx.organization.name,
        [SemanticInternalAttributes.MACHINE_PRESET_NAME]: this.ctx.machine?.name,
        [SemanticInternalAttributes.MACHINE_PRESET_CPU]: this.ctx.machine?.cpu,
        [SemanticInternalAttributes.MACHINE_PRESET_MEMORY]: this.ctx.machine?.memory,
        [SemanticInternalAttributes.MACHINE_PRESET_CENTS_PER_MS]: this.ctx.machine?.centsPerMs,
      };
    }

    return {};
  }

  get workerAttributes(): Attributes {
    if (this.worker) {
      return {
        [SemanticInternalAttributes.WORKER_ID]: this.worker.id,
        [SemanticInternalAttributes.WORKER_VERSION]: this.worker.version,
      };
    }

    return {};
  }

  get contextAttributes(): Attributes {
    if (this.ctx) {
      return {
        [SemanticInternalAttributes.ATTEMPT_NUMBER]: this.ctx.attempt.number,
        [SemanticInternalAttributes.TASK_SLUG]: this.ctx.task.id,
        [SemanticInternalAttributes.TASK_PATH]: this.ctx.task.filePath,
        [SemanticInternalAttributes.QUEUE_NAME]: this.ctx.queue.name,
        [SemanticInternalAttributes.QUEUE_ID]: this.ctx.queue.id,
        [SemanticInternalAttributes.RUN_ID]: this.ctx.run.id,
        [SemanticInternalAttributes.RUN_IS_TEST]: this.ctx.run.isTest,
        [SemanticInternalAttributes.RUN_IS_REPLAY]: this.ctx.run.isReplay,
        [SemanticInternalAttributes.BATCH_ID]: this.ctx.batch?.id,
        [SemanticInternalAttributes.IDEMPOTENCY_KEY]: this.ctx.run.idempotencyKey,
      };
    }

    return {};
  }

  public disable() {
    this._runDisabled = true;
  }

  public setGlobalTaskContext(taskContext: TaskContext): boolean {
    this._runDisabled = false;
    // Each run boot re-registers the global; clear any conversation id
    // left over from a previous run on this warm-restarted process so
    // attributes don't bleed across runs that don't call
    // `setConversationId` themselves.
    this._conversationId = undefined;
    return registerGlobal(API_NAME, taskContext, true);
  }

  #getTaskContext(): TaskContext | undefined {
    return getGlobal(API_NAME);
  }
}