TF_CAPI_EXPORT extern TF_CoordinationServiceAgent*
TF_GetCoordinationServiceAgent(TF_OpKernelContext* ctx);

// Returns true if the coordination service agent has been initialized.
TF_CAPI_EXPORT extern bool TF_CoordinationServiceIsInitialized(
    TF_CoordinationServiceAgent* agent);

TF_CAPI_EXPORT extern void TF_CoordinationServiceInsertKeyValue(
    const char* key, int64_t key_size, const char* value, int64_t value_size,

  auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
  return reinterpret_cast<TF_CoordinationServiceAgent*>(
      cc_ctx->coordination_service_agent());
}

bool TF_CoordinationServiceIsInitialized(TF_CoordinationServiceAgent* agent) {
  if (agent == nullptr) return false;
  auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
  return cc_agent->IsInitialized();
}